You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kafka.apache.org by ce...@apache.org on 2026/04/29 05:38:47 UTC
(kafka) branch metamorphomover-prototype created (now 08d656e584d)
This is an automated email from the ASF dual-hosted git repository.
C0urante pushed a change to branch metamorphomover-prototype
in repository https://gitbox.apache.org/repos/asf/kafka.git
at 08d656e584d [DRAFT] Add pluggable worker assignor, remove Uber module
This branch includes the following new commits:
new e26cf3c8a07 [DRAFT] Prototype workload-based load balancing, Kafka-based workload tracking, and UReplicator3 connector
new 08d656e584d [DRAFT] Add pluggable worker assignor, remove Uber module
The 2 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(kafka) 01/02: [DRAFT] Prototype workload-based load balancing, Kafka-based workload tracking, and UReplicator3 connector
Posted by ce...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
C0urante pushed a commit to branch metamorphomover-prototype
in repository https://gitbox.apache.org/repos/asf/kafka.git
commit e26cf3c8a07068d7e9a60ba3d36eda2c7beda0b5
Author: chrise <ch...@uber.com>
AuthorDate: Tue Mar 3 22:02:00 2026 +0000
[DRAFT] Prototype workload-based load balancing, Kafka-based workload tracking, and UReplicator3 connector
---
build-prototype.sh | 70 ++
build.gradle | 104 +-
checkstyle/import-control.xml | 3 +
checkstyle/suppressions.xml | 6 +
.../connect/distributed/ClusterAssignment.java | 14 +
.../kafka/connect/distributed/ClusterAssignor.java | 12 +
.../connect/distributed/ClusterConfigState.java | 17 +
.../kafka/connect/distributed/ConnectorTaskId.java | 13 +
.../connect/distributed/ConnectorsAndTasks.java | 82 ++
.../kafka/connect/connector/ConnectorContext.java | 14 +
.../kafka/connect/sink/SinkConnectorTest.java | 10 +
.../kafka/connect/source/SourceConnectorTest.java | 10 +
.../apache/kafka/connect/json/JsonConverter.java | 5 +-
.../kafka/connect/runtime/AbstractHerder.java | 21 +-
.../org/apache/kafka/connect/runtime/Herder.java | 13 +-
.../connect/runtime/HerderConnectorContext.java | 21 +
.../kafka/connect/runtime/WorkerConnector.java | 11 +
.../distributed/ClusterConfigStateImpl.java | 39 +
.../distributed/ConnectProtocolCompatibility.java | 14 +
.../runtime/distributed/DistributedConfig.java | 30 +-
.../runtime/distributed/DistributedHerder.java | 49 +-
.../runtime/distributed/UberAssignmentV1.java | 99 ++
.../connect/runtime/distributed/UberAssignor.java | 321 +++++
.../runtime/distributed/UberConnectProtocol.java | 193 +++
.../runtime/distributed/UberWorkerLoad.java | 63 +
.../runtime/distributed/UberWorkerState.java | 16 +
.../runtime/distributed/WorkerCoordinator.java | 44 +-
.../runtime/distributed/WorkerGroupMember.java | 6 +
.../runtime/health/ConnectClusterStateImpl.java | 38 +-
.../runtime/isolation/PluginScanResult.java | 15 +-
.../connect/runtime/isolation/PluginType.java | 4 +-
.../runtime/isolation/ReflectionScanner.java | 3 +-
.../runtime/isolation/ServiceLoaderScanner.java | 3 +-
.../runtime/rest/resources/ConnectorsResource.java | 10 +-
.../runtime/standalone/StandaloneHerder.java | 18 +-
.../apache/kafka/connect/util/ConnectorTaskId.java | 8 +
.../integration/OffsetsApiIntegrationTest.java | 36 +-
.../integration/TestableSourceConnector.java | 35 +-
.../WorkerCoordinatorIncrementalTest.java | 1 +
.../runtime/distributed/WorkerCoordinatorTest.java | 1 +
.../runtime/distributed/WorkerGroupMemberTest.java | 4 +-
.../isolation/DelegatingClassLoaderTest.java | 1 +
.../connect/runtime/isolation/PluginUtilsTest.java | 4 +
.../connect/runtime/isolation/PluginsTest.java | 1 +
.../rest/resources/ConnectorsResourceTest.java | 14 +-
.../connect/util/clusters/EmbeddedConnect.java | 6 +-
.../aakafkaoffsetmgmt/OffsetSnapshotClient.java | 217 ++++
.../aakafkaoffsetmgmt/OffsetSnapshotReporter.java | 207 ++++
.../connect/aakafkaoffsetmgmt/OffsetTracker.java | 32 +
.../aakafkaoffsetmgmt/PartitionOffsetMap.java | 28 +
.../connect/common/ReplicatedTopicPartition.java | 72 ++
.../distributed/assignment/AbstractAssignor.java | 375 ++++++
.../assignment/UberClusterAssignor.java | 195 +++
.../assignment/UberClusterAssignorConfig.java | 98 ++
.../distributed/assignment/UberWorkerLoad.java | 70 ++
.../assignment/WorkerConnectorAssignor.java | 48 +
.../distributed/assignment/WorkerTaskAssignor.java | 81 ++
.../distributed/assignment/WorkloadProvider.java | 39 +
.../connect/rest/UberConnectRestExtension.java | 60 +
.../rest/entities/AssignedPartitionEntity.java | 29 +
.../rest/entities/CachedWorkloadEntity.java | 15 +
.../rest/entities/PartitionWorkloadEntity.java | 59 +
.../entities/ReplicatedTopicPartitionEntity.java | 22 +
.../rest/entities/TaskAssignmentEntity.java | 26 +
.../connect/rest/entities/TaskWorkloadEntity.java | 10 +
.../rest/entities/TopicPartitionEntity.java | 20 +
.../rest/entities/WorkerWorkloadEntity.java | 13 +
.../connect/rest/entities/WorkloadEntity.java | 18 +
.../rest/filter/ContentTypeHeaderFilter.java | 56 +
.../interceptor/ConnectorConfigInterceptor.java | 235 ++++
.../interceptor/ConnectorOffsetsInterceptor.java | 147 +++
.../rest/resources/AssignmentsResource.java | 85 ++
.../connect/rest/resources/WorkloadsResource.java | 316 +++++
.../ureplicator3/TaskPartitionAssignor.java | 112 ++
.../kafka/connect/ureplicator3/TaskPartitions.java | 25 +
.../connect/ureplicator3/UReplicator3Config.java | 251 ++++
.../ureplicator3/UReplicator3Connector.java | 449 +++++++
.../ureplicator3/UReplicator3ConnectorConfig.java | 385 ++++++
.../connect/ureplicator3/UReplicator3Task.java | 367 ++++++
.../ureplicator3/UReplicator3TaskConfig.java | 42 +
.../coordination/AssignedPartition.java | 27 +
.../coordination/AssignmentListener.java | 11 +
.../ureplicator3/coordination/AssignmentStore.java | 368 ++++++
.../coordination/AssignmentStoreConfig.java | 145 +++
.../coordination/AssignmentStores.java | 39 +
.../data/kafka/connect/utils/BrokerResolver.java | 39 +
.../data/kafka/connect/utils/ExecutorUtils.java | 88 ++
.../kafka/connect/utils/M3ReporterProvider.java | 43 +
.../data/kafka/connect/utils/MetricsUtils.java | 37 +
.../kafka/connect/utils/SerializationUtils.java | 78 ++
.../data/kafka/connect/utils/UberConnectUtils.java | 205 ++++
.../connect/workload/common/CachedWorkload.java | 28 +
.../kafka/connect/workload/common/Workload.java | 70 ++
.../connect/workload/common/WorkloadCache.java | 221 ++++
.../connect/workload/common/WorkloadProvider.java | 39 +
.../workload/kafka/KafkaBasedWorkloadProvider.java | 132 +++
.../kafka/KafkaBasedWorkloadProviderConfig.java | 225 ++++
.../workload/kafka/KafkaBasedWorkloadStore.java | 286 +++++
.../KafkaBasedWorkloadTrackingTransformation.java | 15 +
.../workload/tracking/WorkloadCollector.java | 103 ++
.../connect/workload/tracking/WorkloadStore.java | 12 +
.../connect/workload/tracking/WorkloadTracker.java | 132 +++
.../tracking/WorkloadTrackingProcessor.java | 276 +++++
.../WorkloadTrackingTransformation.java | 62 +
.../WorkloadTrackingTransformationConfig.java | 62 +
....data.kafka.connect.distributed.ClusterAssignor | 1 +
....apache.kafka.connect.rest.ConnectRestExtension | 1 +
...org.apache.kafka.connect.source.SourceConnector | 1 +
.../integration/ReplicationIntegrationTest.java | 1243 ++++++++++++++++++++
.../ureplicator3/TaskPartitionAssignorTest.java | 50 +
.../kafka/connect/utils/BrokerResolverTest.java | 43 +
connect/uber/src/test/resources/log4j2.yaml | 47 +
settings.gradle | 1 +
113 files changed, 9776 insertions(+), 60 deletions(-)
diff --git a/build-prototype.sh b/build-prototype.sh
new file mode 100755
index 00000000000..1ab302b0ea8
--- /dev/null
+++ b/build-prototype.sh
@@ -0,0 +1,70 @@
+#! /usr/bin/env bash
+
+if [[ "$#" -ne 1 ]]; then
+ echo "Usage: $0 <tb_path>"
+ exit 1
+fi
+
+TB_PATH="$1"
+
+upload() {
+ echo "Uploading $1 to TerraBlob"
+ if ! [ -z ${FORCE+x} ]; then
+ tb-cli delete "$TB_PATH"/"$1" || true
+ fi
+ tb-cli put --timeout 90s --multipart "$1" "$TB_PATH"/"$1"
+}
+
+set -e
+
+SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd)
+
+if tb-cli ls "$TB_PATH" > /dev/null 2> /dev/null; then
+ if [ -z ${FORCE+x} ] && [[ $(tb-cli ls "$TB_PATH" | wc -l) != 0 ]]; then
+ echo "Path $TB_PATH already exists and is non-empty; cannot use"
+ exit 1
+ elif ! [ -z ${FORCE+x} ]; then
+ echo "Path $TB_PATH is non-empty but FORCE is set; proceeding"
+ else
+ echo "Path $TB_PATH already exists and is empty; can use safely"
+ fi
+else
+ tb-cli mkdir -p "$TB_PATH"
+fi
+
+
+if [ -z ${NOBUILD+x} ]; then
+
+ if [ -z ${NOCLEAN+x} ]; then
+ "$SCRIPT_DIR"/gradlew clean
+ else
+ echo "Skipping clean"
+ fi
+
+ "$SCRIPT_DIR"/gradlew :connect:uber:{spotlessApply,build} releaseTarGz -x test
+
+else
+ echo "Skipping build"
+fi
+
+if ! [ -z ${NOPUSH+x} ]; then
+ echo "Skipping upload"
+ exit 0
+fi
+
+
+pushd "$SCRIPT_DIR"/core/build/distributions > /dev/null
+upload kafka_2.13-4.3.0-SNAPSHOT.tgz
+popd > /dev/null
+
+pushd "$SCRIPT_DIR"/connect/uber/build > /dev/null
+TEMP_DIR="$(mktemp -d)"
+mkdir "$TEMP_DIR"/connect-uber-plugins
+cp libs/* dependant-libs/* "$TEMP_DIR"/connect-uber-plugins
+pushd "$TEMP_DIR" > /dev/null
+tar czf connect-uber-plugins.tgz connect-uber-plugins
+upload connect-uber-plugins.tgz
+popd > /dev/null
+popd > /dev/null
+
+echo "Finished uploading artifacts to $TB_PATH"
diff --git a/build.gradle b/build.gradle
index bb11d6eb3bc..404e1a0ded6 100644
--- a/build.gradle
+++ b/build.gradle
@@ -41,6 +41,9 @@ plugins {
id 'org.scoverage' version '8.1' apply false
id 'com.gradleup.shadow' version '8.3.9' apply false
id 'com.diffplug.spotless' version "8.0.0"
+
+ id 'application'
+
}
ext {
@@ -168,6 +171,14 @@ allprojects {
repositories {
mavenCentral()
+ maven {
+ url 'http://artifactory.uber.internal:4587/artifactory/libs-release-local/'
+ allowInsecureProtocol = true
+ metadataSources {
+ mavenPom()
+ artifact()
+ }
+ }
}
dependencyUpdates {
@@ -3633,6 +3644,10 @@ project(':connect:runtime') {
implementation libs.classgraph
implementation libs.mavenArtifact
implementation libs.swaggerAnnotations
+ // TODO: We should not be hacking the :connect:runtime module with our own dependencies; instead,
+ // we should be adding them to the classpath
+ // This is also causing embedded integration tests to fail ☹️
+ implementation ('com.uber.kafka:kafka-upki-provider:1.0.1-217732-a4c6b554904e-java-monorepo')
compileOnly libs.bndlib
compileOnly libs.spotbugs
@@ -3662,7 +3677,8 @@ project(':connect:runtime') {
testImplementation libs.junitJupiter
testImplementation libs.mockitoCore
testImplementation libs.mockitoJunitJupiter
- testImplementation libs.httpclient
+ // TODO: Conflicts with Kafka UPKI provider
+// testImplementation libs.httpclient
testImplementation testLog4j2Libs
testCompileOnly libs.bndlib
@@ -3969,6 +3985,92 @@ project(':connect:mirror-client') {
}
}
+project(':connect:uber') {
+ base {
+ archivesName = "connect-uber"
+ }
+
+ task genUReplicator3Config(type: JavaExec) {
+ classpath sourceSets.main.runtimeClasspath
+ classpath sourceSets.main.compileClasspath
+ mainClass = 'com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig'
+ if( !generatedDocsDir.exists() ) { generatedDocsDir.mkdirs() }
+ standardOutput = new File(generatedDocsDir, "ureplicator3_config.html").newOutputStream()
+ }
+
+ dependencies {
+ implementation 'com.google.guava:guava:23.0'
+ implementation 'com.uber.m3:tally-core:0.11.1'
+ implementation 'com.uber.m3:tally-m3:0.11.1'
+ implementation('com.uber.data:streaming-common:1.13.1-230620-a54bfb9c5127-java-monorepo') {
+ exclude group: 'jakarta.annotation', module: 'jakarta.annotation-api'
+ exclude group: 'jakarta.ws.rs', module: 'jakarta.ws.rs-api'
+ exclude group: 'jakarta.servlet', module: 'jakarta.servlet-api'
+ exclude group: 'jakarta.validation', module: 'jakarta.validation-api'
+ }
+ implementation 'com.alibaba:fastjson:1.1.41'
+ implementation 'com.google.guava:guava:23.0'
+
+ implementation('io.opentracing.contrib:opentracing-apache-httpclient:0.3.3') {
+ /*
+ opentracing client is using 0.33.0 of opentracing-api and conflict with 0.32 used in javamonorepo, the only way is to downgrade it.
+ Monorepo won't ugprade opentracing to 0.33 due to https://t3.uberinternal.com/browse/JAG-2163
+ More details: https://github.com/opentracing/opentracing-java/blob/master/CHANGELOG.md#v0330-2019-05-06
+ */
+ exclude group: 'io.opentracing', module: 'opentracing-api'
+ exclude group: 'io.opentracing', module: 'opentracing-util'
+ }
+ implementation libs.slf4jApi
+ implementation libs.swaggerAnnotations
+
+ compileOnly libs.jakartaRsApi
+ compileOnly libs.jakartaServletApi
+ compileOnly libs.jaxAnnotationApi
+ compileOnly libs.jerseyContainerServlet
+ compileOnly project(':connect:api')
+ compileOnly project(':clients')
+ compileOnly project(':connect:runtime')
+ compileOnly project(':connect:mirror')
+
+ testImplementation testLog4j2Libs
+ testImplementation libs.jakartaRsApi
+ testImplementation libs.junitJupiter
+ testImplementation libs.bndlib
+ testImplementation libs.mockitoCore
+ testImplementation project(':clients').sourceSets.test.output
+ testImplementation project(':connect:runtime')
+ testImplementation project(':connect:runtime').sourceSets.test.output
+ testImplementation project(':connect:mirror')
+ testImplementation project(':connect:mirror-client')
+ testImplementation project(':core')
+ testImplementation project(':test-common:test-common-runtime')
+ testImplementation project(':raft')
+ testImplementation project(':server')
+ testImplementation project(':server-common')
+
+ testRuntimeOnly project(':connect:runtime')
+ testRuntimeOnly libs.bcpkix
+ testRuntimeOnly runtimeTestLibs
+ }
+
+ javadoc {
+ enabled = false
+ }
+
+ tasks.create(name: "copyDependantLibs", type: Copy) {
+ from (configurations.runtimeClasspath) {
+ exclude('kafka-clients*')
+ exclude('connect-*')
+ }
+ into "${layout.buildDirectory.get().asFile.path}/dependant-libs"
+ duplicatesStrategy 'exclude'
+ }
+
+ jar {
+ dependsOn copyDependantLibs
+ }
+}
+
project(':connect:test-plugins') {
base {
archivesName = "connect-test-plugins"
diff --git a/checkstyle/import-control.xml b/checkstyle/import-control.xml
index 86659a1856d..9e84333d18e 100644
--- a/checkstyle/import-control.xml
+++ b/checkstyle/import-control.xml
@@ -482,6 +482,9 @@
</subpackage>
<subpackage name="connect">
+
+ <allow pkg="com.uber.data.kafka.connect" />
+
<allow pkg="org.apache.kafka.common" />
<allow pkg="org.apache.kafka.connect.data" />
<allow pkg="org.apache.kafka.connect.errors" />
diff --git a/checkstyle/suppressions.xml b/checkstyle/suppressions.xml
index 359334f2cc2..b16dfbf46ff 100644
--- a/checkstyle/suppressions.xml
+++ b/checkstyle/suppressions.xml
@@ -21,6 +21,9 @@
-->
<suppressions>
+ <!-- Disable Checkstyle for our internal Connect module (we may selectively enable some checks in the future) -->
+ <suppress checks="." files="com[\\/]uber[\\/]data[\\/]kafka[\\/]connect[\\/]"/>
+
<!-- Note that [/\\] must be used as the path separator for cross-platform support -->
<!-- Generator -->
@@ -379,4 +382,7 @@
<!-- Add the new suppression rule for JaasTestUtils.java -->
<suppress checks="ImportControl" files="(JaasTestUtils).java" />
+ <!-- Uber -->
+ <suppress checks="Header" files="(Uber.*|WorkloadProvider).java" />
+
</suppressions>
diff --git a/connect/api/src/main/java/com/uber/data/kafka/connect/distributed/ClusterAssignment.java b/connect/api/src/main/java/com/uber/data/kafka/connect/distributed/ClusterAssignment.java
new file mode 100644
index 00000000000..3c1675b99e7
--- /dev/null
+++ b/connect/api/src/main/java/com/uber/data/kafka/connect/distributed/ClusterAssignment.java
@@ -0,0 +1,14 @@
+package com.uber.data.kafka.connect.distributed;
+
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.Map;
+
+public record ClusterAssignment(Map<String, ConnectorsAndTasks> workerAssignments, int scheduledRebalanceDelayMs) {
+
+ public ClusterAssignment(Map<String, ConnectorsAndTasks> workerAssignments, int scheduledRebalanceDelayMs) {
+ this.workerAssignments = Collections.unmodifiableMap(new LinkedHashMap<>(workerAssignments));
+ this.scheduledRebalanceDelayMs = scheduledRebalanceDelayMs;
+ }
+
+}
diff --git a/connect/api/src/main/java/com/uber/data/kafka/connect/distributed/ClusterAssignor.java b/connect/api/src/main/java/com/uber/data/kafka/connect/distributed/ClusterAssignor.java
new file mode 100644
index 00000000000..a7446cce7d2
--- /dev/null
+++ b/connect/api/src/main/java/com/uber/data/kafka/connect/distributed/ClusterAssignor.java
@@ -0,0 +1,12 @@
+package com.uber.data.kafka.connect.distributed;
+
+import org.apache.kafka.common.Configurable;
+import org.apache.kafka.connect.components.Versioned;
+
+import java.util.Map;
+
+public interface ClusterAssignor extends Versioned, Configurable, AutoCloseable {
+
+ ClusterAssignment assign(ClusterConfigState clusterConfigState, Map<String, ConnectorsAndTasks> currentAssignments);
+
+}
diff --git a/connect/api/src/main/java/com/uber/data/kafka/connect/distributed/ClusterConfigState.java b/connect/api/src/main/java/com/uber/data/kafka/connect/distributed/ClusterConfigState.java
new file mode 100644
index 00000000000..49a2456a084
--- /dev/null
+++ b/connect/api/src/main/java/com/uber/data/kafka/connect/distributed/ClusterConfigState.java
@@ -0,0 +1,17 @@
+package com.uber.data.kafka.connect.distributed;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+public interface ClusterConfigState {
+
+ Set<String> connectors();
+
+ List<ConnectorTaskId> tasks(String connector);
+
+ Map<String, String> connectorConfig(String connector);
+
+ Map<String, String> taskConfig(ConnectorTaskId task);
+
+}
diff --git a/connect/api/src/main/java/com/uber/data/kafka/connect/distributed/ConnectorTaskId.java b/connect/api/src/main/java/com/uber/data/kafka/connect/distributed/ConnectorTaskId.java
new file mode 100644
index 00000000000..4e9fa4f052a
--- /dev/null
+++ b/connect/api/src/main/java/com/uber/data/kafka/connect/distributed/ConnectorTaskId.java
@@ -0,0 +1,13 @@
+package com.uber.data.kafka.connect.distributed;
+
+public record ConnectorTaskId(String connector, int task) implements Comparable<ConnectorTaskId> {
+
+ @Override
+ public int compareTo(ConnectorTaskId o) {
+ int connectorCmp = connector.compareTo(o.connector);
+ if (connectorCmp != 0)
+ return connectorCmp;
+ return Integer.compare(task, o.task);
+ }
+
+}
diff --git a/connect/api/src/main/java/com/uber/data/kafka/connect/distributed/ConnectorsAndTasks.java b/connect/api/src/main/java/com/uber/data/kafka/connect/distributed/ConnectorsAndTasks.java
new file mode 100644
index 00000000000..7f54ce0b5d0
--- /dev/null
+++ b/connect/api/src/main/java/com/uber/data/kafka/connect/distributed/ConnectorsAndTasks.java
@@ -0,0 +1,82 @@
+package com.uber.data.kafka.connect.distributed;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.LinkedHashSet;
+import java.util.Set;
+
+public record ConnectorsAndTasks(Set<String> connectors, Set<ConnectorTaskId> tasks) {
+
+ public ConnectorsAndTasks(Set<String> connectors, Set<ConnectorTaskId> tasks) {
+ this.connectors = Collections.unmodifiableSet(new LinkedHashSet<>(connectors));
+ this.tasks = Collections.unmodifiableSet(new LinkedHashSet<>(tasks));
+ }
+
+ public Builder toBuilder() {
+ return builder()
+ .add(this);
+ }
+
+ public static ConnectorsAndTasks of(Collection<String> connectors, Collection<ConnectorTaskId> tasks) {
+ return new ConnectorsAndTasks(Set.copyOf(connectors), Set.copyOf(tasks));
+ }
+
+ public static Builder builder() {
+ return new Builder();
+ }
+
+ public static Builder builder(Collection<String> connectors, Collection<ConnectorTaskId> tasks) {
+ return builder()
+ .addConnectors(connectors)
+ .addTasks(tasks);
+ }
+
+ public static class Builder {
+
+ private final Set<String> connectors;
+ private final Set<ConnectorTaskId> tasks;
+
+ public Builder() {
+ this.connectors = new LinkedHashSet<>();
+ this.tasks = new LinkedHashSet<>();
+ }
+
+ public Builder addConnectors(Collection<String> connectors) {
+ this.connectors.addAll(connectors);
+ return this;
+ }
+
+ public Builder addTasks(Collection<ConnectorTaskId> tasks) {
+ this.tasks.addAll(tasks);
+ return this;
+ }
+
+ public Builder add(ConnectorsAndTasks connectorsAndTasks) {
+ return this
+ .addConnectors(connectorsAndTasks.connectors())
+ .addTasks(connectorsAndTasks.tasks());
+ }
+
+ public Builder removeConnectors(Collection<String> connectors) {
+ this.connectors.removeAll(connectors);
+ return this;
+ }
+
+ public Builder removeTasks(Collection<ConnectorTaskId> tasks) {
+ this.tasks.removeAll(tasks);
+ return this;
+ }
+
+ public Builder remove(ConnectorsAndTasks connectorsAndTasks) {
+ return this
+ .removeConnectors(connectorsAndTasks.connectors())
+ .removeTasks(connectorsAndTasks.tasks());
+ }
+
+ public ConnectorsAndTasks build() {
+ return new ConnectorsAndTasks(connectors, tasks);
+ }
+
+ }
+
+}
diff --git a/connect/api/src/main/java/org/apache/kafka/connect/connector/ConnectorContext.java b/connect/api/src/main/java/org/apache/kafka/connect/connector/ConnectorContext.java
index 09d1a71eb2a..823f8e34f8f 100644
--- a/connect/api/src/main/java/org/apache/kafka/connect/connector/ConnectorContext.java
+++ b/connect/api/src/main/java/org/apache/kafka/connect/connector/ConnectorContext.java
@@ -18,6 +18,9 @@ package org.apache.kafka.connect.connector;
import org.apache.kafka.common.metrics.PluginMetrics;
+import java.util.List;
+import java.util.Map;
+
/**
* ConnectorContext allows {@link Connector}s to proactively interact with the Kafka Connect runtime.
*/
@@ -57,4 +60,15 @@ public interface ConnectorContext {
* @since 4.1
*/
PluginMetrics pluginMetrics();
+
+ /**
+ * @return the current task configs for the connector; may be empty, but never null
+ */
+ List<Map<String, String>> uberTaskConfigs();
+
+ /**
+ * @return the number of workers in the Kafka Connect cluster, or -1
+ * if the number of workers is not known, or 1 if this is a standalone worker
+ */
+ int uberClusterSize();
}
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
index 176ebdf2caf..9a5c9275cf5 100644
--- 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
@@ -60,6 +60,16 @@ public class SinkConnectorTest extends ConnectorTest {
// Unexpected in these tests
throw new UnsupportedOperationException();
}
+
+ @Override
+ public List<Map<String, String>> uberTaskConfigs() {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public int uberClusterSize() {
+ throw new UnsupportedOperationException();
+ }
}
protected static class TestSinkConnector extends SinkConnector implements ConnectorTest.AssertableConnector {
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
index 913c5e2019d..62358977f94 100644
--- 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
@@ -66,6 +66,16 @@ public class SourceConnectorTest extends ConnectorTest {
public OffsetStorageReader offsetStorageReader() {
return null;
}
+
+ @Override
+ public List<Map<String, String>> uberTaskConfigs() {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public int uberClusterSize() {
+ throw new UnsupportedOperationException();
+ }
}
private static class TestSourceConnector extends SourceConnector implements AssertableConnector {
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 dac2ce56741..8e7ae7fb368 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
@@ -53,6 +53,7 @@ import java.util.ArrayList;
import java.util.Collection;
import java.util.EnumMap;
import java.util.HashMap;
+import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Set;
@@ -98,7 +99,9 @@ public class JsonConverter implements Converter, HeaderConverter, Versioned {
// If the map uses strings for keys, it should be encoded in the natural JSON format. If it uses other
// primitive types or a complex type as a key, it will be encoded as a list of pairs. If we don't have a
// schema, we default to encoding in a Map.
- Map<Object, Object> result = new HashMap<>();
+ // Use a LinkedHashMap to preserve ordering of key/value pairs
+ // TODO: Upstream this change
+ Map<Object, Object> result = new LinkedHashMap<>();
if (schema == null || keySchema.type() == Schema.Type.STRING) {
if (!value.isObject())
throw new DataException("Maps with string fields should be encoded as JSON objects, but found " + value.getNodeType());
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 d5d6edf703e..bc25975d210 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
@@ -1241,26 +1241,31 @@ public abstract class AbstractHerder implements Herder, TaskStatus.Listener, Con
}
@Override
- public void alterConnectorOffsets(String connName, Map<Map<String, ?>, Map<String, ?>> offsets, Callback<Message> callback) {
+ public void uberAlterConnectorOffsets(String connName, Map<Map<String, ?>, Map<String, ?>> offsets, boolean uberForce, Callback<Message> callback) {
if (offsets == null || offsets.isEmpty()) {
callback.onCompletion(new ConnectException("The offsets to be altered may not be null or empty"), null);
return;
}
- modifyConnectorOffsets(connName, offsets, callback);
+ uberModifyConnectorOffsets(connName, offsets, uberForce, callback);
}
@Override
public void resetConnectorOffsets(String connName, Callback<Message> callback) {
- modifyConnectorOffsets(connName, null, callback);
+ uberModifyConnectorOffsets(connName, null, false, callback);
+ }
+
+ public void modifyConnectorOffsets(String connName, Map<Map<String, ?>, Map<String, ?>> offsets, Callback<Message> cb) {
+ uberModifyConnectorOffsets(connName, offsets, false, cb);
}
/**
* Service external requests to alter or reset connector offsets.
* @param connName the name of the connector whose offsets are to be modified
* @param offsets the offsets to be written; this should be {@code null} for offsets reset requests
+ * @param uberForce force the offset modification to take place even if the connector is not stopped
* @param cb callback to invoke upon completion
*/
- protected abstract void modifyConnectorOffsets(String connName, Map<Map<String, ?>, Map<String, ?>> offsets, Callback<Message> cb);
+ protected abstract void uberModifyConnectorOffsets(String connName, Map<Map<String, ?>, Map<String, ?>> offsets, boolean uberForce, Callback<Message> cb);
@Override
public LoggerLevel loggerLevel(String logger) {
@@ -1283,4 +1288,12 @@ public abstract class AbstractHerder implements Herder, TaskStatus.Listener, Con
return loggers.setLevel(namespace, normalizedLevel);
}
+
+ public ClusterConfigState uberConfigSnapshot() {
+ return configBackingStore.snapshot();
+ }
+
+ public int uberClusterSize() {
+ return 1;
+ }
}
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 a1a1505c983..ecfa2fae194 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
@@ -336,9 +336,20 @@ public interface Herder {
* Alter a connector's offsets.
* @param connName the name of the connector whose offsets are to be altered
* @param offsets a mapping from partitions to offsets that need to be written
+ * @param uberForce force the offset modification to take place even if the connector is not stopped
* @param cb callback to invoke upon completion
*/
- void alterConnectorOffsets(String connName, Map<Map<String, ?>, Map<String, ?>> offsets, Callback<Message> cb);
+ void uberAlterConnectorOffsets(String connName, Map<Map<String, ?>, Map<String, ?>> offsets, boolean uberForce, Callback<Message> cb);
+
+ /**
+ * Alter a connector's offsets.
+ * @param connName the name of the connector whose offsets are to be altered
+ * @param offsets a mapping from partitions to offsets that need to be written
+ * @param cb callback to invoke upon completion
+ */
+ default void alterConnectorOffsets(String connName, Map<Map<String, ?>, Map<String, ?>> offsets, Callback<Message> cb) {
+ uberAlterConnectorOffsets(connName, offsets, false, cb);
+ }
/**
* Reset a connector's offsets.
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 8568e91f400..0058cd6fce8 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
@@ -20,10 +20,16 @@ import org.apache.kafka.common.metrics.PluginMetrics;
import org.apache.kafka.common.metrics.internals.PluginMetricsImpl;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.storage.ClusterConfigState;
+import org.apache.kafka.connect.util.ConnectorTaskId;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.IntStream;
+
/**
* ConnectorContext for use with a Herder
*/
@@ -78,4 +84,19 @@ public class HerderConnectorContext implements CloseableConnectorContext {
Utils.closeQuietly(pluginMetrics, "Plugin metrics for " + connectorName);
closed = true;
}
+
+ @Override
+ public List<Map<String, String>> uberTaskConfigs() {
+ ClusterConfigState snapshot = herder.uberConfigSnapshot();
+ int numTasks = snapshot.taskCount(connectorName);
+ return IntStream.range(0, numTasks)
+ .mapToObj(i -> new ConnectorTaskId(connectorName, i))
+ .map(snapshot::taskConfig)
+ .toList();
+ }
+
+ @Override
+ public int uberClusterSize() {
+ return herder.uberClusterSize();
+ }
}
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 3faf70f898c..667da7c6199 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
@@ -34,6 +34,7 @@ import org.apache.kafka.connect.util.LoggingContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Objects;
@@ -599,6 +600,16 @@ public class WorkerConnector implements Runnable {
public PluginMetrics pluginMetrics() {
return WorkerConnector.this.ctx.pluginMetrics();
}
+
+ @Override
+ public List<Map<String, String>> uberTaskConfigs() {
+ return WorkerConnector.this.ctx.uberTaskConfigs();
+ }
+
+ @Override
+ public int uberClusterSize() {
+ return WorkerConnector.this.ctx.uberClusterSize();
+ }
}
private class WorkerSinkConnectorContext extends WorkerConnectorContext implements SinkConnectorContext {
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ClusterConfigStateImpl.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ClusterConfigStateImpl.java
new file mode 100644
index 00000000000..def0f344d56
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ClusterConfigStateImpl.java
@@ -0,0 +1,39 @@
+package org.apache.kafka.connect.runtime.distributed;
+
+import com.uber.data.kafka.connect.distributed.ClusterConfigState;
+import com.uber.data.kafka.connect.distributed.ConnectorTaskId;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+public class ClusterConfigStateImpl implements ClusterConfigState {
+
+ private final org.apache.kafka.connect.storage.ClusterConfigState configSnapshot;
+
+ public ClusterConfigStateImpl(org.apache.kafka.connect.storage.ClusterConfigState configSnapshot) {
+ this.configSnapshot = configSnapshot;
+ }
+
+ @Override
+ public Set<String> connectors() {
+ return configSnapshot.connectors();
+ }
+
+ @Override
+ public List<ConnectorTaskId> tasks(String connector) {
+ return configSnapshot.tasks(connector).stream()
+ .map(org.apache.kafka.connect.util.ConnectorTaskId::uberToPublicApi)
+ .toList();
+ }
+
+ @Override
+ public Map<String, String> connectorConfig(String connector) {
+ return configSnapshot.connectorConfig(connector);
+ }
+
+ @Override
+ public Map<String, String> taskConfig(ConnectorTaskId task) {
+ return configSnapshot.taskConfig(new org.apache.kafka.connect.util.ConnectorTaskId(task.connector(), task.task()));
+ }
+}
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
index 84ec2f0a392..4c0214109e0 100644
--- 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
@@ -22,6 +22,7 @@ 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;
+import static org.apache.kafka.connect.runtime.distributed.UberConnectProtocol.UBER_CONNECT_PROTOCOL_V1;
/**
* An enumeration of the modes available to the worker to signal which Connect protocols are
@@ -73,6 +74,18 @@ public enum ConnectProtocolCompatibility {
public short protocolVersion() {
return CONNECT_PROTOCOL_V2;
}
+ },
+
+ UBER_V1 {
+ @Override
+ public String protocol() {
+ return "uber-v1";
+ }
+
+ @Override
+ public short protocolVersion() {
+ return UBER_CONNECT_PROTOCOL_V1;
+ }
};
/**
@@ -103,6 +116,7 @@ public enum ConnectProtocolCompatibility {
case CONNECT_PROTOCOL_V0 -> EAGER;
case CONNECT_PROTOCOL_V1 -> COMPATIBLE;
case CONNECT_PROTOCOL_V2 -> SESSIONED;
+ case UBER_CONNECT_PROTOCOL_V1 -> UBER_V1;
default -> throw new IllegalArgumentException("Unknown Connect protocol version: " + protocolVersion);
};
}
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 ccf33926bf9..391597f664c 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
@@ -16,6 +16,7 @@
*/
package org.apache.kafka.connect.runtime.distributed;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.clients.MetadataRecoveryStrategy;
import org.apache.kafka.clients.producer.ProducerConfig;
@@ -182,7 +183,7 @@ public final class DistributedConfig extends WorkerConfig {
*/
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();
+ public static final String CONNECT_PROTOCOL_DEFAULT = ConnectProtocolCompatibility.UBER_V1.toString();
/**
* <code>scheduled.rebalance.max.delay.ms</code>
@@ -255,6 +256,13 @@ public final class DistributedConfig extends WorkerConfig {
+ "<a href=\"https://kafka.apache.org/documentation.html#connect_exactlyoncesource\">exactly-once source support documentation</a>.";
public static final String EXACTLY_ONCE_SOURCE_SUPPORT_DEFAULT = ExactlyOnceSourceSupport.DISABLED.toString();
+ public static final String UBER_CLUSTER_ASSIGNOR_CONFIG = "uber.cluster.assignor";
+ public static final String UBER_CLUSTER_ASSIGNOR_DOC = "Implementation of the "
+ + ClusterAssignor.class.getName() + " interface that will be used to assign connectors and tasks "
+ + "to workers during rebalances, possibly taking into account the workloads of each connector and task. "
+ + "If left unset, then the default Kafka Connect rebalancing algorithm will be used instead, "
+ + "which does not take workload into account and weighs all tasks equally";
+
private static Object defaultKeyGenerationAlgorithm(Crypto crypto) {
try {
validateKeyAlgorithm(crypto, INTER_WORKER_KEY_GENERATION_ALGORITHM_CONFIG, INTER_WORKER_KEY_GENERATION_ALGORITHM_DEFAULT);
@@ -309,7 +317,7 @@ public final class DistributedConfig extends WorkerConfig {
@SuppressWarnings("unchecked")
private static ConfigDef config(Crypto crypto) {
- return baseConfigDef()
+ return addUberExtensions(baseConfigDef())
.define(GROUP_ID_CONFIG,
ConfigDef.Type.STRING,
ConfigDef.Importance.HIGH,
@@ -587,6 +595,10 @@ public final class DistributedConfig extends WorkerConfig {
return getString(GROUP_ID_CONFIG);
}
+ public String uberClusterAssignor() {
+ return getString(UBER_CLUSTER_ASSIGNOR_CONFIG);
+ }
+
@Override
protected Map<String, Object> postProcessParsedConfig(final Map<String, Object> parsedValues) {
CommonClientConfigs.warnDisablingExponentialBackoff(this);
@@ -603,6 +615,11 @@ public final class DistributedConfig extends WorkerConfig {
this.crypto = crypto;
exactlyOnceSourceSupport = ExactlyOnceSourceSupport.fromProperty(getString(EXACTLY_ONCE_SOURCE_SUPPORT_CONFIG));
validateInterWorkerKeyConfigs();
+ if (ConnectProtocolCompatibility.compatibility(getString(CONNECT_PROTOCOL_CONFIG)) == ConnectProtocolCompatibility.UBER_V1) {
+ if (uberClusterAssignor() == null) {
+ throw new IllegalStateException("Cannot use protocol " + ConnectProtocolCompatibility.UBER_V1.protocol() + " without also specifying a cluster assignor ");
+ }
+ }
}
public static void main(String[] args) {
@@ -722,4 +739,13 @@ public final class DistributedConfig extends WorkerConfig {
return result;
}
+ private static ConfigDef addUberExtensions(ConfigDef baseConfigDef) {
+ return baseConfigDef
+ .define(UBER_CLUSTER_ASSIGNOR_CONFIG,
+ ConfigDef.Type.STRING,
+ null, // Null instead of ConfigDef.NO_DEFAULT_VALUE to signify that the property is optional
+ ConfigDef.Importance.HIGH,
+ UBER_CLUSTER_ASSIGNOR_DOC);
+ }
+
}
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 6c4bed311d3..3b85f561e16 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,6 +16,7 @@
*/
package org.apache.kafka.connect.runtime.distributed;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.common.KafkaFuture;
import org.apache.kafka.common.config.ConfigDef;
@@ -313,10 +314,18 @@ public class DistributedHerder extends AbstractHerder implements Runnable {
LogContext logContext = new LogContext("[Worker clientId=" + clientId + ", groupId=" + this.workerGroupId + "] ");
log = logContext.logger(DistributedHerder.class);
+ String clusterAssignorClass = config.uberClusterAssignor();
+ final ClusterAssignor clusterAssignor;
+ if (clusterAssignorClass != null) {
+ clusterAssignor = plugins().newPlugin(clusterAssignorClass, config, ClusterAssignor.class);
+ } else {
+ clusterAssignor = null;
+ }
+
this.member = member != null
? member
: new WorkerGroupMember(config, restUrl, this.configBackingStore,
- new RebalanceListener(time), time, clientId, logContext);
+ new RebalanceListener(time), clusterAssignor, time, clientId, logContext);
this.herderExecutor = new ThreadPoolExecutor(1, 1, 0L,
TimeUnit.MILLISECONDS,
@@ -1634,16 +1643,20 @@ public class DistributedHerder extends AbstractHerder implements Runnable {
}
@Override
- protected void modifyConnectorOffsets(String connName, Map<Map<String, ?>, Map<String, ?>> offsets, Callback<Message> callback) {
+ protected void uberModifyConnectorOffsets(String connName, Map<Map<String, ?>, Map<String, ?>> offsets, boolean uberForce, Callback<Message> callback) {
log.trace("Submitting {} offsets request for connector '{}'", offsets == null ? "reset" : "alter", connName);
addRequest(() -> {
- if (!modifyConnectorOffsetsChecks(connName, callback)) {
+ if (!modifyConnectorOffsetsChecks(connName, uberForce, callback)) {
return null;
}
// At this point, we should be the leader (the call to modifyConnectorOffsetsChecks makes sure of that) and can safely run
// a zombie fencing request
if (isSourceConnector(connName) && config.exactlyOnceSourceEnabled()) {
+ if (uberForce) {
+ callback.onCompletion(new BadRequestException("Offset alterations cannot be forced for exactly-once source connectors"), null);
+ return null;
+ }
log.debug("Performing a round of zombie fencing before modifying offsets for source connector {} with exactly-once support enabled.", connName);
doFenceZombieSourceTasks(connName, (error, ignored) -> {
if (error != null) {
@@ -1655,7 +1668,7 @@ public class DistributedHerder extends AbstractHerder implements Runnable {
// zombie fencing is done asynchronously and the conditions could have changed since the previous check
addRequest(() -> {
try (TickThreadStage stage = new TickThreadStage("modifying offsets for connector " + connName)) {
- if (modifyConnectorOffsetsChecks(connName, callback)) {
+ if (modifyConnectorOffsetsChecks(connName, false, callback)) {
worker.modifyConnectorOffsets(connName, configState.connectorConfig(connName), offsets, callback);
}
}
@@ -1676,10 +1689,11 @@ public class DistributedHerder extends AbstractHerder implements Runnable {
* This method performs a few checks for external requests to modify (alter or reset) connector offsets and
* completes the callback exceptionally if any check fails.
* @param connName the name of the connector whose offsets are to be modified
+ * @param uberForce force the offset modification to take place even if the connector is not stopped
* @param callback callback to invoke upon completion
* @return true if all the checks passed, false otherwise
*/
- private boolean modifyConnectorOffsetsChecks(String connName, Callback<Message> callback) {
+ private boolean modifyConnectorOffsetsChecks(String connName, boolean uberForce, Callback<Message> callback) {
if (checkRebalanceNeeded(callback)) {
return false;
}
@@ -1703,9 +1717,14 @@ public class DistributedHerder extends AbstractHerder implements Runnable {
// Zombie tasks are handled by a round of zombie fencing for exactly once source connectors. Zombie sink tasks are handled
// naturally because requests to alter consumer group offsets / delete consumer groups will fail if there are still active members
// in the group.
- if (configState.targetState(connName) != TargetState.STOPPED || configState.taskCount(connName) != 0) {
- callback.onCompletion(new BadRequestException("Connectors must be in the STOPPED state before their offsets can be modified. This can be done " +
- "for the specified connector by issuing a 'PUT' request to the '/connectors/" + connName + "/stop' endpoint"), null);
+ if (!uberForce && (configState.targetState(connName) != TargetState.STOPPED || configState.taskCount(connName) != 0)) {
+ String message = "Connectors must be in the STOPPED state before their offsets can be modified. This can be done "
+ + "for the specified connector by issuing a 'PUT' request to the '/connectors/" + connName + "/stop' endpoint. "
+ + "Alternatively, for offset alterations (but not resets), you may force the request to take place via the ?uber-force=true "
+ + "URL query parameter; however, in order for the change to take effect properly, the connector must not be "
+ + "actively processing data for the affected offsets (e.g., if it's replicating from a Kafka topic t1, you must not "
+ + "alter offsets for the t1 topic, but you may alter offsets for topics t2, t3, etc.)";
+ callback.onCompletion(new BadRequestException(message), null);
return false;
}
return true;
@@ -2815,6 +2834,20 @@ public class DistributedHerder extends AbstractHerder implements Runnable {
return result;
}
+ @Override
+ public int uberClusterSize() {
+ if (assignment == null) {
+ log.warn("Worker does not currently have an assignment; cannot provide cluster size to connector");
+ return -1;
+ }
+ if (assignment instanceof UberAssignmentV1 uberAssignment) {
+ return uberAssignment.clusterSize();
+ } else {
+ log.warn("Cluster is not currently running with Uber rebalance protocol; cannot provide cluster size to connector");
+ return -1;
+ }
+ }
+
private synchronized void recordTickThreadStage(String description) {
assert description != null;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberAssignmentV1.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberAssignmentV1.java
new file mode 100644
index 00000000000..eefa8f5b01c
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberAssignmentV1.java
@@ -0,0 +1,99 @@
+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.Collection;
+
+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.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.IncrementalCooperativeConnectProtocol.REVOKED_KEY_NAME;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.SCHEDULED_DELAY_KEY_NAME;
+import static org.apache.kafka.connect.runtime.distributed.UberConnectProtocol.UBER_ASSIGNMENT_V1;
+import static org.apache.kafka.connect.runtime.distributed.UberConnectProtocol.UBER_CLUSTER_SIZE_KEY_NAME;
+
+public class UberAssignmentV1 extends ExtendedAssignment {
+
+ private static final UberAssignmentV1 EMPTY = fromExtended(ExtendedAssignment.empty(), 0);
+
+ private final int clusterSize;
+
+ public UberAssignmentV1(
+ short version,
+ short error,
+ String leader,
+ String leaderUrl,
+ long configOffset,
+ Collection<String> connectorIds,
+ Collection<ConnectorTaskId> taskIds,
+ Collection<String> revokedConnectorIds,
+ Collection<ConnectorTaskId> revokedTaskIds,
+ int delay,
+ int clusterSize) {
+ super(version, error, leader, leaderUrl, configOffset, connectorIds, taskIds, revokedConnectorIds, revokedTaskIds, delay);
+ this.clusterSize = clusterSize;
+ }
+
+ public int clusterSize() {
+ return clusterSize;
+ }
+
+ /**
+ * Return the {@code Struct} that corresponds to this assignment.
+ *
+ * @return the assignment struct
+ */
+ public Struct toUberStruct() {
+ Struct superStruct = super.toStruct();
+ return new Struct(UBER_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, superStruct.get(ASSIGNMENT_KEY_NAME))
+ .set(REVOKED_KEY_NAME, superStruct.get(REVOKED_KEY_NAME))
+ .set(SCHEDULED_DELAY_KEY_NAME, superStruct.get(SCHEDULED_DELAY_KEY_NAME))
+ .set(UBER_CLUSTER_SIZE_KEY_NAME, clusterSize);
+ }
+
+ // We don't override toStruct; it's only used when serializing assignments to send to the leader when
+ // joining a group during rebalance, and the cluster size field isn't used there
+
+ public static UberAssignmentV1 fromExtended(ExtendedAssignment baseAssignment, int clusterSize) {
+ return new UberAssignmentV1(
+ baseAssignment.version(),
+ baseAssignment.error(),
+ baseAssignment.leader(),
+ baseAssignment.leaderUrl(),
+ baseAssignment.offset(),
+ baseAssignment.connectors(),
+ baseAssignment.tasks(),
+ baseAssignment.revokedConnectors(),
+ baseAssignment.revokedTasks(),
+ baseAssignment.delay(),
+ clusterSize
+ );
+ }
+
+ /**
+ * Given a {@code Struct} that encodes an assignment return the assignment object.
+ *
+ * @param struct a struct representing an assignment
+ * @return the assignment
+ */
+ public static UberAssignmentV1 fromStruct(short version, Struct struct) {
+ ExtendedAssignment extendedAssignment = ExtendedAssignment.fromStruct(version, struct);
+ Integer clusterSize = struct.getInt(UBER_CLUSTER_SIZE_KEY_NAME);
+ if (clusterSize == null) {
+ clusterSize = 0;
+ }
+ return fromExtended(extendedAssignment, clusterSize);
+ }
+
+ public static UberAssignmentV1 empty() {
+ return EMPTY;
+ }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberAssignor.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberAssignor.java
new file mode 100644
index 00000000000..4a466575ad3
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberAssignor.java
@@ -0,0 +1,321 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY 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 com.uber.data.kafka.connect.distributed.ClusterAssignor;
+import com.uber.data.kafka.connect.distributed.ConnectorTaskId;
+import com.uber.data.kafka.connect.distributed.ConnectorsAndTasks;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.connect.storage.ClusterConfigState;
+
+import org.slf4j.Logger;
+
+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 java.util.Set;
+import java.util.TreeSet;
+import java.util.stream.Collectors;
+
+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;
+import static org.apache.kafka.connect.util.ConnectUtils.combineCollections;
+import static org.apache.kafka.connect.util.ConnectUtils.transformValues;
+
+/**
+ * An assignor that provides a pluggable interface layer for users to customize the logic for
+ * assigning connectors and tasks to workers, possibly taking into account workload information
+ */
+public class UberAssignor implements ConnectAssignor {
+
+ private final Logger log;
+ private final ClusterAssignor assignor;
+
+ public UberAssignor(
+ LogContext logContext,
+ ClusterAssignor assignor
+ ) {
+ this.log = logContext.logger(UberAssignor.class);
+ this.assignor = assignor;
+ }
+
+ @Override
+ public Map<String, ByteBuffer> performAssignment(String leaderId, ConnectProtocolCompatibility protocol,
+ List<JoinGroupResponseMember> allMemberMetadata,
+ WorkerCoordinator coordinator) {
+ log.debug("Performing task assignment");
+
+ Map<String, ExtendedWorkerState> 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 = protocol.protocolVersion();
+
+ ClusterConfigState leaderSnapshot = ensureLeaderConfig(maxOffset, coordinator);
+ if (leaderSnapshot == null) {
+ Map<String, ExtendedAssignment> assignments = fillAssignments(
+ memberConfigs.keySet(), Assignment.CONFIG_MISMATCH,
+ leaderId, memberConfigs.get(leaderId).url(), maxOffset,
+ ClusterAssignment.EMPTY, protocolVersion);
+ return serializeAssignments(assignments);
+ }
+ return performTaskAssignment(leaderId, leaderSnapshot, memberConfigs, coordinator, protocolVersion);
+ }
+
+ private ClusterConfigState 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;
+ }
+ }
+ return coordinator.configSnapshot();
+ }
+
+ protected Map<String, ByteBuffer> performTaskAssignment(String leaderId, ClusterConfigState configSnapshot,
+ Map<String, ExtendedWorkerState> memberConfigs,
+ WorkerCoordinator coordinator, short protocolVersion) {
+ log.debug("Performing task assignment during generation: {} with memberId: {}",
+ coordinator.generationId(), coordinator.memberId());
+ Map<String, ConnectorsAndTasks> memberAssignments = transformValues(
+ memberConfigs,
+ memberConfig -> ConnectorsAndTasks.of(
+ memberConfig.assignment().connectors(),
+ memberConfig.assignment().tasks().stream()
+ .map(org.apache.kafka.connect.util.ConnectorTaskId::uberToPublicApi)
+ .toList()
+ )
+ );
+ ClusterAssignment clusterAssignment = performTaskAssignment(
+ configSnapshot,
+ memberAssignments
+ );
+
+ coordinator.leaderState(new LeaderState(memberConfigs, clusterAssignment.allAssignedConnectors(), clusterAssignment.allAssignedTasks()));
+
+ Map<String, ExtendedAssignment> assignments =
+ fillAssignments(memberConfigs.keySet(), Assignment.NO_ERROR, leaderId,
+ memberConfigs.get(leaderId).url(), configSnapshot.offset(),
+ clusterAssignment,
+ protocolVersion);
+
+ log.debug("Actual assignments: {}", assignments);
+ return serializeAssignments(assignments);
+ }
+
+ // Visible for testing
+ ClusterAssignment performTaskAssignment(
+ ClusterConfigState configSnapshot,
+ Map<String, ConnectorsAndTasks> currentAssignments
+ ) {
+ log.debug("Current assignments: {}", currentAssignments);
+
+ Set<String> configuredConnectors = new TreeSet<>(configSnapshot.connectors());
+ Set<ConnectorTaskId> configuredTasks = combineCollections(
+ configuredConnectors,
+ connector -> configSnapshot.tasks(connector).stream().map(org.apache.kafka.connect.util.ConnectorTaskId::uberToPublicApi).toList(),
+ Collectors.toSet()
+ );
+
+ // The connectors and tasks that should be running on the cluster, based on the contents of the
+ // config topic (which is the source of truth for the cluster's intended connectors and tasks)
+ ConnectorsAndTasks configured = ConnectorsAndTasks.of(configuredConnectors, configuredTasks);
+ log.debug("Configured assignments: {}", configured);
+
+ com.uber.data.kafka.connect.distributed.ClusterAssignment nextClusterAssignment = assignor.assign(new ClusterConfigStateImpl(configSnapshot), currentAssignments);
+ Map<String, ConnectorsAndTasks> nextAssignments = nextClusterAssignment.workerAssignments();
+ log.debug("Next assignments: {}", nextAssignments);
+
+ Map<String, Collection<String>> nextConnectorAssignments = transformValues(nextAssignments, ConnectorsAndTasks::connectors);
+ Map<String, Collection<ConnectorTaskId>> nextTaskAssignments = transformValues(nextAssignments, ConnectorsAndTasks::tasks);
+
+ Map<String, Collection<String>> currentConnectorAssignments = transformValues(currentAssignments, ConnectorsAndTasks::connectors);
+ Map<String, Collection<ConnectorTaskId>> currentTaskAssignments = transformValues(currentAssignments, ConnectorsAndTasks::tasks);
+
+ // Example:
+ // Current: {C1, C2, C3}
+ // Next: {C2, C3, C4}
+ // Added: {C2, C3, C4} - {C1, C2, C3} = {C4}
+ Map<String, Collection<String>> addedConnectors = diff(nextConnectorAssignments, currentConnectorAssignments);
+ Map<String, Collection<ConnectorTaskId>> addedTasks = diff(nextTaskAssignments, currentTaskAssignments);
+ log.debug("Incremental connector assignments: {}", addedConnectors);
+ log.debug("Incremental task assignments: {}", addedTasks);
+
+ // Example:
+ // Current: {C1, C2, C3}
+ // Next: {C2, C3, C4}
+ // Revoked: {C1, C2, C3} - {C2, C3, C4} = {C1}
+ Map<String, Collection<String>> revokedConnectors = diff(currentConnectorAssignments, nextConnectorAssignments);
+ Map<String, Collection<ConnectorTaskId>> revokedTasks = diff(currentTaskAssignments, nextTaskAssignments);
+ log.debug("Revoked connector assignments: {}", revokedConnectors);
+ log.debug("Revoked task assignments: {}", revokedTasks);
+
+ return new ClusterAssignment(
+ addedConnectors,
+ toInternal(addedTasks),
+ revokedConnectors,
+ toInternal(revokedTasks),
+ nextConnectorAssignments,
+ toInternal(nextTaskAssignments),
+ nextClusterAssignment.scheduledRebalanceDelayMs()
+ );
+ }
+
+ private Map<String, ExtendedAssignment> fillAssignments(Collection<String> members, short error,
+ String leaderId, String leaderUrl, long maxOffset,
+ ClusterAssignment clusterAssignment,
+ short protocolVersion) {
+ Map<String, ExtendedAssignment> groupAssignment = new HashMap<>();
+ for (String member : members) {
+ Collection<String> connectorsToStart = clusterAssignment.newlyAssignedConnectors(member);
+ Collection<org.apache.kafka.connect.util.ConnectorTaskId> tasksToStart = clusterAssignment.newlyAssignedTasks(member);
+ Collection<String> connectorsToStop = clusterAssignment.newlyRevokedConnectors(member);
+ Collection<org.apache.kafka.connect.util.ConnectorTaskId> tasksToStop = clusterAssignment.newlyRevokedTasks(member);
+ ExtendedAssignment assignment =
+ new ExtendedAssignment(protocolVersion, error, leaderId, leaderUrl, maxOffset,
+ connectorsToStart, tasksToStart, connectorsToStop, tasksToStop, clusterAssignment.uberScheduledRebalanceDelayMs());
+ 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<String, ByteBuffer> serializeAssignments(Map<String, ExtendedAssignment> assignments) {
+ int clusterSize = assignments.size();
+ return assignments.entrySet()
+ .stream()
+ .collect(Collectors.toMap(
+ Map.Entry::getKey,
+ e -> {
+ UberAssignmentV1 uberAssignment = UberAssignmentV1.fromExtended(e.getValue(), clusterSize);
+ return UberConnectProtocol.serializeAssignment(uberAssignment);
+ }));
+ }
+
+ private static <T> Map<String, Collection<T>> diff(Map<String, Collection<T>> base,
+ Map<String, Collection<T>> toSubtract) {
+ Map<String, Collection<T>> incremental = new HashMap<>();
+ for (Map.Entry<String, Collection<T>> entry : base.entrySet()) {
+ List<T> values = new ArrayList<>(entry.getValue());
+ values.removeAll(toSubtract.getOrDefault(entry.getKey(), Set.of()));
+ incremental.put(entry.getKey(), values);
+ }
+ return incremental;
+ }
+
+ private static Map<String, Collection<org.apache.kafka.connect.util.ConnectorTaskId>> toInternal(
+ Map<String, Collection<ConnectorTaskId>> publicApi
+ ) {
+ return publicApi.entrySet().stream()
+ .collect(Collectors.toMap(
+ Map.Entry::getKey,
+ e -> toInternal(e.getValue())
+ ));
+ }
+
+ private static Collection<org.apache.kafka.connect.util.ConnectorTaskId> toInternal(Collection<ConnectorTaskId> publicApi) {
+ return publicApi.stream()
+ .map(org.apache.kafka.connect.util.ConnectorTaskId::uberFromPublicApi)
+ .toList();
+ }
+
+ static class ClusterAssignment extends IncrementalCooperativeAssignor.ClusterAssignment {
+
+ public static final ClusterAssignment EMPTY = new ClusterAssignment(
+ Map.of(),
+ Map.of(),
+ Map.of(),
+ Map.of(),
+ Map.of(),
+ Map.of(),
+ 0
+ );
+
+ private final int uberScheduledRebalanceDelayMs;
+
+ public ClusterAssignment(
+ Map<String, Collection<String>> newlyAssignedConnectors,
+ Map<String, Collection<org.apache.kafka.connect.util.ConnectorTaskId>> newlyAssignedTasks,
+ Map<String, Collection<String>> newlyRevokedConnectors,
+ Map<String, Collection<org.apache.kafka.connect.util.ConnectorTaskId>> newlyRevokedTasks,
+ Map<String, Collection<String>> allAssignedConnectors,
+ Map<String, Collection<org.apache.kafka.connect.util.ConnectorTaskId>> allAssignedTasks,
+ int uberScheduledRebalanceDelayMs
+ ) {
+ super(
+ newlyAssignedConnectors,
+ newlyAssignedTasks,
+ newlyRevokedConnectors,
+ newlyRevokedTasks,
+ allAssignedConnectors,
+ allAssignedTasks
+ );
+ this.uberScheduledRebalanceDelayMs = uberScheduledRebalanceDelayMs;
+ }
+
+ public int uberScheduledRebalanceDelayMs() {
+ return uberScheduledRebalanceDelayMs;
+ }
+
+ @Override
+ public String toString() {
+ return "ClusterAssignment{"
+ + "newlyAssignedConnectors=" + newlyAssignedConnectors()
+ + ", newlyAssignedTasks=" + newlyAssignedTasks()
+ + ", newlyRevokedConnectors=" + newlyRevokedConnectors()
+ + ", newlyRevokedTasks=" + newlyRevokedTasks()
+ + ", allAssignedConnectors=" + allAssignedConnectors()
+ + ", allAssignedTasks=" + allAssignedTasks()
+ + ", allWorkers=" + allWorkers()
+ + ", uberScheduledRebalanceDelayMs=" + uberScheduledRebalanceDelayMs
+ + '}';
+ }
+ }
+
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberConnectProtocol.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberConnectProtocol.java
new file mode 100644
index 00000000000..ba7855e9c16
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberConnectProtocol.java
@@ -0,0 +1,193 @@
+package org.apache.kafka.connect.runtime.distributed;
+
+import org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocol;
+import org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequestProtocolCollection;
+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.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.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;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocolCompatibility.UBER_V1;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.ALLOCATION_KEY_NAME;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.ALLOCATION_V1;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.CONFIG_STATE_V1;
+import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeConnectProtocol.CONNECTOR_ASSIGNMENT_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;
+
+public class UberConnectProtocol {
+
+ public static final String UBER_CLUSTER_SIZE_KEY_NAME = "uberClusterSize";
+ public static final short UBER_CONNECT_PROTOCOL_V1 = 1001;
+ public static final boolean TOLERATE_MISSING_FIELDS_WITH_DEFAULTS = true;
+
+ /**
+ * Uber Connect Protocol Header V3:
+ * <pre>
+ * Version => Int16
+ * </pre>
+ * The Uber V3 protocol is identical to Connect V2, <em>except</em> that it also includes the total cluster size
+ * in each member's assignment
+ */
+ private static final Struct UBER_CONNECT_PROTOCOL_HEADER_V1 = new Struct(CONNECT_PROTOCOL_HEADER_SCHEMA)
+ .set(VERSION_KEY_NAME, UBER_CONNECT_PROTOCOL_V1);
+
+ /**
+ * Raw (non versioned) Uber assignment V1:
+ * <pre>
+ * Error => Int16
+ * Leader => [String]
+ * LeaderUrl => [String]
+ * ConfigOffset => Int64
+ * Assignment => [Connector Assignment]
+ * Revoked => [Connector Assignment]
+ * ScheduledDelay => Int32
+ * UberClusterSize => Int32
+ * </pre>
+ */
+ public static final Schema UBER_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),
+ new Field(UBER_CLUSTER_SIZE_KEY_NAME, Type.INT32, null, 0));
+
+ /**
+ * The fields are serialized in sequence as follows:
+ * Subscription V1:
+ * <pre>
+ * Version => Int16
+ * Url => [String]
+ * ConfigOffset => Int64
+ * Current Assignment => [Byte]
+ * </pre>
+ */
+ public static ByteBuffer serializeMetadata(UberWorkerState workerState) {
+ 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 = UBER_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
+ * @return the collection of Connect protocol metadata
+ */
+ public static JoinGroupRequestProtocolCollection metadataRequest(UberWorkerState workerState) {
+ // Order matters in terms of protocol preference
+ List<JoinGroupRequestProtocol> joinGroupRequestProtocols = new ArrayList<>();
+
+ // Metadata for our custom protocol
+ joinGroupRequestProtocols.add(new JoinGroupRequestProtocol()
+ .setName(UBER_V1.protocol())
+ .setMetadata(serializeMetadata(workerState).array())
+ );
+ // Metadata for OSS Connect protocols
+ 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());
+ }
+
+ /**
+ * The fields are serialized in sequence as follows:
+ * Complete Assignment V1:
+ * <pre>
+ * Version => Int16
+ * Error => Int16
+ * Leader => [String]
+ * LeaderUrl => [String]
+ * ConfigOffset => Int64
+ * Assignment => [Connector Assignment]
+ * Revoked => [Connector Assignment]
+ * ScheduledDelay => Int32
+ * UberClusterSize => Int32
+ * </pre>
+ */
+ public static ByteBuffer serializeAssignment(UberAssignmentV1 assignment) {
+ // comparison depends on reference equality for now
+ if (assignment == null || UberAssignmentV1.empty().equals(assignment)) {
+ return null;
+ }
+ Struct struct = assignment.toUberStruct();
+ Struct protocolHeader = UBER_CONNECT_PROTOCOL_HEADER_V1;
+ ByteBuffer buffer = ByteBuffer.allocate(protocolHeader.sizeOf()
+ + UBER_ASSIGNMENT_V1.sizeOf(struct));
+ protocolHeader.writeTo(buffer);
+ UBER_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 UberAssignmentV1 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 = UBER_ASSIGNMENT_V1.read(buffer);
+ return UberAssignmentV1.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/UberWorkerLoad.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberWorkerLoad.java
new file mode 100644
index 00000000000..ae0a92a700b
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberWorkerLoad.java
@@ -0,0 +1,63 @@
+package org.apache.kafka.connect.runtime.distributed;
+
+import com.uber.data.kafka.connect.distributed.ConnectorTaskId;
+import com.uber.data.kafka.connect.distributed.ConnectorsAndTasks;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Objects;
+
+public record UberWorkerLoad(String worker, Collection<String> connectors, Collection<ConnectorTaskId> tasks) {
+
+ public UberWorkerLoad(String worker, ConnectorsAndTasks connectorsAndTasks) {
+ this(worker, connectorsAndTasks.connectors(), connectorsAndTasks.tasks());
+ }
+
+ public UberWorkerLoad(
+ String worker,
+ Collection<String> connectors,
+ Collection<ConnectorTaskId> tasks
+ ) {
+ this.worker = worker;
+ this.connectors = connectors != null ? new ArrayList<>(connectors) : new ArrayList<>();
+ this.tasks = tasks != null ? new ArrayList<>(tasks) : new ArrayList<>();
+ }
+
+ public void assign(String connector) {
+ connectors.add(connector);
+ }
+
+ public void assign(ConnectorTaskId task) {
+ tasks.add(task);
+ }
+
+ @Override
+ public String toString() {
+ return "UberWorkerLoad{ worker=" + worker + ", connectorIds=" + connectors + ", taskIds=" + tasks + '}';
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (!(o instanceof UberWorkerLoad that)) {
+ return false;
+ }
+ return Objects.equals(this.worker, that.worker);
+ }
+
+ /**
+ * <b>IMPORTANT:</b> we do not hash on the sets of connectors or tasks, for two vital reasons:
+ * <ul>
+ * <li>These collections are mutable and expected to change over the lifetime of the object, causing its hash code to change</li>
+ * <li>These collections may grow fairly large and the time to compute their hash code should remain as low as possible since they
+ * may be used as keys in maps for, e.g., memoization of total worker load</li>
+ * </ul>
+ * @return the hash code of the worker load, derived solely from the worker ID
+ */
+ @Override
+ public int hashCode() {
+ return Objects.hash(worker);
+ }
+}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberWorkerState.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberWorkerState.java
new file mode 100644
index 00000000000..ffe62385af2
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberWorkerState.java
@@ -0,0 +1,16 @@
+package org.apache.kafka.connect.runtime.distributed;
+
+public class UberWorkerState extends ExtendedWorkerState {
+
+ private final UberAssignmentV1 assignment;
+
+ public UberWorkerState(String url, long offset, UberAssignmentV1 assignment) {
+ super(url, offset, assignment);
+ this.assignment = assignment;
+ }
+
+ @Override
+ public UberAssignmentV1 assignment() {
+ return assignment;
+ }
+}
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 edb174ab52b..192f035c6b7 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
@@ -16,6 +16,7 @@
*/
package org.apache.kafka.connect.runtime.distributed;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import org.apache.kafka.clients.GroupRebalanceConfig;
import org.apache.kafka.clients.consumer.CloseOptions;
import org.apache.kafka.clients.consumer.internals.AbstractCoordinator;
@@ -49,6 +50,7 @@ import static org.apache.kafka.common.message.JoinGroupRequestData.JoinGroupRequ
import static org.apache.kafka.common.message.JoinGroupResponseData.JoinGroupResponseMember;
import static org.apache.kafka.common.utils.Utils.UncheckedCloseable;
import static org.apache.kafka.connect.runtime.distributed.ConnectProtocolCompatibility.EAGER;
+import static org.apache.kafka.connect.runtime.distributed.ConnectProtocolCompatibility.UBER_V1;
/**
@@ -59,7 +61,7 @@ public class WorkerCoordinator extends AbstractCoordinator implements Closeable
private final Logger log;
private final String restUrl;
private final ConfigBackingStore configStorage;
- private volatile ExtendedAssignment assignmentSnapshot;
+ private volatile UberAssignmentV1 assignmentSnapshot;
private ClusterConfigState configSnapshot;
private final WorkerRebalanceListener listener;
private final ConnectProtocolCompatibility protocolCompatibility;
@@ -70,6 +72,7 @@ public class WorkerCoordinator extends AbstractCoordinator implements Closeable
private volatile int lastCompletedGenerationId;
private final ConnectAssignor eagerAssignor;
private final ConnectAssignor incrementalAssignor;
+ private final ConnectAssignor uberAssignor;
private final int coordinatorDiscoveryTimeoutMs;
/**
@@ -84,6 +87,7 @@ public class WorkerCoordinator extends AbstractCoordinator implements Closeable
String restUrl,
ConfigBackingStore configStorage,
WorkerRebalanceListener listener,
+ ClusterAssignor clusterAssignor,
ConnectProtocolCompatibility protocolCompatibility,
int maxDelay) {
super(config,
@@ -100,6 +104,17 @@ public class WorkerCoordinator extends AbstractCoordinator implements Closeable
this.listener = listener;
this.rejoinRequested = false;
this.protocolCompatibility = protocolCompatibility;
+
+ if (clusterAssignor == null) {
+ log.warn("No cluster assignor given; the default Kafka Connect rebalance protocol will be used and workload-based load balancing will be disabled");
+ this.uberAssignor = null;
+ } else {
+ this.uberAssignor = new UberAssignor(
+ logContext,
+ clusterAssignor
+ );
+ }
+
this.incrementalAssignor = new IncrementalCooperativeAssignor(logContext, time, maxDelay);
this.eagerAssignor = new EagerAssignor(logContext);
this.currentConnectProtocol = protocolCompatibility;
@@ -181,18 +196,19 @@ public class WorkerCoordinator extends AbstractCoordinator implements Closeable
@Override
public JoinGroupRequestProtocolCollection metadata() {
configSnapshot = configStorage.snapshot();
- final ExtendedAssignment localAssignmentSnapshot = assignmentSnapshot;
- ExtendedWorkerState workerState = new ExtendedWorkerState(restUrl, configSnapshot.offset(), localAssignmentSnapshot);
+ final UberAssignmentV1 localAssignmentSnapshot = assignmentSnapshot;
+ UberWorkerState workerState = new UberWorkerState(restUrl, configSnapshot.offset(), localAssignmentSnapshot);
return switch (protocolCompatibility) {
case EAGER -> ConnectProtocol.metadataRequest(workerState);
case COMPATIBLE -> IncrementalCooperativeConnectProtocol.metadataRequest(workerState, false);
case SESSIONED -> IncrementalCooperativeConnectProtocol.metadataRequest(workerState, true);
+ case UBER_V1 -> UberConnectProtocol.metadataRequest(workerState);
};
}
@Override
protected void onJoinComplete(int generation, String memberId, String protocol, ByteBuffer memberAssignment) {
- ExtendedAssignment newAssignment = IncrementalCooperativeConnectProtocol.deserializeAssignment(memberAssignment);
+ UberAssignmentV1 newAssignment = UberConnectProtocol.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
@@ -229,9 +245,19 @@ public class WorkerCoordinator extends AbstractCoordinator implements Closeable
throw new IllegalStateException("Can't skip assignment because Connect does not support static membership.");
ConnectProtocolCompatibility protocolCompatibility = ConnectProtocolCompatibility.fromProtocol(protocol);
- return protocolCompatibility == EAGER
- ? eagerAssignor.performAssignment(leaderId, protocolCompatibility, allMemberMetadata, this)
- : incrementalAssignor.performAssignment(leaderId, protocolCompatibility, allMemberMetadata, this);
+
+ final ConnectAssignor assignor;
+ if (protocolCompatibility == EAGER) {
+ assignor = eagerAssignor;
+ } else if (protocolCompatibility == UBER_V1) {
+ if (uberAssignor == null)
+ throw new IllegalStateException("Cannot perform assignment using protocol " + UBER_V1.protocol() + " when no cluster assignor has been configured");
+ assignor = uberAssignor;
+ } else {
+ assignor = incrementalAssignor;
+ }
+
+ return assignor.performAssignment(leaderId, protocolCompatibility, allMemberMetadata, this);
}
@Override
@@ -469,6 +495,10 @@ public class WorkerCoordinator extends AbstractCoordinator implements Closeable
this.tasks = tasks;
}
+ public static ConnectorsAndTasks of(Collection<String> connectors, Collection<ConnectorTaskId> tasks) {
+ return new Builder().with(connectors, tasks).build();
+ }
+
public static class Builder {
private Set<String> withConnectors = new LinkedHashSet<>();
private Set<ConnectorTaskId> withTasks = new LinkedHashSet<>();
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 a283ce7cf38..c3dc2acdd3f 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
@@ -16,6 +16,7 @@
*/
package org.apache.kafka.connect.runtime.distributed;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import org.apache.kafka.clients.ApiVersions;
import org.apache.kafka.clients.ClientUtils;
import org.apache.kafka.clients.CommonClientConfigs;
@@ -68,6 +69,7 @@ public class WorkerGroupMember {
private final String clientId;
private final ConsumerNetworkClient client;
private final Metrics metrics;
+ private final ClusterAssignor clusterAssignor;
private final WorkerCoordinator coordinator;
private boolean stopped = false;
@@ -76,6 +78,7 @@ public class WorkerGroupMember {
String restUrl,
ConfigBackingStore configStorage,
WorkerRebalanceListener listener,
+ ClusterAssignor clusterAssignor,
Time time,
String clientId,
LogContext logContext) {
@@ -133,6 +136,7 @@ public class WorkerGroupMember {
retryBackoffMs,
config.getInt(CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG),
Integer.MAX_VALUE);
+ this.clusterAssignor = clusterAssignor;
this.coordinator = new WorkerCoordinator(
new GroupRebalanceConfig(config, GroupRebalanceConfig.ProtocolType.CONNECT),
logContext,
@@ -143,6 +147,7 @@ public class WorkerGroupMember {
restUrl,
configStorage,
listener,
+ clusterAssignor,
ConnectProtocolCompatibility.compatibility(config.getString(DistributedConfig.CONNECT_PROTOCOL_CONFIG)),
config.getInt(DistributedConfig.SCHEDULED_REBALANCE_MAX_DELAY_MS_CONFIG));
@@ -230,6 +235,7 @@ public class WorkerGroupMember {
Utils.closeQuietly(coordinator, "coordinator", firstException);
Utils.closeQuietly(metrics, "consumer metrics", firstException);
Utils.closeQuietly(client, "consumer network client", firstException);
+ Utils.closeQuietly(clusterAssignor, "cluster assignor", firstException);
AppInfoParser.unregisterAppInfo(JMX_PREFIX, clientId, metrics);
if (firstException.get() != null && !swallowException)
throw new KafkaException("Failed to stop the Connect group member", firstException.get());
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImpl.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImpl.java
index c3af78b78b8..3f2ce96766b 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImpl.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImpl.java
@@ -17,6 +17,8 @@
package org.apache.kafka.connect.runtime.health;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Timer;
import org.apache.kafka.connect.errors.ConnectException;
import org.apache.kafka.connect.health.ConnectClusterDetails;
import org.apache.kafka.connect.health.ConnectClusterState;
@@ -25,18 +27,26 @@ 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.distributed.RebalanceNeededException;
import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo;
+import org.apache.kafka.connect.runtime.rest.entities.TaskInfo;
import org.apache.kafka.connect.util.FutureCallback;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
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.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
public class ConnectClusterStateImpl implements ConnectClusterState {
+
+ private static final Logger log = LoggerFactory.getLogger(ConnectClusterStateImpl.class);
private final long herderRequestTimeoutMs;
private final ConnectClusterDetails clusterDetails;
@@ -84,14 +94,8 @@ public class ConnectClusterStateImpl implements ConnectClusterState {
public Map<String, String> connectorConfig(String connName) {
FutureCallback<Map<String, String>> 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
- );
- }
+ String operation = "retrieve configuration for connector '" + connName + "'";
+ return getWithRetry(connectorConfigCallback, operation);
}
@Override
@@ -99,6 +103,24 @@ public class ConnectClusterStateImpl implements ConnectClusterState {
return clusterDetails;
}
+ // TODO: This patch can definitely be contributed upstream
+ private <T> T getWithRetry(Future<T> future, String operation) {
+ Timer timer = Time.SYSTEM.timer(herderRequestTimeoutMs);
+ while (true) {
+ try {
+ timer.update();
+ long remainingMs = Math.max(1, timer.remainingMs());
+ return future.get(remainingMs, TimeUnit.MILLISECONDS);
+ } catch (InterruptedException | TimeoutException | ExecutionException e) {
+ if (e instanceof ExecutionException && e.getCause() instanceof RebalanceNeededException cause) {
+ log.debug("Failed to {} due to pending rebalance; will retry", operation, cause);
+ continue;
+ }
+ throw new ConnectException("Failed to " + operation, e);
+ }
+ }
+ }
+
private Map<Integer, TaskState> taskStates(List<ConnectorStateInfo.TaskState> states) {
Map<Integer, TaskState> taskStates = new HashMap<>();
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 04be35f2d26..3c6a193a4bf 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,6 +16,7 @@
*/
package org.apache.kafka.connect.runtime.isolation;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import org.apache.kafka.common.config.provider.ConfigProvider;
import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
import org.apache.kafka.connect.rest.ConnectRestExtension;
@@ -42,6 +43,7 @@ public class PluginScanResult {
private final SortedSet<PluginDesc<ConfigProvider>> configProviders;
private final SortedSet<PluginDesc<ConnectRestExtension>> restExtensions;
private final SortedSet<PluginDesc<ConnectorClientConfigOverridePolicy>> connectorClientConfigPolicies;
+ private final SortedSet<PluginDesc<ClusterAssignor>> clusterAssignors;
private final List<SortedSet<? extends PluginDesc<?>>> allPlugins;
@@ -54,7 +56,8 @@ public class PluginScanResult {
SortedSet<PluginDesc<Predicate<?>>> predicates,
SortedSet<PluginDesc<ConfigProvider>> configProviders,
SortedSet<PluginDesc<ConnectRestExtension>> restExtensions,
- SortedSet<PluginDesc<ConnectorClientConfigOverridePolicy>> connectorClientConfigPolicies
+ SortedSet<PluginDesc<ConnectorClientConfigOverridePolicy>> connectorClientConfigPolicies,
+ SortedSet<PluginDesc<ClusterAssignor>> clusterAssignors
) {
this.sinkConnectors = sinkConnectors;
this.sourceConnectors = sourceConnectors;
@@ -65,9 +68,10 @@ public class PluginScanResult {
this.configProviders = configProviders;
this.restExtensions = restExtensions;
this.connectorClientConfigPolicies = connectorClientConfigPolicies;
+ this.clusterAssignors = clusterAssignors;
this.allPlugins =
List.of(sinkConnectors, sourceConnectors, converters, headerConverters, transformations, predicates,
- configProviders, restExtensions, connectorClientConfigPolicies);
+ configProviders, restExtensions, connectorClientConfigPolicies, clusterAssignors);
}
/**
@@ -83,7 +87,8 @@ public class PluginScanResult {
merge(results, PluginScanResult::predicates),
merge(results, PluginScanResult::configProviders),
merge(results, PluginScanResult::restExtensions),
- merge(results, PluginScanResult::connectorClientConfigPolicies)
+ merge(results, PluginScanResult::connectorClientConfigPolicies),
+ merge(results, PluginScanResult::clusterAssignors)
);
}
@@ -131,6 +136,10 @@ public class PluginScanResult {
return connectorClientConfigPolicies;
}
+ public SortedSet<PluginDesc<ClusterAssignor>> clusterAssignors() {
+ return clusterAssignors;
+ }
+
public void forEach(Consumer<PluginDesc<?>> consumer) {
allPlugins.forEach(plugins -> plugins.forEach(consumer));
}
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 68e3cac36bd..05a4766df71 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,6 +16,7 @@
*/
package org.apache.kafka.connect.runtime.isolation;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import org.apache.kafka.common.config.provider.ConfigProvider;
import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
import org.apache.kafka.connect.rest.ConnectRestExtension;
@@ -39,7 +40,8 @@ public enum PluginType {
PREDICATE(Predicate.class),
CONFIGPROVIDER(ConfigProvider.class),
REST_EXTENSION(ConnectRestExtension.class),
- CONNECTOR_CLIENT_CONFIG_OVERRIDE_POLICY(ConnectorClientConfigOverridePolicy.class);
+ CONNECTOR_CLIENT_CONFIG_OVERRIDE_POLICY(ConnectorClientConfigOverridePolicy.class),
+ CLUSTER_ASSIGNOR(ClusterAssignor.class);
private final Class<?> klass;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/ReflectionScanner.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/ReflectionScanner.java
index 85f514d5f1e..2927b0844ed 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/ReflectionScanner.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/ReflectionScanner.java
@@ -92,7 +92,8 @@ public class ReflectionScanner extends PluginScanner {
getPredicatePluginDesc(source, classGraph),
getServiceLoaderPluginDesc(PluginType.CONFIGPROVIDER, source),
getServiceLoaderPluginDesc(PluginType.REST_EXTENSION, source),
- getServiceLoaderPluginDesc(PluginType.CONNECTOR_CLIENT_CONFIG_OVERRIDE_POLICY, source)
+ getServiceLoaderPluginDesc(PluginType.CONNECTOR_CLIENT_CONFIG_OVERRIDE_POLICY, source),
+ getServiceLoaderPluginDesc(PluginType.CLUSTER_ASSIGNOR, source)
);
}
}
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/ServiceLoaderScanner.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/ServiceLoaderScanner.java
index 9f36a5ad8e0..bab499121a4 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/ServiceLoaderScanner.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/ServiceLoaderScanner.java
@@ -64,7 +64,8 @@ public class ServiceLoaderScanner extends PluginScanner {
getPredicatePluginDesc(source),
getServiceLoaderPluginDesc(PluginType.CONFIGPROVIDER, source),
getServiceLoaderPluginDesc(PluginType.REST_EXTENSION, source),
- getServiceLoaderPluginDesc(PluginType.CONNECTOR_CLIENT_CONFIG_OVERRIDE_POLICY, source)
+ getServiceLoaderPluginDesc(PluginType.CONNECTOR_CLIENT_CONFIG_OVERRIDE_POLICY, source),
+ getServiceLoaderPluginDesc(PluginType.CLUSTER_ASSIGNOR, source)
);
}
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 148e96a4cee..bef1fe2cb8c 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
@@ -75,6 +75,8 @@ import static org.apache.kafka.connect.runtime.rest.HerderRequestHandler.Transla
@Produces(MediaType.APPLICATION_JSON)
@Consumes(MediaType.APPLICATION_JSON)
public class ConnectorsResource {
+
+ public static final String UBER_FORCE_OFFSET_PATCH = "uber-force";
private static final Logger log = LoggerFactory.getLogger(ConnectorsResource.class);
private final Herder herder;
@@ -361,14 +363,16 @@ public class ConnectorsResource {
@Operation(summary = "Alter the offsets for the specified connector")
public Response alterConnectorOffsets(final @Parameter(hidden = true) @QueryParam("forward") Boolean forward,
final @Context HttpHeaders headers, final @PathParam("connector") String connector,
- final ConnectorOffsets offsets) throws Throwable {
+ final ConnectorOffsets offsets, @QueryParam(UBER_FORCE_OFFSET_PATCH) Boolean uberForce) throws Throwable {
if (offsets.offsets() == null || offsets.offsets().isEmpty()) {
throw new BadRequestException("Partitions / offsets need to be provided for an alter offsets request");
}
FutureCallback<Message> cb = new FutureCallback<>();
- herder.alterConnectorOffsets(connector, offsets.toMap(), cb);
- Message msg = requestHandler.completeOrForwardRequest(cb, "/connectors/" + connector + "/offsets", "PATCH", headers, offsets,
+ boolean uberForceOrDefault = uberForce != null ? uberForce : false;
+ herder.uberAlterConnectorOffsets(connector, offsets.toMap(), uberForceOrDefault, cb);
+ String path = "/connectors/" + connector + "/offsets" + (uberForce != null ? "?" + UBER_FORCE_OFFSET_PATCH + "=" + uberForce : "");
+ Message msg = requestHandler.completeOrForwardRequest(cb, path, "PATCH", headers, offsets,
new TypeReference<>() { }, new IdentityTranslator<>(), forward);
return Response.ok().entity(msg).build();
}
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 e654ded2239..9dbdce883bd 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
@@ -424,8 +424,8 @@ public final class StandaloneHerder extends AbstractHerder {
}
@Override
- protected synchronized void modifyConnectorOffsets(String connName, Map<Map<String, ?>, Map<String, ?>> offsets, Callback<Message> cb) {
- if (!modifyConnectorOffsetsChecks(connName, cb)) {
+ protected synchronized void uberModifyConnectorOffsets(String connName, Map<Map<String, ?>, Map<String, ?>> offsets, boolean uberForce, Callback<Message> cb) {
+ if (!modifyConnectorOffsetsChecks(connName, uberForce, cb)) {
return;
}
@@ -436,18 +436,24 @@ public final class StandaloneHerder extends AbstractHerder {
* This method performs a few checks for external requests to modify (alter or reset) connector offsets and
* completes the callback exceptionally if any check fails.
* @param connName the name of the connector whose offsets are to be modified
+ * @param uberForce force the offset modification to take place even if the connector is not stopped
* @param cb callback to invoke upon completion
* @return true if all the checks passed, false otherwise
*/
- private boolean modifyConnectorOffsetsChecks(String connName, Callback<Message> cb) {
+ private boolean modifyConnectorOffsetsChecks(String connName, boolean uberForce, Callback<Message> cb) {
if (!configState.contains(connName)) {
cb.onCompletion(new NotFoundException("Connector " + connName + " not found", null), null);
return false;
}
- if (configState.targetState(connName) != TargetState.STOPPED || configState.taskCount(connName) != 0) {
- cb.onCompletion(new BadRequestException("Connectors must be in the STOPPED state before their offsets can be modified. This can be done " +
- "for the specified connector by issuing a 'PUT' request to the '/connectors/" + connName + "/stop' endpoint"), null);
+ if (!uberForce && (configState.targetState(connName) != TargetState.STOPPED || configState.taskCount(connName) != 0)) {
+ String message = "Connectors must be in the STOPPED state before their offsets can be modified. This can be done "
+ + "for the specified connector by issuing a 'PUT' request to the '/connectors/" + connName + "/stop' endpoint. "
+ + "Alternatively, for offset alterations (but not resets), you may force the request to take place via the ?uber-force=true "
+ + "URL query parameter; however, in order for the change to take effect properly, the connector must not be "
+ + "actively processing data for the affected offsets (e.g., if it's replicating from a Kafka topic t1, you must not "
+ + "alter offsets for the t1 topic, but you may alter offsets for topics t2, t3, etc.)";
+ cb.onCompletion(new BadRequestException(message), null);
return false;
}
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 613caf1dbad..360ad465fa7 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
@@ -56,4 +56,12 @@ public record ConnectorTaskId(String connector, int task) implements Serializabl
return connectorCmp;
return Integer.compare(task, o.task);
}
+
+ public com.uber.data.kafka.connect.distributed.ConnectorTaskId uberToPublicApi() {
+ return new com.uber.data.kafka.connect.distributed.ConnectorTaskId(connector, task);
+ }
+
+ public static ConnectorTaskId uberFromPublicApi(com.uber.data.kafka.connect.distributed.ConnectorTaskId connectorTaskId) {
+ return new ConnectorTaskId(connectorTaskId.connector(), connectorTaskId.task());
+ }
}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/OffsetsApiIntegrationTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/OffsetsApiIntegrationTest.java
index c9edd6093bd..f30adc0f002 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/OffsetsApiIntegrationTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/OffsetsApiIntegrationTest.java
@@ -46,6 +46,7 @@ import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
+import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
@@ -600,7 +601,7 @@ public class OffsetsApiIntegrationTest {
// The TestableSourceConnector has a source partition per task
for (int i = 0; i < NUM_TASKS; i++) {
offsetsToAlter.add(
- new ConnectorOffset(Map.of("task.id", connectorName + "-" + i), Map.of("saved", 5))
+ new ConnectorOffset(sourcePartition(i), Map.of("saved", 5))
);
}
@@ -620,7 +621,7 @@ public class OffsetsApiIntegrationTest {
// The TestableSourceConnector has a source partition per task
for (int i = 0; i < NUM_TASKS; i++) {
offsetsToAlter.add(
- new ConnectorOffset(Map.of("task.id", connectorName + "-" + i), Map.of("saved", 7))
+ new ConnectorOffset(sourcePartition(i), Map.of("saved", 7))
);
}
@@ -1082,4 +1083,35 @@ public class OffsetsApiIntegrationTest {
return offsets.offsets().isEmpty();
}, OFFSET_READ_TIMEOUT_MS, "Connector offsets should be empty after resetting offsets");
}
+
+ private Map<String, Object> sourcePartition(int taskId) {
+ Map<String, Object> result = new LinkedHashMap<>();
+ result.put("task.id", connectorName + "-" + taskId);
+ result.put("property-a", "a");
+ result.put("property-b", "b");
+ result.put("property-c", "c");
+ result.put("property-d", "d");
+ result.put("property-e", "e");
+ result.put("property-f", "f");
+ result.put("property-g", "g");
+ result.put("property-h", "h");
+ result.put("property-i", "i");
+ result.put("property-j", "j");
+ result.put("property-k", "k");
+ result.put("property-l", "l");
+ result.put("property-m", "m");
+ result.put("property-n", "n");
+ result.put("property-o", "o");
+ result.put("property-p", "p");
+ result.put("property-q", "q");
+ result.put("property-r", "r");
+ result.put("property-s", "s");
+ result.put("property-t", "t");
+ result.put("property-u", "u");
+ result.put("property-v", "v");
+ result.put("property-x", "x");
+ result.put("property-y", "y");
+ result.put("property-z", "z");
+ return result;
+ }
}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TestableSourceConnector.java b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TestableSourceConnector.java
index dec94210589..1e8227dd85f 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TestableSourceConnector.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/integration/TestableSourceConnector.java
@@ -34,6 +34,7 @@ import org.slf4j.LoggerFactory;
import java.util.ArrayList;
import java.util.HashMap;
+import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
@@ -293,7 +294,39 @@ public class TestableSourceConnector extends SampleSourceConnector {
}
public static Map<String, Object> sourcePartition(String taskId) {
- return Map.of("task.id", taskId);
+ // LinkedHashMap guarantees iteration order, which matters during serialization
+ Map<String, Object> result = new LinkedHashMap<>();
+ result.put("task.id", taskId);
+ // Garbage key/value pairs to ensure that serialization and deserialization of
+ // source partitions with multiple fields works
+ // We put task ID first to make sure that the actual order of key/value pairs
+ // is preserved, instead of something else (like sorting them lexicographically)
+ result.put("property-a", "a");
+ result.put("property-b", "b");
+ result.put("property-c", "c");
+ result.put("property-d", "d");
+ result.put("property-e", "e");
+ result.put("property-f", "f");
+ result.put("property-g", "g");
+ result.put("property-h", "h");
+ result.put("property-i", "i");
+ result.put("property-j", "j");
+ result.put("property-k", "k");
+ result.put("property-l", "l");
+ result.put("property-m", "m");
+ result.put("property-n", "n");
+ result.put("property-o", "o");
+ result.put("property-p", "p");
+ result.put("property-q", "q");
+ result.put("property-r", "r");
+ result.put("property-s", "s");
+ result.put("property-t", "t");
+ result.put("property-u", "u");
+ result.put("property-v", "v");
+ result.put("property-x", "x");
+ result.put("property-y", "y");
+ result.put("property-z", "z");
+ return result;
}
public static Map<String, Object> sourceOffset(long seqno) {
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
index 047a3b5d4e9..2aecb451664 100644
--- 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
@@ -159,6 +159,7 @@ public class WorkerCoordinatorIncrementalTest {
expectedUrl(leaderId),
configStorage,
rebalanceListener,
+ null,
compatibility,
rebalanceDelay);
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 46a5b076a33..eb943f6a544 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
@@ -150,6 +150,7 @@ public class WorkerCoordinatorTest {
LEADER_URL,
configStorage,
rebalanceListener,
+ null,
compatibility,
0);
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
index 4886431869c..c67e90e7e0b 100644
--- 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
@@ -72,7 +72,7 @@ public class WorkerGroupMemberTest {
doReturn("cluster-1").when(config).kafkaClusterId();
LogContext logContext = new LogContext("[Worker clientId=client-1 + groupId= group-1]");
- member = new WorkerGroupMember(config, "", configBackingStore, null, Time.SYSTEM, "client-1", logContext);
+ member = new WorkerGroupMember(config, "", configBackingStore, null, null, Time.SYSTEM, "client-1", logContext);
verify(config, atLeastOnce()).kafkaClusterId();
boolean foundMockReporter = false;
@@ -115,7 +115,7 @@ public class WorkerGroupMemberTest {
doReturn("cluster-1").when(config).kafkaClusterId();
LogContext logContext = new LogContext("[Worker clientId=client-1 + groupId= group-1]");
- member = new WorkerGroupMember(config, "", configBackingStore, null, Time.SYSTEM, "client-1", logContext);
+ member = new WorkerGroupMember(config, "", configBackingStore, null, null, Time.SYSTEM, "client-1", logContext);
verify(config, atLeastOnce()).kafkaClusterId();
assertTrue(member.metrics().reporters().isEmpty());
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/DelegatingClassLoaderTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/DelegatingClassLoaderTest.java
index 76bad03fb01..e3a448a6287 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/DelegatingClassLoaderTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/DelegatingClassLoaderTest.java
@@ -128,6 +128,7 @@ public class DelegatingClassLoaderTest {
new TreeSet<>(),
new TreeSet<>(),
new TreeSet<>(),
+ new TreeSet<>(),
new TreeSet<>()
);
}
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 24ef0d535b8..ed0029f2c95 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
@@ -516,6 +516,7 @@ public class PluginUtilsTest {
new TreeSet<>(),
new TreeSet<>(),
new TreeSet<>(),
+ new TreeSet<>(),
new TreeSet<>()
);
Map<String, String> actualAliases = PluginUtils.computeAliases(result);
@@ -545,6 +546,7 @@ public class PluginUtilsTest {
new TreeSet<>(),
new TreeSet<>(),
new TreeSet<>(),
+ new TreeSet<>(),
new TreeSet<>()
);
Map<String, String> actualAliases = PluginUtils.computeAliases(result);
@@ -569,6 +571,7 @@ public class PluginUtilsTest {
new TreeSet<>(),
new TreeSet<>(),
new TreeSet<>(),
+ new TreeSet<>(),
new TreeSet<>()
);
Map<String, String> actualAliases = PluginUtils.computeAliases(result);
@@ -594,6 +597,7 @@ public class PluginUtilsTest {
new TreeSet<>(),
new TreeSet<>(),
new TreeSet<>(),
+ new TreeSet<>(),
new TreeSet<>()
);
Map<String, String> actualAliases = PluginUtils.computeAliases(result);
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/PluginsTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/PluginsTest.java
index 4b144288a54..ddc8d2773a4 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/PluginsTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/isolation/PluginsTest.java
@@ -115,6 +115,7 @@ public class PluginsTest {
new TreeSet<>(),
new TreeSet<>(),
new TreeSet<>(),
+ new TreeSet<>(),
new TreeSet<>()
);
empty = new PluginScanResult(List.of());
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 26a46658248..505b74a4446 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
@@ -869,7 +869,7 @@ public class ConnectorsResourceTest {
@Test
public void testAlterOffsetsEmptyOffsets() {
assertThrows(BadRequestException.class, () -> connectorsResource.alterConnectorOffsets(
- false, NULL_HEADERS, CONNECTOR_NAME, new ConnectorOffsets(List.of())));
+ false, NULL_HEADERS, CONNECTOR_NAME, new ConnectorOffsets(List.of()), false));
}
@Test
@@ -880,11 +880,11 @@ public class ConnectorsResourceTest {
ConnectorOffsets body = new ConnectorOffsets(List.of(connectorOffset));
final ArgumentCaptor<Callback<Message>> cb = ArgumentCaptor.forClass(Callback.class);
- expectAndCallbackNotLeaderException(cb).when(herder).alterConnectorOffsets(eq(CONNECTOR_NAME), eq(body.toMap()), cb.capture());
+ expectAndCallbackNotLeaderException(cb).when(herder).uberAlterConnectorOffsets(eq(CONNECTOR_NAME), eq(body.toMap()), eq(false), cb.capture());
when(restClient.httpRequest(eq(LEADER_URL + "connectors/" + CONNECTOR_NAME + "/offsets?forward=true"), eq("PATCH"), isNull(), eq(body), any()))
.thenReturn(new RestClient.HttpResponse<>(200, new HashMap<>(), new Message("")));
- connectorsResource.alterConnectorOffsets(null, NULL_HEADERS, CONNECTOR_NAME, body);
+ connectorsResource.alterConnectorOffsets(null, NULL_HEADERS, CONNECTOR_NAME, body, false);
}
@Test
@@ -895,9 +895,9 @@ public class ConnectorsResourceTest {
ConnectorOffsets body = new ConnectorOffsets(List.of(connectorOffset));
final ArgumentCaptor<Callback<Message>> cb = ArgumentCaptor.forClass(Callback.class);
expectAndCallbackException(cb, new NotFoundException("Connector not found"))
- .when(herder).alterConnectorOffsets(eq(CONNECTOR_NAME), eq(body.toMap()), cb.capture());
+ .when(herder).uberAlterConnectorOffsets(eq(CONNECTOR_NAME), eq(body.toMap()), eq(false), cb.capture());
- assertThrows(NotFoundException.class, () -> connectorsResource.alterConnectorOffsets(null, NULL_HEADERS, CONNECTOR_NAME, body));
+ assertThrows(NotFoundException.class, () -> connectorsResource.alterConnectorOffsets(null, NULL_HEADERS, CONNECTOR_NAME, body, false));
}
@Test
@@ -912,8 +912,8 @@ public class ConnectorsResourceTest {
doAnswer(invocation -> {
cb.getValue().onCompletion(null, msg);
return null;
- }).when(herder).alterConnectorOffsets(eq(CONNECTOR_NAME), eq(body.toMap()), cb.capture());
- Response response = connectorsResource.alterConnectorOffsets(null, NULL_HEADERS, CONNECTOR_NAME, body);
+ }).when(herder).uberAlterConnectorOffsets(eq(CONNECTOR_NAME), eq(body.toMap()), eq(false), cb.capture());
+ Response response = connectorsResource.alterConnectorOffsets(null, NULL_HEADERS, CONNECTOR_NAME, body, false);
assertEquals(200, response.getStatus());
assertEquals(msg, response.getEntity());
}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedConnect.java b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedConnect.java
index ddd7eab4e8b..da9ca69392c 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedConnect.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/clusters/EmbeddedConnect.java
@@ -717,7 +717,11 @@ abstract class EmbeddedConnect {
* @return the API response as a {@link java.lang.String}
*/
public String alterConnectorOffsets(String connectorName, ConnectorOffsets offsets) {
- String url = endpointForResource(String.format("connectors/%s/offsets", connectorName));
+ return alterConnectorOffsets(connectorName, offsets, false);
+ }
+
+ public String alterConnectorOffsets(String connectorName, ConnectorOffsets offsets, Boolean force) {
+ String url = endpointForResource(String.format("connectors/%s/offsets%s", connectorName, force != null ? "?uber-force=" + force : ""));
ObjectMapper mapper = new ObjectMapper();
String content;
try {
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/OffsetSnapshotClient.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/OffsetSnapshotClient.java
new file mode 100644
index 00000000000..0a9b3111a4c
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/OffsetSnapshotClient.java
@@ -0,0 +1,217 @@
+package com.uber.data.kafka.connect.aakafkaoffsetmgmt;
+
+import org.apache.kafka.common.TopicPartition;
+
+import com.alibaba.fastjson.JSON;
+import com.alibaba.fastjson.JSONArray;
+import com.alibaba.fastjson.JSONObject;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Charsets;
+
+import org.apache.commons.io.IOUtils;
+import org.apache.commons.lang3.tuple.Pair;
+import org.apache.http.client.ResponseHandler;
+import org.apache.http.client.config.RequestConfig;
+import org.apache.http.client.methods.CloseableHttpResponse;
+import org.apache.http.client.methods.HttpGet;
+import org.apache.http.client.methods.HttpPost;
+import org.apache.http.conn.HttpClientConnectionManager;
+import org.apache.http.entity.ByteArrayEntity;
+import org.apache.http.impl.client.CloseableHttpClient;
+import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.net.SocketTimeoutException;
+import java.util.Map;
+import java.util.Optional;
+
+import io.opentracing.contrib.apache.http.client.TracingHttpClientBuilder;
+import io.opentracing.util.GlobalTracer;
+
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+public class OffsetSnapshotClient implements Closeable {
+
+ private static final Logger log = LoggerFactory.getLogger(OffsetSnapshotClient.class);
+
+ private static final String getUrlFormatString = "http://127.0.0.1:18668/snapshot/%s/%s/%s/%d/%d?srcdst=src";
+ private static final int DEFAULT_MAX_CONNECTION_PER_ROUTE = 8;
+ private static final String DEFAULT_UBER_SOURCE_VALUE = "ureplicator2-worker";
+ private static final String X_UBER_SOURCE_HEADER = "x-uber-source";
+
+ private final String snapshotPostUrl;
+ private final CloseableHttpClient httpClient;
+ private final RequestConfig requestConfig;
+// protected final Meter snapshotReportCount = new Meter();
+// protected final Meter snapshotReportFailure = new Meter();
+// protected final Timer snapshotReportResponse = new Timer();
+// protected final Meter snapshotFetchCount = new Meter();
+// protected final Meter snapshotFetchFailure = new Meter();
+// protected final Timer snapshotFetchResponse = new Timer();
+
+ public OffsetSnapshotClient(String snapshotPostUrl, int requestTimeOutInMs) {
+ this(snapshotPostUrl, requestTimeOutInMs, buildDefaultConnectionManager());
+ }
+
+ private static HttpClientConnectionManager buildDefaultConnectionManager() {
+ PoolingHttpClientConnectionManager result = new PoolingHttpClientConnectionManager();
+ result.setDefaultMaxPerRoute(DEFAULT_MAX_CONNECTION_PER_ROUTE);
+ return result;
+ }
+
+ @VisibleForTesting
+ OffsetSnapshotClient(
+ String snapshotPostUrl,
+ int requestTimeOutInMs,
+ HttpClientConnectionManager connMgrForHttpClient
+ ) {
+ this.snapshotPostUrl = snapshotPostUrl;
+ // requestConfig is immutable. These three timeouts are for
+ // 1. getting connection from connection manager;
+ // 2. establishing connection with server;
+ // 3. getting next data snippet from server.
+ this.requestConfig = RequestConfig.custom()
+ .setConnectionRequestTimeout(requestTimeOutInMs)
+ .setConnectTimeout(requestTimeOutInMs)
+ .setSocketTimeout(requestTimeOutInMs)
+ .build();
+
+ TracingHttpClientBuilder builder = TracingHttpClientBuilder.create()
+ .withTracer(GlobalTracer.get());
+ this.httpClient = builder.setConnectionManager(connMgrForHttpClient).build();
+// KafkaUReplicatorMetricsReporter.get().registerMetric("snapshot.report-count", snapshotReportCount);
+// KafkaUReplicatorMetricsReporter.get().registerMetric("snapshot.report-failure", snapshotReportFailure);
+// KafkaUReplicatorMetricsReporter.get().registerMetric("snapshot.report-time", snapshotReportResponse);
+// KafkaUReplicatorMetricsReporter.get().registerMetric("snapshot.fetch-count", snapshotFetchCount);
+// KafkaUReplicatorMetricsReporter.get().registerMetric("snapshot.fetch-failure", snapshotFetchFailure);
+// KafkaUReplicatorMetricsReporter.get().registerMetric("snapshot.fetch-time", snapshotFetchResponse);
+ }
+
+ public Optional<Long> convertOffsetBySourceCluster(
+ TopicPartition topicPartition,
+ String srcCluster,
+ String dstCluster,
+ long offsetInSrcCluster
+ ) {
+// snapshotFetchCount.mark();
+ // Rule is that the last 3 params are all for the same cluster
+ String urlString = String.format(getUrlFormatString, topicPartition.topic(),
+ srcCluster, dstCluster, topicPartition.partition(), offsetInSrcCluster);
+ log.info("Converting offset via url {}", urlString);
+ HttpGet httpGet = new HttpGet(urlString);
+ httpGet.setConfig(requestConfig);
+ httpGet.addHeader(X_UBER_SOURCE_HEADER, DEFAULT_UBER_SOURCE_VALUE);
+// return snapshotFetchResponse.timeSupplier(() -> {
+ try {
+ CloseableHttpResponse resp = httpClient.execute(httpGet);
+ String retVal = IOUtils.toString(resp.getEntity().getContent(), UTF_8);
+ resp.close();
+ JSONObject parsed = JSON.parseObject(retVal);
+ JSONArray offsets = parsed.getJSONArray(String.valueOf(topicPartition.partition()));
+ if (offsets != null && !offsets.isEmpty()) {
+ Long offset = offsets.getLong(0);
+ return Optional.of(offset);
+ } else {
+ log.warn("Failed to fetch offset from response {}", retVal);
+ }
+ }catch(IOException ex) {
+// snapshotReportFailure.mark();
+ log.error("Failed to GET via {}", urlString, ex);
+ }
+ return Optional.empty();
+// });
+ }
+
+ private static boolean isSuccessStatus(final int status) {
+ // 2xx indicates success
+ return status >= 200 && status < 300;
+ }
+
+ public void reportSnapshot(
+ Map<String, PartitionOffsetMap> topicPartitionOffsetMap,
+ String srcClusterName,
+ String dstClusterName
+ ) {
+ long startInMs = System.currentTimeMillis();
+ log.info("Start to report offset snapshot at={}", startInMs);
+ HttpPost httpPost = createHttpPost(topicPartitionOffsetMap, srcClusterName, dstClusterName);
+ if (httpPost == null) {
+ log.info("Skipped posting offset because there is no data");
+ } else {
+// snapshotReportCount.mark();
+// snapshotReportResponse.time(() -> {
+ try {
+ int respStatus = httpClient.execute(httpPost, createResponseCodeExtractor());
+ if (isSuccessStatus(respStatus)) {
+ long endInMs = System.currentTimeMillis();
+ log.info("Report offset snapshot done at={} tookMs={}", endInMs, endInMs - startInMs);
+ } else {
+// snapshotReportFailure.mark();
+ log.info("Got error respStatus={} to post offset snapshot", respStatus);
+ }
+ } catch (SocketTimeoutException e) {
+// snapshotReportFailure.mark();
+ log.warn("Got timeout exception to post offset snapshot, {}", e.getMessage());
+ } catch (Throwable t) {
+// snapshotReportFailure.mark();
+ log.warn("Got unknown exception to post offset snapshot", t);
+ }
+// });
+ }
+ }
+
+ public void markSnapshotReportFailure() {
+// snapshotReportFailure.mark();
+ }
+
+ private HttpPost createHttpPost(
+ Map<String, PartitionOffsetMap> topicPartitionOffsetMap,
+ String srcClusterName,
+ String dstClusterName
+ ) {
+ if (topicPartitionOffsetMap.isEmpty()) {
+ return null;
+ }
+ JSONArray postBody = new JSONArray();
+ for (Map.Entry<String, PartitionOffsetMap> entry : topicPartitionOffsetMap.entrySet()) {
+ String topicName = entry.getKey();
+ Map<Pair<Integer, Integer>, Pair<Long, Long>> partitionOffsetMap = entry.getValue().getPartitionOffsetMap();
+ for (Map.Entry<Pair<Integer, Integer>, Pair<Long, Long>> mapping : partitionOffsetMap.entrySet()) {
+ JSONObject mappingInJson = new JSONObject();
+ mappingInJson.put("topicName", topicName);
+ mappingInJson.put("srcClusterName", srcClusterName);
+ mappingInJson.put("dstClusterName", dstClusterName);
+ mappingInJson.put("dstPartitionId", mapping.getKey().getLeft());
+ mappingInJson.put("dstOffset", mapping.getValue().getLeft());
+
+ JSONObject srcOffsetInJson = new JSONObject();
+ srcOffsetInJson.put("" + mapping.getKey().getRight(), mapping.getValue().getRight());
+ mappingInJson.put("srcPartitionOffsets", srcOffsetInJson);
+ postBody.add(mappingInJson);
+ }
+ }
+
+ // there is no strict synchronization, so few offset mapping might be dropped
+ // this is negligible, since we overwrite offset mappings during post interval anyway
+ topicPartitionOffsetMap.clear();
+
+ HttpPost httpPost = new HttpPost(snapshotPostUrl);
+ // since the post interval is pretty long, thus close conn to release resource
+ httpPost.addHeader("Connection", "close");
+ httpPost.addHeader("Content-Type", "application/json; charset=UTF-8");
+ httpPost.setEntity(new ByteArrayEntity(postBody.toString().getBytes(Charsets.UTF_8)));
+ httpPost.setConfig(requestConfig);
+ return httpPost;
+ }
+
+ private static ResponseHandler<Integer> createResponseCodeExtractor() {
+ return response -> response.getStatusLine().getStatusCode();
+ }
+
+ public void close() throws IOException {
+ httpClient.close();
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/OffsetSnapshotReporter.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/OffsetSnapshotReporter.java
new file mode 100644
index 00000000000..3f69529fb98
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/OffsetSnapshotReporter.java
@@ -0,0 +1,207 @@
+package com.uber.data.kafka.connect.aakafkaoffsetmgmt;
+
+import org.apache.kafka.common.TopicPartition;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.util.concurrent.ThreadFactoryBuilder;
+import com.uber.data.kafka.connect.utils.ExecutorUtils;
+import com.uber.jfx.core.UberEnvironment;
+import com.uber.m3.tally.RootScopeBuilder;
+import com.uber.m3.tally.Scope;
+import com.uber.m3.util.Duration;
+
+import org.apache.commons.lang3.tuple.Pair;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+/**
+ * Report offset snapshot to AAKafkaOffsetMgmt backend to persist the offset mapping information
+ */
+public class OffsetSnapshotReporter {
+
+ private static final Logger log = LoggerFactory.getLogger(OffsetSnapshotReporter.class);
+
+ private static final String PELOTON_INSTANCE_ID_KEY = "PELOTON_INSTANCE_ID";
+
+ private final String srcClusterName;
+ private final String dstClusterName;
+ private final Map<String, PartitionOffsetMap> topicPartitionOffsetMap = new ConcurrentHashMap<>();
+ private final ScheduledExecutorService reportExecutor;
+ private final AtomicBoolean isStarted = new AtomicBoolean(false);
+ private final int postIntervalInMs;
+ private final Map<TopicPartition, Long> lastCommittedSrcOffset = new ConcurrentHashMap<>();
+ private final OffsetTracker offsetTracker;
+ private final OffsetSnapshotClient snapshotClient;
+ private final Scope scope;
+ private final String routeId;
+
+ public OffsetSnapshotReporter(
+ String srcClusterName,
+ String dstClusterName,
+ int postIntervalInMs,
+ OffsetSnapshotClient snapshotClient,
+ String routeId
+ ) {
+ this(srcClusterName, dstClusterName, postIntervalInMs, snapshotClient, routeId, new OffsetTracker(), buildDefaultScope());
+ }
+
+ private static Scope buildDefaultScope() {
+ return new RootScopeBuilder().reportEvery(Duration.ofSeconds(10));
+ }
+
+ @VisibleForTesting
+ OffsetSnapshotReporter(
+ String srcClusterName,
+ String dstClusterName,
+ int postIntervalInMs,
+ OffsetSnapshotClient snapshotClient,
+ String routeId,
+ OffsetTracker offsetTracker,
+ Scope scope
+ ) {
+ this.srcClusterName = srcClusterName;
+ this.dstClusterName = dstClusterName;
+ this.postIntervalInMs = postIntervalInMs;
+ this.offsetTracker = offsetTracker;
+ this.snapshotClient = snapshotClient;
+
+ this.reportExecutor = Executors.newSingleThreadScheduledExecutor(
+ new ThreadFactoryBuilder().setNameFormat("offset-snapshot-reporter-%d").build());
+ Map<String, String> envVars = System.getenv();
+
+ @SuppressWarnings("deprecation")
+ String deployment = envVars.getOrDefault(UberEnvironment.DEPLOYMENT_KEY, "unknown");
+ String pelotonInstanceId = envVars.getOrDefault(PELOTON_INSTANCE_ID_KEY, "unknown");
+ this.scope = scope.tagged(ImmutableMap.of(
+ "deployment", deployment,
+ "peloton_instance_id", pelotonInstanceId));
+ this.routeId = routeId;
+ }
+
+ /**
+ * Starts to report offset
+ */
+ public void start() {
+ int delayInMs = ThreadLocalRandom.current().nextInt();
+ log.info("Start Offset Snapshot Reporter with srcClusterName={}, dstClusterName={}, "
+ + "delayInMs={} and postIntervalInMs={}", srcClusterName, dstClusterName, delayInMs,
+ postIntervalInMs);
+ reportExecutor.scheduleWithFixedDelay(createReportSnapshotTask(), delayInMs, postIntervalInMs, TimeUnit.MILLISECONDS);
+ isStarted.set(true);
+ }
+
+ public void record(String srcTopic, int srcPartition, long srcOffset, int dstPartition, long dstOffset) {
+ PartitionOffsetMap partitionOffsetMap = topicPartitionOffsetMap.computeIfAbsent(srcTopic, k -> new PartitionOffsetMap());
+ partitionOffsetMap.put(dstPartition, srcPartition, dstOffset, srcOffset);
+ }
+
+ void updateTopicPartitionOffset(TopicPartition topicPartition) {
+ long lastCommitted = lastCommittedSrcOffset.getOrDefault(topicPartition, -1L);
+ long lastOffsetCopied = offsetTracker.getHighestCopied().getOrDefault(topicPartition, -1L);
+ long lastSkipped = offsetTracker.getHighestSkipped().getOrDefault(topicPartition, -1L);
+ if (lastOffsetCopied == -1 && lastCommitted == -1L) {
+ log.info(
+ "TopicPartition {} has skipped range of {}->{}, last committed {}. Fetching from AA..",
+ topicPartition, lastOffsetCopied, lastSkipped, lastCommitted);
+
+ Optional<Long> reversedLookup = fetchSnapshotFromAA(topicPartition, lastSkipped);
+ if (reversedLookup.isPresent()) {
+ PartitionOffsetMap offsetMap = topicPartitionOffsetMap
+ .computeIfAbsent(topicPartition.topic(), k -> new PartitionOffsetMap());
+ log.info(
+ "Fetched from AA. Snapshot offset {}, {} for {}->{}. Faking snapshot of {}->{} ",
+ reversedLookup.get(), topicPartition, dstClusterName, srcClusterName, lastSkipped,
+ reversedLookup.get());
+ offsetMap.put(topicPartition.partition(), topicPartition.partition(), reversedLookup.get(),
+ lastSkipped);
+ } else {
+ log.info("Failed to fetch snapshot from AA {}", topicPartition);
+ }
+ }
+ }
+
+ private Runnable createReportSnapshotTask() {
+ return () -> {
+ try{
+ reportSnapshot();
+ scope.tagged(ImmutableMap.of("routeid", routeId))
+ .counter("computemapping")
+ .inc(1);
+ } catch (Exception e) {
+ log.error("Failed to report snapshot", e);
+ snapshotClient.markSnapshotReportFailure();
+ }
+ clearAfterReport();
+ };
+ }
+
+ @VisibleForTesting
+ void reportSnapshot() {
+ mergeToLastCommitedMap();
+ offsetTracker.getHighestSkipped().keySet().forEach(this::updateTopicPartitionOffset);
+ snapshotClient.reportSnapshot(topicPartitionOffsetMap, srcClusterName, dstClusterName);
+ }
+
+ @VisibleForTesting
+ void clearAfterReport() {
+ lastCommittedSrcOffset.clear();
+ topicPartitionOffsetMap.clear();
+ offsetTracker.clear();
+ }
+
+ @VisibleForTesting
+ Optional<Long> fetchSnapshotFromAA(TopicPartition topicPartition, long lastSkipped) {
+ return snapshotClient.convertOffsetBySourceCluster(topicPartition, dstClusterName, srcClusterName, lastSkipped);
+ }
+
+ @VisibleForTesting
+ void mergeToLastCommitedMap() {
+ topicPartitionOffsetMap.keySet().forEach( topic -> {
+ PartitionOffsetMap offsetMap = topicPartitionOffsetMap.get(topic);
+ Map<Pair<Integer, Integer>, Pair<Long, Long>> partitionOffsetMap = offsetMap.getPartitionOffsetMap();
+ for (Map.Entry<Pair<Integer, Integer>, Pair<Long, Long>> mapping : partitionOffsetMap.entrySet()) {
+ int paritionId = mapping.getKey().getRight();
+ long offset = mapping.getValue().getRight();
+ TopicPartition topicPartition = new TopicPartition(topic, paritionId);
+ lastCommittedSrcOffset.put(topicPartition, offset);
+ }
+ });
+ }
+
+ @VisibleForTesting
+ Optional<Long> getLastCommittedOffset(TopicPartition tp) {
+ return Optional.ofNullable(lastCommittedSrcOffset.get(tp));
+ }
+
+ /**
+ * Shuts down the reporter
+ */
+ public void shutdown() {
+ log.info("Shutdown Offset Snapshot Reporter");
+ if (isStarted.compareAndSet(true, false)) {
+ ExecutorUtils.shutdown(reportExecutor, "offset snapshot reporter thread", 1, TimeUnit.SECONDS);
+ try {
+ snapshotClient.close();
+ } catch (IOException ex) {
+ log.error("Failed to close snapshot client");
+ }
+ topicPartitionOffsetMap.clear();
+ }
+ }
+
+ public boolean isStarted() {
+ return isStarted.get();
+ }
+
+}
\ No newline at end of file
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/OffsetTracker.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/OffsetTracker.java
new file mode 100644
index 00000000000..aa25a1d562d
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/OffsetTracker.java
@@ -0,0 +1,32 @@
+package com.uber.data.kafka.connect.aakafkaoffsetmgmt;
+
+import org.apache.kafka.common.TopicPartition;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+public class OffsetTracker {
+ private final Map<TopicPartition, Long> highestOffsetCopied = new ConcurrentHashMap<>();
+ private final Map<TopicPartition, Long> highestOffsetSkipped = new ConcurrentHashMap<>();
+
+ public void trackSkipped(TopicPartition tp, long offset) {
+ highestOffsetSkipped.put(tp, offset);
+ }
+
+ public void trackCopied(TopicPartition tp, long offset) {
+ highestOffsetCopied.put(tp, offset);
+ }
+
+ public Map<TopicPartition, Long> getHighestSkipped() {
+ return highestOffsetSkipped;
+ }
+
+ public Map<TopicPartition, Long> getHighestCopied() {
+ return highestOffsetCopied;
+ }
+
+ public void clear() {
+ highestOffsetCopied.clear();
+ highestOffsetSkipped.clear();
+ }
+}
\ No newline at end of file
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/PartitionOffsetMap.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/PartitionOffsetMap.java
new file mode 100644
index 00000000000..674d4ab0024
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/PartitionOffsetMap.java
@@ -0,0 +1,28 @@
+package com.uber.data.kafka.connect.aakafkaoffsetmgmt;
+
+import com.google.common.collect.ImmutableMap;
+
+import org.apache.commons.lang3.tuple.Pair;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+// helper class to keep track of offset mapping for topic partitions
+public class PartitionOffsetMap {
+
+ // first Pair is <dst partition, src partition>, second Pair is <dst offset, source offset>
+ private final Map<Pair<Integer, Integer>, Pair<Long, Long>> partitionOffsetMap = new ConcurrentHashMap<>();
+
+ public void put(int dstPartition, int srcPartition, long dstOffset, long srcOffset) {
+ partitionOffsetMap.put(Pair.of(dstPartition, srcPartition), Pair.of(dstOffset, srcOffset));
+ }
+
+ public void removePartition(int partition) {
+ partitionOffsetMap.remove(Pair.of(partition, partition));
+ }
+
+ public Map<Pair<Integer, Integer>, Pair<Long, Long>> getPartitionOffsetMap() {
+ return ImmutableMap.copyOf(partitionOffsetMap);
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/common/ReplicatedTopicPartition.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/common/ReplicatedTopicPartition.java
new file mode 100644
index 00000000000..5913296a216
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/common/ReplicatedTopicPartition.java
@@ -0,0 +1,72 @@
+package com.uber.data.kafka.connect.common;
+
+import org.apache.kafka.common.TopicPartition;
+
+import java.util.Objects;
+
+/**
+ * POJO containing a topic partition and a source and destination Kafka cluster
+ */
+public class ReplicatedTopicPartition {
+
+ private final String srcKafka;
+ private final String dstKafka;
+ private final String topic;
+ private final int partition;
+
+ public ReplicatedTopicPartition(String srcKafka, String dstKafka, String topic, int partition) {
+ this.srcKafka = srcKafka;
+ this.dstKafka = dstKafka;
+ this.topic = topic;
+ this.partition = partition;
+ }
+
+ public ReplicatedTopicPartition(String srcKafka, String dstKafka, TopicPartition topicPartition) {
+ this(srcKafka, dstKafka, topicPartition.topic(), topicPartition.partition());
+ }
+
+ public String srcKafka() {
+ return srcKafka;
+ }
+
+ public String dstKafka() {
+ return dstKafka;
+ }
+
+ public String topic() {
+ return topic;
+ }
+
+ public int partition() {
+ return partition;
+ }
+
+ public TopicPartition topicPartition() {
+ return new TopicPartition(topic, partition);
+ }
+
+ public boolean matchesPipeline(String srcKafka, String dstKafka) {
+ return Objects.equals(srcKafka, this.srcKafka)
+ && Objects.equals(dstKafka, this.dstKafka);
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (o == null || getClass() != o.getClass()) return false;
+ ReplicatedTopicPartition that = (ReplicatedTopicPartition) o;
+ return partition == that.partition
+ && Objects.equals(srcKafka, that.srcKafka)
+ && Objects.equals(dstKafka, that.dstKafka)
+ && Objects.equals(topic, that.topic);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(srcKafka, dstKafka, topic, partition);
+ }
+
+ @Override
+ public String toString() {
+ return "@" + srcKafka + "@" + dstKafka + "@" + topicPartition();
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/AbstractAssignor.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/AbstractAssignor.java
new file mode 100644
index 00000000000..134c89ee758
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/AbstractAssignor.java
@@ -0,0 +1,375 @@
+package com.uber.data.kafka.connect.distributed.assignment;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.PriorityQueue;
+import java.util.Set;
+import java.util.function.Predicate;
+
+// Hey just FYI this class is absolutely not thread-safe and you're deranged if you try to use it that way
+public abstract class AbstractAssignor<G, J> {
+
+ private static final Logger log = LoggerFactory.getLogger(AbstractAssignor.class);
+
+ private final Comparator<G> groupComparator;
+ private final Comparator<J> jobComparator;
+ private final double groupLoadThreshold;
+ private final double defaultJobLoad;
+ private final int groupJobLimit;
+ private final Map<G, Double> cachedGroupLoads;
+ private final Map<J, Double> cachedJobLoads;
+
+ @SuppressWarnings("this-escape")
+ protected AbstractAssignor(double groupLoadThreshold, double defaultJobLoad, int groupJobLimit) {
+ this.groupComparator = Comparator.comparing(this::cachedGroupLoad);
+ this.jobComparator = Comparator.comparing(this::cachedJobLoad);
+ this.groupLoadThreshold = groupLoadThreshold;
+ this.defaultJobLoad = defaultJobLoad;
+ this.groupJobLimit = groupJobLimit;
+ this.cachedGroupLoads = new HashMap<>();
+ this.cachedJobLoads = new HashMap<>();
+ }
+
+ protected abstract Double jobLoad(J job) throws Exception;
+
+ protected abstract void assign(G group, J job);
+ protected abstract void revoke(G group, J job);
+ protected abstract Collection<J> jobs(G group);
+
+ protected abstract String groupType();
+ protected abstract String jobType();
+
+ // Can be overridden by subclasses if the lawless chaos of the English language demands it
+ protected String groupTypePlural() {
+ return groupType() + "(s)";
+ }
+
+ // Can be overridden by subclasses if the lawless chaos of the English language demands it
+ protected String jobTypePlural() {
+ return jobType() + "(s)";
+ }
+
+ protected String jobName(J job) {
+ return job.toString();
+ }
+
+ protected String groupName(G group) {
+ return group.toString();
+ }
+
+ // toAdd contains only new jobs that aren't already assigned and which should
+ // be assigned to a group if possible
+ public void assign(Collection<G> groups, Collection<J> toAdd) {
+ log.info("Rebalancing {} {} with {} new {} to assign", groups.size(), groupTypePlural(), toAdd.size(), jobTypePlural());
+
+ cachedGroupLoads.clear();
+ cachedJobLoads.clear();
+
+ Collection<J> toAssign = new HashSet<>(toAdd);
+ Collection<J> removed = enforceGroupJobLimit(groups);
+ toAssign.addAll(removed);
+
+ assignAll(groups, toAssign);
+
+ applyLoadThreshold(groups);
+
+ log.info("Finished rebalancing {} {} with {} new {} to assign", groups.size(), groupTypePlural(), toAdd.size(), jobTypePlural());
+ }
+
+ private Collection<J> enforceGroupJobLimit(Collection<G> groups) {
+ if (groupJobLimit <= 0) {
+ log.debug("Skipping limit enforcement for {} per {} since limit is non-positive", jobTypePlural(), groupType());
+ return Set.of();
+ }
+
+ log.debug("Enforcing limit for {} per {} across {} {}", jobTypePlural(), groupType(), groups.size(), groupTypePlural());
+
+ Collection<J> result = new ArrayList<>();
+ for (G group : groups) {
+ List<J> jobs = new ArrayList<>(jobs(group));
+ int numJobs = jobs.size();
+ int numRevoked = numJobs - groupJobLimit;
+ if (numRevoked <= 0) {
+ log.trace("{} {} has {} {}, which is within the limit of {}",
+ capitalize(groupType()), groupName(group), numJobs, jobTypePlural(), groupJobLimit
+ );
+ continue;
+ }
+
+ log.debug("{} {} has {} {}, which is above the limit of {}; will revoke {} {}",
+ capitalize(groupType()), groupName(group), numJobs, jobTypePlural(), groupJobLimit, numRevoked, jobTypePlural()
+ );
+ // Revoke smallest jobs in group
+ jobs.sort(jobComparator);
+ for (int i = 0; i < numRevoked; i++) {
+ J revoked = jobs.get(i);
+ log.trace("Revoking {} {} from {} {}", jobType(), jobName(revoked), groupType(), groupName(group));
+ revokeAndUpdateLoad(group, revoked);
+ result.add(revoked);
+ }
+ }
+
+ log.debug("Finished enforcing limit for {} per {} across {} {}", jobTypePlural(), groupType(), groups.size(), groupTypePlural());
+
+ return result;
+ }
+
+ private void assignAll(Collection<G> groups, Collection<J> toAssign) {
+ assert !groups.isEmpty();
+
+ if (toAssign.isEmpty()) {
+ log.debug("No new {} to assign", jobTypePlural());
+ return;
+ }
+
+ log.debug("Assigning {} new {} across {} {}", toAssign.size(), jobTypePlural(), groups.size(), groupTypePlural());
+
+ // Assign largest jobs first
+ List<J> sortedJobs = new ArrayList<>(toAssign);
+ sortedJobs.sort(jobComparator);
+
+ // Attempt to assign to least-loaded groups first, but only to groups
+ // that are under the job limit
+ PriorityQueue<G> eligibleGroups = new PriorityQueue<>(groupComparator);
+ groups.forEach(group -> {
+ if (beneathJobLimit(group)) {
+ eligibleGroups.add(group);
+ }
+ });
+
+ int assigned = 0;
+ int assignedPercent = 0;
+
+ for (J job : sortedJobs) {
+ G group = eligibleGroups.poll();
+
+ if (group == null) {
+ if (groupJobLimit > 0) {
+ // TODO: Metrics
+ log.warn("No {} found to accept {} {}; all {} may be above {} limit of {}",
+ groupTypePlural(), jobType(), jobName(job), groupTypePlural(), jobType(), groupJobLimit
+ );
+ } else {
+ log.warn(
+ "No {} found to accept {} {}; this should never happen",
+ groupTypePlural(), jobType(), jobName(job)
+ );
+ assert false; // Fail if we're in a testing environment
+ }
+ continue;
+ }
+
+ log.trace("Assigning {} {} to {} {}", jobType(), jobName(job), groupType(), groupName(group));
+ assignAndUpdateLoad(group, job);
+
+ if (beneathJobLimit(group)) {
+ boolean added = eligibleGroups.add(group);
+
+ if (!added) {
+ log.warn(
+ "Failed to re-insert {} {} into queue after assigning {} {}; this should never happen",
+ groupType(), groupName(group), jobType(), jobName(job)
+ );
+ assert false; // Fail if we're in a testing environment
+ }
+ } else {
+ log.debug("{} {} has now reached limit of {}; will not assign any more {} to it",
+ capitalize(groupType()), groupName(group), groupJobLimit, jobTypePlural()
+ );
+ }
+
+ int newAssignedPercent = ++assigned * 100 / toAssign.size();
+ if (newAssignedPercent != assignedPercent) {
+ assignedPercent = newAssignedPercent;
+ log.trace("Assigned {}% of total {} ({} / {})", assignedPercent, jobTypePlural(), assigned, toAssign.size());
+ }
+ }
+
+ log.debug("Finished assigning {} new {} across {} {}", toAssign.size(), jobTypePlural(), groups.size(), groupTypePlural());
+ }
+
+ // TODO: Dynamic load threshold based on ratio of least-loaded group to others?
+ // For example, if loads are 3, 5, 6, 7, 8, and 9, and the threshold is 2, we might
+ // still choose to do some balancing even though all groups are above the threshold;
+ // the heuristic we could use is to try to balance when groups have 1.5x, 2x, 3x, 5x, etc.
+ // the lowest load in the group; if the factor were 2x, we'd try to balance out the groups
+ // with loads 6, 7, 8, and 9; if the factor were 2.5x, we'd try to balance out the groups with
+ // loads 8 and 9; if it were 3x, we'd try to balance out the group with load 9, and if it were
+ // higher, we wouldn't try to do any balancing
+
+ private void applyLoadThreshold(Collection<G> groups) {
+ assert !groups.isEmpty();
+
+ log.debug("Applying load threshold across {} {}", groups.size(), groupTypePlural());
+
+ /*
+ - Load-balancing reassignment process:
+ - Sort groups by load
+ - Iterating from most-loaded to least-loaded group:
+ - While group is overloaded:
+ - Iterating from least-heavy to most-heavy job in group:
+ - Attempt to reassign job to least-loaded group that can accept job
+ without exceeding threshold
+ - If job cannot be reassigned without exceeding threshold, stop trying
+ to take load off of group (and move on to next overloaded group,
+ if there is one)
+ */
+ List<G> overloadedGroups = groups.stream()
+ .filter(this::isOverloaded)
+ .sorted(groupComparator.reversed())
+ .toList();
+
+ int numOverloaded = overloadedGroups.size();
+ log.debug("Found {} overloaded {}", numOverloaded, groupTypePlural());
+
+ if (overloadedGroups.isEmpty())
+ return;
+
+ PriorityQueue<G> eligibleGroups = new PriorityQueue<>(groupComparator);
+ groups.stream()
+ .filter(this::beneathJobLimit)
+ .filter(Predicate.not(this::isOverloaded))
+ .forEach(eligibleGroups::add);
+
+ log.debug("Found {} eligible {} to receive new {} for load balancing",
+ eligibleGroups.size(), groupTypePlural(), jobTypePlural()
+ );
+
+ for (G overloaded : overloadedGroups) {
+ List<J> sortedJobs = new ArrayList<>(jobs(overloaded));
+ sortedJobs.sort(jobComparator);
+
+ for (J job : sortedJobs) {
+ G leastLoaded = eligibleGroups.peek();
+ if (leastLoaded == null) {
+ return; // Nothing else we can do, every group is overloaded
+ }
+
+ if (!canAccept(leastLoaded, job)) {
+ log.trace("Cannot reassign {} {} to {} {} (would become overloaded)",
+ jobType(), jobName(job), groupType(), groupName(leastLoaded));
+ // Least-heavy job in this group cannot be assigned to the least-loaded group;
+ // nothing left we can do
+ break;
+ }
+
+ log.trace("Reassigning {} {} from {} {} to {} {}",
+ jobType(), jobName(job), groupType(), groupName(overloaded), groupType(), groupName(leastLoaded));
+ // Reassign from the overloaded group to the non-overloaded one
+ revokeAndUpdateLoad(overloaded, job);
+ assignAndUpdateLoad(leastLoaded, job);
+
+ if (beneathJobLimit(leastLoaded)) {
+ // Re-insert into the queue to reflect new workload
+ eligibleGroups.poll();
+ eligibleGroups.add(leastLoaded);
+ } else {
+ log.debug("{} {} is at or above {} limit of {}; will not assign any more {} to it",
+ capitalize(groupType()), groupName(leastLoaded), jobType(), groupJobLimit, jobTypePlural()
+ );
+ }
+
+ // Check to see if the group is still overloaded
+ if (!isOverloaded(overloaded)) {
+ numOverloaded--;
+ log.debug("{} {} is no longer overloaded", capitalize(groupType()), groupName(overloaded));
+ // If the group is no longer overloaded, we can stop trying to transfer
+ // jobs off of it, and add it to the non-overloaded queue
+ // so that it can potentially receive new jobs from other overloaded groups
+ if (beneathJobLimit(overloaded)) {
+ eligibleGroups.add(overloaded);
+ } else {
+ log.debug("{} {} is at or above {} limit of {}; will not assign any more {} to it",
+ capitalize(groupType()), groupName(overloaded), jobType(), groupJobLimit, jobTypePlural()
+ );
+ }
+ break;
+ } else {
+ log.trace("{} {} is still overloaded", capitalize(groupType()), groupName(overloaded));
+ }
+ }
+ }
+
+ log.debug("Finished applying load threshold across {} {}; there are {} overloaded {} remaining",
+ groups.size(), groupTypePlural(), numOverloaded, groupTypePlural()
+ );
+ }
+
+ private boolean canAccept(G group, J job) {
+ double newLoad = cachedGroupLoad(group) + cachedJobLoad(job);
+ return !isOverloaded(newLoad);
+ }
+
+ private boolean isOverloaded(G group) {
+ return isOverloaded(cachedGroupLoad(group));
+ }
+
+ private double cachedGroupLoad(G group) {
+ return cachedGroupLoads.computeIfAbsent(group, g ->
+ jobs(g).stream().mapToDouble(this::cachedJobLoad).sum()
+ );
+ }
+
+ private boolean isOverloaded(double load) {
+ return load > groupLoadThreshold;
+ }
+
+ private double cachedJobLoad(J job) {
+ return cachedJobLoads.computeIfAbsent(job, j -> {
+ Double load;
+ try {
+ load = jobLoad(job);
+ } catch (Exception e) {
+ log.error("Failed to retrieve load for {} {}; falling back on default load of {}", jobType(), jobName(job), defaultJobLoad);
+ return defaultJobLoad;
+ }
+
+ if (load == null) {
+ log.debug("No load available for {} {}; falling back on default load of {}", jobType(), jobName(job), defaultJobLoad);
+ return defaultJobLoad;
+ }
+
+ if (load < 0) {
+ log.warn("Negative load provided for {} {}; will fall back on default load of {}", jobType(), jobName(job), defaultJobLoad);
+ return defaultJobLoad;
+ }
+
+ log.trace("Found load of {} for {} {}", load, jobType(), jobName(job));
+
+ return load;
+ });
+ }
+
+ private void revokeAndUpdateLoad(G group, J job) {
+ revoke(group, job);
+ double newLoad = cachedGroupLoad(group) - cachedJobLoad(job);
+ cachedGroupLoads.put(group, newLoad);
+ }
+
+ private void assignAndUpdateLoad(G group, J job) {
+ assign(group, job);
+ double newLoad = cachedGroupLoad(group) + cachedJobLoad(job);
+ cachedGroupLoads.put(group, newLoad);
+ }
+
+ private int numJobs(G group) {
+ return jobs(group).size();
+ }
+
+ private boolean beneathJobLimit(G group) {
+ return groupJobLimit <= 0 || numJobs(group) < groupJobLimit;
+ }
+
+ private static String capitalize(String s) {
+ return (s.substring(0, 1).toUpperCase(Locale.ROOT)) + s.substring(1);
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/UberClusterAssignor.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/UberClusterAssignor.java
new file mode 100644
index 00000000000..ed99b72b646
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/UberClusterAssignor.java
@@ -0,0 +1,195 @@
+package com.uber.data.kafka.connect.distributed.assignment;
+
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.apache.kafka.connect.runtime.distributed.UberWorkerLoad;
+import org.apache.kafka.connect.util.ConnectUtils;
+
+import com.uber.data.kafka.connect.distributed.ClusterAssignment;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
+import com.uber.data.kafka.connect.distributed.ClusterConfigState;
+import com.uber.data.kafka.connect.distributed.ConnectorTaskId;
+import com.uber.data.kafka.connect.distributed.ConnectorsAndTasks;
+import com.uber.data.kafka.connect.workload.common.WorkloadProvider;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.function.Function;
+import java.util.stream.Collector;
+import java.util.stream.Collectors;
+
+import static org.apache.kafka.common.utils.Utils.closeQuietly;
+
+public class UberClusterAssignor implements ClusterAssignor {
+
+ private static final Logger log = LoggerFactory.getLogger(UberClusterAssignor.class);
+
+ private WorkloadProvider workloadProvider;
+ private double workerLoadThreshold;
+ private double defaultTaskLoad;
+ private int periodicRebalanceIntervalMs;
+ private int workerTaskLimit;
+
+ @Override
+ public void configure(Map<String, ?> props) {
+ UberClusterAssignorConfig config = new UberClusterAssignorConfig(props);
+ this.workloadProvider = config.workloadProvider();
+ this.workerLoadThreshold = config.uberWorkerLoadThreshold();
+ this.defaultTaskLoad = config.uberDefaultTaskLoad();
+ this.periodicRebalanceIntervalMs = config.uberPeriodicRebalanceIntervalMs();
+ this.workerTaskLimit = 10; // TODO: Don't even bother making this configurable, just put it into Flipr with a sane default
+ }
+
+ @Override
+ public String version() {
+ return AppInfoParser.getVersion();
+ }
+
+ @Override
+ public void close() throws Exception {
+ closeQuietly(workloadProvider, "Workload provider for Uber cluster assignor");
+ }
+
+ @Override
+ public ClusterAssignment assign(ClusterConfigState clusterConfigState, Map<String, ConnectorsAndTasks> currentAssignment) {
+ Set<String> configuredConnectors = new TreeSet<>(clusterConfigState.connectors());
+ Set<ConnectorTaskId> configuredTasks = combineCollections(configuredConnectors, clusterConfigState::tasks, Collectors.toSet());
+
+ // The connectors and tasks that should be running on the cluster, based on the contents of the
+ // config topic (which is the source of truth for the cluster's intended connectors and tasks)
+ ConnectorsAndTasks configured = ConnectorsAndTasks.of(configuredConnectors, configuredTasks);
+
+ // The connectors and tasks currently running on each member that's participating in this rebalance
+ ConnectorsAndTasks activeAssignments = assignment(currentAssignment);
+ log.debug("Active assignments: {}", activeAssignments);
+
+ // The complete set of connectors and tasks that should be newly-assigned during this round,
+ // either because they were just created, or the worker they were previously assigned to
+ // is no longer running
+ ConnectorsAndTasks toAssign = diff(configured, activeAssignments);
+ log.debug("Unassigned: {}", toAssign);
+
+ // The set of deleted connectors and tasks that members of the cluster are currently running
+ // and which should now be explicitly revoked
+ ConnectorsAndTasks deleted = diff(activeAssignments, configured);
+ log.debug("Deleted assignments: {}", deleted);
+
+ // The connectors and tasks that are currently running on more than one worker each
+ ConnectorsAndTasks duplicated = duplicatedAssignments(currentAssignment);
+ log.debug("Duplicated assignments: {}", duplicated);
+
+ // The connectors and tasks that need to be completely revoked from all workers
+ // (A follow-up rebalance is automatically triggered whenever a revocation takes place,
+ // during which the duplicated connectors and tasks revoked in this round can be safely
+ // reassigned)
+ ConnectorsAndTasks toRemove = deleted.toBuilder()
+ .add(duplicated)
+ .build();
+
+ Collection<UberWorkerLoad> workerLoads = workerLoads(currentAssignment);
+
+ revokeAll(workerLoads, toRemove);
+
+ int connectorsPerWorker = (int) Math.ceil(configured.connectors().size() / (double) workerLoads.size());
+ AbstractAssignor<UberWorkerLoad, String> connectorAssignor = new WorkerConnectorAssignor(connectorsPerWorker);
+ connectorAssignor.assign(workerLoads, toAssign.connectors());
+
+ AbstractAssignor<UberWorkerLoad, ConnectorTaskId> taskAssignor =
+ new WorkerTaskAssignor(workloadProvider, clusterConfigState, workerLoadThreshold, defaultTaskLoad, workerTaskLimit);
+ taskAssignor.assign(workerLoads, toAssign.tasks());
+
+ Map<String, ConnectorsAndTasks.Builder> resultBuilder = new HashMap<>();
+ workerLoads.forEach(worker ->
+ resultBuilder.computeIfAbsent(worker.worker(), w -> ConnectorsAndTasks.builder())
+ .addConnectors(worker.connectors())
+ .addTasks(worker.tasks())
+ );
+
+ Map<String, ConnectorsAndTasks> workerAssignments = transformValues(resultBuilder, ConnectorsAndTasks.Builder::build);
+ return new ClusterAssignment(workerAssignments, periodicRebalanceIntervalMs);
+ }
+
+ private ConnectorsAndTasks assignment(Map<String, ConnectorsAndTasks> memberAssignments) {
+ log.debug("Received assignments: {}", memberAssignments);
+ return ConnectorsAndTasks.of(
+ ConnectUtils.combineCollections(memberAssignments.values(), ConnectorsAndTasks::connectors),
+ ConnectUtils.combineCollections(memberAssignments.values(), ConnectorsAndTasks::tasks)
+ );
+ }
+
+ private ConnectorsAndTasks duplicatedAssignments(Map<String, ConnectorsAndTasks> memberAssignments) {
+ Set<String> duplicatedConnectors = duplicatedElements(memberAssignments, ConnectorsAndTasks::connectors);
+ Set<ConnectorTaskId> duplicatedTasks = duplicatedElements(memberAssignments, ConnectorsAndTasks::tasks);
+
+ return ConnectorsAndTasks.of(duplicatedConnectors, duplicatedTasks);
+ }
+
+ private static <E> Set<E> duplicatedElements(
+ Map<String, ConnectorsAndTasks> memberAssignments,
+ Function<ConnectorsAndTasks, Collection<E>> assignmentExtractor
+ ) {
+ Map<E, Long> elementCounts = combineCollections(
+ memberAssignments.values(),
+ assignmentExtractor,
+ Collectors.groupingBy(Function.identity(), Collectors.counting())
+ );
+ return elementCounts
+ .entrySet().stream()
+ .filter(entry -> entry.getValue() > 1L)
+ .map(Entry::getKey)
+ .collect(Collectors.toSet());
+ }
+
+ /**
+ * Revoke all of the specified connectors and tasks from all workers
+ * @param workerLoads the workers; may be empty, but not null
+ * @param toRevoke the connectors and tasks to revoke; may be empty, but not null
+ */
+ private void revokeAll(Collection<UberWorkerLoad> workerLoads, ConnectorsAndTasks toRevoke) {
+ workerLoads.forEach(w -> {
+ w.connectors().removeAll(toRevoke.connectors());
+ w.tasks().removeAll(toRevoke.tasks());
+ });
+ }
+
+ private static List<UberWorkerLoad> workerLoads(Map<String, ConnectorsAndTasks> memberAssignments) {
+ return memberAssignments.entrySet().stream()
+ .map(e -> new UberWorkerLoad(e.getKey(), e.getValue()))
+ .collect(Collectors.toList());
+ }
+
+ private static <K, I, O> Map<K, O> transformValues(Map<K, I> map, Function<I, O> transformation) {
+ return map.entrySet().stream().collect(Collectors.toMap(
+ Map.Entry::getKey,
+ transformation.compose(Map.Entry::getValue)
+ ));
+ }
+
+ private static ConnectorsAndTasks diff(ConnectorsAndTasks base, ConnectorsAndTasks... toSubtract) {
+ Set<String> connectors = new TreeSet<>(base.connectors());
+ Set<ConnectorTaskId> tasks = new TreeSet<>(base.tasks());
+ for (ConnectorsAndTasks sub : toSubtract) {
+ connectors.removeAll(sub.connectors());
+ tasks.removeAll(sub.tasks());
+ }
+ return new ConnectorsAndTasks(connectors, tasks);
+ }
+
+ private static <I, T, C> C combineCollections(
+ Collection<I> collection,
+ Function<I, Collection<T>> extractCollection,
+ Collector<T, ?, C> collector
+ ) {
+ return collection.stream()
+ .map(extractCollection)
+ .flatMap(Collection::stream)
+ .collect(collector);
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/UberClusterAssignorConfig.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/UberClusterAssignorConfig.java
new file mode 100644
index 00000000000..feecfde4dfe
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/UberClusterAssignorConfig.java
@@ -0,0 +1,98 @@
+package com.uber.data.kafka.connect.distributed.assignment;
+
+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 com.uber.data.kafka.connect.workload.common.WorkloadProvider;
+
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+
+import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
+
+public class UberClusterAssignorConfig extends AbstractConfig {
+
+ // TODO: Does this still need to be configurable?
+ public static final String UBER_WORKLOAD_PROVIDER_CONFIG = "uber.workload.provider";
+ public static final String UBER_WORKLOAD_PROVIDER_DOC = "Implementation of the "
+ + WorkloadProvider.class.getName() + " interface that will be used to determine the total workload "
+ + "of each worker. If left unset, then the default Kafka Connect rebalancing algorithm will be used instead, "
+ + "which does not take workload into account and weighs all tasks equally";
+
+ // TODO: This can be moved to Flipr
+ public static final String UBER_WORKER_LOAD_THRESHOLD_CONFIG = "uber.worker.load.threshold";
+ public static final double UBER_WORKER_LOAD_THRESHOLD_DEFAULT = 1;
+ public static final String UBER_WORKER_LOAD_THRESHOLD_DOC = "The maximum total load "
+ + "that a single worker should be assigned during rebalance. This will be enforced on a best-effort basis; "
+ + "in some cases (such as if the cluster is underprovisioned), this threshold will be breached instead of "
+ + "leaving tasks unassigned";
+
+ // TODO: Flipr
+ public static final String UBER_DEFAULT_TASK_LOAD_CONFIG = "uber.default.task.load";
+ public static final double UBER_DEFAULT_TASK_LOAD_DEFAULT = 0.1; // Allow up to 10 unknown tasks per worker
+ public static final String UBER_DEFAULT_TASK_LOAD_DOC = "The default load to assume "
+ + "for tasks that do not have a known load according to the configured workload provider";
+
+ // TODO: Flipr
+ public static final String UBER_PERIODIC_REBALANCE_INTERVAL_MS_CONFIG = "uber.periodic.rebalance.interval.ms";
+ public static final int UBER_PERIODIC_REBALANCE_INTERVAL_MS_DEFAULT = (int) TimeUnit.MINUTES.toMillis(5);
+ public static final String UBER_PERIODIC_REBALANCE_INTERVAL_MS_DOC = "How frequently to trigger a "
+ + "cluster rebalance, even if there has been no change in the number of workers, connectors, or tasks. "
+ + "This can be used to respond to changes in traffic and balance workloads across workers. "
+ + "Set to zero or a negative value to disable periodic rebalancing";
+
+ private final WorkloadProvider workloadProvider;
+
+ @SuppressWarnings("this-escape")
+ public UberClusterAssignorConfig(Map<?, ?> props) {
+ super(config(), props);
+ WorkloadProvider configuredWorkloadProvider = getConfiguredInstance(UBER_WORKLOAD_PROVIDER_CONFIG, WorkloadProvider.class);
+ this.workloadProvider = configuredWorkloadProvider != null ? configuredWorkloadProvider : WorkloadProvider.noOp();
+ }
+
+ public static ConfigDef config() {
+ return new ConfigDef()
+ .define(UBER_WORKLOAD_PROVIDER_CONFIG,
+ ConfigDef.Type.CLASS,
+ null, // Null instead of ConfigDef.NO_DEFAULT_VALUE to signify that the property is optional
+ ConfigDef.Importance.HIGH,
+ UBER_WORKLOAD_PROVIDER_DOC)
+ .define(UBER_WORKER_LOAD_THRESHOLD_CONFIG,
+ ConfigDef.Type.DOUBLE,
+ UBER_WORKER_LOAD_THRESHOLD_DEFAULT,
+ atLeast(0),
+ Importance.MEDIUM,
+ UBER_WORKER_LOAD_THRESHOLD_DOC)
+ .define(UBER_DEFAULT_TASK_LOAD_CONFIG,
+ ConfigDef.Type.DOUBLE,
+ UBER_DEFAULT_TASK_LOAD_DEFAULT,
+ atLeast(0),
+ Importance.MEDIUM,
+ UBER_DEFAULT_TASK_LOAD_DOC
+ ).define(UBER_PERIODIC_REBALANCE_INTERVAL_MS_CONFIG,
+ Type.INT,
+ UBER_PERIODIC_REBALANCE_INTERVAL_MS_DEFAULT,
+ Importance.MEDIUM,
+ UBER_PERIODIC_REBALANCE_INTERVAL_MS_DOC
+ );
+ }
+
+ public WorkloadProvider workloadProvider() {
+ return workloadProvider;
+ }
+
+ public double uberWorkerLoadThreshold() {
+ return getDouble(UBER_WORKER_LOAD_THRESHOLD_CONFIG);
+ }
+
+ public double uberDefaultTaskLoad() {
+ return getDouble(UBER_DEFAULT_TASK_LOAD_CONFIG);
+ }
+
+ public int uberPeriodicRebalanceIntervalMs() {
+ return getInt(UBER_PERIODIC_REBALANCE_INTERVAL_MS_CONFIG);
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/UberWorkerLoad.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/UberWorkerLoad.java
new file mode 100644
index 00000000000..fc081e64d55
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/UberWorkerLoad.java
@@ -0,0 +1,70 @@
+package com.uber.data.kafka.connect.distributed.assignment;
+
+import com.uber.data.kafka.connect.distributed.ConnectorTaskId;
+import com.uber.data.kafka.connect.distributed.ConnectorsAndTasks;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Objects;
+
+public class UberWorkerLoad {
+
+ private final String worker;
+ private final Collection<String> connectors;
+ private final Collection<ConnectorTaskId> tasks;
+
+ public UberWorkerLoad(String worker, ConnectorsAndTasks connectorsAndTasks) {
+ this(worker, connectorsAndTasks.connectors(), connectorsAndTasks.tasks());
+ }
+
+ public UberWorkerLoad(
+ String worker,
+ Collection<String> connectors,
+ Collection<ConnectorTaskId> tasks
+ ) {
+ this.worker = worker;
+ this.connectors = connectors != null ? new ArrayList<>(connectors) : new ArrayList<>();
+ this.tasks = tasks != null ? new ArrayList<>(tasks) : new ArrayList<>();
+ }
+
+ public String worker() {
+ return worker;
+ }
+
+ public Collection<String> connectors() {
+ return connectors;
+ }
+
+ public Collection<ConnectorTaskId> tasks() {
+ return tasks;
+ }
+
+ public void assign(String connector) {
+ connectors.add(connector);
+ }
+
+ public void assign(ConnectorTaskId task) {
+ tasks.add(task);
+ }
+
+ @Override
+ public String toString() {
+ return "UberWorkerLoad{ worker=" + worker + ", connectorIds=" + connectors + ", taskIds=" + tasks + '}';
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o) {
+ return true;
+ }
+ if (!(o instanceof UberWorkerLoad that)) {
+ return false;
+ }
+ return Objects.equals(this.worker, that.worker);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(worker);
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/WorkerConnectorAssignor.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/WorkerConnectorAssignor.java
new file mode 100644
index 00000000000..47b66658936
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/WorkerConnectorAssignor.java
@@ -0,0 +1,48 @@
+package com.uber.data.kafka.connect.distributed.assignment;
+
+import org.apache.kafka.connect.runtime.distributed.UberWorkerLoad;
+
+import java.util.Collection;
+
+public class WorkerConnectorAssignor extends AbstractAssignor<UberWorkerLoad, String> {
+
+ public WorkerConnectorAssignor(int connectorWorkerThreshold) {
+ // Try to have at most one connector per worker, and permit at most 10 at a time
+ super(connectorWorkerThreshold, 1, 10);
+ }
+
+ @Override
+ protected Double jobLoad(String job) {
+ return 1.0;
+ }
+
+ @Override
+ protected void assign(UberWorkerLoad worker, String connector) {
+ worker.assign(connector);
+ }
+
+ @Override
+ protected void revoke(UberWorkerLoad worker, String connector) {
+ worker.connectors().remove(connector);
+ }
+
+ @Override
+ protected Collection<String> jobs(UberWorkerLoad worker) {
+ return worker.connectors();
+ }
+
+ @Override
+ protected String groupType() {
+ return "worker";
+ }
+
+ @Override
+ protected String jobType() {
+ return "connector";
+ }
+
+ @Override
+ protected String groupName(UberWorkerLoad worker) {
+ return worker.worker();
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/WorkerTaskAssignor.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/WorkerTaskAssignor.java
new file mode 100644
index 00000000000..ad886b59894
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/WorkerTaskAssignor.java
@@ -0,0 +1,81 @@
+package com.uber.data.kafka.connect.distributed.assignment;
+
+import org.apache.kafka.connect.runtime.distributed.UberWorkerLoad;
+
+import com.uber.data.kafka.connect.distributed.ClusterConfigState;
+import com.uber.data.kafka.connect.distributed.ConnectorTaskId;
+import com.uber.data.kafka.connect.workload.common.WorkloadProvider;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Collection;
+import java.util.Map;
+
+public class WorkerTaskAssignor extends AbstractAssignor<UberWorkerLoad, ConnectorTaskId> {
+
+ private static final Logger log = LoggerFactory.getLogger(WorkerTaskAssignor.class);
+
+ private final WorkloadProvider workloadProvider;
+ private final ClusterConfigState clusterConfigState;
+
+ public WorkerTaskAssignor(
+ WorkloadProvider workloadProvider,
+ ClusterConfigState clusterConfigState,
+ double workerLoadThreshold,
+ double defaultTaskLoad,
+ int workerTaskLimit
+ ) {
+ super(workerLoadThreshold, defaultTaskLoad, workerTaskLimit);
+ this.workloadProvider = workloadProvider;
+ this.clusterConfigState = clusterConfigState;
+ }
+
+ @Override
+ protected Double jobLoad(ConnectorTaskId task) {
+ Map<String, String> taskConfig = clusterConfigState.taskConfig(task);
+ if (taskConfig == null) {
+ log.warn("Unable to find config for task {}", task);
+ return null;
+ }
+
+ String connector = task.connector();
+ Map<String, String> connectorConfig = clusterConfigState.connectorConfig(connector);
+ if (connectorConfig == null) {
+ log.warn("Unable to find config for connector {}", connector);
+ return null;
+ }
+
+ return workloadProvider.taskLoad(connector, task.task(), connectorConfig, taskConfig);
+ }
+
+ @Override
+ protected void assign(UberWorkerLoad worker, ConnectorTaskId task) {
+ worker.assign(task);
+ }
+
+ @Override
+ protected void revoke(UberWorkerLoad worker, ConnectorTaskId task) {
+ worker.tasks().remove(task);
+ }
+
+ @Override
+ protected Collection<ConnectorTaskId> jobs(UberWorkerLoad worker) {
+ return worker.tasks();
+ }
+
+ @Override
+ protected String groupType() {
+ return "worker";
+ }
+
+ @Override
+ protected String jobType() {
+ return "task";
+ }
+
+ @Override
+ protected String groupName(UberWorkerLoad worker) {
+ return worker.worker();
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/WorkloadProvider.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/WorkloadProvider.java
new file mode 100644
index 00000000000..8d22d41f2c9
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/WorkloadProvider.java
@@ -0,0 +1,39 @@
+package com.uber.data.kafka.connect.distributed.assignment;
+
+import org.apache.kafka.common.Configurable;
+import org.apache.kafka.connect.components.Versioned;
+
+import java.io.Closeable;
+import java.util.Map;
+
+public interface WorkloadProvider extends Versioned, Configurable, Closeable {
+
+ /**
+ * Determine and return the load for the given task, if it is known
+ * @return the load for the task, or null if none is known
+ */
+ Double taskLoad(String connector, int taskId, Map<String, String> connectorConfig, Map<String, String> taskConfig);
+
+ static com.uber.data.kafka.connect.workload.common.WorkloadProvider noOp() {
+ return new com.uber.data.kafka.connect.workload.common.WorkloadProvider() {
+ @Override
+ public Double taskLoad(String connector, int taskId, Map<String, String> connectorConfig, Map<String, String> taskConfig) {
+ return null;
+ }
+
+ @Override
+ public void close() {
+ }
+
+ @Override
+ public void configure(Map<String, ?> configs) {
+ }
+
+ @Override
+ public String version() {
+ return "no-op";
+ }
+ };
+ }
+
+}
\ No newline at end of file
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/UberConnectRestExtension.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/UberConnectRestExtension.java
new file mode 100644
index 00000000000..77abd774dda
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/UberConnectRestExtension.java
@@ -0,0 +1,60 @@
+package com.uber.data.kafka.connect.rest;
+
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.apache.kafka.connect.health.ConnectClusterState;
+import org.apache.kafka.connect.rest.ConnectRestExtension;
+import org.apache.kafka.connect.rest.ConnectRestExtensionContext;
+
+import com.uber.data.kafka.connect.rest.filter.ContentTypeHeaderFilter;
+import com.uber.data.kafka.connect.rest.interceptor.ConnectorConfigInterceptor;
+import com.uber.data.kafka.connect.rest.interceptor.ConnectorOffsetsInterceptor;
+import com.uber.data.kafka.connect.rest.resources.AssignmentsResource;
+import com.uber.data.kafka.connect.rest.resources.WorkloadsResource;
+
+import org.glassfish.jersey.internal.inject.AbstractBinder;
+
+import java.util.Map;
+
+import jakarta.ws.rs.core.Configurable;
+
+public class UberConnectRestExtension implements ConnectRestExtension {
+
+ @Override
+ public void configure(Map<String, ?> configs) {
+ }
+
+ @Override
+ public void register(ConnectRestExtensionContext restPluginContext) {
+ Configurable<?> configurable = restPluginContext.configurable();
+ configurable.register(ContentTypeHeaderFilter.class);
+ configurable.register(new Binder(restPluginContext.clusterState()));
+ configurable.register(ConnectorConfigInterceptor.class);
+ configurable.register(ConnectorOffsetsInterceptor.class);
+ configurable.register(WorkloadsResource.class);
+ configurable.register(AssignmentsResource.class);
+ }
+
+ @Override
+ public void close() {
+ }
+
+ @Override
+ public String version() {
+ return AppInfoParser.getVersion();
+ }
+
+ private static class Binder extends AbstractBinder {
+ private final ConnectClusterState clusterState;
+
+ public Binder(ConnectClusterState clusterState) {
+ this.clusterState = clusterState;
+ }
+
+ @Override
+ protected void configure() {
+ bind(clusterState).to(ConnectClusterState.class);
+ }
+ }
+
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/AssignedPartitionEntity.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/AssignedPartitionEntity.java
new file mode 100644
index 00000000000..859b6523b6a
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/AssignedPartitionEntity.java
@@ -0,0 +1,29 @@
+package com.uber.data.kafka.connect.rest.entities;
+
+import com.fasterxml.jackson.annotation.JsonInclude;
+import com.fasterxml.jackson.annotation.JsonInclude.Include;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.uber.data.kafka.connect.ureplicator3.coordination.AssignedPartition;
+
+@JsonInclude(Include.NON_NULL)
+public record AssignedPartitionEntity(
+ @JsonProperty String topic,
+ @JsonProperty int partition,
+ @JsonProperty String dstTopic,
+ @JsonProperty Integer dstPartition
+) {
+
+ public static AssignedPartitionEntity of(AssignedPartition assignedPartition) {
+ return new AssignedPartitionEntity(
+ assignedPartition.srcTopic(),
+ assignedPartition.srcPartition(),
+ assignedPartition.dstTopic(),
+ assignedPartition.dstPartition()
+ );
+ }
+
+ public AssignedPartition assignedPartition() {
+ return new AssignedPartition(topic, partition, dstTopic, dstPartition);
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/CachedWorkloadEntity.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/CachedWorkloadEntity.java
new file mode 100644
index 00000000000..9bb16ca8d34
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/CachedWorkloadEntity.java
@@ -0,0 +1,15 @@
+package com.uber.data.kafka.connect.rest.entities;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.uber.data.kafka.connect.workload.common.CachedWorkload;
+
+public record CachedWorkloadEntity(
+ @JsonProperty long lastUpdated,
+ @JsonProperty WorkloadEntity workload
+) {
+
+ public static CachedWorkloadEntity of(CachedWorkload cachedWorkload) {
+ WorkloadEntity workload = WorkloadEntity.of(cachedWorkload.workload());
+ return new CachedWorkloadEntity(cachedWorkload.lastUpdated(), workload);
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/PartitionWorkloadEntity.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/PartitionWorkloadEntity.java
new file mode 100644
index 00000000000..6dada99a8c2
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/PartitionWorkloadEntity.java
@@ -0,0 +1,59 @@
+package com.uber.data.kafka.connect.rest.entities;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
+import com.uber.data.kafka.connect.workload.common.Workload;
+
+import java.util.Comparator;
+
+// TODO: Property name overrides aren't being used in REST responses for some reason
+public record PartitionWorkloadEntity(
+ @JsonProperty("replicated_partition") ReplicatedTopicPartitionEntity replicatedPartition,
+ @JsonProperty("bytes_per_second") double bytesPerSecond,
+ @JsonProperty("messages_per_second") double messagesPerSecond
+) implements Comparable<PartitionWorkloadEntity> {
+
+ public static PartitionWorkloadEntity of(ReplicatedTopicPartition replicatedPartition, Workload workload) {
+ return new PartitionWorkloadEntity(
+ ReplicatedTopicPartitionEntity.of(replicatedPartition),
+ workload.getBytesPerSecond(),
+ workload.getMessagesPerSecond()
+ );
+ }
+
+ public static Comparator<PartitionWorkloadEntity> lexicographicComparator() {
+ return Comparator.<PartitionWorkloadEntity, String>comparing(w -> w.replicatedPartition().sourceKafka())
+ .thenComparing(w -> w.replicatedPartition().targetKafka())
+ .thenComparing(w -> w.replicatedPartition().topic())
+ .thenComparing(w -> w.replicatedPartition().partition())
+ .thenComparing(Comparator.naturalOrder());
+ }
+
+ public Workload workload() {
+ return new Workload(bytesPerSecond, messagesPerSecond);
+ }
+
+ public ReplicatedTopicPartition replicatedTopicPartition() {
+ return new ReplicatedTopicPartition(
+ replicatedPartition.sourceKafka(),
+ replicatedPartition.targetKafka(),
+ replicatedPartition.topic(),
+ replicatedPartition.partition()
+ );
+ }
+
+ @Override
+ public int compareTo(PartitionWorkloadEntity that) {
+ if (that == null)
+ return 1;
+
+ if (this == that)
+ return 0;
+
+ int result = Double.compare(this.bytesPerSecond, that.bytesPerSecond);
+ if (result != 0)
+ return result;
+
+ return Double.compare(this.messagesPerSecond, that.messagesPerSecond);
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/ReplicatedTopicPartitionEntity.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/ReplicatedTopicPartitionEntity.java
new file mode 100644
index 00000000000..d79ef6ca9de
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/ReplicatedTopicPartitionEntity.java
@@ -0,0 +1,22 @@
+package com.uber.data.kafka.connect.rest.entities;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
+
+public record ReplicatedTopicPartitionEntity(
+ @JsonProperty("src_kafka") String sourceKafka,
+ @JsonProperty("dst_kafka") String targetKafka,
+ @JsonProperty("topic") String topic,
+ @JsonProperty("partition") int partition
+) {
+
+ public static ReplicatedTopicPartitionEntity of(ReplicatedTopicPartition replicatedTopicPartition) {
+ return new ReplicatedTopicPartitionEntity(
+ replicatedTopicPartition.srcKafka(),
+ replicatedTopicPartition.dstKafka(),
+ replicatedTopicPartition.topic(),
+ replicatedTopicPartition.partition()
+ );
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/TaskAssignmentEntity.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/TaskAssignmentEntity.java
new file mode 100644
index 00000000000..e3d4288f1aa
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/TaskAssignmentEntity.java
@@ -0,0 +1,26 @@
+package com.uber.data.kafka.connect.rest.entities;
+
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.uber.data.kafka.connect.ureplicator3.coordination.AssignedPartition;
+
+import java.util.Collection;
+import java.util.List;
+
+import static com.uber.data.kafka.connect.utils.UberConnectUtils.compareLexicographically;
+
+public record TaskAssignmentEntity(
+ @JsonProperty String connector,
+ @JsonProperty int task,
+ @JsonProperty List<AssignedPartitionEntity> partitions
+) {
+
+ public static TaskAssignmentEntity of(String connector, int task, Collection<AssignedPartition> partitions) {
+ List<AssignedPartitionEntity> sortedPartitions = partitions.stream()
+ .sorted((ap1, ap2) -> compareLexicographically(ap1.source(), ap2.source()))
+ .map(AssignedPartitionEntity::of)
+ .toList();
+
+ return new TaskAssignmentEntity(connector, task, sortedPartitions);
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/TaskWorkloadEntity.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/TaskWorkloadEntity.java
new file mode 100644
index 00000000000..4a0d53b527f
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/TaskWorkloadEntity.java
@@ -0,0 +1,10 @@
+package com.uber.data.kafka.connect.rest.entities;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+public record TaskWorkloadEntity(
+ @JsonProperty String connector,
+ @JsonProperty int taskId,
+ @JsonProperty WorkloadEntity workload
+) {
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/TopicPartitionEntity.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/TopicPartitionEntity.java
new file mode 100644
index 00000000000..de42622d360
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/TopicPartitionEntity.java
@@ -0,0 +1,20 @@
+package com.uber.data.kafka.connect.rest.entities;
+
+import org.apache.kafka.common.TopicPartition;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+public record TopicPartitionEntity(
+ @JsonProperty String topic,
+ @JsonProperty int partition
+) {
+
+ public static TopicPartitionEntity of(TopicPartition topicPartition) {
+ return new TopicPartitionEntity(topicPartition.topic(), topicPartition.partition());
+ }
+
+ public TopicPartition topicPartition() {
+ return new TopicPartition(topic, partition);
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/WorkerWorkloadEntity.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/WorkerWorkloadEntity.java
new file mode 100644
index 00000000000..6aa4d0ce366
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/WorkerWorkloadEntity.java
@@ -0,0 +1,13 @@
+package com.uber.data.kafka.connect.rest.entities;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.List;
+
+
+public record WorkerWorkloadEntity(
+ @JsonProperty String worker,
+ @JsonProperty("total_workload") WorkloadEntity totalWorkload,
+ @JsonProperty("task_workloads") List<TaskWorkloadEntity> taskWorkloads
+) {
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/WorkloadEntity.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/WorkloadEntity.java
new file mode 100644
index 00000000000..b2396e187c6
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/WorkloadEntity.java
@@ -0,0 +1,18 @@
+package com.uber.data.kafka.connect.rest.entities;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.uber.data.kafka.connect.workload.common.Workload;
+
+public record WorkloadEntity(
+ @JsonProperty("bytes_per_second") double bytesPerSecond,
+ @JsonProperty("messages_per_second") double messagesPerSecond
+) {
+
+ public static WorkloadEntity of(Workload workload) {
+ return new WorkloadEntity(workload.getBytesPerSecond(), workload.getMessagesPerSecond());
+ }
+
+ public Workload workload() {
+ return new Workload(bytesPerSecond, messagesPerSecond);
+ }
+}
\ No newline at end of file
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/filter/ContentTypeHeaderFilter.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/filter/ContentTypeHeaderFilter.java
new file mode 100644
index 00000000000..e299c23f3ad
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/filter/ContentTypeHeaderFilter.java
@@ -0,0 +1,56 @@
+package com.uber.data.kafka.connect.rest.filter;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Set;
+
+import jakarta.ws.rs.container.ContainerRequestContext;
+import jakarta.ws.rs.container.ContainerRequestFilter;
+import jakarta.ws.rs.container.PreMatching;
+import jakarta.ws.rs.core.MediaType;
+
+import static jakarta.ws.rs.core.HttpHeaders.CONTENT_TYPE;
+import static jakarta.ws.rs.core.MediaType.APPLICATION_FORM_URLENCODED_TYPE;
+import static jakarta.ws.rs.core.MediaType.APPLICATION_JSON;
+
+@PreMatching
+public class ContentTypeHeaderFilter implements ContainerRequestFilter {
+
+ private static final Logger log = LoggerFactory.getLogger(ContentTypeHeaderFilter.class);
+
+ private static final String SKIP_PARAM = "keepContentType";
+
+ // Media types that we override unless explicitly instructed not to via the skip URL param
+ // TODO: We could technically make this configurable but it's probably not worth the effort
+ private static final Set<MediaType> STUPID_MEDIA_TYPES = Set.of(
+ APPLICATION_FORM_URLENCODED_TYPE // Automatically injected by curl (this fucking shit took way too long to debug)
+ );
+
+ @Override
+ public void filter(ContainerRequestContext containerRequestContext) {
+ if (containerRequestContext.getUriInfo().getQueryParameters().containsKey(SKIP_PARAM)) {
+ log.trace("Skipping {} header injection since {} URL param is specified", CONTENT_TYPE, SKIP_PARAM);
+ return;
+ }
+
+ if (!containerRequestContext.hasEntity()) {
+ log.trace("Skipping {} header injection since request has no body", CONTENT_TYPE);
+ return;
+ }
+
+ MediaType requestMediaType = containerRequestContext.getMediaType();
+ if (requestMediaType == null) {
+ log.trace("Request does not have media type; will inject {} header", CONTENT_TYPE);
+ } else if (STUPID_MEDIA_TYPES.contains(requestMediaType)) {
+ log.trace("Request has stupid media type {}; will override by injecting {} header", requestMediaType, CONTENT_TYPE);
+ } else {
+ log.trace("Skipping {} header injection since request already has non-stupid media type {}", CONTENT_TYPE, requestMediaType);
+ return;
+ }
+
+ log.debug("Adding {} header with value {} to request", CONTENT_TYPE, APPLICATION_JSON);
+ containerRequestContext.getHeaders().putSingle(CONTENT_TYPE, APPLICATION_JSON);
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/interceptor/ConnectorConfigInterceptor.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/interceptor/ConnectorConfigInterceptor.java
new file mode 100644
index 00000000000..eca79cfa970
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/interceptor/ConnectorConfigInterceptor.java
@@ -0,0 +1,235 @@
+package com.uber.data.kafka.connect.rest.interceptor;
+
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.connect.health.ConnectClusterState;
+import org.apache.kafka.connect.runtime.rest.entities.CreateConnectorRequest;
+import org.apache.kafka.connect.runtime.rest.errors.ConnectRestException;
+
+import com.uber.data.kafka.connect.utils.BrokerResolver;
+import com.uber.data.kafka.exception.ResolverException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.function.Function;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+import jakarta.inject.Inject;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.ws.rs.HttpMethod;
+import jakarta.ws.rs.WebApplicationException;
+import jakarta.ws.rs.container.ResourceInfo;
+import jakarta.ws.rs.core.Context;
+import jakarta.ws.rs.core.Response.Status;
+import jakarta.ws.rs.core.UriInfo;
+import jakarta.ws.rs.ext.ReaderInterceptor;
+import jakarta.ws.rs.ext.ReaderInterceptorContext;
+
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Config.CLIENTS_PREFIX;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Config.TARGET_CLIENTS_PREFIX;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Config.TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Config.TARGET_CLUSTER_CONFIG;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Config.TARGET_PRODUCER_PREFIX;
+import static com.uber.data.kafka.connect.utils.UberConnectUtils.isURep3;
+
+public class ConnectorConfigInterceptor implements ReaderInterceptor {
+
+ public static final String SKIP_QUERY_PARAM = "uberSkipInterception";
+
+ private static final Logger log = LoggerFactory.getLogger(ConnectorConfigInterceptor.class);
+
+ private static final String PRODUCER_OVERRIDE_PREFIX = "producer.override.";
+ private static final String PRODUCER_OVERRIDE_BOOTSTRAP_CONFIG =
+ PRODUCER_OVERRIDE_PREFIX + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG;
+
+ private final ConnectClusterState clusterState;
+
+ @Inject
+ public ConnectorConfigInterceptor(ConnectClusterState clusterState) {
+ this.clusterState = clusterState;
+ }
+
+ @Context
+ private UriInfo uriInfo;
+
+ @Context
+ private ResourceInfo resourceInfo;
+
+ @Context
+ private HttpServletRequest httpRequest;
+
+ @Override
+ public Object aroundReadFrom(ReaderInterceptorContext readerInterceptorContext) throws IOException, WebApplicationException {
+ Object entity = readerInterceptorContext.proceed();
+ if (shouldSkip()) {
+ return entity;
+ } else if (entity == null) {
+ log.debug("Request has no body; skipping interception");
+ return null;
+ }
+
+ if (isConnectorPostRequest()) {
+ if (entity instanceof CreateConnectorRequest createRequest) {
+ return interceptConnectorConfig(
+ createRequest.config(),
+ createRequest,
+ CreateConnectorRequest::config,
+ intercepted -> new CreateConnectorRequest(createRequest.name(), intercepted, createRequest.initialState())
+ );
+ } else {
+ log.error("Unable to intercept connector create request");
+ return entity;
+ }
+ } else if (isConnectorConfigPutRequest()) {
+ if (entity instanceof Map<?, ?> rawConfig) {
+ @SuppressWarnings("unchecked")
+ Map<String, String> originalConfig = (Map<String, String>) rawConfig;
+ return interceptConnectorConfig(
+ originalConfig,
+ originalConfig,
+ Function.identity(),
+ Function.identity()
+ );
+ } else {
+ log.error("Unable to intercept connector config put request");
+ return entity;
+ }
+ } else if (isConnectorConfigPatchRequest()) {
+ String path = uriInfo.getPath();
+ Matcher m = Pattern.compile("/*connectors/+([^/]+)/+config").matcher(path);
+ if (!m.matches()) {
+ log.error("Unable to intercept connector config patch request: unrecognized URL path");
+ return entity;
+ }
+ String connectorName = m.group(1);
+ if (connectorName == null) {
+ log.error("Unable to intercept connector config patch request: could not deduce connector name from URL path");
+ return entity;
+ }
+ // This is technically prone to race conditions since the connector config may be
+ // modified in between us reading it and when we allow the patch request to go through
+ // Changing a connector's type is highly unusual and doing it with the kind of frequency
+ // necessary to make that kind of race condition actually happen is not a usage pattern
+ // worth accommodating right now
+ Map<String, String> connectorConfig = clusterState.connectorConfig(connectorName);
+ if (connectorConfig == null) {
+ return entity;
+ }
+ if (entity instanceof Map<?, ?> rawPatch) {
+ @SuppressWarnings("unchecked")
+ Map<String, String> originalPatch = (Map<String, String>) rawPatch;
+ return interceptConnectorConfig(
+ connectorConfig,
+ originalPatch,
+ Function.identity(),
+ Function.identity()
+ );
+ } else {
+ log.error("Unable to intercept connector config patch request: unrecognized request body type {}", entity);
+ return entity;
+ }
+ } else {
+ return entity;
+ }
+ }
+
+ private boolean shouldSkip() {
+ if(httpRequest.getParameter(SKIP_QUERY_PARAM) != null) {
+ log.debug("Skipping interception (requested by {} URL query parameter)", SKIP_QUERY_PARAM);
+ return true;
+ }
+ return false;
+ }
+
+ private boolean isConnectorPostRequest() {
+ return isConnectorConfigRequest("createConnector", HttpMethod.POST);
+ }
+
+ private boolean isConnectorConfigPutRequest() {
+ return isConnectorConfigRequest("putConnectorConfig", HttpMethod.PUT);
+ }
+
+ private boolean isConnectorConfigPatchRequest() {
+ return isConnectorConfigRequest("patchConnectorConfig", HttpMethod.PATCH);
+ }
+
+ private boolean isConnectorConfigRequest(String resourceMethod, String httpMethod) {
+ return isConnectorRequest()
+ && httpMethod.equals(httpRequest.getMethod())
+ && resourceMethod.equals(resourceInfo.getResourceMethod().getName());
+ }
+
+ private boolean isConnectorRequest() {
+ return resourceInfo.getResourceClass().getSimpleName().equals("ConnectorsResource");
+ }
+
+ private <T> T interceptConnectorConfig(
+ Map<String, String> completeCurrentConnectorConfig,
+ T originalEntity,
+ Function<T, Map<String, String>> configExtractor,
+ Function<Map<String, String>, T> entityReconstructor
+ ) {
+ if (!isURep3(completeCurrentConnectorConfig)) {
+ log.debug("Ignoring config for non-uReplicator3 connector");
+ return originalEntity;
+ }
+
+ Map<String, String> originalConfig = configExtractor.apply(originalEntity);
+ Map<String, String> interceptedConfig = interceptConnectorConfig(originalConfig);
+ return entityReconstructor.apply(interceptedConfig);
+ }
+
+ private Map<String, String> interceptConnectorConfig(Map<String, String> connectorConfig) {
+ Map<String, String> result = new HashMap<>(connectorConfig);
+
+ String targetCluster = connectorConfig.get(TARGET_CLUSTER_CONFIG);
+ if (targetCluster != null) {
+ if (!connectorConfig.containsKey(PRODUCER_OVERRIDE_BOOTSTRAP_CONFIG)) {
+ String targetBootstrap = connectorConfig.get(TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG);
+ if (targetBootstrap == null)
+ targetBootstrap = resolveBootstrap(targetCluster);
+ result.put(PRODUCER_OVERRIDE_BOOTSTRAP_CONFIG, targetBootstrap);
+ }
+ }
+
+ translatePrefix(connectorConfig, result, CLIENTS_PREFIX, PRODUCER_OVERRIDE_PREFIX);
+ translatePrefix(connectorConfig, result, TARGET_CLIENTS_PREFIX, PRODUCER_OVERRIDE_PREFIX);
+ translatePrefix(connectorConfig, result, TARGET_PRODUCER_PREFIX, PRODUCER_OVERRIDE_PREFIX);
+
+ // TODO: Auto-inject default partition byte and message rates,
+ // and possibly max per-task byte and message rates
+
+ return result;
+ }
+
+ private void translatePrefix(Map<String, String> from, Map<String, String> to, String fromPrefix, String toPrefix) {
+ from.forEach((prefixedProp, value) -> {
+ if (prefixedProp.startsWith(fromPrefix)) {
+ String strippedProp = prefixedProp.substring(fromPrefix.length());
+ if (strippedProp.isBlank()) {
+ log.warn("Ignoring property {} with value {}", prefixedProp, value);
+ return;
+ }
+
+ to.put(toPrefix + strippedProp, value);
+ }
+ });
+ }
+
+ private String resolveBootstrap(String cluster) {
+ try {
+ return BrokerResolver.resolveBootstrap(cluster);
+ } catch (ResolverException e) {
+ throw new ConnectRestException(
+ Status.BAD_REQUEST,
+ "Failed to resolve bootstrap servers for cluster " + cluster,
+ e
+ );
+ }
+ }
+
+}
\ No newline at end of file
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/interceptor/ConnectorOffsetsInterceptor.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/interceptor/ConnectorOffsetsInterceptor.java
new file mode 100644
index 00000000000..721e38cf95b
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/interceptor/ConnectorOffsetsInterceptor.java
@@ -0,0 +1,147 @@
+package com.uber.data.kafka.connect.rest.interceptor;
+
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.config.ConfigDef.Type;
+import org.apache.kafka.connect.health.ConnectClusterState;
+import org.apache.kafka.connect.runtime.rest.entities.ConnectorOffsets;
+import org.apache.kafka.connect.runtime.rest.errors.BadRequestException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+import jakarta.inject.Inject;
+import jakarta.servlet.http.HttpServletRequest;
+import jakarta.ws.rs.HttpMethod;
+import jakarta.ws.rs.WebApplicationException;
+import jakarta.ws.rs.container.ResourceInfo;
+import jakarta.ws.rs.core.Context;
+import jakarta.ws.rs.core.UriInfo;
+import jakarta.ws.rs.ext.ReaderInterceptor;
+import jakarta.ws.rs.ext.ReaderInterceptorContext;
+
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.TOPICS_CONFIG;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Task.SOURCE_PARTITION_KAFKA_TOPIC;
+import static com.uber.data.kafka.connect.utils.UberConnectUtils.isURep3;
+import static org.apache.kafka.connect.runtime.rest.resources.ConnectorsResource.UBER_FORCE_OFFSET_PATCH;
+
+public class ConnectorOffsetsInterceptor implements ReaderInterceptor {
+
+ public static final String SKIP_QUERY_PARAM = "uberSkipInterception";
+
+ private static final Logger log = LoggerFactory.getLogger(ConnectorOffsetsInterceptor.class);
+
+ private final ConnectClusterState clusterState;
+
+ @Inject
+ public ConnectorOffsetsInterceptor(ConnectClusterState clusterState) {
+ this.clusterState = clusterState;
+ }
+
+ @Context
+ private UriInfo uriInfo;
+
+ @Context
+ private ResourceInfo resourceInfo;
+
+ @Context
+ private HttpServletRequest httpRequest;
+
+ @Override
+ public Object aroundReadFrom(ReaderInterceptorContext readerInterceptorContext) throws IOException, WebApplicationException {
+ Object entity = readerInterceptorContext.proceed();
+ if (shouldSkip()) {
+ return entity;
+ } else if (entity == null) {
+ log.debug("Request has no body; skipping interception");
+ return null;
+ } else if (!Boolean.parseBoolean(uriInfo.getQueryParameters().getFirst(UBER_FORCE_OFFSET_PATCH))) {
+ log.debug("Request will not be forced if connector is not stopped; skipping interception since no safeguard is necessary");
+ return entity;
+ }
+
+ if (!isConnectorOffsetPatchRequest()) {
+ return entity;
+ }
+
+ String path = uriInfo.getPath();
+ Matcher m = Pattern.compile("/*connectors/+([^/]+)/+offsets").matcher(path);
+ if (!m.matches()) {
+ log.error("Unable to intercept connector offset patch request: unrecognized URL path");
+ return entity;
+ }
+ String connectorName = m.group(1);
+ if (connectorName == null) {
+ log.error("Unable to intercept connector offset patch request: could not deduce connector name from URL path");
+ return entity;
+ }
+
+ Map<String, String> connectorConfig = clusterState.connectorConfig(connectorName);
+ if (connectorConfig == null) {
+ log.warn("Unable to locate config for connector {}; will not be able to safeguard offset patch request", connectorName);
+ return entity;
+ }
+
+ if (!isURep3(connectorConfig)) {
+ log.warn("Skipping interception for connector {} since it is not uReplicator3", connectorName);
+ return entity;
+ }
+
+ if (!(entity instanceof ConnectorOffsets offsets)) {
+ log.warn("Unrecognized entity type {} for connector {}; will not be able to safeguard offset patch request", entity.getClass(), connectorName);
+ return entity;
+ }
+
+ String topicsValue = connectorConfig.get(TOPICS_CONFIG);
+ @SuppressWarnings("unchecked")
+ List<String> topicsList = (List<String>) ConfigDef.parseType(TOPICS_CONFIG, topicsValue, Type.LIST);
+ if (topicsList == null) {
+ log.warn("Unable to parse topics list in config for connector {}; will not be able to safeguard offset patch request", connectorName);
+ return entity;
+ }
+
+ Set<String> invalidTopics = new TreeSet<>();
+ Set<String> topicsSet = new HashSet<>(topicsList); // We could also potentially fetch this from the assignment store, which might be more accurate
+ offsets.offsets().forEach(offset -> {
+ Object rawTopic = offset.partition().get(SOURCE_PARTITION_KAFKA_TOPIC);
+ // TODO: Technically this is too broad; we should also narrow down by src/dst Kafka cluster as well
+ // I'm too lazy to implement that right now and am currently betting it will not come back to bite me
+ // (if you are reading this I was probably wrong, sorry 🥲)
+ if (rawTopic instanceof String topic && topicsSet.contains(topic))
+ invalidTopics.add(topic);
+ });
+
+ if (!invalidTopics.isEmpty()) {
+ String message = "Cannot forcibly alter offsets for connector " + connectorName + " since it is still configured "
+ + "to replicate these topics: " + invalidTopics
+ + "\nEither stop the connector before modifying its offsets, or reconfigure the connector to no longer "
+ + "replicate from these topics before trying again";
+ throw new BadRequestException(message);
+ }
+
+ return entity;
+ }
+
+ private boolean shouldSkip() {
+ if(httpRequest.getParameter(SKIP_QUERY_PARAM) != null) {
+ log.debug("Skipping interception (requested by {} URL query parameter)", SKIP_QUERY_PARAM);
+ return true;
+ }
+ return false;
+ }
+
+ private boolean isConnectorOffsetPatchRequest() {
+ return "ConnectorsResource".equals(resourceInfo.getResourceClass().getSimpleName())
+ && HttpMethod.PATCH.equals(httpRequest.getMethod())
+ && "alterConnectorOffsets".equals(resourceInfo.getResourceMethod().getName());
+ }
+
+}
\ No newline at end of file
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/resources/AssignmentsResource.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/resources/AssignmentsResource.java
new file mode 100644
index 00000000000..c263daa3933
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/resources/AssignmentsResource.java
@@ -0,0 +1,85 @@
+package com.uber.data.kafka.connect.rest.resources;
+
+import org.apache.kafka.connect.health.ConnectClusterState;
+
+import com.uber.data.kafka.connect.rest.entities.TaskAssignmentEntity;
+import com.uber.data.kafka.connect.ureplicator3.coordination.AssignedPartition;
+import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStores;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.IntStream;
+
+import jakarta.inject.Inject;
+import jakarta.ws.rs.GET;
+import jakarta.ws.rs.Path;
+import jakarta.ws.rs.PathParam;
+import jakarta.ws.rs.Produces;
+import jakarta.ws.rs.core.MediaType;
+import jakarta.ws.rs.core.Response;
+import jakarta.ws.rs.core.Response.Status;
+
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Config.SOURCE_CLUSTER_CONFIG;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Config.TARGET_CLUSTER_CONFIG;
+
+@Path("/uber/assignments")
+@Produces(MediaType.APPLICATION_JSON)
+public class AssignmentsResource {
+
+ private static final Logger log = LoggerFactory.getLogger(AssignmentsResource.class);
+
+ private final ConnectClusterState clusterState;
+
+ @Inject
+ public AssignmentsResource(ConnectClusterState clusterState) {
+ this.clusterState = clusterState;
+ }
+
+ @Path("/{connector}/all")
+ @GET
+ public Response getConnectorAssignments(@PathParam("connector") String connector) {
+ Map<String, String> connectorConfig = clusterState.connectorConfig(connector);
+ if (connectorConfig == null) {
+ return Response.status(Status.NOT_FOUND)
+ .entity("Connector '" + connector + " does not appear to exist")
+ .build();
+ }
+
+ String srcKafka = connectorConfig.get(SOURCE_CLUSTER_CONFIG);
+ if (srcKafka == null) {
+ return Response.serverError()
+ .entity("Could not determine source cluster for connector '" + connector + "'")
+ .build();
+ }
+
+ String dstKafka = connectorConfig.get(TARGET_CLUSTER_CONFIG);
+ if (dstKafka == null) {
+ return Response.serverError()
+ .entity("Could not determine target cluster for connector '" + connector + "'")
+ .build();
+ }
+
+ Map<Integer, ? extends Collection<AssignedPartition>> taskAssignments =
+ AssignmentStores.get().getAssignments(connector, srcKafka, dstKafka);
+
+ int highestTask = Collections.max(taskAssignments.keySet());
+
+ List<TaskAssignmentEntity> body = IntStream.rangeClosed(0, highestTask)
+ .mapToObj(i -> {
+ Collection<AssignedPartition> assignment = taskAssignments.get(i);
+ if (assignment == null) {
+ assignment = Set.of();
+ }
+ return TaskAssignmentEntity.of(connector, i, assignment);
+ }).toList();
+
+ return Response.ok(body).build();
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/resources/WorkloadsResource.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/resources/WorkloadsResource.java
new file mode 100644
index 00000000000..28b1fd28123
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/resources/WorkloadsResource.java
@@ -0,0 +1,316 @@
+package com.uber.data.kafka.connect.rest.resources;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.connect.health.ConnectClusterState;
+import org.apache.kafka.connect.health.ConnectorHealth;
+import org.apache.kafka.connect.health.TaskState;
+import org.apache.kafka.connect.util.ConnectorTaskId;
+
+import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
+import com.uber.data.kafka.connect.rest.entities.CachedWorkloadEntity;
+import com.uber.data.kafka.connect.rest.entities.PartitionWorkloadEntity;
+import com.uber.data.kafka.connect.rest.entities.TaskWorkloadEntity;
+import com.uber.data.kafka.connect.rest.entities.WorkerWorkloadEntity;
+import com.uber.data.kafka.connect.rest.entities.WorkloadEntity;
+import com.uber.data.kafka.connect.workload.common.CachedWorkload;
+import com.uber.data.kafka.connect.workload.common.Workload;
+import com.uber.data.kafka.connect.workload.kafka.KafkaBasedWorkloadProvider;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.TreeMap;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import jakarta.inject.Inject;
+import jakarta.ws.rs.DELETE;
+import jakarta.ws.rs.DefaultValue;
+import jakarta.ws.rs.GET;
+import jakarta.ws.rs.POST;
+import jakarta.ws.rs.Path;
+import jakarta.ws.rs.PathParam;
+import jakarta.ws.rs.Produces;
+import jakarta.ws.rs.QueryParam;
+import jakarta.ws.rs.core.MediaType;
+import jakarta.ws.rs.core.Response;
+import jakarta.ws.rs.core.Response.Status;
+
+import static com.uber.data.kafka.connect.utils.UberConnectUtils.assignedPartitions;
+import static com.uber.data.kafka.connect.utils.UberConnectUtils.awaitAll;
+import static com.uber.data.kafka.connect.utils.UberConnectUtils.createErrorMessage;
+
+@Path("/uber/workloads")
+@Produces(MediaType.APPLICATION_JSON)
+public class WorkloadsResource {
+
+ private static final Logger log = LoggerFactory.getLogger(WorkloadsResource.class);
+
+ private final ConnectClusterState clusterState;
+
+ @Inject
+ public WorkloadsResource(ConnectClusterState clusterState) {
+ this.clusterState = clusterState;
+ }
+
+ @GET
+ @Path("/topics/all")
+ public List<PartitionWorkloadEntity> getAllTopicWorkloads(
+ @QueryParam("src") String srcKafka,
+ @QueryParam("dst") String dstKafka,
+ @QueryParam("topic") String topic,
+ @QueryParam("refresh") @DefaultValue("false") boolean refresh
+ ) {
+ if (refresh)
+ KafkaBasedWorkloadProvider.refreshWorkloadCache();
+
+ return KafkaBasedWorkloadProvider.getAllWorkloads().entrySet().stream()
+ .filter(e -> matches(e.getKey(), srcKafka, dstKafka, topic))
+ .map(e -> PartitionWorkloadEntity.of(e.getKey(), e.getValue()))
+ .sorted(PartitionWorkloadEntity.lexicographicComparator())
+ .collect(Collectors.toList());
+ }
+
+ @GET
+ @Path("/topics/historical/{src}/{dst}/{topic}/{partition}")
+ public Response getHistoricalWorkloads(
+ @PathParam("src") String srcKafka,
+ @PathParam("dst") String dstKafka,
+ @PathParam("topic") String topic,
+ @PathParam("partition") int partition,
+ @QueryParam("refresh") @DefaultValue("false") boolean refresh
+ ) {
+ if (refresh)
+ KafkaBasedWorkloadProvider.refreshWorkloadCache();
+
+ TopicPartition topicPartition = new TopicPartition(topic, partition);
+ List<CachedWorkload> cachedWorkloads = KafkaBasedWorkloadProvider.getHistoricalWorkloads(srcKafka, dstKafka, topicPartition);
+ if (cachedWorkloads == null)
+ return Response.status(Status.NOT_FOUND).entity("no workloads found for the requested topic").build();
+
+ List<CachedWorkloadEntity> body = cachedWorkloads.stream()
+ .map(CachedWorkloadEntity::of)
+ .toList();
+ return Response.ok(body).build();
+ }
+
+ @POST
+ @Path("/topics/overrides/{src}/{dst}/{topic}/{partition}")
+ public Response publishTopicWorkload(
+ @QueryParam("awaitAck") @DefaultValue("true") boolean awaitAck,
+ @PathParam("src") String srcKafka,
+ @PathParam("dst") String dstKafka,
+ @PathParam("topic") String topic,
+ @PathParam("partition") int partition,
+ WorkloadEntity workload
+ ) {
+ return publishWorkload(srcKafka, dstKafka, topic, partition, workload.workload(), awaitAck);
+ }
+
+ @DELETE
+ @Path("/topics/overrides/{src}/{dst}/{topic}")
+ public Response deleteTopicWorkloads(
+ @QueryParam("awaitAck") @DefaultValue("true") boolean awaitAck,
+ @QueryParam("partitions") int partitions,
+ @PathParam("src") String srcKafka,
+ @PathParam("dst") String dstKafka,
+ @PathParam("topic") String topic
+ ) {
+ List<? extends Future<?>> produceFutures = IntStream.range(0, partitions)
+ .mapToObj(partition -> KafkaBasedWorkloadProvider.publishWorkload(
+ srcKafka,
+ dstKafka,
+ new TopicPartition(topic, partition),
+ null
+ )).toList();
+
+ if (awaitAck) {
+ try {
+ awaitAll(produceFutures, Duration.ofMinutes(1));
+ } catch (InterruptedException e) {
+ return overrideErrorResponse(e);
+ } catch (TimeoutException e) {
+ return overrideErrorResponse(e);
+ } catch (ExecutionException e) {
+ return overrideErrorResponse(e);
+ }
+ } else {
+ log.debug("Not awaiting producer ack for workload override on {} partitions of topic {} in pipeline @{}@{}",
+ partitions, topic, srcKafka, dstKafka
+ );
+ }
+
+ return Response.noContent().build();
+ }
+
+ @DELETE
+ @Path("/topics/overrides/{src}/{dst}/{topic}/{partition}")
+ public Response deleteTopicWorkload(
+ @QueryParam("awaitAck") @DefaultValue("true") boolean awaitAck,
+ @PathParam("src") String srcKafka,
+ @PathParam("dst") String dstKafka,
+ @PathParam("topic") String topic,
+ @PathParam("partition") int partition
+ ) {
+ return publishWorkload(srcKafka, dstKafka, topic, partition, null, awaitAck);
+ }
+
+ private Response publishWorkload(
+ String srcKafka,
+ String dstKafka,
+ String topic,
+ int partition,
+ Workload workload,
+ boolean awaitAck
+ ) {
+ Future<?> produceFuture = KafkaBasedWorkloadProvider.publishWorkload(
+ srcKafka,
+ dstKafka,
+ new TopicPartition(topic, partition),
+ workload
+ );
+ if (awaitAck) {
+ try {
+ produceFuture.get(1, TimeUnit.MINUTES);
+ } catch (InterruptedException e) {
+ return overrideErrorResponse(e);
+ } catch (TimeoutException e) {
+ return overrideErrorResponse(e);
+ } catch (ExecutionException e) {
+ return overrideErrorResponse(e);
+ }
+ } else {
+ ReplicatedTopicPartition loggedPartition = new ReplicatedTopicPartition(srcKafka, dstKafka, topic, partition);
+ log.debug("Not awaiting producer ack for workload override on {}", loggedPartition);
+ }
+
+ return Response.noContent().build();
+ }
+
+ @GET
+ @Path("/workers/all")
+ public List<WorkerWorkloadEntity> getAllWorkerWorkloads(@QueryParam("refresh") @DefaultValue("false") boolean refresh) {
+ Collection<String> connectors = clusterState.connectors();
+ Map<String, Set<ConnectorTaskId>> workerTasks = new HashMap<>();
+
+ for (String connector : connectors) {
+ ConnectorHealth connectorHealth = clusterState.connectorHealth(connector);
+ if (connectorHealth == null) {
+ // The connector may have just been deleted
+ continue;
+ }
+ for (TaskState taskState : connectorHealth.tasksState().values()) {
+ workerTasks.computeIfAbsent(taskState.workerId(), k -> new HashSet<>())
+ .add(new ConnectorTaskId(connector, taskState.taskId()));
+ }
+ }
+
+ if (refresh)
+ KafkaBasedWorkloadProvider.refreshWorkloadCache();
+ Map<ReplicatedTopicPartition, Workload> workloads = KafkaBasedWorkloadProvider.getAllWorkloads();
+
+ Map<String, Map<ConnectorTaskId, Workload>> workerTaskWorkloads = new HashMap<>();
+ Map<String, Map<String, String>> connectorConfigs = new HashMap<>();
+ workerTasks.forEach((worker, tasks) ->
+ tasks.forEach(task -> {
+ Map<String, String> connectorConfig = connectorConfigs.computeIfAbsent(task.connector(), clusterState::connectorConfig);
+ if (connectorConfig == null) {
+ // The connector may have just been deleted
+ return;
+ }
+
+ Set<ReplicatedTopicPartition> taskPartitions = assignedPartitions(task.connector(), task.task(), connectorConfig);
+
+ if (taskPartitions == null) {
+ return;
+ }
+
+ for (ReplicatedTopicPartition replicatedTopicPartition : taskPartitions) {
+ Workload workload = workloads.get(replicatedTopicPartition);
+ if (workload == null) {
+ continue;
+ }
+ workerTaskWorkloads.computeIfAbsent(
+ worker,
+ k -> new TreeMap<>(Comparator.comparing(ConnectorTaskId::connector).thenComparing(ConnectorTaskId::task))
+ ).compute(
+ task,
+ (k, existingWorkload) -> workload.plus(existingWorkload)
+ );
+ }
+ })
+ );
+
+ return workerTaskWorkloads.entrySet().stream()
+ .sorted(Map.Entry.comparingByKey())
+ .map(e -> {
+ String worker = e.getKey();
+ Map<ConnectorTaskId, Workload> taskWorkloads = e.getValue();
+
+ List<TaskWorkloadEntity> taskWorkloadEntities = new ArrayList<>();
+ double totalByteRate = 0;
+ double totalMessageRate = 0;
+
+ // TODO: We might want to order this somehow (probably lexicographically by connector name and then task ID)
+ for (Map.Entry<ConnectorTaskId, Workload> taskWorkloadEntry : taskWorkloads.entrySet()) {
+ totalByteRate += taskWorkloadEntry.getValue().getBytesPerSecond();
+ totalMessageRate += taskWorkloadEntry.getValue().getMessagesPerSecond();;
+
+ TaskWorkloadEntity taskWorkloadEntity = new TaskWorkloadEntity(
+ taskWorkloadEntry.getKey().connector(),
+ taskWorkloadEntry.getKey().task(),
+ WorkloadEntity.of(taskWorkloadEntry.getValue())
+ );
+ taskWorkloadEntities.add(taskWorkloadEntity);
+ }
+
+ WorkloadEntity totalWorkload = new WorkloadEntity(totalByteRate, totalMessageRate);
+ return new WorkerWorkloadEntity(worker, totalWorkload, taskWorkloadEntities);
+ }).toList();
+ }
+
+ private boolean matches(ReplicatedTopicPartition rtp, String srcKafka, String dstKafka, String topic) {
+ if (srcKafka != null && !srcKafka.equals(rtp.srcKafka()))
+ return false;
+
+ if (dstKafka != null && !dstKafka.equals(rtp.dstKafka()))
+ return false;
+
+ if (topic != null && !topic.equals(rtp.topic()))
+ return false;
+
+ return true;
+ }
+
+ private Response overrideErrorResponse(InterruptedException e) {
+ return Response.serverError()
+ .entity(createErrorMessage("Interrupted while waiting for producer ack", e))
+ .build();
+ }
+
+ private Response overrideErrorResponse(TimeoutException e) {
+ return Response.serverError()
+ .entity(createErrorMessage("Timed out waiting for producer ack", e))
+ .build();
+ }
+
+ private Response overrideErrorResponse(ExecutionException e) {
+ return Response.serverError()
+ .entity(createErrorMessage("Failed to publish workload override", e))
+ .build();
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/TaskPartitionAssignor.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/TaskPartitionAssignor.java
new file mode 100644
index 00000000000..d2a51bc60b8
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/TaskPartitionAssignor.java
@@ -0,0 +1,112 @@
+package com.uber.data.kafka.connect.ureplicator3;
+
+import org.apache.kafka.common.TopicPartition;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
+import com.uber.data.kafka.connect.distributed.assignment.AbstractAssignor;
+import com.uber.data.kafka.connect.workload.common.Workload;
+import com.uber.data.kafka.connect.workload.kafka.KafkaBasedWorkloadProvider;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Collection;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+public class TaskPartitionAssignor extends AbstractAssignor<TaskPartitions, TopicPartition> {
+
+ private static final Logger log = LoggerFactory.getLogger(TaskPartitionAssignor.class);
+
+ private final Map<TopicPartition, Double> partitionLoads;
+
+ public TaskPartitionAssignor(
+ String srcKafka,
+ String dstKafka,
+ double defaultPartitionLoad,
+ double taskByteRateThreshold,
+ double taskMessageRateThreshold
+ ) {
+ // TODO: Should we have a limit of partitions per task?
+ this(
+ defaultPartitionLoad,
+ retrievePartitionLoads(srcKafka, dstKafka, taskByteRateThreshold, taskMessageRateThreshold)
+ );
+ }
+
+ @VisibleForTesting
+ TaskPartitionAssignor(
+ double defaultPartitionLoad,
+ Map<TopicPartition, Double> partitionLoads
+ ) {
+ super(1, defaultPartitionLoad, 0);
+ this.partitionLoads = partitionLoads;
+ }
+
+ @Override
+ protected Double jobLoad(TopicPartition topicPartition) {
+ return partitionLoads.get(topicPartition);
+ }
+
+ @Override
+ protected void assign(TaskPartitions task, TopicPartition topicPartition) {
+ task.partitions().add(topicPartition);
+ }
+
+ @Override
+ protected void revoke(TaskPartitions task, TopicPartition topicPartition) {
+ task.partitions().remove(topicPartition);
+ }
+
+ @Override
+ protected Collection<TopicPartition> jobs(TaskPartitions task) {
+ return task.partitions();
+ }
+
+ @Override
+ protected String groupType() {
+ return "task";
+ }
+
+ @Override
+ protected String jobType() {
+ return "partition";
+ }
+
+ private static Map<TopicPartition, Double> retrievePartitionLoads(
+ String srcKafka,
+ String dstKafka,
+ double byteRateThreshold,
+ double messageRateThreshold
+ ) {
+ log.debug("Retrieving partition loads in preparation for task partition assignment");
+
+ Map<ReplicatedTopicPartition, Workload> allWorkloads;
+ try {
+ allWorkloads = KafkaBasedWorkloadProvider.getAllWorkloads();
+ } catch (IllegalStateException e) {
+ log.warn(
+ "Worker does not appear to be configured with Kafka-based workload provider; "
+ + "will not be able to balance task assignment based on partition workloads"
+ );
+ return Map.of();
+ }
+
+ Map<TopicPartition, Double> result = allWorkloads.entrySet().stream()
+ .filter(e -> srcKafka.equals(e.getKey().srcKafka()))
+ .filter(e -> dstKafka.equals(e.getKey().dstKafka()))
+ .collect(Collectors.toMap(
+ e -> e.getKey().topicPartition(),
+ e -> {
+ Workload workload = e.getValue();
+ double byteRateScale = workload.getBytesPerSecond() / byteRateThreshold;
+ double messageRateScale = workload.getMessagesPerSecond() / messageRateThreshold;
+ return Math.max(byteRateScale, messageRateScale);
+ }
+ ));
+
+ log.debug("Finished retrieving partition loads in preparation for task partition assignment");
+ return result;
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/TaskPartitions.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/TaskPartitions.java
new file mode 100644
index 00000000000..4935a907615
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/TaskPartitions.java
@@ -0,0 +1,25 @@
+package com.uber.data.kafka.connect.ureplicator3;
+
+import org.apache.kafka.common.TopicPartition;
+
+import java.util.HashSet;
+import java.util.Objects;
+import java.util.Set;
+
+public record TaskPartitions(String connectorName, int taskId, Set<TopicPartition> partitions) {
+
+ @Override
+ public String toString() {
+ return connectorName + "-" + taskId;
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(connectorName, taskId);
+ }
+
+ public TaskPartitions copy() {
+ return new TaskPartitions(connectorName, taskId, new HashSet<>(partitions));
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Config.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Config.java
new file mode 100644
index 00000000000..14a415e68c7
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Config.java
@@ -0,0 +1,251 @@
+package com.uber.data.kafka.connect.ureplicator3;
+
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+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.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.serialization.ByteArrayDeserializer;
+import org.apache.kafka.connect.runtime.ConnectorConfig;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.uber.data.kafka.connect.utils.BrokerResolver;
+import com.uber.data.kafka.exception.ResolverException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.time.Duration;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+
+import static com.uber.data.kafka.connect.utils.UberConnectUtils.createErrorMessage;
+import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
+
+public abstract class UReplicator3Config extends AbstractConfig {
+
+ private static final Logger log = LoggerFactory.getLogger(UReplicator3Config.class);
+
+ public static final String SOURCE_CLUSTER_CONFIG = "source.cluster";
+ private static final String SOURCE_CLUSTER_DOC = "Name of source cluster to consume from (e.g., phx2-dev1)";
+
+ public static final String TARGET_CLUSTER_CONFIG = "target.cluster";
+ private static final String TARGET_CLUSTER_DOC = "Name of target cluster to produce to (e.g., dca1-dev1)";
+
+ public static final String CONSUMER_POLL_TIMEOUT_MS_CONFIG = "consumer.poll.timeout.ms";
+ public static final long CONSUMER_POLL_TIMEOUT_MS_DEFAULT = 1000;
+ private static final String CONSUMER_POLL_TIMEOUT_MS_DOC = "The timeout to use when invoking Consumer::poll";
+
+ public static final String CLIENTS_PREFIX = "all.kafka.clients.";
+ public static final String SOURCE_CLIENTS_PREFIX = "source.kafka.clients.";
+ public static final String TARGET_CLIENTS_PREFIX = "target.kafka.clients.";
+ public static final String SOURCE_CONSUMER_PREFIX = "source.consumer.";
+ // This prefix controls "phantom" properties: ones not actually used by our connector code
+ // but instead, picked up by our REST extension and translated into properties recognized by
+ // the Connect runtime (i.e., "producer.override.<property>")
+ public static final String TARGET_PRODUCER_PREFIX = "target.producer.";
+ public static final String SOURCE_ADMIN_PREFIX = "source.admin.";
+ public static final String TARGET_ADMIN_PREFIX = "target.admin.";
+
+ public static final String WORKLOAD_REPORT_INTERVAL_MS_CONFIG = "workload.report.interval.ms";
+ private static final long WORKLOAD_REPORT_INTERVAL_MS_DEFAULT = TimeUnit.MINUTES.toMillis(10);
+
+ public static final String OFFSET_REPORTER_ENABLED_CONFIG = "offset.reporter.enabled";
+ private static final boolean OFFSET_REPORTER_ENABLED_DEFAULT = true;
+
+ public static final String SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG = SOURCE_CLIENTS_PREFIX + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG;
+ public static final String TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG = TARGET_CLIENTS_PREFIX + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG;
+
+ @SuppressWarnings("this-escape")
+ protected UReplicator3Config(ConfigDef config, Map<String, ?> props) {
+ super(config, props);
+ }
+
+ protected static ConfigDef baseConfig() {
+ return new ConfigDef()
+ .define(
+ SOURCE_CLUSTER_CONFIG,
+ Type.STRING,
+ ConfigDef.NO_DEFAULT_VALUE,
+ Importance.HIGH,
+ SOURCE_CLUSTER_DOC
+ ).define(
+ TARGET_CLUSTER_CONFIG,
+ Type.STRING,
+ ConfigDef.NO_DEFAULT_VALUE,
+ Importance.HIGH,
+ TARGET_CLUSTER_DOC
+ ).define(
+ CONSUMER_POLL_TIMEOUT_MS_CONFIG,
+ Type.LONG,
+ CONSUMER_POLL_TIMEOUT_MS_DEFAULT,
+ atLeast(1),
+ Importance.MEDIUM,
+ CONSUMER_POLL_TIMEOUT_MS_DOC
+ ).defineInternal(
+ WORKLOAD_REPORT_INTERVAL_MS_CONFIG,
+ Type.LONG,
+ WORKLOAD_REPORT_INTERVAL_MS_DEFAULT,
+ Importance.LOW
+ ).defineInternal(
+ OFFSET_REPORTER_ENABLED_CONFIG,
+ Type.BOOLEAN,
+ OFFSET_REPORTER_ENABLED_DEFAULT,
+ Importance.LOW
+ ).defineInternal(
+ ConnectorConfig.NAME_CONFIG,
+ Type.STRING,
+ ConfigDef.NO_DEFAULT_VALUE,
+ Importance.LOW
+ );
+ }
+
+ protected abstract String role();
+
+ public String sourceCluster() {
+ return getString(SOURCE_CLUSTER_CONFIG);
+ }
+
+ public String targetCluster() {
+ return getString(TARGET_CLUSTER_CONFIG);
+ }
+
+ public Duration consumerPollTimeout() {
+ long timeoutMs = getLong(CONSUMER_POLL_TIMEOUT_MS_CONFIG);
+ return Duration.ofMillis(timeoutMs);
+ }
+
+ public String connectorName() {
+ return getString(ConnectorConfig.NAME_CONFIG);
+ }
+
+ public Duration workloadReportInterval() {
+ Long durationMs = getLong(WORKLOAD_REPORT_INTERVAL_MS_CONFIG);
+ return durationMs == null || durationMs <= 0
+ ? null
+ : Duration.ofMillis(durationMs);
+ }
+
+ public boolean offsetReporterEnabled() {
+ return getBoolean(OFFSET_REPORTER_ENABLED_CONFIG);
+ }
+
+ public Admin sourceAdmin() {
+ return admin(
+ sourceBootstrap(),
+ clientId(true, "admin"),
+ CLIENTS_PREFIX,
+ SOURCE_CLIENTS_PREFIX,
+ SOURCE_ADMIN_PREFIX
+ );
+ }
+
+ public Admin targetAdmin() {
+ return admin(
+ targetBootstrap(),
+ clientId(false, "admin"),
+ CLIENTS_PREFIX,
+ TARGET_CLIENTS_PREFIX,
+ TARGET_ADMIN_PREFIX
+ );
+ }
+
+ public Consumer<byte[], byte[]> sourceConsumer() {
+ return consumer(
+ sourceBootstrap(),
+ clientId(true, "consumer"),
+ CLIENTS_PREFIX,
+ SOURCE_CLIENTS_PREFIX,
+ SOURCE_CONSUMER_PREFIX
+ );
+ }
+
+ private String sourceBootstrap() {
+ String override = sourceBootstrapOverride();
+ if (override != null) {
+ log.info("Using overridden bootstrap servers for source cluster: {}", override);
+ return override;
+ }
+
+ return resolveBootstrap(SOURCE_CLUSTER_CONFIG, sourceCluster());
+ }
+
+ protected String sourceBootstrapOverride() {
+ return originalsStrings().get(SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG);
+ }
+
+ private String targetBootstrap() {
+ String override = targetBootstrapOverride();
+ if (override != null) {
+ log.info("Using overridden bootstrap servers for target cluster: {}", override);
+ return override;
+ }
+
+ return resolveBootstrap(TARGET_CLUSTER_CONFIG, targetCluster());
+ }
+
+ protected String targetBootstrapOverride() {
+ return originalsStrings().get(TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG);
+ }
+
+ // TODO: Maybe align this with the IDs assigned by the Connect runtime to the
+ // clients it instantiates on behalf of connectors and tasks?
+ private String clientId(boolean source, String clientType) {
+ return "ureplicator3-%s-%s-%s".formatted(
+ role(),
+ source ? "source" : "target",
+ clientType
+ );
+ }
+
+ private Admin admin(String bootstrapServers, String clientId, String... overridesPrefixes) {
+ Map<String, Object> adminProps = new HashMap<>(Map.of(
+ CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers,
+ CommonClientConfigs.CLIENT_ID_CONFIG, clientId
+ ));
+
+ for (String overridesPrefix : overridesPrefixes)
+ adminProps.putAll(originalsWithPrefix(overridesPrefix));
+
+ return Admin.create(adminProps);
+ }
+
+ private Consumer<byte[], byte[]> consumer(String bootstrapServers, String clientId, String... overridesPrefixes) {
+ Map<String, Object> consumerProps = new HashMap<>(Map.of(
+ CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers,
+ CommonClientConfigs.CLIENT_ID_CONFIG, clientId,
+ ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest",
+ ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class,
+ ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class
+ ));
+
+ for (String overridesPrefix : overridesPrefixes)
+ consumerProps.putAll(originalsWithPrefix(overridesPrefix));
+
+ return new KafkaConsumer<>(consumerProps);
+ }
+
+ @VisibleForTesting
+ static String resolveBootstrap(String prop, String cluster) {
+ try {
+ return BrokerResolver.resolveBootstrap(cluster);
+ } catch (IllegalStateException e) {
+ throw new IllegalStateException(
+ "Failed to initialize broker resolver. If testing locally, manually specify the bootstrap servers "
+ + "for the source and target cluster using the " + SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG + " "
+ + "and " + TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG + " properties"
+ );
+ } catch (ResolverException e) {
+ log.error("Failed to resolve bootstrap servers for {}", cluster, e);
+ String errorMessage = createErrorMessage("failed to resolve bootstrap servers for cluster", e);
+ throw new ConfigException(prop, cluster, errorMessage);
+ }
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Connector.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Connector.java
new file mode 100644
index 00000000000..a5aba091d08
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Connector.java
@@ -0,0 +1,449 @@
+package com.uber.data.kafka.connect.ureplicator3;
+
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.config.Config;
+import org.apache.kafka.common.config.ConfigDef;
+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 com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
+import com.uber.data.kafka.connect.distributed.assignment.AbstractAssignor;
+import com.uber.data.kafka.connect.ureplicator3.coordination.AssignedPartition;
+import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStore;
+import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStores;
+import com.uber.data.kafka.connect.utils.ExecutorUtils;
+import com.uber.data.kafka.connect.utils.UberConnectUtils;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.time.Duration;
+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.LinkedHashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import static com.uber.data.kafka.connect.utils.UberConnectUtils.awaitAll;
+import static com.uber.data.kafka.connect.utils.UberConnectUtils.diff;
+
+public class UReplicator3Connector extends SourceConnector {
+
+ private static final Logger log = LoggerFactory.getLogger(UReplicator3Connector.class);
+
+ public static final String SKIP_CONFIG_VALIDATION = "skip.config.validation";
+
+ private UReplicator3ConnectorConfig config;
+ private AbstractAssignor<TaskPartitions, TopicPartition> assignor;
+ private ScheduledExecutorService executor;
+ private AssignmentStore assignmentStore;
+
+ // The values for these maps may be changed at any time, but their contents of any single map will never be mutated
+ private volatile Map<String, Integer> sourcePartitionCounts;
+ private volatile Map<String, Integer> targetPartitionCounts;
+ private volatile int clusterSize;
+
+ @Override
+ public ConfigDef config() {
+ return UReplicator3ConnectorConfig.config();
+ }
+
+ @Override
+ public String version() {
+ return AppInfoParser.getVersion();
+ }
+
+ @Override
+ public Class<? extends Task> taskClass() {
+ return UReplicator3Task.class;
+ }
+
+ @Override
+ public void start(Map<String, String> props) {
+ log.info("Starting connector");
+
+ config = new UReplicator3ConnectorConfig(props);
+
+ assignmentStore = AssignmentStores.get();
+
+ assignor = new TaskPartitionAssignor(
+ config.sourceCluster(),
+ config.targetCluster(),
+ 0.1, // This should never really be used so we hardcode it for now
+ config.taskByteRateThreshold(),
+ config.taskMessageRateThreshold()
+ );
+
+ // Synchronously collect partition counts the first time around
+ // to make sure they're available before we generate task configs
+ doRefreshClusterMetadata();
+
+ long refreshIntervalMs = config.refreshIntervalMs();
+ executor = ExecutorUtils.singleThreadedScheduled("executor-" + config.connectorName());
+ executor.scheduleWithFixedDelay(this::refreshClusterMetadata, refreshIntervalMs, refreshIntervalMs, TimeUnit.MILLISECONDS);
+ // TODO: Should be two separate intervals for these operations
+ executor.scheduleWithFixedDelay(this::refreshAssignments, refreshIntervalMs, refreshIntervalMs, TimeUnit.MILLISECONDS);
+
+ log.info("Finished starting connector");
+ }
+
+ @Override
+ public List<Map<String, String>> taskConfigs(int maxTasks) {
+ log.info("Generating task configs");
+
+ Map<String, Integer> sourcePartitionCountsSnapshot = sourcePartitionCounts;
+ Map<String, Integer> targetPartitionCountsSnapshot = targetPartitionCounts;
+
+ List<TopicPartition> topicPartitions = replicatedPartitions(sourcePartitionCountsSnapshot, targetPartitionCountsSnapshot);
+
+ clusterSize = context.uberClusterSize();
+ int numTasks = numTasks(maxTasks, topicPartitions.size(), clusterSize);
+
+ List<Map<String, String>> result = IntStream.range(0, numTasks)
+ .mapToObj(i -> config.taskConfig(i))
+ .collect(Collectors.toList());
+
+ // Make sure partition assignments are updated in line with the expected number of tasks
+ // Run on a separate thread since we really, really don't want to block too long in this method
+ // (it's invoked from the Connect worker's herder thread, which is basically the central nervous
+ // system of the control plane)
+ executor.execute(() -> refreshAssignments(numTasks));
+
+ log.info("Finished generating task configs");
+
+ return result;
+ }
+
+ @Override
+ public void stop() {
+ log.info("Stopping connector");
+ ExecutorUtils.shutdown(executor, "connector executor", 2, TimeUnit.SECONDS);
+ log.info("Finished stopping connector");
+ }
+
+ @Override
+ public Config validate(Map<String, String> props) {
+ // Breakglass logic: allow the connector config to be updated even if it's failing
+ // Should only be used in emergencies
+ String skipValidationValue = props.get(SKIP_CONFIG_VALIDATION);
+ if (Boolean.parseBoolean(skipValidationValue)) {
+ log.warn("Skipping config validation since {} is set to {}", SKIP_CONFIG_VALIDATION, skipValidationValue);
+ return new Config(List.of());
+ }
+
+ Config result = super.validate(props);
+ UReplicator3ConnectorConfig config;
+
+ try {
+ config = new UReplicator3ConnectorConfig(props);
+ } catch (Exception e) {
+ boolean alreadyHasErrors = result.configValues().stream()
+ .anyMatch(configValue -> !configValue.errorMessages().isEmpty());
+
+ if (alreadyHasErrors) {
+ // We can just log the error; preflight validation will fail without any intervention on our part
+ log.debug("Unable to fully validate connector config", e);
+ return result;
+ } else {
+ // This shouldn't really happen, make sure to log at error level if it does so we can investigate
+ log.error("Unable to fully validate connector config", e);
+ // Force preflight validation to fail
+ throw new ConnectException("Unable to fully validate connector config", e);
+ }
+ }
+
+ config.validate(result);
+
+ return result;
+ }
+
+ @Override
+ public boolean alterOffsets(Map<String, String> connectorConfig, Map<Map<String, ?>, Map<String, ?>> offsets) {
+ return true;
+ }
+
+ private void refreshClusterMetadata() {
+ log.info("Refreshing cluster metadata");
+ try {
+ doRefreshClusterMetadata();
+ log.info("Finished refreshing cluster metadata");
+ } catch (Exception e) {
+ // TODO: Emit metric
+ log.error("Failed to refresh cluster metadata", e);
+ }
+ }
+
+ private void doRefreshClusterMetadata() {
+ Map<String, Integer> oldSourcePartitionCounts = sourcePartitionCounts;
+ try (Admin sourceAdmin = config.sourceAdmin()) {
+ sourcePartitionCounts = partitionCounts(sourceAdmin, config.sourceTopics(), "source");
+ }
+
+ Map<String, Integer> oldTargetPartitionCounts = targetPartitionCounts;
+ try (Admin targetAdmin = config.targetAdmin()) {
+ targetPartitionCounts = partitionCounts(targetAdmin, config.targetTopics(), "target");
+ }
+
+ if (oldSourcePartitionCounts != null && !oldSourcePartitionCounts.equals(sourcePartitionCounts)) {
+ context.requestTaskReconfiguration();
+ }
+ if (oldTargetPartitionCounts != null && !oldTargetPartitionCounts.equals(targetPartitionCounts)) {
+ context.requestTaskReconfiguration();
+ }
+
+ int currentClusterSize = context.uberClusterSize();
+ if (currentClusterSize > 0 && clusterSize > 0) {
+ double clusterSizeRatio = 100.0 * Math.abs(clusterSize - currentClusterSize) / clusterSize;
+ log.trace("Current cluster size: {}; last-used cluster size: {}", currentClusterSize, clusterSize);
+ double reconfigurationThreshold = config.clusterSizeReconfigurationThreshold();
+ if (clusterSizeRatio >= reconfigurationThreshold) {
+ log.debug(
+ "Ratio of current cluster size to last-used cluster size exceeds reconfiguration threshold ({}); generating new task configs",
+ reconfigurationThreshold
+ );
+ context.requestTaskReconfiguration();
+ } else if (currentClusterSize == clusterSize) {
+ log.trace("No change in cluster size");
+ } else {
+ log.trace("Ratio of current cluster size to last-used cluster size is within reconfiguration threshold; will not generate new task configs");
+ }
+ }
+ }
+
+ private static Map<String, Integer> partitionCounts(Admin admin, Set<String> topics, String clusterType) {
+ try {
+ return UberConnectUtils.listPartitions(topics, admin, 1, TimeUnit.MINUTES);
+ } catch (Exception e) {
+ throw new ConnectException("Failed to list partition counts for topics on " + clusterType + " cluster", e);
+ }
+ }
+
+ private List<TopicPartition> replicatedPartitions(
+ Map<String, Integer> sourcePartitionCountsSnapshot,
+ Map<String, Integer> targetPartitionCountsSnapshot
+ ) {
+ Comparator<Map.Entry<String, Integer>> topicPartitionCountEntryComparator =
+ Map.Entry.<String, Integer>comparingByValue()
+ .reversed()
+ .thenComparing(
+ Map.Entry.comparingByKey()
+ );
+
+ return sourcePartitionCountsSnapshot.entrySet().stream()
+ .filter(e -> {
+ String sourceTopic = e.getKey();
+ String targetTopic = config.targetTopic(sourceTopic);
+
+ if (targetPartitionCountsSnapshot.containsKey(targetTopic)) {
+ return true;
+ } else {
+ log.warn(
+ "Not assigning topic {} to tasks since its target partition count is unknown "
+ + "(topic may have been deleted)",
+ sourceTopic);
+ return false;
+ }
+ })
+ .sorted(topicPartitionCountEntryComparator)
+ .flatMap(e -> IntStream.range(0, e.getValue())
+ .mapToObj(partition -> new TopicPartition(e.getKey(), partition))
+ )
+ .toList();
+ }
+
+ private int numTasks(int maxTasks, int numTopicPartitions, int clusterSize) {
+ int numTasks = Math.min(maxTasks, numTopicPartitions);
+ if (clusterSize > 0) {
+ int workerTaskThreshold = config.workerTaskLimit();
+ if (workerTaskThreshold > 0) {
+ int maxWorkerTasks = workerTaskThreshold * clusterSize;
+ log.debug("Total cluster-based task limit: {} (cluster size: {}, tasks per worker: {})",
+ maxWorkerTasks, clusterSize, workerTaskThreshold);
+ numTasks = Math.min(maxWorkerTasks, numTasks);
+ } else {
+ log.debug("Ignoring cluster size since worker task limit is disabled");
+ }
+ } else {
+ log.warn("Cluster size not available; will generate at most {} tasks", numTasks);
+ }
+ return numTasks;
+ }
+
+ private void refreshAssignments() {
+ refreshAssignments(context.uberTaskConfigs().size());
+ }
+
+ private void refreshAssignments(int numTasks) {
+ log.info("Refreshing task partition assignments");
+ try {
+ doRefreshAssignments(numTasks);
+ log.info("Finished refreshing task partition assignments");
+ } catch (Exception e) {
+ // TODO: Emit metric
+ log.error("Failed to refresh task partition assignments", e);
+ }
+ }
+
+ private void doRefreshAssignments(int numTasks) {
+ Map<String, Integer> sourcePartitionCountsSnapshot = sourcePartitionCounts;
+ Map<String, Integer> targetPartitionCountsSnapshot = targetPartitionCounts;
+
+ List<TopicPartition> topicPartitions = replicatedPartitions(sourcePartitionCountsSnapshot, targetPartitionCountsSnapshot);
+
+ Map<Integer, ? extends Collection<AssignedPartition>> existingStoredAssignments = assignmentStore.getAssignments(
+ config.connectorName(),
+ config.sourceCluster(),
+ config.targetCluster()
+ );
+
+ List<TaskPartitions> existingAssignments = existingAssignments(existingStoredAssignments);
+ Collection<TopicPartition> unassignedPartitions = unassignedPartitions(numTasks, existingAssignments, topicPartitions);
+
+ List<TaskPartitions> newAssignments = IntStream.range(0, numTasks)
+ .mapToObj(i -> i < existingAssignments.size()
+ ? existingAssignments.get(i).copy()
+ : new TaskPartitions(config.connectorName(), i, new HashSet<>())
+ ).toList();
+ assignor.assign(newAssignments, unassignedPartitions);
+
+ Map<Integer, Collection<AssignedPartition>> newStoredAssignments = newStoredAssignments(
+ targetPartitionCountsSnapshot,
+ newAssignments
+ );
+
+ Collection<Future<?>> storeFutures = new ArrayList<>();
+
+ Set<TopicPartition> toExplicitlyRevoke = existingAssignments.stream()
+ .map(TaskPartitions::partitions)
+ .flatMap(Collection::stream)
+ .collect(Collectors.toCollection(HashSet::new));
+ toExplicitlyRevoke.removeAll(new HashSet<>(topicPartitions));
+
+ if (!toExplicitlyRevoke.isEmpty()) {
+ log.debug("Found {} partitions to explicitly revoke", toExplicitlyRevoke.size());
+ } else {
+ log.trace("No partitions to explicitly revoke");
+ }
+
+ toExplicitlyRevoke.stream()
+ .map(this::replicatedTopicPartition)
+ .map(rtp -> assignmentStore.revokePartition(config.connectorName(), rtp))
+ .forEach(storeFutures::add);
+
+ int tasksWithNewPartitions = 0;
+ for (int i = 0; i < newAssignments.size(); i++) {
+ final int taskId = i;
+ Collection<AssignedPartition> existingPartitions = existingStoredAssignments.get(taskId);
+ Collection<AssignedPartition> newPartitions = newStoredAssignments.get(taskId);
+
+ Set<AssignedPartition> added = newPartitions != null ? diff(newPartitions, existingPartitions) : Set.of();
+
+ if (!added.isEmpty()) {
+ log.debug("Found {} partitions to add for task {}-{}", added.size(), config.connectorName(), taskId);
+ tasksWithNewPartitions++;
+ } else {
+ log.trace("No partitions to add for task {}-{}", config.connectorName(), taskId);
+ }
+
+ added.stream()
+ .map(ap -> {
+ ReplicatedTopicPartition rtp = replicatedTopicPartition(ap.source());
+ return assignmentStore.assignPartition(config.connectorName(), rtp, taskId, ap.dstTopic(), ap.dstPartition());
+ })
+ .forEach(storeFutures::add);
+ }
+
+ if (tasksWithNewPartitions > 0) {
+ log.debug("Assigned new partitions to {} total tasks", tasksWithNewPartitions);
+ } else {
+ log.debug("Assigned no new partitions to any tasks");
+ }
+
+ // We may want to make this duration configurable in the future
+ // TODO: Should also add metrics on write failures
+ try {
+ awaitAll(storeFutures, Duration.ofMinutes(1));
+ } catch (InterruptedException | TimeoutException | ExecutionException e) {
+ log.error("Failed to publish task partition assignment updates", e);
+ }
+ }
+
+ private List<TaskPartitions> existingAssignments(Map<Integer, ? extends Collection<AssignedPartition>> storedAssignments) {
+ if (storedAssignments.isEmpty())
+ return List.of();
+
+ int highestTask = Collections.max(storedAssignments.keySet());
+ return IntStream.rangeClosed(0, highestTask) // rangeClosed because tasks are zero-indexed
+ .mapToObj(i -> {
+ Collection<AssignedPartition> partitions = storedAssignments.get(i);
+ Set<TopicPartition> partitionsSet = new HashSet<>();
+ if (partitions != null) {
+ partitions.stream()
+ .map(AssignedPartition::source)
+ .forEach(partitionsSet::add);
+ }
+ return new TaskPartitions(config.connectorName(), i, partitionsSet);
+ }).toList();
+ }
+
+ private Set<TopicPartition> unassignedPartitions(int numTasks, List<TaskPartitions> taskPartitions, Collection<TopicPartition> topicPartitions) {
+ Set<TopicPartition> result = new LinkedHashSet<>(topicPartitions);
+
+ for (int i = 0; i < numTasks && i < taskPartitions.size(); i++) {
+ result.removeAll(taskPartitions.get(i).partitions());
+ }
+
+ return result;
+ }
+
+ private ReplicatedTopicPartition replicatedTopicPartition(TopicPartition tp) {
+ return new ReplicatedTopicPartition(config.sourceCluster(), config.targetCluster(), tp);
+ }
+
+ private Map<Integer, Collection<AssignedPartition>> newStoredAssignments(
+ Map<String, Integer> targetPartitionCounts,
+ List<TaskPartitions> newAssignments
+ ) {
+ Map<Integer, Collection<AssignedPartition>> result = new HashMap<>();
+ for (int i = 0; i < newAssignments.size(); i++) {
+ TaskPartitions newAssignment = newAssignments.get(i);
+ Collection<AssignedPartition> assignment = new HashSet<>();
+ for (TopicPartition source : newAssignment.partitions()) {
+ String dstTopic = config.targetTopic(source.topic());
+ Integer targetPartitionCount = targetPartitionCounts.get(dstTopic);
+ if (targetPartitionCount == null) {
+ log.warn("Could not find partition count for topic {} on target cluster; will not replicate topic partition {} (this should never happen)",
+ dstTopic, source);
+ continue;
+ }
+ int dstPartition = source.partition() % targetPartitionCount;
+ AssignedPartition assignedPartition = new AssignedPartition(
+ source.topic(),
+ source.partition(),
+ !Objects.equals(dstTopic, source.topic()) ? dstTopic : null,
+ dstPartition != source.partition() ? dstPartition : null
+ );
+ assignment.add(assignedPartition);
+ }
+ result.put(i, assignment);
+ }
+ return result;
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3ConnectorConfig.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3ConnectorConfig.java
new file mode 100644
index 00000000000..3f30cec9e9e
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3ConnectorConfig.java
@@ -0,0 +1,385 @@
+package com.uber.data.kafka.connect.ureplicator3;
+
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.consumer.Consumer;
+import org.apache.kafka.common.config.Config;
+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 com.google.common.annotations.VisibleForTesting;
+import com.uber.data.kafka.connect.workload.kafka.KafkaBasedWorkloadProvider;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3TaskConfig.TASK_ID_CONFIG;
+import static com.uber.data.kafka.connect.utils.UberConnectUtils.addError;
+import static com.uber.data.kafka.connect.utils.UberConnectUtils.diff;
+import static com.uber.data.kafka.connect.utils.UberConnectUtils.listTopics;
+import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.TASKS_MAX_CONFIG;
+
+public class UReplicator3ConnectorConfig extends UReplicator3Config {
+
+ private static final Logger log = LoggerFactory.getLogger(UReplicator3ConnectorConfig.class);
+
+ public static final String TOPICS_CONFIG = "topics";
+ private static final String TOPICS_DOC = "Comma-separated list of topics to replicate";
+
+ public static final String WORKER_TASK_LIMIT_CONFIG = "worker.task.limit";
+ public static final int WORKER_TASK_LIMIT_DEFAULT = 2;
+ private static final String WORKER_TASK_LIMIT_DOC = "Maximum number of tasks to generate per worker. "
+ + "To allow any number of tasks per worker, set to a non-positive value";
+
+ public static final String TASK_BYTE_RATE_THRESHOLD_CONFIG = "task.byte.rate.threshold";
+ public static final double TASK_BYTE_RATE_THRESHOLD_DEFAULT = 70 * 1024 * 1024;
+ private static final String TASK_BYTE_RATE_THRESHOLD_DOC = "Maximum total byte rate for partitions assigned "
+ + "to each task before attempting to revoke partitions and redistribute to other tasks";
+
+ public static final String TASK_MESSAGE_RATE_THRESHOLD_CONFIG = "task.message.rate.threshold";
+ public static final double TASK_MESSAGE_RATE_THRESHOLD_DEFAULT = 70_000;
+ private static final String TASK_MESSAGE_RATE_THRESHOLD_DOC = "Maximum total message rate for partitions assigned "
+ + "to each task before attempting to revoke partitions and redistribute to other workers";
+
+ public static final String REFRESH_INTERVAL_MS = "refresh.interval.ms";
+ public static final long REFRESH_INTERVAL_MS_DEFAULT = TimeUnit.MINUTES.toMillis(1);
+ private static final String REFRESH_INTERVAL_MS_DOC = "How often to check for changes in the partition counts of source and target topics, "
+ + "the size of the Kafka Connect cluster, and imbalances in the partitions assigned to each task, "
+ + "possibly generating new task configs and/or reassigning partitions in response";
+
+ public static final String CLUSTER_SIZE_RECONFIGURATION_THRESHOLD_CONFIG = "cluster.size.reconfiguration.threshold";
+ public static final double CLUSTER_SIZE_RECONFIGURATION_THRESHOLD_DEFAULT = 10;
+ private static final String CLUSTER_SIZE_RECONFIGURATION_THRESHOLD_DOC = "Minimum percentage difference "
+ + "in cluster size (i.e., number of workers) before new task configs are generated";
+
+ public static final String TOPIC_RENAME_PREFIX = "topic.rename.";
+
+ private final Map<String, String> topicRenames;
+ private final Set<String> sourceTopics;
+ private final Set<String> targetTopics;
+
+ @SuppressWarnings("this-escape")
+ public UReplicator3ConnectorConfig(Map<String, ?> props) {
+ super(config(), props);
+ this.topicRenames = Collections.unmodifiableMap(parseTopicRenames(originalsStrings(), null));
+ this.sourceTopics = Set.copyOf(getList(TOPICS_CONFIG));
+ this.targetTopics = Collections.unmodifiableSet(computeTargetTopics(sourceTopics, topicRenames));
+ }
+
+ protected static ConfigDef config() {
+ return baseConfig()
+ .define(
+ TOPICS_CONFIG,
+ Type.LIST,
+ ConfigDef.NO_DEFAULT_VALUE,
+ ConfigDef.ValidList.anyNonDuplicateValues(false, false),
+ Importance.HIGH,
+ TOPICS_DOC
+ ).define(
+ WORKER_TASK_LIMIT_CONFIG,
+ Type.INT,
+ WORKER_TASK_LIMIT_DEFAULT,
+ Importance.MEDIUM,
+ WORKER_TASK_LIMIT_DOC
+ ).define(
+ TASK_BYTE_RATE_THRESHOLD_CONFIG,
+ Type.DOUBLE,
+ TASK_BYTE_RATE_THRESHOLD_DEFAULT,
+ atLeast(1),
+ Importance.MEDIUM,
+ TASK_BYTE_RATE_THRESHOLD_DOC
+ ).define(
+ TASK_MESSAGE_RATE_THRESHOLD_CONFIG,
+ Type.DOUBLE,
+ TASK_MESSAGE_RATE_THRESHOLD_DEFAULT,
+ atLeast(1),
+ Importance.MEDIUM,
+ TASK_MESSAGE_RATE_THRESHOLD_DOC
+ ).define(
+ REFRESH_INTERVAL_MS,
+ Type.LONG,
+ REFRESH_INTERVAL_MS_DEFAULT,
+ atLeast(1),
+ Importance.MEDIUM,
+ REFRESH_INTERVAL_MS_DOC
+ ).define(
+ CLUSTER_SIZE_RECONFIGURATION_THRESHOLD_CONFIG,
+ Type.DOUBLE,
+ CLUSTER_SIZE_RECONFIGURATION_THRESHOLD_DEFAULT,
+ atLeast(1),
+ Importance.LOW,
+ CLUSTER_SIZE_RECONFIGURATION_THRESHOLD_DOC
+ );
+ }
+
+ @Override
+ protected String role() {
+ return "connector-" + connectorName();
+ }
+
+ public void validate(Config configValidationResult) {
+ parseTopicRenames(originalsStrings(), configValidationResult);
+
+ if (!validateBootstrapServersOverrides(configValidationResult)) {
+ // Don't try to validate more since we might end up contacting the wrong Kafka clusters
+ return;
+ }
+
+ validateCluster(configValidationResult, true);
+ validateCluster(configValidationResult, false);
+
+ validateWorkloadPublishing(configValidationResult);
+ }
+
+ public Set<String> sourceTopics() {
+ return sourceTopics;
+ }
+
+ public Set<String> targetTopics() {
+ return targetTopics;
+ }
+
+ public String targetTopic(String sourceTopic) {
+ return topicRenames.getOrDefault(sourceTopic, sourceTopic);
+ }
+
+ public int workerTaskLimit() {
+ return getInt(WORKER_TASK_LIMIT_CONFIG);
+ }
+
+ public double taskByteRateThreshold() {
+ return getDouble(TASK_BYTE_RATE_THRESHOLD_CONFIG);
+ }
+
+ public double taskMessageRateThreshold() {
+ return getDouble(TASK_MESSAGE_RATE_THRESHOLD_CONFIG);
+ }
+
+ public long refreshIntervalMs() {
+ return getLong(REFRESH_INTERVAL_MS);
+ }
+
+ public double clusterSizeReconfigurationThreshold() {
+ return getDouble(CLUSTER_SIZE_RECONFIGURATION_THRESHOLD_CONFIG);
+ }
+
+ public Map<String, String> taskConfig(int taskId) {
+ Map<String, String> result = new HashMap<>(originalsStrings());
+
+ // If we don't do this, every time a topic is added or removed, a task reconfiguration will be triggered
+ result.remove(TOPICS_CONFIG);
+ // Same for topic renames
+ result.keySet().removeIf(key -> key.startsWith(TOPIC_RENAME_PREFIX));
+ // Same for horizontal scaling properties that are only used by the connector and not its tasks
+ result.remove(TASKS_MAX_CONFIG);
+ result.remove(WORKER_TASK_LIMIT_CONFIG);
+ result.remove(CLUSTER_SIZE_RECONFIGURATION_THRESHOLD_CONFIG);
+ // And other properties that are also not used by tasks
+ result.remove(REFRESH_INTERVAL_MS);
+ result.remove(TASK_BYTE_RATE_THRESHOLD_CONFIG);
+ result.remove(TASK_MESSAGE_RATE_THRESHOLD_CONFIG);
+
+ result.put(TASK_ID_CONFIG, Integer.toString(taskId));
+
+ return result;
+ }
+
+ private boolean validateBootstrapServersOverrides(Config configValidationResult) {
+ AtomicBoolean result = new AtomicBoolean(true);
+ Stream.of(SOURCE_CONSUMER_PREFIX, SOURCE_ADMIN_PREFIX, TARGET_ADMIN_PREFIX)
+ .map(prefix -> prefix + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG)
+ .forEach(prop -> {
+ Object forbiddenOverride = originals().get(prop);
+ if (forbiddenOverride != null) {
+ String errorMessage = "cannot override bootstrap servers for individual clients; instead, use the "
+ + SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG + " and/or "
+ + TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG + " properties";
+ addError(configValidationResult, prop, forbiddenOverride, errorMessage);
+ result.set(false);
+ }
+ });
+ return result.get();
+ }
+
+ private void validateCluster(Config configValidationResult, boolean source) {
+ String sourceTarget = source ? "source" : "target";
+ String bootstrapOverride = source ? sourceBootstrapOverride() : targetBootstrapOverride();
+
+ String prop;
+ String value;
+ if (bootstrapOverride != null) {
+ prop = source ? SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG : TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG;
+ value = bootstrapOverride;
+ } else {
+ String clusterName = source ? sourceCluster() : targetCluster();
+ prop = source ? SOURCE_CLUSTER_CONFIG : TARGET_CLUSTER_CONFIG;
+ try {
+ resolveBootstrap(prop, clusterName);
+ } catch (ConfigException e) {
+ addError(configValidationResult, prop, clusterName, "failed to resolve bootstrap servers for " + sourceTarget + " cluster", e);
+ return;
+ }
+ value = clusterName;
+ }
+
+ boolean clusterReachable;
+ Set<String> topics = source ? sourceTopics() : targetTopics();
+ try (Admin admin = source ? sourceAdmin() : targetAdmin()) {
+ clusterReachable = true;
+ validateTopics(configValidationResult, topics, source, admin);
+ } catch (Exception e) {
+ clusterReachable = false;
+ log.error("Failed to create admin client for {} cluster", sourceTarget, e);
+ addError(configValidationResult, prop, value, "failed to create admin client for " + sourceTarget + " cluster", e);
+ }
+
+ // Only bother validating the source consumer if the Kafka cluster can actually be reached
+ if (clusterReachable && source) {
+ validateSourceConsumer(configValidationResult);
+ }
+ }
+
+ private void validateTopics(Config configValidationResult, Set<String> configuredTopics, boolean source, Admin admin) {
+ String sourceTarget = source ? "source" : "target";
+
+ Set<String> presentTopics;
+ try {
+ presentTopics = listTopics(admin, 1, TimeUnit.MINUTES);
+ } catch (InterruptedException | TimeoutException | ExecutionException e) {
+ log.error("Failed to list topics on {} cluster", sourceTarget, e);
+
+ String clusterProp = source ? SOURCE_CLUSTER_CONFIG : TARGET_CLUSTER_CONFIG;
+ addError(configValidationResult, clusterProp, null, "failed to list topics on " + sourceTarget + " cluster", e);
+
+ return;
+ }
+
+ Set<String> missingTopics = diff(configuredTopics, presentTopics);
+
+ if (!missingTopics.isEmpty()) {
+ if (source) {
+ String errorMessage = "the following topics are missing on the source cluster: " + String.join(", ", missingTopics);
+ addError(configValidationResult, TOPICS_CONFIG, null, errorMessage);
+ } else {
+ Map<String, Set<String>> reversedTopicRenames = topicRenames.entrySet().stream()
+ .collect(Collectors.groupingBy(
+ Map.Entry::getValue,
+ Collectors.mapping(Map.Entry::getKey, Collectors.toSet()
+ )));
+ Set<String> missingNonRenamedTopics = new HashSet<>();
+ for (String targetTopic : missingTopics) {
+ Set<String> renamedSourceTopics = reversedTopicRenames.get(targetTopic);
+ if (renamedSourceTopics != null) {
+ renamedSourceTopics.forEach(sourceTopic -> {
+ // If a missing target topic is specified via a topic rename, we should raise that to the user by
+ // attaching an error message to the topic.rename.<source_topic> property
+ addError(
+ configValidationResult,
+ TOPIC_RENAME_PREFIX + sourceTopic,
+ targetTopic,
+ "renamed topic is missing on the target cluster"
+ );
+ });
+ }
+
+ if (sourceTopics.contains(targetTopic) && !topicRenames.containsKey(targetTopic))
+ missingNonRenamedTopics.add(targetTopic);
+ }
+ if (!missingNonRenamedTopics.isEmpty()) {
+ String errorMessage = "the following topics are missing on the target cluster: " + String.join(", ", missingNonRenamedTopics);
+ addError(configValidationResult, TOPICS_CONFIG, null, errorMessage);
+ }
+ }
+ }
+ }
+
+ private void validateSourceConsumer(Config configValidationResult) {
+ try (Consumer<?, ?> consumer = sourceConsumer()) {
+ } catch (Exception e) {
+ // This error could be caused by any number of custom consumer properties and we don't really
+ // have a way of knowing which one(s) caused the issue, so we just report the error with a detailed
+ // message on the source cluster property and hope that there's enough info there for the user to
+ // pinpoint the problem
+ addError(configValidationResult, SOURCE_CLUSTER_CONFIG, sourceCluster(), "failed to create consumer for source cluster", e);
+ }
+ }
+
+ private void validateWorkloadPublishing(Config configValidationResult) {
+ if (workloadReportInterval() == null) {
+ log.debug("Workload publishing is disabled; skipping validation");
+ return;
+ }
+
+ try {
+ KafkaBasedWorkloadProvider.assertStoreInitialized();
+ } catch (IllegalStateException e) {
+ log.debug("Kafka workload store does not appear to be initialized", e);
+ String errorMessage = "This worker does not appear to be configured to use the Kafka-based workload provider, "
+ + "which means that uReplicator3's workload publishing feature must be disabled. Set the "
+ + WORKLOAD_REPORT_INTERVAL_MS_CONFIG + " property to zero or a negative value to disable, or reconfigure "
+ + "the workers in your Kafka Connect cluster to use the Kafka-based workload provider";
+
+ addError(configValidationResult, WORKLOAD_REPORT_INTERVAL_MS_CONFIG, null, errorMessage);
+ }
+ }
+
+ @VisibleForTesting
+ static Map<String, String> parseTopicRenames(Map<String, String> originals, Config configValidationResult) {
+ Map<String, String> result = new HashMap<>();
+ originals.forEach((key, value) -> {
+ if (!key.startsWith(TOPIC_RENAME_PREFIX))
+ return;
+
+ String sourceTopic = key.substring(TOPIC_RENAME_PREFIX.length());
+
+ if (sourceTopic.isBlank()) {
+ if (configValidationResult != null) {
+ addError(configValidationResult, key, value, "source topic is empty or blank");
+ } else {
+ log.warn("Ignoring invalid property {} with value {}", key, value);
+ }
+ return;
+ }
+
+ if (value == null || value.isBlank()) {
+ if (configValidationResult != null) {
+ addError(configValidationResult, key, value, "target topic is null, empty, or blank");
+ } else {
+ log.warn("Ignoring property {} with invalid value {}", key, value);
+ }
+ return;
+ }
+
+ result.putIfAbsent(sourceTopic, value);
+ });
+
+ return result;
+ }
+
+ private static Set<String> computeTargetTopics(Set<String> sourceTopics, Map<String, String> topicRenames) {
+ return sourceTopics.stream()
+ .map(t -> topicRenames.getOrDefault(t, t))
+ .collect(Collectors.toSet());
+ }
+
+ // Run with `./gradlew :connect:uber:genUReplicator3Config`
+ public static void main(String[] args) {
+ System.out.println(config().toHtml());
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Task.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Task.java
new file mode 100644
index 00000000000..df70b0c1442
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Task.java
@@ -0,0 +1,367 @@
+package com.uber.data.kafka.connect.ureplicator3;
+
+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.producer.RecordMetadata;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.header.Header;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.header.ConnectHeaders;
+import org.apache.kafka.connect.header.Headers;
+import org.apache.kafka.connect.source.SourceRecord;
+import org.apache.kafka.connect.source.SourceTask;
+
+import com.uber.data.kafka.connect.aakafkaoffsetmgmt.OffsetSnapshotClient;
+import com.uber.data.kafka.connect.aakafkaoffsetmgmt.OffsetSnapshotReporter;
+import com.uber.data.kafka.connect.ureplicator3.coordination.AssignedPartition;
+import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStore;
+import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStores;
+import com.uber.data.kafka.connect.workload.kafka.KafkaBasedWorkloadProvider;
+import com.uber.data.kafka.connect.workload.tracking.WorkloadTrackingProcessor;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
+
+public class UReplicator3Task extends SourceTask {
+
+ private static final Logger log = LoggerFactory.getLogger(UReplicator3Task.class);
+
+ public static final String SOURCE_PARTITION_SOURCE_CLUSTER = "source.cluster";
+ public static final String SOURCE_PARTITION_TARGET_CLUSTER = "target.cluster";
+ public static final String SOURCE_PARTITION_KAFKA_TOPIC = "topic";
+ public static final String SOURCE_PARTITION_KAFKA_PARTITION = "partition";
+ public static final String SOURCE_OFFSET_KAFKA_OFFSET = "offset";
+
+ private Time time;
+ private UReplicator3TaskConfig config;
+ private Consumer<byte[], byte[]> sourceConsumer;
+ private WorkloadTrackingProcessor workloadTracker;
+ private AssignmentStore partitionAssignmentStore;
+ private volatile boolean assignmentChanged;
+ private Map<TopicPartition, TopicPartition> overrides;
+ private OffsetSnapshotReporter offsetSnapshotReporter;
+
+ @Override
+ public void start(Map<String, String> props) {
+ log.info("Starting task");
+
+ time = Time.SYSTEM;
+
+ config = new UReplicator3TaskConfig(props);
+
+ overrides = Map.of();
+
+ partitionAssignmentStore = AssignmentStores.get();
+ partitionAssignmentStore.addListener(config.connectorName(), config.taskId(), this::onAssignmentUpdate);
+
+ sourceConsumer = config.sourceConsumer();
+
+ // Force us to update the consumer assignment before polling it for the first time,
+ // since right now it's not assigned anything
+ assignmentChanged = true;
+
+ initializeWorkloadTracker();
+ initializeOffsetSnapshotReporter();
+
+ log.info("Finished starting task");
+ }
+
+ @Override
+ public List<SourceRecord> poll() {
+ log.trace("Starting to poll for records"); // This gets called a lot; can't put it at INFO like the others
+
+ // Consumer isn't thread-safe, so we make sure to update its partition assignment on the same thread
+ // that we use to poll it for records
+ if (assignmentChanged) {
+ // Reset the flag before updating the assignment; that way, any subsequent assignment updates
+ // (which will be picked up on a separate thread and may be observed concurrently with our consumer update)
+ // won't be skipped
+ assignmentChanged = false;
+ updateAssignment();
+ }
+
+ if (sourceConsumer.assignment().isEmpty()) {
+ log.trace("Task is not assigned any partitions; will idle and return empty list of records");
+ // Idle here so that we don't enter a tight loop where no (real) work is done. We can't block indefinitely though,
+ // since the task may be scheduled for shutdown and the same thread that calls this method is also responsible
+ // for calling Task::stop
+ // We may want to make this configurable in the future but for now this should be fine
+ time.sleep(100);
+ return List.of();
+ }
+
+ ConsumerRecords<byte[], byte[]> records = sourceConsumer.poll(config.consumerPollTimeout());
+
+ List<SourceRecord> result = new ArrayList<>(records.count());
+ for (ConsumerRecord<byte[], byte[]> record : records) {
+ SourceRecord sourceRecord = convertRecord(record);
+
+ if (workloadTracker != null)
+ workloadTracker.track(sourceRecord);
+
+ result.add(sourceRecord);
+ }
+
+ log.trace("Finished polling for records");
+ return result;
+ }
+
+ @Override
+ public void stop() {
+ log.info("Stopping task");
+ Utils.closeQuietly(sourceConsumer, "source consumer");
+ Utils.closeQuietly(workloadTracker, "workload tracker");
+ if (offsetSnapshotReporter != null)
+ Utils.closeQuietly(offsetSnapshotReporter::shutdown, "offset snapshot reporter");
+ log.info("Finished stopping task");
+ }
+
+ @Override
+ public String version() {
+ return AppInfoParser.getVersion();
+ }
+
+ @Override
+ public void commitRecord(SourceRecord record, RecordMetadata metadata) {
+ if (offsetSnapshotReporter != null) {
+ TopicPartition dstTopicPartition = new TopicPartition(metadata.topic(), metadata.partition());
+ if (!metadata.hasOffset()) {
+ log.warn("Record from topic partition {} does not have offset; cannot report offsets for it", dstTopicPartition);
+ return;
+ }
+ long dstOffset = metadata.offset();
+
+ TopicPartition srcTopicPartition = decodeSourcePartition(record.sourcePartition());
+ if (srcTopicPartition == null) {
+ log.warn("Could not decode source partition for record sent to topic partition {}; cannot report offsets for it", dstTopicPartition);
+ return;
+ }
+
+ Long srcOffset = decodeSourceOffset(srcTopicPartition, record.sourceOffset());
+ if (srcOffset == null) {
+ log.warn("Could not decode source offset for record sent to topic partition {}; cannot report offsets for it", dstTopicPartition);
+ return;
+ }
+
+ offsetSnapshotReporter.record(srcTopicPartition.topic(), srcTopicPartition.partition(), srcOffset, dstTopicPartition.partition(), dstOffset);
+ }
+ }
+
+ private Collection<AssignedPartition> assignedPartitions() {
+ return partitionAssignmentStore.getAssignment(
+ config.connectorName(),
+ config.taskId(),
+ config.sourceCluster(),
+ config.targetCluster()
+ );
+ }
+
+ private static Set<TopicPartition> sourcePartitions(Collection<AssignedPartition> partitions) {
+ return partitions.stream().map(AssignedPartition::source).collect(Collectors.toSet());
+ }
+
+ private void updateAssignment() {
+ log.debug("Updating consumer assignment");
+
+ Set<TopicPartition> existingAssignment = sourceConsumer.assignment();
+ Collection<AssignedPartition> newAssignment = assignedPartitions();
+ Set<TopicPartition> sourceAssignment = sourcePartitions(newAssignment);
+ sourceConsumer.assign(sourceAssignment);
+
+ List<TopicPartition> newPartitions = sourceAssignment.stream()
+ .filter(Predicate.not(existingAssignment::contains))
+ .toList();
+ List<Map<String, Object>> newSourcePartitions = newPartitions.stream()
+ .map(tp -> encodeSourcePartition(tp.topic(), tp.partition()))
+ .toList();
+
+ if (!newPartitions.isEmpty()) {
+ log.debug("Assigned new partition(s): {}", newPartitions);
+ } else {
+ log.debug("Assigned no new partitions; revocation or topic/partition override may have taken place");
+ }
+
+ context.offsetStorageReader().offsets(newSourcePartitions).forEach((sourcePartition, sourceOffset) -> {
+ TopicPartition topicPartition = decodeSourcePartition(sourcePartition);
+ if (topicPartition == null) {
+ return;
+ }
+
+ Long offset = decodeSourceOffset(topicPartition, sourceOffset);
+ if (offset == null) {
+ log.trace("Found no committed offset for newly-assigned partition {}", topicPartition);
+ return;
+ }
+
+ log.trace("Found committed offset {} for newly-assigned partition {}", offset, topicPartition);
+ // Stored offset is what we were last able to replicate; we should seek to one beyond that
+ sourceConsumer.seek(topicPartition, offset + 1);
+ });
+
+ overrides = newAssignment.stream()
+ .filter(AssignedPartition::hasOverride)
+ .collect(Collectors.toMap(
+ AssignedPartition::source,
+ AssignedPartition::target
+ ));
+ }
+
+ private void initializeWorkloadTracker() {
+ Duration workloadReportInterval = config.workloadReportInterval();
+ if (workloadReportInterval == null) {
+ log.warn("Workload tracking is disabled");
+ return;
+ }
+
+ log.debug("Initializing workload tracker");
+ workloadTracker = new WorkloadTrackingProcessor(
+ config.sourceCluster(),
+ config.targetCluster(),
+ // TODO: Enable workload tracking to be performed even if the workload provider isn't set up on this worker
+ KafkaBasedWorkloadProvider::publishWorkload,
+ () -> sourcePartitions(assignedPartitions()),
+ workloadReportInterval
+ );
+ workloadTracker.start();
+ log.debug("Finished initializing workload tracker");
+ }
+
+ private void initializeOffsetSnapshotReporter() {
+ if (!config.offsetReporterEnabled()) {
+ log.warn("Offset reporter is disabled");
+ return;
+ }
+
+ log.debug("Initializing offset reporter");
+ // We may want to make some of these values configurable in the future, but for now
+ // this is fine since they're the values used across the entire uReplicator2 fleet
+ int offsetReportIntervalMs = 600_000;
+ String offsetSnapshotPostUrl = "http://127.0.0.1:18668/snapshot";
+ int offsetSnapshotRequestTimeoutMs = 10_000;
+ OffsetSnapshotClient offsetSnapshotClient = new OffsetSnapshotClient(offsetSnapshotPostUrl, offsetSnapshotRequestTimeoutMs);
+
+ // 🤷 Seems as reasonable as anything I guess
+ String routeId = config.connectorName() + "-" + config.sourceCluster() + "-" + config.targetCluster();
+ offsetSnapshotReporter = new OffsetSnapshotReporter(
+ config.sourceCluster(),
+ config.targetCluster(),
+ offsetReportIntervalMs,
+ offsetSnapshotClient,
+ routeId
+ );
+ offsetSnapshotReporter.start();
+ log.debug("Initialized offset reporter");
+ }
+
+ private SourceRecord convertRecord(ConsumerRecord<byte[], byte[]> record) {
+ TopicPartition source = new TopicPartition(record.topic(), record.partition());
+ TopicPartition target = overrides.getOrDefault(source, source);
+
+ return new SourceRecord(
+ encodeSourcePartition(source.topic(), source.partition()),
+ encodeSourceOffset(record),
+ target.topic(),
+ target.partition(),
+ Schema.OPTIONAL_BYTES_SCHEMA,
+ record.key(),
+ Schema.OPTIONAL_BYTES_SCHEMA,
+ record.value(),
+ record.timestamp(),
+ convertHeaders(record.headers())
+ );
+ }
+
+ private Map<String, Object> encodeSourcePartition(String topic, int partition) {
+ return encodeSourcePartition(config.sourceCluster(), config.targetCluster(), topic, partition);
+ }
+
+ public static Map<String, Object> encodeSourcePartition(String srcKafka, String dstKafka, String topic, int partition) {
+ // LinkedHashMap to guarantee order of key/value pairs across serialization
+ Map<String, Object> result = new LinkedHashMap<>();
+ result.put(SOURCE_PARTITION_SOURCE_CLUSTER, srcKafka);
+ result.put(SOURCE_PARTITION_TARGET_CLUSTER, dstKafka);
+ result.put(SOURCE_PARTITION_KAFKA_TOPIC, topic);
+ result.put(SOURCE_PARTITION_KAFKA_PARTITION, partition);
+ return result;
+ }
+
+ private Map<String, Object> encodeSourceOffset(ConsumerRecord<?, ?> record) {
+ return encodeSourceOffset(record.offset());
+ }
+
+ public static Map<String, Object> encodeSourceOffset(long offset) {
+ return Map.of(SOURCE_OFFSET_KAFKA_OFFSET, offset);
+ }
+
+ private TopicPartition decodeSourcePartition(Map<String, ?> sourcePartition) {
+ Object rawTopic = sourcePartition.get(SOURCE_PARTITION_KAFKA_TOPIC);
+ if (rawTopic == null) {
+ log.warn("Ignoring invalid source partition {}; missing {} field", sourcePartition, SOURCE_PARTITION_KAFKA_TOPIC);
+ return null;
+ }
+ if (!(rawTopic instanceof String)) {
+ log.warn("Ignoring invalid source partition {}; field {} is not a string", sourcePartition, SOURCE_PARTITION_KAFKA_TOPIC);
+ return null;
+ }
+
+ Object rawPartition = sourcePartition.get(SOURCE_PARTITION_KAFKA_PARTITION);
+ if (rawPartition == null) {
+ log.warn("Ignoring invalid source partition {}; missing {} field", sourcePartition, SOURCE_PARTITION_KAFKA_PARTITION);
+ return null;
+ }
+ if (!(rawPartition instanceof Number)) {
+ log.warn("Ignoring invalid source partition {}; field {} is not a number", sourcePartition, SOURCE_PARTITION_KAFKA_PARTITION);
+ return null;
+ }
+
+ String topic = (String) rawTopic;
+ int partition = ((Number) rawPartition).intValue();
+ return new TopicPartition(topic, partition);
+ }
+
+ private Long decodeSourceOffset(TopicPartition topicPartition, Map<String, ?> sourceOffset) {
+ if (sourceOffset == null) {
+ log.debug("No source offset found for {}; this partition is probably being replicated for the first time", topicPartition);
+ return null;
+ }
+
+ Object rawOffset = sourceOffset.get(SOURCE_OFFSET_KAFKA_OFFSET);
+ if (rawOffset == null) {
+ log.warn("Ignoring invalid source offset {}; missing {} field", sourceOffset, SOURCE_OFFSET_KAFKA_OFFSET);
+ return null;
+ }
+ if (!(rawOffset instanceof Number)) {
+ log.warn("Ignoring invalid source offset {}; field {} is not a number", sourceOffset, SOURCE_OFFSET_KAFKA_OFFSET);
+ return null;
+ }
+
+ return ((Number) rawOffset).longValue();
+ }
+
+ private static Headers convertHeaders(org.apache.kafka.common.header.Headers headers) {
+ ConnectHeaders result = new ConnectHeaders();
+ for (Header header : headers) {
+ result.addBytes(header.key(), header.value());
+ }
+ return result;
+ }
+
+ private void onAssignmentUpdate() {
+ this.assignmentChanged = true;
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3TaskConfig.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3TaskConfig.java
new file mode 100644
index 00000000000..749a5a2e1d9
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3TaskConfig.java
@@ -0,0 +1,42 @@
+package com.uber.data.kafka.connect.ureplicator3;
+
+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.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Map;
+
+public class UReplicator3TaskConfig extends UReplicator3Config {
+
+ private static final Logger log = LoggerFactory.getLogger(UReplicator3TaskConfig.class);
+
+ public static final String TASK_ID_CONFIG = "task.id";
+
+ @SuppressWarnings("this-escape")
+ public UReplicator3TaskConfig(Map<String, ?> props) {
+ super(config(), props);
+ }
+
+ protected static ConfigDef config() {
+ return baseConfig()
+ .defineInternal(
+ TASK_ID_CONFIG,
+ Type.INT,
+ ConfigDef.NO_DEFAULT_VALUE,
+ Importance.HIGH
+ );
+ }
+
+ @Override
+ protected String role() {
+ return "task-" + connectorName() + "-" + taskId();
+ }
+
+ public int taskId() {
+ return getInt(TASK_ID_CONFIG);
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignedPartition.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignedPartition.java
new file mode 100644
index 00000000000..5cc6c77604b
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignedPartition.java
@@ -0,0 +1,27 @@
+package com.uber.data.kafka.connect.ureplicator3.coordination;
+
+import org.apache.kafka.common.TopicPartition;
+
+public record AssignedPartition(
+ String srcTopic,
+ int srcPartition,
+ String dstTopic,
+ Integer dstPartition
+) {
+
+ public TopicPartition source() {
+ return new TopicPartition(srcTopic, srcPartition);
+ }
+
+ public TopicPartition target() {
+ return new TopicPartition(
+ dstTopic != null ? dstTopic : srcTopic,
+ dstPartition != null ? dstPartition : srcPartition
+ );
+ }
+
+ public boolean hasOverride() {
+ return !source().equals(target());
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentListener.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentListener.java
new file mode 100644
index 00000000000..3058d3d45ad
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentListener.java
@@ -0,0 +1,11 @@
+package com.uber.data.kafka.connect.ureplicator3.coordination;
+
+public interface AssignmentListener {
+
+ /**
+ * Invoked when the assignment for the task has been updated. Long-running or blocking
+ * operations should <strong>not</strong> be performed inside this method
+ */
+ void onUpdate();
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentStore.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentStore.java
new file mode 100644
index 00000000000..2e2c413a537
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentStore.java
@@ -0,0 +1,368 @@
+package com.uber.data.kafka.connect.ureplicator3.coordination;
+
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.internals.KafkaCompletableFuture;
+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.common.utils.Utils;
+import org.apache.kafka.connect.util.KafkaBasedLog;
+import org.apache.kafka.connect.util.TopicAdmin;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.LinkedHashMap;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
+import java.util.stream.Collectors;
+
+import static com.uber.data.kafka.connect.utils.SerializationUtils.parseMapField;
+import static com.uber.data.kafka.connect.utils.SerializationUtils.parseNumberField;
+import static com.uber.data.kafka.connect.utils.SerializationUtils.parseOptionalNumberField;
+import static com.uber.data.kafka.connect.utils.SerializationUtils.parseOptionalStringField;
+import static com.uber.data.kafka.connect.utils.SerializationUtils.parseStringField;
+
+// This class is thread-safe
+public class AssignmentStore implements Closeable {
+
+ private static final Logger log = LoggerFactory.getLogger(AssignmentStore.class);
+ private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+ private static final String ASSIGNMENT_FIELD = "assignment-v1";
+ private static final String CONNECTOR_FIELD = "connector";
+ private static final String SRC_KAFKA_FIELD = "src";
+ private static final String DST_KAFKA_FIELD = "dst";
+ private static final String TOPIC_FIELD = "topic";
+ private static final String PARTITION_FIELD = "partition";
+
+ private static final String TASK_FIELD = "task";
+ private static final String DST_TOPIC_FIELD = "dst-topic";
+ private static final String DST_PARTITION_FIELD = "dst-partition";
+
+ private final String topic;
+ private final int topicPartitions;
+ private final short topicReplicationFactor;
+ private final Map<String, Object> topicSettings;
+ private final Map<String, Object> producerProps;
+ private final Map<String, Object> consumerProps;
+ private final Map<String, Object> adminProps;
+ private final Map<String, Map<ReplicatedTopicPartition, AssignmentValue>> taskPartitionAssignments;
+ private final Map<String, Map<Integer, AssignmentListener>> assignmentListeners;
+
+ private TopicAdmin topicAdmin;
+ private KafkaBasedLog<byte[], byte[]> kafkaLog;
+
+ public AssignmentStore(AssignmentStoreConfig config) {
+ this.topic = config.storageTopic();
+ this.topicPartitions = config.storageTopicPartitions();
+ this.topicReplicationFactor = config.storageTopicReplicationFactor();
+ this.topicSettings = config.storageTopicSettings();
+ this.producerProps = producerProps(config.producerConfig());
+ this.consumerProps = consumerProps(config.consumerConfig());
+ this.adminProps = config.adminConfig();
+ this.taskPartitionAssignments = new HashMap<>();
+ this.assignmentListeners = new HashMap<>();
+ }
+
+ public void start() {
+ this.topicAdmin = new TopicAdmin(adminProps);
+ this.kafkaLog = new KafkaBasedLog<>(
+ topic,
+ producerProps,
+ consumerProps,
+ () -> topicAdmin,
+ this::onConsume,
+ Time.SYSTEM,
+ this::createAndValidateTopic
+ );
+ kafkaLog.start();
+ }
+
+ private static Map<String, Object> producerProps(Map<String, Object> baseProps) {
+ baseProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
+ baseProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
+ baseProps.put(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG, Integer.MAX_VALUE);
+ return baseProps;
+ }
+
+ private static Map<String, Object> consumerProps(Map<String, Object> baseProps) {
+ baseProps.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
+ baseProps.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
+ return baseProps;
+ }
+
+ private void createAndValidateTopic(TopicAdmin topicAdmin) {
+ NewTopic newTopic = TopicAdmin.defineTopic(topic)
+ .config(topicSettings)
+ .compacted()
+ .partitions(topicPartitions)
+ .replicationFactor(topicReplicationFactor)
+ .build();
+ boolean created = topicAdmin.createTopic(newTopic);
+ if (!created) {
+ // Ensure existing topic is compacted
+ boolean verified = topicAdmin.verifyTopicCleanupPolicyOnlyCompact(
+ topic,
+ AssignmentStoreConfig.STORAGE_TOPIC_CONFIG,
+ "task partition assignments"
+ );
+ // Should never, NEVER happen, but just in case
+ if (!verified) {
+ throw new KafkaException("Failed to verify that task partition assignment storage topic " + topic + " is compacted");
+ }
+ }
+ }
+
+ @Override
+ public void close() {
+ if (kafkaLog != null) {
+ Utils.closeQuietly(kafkaLog::stop, "Kafka-based log");
+ }
+ Utils.closeQuietly(topicAdmin, "topic admin");
+ }
+
+ public Future<?> assignPartition(String connector, ReplicatedTopicPartition replicatedTopicPartition, int taskId, String dstTopic, Integer dstPartition) {
+ return publishAssignment(connector, replicatedTopicPartition, dstTopic, dstPartition, taskId);
+ }
+
+ public Future<?> revokePartition(String connector, ReplicatedTopicPartition replicatedTopicPartition) {
+ return publishAssignment(connector, replicatedTopicPartition, null, null, null);
+ }
+
+ public Map<Integer, ? extends Collection<AssignedPartition>> getAssignments(String connector, String srcKafka, String dstKafka) {
+ Objects.requireNonNull(connector);
+ Objects.requireNonNull(srcKafka);
+ Objects.requireNonNull(dstKafka);
+
+ Map<ReplicatedTopicPartition, AssignmentValue> connectorSnapshot;
+ synchronized (this) {
+ connectorSnapshot = new HashMap<>(taskPartitionAssignments.getOrDefault(connector, Map.of()));
+ }
+
+ return connectorSnapshot.entrySet().stream()
+ .filter(e -> e.getKey().matchesPipeline(srcKafka, dstKafka))
+ .collect(Collectors.groupingBy(
+ e -> e.getValue().taskId(),
+ Collectors.mapping(AssignmentStore::assignedPartition, Collectors.toSet())
+ ));
+ }
+
+ public Collection<AssignedPartition> getAssignment(String connector, int taskId, String srcKafka, String dstKafka) {
+ synchronized (this) {
+ return taskPartitionAssignments.getOrDefault(connector, Map.of()).entrySet().stream()
+ .filter(e -> e.getKey().matchesPipeline(srcKafka, dstKafka))
+ .filter(e -> taskId == e.getValue().taskId())
+ .map(AssignmentStore::assignedPartition)
+ .collect(Collectors.toSet());
+ }
+ }
+
+ // TODO: Add an API to remove listeners as well (have to be careful not to overwrite listeners on slow task shutdown, though)
+ public void addListener(String connector, int taskId, AssignmentListener listener) {
+ synchronized (this) {
+ assignmentListeners.computeIfAbsent(connector, k -> new HashMap<>())
+ .put(taskId, listener);
+ }
+ }
+
+ public void refresh() {
+ try {
+ kafkaLog.readToEnd().get();
+ } catch (InterruptedException | ExecutionException e) {
+ throw new KafkaException("Failed to read to end of task partition assignment log", e);
+ }
+ }
+
+ // For testing only
+ public void clear() {
+ taskPartitionAssignments.clear();
+ }
+
+ private Future<?> publishAssignment(String connector, ReplicatedTopicPartition replicatedTopicPartition, String dstTopic, Integer dstPartition, Integer taskId) {
+ if (connector == null) {
+ log.error("Failed to serialize assignment key with null connector name");
+ return KafkaCompletableFuture.failedFuture(new NullPointerException("Connector may not be null"));
+ }
+ if (replicatedTopicPartition == null) {
+ log.error("Failed to serialize assignment key with null replicated topic partition");
+ return KafkaCompletableFuture.failedFuture(new NullPointerException("Replicated topic partition may not be null"));
+ }
+
+ byte[] key;
+ try {
+ key = serializeAssignmentKey(connector, replicatedTopicPartition);
+ } catch (JsonProcessingException e) {
+ log.error("Failed to serialize assignment key for connector {} and replicated topic partition {} as JSON", connector, replicatedTopicPartition, e);
+ return KafkaCompletableFuture.failedFuture(e);
+ }
+
+ byte[] value;
+ if (taskId == null) {
+ value = null;
+ } else {
+ try {
+ value = serializeAssignmentValue(taskId, dstTopic, dstPartition);
+ } catch (JsonProcessingException e) {
+ log.error(
+ "Failed to serialize assignment value (task={}, dstTopic={}, dstPartition={}) for connector {} and replicated topic partition {} as JSON",
+ taskId, dstTopic, dstPartition, connector, replicatedTopicPartition, e);
+ return KafkaCompletableFuture.failedFuture(e);
+ }
+ }
+
+ // TODO: Error handling (probably just log and emit a metric)
+ return kafkaLog.sendWithReceipt(key, value);
+ }
+
+ private void onConsume(Throwable error, ConsumerRecord<byte[], byte[]> record) {
+ if (error != null) {
+ log.error("Error occurred while consuming from task partition assignment topic", error);
+ return;
+ }
+
+ int recordPartition = record.partition();
+ long recordOffset = record.offset();
+
+ if (record.key() == null) {
+ log.error("Ignoring record with null key {} in partition {} at offset {}", record.key(), recordPartition, recordOffset);
+ return;
+ }
+
+ AssignmentKey assignmentKey;
+ try {
+ assignmentKey = deserializeAssignmentKey(record.key());
+ } catch (Exception e) {
+ log.error("Failed to deserialize task partition assignment key in partition {} at offset {} as JSON", recordPartition, recordOffset, e);
+ return;
+ }
+
+ AssignmentValue assignmentValue;
+ try {
+ assignmentValue = deserializeAssignmentValue(record.value());
+ } catch (IOException e) {
+ log.error("Failed to deserialize task partition assignment value in partition {} at offset {} for assignment key {} as JSON", recordPartition, recordOffset, assignmentKey, e);
+ return;
+ }
+
+ // TODO: Notify listeners
+ String connector = assignmentKey.connector();
+ ReplicatedTopicPartition replicatedTopicPartition = assignmentKey.replicatedTopicPartition();
+ if (assignmentValue == null) {
+ log.trace("Picked up revocation of {} for connector {}", replicatedTopicPartition, connector);
+ synchronized (this) {
+ Map<ReplicatedTopicPartition, AssignmentValue> connectorAssignments = taskPartitionAssignments.get(connector);
+ if (connectorAssignments == null) {
+ log.debug("No cached assignments for connector {}", connector);
+ } else {
+ AssignmentValue priorAssignment = connectorAssignments.remove(replicatedTopicPartition);
+ if (priorAssignment != null) {
+ notifyListener(connector, priorAssignment.taskId());
+ }
+ }
+ }
+ } else {
+ log.trace("Picked up assignment of {} with destination topic {} and destination partition {} to task {} for connector {}",
+ replicatedTopicPartition, assignmentValue.dstTopic(), assignmentValue.dstPartition(), assignmentValue.taskId(), connector);
+ synchronized (this) {
+ Map<ReplicatedTopicPartition, AssignmentValue> connectorAssignments = taskPartitionAssignments.computeIfAbsent(connector, k -> new HashMap<>());
+ AssignmentValue priorAssignment = connectorAssignments.put(replicatedTopicPartition, assignmentValue);
+ if (priorAssignment != null) {
+ notifyListener(connector, priorAssignment.taskId());
+ }
+ notifyListener(connector, assignmentValue.taskId());
+ }
+ }
+ }
+
+ private byte[] serializeAssignmentKey(String connector, ReplicatedTopicPartition replicatedTopicPartition) throws JsonProcessingException {
+ // LinkedHashMap preserves order and allows us to deterministically generate a serialized key from a replicated topic partition
+ Map<String, Object> innerMap = new LinkedHashMap<>();
+ innerMap.put(CONNECTOR_FIELD, connector);
+ innerMap.put(SRC_KAFKA_FIELD, replicatedTopicPartition.srcKafka());
+ innerMap.put(DST_KAFKA_FIELD, replicatedTopicPartition.dstKafka());
+ innerMap.put(TOPIC_FIELD, replicatedTopicPartition.topicPartition().topic());
+ innerMap.put(PARTITION_FIELD, replicatedTopicPartition.topicPartition().partition());
+ Map<String, Object> wrapped = Map.of(ASSIGNMENT_FIELD, innerMap);
+ return OBJECT_MAPPER.writeValueAsBytes(wrapped);
+ }
+
+ private byte[] serializeAssignmentValue(int taskId, String dstTopic, Integer dstPartition) throws JsonProcessingException {
+ Map<String, Object> map = new LinkedHashMap<>();
+ map.put(TASK_FIELD, taskId);
+ if (dstTopic != null)
+ map.put(DST_TOPIC_FIELD, dstTopic);
+ if (dstPartition != null)
+ map.put(DST_PARTITION_FIELD, dstPartition);
+ return OBJECT_MAPPER.writeValueAsBytes(map);
+ }
+
+ private AssignmentKey deserializeAssignmentKey(byte[] key) throws IOException {
+ if (key == null)
+ throw new IOException("Key should not be null");
+
+ Map<String, Object> map = OBJECT_MAPPER.readValue(key, new TypeReference<>() { });
+ Map<?, ?> wrapped = parseMapField(true, map, ASSIGNMENT_FIELD);
+
+ String connector = parseStringField(true, wrapped, CONNECTOR_FIELD);
+ String srcKafka = parseStringField(true, wrapped, SRC_KAFKA_FIELD);
+ String dstKafka = parseStringField(true, wrapped, DST_KAFKA_FIELD);
+ String topic = parseStringField(true, wrapped, TOPIC_FIELD);
+ Number partition = parseNumberField(true, wrapped, PARTITION_FIELD);
+ ReplicatedTopicPartition replicatedTopicPartition = new ReplicatedTopicPartition(srcKafka, dstKafka, topic, partition.intValue());
+
+ return new AssignmentKey(connector, replicatedTopicPartition);
+ }
+
+ private AssignmentValue deserializeAssignmentValue(byte[] value) throws IOException {
+ // Permit null values to signify revocation of the partition
+ if (value == null)
+ return null;
+
+ Map<String, Object> map = OBJECT_MAPPER.readValue(value, new TypeReference<>() { });
+ int taskId = parseNumberField(false, map, TASK_FIELD).intValue();
+ String dstTopic = parseOptionalStringField(false, map, DST_TOPIC_FIELD);
+ Number rawDstPartition = parseOptionalNumberField(false, map, DST_PARTITION_FIELD);
+ Integer dstPartition = rawDstPartition != null ? rawDstPartition.intValue() : null;
+ return new AssignmentValue(taskId, dstTopic, dstPartition);
+ }
+
+ private void notifyListener(String connector, int taskId) {
+ AssignmentListener listener = assignmentListeners.getOrDefault(connector, Map.of())
+ .get(taskId);
+
+ if (listener == null)
+ return;
+
+ try {
+ listener.onUpdate();
+ } catch (Throwable t) {
+ log.error("Listener for {}-{} threw an exception", connector, taskId, t);
+ }
+ }
+
+ private static AssignedPartition assignedPartition(Map.Entry<ReplicatedTopicPartition, AssignmentValue> entry) {
+ return new AssignedPartition(entry.getKey().topic(), entry.getKey().partition(), entry.getValue().dstTopic(), entry.getValue().dstPartition);
+ }
+
+ private record AssignmentKey(String connector, ReplicatedTopicPartition replicatedTopicPartition) {
+ }
+
+ private record AssignmentValue(int taskId, String dstTopic, Integer dstPartition) {
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentStoreConfig.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentStoreConfig.java
new file mode 100644
index 00000000000..27165db2860
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentStoreConfig.java
@@ -0,0 +1,145 @@
+package com.uber.data.kafka.connect.ureplicator3.coordination;
+
+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.TopicConfig;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
+
+public class AssignmentStoreConfig extends AbstractConfig {
+
+ private static final Logger log = LoggerFactory.getLogger(AssignmentStoreConfig.class);
+
+ private static final String CONFIG_PREFIX = "uber.ureplicator3.assignments.";
+
+ public static final String STORAGE_TOPIC_CONFIG = CONFIG_PREFIX + "storage.topic";
+ public static final String STORAGE_TOPIC_PREFIX = STORAGE_TOPIC_CONFIG + ".";
+ private static final String STORAGE_TOPIC_DOC = "Topic used to store assignments. "
+ + "This topic must be compacted (i.e., have the 'compacted' cleanup policy) and can have "
+ + "any number of partitions, but it's recommended to use a low number (e.g., 5 or 10). "
+ + "If the topic does not exist, it will be automatically created on startup. "
+ + "Additional properties for the topic can be specified using the " + STORAGE_TOPIC_PREFIX + " "
+ + "prefix";
+
+ private static final String PARTITIONS_SUFFIX = "partitions";
+ public static final String STORAGE_TOPIC_PARTITIONS_CONFIG = STORAGE_TOPIC_PREFIX + PARTITIONS_SUFFIX;
+ private static final int STORAGE_TOPIC_PARTITIONS_DEFAULT = 10;
+ private static final String STORAGE_TOPIC_PARTITIONS_DOC = "The number of partitions to create "
+ + "the assignments topic with, if it does not already exist. If the topic already exists this property "
+ + "will be ignored, even if the topic has a different number of partitions";
+
+ private static final String REPLICATION_FACTOR_SUFFIX = "replication.factor";
+ public static final String STORAGE_TOPIC_REPLICATION_FACTOR_CONFIG = STORAGE_TOPIC_PREFIX + REPLICATION_FACTOR_SUFFIX;
+ private static final short STORAGE_TOPIC_REPLICATION_FACTOR_DEFAULT = (short) 3;
+ private static final String STORAGE_TOPIC_REPLICATION_FACTOR_DOC = "The replication factor to create "
+ + "the assignments topic with, if it does not already exist. If the topic already exists this property "
+ + "will be ignored, even if the topic has a different replication factor";
+
+ private static final String SEGMENT_BYTES_SUFFIX = TopicConfig.SEGMENT_BYTES_CONFIG;
+ public static final String STORAGE_TOPIC_SEGMENT_BYTES_CONFIG = STORAGE_TOPIC_PREFIX + SEGMENT_BYTES_SUFFIX;
+ private static final int STORAGE_TOPIC_SEGMENT_BYTES_DEFAULT = 1024 * 1024 * 3; // Lower default segment size in order to enable more-frequent compaction
+ private static final String STORAGE_TOPIC_SEGMENT_BYTES_DOC = "The per-partition segment size "
+ + "for the assignments topic. This should be relatively low in order to enable topic compaction to take place, which "
+ + "will help keep the topic from growing too large and causing slow startup time for workers. Like other topic "
+ + "properties, if the topic already exists, this property will be ignored, even if the topic has a different segment size";
+
+ public static final String CLIENTS_PREFIX = CONFIG_PREFIX + "kafka.clients";
+ public static final String CONSUMER_PREFIX = CONFIG_PREFIX + "consumer";
+ public static final String PRODUCER_PREFIX = CONFIG_PREFIX + "producer";
+ public static final String ADMIN_PREFIX = CONFIG_PREFIX + "admin";
+
+ public AssignmentStoreConfig(Map<String, ?> props) {
+ super(config(), props);
+ }
+
+ public static ConfigDef config() {
+ return addToConfig(new ConfigDef());
+ }
+
+ public static ConfigDef addToConfig(ConfigDef base) {
+ return base
+ .define(
+ STORAGE_TOPIC_CONFIG,
+ Type.STRING,
+ Importance.HIGH,
+ STORAGE_TOPIC_DOC
+ ).define(
+ STORAGE_TOPIC_PARTITIONS_CONFIG,
+ Type.INT,
+ STORAGE_TOPIC_PARTITIONS_DEFAULT,
+ atLeast(1),
+ Importance.LOW,
+ STORAGE_TOPIC_PARTITIONS_DOC
+ ).define(
+ STORAGE_TOPIC_REPLICATION_FACTOR_CONFIG,
+ Type.SHORT,
+ STORAGE_TOPIC_REPLICATION_FACTOR_DEFAULT,
+ atLeast(1),
+ Importance.LOW,
+ STORAGE_TOPIC_REPLICATION_FACTOR_DOC
+ ).define(
+ STORAGE_TOPIC_SEGMENT_BYTES_CONFIG,
+ Type.INT,
+ STORAGE_TOPIC_SEGMENT_BYTES_DEFAULT,
+ atLeast(1),
+ Importance.LOW,
+ STORAGE_TOPIC_SEGMENT_BYTES_DOC
+ );
+ }
+
+ public Map<String, Object> adminConfig() {
+ return clientConfig(ADMIN_PREFIX);
+ }
+
+ public Map<String, Object> producerConfig() {
+ return clientConfig(PRODUCER_PREFIX);
+ }
+
+ public Map<String, Object> consumerConfig() {
+ return clientConfig(CONSUMER_PREFIX);
+ }
+
+ public String storageTopic() {
+ return getString(STORAGE_TOPIC_CONFIG);
+ }
+
+ public int storageTopicPartitions() {
+ return getInt(STORAGE_TOPIC_PARTITIONS_CONFIG);
+ }
+
+ public short storageTopicReplicationFactor() {
+ return getShort(STORAGE_TOPIC_REPLICATION_FACTOR_CONFIG);
+ }
+
+ public Map<String, Object> storageTopicSettings() {
+ Map<String, Object> result = originalsWithPrefix(STORAGE_TOPIC_PREFIX);
+
+ // AbstractConfig::originalsWithPrefix doesn't include default values, we apply that manually here
+ result.computeIfAbsent(TopicConfig.SEGMENT_BYTES_CONFIG, k -> STORAGE_TOPIC_SEGMENT_BYTES_DEFAULT);
+
+ Object removedPolicy = result.remove(TopicConfig.CLEANUP_POLICY_CONFIG);
+ if (removedPolicy != null) {
+ log.warn("Ignoring '{}cleanup.policy={}' setting, since compaction is always used", STORAGE_TOPIC_PREFIX, removedPolicy);
+ }
+
+ result.remove(REPLICATION_FACTOR_SUFFIX);
+ result.remove(PARTITIONS_SUFFIX);
+
+ return result;
+ }
+
+ private Map<String, Object> clientConfig(String specialPrefix) {
+ Map<String, Object> result = new HashMap<>(originals());
+ result.putAll(originalsWithPrefix(CLIENTS_PREFIX));
+ result.putAll(originalsWithPrefix(specialPrefix));
+ return result;
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentStores.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentStores.java
new file mode 100644
index 00000000000..f58ebc830c9
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentStores.java
@@ -0,0 +1,39 @@
+package com.uber.data.kafka.connect.ureplicator3.coordination;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class AssignmentStores {
+
+ private static final Logger log = LoggerFactory.getLogger(AssignmentStores.class);
+
+ private static AssignmentStore store;
+
+ public static synchronized void initialize(AssignmentStoreConfig config) {
+ if (store == null)
+ store = initializeStore(config);
+ else
+ log.warn("Store is already initialized; this should never happen outside an embedded testing environment");
+ }
+
+ // TODO: Instead of exposing the whole store, we could expose a safe set of operations
+ // like assign, revoke, and addListener here (this would prevent callers from doing
+ // something stupid like trying to close the store)
+ public static AssignmentStore get() {
+ if (store == null) {
+ throw new IllegalStateException("Store has not been initialized yet");
+ }
+ return store;
+ }
+
+ // TODO: private void destroyStore()
+
+ private static AssignmentStore initializeStore(AssignmentStoreConfig config) {
+ log.info("Initializing assignment store");
+ AssignmentStore result = new AssignmentStore(config);
+ result.start();
+ log.info("Finished initializing assignment store");
+ return result;
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/BrokerResolver.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/BrokerResolver.java
new file mode 100644
index 00000000000..0933ebbe7e9
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/BrokerResolver.java
@@ -0,0 +1,39 @@
+package com.uber.data.kafka.connect.utils;
+
+
+import com.uber.data.kafka.common.ClusterNameHelper;
+import com.uber.data.kafka.common.ClusterNameHelper.ClusterNameParseResult;
+import com.uber.data.kafka.exception.ResolverException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+// TODO: Singleton pattern and fault injection in dependent classes
+public final class BrokerResolver {
+
+ private static final Logger log = LoggerFactory.getLogger(BrokerResolver.class);
+
+ // TODO: We should expose normalization logic so that different aliases for the same Kafka
+ // cluster get treated the same (e.g., when storing offsets or logging replication pipelines)
+
+ public static String resolveBootstrap(String cluster) throws ResolverException {
+ if (cluster == null) {
+ throw new ResolverException("Cannot parse null cluster name");
+ }
+
+ ClusterNameParseResult parseResult;
+ try {
+ parseResult = ClusterNameHelper.parseClusterName(cluster);
+ } catch (IllegalArgumentException e) {
+ throw new ResolverException("Failed to parse cluster name '" + cluster + "'", e);
+ }
+ // TODO: There has to be a better way 💀
+ return String.format(
+ "kafka-%s.%s.uber.internal:9092",
+ parseResult.functionalGroup,
+ parseResult.region
+ );
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/ExecutorUtils.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/ExecutorUtils.java
new file mode 100644
index 00000000000..74f1517550e
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/ExecutorUtils.java
@@ -0,0 +1,88 @@
+package com.uber.data.kafka.connect.utils;
+
+import com.google.common.util.concurrent.ThreadFactoryBuilder;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ThreadFactory;
+import java.util.concurrent.TimeUnit;
+
+
+/**
+ * This is a simple wrapper around the standard library and Guava's thread utils that can
+ * be used to create and work with executors.
+ *
+ * @see Executors
+ * @see ExecutorService
+ * @see ScheduledExecutorService
+ * @see ThreadFactory
+ * @see ThreadFactoryBuilder
+ */
+public class ExecutorUtils {
+ private static final Logger log = LoggerFactory.getLogger(ExecutorUtils.class);
+
+ private ExecutorUtils() {
+ }
+
+ /**
+ * Create a new executor that uses the given thread name format and for which
+ * all created threads are daemons
+ *
+ * @param nameFormat the name format for executor threads
+ * @return the executor; never null
+ */
+ public static ScheduledExecutorService singleThreadedScheduled(String nameFormat) {
+ return Executors.newSingleThreadScheduledExecutor(threadFactory(nameFormat));
+ }
+
+ /**
+ * Shut down the executor, waiting for the specified timeout for all running
+ * and scheduled tasks to complete
+ *
+ * @param executor the executor to shut down; may not be null
+ * @param executorDescription a description of the executor to use in log messages;
+ * for example, "offset monitor" or "topic watcher service"
+ * @param timeout the time to wait for all tasks to complete
+ * @param unit the unit for the timeout
+ */
+ public static void shutdown(
+ ExecutorService executor,
+ String executorDescription,
+ long timeout,
+ TimeUnit unit
+ ) {
+ executor.shutdown();
+ boolean shutdownSucceeded = false;
+ boolean interrupted = false;
+ try {
+ shutdownSucceeded = executor.awaitTermination(timeout, unit);
+ if (!shutdownSucceeded) {
+ log.warn("Could not shut down {} in time", executorDescription);
+ } else {
+ log.info("Gracefully shut down {}", executorDescription);
+ }
+ } catch (InterruptedException e) {
+ log.warn("Interrupted while shutting down {}", executorDescription, e);
+ interrupted = true;
+ } finally {
+ if (!shutdownSucceeded) {
+ executor.shutdownNow();
+ }
+
+ // Preserve the interrupt status when control is returned to the caller
+ if (interrupted)
+ Thread.currentThread().interrupt();
+ }
+ }
+
+ private static ThreadFactory threadFactory(String nameFormat) {
+ return new ThreadFactoryBuilder()
+ .setNameFormat(nameFormat)
+ .setDaemon(true)
+ .build();
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/M3ReporterProvider.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/M3ReporterProvider.java
new file mode 100644
index 00000000000..1538ae9d42b
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/M3ReporterProvider.java
@@ -0,0 +1,43 @@
+package com.uber.data.kafka.connect.utils;
+
+import com.uber.m3.tally.RootScopeBuilder;
+import com.uber.m3.tally.Scope;
+import com.uber.m3.tally.m3.M3Reporter;
+import com.uber.m3.util.Duration;
+import com.uber.m3.util.ImmutableMap;
+
+import java.net.InetSocketAddress;
+import java.net.SocketAddress;
+
+/**
+ * Singleton class to provide M3 reporter object
+ */
+public class M3ReporterProvider {
+
+ public static final SocketAddress SOCKET_ADDRESS = new InetSocketAddress("localhost", 9052);
+ public static final String RUNTIME_ENV_TAG = "runtime_env";
+ public static final String TARGET_ENVIRONMENT = "target_environment";
+
+ public static final Scope SCOPE = new RootScopeBuilder()
+ .reporter(new M3Reporter.Builder(SOCKET_ADDRESS)
+ .includeHost(false)
+ .commonTags(buildCommonTags())
+ .build())
+ .reportEvery(Duration.ofSeconds(10));
+
+ public static final Scope HOST_SCOPE = new RootScopeBuilder()
+ .reporter(new M3Reporter.Builder(SOCKET_ADDRESS)
+ .includeHost(true)
+ .commonTags(buildCommonTags())
+ .build())
+ .reportEvery(Duration.ofSeconds(10));
+
+ private static ImmutableMap<String, String> buildCommonTags() {
+ ImmutableMap.Builder<String, String> tags = new ImmutableMap.Builder<>();
+ tags.put(M3Reporter.SERVICE_TAG, System.getenv("UDEPLOY_SERVICE_NAME"));
+ tags.put(M3Reporter.ENV_TAG, System.getenv("UBER_ENVIRONMENT"));
+ tags.put(RUNTIME_ENV_TAG, System.getenv("UBER_RUNTIME_ENVIRONMENT"));
+ tags.put(TARGET_ENVIRONMENT, System.getenv("UBER_RUNTIME_ENVIRONMENT"));
+ return tags.build();
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/MetricsUtils.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/MetricsUtils.java
new file mode 100644
index 00000000000..064e12b43fe
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/MetricsUtils.java
@@ -0,0 +1,37 @@
+package com.uber.data.kafka.connect.utils;
+
+import org.apache.kafka.common.TopicPartition;
+
+import com.google.common.collect.ImmutableMap;
+import com.uber.m3.tally.Scope;
+
+import java.util.Map;
+
+public class MetricsUtils {
+
+ public static final MetricsUtils INSTANCE = new MetricsUtils();
+
+ private static final String TOPIC = "topic";
+ private static final String PARTITION = "partition";
+ private static final String SRC_CLUSTER = "srcCluster";
+ private static final String DST_CLUSTER = "dstCluster";
+
+ private MetricsUtils() {
+ }
+
+ public Scope routeScope(String srcKafka, String dstKafka) {
+ return M3ReporterProvider.SCOPE
+ .tagged(ImmutableMap.of(
+ SRC_CLUSTER, srcKafka,
+ DST_CLUSTER, dstKafka
+ ));
+ }
+
+ public Map<String, String> tags(TopicPartition topicPartition) {
+ return ImmutableMap.of(
+ TOPIC, topicPartition.topic(),
+ PARTITION, Integer.toString(topicPartition.partition())
+ );
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/SerializationUtils.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/SerializationUtils.java
new file mode 100644
index 00000000000..60cd5c91418
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/SerializationUtils.java
@@ -0,0 +1,78 @@
+package com.uber.data.kafka.connect.utils;
+
+import java.io.IOException;
+import java.util.Map;
+
+public final class SerializationUtils {
+
+ public static Map<?, ?> parseMapField(boolean isKey, Map<?, ?> map, String field) throws IOException {
+ Object fieldValue = map.get(field);
+ if (fieldValue == null)
+ throw new MissingFieldException(isKey, field);
+
+ if (!(fieldValue instanceof Map))
+ throw new InvalidFieldTypeException(isKey, field, "map", fieldValue);
+
+ return (Map<?, ?>) fieldValue;
+ }
+
+ public static String parseStringField(boolean isKey, Map<?, ?> map, String field) throws IOException {
+ String result = parseOptionalStringField(isKey, map, field);
+
+ if (result == null)
+ throw new MissingFieldException(isKey, field);
+
+ return result;
+ }
+
+ public static String parseOptionalStringField(boolean isKey, Map<?, ?> map, String field) throws IOException {
+ Object fieldValue = map.get(field);
+ if (fieldValue == null)
+ return null;
+
+ if (!(fieldValue instanceof String))
+ throw new InvalidFieldTypeException(isKey, field, "string", fieldValue);
+
+ return (String) fieldValue;
+ }
+
+ public static Number parseNumberField(boolean isKey, Map<?, ?> map, String field) throws IOException {
+ Number result = parseOptionalNumberField(isKey, map, field);
+
+ if (result == null)
+ throw new MissingFieldException(isKey, field);
+
+ return result;
+ }
+
+ public static Number parseOptionalNumberField(boolean isKey, Map<?, ?> map, String field) throws IOException {
+ Object fieldValue = map.get(field);
+ if (fieldValue == null)
+ return null;
+
+ if (!(fieldValue instanceof Number))
+ throw new InvalidFieldTypeException(isKey, field, "string", fieldValue);
+
+ return (Number) fieldValue;
+ }
+
+ public static class MissingFieldException extends IOException {
+ public MissingFieldException(boolean isKey, String field) {
+ super(
+ "record has invalid " + (isKey ? "key" : "value")
+ + ": missing '" + field + "' field"
+ );
+ }
+ }
+
+ public static class InvalidFieldTypeException extends IOException {
+ public InvalidFieldTypeException(boolean isKey, String field, String expectedType, Object actualValue) {
+ super(
+ "record has invalid " + (isKey ? "key" : "value")
+ + ": field '" + field + "' should be a(n) " + expectedType
+ + "; was " + (actualValue == null ? "null" : actualValue.getClass().getSimpleName())
+ );
+ }
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/UberConnectUtils.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/UberConnectUtils.java
new file mode 100644
index 00000000000..c2e7eb0af22
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/UberConnectUtils.java
@@ -0,0 +1,205 @@
+package com.uber.data.kafka.connect.utils;
+
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.DescribeTopicsOptions;
+import org.apache.kafka.clients.admin.TopicDescription;
+import org.apache.kafka.clients.admin.TopicListing;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.config.Config;
+import org.apache.kafka.common.config.ConfigValue;
+import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.common.utils.Timer;
+
+import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
+import com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig;
+import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStore;
+import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStores;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collection;
+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.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
+
+import static org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLASS_CONFIG;
+
+public final class UberConnectUtils {
+
+ private static final Logger log = LoggerFactory.getLogger(UberConnectUtils.class);
+
+
+ @SafeVarargs
+ public static <E> Set<E> diff(Collection<E> base, Collection<E>... toSubtract) {
+ Set<E> result = new HashSet<>(base);
+ for (Collection<E> subtracted : toSubtract) {
+ if (subtracted != null)
+ result.removeAll(subtracted);
+ }
+ return result;
+ }
+
+ public static Set<String> listTopics(
+ Admin admin, long timeout, TimeUnit timeUnit
+ ) throws InterruptedException, TimeoutException, ExecutionException {
+ return admin.listTopics().listings().get(timeout, timeUnit).stream()
+ .filter(Predicate.not(TopicListing::isInternal))
+ .map(TopicListing::name)
+ .collect(Collectors.toSet());
+ }
+
+ public static Map<String, Integer> listPartitions(Collection<String> topics, Admin admin, long timeout, TimeUnit timeUnit) {
+ int timeoutMs = (int) Math.max(Integer.MAX_VALUE, timeUnit.toMillis(timeout));
+ DescribeTopicsOptions options = new DescribeTopicsOptions()
+ .includeAuthorizedOperations(false) // Don't care about ACLS
+ .timeoutMs(timeoutMs);
+
+ Map<String, Integer> result = new HashMap<>();
+ admin.describeTopics(topics, options).topicNameValues().forEach((topic, topicDescriptionFuture) -> {
+ TopicDescription topicDescription;
+ try {
+ topicDescription = topicDescriptionFuture.get();
+ } catch (InterruptedException | ExecutionException e) {
+ log.warn("Unable to describe topic {}", topic, e);
+ return;
+ }
+
+ result.put(topic, topicDescription.partitions().size());
+ });
+
+ return result;
+ }
+
+ /**
+ * Extract and return a record's {@link ConsumerRecord#timestamp() timestamp}, if it has one
+ * @param record the record; may not be null
+ * @return an {@link Optional} containing the record's timestamp, or {@link Optional#empty()}
+ * if the record does not have a timestamp
+ */
+ public static Optional<Long> extractTimestamp(ConsumerRecord<?, ?> record) {
+ return Optional.of(record)
+ .filter(r -> !TimestampType.NO_TIMESTAMP_TYPE.equals(record.timestampType()))
+ .map(ConsumerRecord::timestamp)
+ .filter(ts -> ts >= 0);
+ }
+
+ public static String createErrorMessage(String base, Throwable t) {
+ String separator = ": ";
+ String result = base;
+ while (t != null) {
+ if (t.getMessage() != null)
+ result += separator + t.getMessage();
+ t = t.getCause();
+ }
+
+ return result;
+ }
+
+ public static void addError(Config configValidationResult, String prop, Object value, String baseErrorMessage, Throwable error) {
+ addError(configValidationResult, prop, value, createErrorMessage(baseErrorMessage, error));
+ }
+
+ public static void addError(Config configValidationResult, String prop, Object value, String errorMessage) {
+ findOrCreate(configValidationResult, prop, value)
+ .addErrorMessage(errorMessage);
+ }
+
+ public static ConfigValue findOrCreate(Config configValidationResult, String prop, Object value) {
+ for (ConfigValue definedValue : configValidationResult.configValues()) {
+ if (Objects.equals(prop, definedValue.name())) {
+ return definedValue;
+ }
+ }
+
+ ConfigValue newValue = value == null ? new ConfigValue(prop) : new ConfigValue(prop, value, new ArrayList<>(), new ArrayList<>());
+ configValidationResult.configValues().add(newValue);
+ return newValue;
+ }
+
+ public static Set<ReplicatedTopicPartition> assignedPartitions(String connector, int taskId, Map<String, String> connectorConfig) {
+ String connectorClass = connectorConfig.get(CONNECTOR_CLASS_CONFIG);
+ if (connectorClass == null) {
+ log.warn("No connector class found for connector {}", connector);
+ return null;
+ }
+
+ if (!isURep3(connectorClass)) {
+ log.warn("Ignoring task {}-{} with unsupported connector type {}", connector, taskId, connectorClass);
+ return null;
+ }
+
+ String srcKafka = connectorConfig.get(UReplicator3ConnectorConfig.SOURCE_CLUSTER_CONFIG);
+ String dstKafka = connectorConfig.get(UReplicator3ConnectorConfig.TARGET_CLUSTER_CONFIG);
+ return uRep3Assignment(connector, taskId, srcKafka, dstKafka);
+ }
+
+ public static boolean isURep3(Map<String, String> connectorConfig) {
+ String connectorClass = connectorConfig.get(CONNECTOR_CLASS_CONFIG);
+ if (connectorClass == null) {
+ log.warn("No connector class found for connector");
+ return false;
+ }
+
+ return isURep3(connectorClass);
+ }
+
+ public static boolean isURep3(String connectorClass) {
+ return Set.of(
+ "com.uber.data.kafka.connect.ureplicator3.UReplicator3Connector",
+ "UReplicator3Connector",
+ "UReplicator3"
+ ).contains(connectorClass);
+ }
+
+ public static Set<ReplicatedTopicPartition> uRep3Assignment(
+ String connector,
+ int taskId,
+ String srcKafka,
+ String dstKafka
+ ) {
+ AssignmentStore store = AssignmentStores.get();
+ return store.getAssignment(connector, taskId, srcKafka, dstKafka)
+ .stream()
+ .map(ap -> new ReplicatedTopicPartition(srcKafka, dstKafka, ap.source()))
+ .collect(Collectors.toSet());
+ }
+
+ public static void awaitAll(
+ Collection<? extends Future<?>> futures, Duration duration
+ ) throws InterruptedException, TimeoutException, ExecutionException {
+ Timer timer = Time.SYSTEM.timer(duration);
+ for (Future<?> future : futures) {
+ timer.update();
+ long remaining = Math.max(1, timer.remainingMs());
+ future.get(remaining, TimeUnit.MILLISECONDS);
+ }
+ }
+
+ public static int compareLexicographically(TopicPartition first, TopicPartition second) {
+ if (first == null)
+ return second == null ? 0 : -1;
+
+ if (second == null)
+ return 1;
+
+ int result = first.topic().compareTo(second.topic());
+ if (result != 0)
+ return result;
+
+ return Integer.compare(first.partition(), second.partition());
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/CachedWorkload.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/CachedWorkload.java
new file mode 100644
index 00000000000..0a15f453a5d
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/CachedWorkload.java
@@ -0,0 +1,28 @@
+package com.uber.data.kafka.connect.workload.common;
+
+import java.util.Objects;
+
+public record CachedWorkload(long lastUpdated, Workload workload) implements Comparable<CachedWorkload> {
+
+ public CachedWorkload(long lastUpdated, Workload workload) {
+ this.lastUpdated = lastUpdated;
+ this.workload = Objects.requireNonNull(workload, "workload may not be null");
+ }
+
+ @Override
+ public int compareTo(CachedWorkload that) {
+ if (this == that)
+ return 0;
+
+ if (that == null)
+ return 1;
+
+ int result = this.workload.compareTo(that.workload);
+ if (result != 0)
+ return result;
+
+ return Long.compare(this.lastUpdated, that.lastUpdated);
+ }
+
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/Workload.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/Workload.java
new file mode 100644
index 00000000000..a6faaac8545
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/Workload.java
@@ -0,0 +1,70 @@
+package com.uber.data.kafka.connect.workload.common;
+
+import java.util.Objects;
+
+/**
+ * POJO that tracks two generic workload statistics: bytes per second, and messages per second.
+ */
+public class Workload implements Comparable<Workload> {
+
+ public static final Workload EMPTY = new Workload(0, 0);
+
+ private final double bytesPerSecond;
+ private final double messagesPerSecond;
+
+ public Workload(double bytesPerSecond, double messagesPerSecond) {
+ this.bytesPerSecond = bytesPerSecond;
+ this.messagesPerSecond = messagesPerSecond;
+ }
+
+ public double getBytesPerSecond() {
+ return bytesPerSecond;
+ }
+
+ public double getMessagesPerSecond() {
+ return messagesPerSecond;
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (o == null || getClass() != o.getClass()) return false;
+ Workload workload = (Workload) o;
+ return Double.compare(bytesPerSecond, workload.bytesPerSecond) == 0 && Double.compare(messagesPerSecond, workload.messagesPerSecond) == 0;
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(bytesPerSecond, messagesPerSecond);
+ }
+
+ @Override
+ public String toString() {
+ return "bytesPerSecond=" + bytesPerSecond + ",messagesPerSecond=" + messagesPerSecond;
+ }
+
+ @Override
+ public int compareTo(Workload that) {
+ if (this == that)
+ return 0;
+
+ if (that == null)
+ return 1;
+
+ int cmp = Double.compare(this.bytesPerSecond, that. bytesPerSecond);
+ if (cmp != 0)
+ return cmp;
+
+ cmp = Double.compare(this.messagesPerSecond, that.messagesPerSecond);
+ return cmp;
+ }
+
+ public Workload plus(Workload that) {
+ if (that == null)
+ return this;
+
+ return new Workload(
+ this.bytesPerSecond + that.bytesPerSecond,
+ this.messagesPerSecond + that.messagesPerSecond
+ );
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/WorkloadCache.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/WorkloadCache.java
new file mode 100644
index 00000000000..f1eada2cb82
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/WorkloadCache.java
@@ -0,0 +1,221 @@
+package com.uber.data.kafka.connect.workload.common;
+
+import org.apache.kafka.common.utils.Time;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+
+
+/**
+ * Cache for topic workloads, with logic for automatic eviction of stale workloads and calculating
+ * aggregate workloads when multiple samples from different time windows are available.
+ * <p>
+ * This class is fully thread-safe: all methods may be safely invoked concurrently.
+ */
+public class WorkloadCache {
+
+ private static final Logger log = LoggerFactory.getLogger(WorkloadCache.class);
+
+ private final Time time;
+ private final long maxAgeMs;
+ private final long estimationLookbackWindowMs;
+ private final Map<ReplicatedTopicPartition, LinkedList<CachedWorkload>> cachedWorkloads;
+
+ /**
+ * Create a new cache.
+ * @param maxAgeMs the maximum age of any workloads before they are expired from the cache
+ * @param estimationLookbackWindowMs the preferred age of workloads when computing the
+ * aggregate workload for a topic or topic partition
+ */
+ public WorkloadCache(
+ long maxAgeMs,
+ long estimationLookbackWindowMs
+ ) {
+ this(maxAgeMs, estimationLookbackWindowMs, Time.SYSTEM);
+ }
+
+ /**
+ * Create a new cache.
+ * @param maxAgeMs the maximum age of any workloads before they are expired from the cache
+ * @param estimationLookbackWindowMs the preferred age of workloads when computing the
+ * aggregate workload for a topic or topic partition
+ * @param time pluggable time interface; useful for mocking; may not be null
+ */
+ @VisibleForTesting
+ WorkloadCache(
+ long maxAgeMs,
+ long estimationLookbackWindowMs,
+ Time time
+ ) {
+ this.time = time;
+ this.maxAgeMs = maxAgeMs;
+ this.estimationLookbackWindowMs = estimationLookbackWindowMs;
+ this.cachedWorkloads = new HashMap<>();
+ }
+
+ /**
+ * Add a workload for a replicated topic to the cache. Stale workloads for the topic submitted here will
+ * also be automatically expired from the cache.
+ * @param replicatedTopicPartition the replicated topic partition; may not be null
+ * @param workload the topic workload; may not be null
+ * @param lastUpdated the time when the workload was recorded
+ */
+ public void put(ReplicatedTopicPartition replicatedTopicPartition, Workload workload, long lastUpdated) {
+ CachedWorkload cachedWorkload = new CachedWorkload(lastUpdated, workload);
+ LinkedList<CachedWorkload> workloads;
+
+ if (isStale(cachedWorkload)) {
+ // Don't bother inserting the workload if it's already stale
+ synchronized (this) {
+ workloads = cachedWorkloads.get(replicatedTopicPartition);
+ }
+ } else {
+ synchronized (this) {
+ workloads = cachedWorkloads.computeIfAbsent(replicatedTopicPartition, ignored -> new LinkedList<>());
+ }
+ synchronized (workloads) {
+ if (workloads.isEmpty() || workloads.getLast().lastUpdated() < cachedWorkload.lastUpdated()) {
+ workloads.add(cachedWorkload);
+ }
+ }
+ }
+
+ if (workloads == null) {
+ return;
+ }
+
+ synchronized (workloads) {
+ // purge the data points out of the valid window
+ while (!workloads.isEmpty() && isStale(workloads.getFirst())) {
+ long workloadTimestamp = workloads.getFirst().lastUpdated();
+ log.trace("Expiring old workload for {} with timestamp of {}", replicatedTopicPartition, workloadTimestamp);
+ workloads.removeFirst();
+ }
+ }
+ }
+
+ /**
+ * Gets the workload for a replicated topic partition
+ * @param replicatedTopicPartition the replicated topic partition; may not be null
+ * @return the cached and aggregated workload for the replicated topic partition, if one can be found
+ */
+ public Optional<Workload> get(ReplicatedTopicPartition replicatedTopicPartition) {
+ Objects.requireNonNull(replicatedTopicPartition, "replicated topic partition may not be null");
+ LinkedList<CachedWorkload> allCachedWorkloads = workloadSnapshot(replicatedTopicPartition);
+ if (allCachedWorkloads == null) {
+ log.debug("There is no workload for {}.", replicatedTopicPartition);
+ return Optional.empty();
+ }
+
+ // we will return the max workload during the valid window
+ long current = time.milliseconds();
+ long lookbackWindow = lookbackWindow(current, allCachedWorkloads);
+
+ CachedWorkload maxWorkload = null;
+ for (CachedWorkload workload : allCachedWorkloads) {
+ if (current - workload.lastUpdated() > lookbackWindow) {
+ continue;
+ }
+
+ if (maxWorkload == null || maxWorkload.compareTo(workload) < 0) {
+ maxWorkload = workload;
+ }
+ }
+
+ return Optional.ofNullable(maxWorkload).map(CachedWorkload::workload);
+ }
+
+ /**
+ * Gets the workloads of all known replicated topic partitions
+ * @return a map from replicated topic partition to workload; may be empty, but never null
+ */
+ public Map<ReplicatedTopicPartition, Workload> getAll() {
+ Set<ReplicatedTopicPartition> allPartitions;
+ synchronized (this) {
+ allPartitions = new HashSet<>(cachedWorkloads.keySet());
+ }
+
+ Map<ReplicatedTopicPartition, Workload> result = new HashMap<>();
+ for (ReplicatedTopicPartition replicatedTopicPartition : allPartitions) {
+ Optional<Workload> workload = get(replicatedTopicPartition);
+ workload.ifPresent(w -> result.put(replicatedTopicPartition, w));
+ }
+ return result;
+ }
+
+ /**
+ * Gets all cached workloads for a replicated topic partition
+ * @param replicatedTopicPartition the replicated topic partition; may not be null
+ * @return all cached workloads for the replicated topic partition; may be null, and may be empty
+ */
+ public List<CachedWorkload> getHistorical(ReplicatedTopicPartition replicatedTopicPartition) {
+ return workloadSnapshot(replicatedTopicPartition);
+ }
+
+ /**
+ * Removes all cached workloads for a replicated topic partition
+ * @param replicatedTopicPartition the replicated topic partition; may not be null
+ * @return whether any workloads were removed for the specified partition
+ */
+ public boolean remove(ReplicatedTopicPartition replicatedTopicPartition) {
+ synchronized (this) {
+ return cachedWorkloads.remove(replicatedTopicPartition) != null;
+ }
+ }
+
+ /**
+ * Clear all entries from the cache
+ */
+ public void clear() {
+ synchronized (this) {
+ cachedWorkloads.clear();
+ }
+ }
+
+ private long lookbackWindow(long current, Collection<CachedWorkload> workloads) {
+ // if there are workloads inside the estimation lookback window (e.g., 1 hour), then we only consider those
+ boolean recentWorkloadsAvailable = workloads.stream()
+ .map(CachedWorkload::lastUpdated)
+ .anyMatch(lastUpdated -> current - lastUpdated < estimationLookbackWindowMs);
+ return recentWorkloadsAvailable ? estimationLookbackWindowMs : maxAgeMs;
+ }
+
+ private boolean isStale(CachedWorkload workload) {
+ return time.milliseconds() - workload.lastUpdated() > maxAgeMs;
+ }
+
+ /**
+ * Take a snapshot of the latest-known workloads for a replicated topic partition. This snapshot will
+ * belong solely to the caller and is guaranteed not to be modified by other threads or side effects
+ * of other modifications made to the cache.
+ * @param replicatedTopicPartition the replicated topic partition; may not be null
+ * @return the workload snapshot for the partition; may be empty, may be null
+ */
+ private LinkedList<CachedWorkload> workloadSnapshot(ReplicatedTopicPartition replicatedTopicPartition) {
+ List<CachedWorkload> cachedWorkloads;
+ synchronized (this) {
+ cachedWorkloads = this.cachedWorkloads.get(replicatedTopicPartition);
+ if (cachedWorkloads == null) {
+ return null;
+ }
+ }
+
+ synchronized (cachedWorkloads) {
+ return new LinkedList<>(cachedWorkloads);
+ }
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/WorkloadProvider.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/WorkloadProvider.java
new file mode 100644
index 00000000000..bab285169a0
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/WorkloadProvider.java
@@ -0,0 +1,39 @@
+package com.uber.data.kafka.connect.workload.common;
+
+import org.apache.kafka.common.Configurable;
+import org.apache.kafka.connect.components.Versioned;
+
+import java.io.Closeable;
+import java.util.Map;
+
+public interface WorkloadProvider extends Versioned, Configurable, Closeable {
+
+ /**
+ * Determine and return the load for the given task, if it is known
+ * @return the load for the task, or null if none is known
+ */
+ Double taskLoad(String connector, int taskId, Map<String, String> connectorConfig, Map<String, String> taskConfig);
+
+ static WorkloadProvider noOp() {
+ return new WorkloadProvider() {
+ @Override
+ public Double taskLoad(String connector, int taskId, Map<String, String> connectorConfig, Map<String, String> taskConfig) {
+ return null;
+ }
+
+ @Override
+ public void close() {
+ }
+
+ @Override
+ public void configure(Map<String, ?> configs) {
+ }
+
+ @Override
+ public String version() {
+ return "0.1.0";
+ }
+ };
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadProvider.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadProvider.java
new file mode 100644
index 00000000000..e555de970ba
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadProvider.java
@@ -0,0 +1,132 @@
+package com.uber.data.kafka.connect.workload.kafka;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.utils.AppInfoParser;
+import org.apache.kafka.common.utils.Utils;
+
+import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
+import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStores;
+import com.uber.data.kafka.connect.workload.common.CachedWorkload;
+import com.uber.data.kafka.connect.workload.common.Workload;
+import com.uber.data.kafka.connect.workload.common.WorkloadProvider;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import static com.uber.data.kafka.connect.utils.UberConnectUtils.assignedPartitions;
+
+public class KafkaBasedWorkloadProvider implements WorkloadProvider {
+
+ private static final AtomicBoolean DESTROY_ON_CLOSE = new AtomicBoolean(false);
+ private static final Logger log = LoggerFactory.getLogger(KafkaBasedWorkloadProvider.class);
+
+ private static KafkaBasedWorkloadStore workloadStore;
+
+ private KafkaBasedWorkloadProviderConfig config;
+
+ @Override
+ public void configure(Map<String, ?> props) {
+ log.info("Starting provider");
+ config = new KafkaBasedWorkloadProviderConfig(props);
+ initializeStore(config);
+ AssignmentStores.initialize(config.assignmentStoreConfig());
+ log.info("Finished Starting provider");
+ }
+
+ @Override
+ public Double taskLoad(String connector, int taskId, Map<String, String> connectorConfig, Map<String, String> taskConfig) {
+ Set<ReplicatedTopicPartition> replicatedTopicPartitions = assignedPartitions(connector, taskId, connectorConfig);
+ if (replicatedTopicPartitions == null) {
+ return null;
+ }
+
+ return replicatedTopicPartitions.stream()
+ .map(rtp -> workloadStore.getWorkload(rtp.srcKafka(), rtp.dstKafka(), rtp.topicPartition()))
+ .map(w -> w != null ? w : config.defaultPartitionWorkload())
+ .mapToDouble(this::scaleWorkload)
+ .sum();
+ }
+
+ @Override
+ public void close() {
+ if (!DESTROY_ON_CLOSE.get()) {
+ log.debug("Skipping teardown of workload store on close; this should only be done in testing environments");
+ return;
+ }
+ log.info("Shutting down provider");
+ destroyStore();
+ log.info("Finished shutting down provider");
+ }
+
+ @Override
+ public String version() {
+ return AppInfoParser.getVersion();
+ }
+
+ public static Future<?> publishWorkload(String srcKafka, String dstKafka, TopicPartition topicPartition, Workload workload) {
+ assertStoreInitialized();
+
+ return workloadStore.putWorkload(srcKafka, dstKafka, topicPartition, workload);
+ }
+
+ public static void refreshWorkloadCache() {
+ assertStoreInitialized();
+
+ AssignmentStores.get().refresh();
+ workloadStore.refresh();
+ }
+
+ public static Map<ReplicatedTopicPartition, Workload> getAllWorkloads() {
+ assertStoreInitialized();
+
+ return workloadStore.getAllWorkloads();
+ }
+
+ public static List<CachedWorkload> getHistoricalWorkloads(String srcKafka, String dstKafka, TopicPartition topicPartition) {
+ assertStoreInitialized();
+
+ return workloadStore.getHistoricalWorkloads(srcKafka, dstKafka, topicPartition);
+ }
+
+ public static void clearCache() {
+ assertStoreInitialized();
+
+ workloadStore.clear();
+ }
+
+ public static void assertStoreInitialized() throws IllegalStateException {
+ if (workloadStore == null)
+ throw new IllegalStateException("Workload store has not been initialized yet, or has already been shut down");
+ }
+
+ public static void destroyOnClose(boolean destroy) {
+ DESTROY_ON_CLOSE.set(destroy);
+ }
+
+ private static synchronized void initializeStore(KafkaBasedWorkloadProviderConfig config) {
+ if (workloadStore != null) {
+ log.warn("Provider appears to have been configured multiple times on the same JVM. This should only happen during testing");
+ return;
+ }
+ workloadStore = new KafkaBasedWorkloadStore(config);
+ workloadStore.start();
+ }
+
+ private static synchronized void destroyStore() {
+ Utils.closeQuietly(workloadStore, "workload store");
+ workloadStore = null;
+ }
+
+ private double scaleWorkload(Workload taskWorkload) {
+ double byteRateLoad = taskWorkload.getBytesPerSecond() / config.maxWorkerWorkload().getBytesPerSecond();
+ double messageRateLoad = taskWorkload.getMessagesPerSecond() / config.maxWorkerWorkload().getMessagesPerSecond();
+ return Math.max(byteRateLoad, messageRateLoad);
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadProviderConfig.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadProviderConfig.java
new file mode 100644
index 00000000000..acbd0d181fb
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadProviderConfig.java
@@ -0,0 +1,225 @@
+package com.uber.data.kafka.connect.workload.kafka;
+
+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.TopicConfig;
+
+import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStoreConfig;
+import com.uber.data.kafka.connect.workload.common.Workload;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
+
+public class KafkaBasedWorkloadProviderConfig extends AbstractConfig {
+
+ private static final Logger log = LoggerFactory.getLogger(KafkaBasedWorkloadProviderConfig.class);
+
+ private static final String CONFIG_PREFIX = "uber.workload.provider.";
+
+ // TODO: Flipr
+ public static final String MAX_WORKER_BYTE_RATE_CONFIG = CONFIG_PREFIX + "max.worker.byte.rate";
+ public static final double MAX_WORKER_BYTE_RATE_DEFAULT = 70 * 1024 * 1024;
+ private static final String MAX_WORKER_BYTE_RATE_DOC = "Maximum total byte rate for workloads assigned "
+ + "to each worker before attempting to revoke workloads and redistribute to other workers";
+
+ // TODO: Flipr
+ public static final String MAX_WORKER_MESSAGE_RATE_CONFIG = CONFIG_PREFIX + "max.worker.message.rate";
+ public static final double MAX_WORKER_MESSAGE_RATE_DEFAULT = 70_000;
+ private static final String MAX_WORKER_MESSAGE_RATE_DOC = "Maximum total message rate for workloads assigned "
+ + "to each worker before attempting to revoke workloads and redistribute to other workers";
+
+ // TODO: Flipr
+ public static final String DEFAULT_PARTITION_BYTE_RATE_CONFIG = CONFIG_PREFIX + "default.partition.byte.rate";
+ public static final double DEFAULT_PARTITION_BYTE_RATE_DEFAULT = 1024;
+ private static final String DEFAULT_PARTITION_BYTE_RATE_DOC = "The default byte rate to assume for partitions "
+ + "that have no workload information available (likely because they have only recently started being replicated)";
+
+ // TODO: Flipr
+ public static final String DEFAULT_PARTITION_MESSAGE_RATE_CONFIG = CONFIG_PREFIX + "default.partition.message.rate";
+ public static final double DEFAULT_PARTITION_MESSAGE_RATE_DEFAULT = 10;
+ private static final String DEFAULT_PARTITION_MESSAGE_RATE_DOC = "The default message rate to assume for partitions "
+ + "that have no workload information available (likely because they have only recently started being replicated)";
+
+ public static final String STORAGE_TOPIC_CONFIG = CONFIG_PREFIX + "storage.topic";
+ public static final String STORAGE_TOPIC_PREFIX = STORAGE_TOPIC_CONFIG + ".";
+ private static final String STORAGE_TOPIC_DOC = "Topic used to store workload information. "
+ + "This topic must be compacted (i.e., have the 'compacted' cleanup policy) and can have "
+ + "any number of partitions, but it's recommended to use a high number (e.g., 25 or 50). "
+ + "If the topic does not exist, it will be automatically created on startup. "
+ + "Additional properties for the topic can be specified using the " + STORAGE_TOPIC_PREFIX + " "
+ + "prefix";
+
+ private static final String PARTITIONS_SUFFIX = "partitions";
+ public static final String STORAGE_TOPIC_PARTITIONS_CONFIG = STORAGE_TOPIC_PREFIX + PARTITIONS_SUFFIX;
+ private static final int STORAGE_TOPIC_PARTITIONS_DEFAULT = 50;
+ private static final String STORAGE_TOPIC_PARTITIONS_DOC = "The number of partitions to create "
+ + "the workload topic with, if it does not already exist. If the topic already exists this property "
+ + "will be ignored, even if the topic has a different number of partitions";
+
+ private static final String REPLICATION_FACTOR_SUFFIX = "replication.factor";
+ public static final String STORAGE_TOPIC_REPLICATION_FACTOR_CONFIG = STORAGE_TOPIC_PREFIX + REPLICATION_FACTOR_SUFFIX;
+ private static final short STORAGE_TOPIC_REPLICATION_FACTOR_DEFAULT = (short) 3;
+ private static final String STORAGE_TOPIC_REPLICATION_FACTOR_DOC = "The replication factor to create "
+ + "the workload topic with, if it does not already exist. If the topic already exists this property "
+ + "will be ignored, even if the topic has a different replication factor";
+
+ private static final String SEGMENT_BYTES_SUFFIX = TopicConfig.SEGMENT_BYTES_CONFIG;
+ public static final String STORAGE_TOPIC_SEGMENT_BYTES_CONFIG = STORAGE_TOPIC_PREFIX + SEGMENT_BYTES_SUFFIX;
+ private static final int STORAGE_TOPIC_SEGMENT_BYTES_DEFAULT = 1024 * 1024 * 3; // Lower default segment size in order to enable more-frequent compaction
+ private static final String STORAGE_TOPIC_SEGMENT_BYTES_DOC = "The per-partition segment size "
+ + "for the workload topic. This should be relatively low in order to enable topic compaction to take place, which "
+ + "will help keep the topic from growing too large and causing slow startup time for workers. Like other topic "
+ + "properties, if the topic already exists, this property will be ignored, even if the topic has a different segment size";
+
+ public static final String CLIENTS_PREFIX = CONFIG_PREFIX + "kafka.clients";
+ public static final String CONSUMER_PREFIX = CONFIG_PREFIX + "consumer";
+ public static final String PRODUCER_PREFIX = CONFIG_PREFIX + "producer";
+ public static final String ADMIN_PREFIX = CONFIG_PREFIX + "admin";
+
+ private final Workload maxWorkerWorkload;
+ private final Workload defaultPartitionWorkload;
+
+ private final AssignmentStoreConfig assignmentStoreConfig;
+
+ @SuppressWarnings("this-escape")
+ public KafkaBasedWorkloadProviderConfig(Map<String, ?> props) {
+ super(config(), props);
+ this.maxWorkerWorkload = new Workload(
+ getDouble(MAX_WORKER_BYTE_RATE_CONFIG),
+ getDouble(MAX_WORKER_MESSAGE_RATE_CONFIG)
+ );
+ this.defaultPartitionWorkload = new Workload(
+ getDouble(DEFAULT_PARTITION_BYTE_RATE_CONFIG),
+ getDouble(DEFAULT_PARTITION_MESSAGE_RATE_CONFIG)
+ );
+ this.assignmentStoreConfig = new AssignmentStoreConfig(props);
+ }
+
+ private static ConfigDef config() {
+ ConfigDef result = new ConfigDef()
+ .define(
+ STORAGE_TOPIC_CONFIG,
+ Type.STRING,
+ Importance.HIGH,
+ STORAGE_TOPIC_DOC
+ ).define(
+ MAX_WORKER_BYTE_RATE_CONFIG,
+ Type.DOUBLE,
+ MAX_WORKER_BYTE_RATE_DEFAULT,
+ atLeast(1),
+ Importance.MEDIUM,
+ MAX_WORKER_BYTE_RATE_DOC
+ ).define(
+ MAX_WORKER_MESSAGE_RATE_CONFIG,
+ Type.DOUBLE,
+ MAX_WORKER_MESSAGE_RATE_DEFAULT,
+ atLeast(1),
+ Importance.MEDIUM,
+ MAX_WORKER_MESSAGE_RATE_DOC
+ ).define(
+ DEFAULT_PARTITION_BYTE_RATE_CONFIG,
+ Type.DOUBLE,
+ DEFAULT_PARTITION_BYTE_RATE_DEFAULT,
+ atLeast(0),
+ Importance.MEDIUM,
+ DEFAULT_PARTITION_BYTE_RATE_DOC
+ ).define(
+ DEFAULT_PARTITION_MESSAGE_RATE_CONFIG,
+ Type.DOUBLE,
+ DEFAULT_PARTITION_MESSAGE_RATE_DEFAULT,
+ atLeast(0),
+ Importance.MEDIUM,
+ DEFAULT_PARTITION_MESSAGE_RATE_DOC
+ ).define(
+ STORAGE_TOPIC_PARTITIONS_CONFIG,
+ Type.INT,
+ STORAGE_TOPIC_PARTITIONS_DEFAULT,
+ atLeast(1),
+ Importance.LOW,
+ STORAGE_TOPIC_PARTITIONS_DOC
+ ).define(
+ STORAGE_TOPIC_REPLICATION_FACTOR_CONFIG,
+ Type.SHORT,
+ STORAGE_TOPIC_REPLICATION_FACTOR_DEFAULT,
+ atLeast(1),
+ Importance.LOW,
+ STORAGE_TOPIC_REPLICATION_FACTOR_DOC
+ ).define(
+ STORAGE_TOPIC_SEGMENT_BYTES_CONFIG,
+ Type.INT,
+ STORAGE_TOPIC_SEGMENT_BYTES_DEFAULT,
+ atLeast(1),
+ Importance.LOW,
+ STORAGE_TOPIC_SEGMENT_BYTES_DOC
+ );
+ AssignmentStoreConfig.addToConfig(result);
+ return result;
+ }
+
+ public AssignmentStoreConfig assignmentStoreConfig() {
+ return assignmentStoreConfig;
+ }
+
+ public Map<String, Object> adminConfig() {
+ return clientConfig(ADMIN_PREFIX);
+ }
+
+ public Map<String, Object> producerConfig() {
+ return clientConfig(PRODUCER_PREFIX);
+ }
+
+ public Map<String, Object> consumerConfig() {
+ return clientConfig(CONSUMER_PREFIX);
+ }
+
+ public String storageTopic() {
+ return getString(STORAGE_TOPIC_CONFIG);
+ }
+
+ public Workload maxWorkerWorkload() {
+ return maxWorkerWorkload;
+ }
+
+ public Workload defaultPartitionWorkload() {
+ return defaultPartitionWorkload;
+ }
+
+ public int storageTopicPartitions() {
+ return getInt(STORAGE_TOPIC_PARTITIONS_CONFIG);
+ }
+
+ public short storageTopicReplicationFactor() {
+ return getShort(STORAGE_TOPIC_REPLICATION_FACTOR_CONFIG);
+ }
+
+ public Map<String, Object> storageTopicSettings() {
+ Map<String, Object> result = originalsWithPrefix(STORAGE_TOPIC_PREFIX);
+
+ // AbstractConfig::originalsWithPrefix doesn't include default values, we apply that manually here
+ result.computeIfAbsent(TopicConfig.SEGMENT_BYTES_CONFIG, k -> STORAGE_TOPIC_SEGMENT_BYTES_DEFAULT);
+
+ Object removedPolicy = result.remove(TopicConfig.CLEANUP_POLICY_CONFIG);
+ if (removedPolicy != null) {
+ log.warn("Ignoring '{}cleanup.policy={}' setting, since compaction is always used", STORAGE_TOPIC_PREFIX, removedPolicy);
+ }
+
+ result.remove(REPLICATION_FACTOR_SUFFIX);
+ result.remove(PARTITIONS_SUFFIX);
+
+ return result;
+ }
+
+ private Map<String, Object> clientConfig(String specialPrefix) {
+ Map<String, Object> result = new HashMap<>(originals());
+ result.putAll(originalsWithPrefix(CLIENTS_PREFIX));
+ result.putAll(originalsWithPrefix(specialPrefix));
+ return result;
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadStore.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadStore.java
new file mode 100644
index 00000000000..92d8246b95f
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadStore.java
@@ -0,0 +1,286 @@
+package com.uber.data.kafka.connect.workload.kafka;
+
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.internals.KafkaCompletableFuture;
+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.common.utils.Utils;
+import org.apache.kafka.connect.util.KafkaBasedLog;
+import org.apache.kafka.connect.util.TopicAdmin;
+
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.annotations.VisibleForTesting;
+import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
+import com.uber.data.kafka.connect.utils.UberConnectUtils;
+import com.uber.data.kafka.connect.workload.common.CachedWorkload;
+import com.uber.data.kafka.connect.workload.common.Workload;
+import com.uber.data.kafka.connect.workload.common.WorkloadCache;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+
+import static com.uber.data.kafka.connect.utils.SerializationUtils.parseMapField;
+import static com.uber.data.kafka.connect.utils.SerializationUtils.parseNumberField;
+import static com.uber.data.kafka.connect.utils.SerializationUtils.parseStringField;
+
+class KafkaBasedWorkloadStore implements Closeable {
+
+ private static final Logger log = LoggerFactory.getLogger(KafkaBasedWorkloadStore.class);
+ private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
+
+ @VisibleForTesting
+ static final long MAX_VALID_TIME_WINDOW_IN_MS = TimeUnit.HOURS.toMillis(25);
+ @VisibleForTesting
+ static final long ESTIMATION_LOOKBACK_WINDOW_IN_MS = TimeUnit.HOURS.toMillis(1);
+
+ private static final String WORKLOAD_FIELD = "workload-v1";
+ private static final String SRC_KAFKA_FIELD = "src";
+ private static final String DST_KAFKA_FIELD = "dst";
+ private static final String TOPIC_FIELD = "topic";
+ private static final String PARTITION_FIELD = "partition";
+ private static final String BYTE_RATE_FIELD = "bytes";
+ private static final String MESSAGE_RATE_FIELD = "messages";
+
+ private final String topic;
+ private final int topicPartitions;
+ private final short topicReplicationFactor;
+ private final Map<String, Object> topicSettings;
+ private final Map<String, Object> producerProps;
+ private final Map<String, Object> consumerProps;
+ private final Map<String, Object> adminProps;
+ private final WorkloadCache workloadsCache;
+ private final Time time;
+
+ private TopicAdmin topicAdmin;
+ private KafkaBasedLog<byte[], byte[]> kafkaLog;
+
+ public KafkaBasedWorkloadStore(KafkaBasedWorkloadProviderConfig config) {
+ this(config, Time.SYSTEM);
+ }
+
+ @VisibleForTesting
+ KafkaBasedWorkloadStore(KafkaBasedWorkloadProviderConfig config, Time time) {
+ this.topic = config.storageTopic();
+ this.topicPartitions = config.storageTopicPartitions();
+ this.topicReplicationFactor = config.storageTopicReplicationFactor();
+ this.topicSettings = config.storageTopicSettings();
+ this.producerProps = producerProps(config.producerConfig());
+ this.consumerProps = consumerProps(config.consumerConfig());
+ this.adminProps = config.adminConfig();
+ this.workloadsCache = new WorkloadCache(MAX_VALID_TIME_WINDOW_IN_MS, ESTIMATION_LOOKBACK_WINDOW_IN_MS);
+ this.time = time;
+ }
+
+ public void start() {
+ this.topicAdmin = new TopicAdmin(adminProps);
+ this.kafkaLog = new KafkaBasedLog<>(
+ topic,
+ producerProps,
+ consumerProps,
+ () -> topicAdmin,
+ this::onConsume,
+ Time.SYSTEM,
+ this::createAndValidateTopic
+ );
+ kafkaLog.start();
+ }
+
+ private static Map<String, Object> producerProps(Map<String, Object> baseProps) {
+ baseProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
+ baseProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
+ baseProps.put(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG, Integer.MAX_VALUE);
+ return baseProps;
+ }
+
+ private static Map<String, Object> consumerProps(Map<String, Object> baseProps) {
+ baseProps.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
+ baseProps.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
+ return baseProps;
+ }
+
+ private void createAndValidateTopic(TopicAdmin topicAdmin) {
+ NewTopic newTopic = TopicAdmin.defineTopic(topic)
+ .config(topicSettings)
+ .compacted()
+ .partitions(topicPartitions)
+ .replicationFactor(topicReplicationFactor)
+ .build();
+ boolean created = topicAdmin.createTopic(newTopic);
+ if (!created) {
+ // Ensure existing topic is compacted
+ boolean verified = topicAdmin.verifyTopicCleanupPolicyOnlyCompact(
+ topic,
+ KafkaBasedWorkloadProviderConfig.STORAGE_TOPIC_CONFIG,
+ "workloads"
+ );
+ // Should never, NEVER happen, but just in case
+ if (!verified) {
+ throw new KafkaException("Failed to verify that workload storage topic " + topic + " is compacted");
+ }
+ }
+ }
+
+ @Override
+ public void close() {
+ if (kafkaLog != null) {
+ Utils.closeQuietly(kafkaLog::stop, "Kafka-based log");
+ }
+ Utils.closeQuietly(topicAdmin, "topic admin");
+ }
+
+ public Future<?> putWorkload(String srcKafka, String dstKafka, TopicPartition topicPartition, Workload workload) {
+ ReplicatedTopicPartition replicatedTopicPartition = new ReplicatedTopicPartition(srcKafka, dstKafka, topicPartition);
+
+ byte[] key;
+ try {
+ key = serializeWorkloadKey(replicatedTopicPartition);
+ } catch (JsonProcessingException e) {
+ log.error("Failed to serialize workload key for replicated topic partition {} as JSON", replicatedTopicPartition, e);
+ return KafkaCompletableFuture.failedFuture(e);
+ }
+
+ byte[] value;
+ try {
+ value = serializeWorkloadValue(workload);
+ } catch (JsonProcessingException e) {
+ log.error("Failed to serialize workload value {} for replicated topic partition {} as JSON", replicatedTopicPartition, e);
+ return KafkaCompletableFuture.failedFuture(e);
+ }
+
+ // TODO: Error handling (probably just log and emit a metric)
+ return kafkaLog.sendWithReceipt(key, value);
+ }
+
+ public Workload getWorkload(String srcKafka, String dstKafka, TopicPartition topicPartition) {
+ ReplicatedTopicPartition replicatedTopicPartition = new ReplicatedTopicPartition(srcKafka, dstKafka, topicPartition);
+ return workloadsCache.get(replicatedTopicPartition).orElse(null);
+ }
+
+ public Map<ReplicatedTopicPartition, Workload> getAllWorkloads() {
+ return workloadsCache.getAll();
+ }
+
+ public List<CachedWorkload> getHistoricalWorkloads(String srcKafka, String dstKafka, TopicPartition topicPartition) {
+ ReplicatedTopicPartition replicatedTopicPartition = new ReplicatedTopicPartition(srcKafka, dstKafka, topicPartition);
+ return workloadsCache.getHistorical(replicatedTopicPartition);
+ }
+
+ public void refresh() {
+ try {
+ kafkaLog.readToEnd().get();
+ } catch (InterruptedException | ExecutionException e) {
+ throw new KafkaException("Failed to read to end of workload log", e);
+ }
+ }
+
+ public void clear() {
+ workloadsCache.clear();
+ }
+
+ private void onConsume(Throwable error, ConsumerRecord<byte[], byte[]> record) {
+ if (error != null) {
+ log.error("Error occurred while consuming from workload topic", error);
+ return;
+ }
+
+ int recordPartition = record.partition();
+ long recordOffset = record.offset();
+
+ ReplicatedTopicPartition replicatedTopicPartition;
+ try {
+ replicatedTopicPartition = deserializeWorkloadKey(record.key());
+ } catch (IOException e) {
+ log.error("Failed to deserialize replicated topic partition in partition {} at offset {} as JSON", recordPartition, recordOffset, e);
+ return;
+ }
+
+ Workload workload;
+ try {
+ workload = deserializeWorkloadValue(record.value());
+ } catch (IOException e) {
+ log.error("Failed to deserialize workload in partition {} at offset {} for replicated topic partition {} as JSON", recordPartition, recordOffset, replicatedTopicPartition, e);
+ return;
+ }
+
+ if (workload != null) {
+ long timestamp = UberConnectUtils.extractTimestamp(record)
+ .orElseGet(() -> {
+ log.warn("No timestamp found for record in partition {} at offset {}; using current wall-clock time instead", recordPartition, recordOffset);
+ return time.milliseconds();
+ });
+
+ workloadsCache.put(replicatedTopicPartition, workload, timestamp);
+ } else {
+ workloadsCache.remove(replicatedTopicPartition);
+ }
+ }
+
+ private byte[] serializeWorkloadKey(ReplicatedTopicPartition replicatedTopicPartition) throws JsonProcessingException {
+ // LinkedHashMap preserves order and allows us to deterministically generate a serialized key from a replicated topic partition
+ Map<String, Object> innerMap = new LinkedHashMap<>();
+ innerMap.put(SRC_KAFKA_FIELD, replicatedTopicPartition.srcKafka());
+ innerMap.put(DST_KAFKA_FIELD, replicatedTopicPartition.dstKafka());
+ innerMap.put(TOPIC_FIELD, replicatedTopicPartition.topicPartition().topic());
+ innerMap.put(PARTITION_FIELD, replicatedTopicPartition.topicPartition().partition());
+ Map<String, Object> wrapped = Map.of(WORKLOAD_FIELD, innerMap);
+ return OBJECT_MAPPER.writeValueAsBytes(wrapped);
+ }
+
+ private ReplicatedTopicPartition deserializeWorkloadKey(byte[] key) throws IOException {
+ if (key == null)
+ throw new IOException("record has no key");
+
+ Map<String, Object> map = OBJECT_MAPPER.readValue(key, new TypeReference<>() { });
+ Map<?, ?> wrapped = parseMapField(true, map, WORKLOAD_FIELD);
+
+ String srcKafka = parseStringField(true, wrapped, SRC_KAFKA_FIELD);
+ String dstKafka = parseStringField(true, wrapped, DST_KAFKA_FIELD);
+ String topic = parseStringField(true, wrapped, TOPIC_FIELD);
+ Number partition = parseNumberField(true, wrapped, PARTITION_FIELD);
+
+ return new ReplicatedTopicPartition(srcKafka, dstKafka, topic, partition.intValue());
+ }
+
+ private byte[] serializeWorkloadValue(Workload workload) throws JsonProcessingException {
+ // This allows us to write tombstone records for replicated topic partitions
+ if (workload == null)
+ return null;
+
+ Map<String, Object> map = Map.of(
+ BYTE_RATE_FIELD, workload.getBytesPerSecond(),
+ MESSAGE_RATE_FIELD, workload.getMessagesPerSecond()
+ );
+ return OBJECT_MAPPER.writeValueAsBytes(map);
+ }
+
+ private Workload deserializeWorkloadValue(byte[] value) throws IOException {
+ // Permit null values (tombstones) that indicate we should clear the cache for this replicated topic partition
+ if (value == null)
+ return null;
+
+ Map<String, Object> map = OBJECT_MAPPER.readValue(value, new TypeReference<>() { });
+
+ Number byteRate = parseNumberField(false, map, BYTE_RATE_FIELD);
+ Number messageRate = parseNumberField(false, map, MESSAGE_RATE_FIELD);
+
+ return new Workload(byteRate.doubleValue(), messageRate.doubleValue());
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadTrackingTransformation.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadTrackingTransformation.java
new file mode 100644
index 00000000000..670a068252a
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadTrackingTransformation.java
@@ -0,0 +1,15 @@
+package com.uber.data.kafka.connect.workload.kafka;
+
+import org.apache.kafka.connect.connector.ConnectRecord;
+
+import com.uber.data.kafka.connect.workload.tracking.WorkloadStore;
+import com.uber.data.kafka.connect.workload.transformation.WorkloadTrackingTransformation;
+
+public class KafkaBasedWorkloadTrackingTransformation<R extends ConnectRecord<R>> extends WorkloadTrackingTransformation<R> {
+
+ @Override
+ protected WorkloadStore createStore() {
+ return KafkaBasedWorkloadProvider::publishWorkload;
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadCollector.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadCollector.java
new file mode 100644
index 00000000000..577ae78d7a4
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadCollector.java
@@ -0,0 +1,103 @@
+package com.uber.data.kafka.connect.workload.tracking;
+
+import org.apache.kafka.common.utils.Time;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.uber.data.kafka.connect.workload.common.Workload;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.time.Duration;
+import java.util.Objects;
+import java.util.Optional;
+
+/**
+ * Tracks segments for a single workload. Segments will be considered in-progress until a
+ * user-specified duration has elapsed, at which point the segment will be considered complete, and
+ * a new active segment will be started.
+ */
+class WorkloadCollector {
+
+ private static final Logger log = LoggerFactory.getLogger(WorkloadCollector.class);
+
+ private final Time time;
+ private final Duration maxSegmentAge;
+
+ private long startMs;
+ private long totalBytes;
+ private long totalMessages;
+
+ WorkloadCollector(Duration maxSegmentAge, Time time) {
+ this.maxSegmentAge = Objects.requireNonNull(maxSegmentAge, "max segment age may not be null");
+ this.time = Objects.requireNonNull(time, "time may not be null");
+ reset();
+ }
+
+ /**
+ * Record the workload for zero or more messages, and if there is an in-progress segment whose
+ * age has exceeded the maximum segment age, return that segment's workload and begin a new
+ * segment.
+ * <p>
+ * Note that if a {@link Optional#isPresent() present Optional} is returned, the
+ * workload information contained in it will be dropped by this collector after the method
+ * returns. It is the responsibility of the caller to store the workload information for as long as it is needed.
+ * @param bytes the total bytes consumed for the messages; may not be negative
+ * @param messages the total messages consumed; may not be negative
+ * @return the workload for the previous segment, if one was just completed
+ */
+ public synchronized Optional<Workload> record(long bytes, long messages) {
+ if (bytes < 0) {
+ log.warn("Ignoring invalid workload with negative number of bytes: {}", bytes);
+ return Optional.empty();
+ }
+ if (messages < 0) {
+ log.warn("Ignoring invalid workload with negative number of messages: {}", messages);
+ return Optional.empty();
+ }
+
+ Optional<Workload> result;
+ Duration age = age();
+
+ if (age.compareTo(maxSegmentAge) >= 0) {
+ result = Optional.ofNullable(collect(totalBytes, totalMessages, age));
+ reset();
+ } else {
+ result = Optional.empty();
+ }
+
+ this.totalBytes += bytes;
+ this.totalMessages += messages;
+
+ return result;
+ }
+
+ private Duration age() {
+ return Duration.ofMillis(time.milliseconds() - startMs);
+ }
+
+ @VisibleForTesting
+ static Workload collect(long totalBytes, long totalMessages, Duration age) {
+ double durationSeconds = Math.ceil(age.toMillis() / 1000.0);
+ if (durationSeconds <= 0) {
+ log.warn(
+ "Tried to collect workload from segment whose age in seconds ({}) is " +
+ "non-positive; discarding this segment and starting fresh",
+ durationSeconds
+ );
+ return null;
+ }
+
+ return new Workload(
+ totalBytes / durationSeconds,
+ totalMessages / durationSeconds
+ );
+ }
+
+ private void reset() {
+ this.startMs = time.milliseconds();
+ this.totalBytes = 0;
+ this.totalMessages = 0;
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadStore.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadStore.java
new file mode 100644
index 00000000000..816b02867df
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadStore.java
@@ -0,0 +1,12 @@
+package com.uber.data.kafka.connect.workload.tracking;
+
+import org.apache.kafka.common.TopicPartition;
+
+import com.uber.data.kafka.connect.workload.common.Workload;
+
+@FunctionalInterface
+public interface WorkloadStore {
+
+ void store(String srcKafka, String dstKafka, TopicPartition topicPartition, Workload workload);
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadTracker.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadTracker.java
new file mode 100644
index 00000000000..00edf3581f6
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadTracker.java
@@ -0,0 +1,132 @@
+package com.uber.data.kafka.connect.workload.tracking;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.utils.Time;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Ticker;
+import com.google.common.cache.CacheBuilder;
+import com.google.common.cache.CacheLoader;
+import com.google.common.cache.LoadingCache;
+import com.google.common.cache.RemovalCause;
+import com.google.common.cache.RemovalListener;
+import com.google.common.cache.RemovalNotification;
+import com.uber.data.kafka.connect.workload.common.Workload;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.time.Duration;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * Tracks {@link Workload workloads} on a per-topic-partition basis, allowing users to periodically
+ * {@link #record(TopicPartition, long, long) record} new workload information for partitions and,
+ * once workload information for a topic partition has matured, retrieve a complete snapshot of the
+ * average per-second workload of that partition.
+ */
+class WorkloadTracker {
+
+ private static final Logger log = LoggerFactory.getLogger(WorkloadTracker.class);
+
+ private final Duration maxSegmentAge;
+ private final Time time;
+ private final LoadingCache<TopicPartition, WorkloadCollector> activeSegments;
+ private final Map<TopicPartition, Workload> staleWorkloads;
+
+ public WorkloadTracker(Duration maxSegmentAge) {
+ this(
+ maxSegmentAge,
+ Time.SYSTEM,
+ Ticker.systemTicker()
+ );
+ }
+
+ @VisibleForTesting
+ WorkloadTracker(
+ Duration maxSegmentAge,
+ Time time,
+ Ticker ticker
+ ) {
+ this.maxSegmentAge = maxSegmentAge;
+ this.time = time;
+
+ this.activeSegments = CacheBuilder.newBuilder()
+ .ticker(ticker)
+ .expireAfterAccess(maxSegmentAge.toMillis(), TimeUnit.MILLISECONDS)
+ .removalListener(new ExpiredCollectorListener())
+ .build(new WorkloadCollectorCacheLoader());
+ this.staleWorkloads = new HashMap<>();
+ }
+
+ /**
+ * Record the throughput for a single record, and if the maximum segment age has been reached
+ * for the given topic partition, return the workload for its just-completed segment
+ * @param topicPartition the topic partition that the record came from; may not be null
+ * @param bytes the size of the record in bytes; may not be negative
+ * @param messages the number of records; may not be negative
+ * @return the workload for the current segment, if it has reached the maximum segment age
+ */
+ public Optional<Workload> record(TopicPartition topicPartition, long bytes, long messages) {
+ return activeSegments.getUnchecked(topicPartition)
+ .record(bytes, messages);
+ }
+
+ /**
+ * Returns and then discards all stale workloads; that is, workloads for segments that have
+ * reached or exceeded the maximum segment age, and for which no throughputs have recently been
+ * recorded. Should be invoked periodically in order to prevent memory leaks.
+ * @return all known stale workloads; may be empty, but never null
+ */
+ public Map<TopicPartition, Workload> collectStaleWorkloads() {
+ activeSegments.cleanUp();
+
+ Map<TopicPartition, Workload> result;
+ synchronized (staleWorkloads) {
+ result = new HashMap<>(staleWorkloads);
+ staleWorkloads.clear();
+ }
+
+ return result;
+ }
+
+ /**
+ * Clear all workloads, both normal and stale.
+ */
+ public void clear() {
+ activeSegments.asMap().clear();
+ staleWorkloads.clear();
+ }
+
+ private class WorkloadCollectorCacheLoader extends CacheLoader<TopicPartition, WorkloadCollector> {
+
+ @Override
+ public WorkloadCollector load(TopicPartition topicPartition) throws Exception {
+ log.debug("Started recording workload for topic partition {}", topicPartition);
+ return new WorkloadCollector(maxSegmentAge, time);
+ }
+ }
+
+ private class ExpiredCollectorListener implements RemovalListener<TopicPartition, WorkloadCollector> {
+ @Override
+ public void onRemoval(RemovalNotification<TopicPartition, WorkloadCollector> removal) {
+ if (!RemovalCause.EXPIRED.equals(removal.getCause())) {
+ return;
+ }
+
+ // We just simulate recording a throughput of 0 bytes/0 messages; this forces a check
+ // to see if the segment is complete, but doesn't have any effect on the throughputs
+ // for the topic partition
+ Optional<Workload> finalSegmentWorkload = removal.getValue().record(0, 0);
+ finalSegmentWorkload.ifPresent(workload -> {
+ synchronized (staleWorkloads) {
+ staleWorkloads.put(removal.getKey(), workload);
+ }
+ });
+ }
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadTrackingProcessor.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadTrackingProcessor.java
new file mode 100644
index 00000000000..83c01cf06a2
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadTrackingProcessor.java
@@ -0,0 +1,276 @@
+package com.uber.data.kafka.connect.workload.tracking;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.connect.connector.ConnectRecord;
+import org.apache.kafka.connect.header.Header;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.uber.data.kafka.connect.utils.ExecutorUtils;
+import com.uber.data.kafka.connect.utils.MetricsUtils;
+import com.uber.data.kafka.connect.workload.common.Workload;
+import com.uber.data.kafka.connect.workload.transformation.WorkloadTrackingTransformation;
+import com.uber.m3.tally.Scope;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.time.Duration;
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Supplier;
+
+public class WorkloadTrackingProcessor implements AutoCloseable {
+
+ @VisibleForTesting
+ static final String WORKLOAD_BYTE_RATE_METRIC = "worker.workloads.byte.rate";
+ @VisibleForTesting
+ static final String WORKLOAD_MESSAGE_RATE_METRIC = "worker.workloads.message.rate";
+ @VisibleForTesting
+ static final String WORKLOAD_EMPTY_PARTITIONS_METRIC = "worker.workloads.empty.partitions";
+ @VisibleForTesting
+ static final String WORKLOAD_STALE_FLUSH_FAILURE_METRIC = "worker.workloads.stale.flush.failures";
+ @VisibleForTesting
+ static final String WORKLOAD_EMPTY_PUBLISH_FAILURE_METRIC = "worker.workloads.empty.publish.failures";
+
+ private static final Logger log = LoggerFactory.getLogger(WorkloadTrackingTransformation.class);
+
+ private final WorkloadStore workloadStore;
+ private final String srcKafka;
+ private final String dstKafka;
+ private final Supplier<? extends Collection<TopicPartition>> assignedPartitions;
+ private final Set<TopicPartition> emptyPartitions;
+ private final WorkloadTracker topicWorkloadTracker;
+ private final ScheduledExecutorService executor;
+ private final Scope scope;
+ private final Duration maxSegmentAge;
+
+ public WorkloadTrackingProcessor(
+ String srcKafka,
+ String dstKafka,
+ WorkloadStore workloadStore,
+ Supplier<? extends Collection<TopicPartition>> assignedPartitions,
+ Duration maxSegmentAge
+ ) {
+ this(
+ srcKafka,
+ dstKafka,
+ assignedPartitions,
+ MetricsUtils.INSTANCE.routeScope(srcKafka, dstKafka),
+ workloadStore,
+ new WorkloadTracker(maxSegmentAge),
+ ExecutorUtils.singleThreadedScheduled("workload-tracking-message-processor-flush"),
+ maxSegmentAge
+ );
+ }
+
+ @VisibleForTesting
+ WorkloadTrackingProcessor(
+ String srcKafka,
+ String dstKafka,
+ Supplier<? extends Collection<TopicPartition>> assignedPartitions,
+ Scope scope,
+ WorkloadStore workloadStore,
+ WorkloadTracker topicWorkloadTracker,
+ ScheduledExecutorService executor,
+ Duration maxSegmentAge
+ ) {
+ this.srcKafka = Objects.requireNonNull(
+ srcKafka,
+ "source Kafka cluster name may not be null"
+ );
+ this.dstKafka = Objects.requireNonNull(
+ dstKafka,
+ "target Kafka cluster name may not be null"
+ );
+ this.assignedPartitions = Objects.requireNonNull(
+ assignedPartitions,
+ "assigned topics supplier may not be null"
+ );
+ this.scope = Objects.requireNonNull(
+ scope,
+ "metrics scope may not be null"
+ );
+ this.workloadStore = Objects.requireNonNull(
+ workloadStore,
+ "workload store may not be nul"
+ );
+ this.topicWorkloadTracker = Objects.requireNonNull(
+ topicWorkloadTracker,
+ "topic workload tracker may not be null"
+ );
+ this.executor = Objects.requireNonNull(
+ executor,
+ "flush executor may not be null"
+ );
+ this.maxSegmentAge = Objects.requireNonNull(
+ maxSegmentAge,
+ "max segment age may not be null"
+ );
+
+ this.emptyPartitions = new HashSet<>();
+ }
+
+ public void start() {
+ executor.scheduleAtFixedRate(
+ this::handleStaleAndEmptyWorkloads,
+ maxSegmentAge.toMillis(),
+ maxSegmentAge.toMillis(),
+ TimeUnit.MILLISECONDS
+ );
+ }
+
+ public void track(ConnectRecord<?> record) {
+ track(record.topic(), record.kafkaPartition(), numRecordBytes(record), 1);
+ }
+
+ @VisibleForTesting
+ void track(String topic, int partition, long bytes, long messages) {
+ TopicPartition topicPartition = new TopicPartition(topic, partition);
+ synchronized (emptyPartitions) {
+ emptyPartitions.remove(topicPartition);
+ }
+ topicWorkloadTracker.record(topicPartition, bytes, messages)
+ .ifPresent(workload -> publishWorkload(topicPartition, workload));
+ }
+
+ @Override
+ public void close() {
+ ExecutorUtils.shutdown(
+ executor,
+ "executor for workload tracker",
+ 5,
+ TimeUnit.SECONDS
+ );
+
+ // Final attempt to flush complete workloads
+ flushStaleWorkloads();
+ }
+
+ @VisibleForTesting
+ void handleStaleAndEmptyWorkloads() {
+ try {
+ flushStaleWorkloads();
+ } catch (Throwable t) {
+ log.error("Failed to flush stale workloads", t);
+ scope.counter(WORKLOAD_STALE_FLUSH_FAILURE_METRIC).inc(1);
+ }
+
+ try {
+ publishEmptyWorkloads();
+ } catch (Throwable t) {
+ log.error("Failed to publish empty workloads", t);
+ scope.counter(WORKLOAD_EMPTY_PUBLISH_FAILURE_METRIC).inc(1);
+ }
+ }
+
+ /**
+ * Check for stale segments (i.e., ones that are old enough to complete, but for which no
+ * throughput has been recorded after they became old enough), complete them, and flush them.
+ */
+ @VisibleForTesting
+ void flushStaleWorkloads() {
+ Map<TopicPartition, Workload> staleWorkloads = topicWorkloadTracker.collectStaleWorkloads();
+ if (staleWorkloads.isEmpty()) {
+ return;
+ }
+
+ synchronized (emptyPartitions) {
+ emptyPartitions.removeAll(staleWorkloads.keySet());
+ }
+
+ log.debug(
+ "Flushing {} stale workloads for topic partitions {}",
+ staleWorkloads.size(),
+ staleWorkloads.keySet()
+ );
+ staleWorkloads.forEach(this::publishWorkload);
+ }
+
+ /**
+ * Check for empty topic partitions (i.e., ones that haven't received any traffic since the
+ * last time we checked), and publish a zero-throughput workload for them.
+ * <p>
+ * This explicit step is necessary because all other workloads are only published after some
+ * traffic is received (via one or more calls to {@link #track(String, int, long, long)});
+ * for empty partitions, we never see traffic, and instead have to explicitly check for them
+ * and publish workloads.
+ */
+ @VisibleForTesting
+ void publishEmptyWorkloads() {
+ Set<TopicPartition> partitions = new HashSet<>(assignedPartitions.get());
+
+ // Take a snapshot of the current set of empty partitions so that we don't perform writes
+ // to ZooKeeper while holding the lock
+ Set<TopicPartition> emptyPartitionsSnapshot;
+ synchronized (this.emptyPartitions) {
+ emptyPartitionsSnapshot = new HashSet<>(this.emptyPartitions);
+
+ // Assume every partition is empty until we see traffic for it
+ this.emptyPartitions.clear();
+ this.emptyPartitions.addAll(partitions);
+ }
+
+ // Don't publish empty workloads for partitions that are no longer assigned to us
+ emptyPartitionsSnapshot.retainAll(partitions);
+
+ scope.gauge(WORKLOAD_EMPTY_PARTITIONS_METRIC)
+ .update(emptyPartitionsSnapshot.size());
+
+ emptyPartitionsSnapshot.forEach(tp -> publishWorkload(tp, Workload.EMPTY));
+ }
+
+ private void publishWorkload(TopicPartition topicPartition, Workload workload) {
+ log.trace(
+ "Publishing workload of {} for topic partition {}",
+ workload,
+ topicPartition
+ );
+
+ workloadStore.store(srcKafka, dstKafka, topicPartition, workload);
+
+ scope.tagged(MetricsUtils.INSTANCE.tags(topicPartition))
+ .gauge(WORKLOAD_BYTE_RATE_METRIC)
+ .update(workload.getBytesPerSecond());
+ scope.tagged(MetricsUtils.INSTANCE.tags(topicPartition))
+ .gauge(WORKLOAD_MESSAGE_RATE_METRIC)
+ .update(workload.getMessagesPerSecond());
+ }
+
+ @VisibleForTesting
+ static int numRecordBytes(ConnectRecord<?> record) {
+ String topic = record.topic();
+ int recordBytes = 0;
+
+ if (record.key() instanceof byte[] keyBytes) {
+ recordBytes += keyBytes.length;
+ } else {
+ log.warn("Key for record in topic {} is not byte array; ignoring", topic);
+ }
+
+ if (record.value() instanceof byte[] valueBytes) {
+ recordBytes += valueBytes.length;
+ } else {
+ log.warn("Value for record in topic {} is not byte array; ignoring", topic);
+ }
+
+ for (Header header : record.headers()) {
+ recordBytes += header.key().length();
+
+ if (header.value() instanceof String stringHeader) {
+ recordBytes += stringHeader.length();
+ } else if (header.value() instanceof byte[] bytesHeader) {
+ recordBytes += bytesHeader.length;
+ } else {
+ log.debug("Value for header {} on record in topic {} is not byte array; ignoring", header.key(), topic);
+ }
+ }
+
+ return recordBytes;
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/transformation/WorkloadTrackingTransformation.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/transformation/WorkloadTrackingTransformation.java
new file mode 100644
index 00000000000..1cb2bc69b0d
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/transformation/WorkloadTrackingTransformation.java
@@ -0,0 +1,62 @@
+package com.uber.data.kafka.connect.workload.transformation;
+
+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.transforms.Transformation;
+
+import com.uber.data.kafka.connect.workload.tracking.WorkloadStore;
+import com.uber.data.kafka.connect.workload.tracking.WorkloadTrackingProcessor;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * {@link org.apache.kafka.connect.transforms.Transformation} implementation that tracks workloads on a per-topic-partition basis as
+ * records are consumed from Kafka, and reports those workloads periodically to a pluggable {@link WorkloadStore}
+ */
+public abstract class WorkloadTrackingTransformation<R extends ConnectRecord<R>> implements Transformation<R> {
+
+ private static final Logger log = LoggerFactory.getLogger(WorkloadTrackingTransformation.class);
+
+ private WorkloadTrackingProcessor processor;
+
+ protected abstract WorkloadStore createStore();
+
+ @Override
+ public ConfigDef config() {
+ return WorkloadTrackingTransformationConfig.config();
+ }
+
+ @Override
+ public void configure(Map<String, ?> props) {
+ WorkloadTrackingTransformationConfig config = new WorkloadTrackingTransformationConfig(props);
+
+ this.processor = new WorkloadTrackingProcessor(
+ config.srcKafkaName(),
+ config.dstKafkaName(),
+ createStore(),
+ Set::of, // TODO: Is there any non-insane way to get the set of topic partitions assigned to the task here?
+ config.reportInterval()
+ );
+
+ this.processor.start();
+
+ log.info("Topic workload tracking is now enabled");
+ }
+
+ @Override
+ public R apply(R record) {
+ processor.track(record);
+ return record;
+ }
+
+ @Override
+ public void close() {
+ Utils.closeQuietly(processor, "workload processor");
+ }
+
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/transformation/WorkloadTrackingTransformationConfig.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/transformation/WorkloadTrackingTransformationConfig.java
new file mode 100644
index 00000000000..4c4e83ea745
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/transformation/WorkloadTrackingTransformationConfig.java
@@ -0,0 +1,62 @@
+package com.uber.data.kafka.connect.workload.transformation;
+
+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.time.Duration;
+import java.util.Map;
+
+import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
+
+public class WorkloadTrackingTransformationConfig extends AbstractConfig {
+
+ public static final String SRC_CLUSTER_NAME_CONFIG = "source.cluster.name";
+ private static final String SRC_CLUSTER_NAME_DOC = "Name of the source Kafka cluster";
+
+ public static final String DST_CLUSTER_NAME_CONFIG = "target.cluster.name";
+ private static final String DST_CLUSTER_NAME_DOC = "Name of the target Kafka cluster";
+
+ public static final String REPORT_INTERVAL_MS_CONFIG = "report.interval.ms";
+ private static final String REPORT_INTERVAL_MS_DOC = "How long to record traffic for in between reporting recorded values";
+
+ public WorkloadTrackingTransformationConfig(Map<String, ?> props) {
+ super(config(), props);
+ }
+
+ static ConfigDef config() {
+ return new ConfigDef()
+ .define(
+ SRC_CLUSTER_NAME_CONFIG,
+ Type.STRING,
+ Importance.HIGH,
+ SRC_CLUSTER_NAME_DOC
+ ).define(
+ DST_CLUSTER_NAME_CONFIG,
+ Type.STRING,
+ Importance.HIGH,
+ DST_CLUSTER_NAME_DOC
+ ).define(
+ REPORT_INTERVAL_MS_CONFIG,
+ Type.LONG,
+ Duration.ofMinutes(10).toMillis(),
+ atLeast(Duration.ofSeconds(1).toMillis()),
+ Importance.LOW,
+ REPORT_INTERVAL_MS_DOC
+ );
+ }
+
+ public String srcKafkaName() {
+ return getString(SRC_CLUSTER_NAME_CONFIG);
+ }
+
+ public String dstKafkaName() {
+ return getString(DST_CLUSTER_NAME_CONFIG);
+ }
+
+ public Duration reportInterval() {
+ return Duration.ofMillis(getLong(REPORT_INTERVAL_MS_CONFIG));
+ }
+
+}
diff --git a/connect/uber/src/main/resources/META-INF/services/com.uber.data.kafka.connect.distributed.ClusterAssignor b/connect/uber/src/main/resources/META-INF/services/com.uber.data.kafka.connect.distributed.ClusterAssignor
new file mode 100644
index 00000000000..3d72df7440a
--- /dev/null
+++ b/connect/uber/src/main/resources/META-INF/services/com.uber.data.kafka.connect.distributed.ClusterAssignor
@@ -0,0 +1 @@
+ com.uber.data.kafka.connect.distributed.assignment.UberClusterAssignor
\ No newline at end of file
diff --git a/connect/uber/src/main/resources/META-INF/services/org.apache.kafka.connect.rest.ConnectRestExtension b/connect/uber/src/main/resources/META-INF/services/org.apache.kafka.connect.rest.ConnectRestExtension
new file mode 100644
index 00000000000..b12936f3e82
--- /dev/null
+++ b/connect/uber/src/main/resources/META-INF/services/org.apache.kafka.connect.rest.ConnectRestExtension
@@ -0,0 +1 @@
+com.uber.data.kafka.connect.rest.UberConnectRestExtension
\ No newline at end of file
diff --git a/connect/uber/src/main/resources/META-INF/services/org.apache.kafka.connect.source.SourceConnector b/connect/uber/src/main/resources/META-INF/services/org.apache.kafka.connect.source.SourceConnector
new file mode 100644
index 00000000000..5822942e2b8
--- /dev/null
+++ b/connect/uber/src/main/resources/META-INF/services/org.apache.kafka.connect.source.SourceConnector
@@ -0,0 +1 @@
+com.uber.data.kafka.connect.ureplicator3.UReplicator3Connector
\ No newline at end of file
diff --git a/connect/uber/src/test/java/com/uber/data/kafka/connect/integration/ReplicationIntegrationTest.java b/connect/uber/src/test/java/com/uber/data/kafka/connect/integration/ReplicationIntegrationTest.java
new file mode 100644
index 00000000000..790f082010b
--- /dev/null
+++ b/connect/uber/src/test/java/com/uber/data/kafka/connect/integration/ReplicationIntegrationTest.java
@@ -0,0 +1,1243 @@
+package com.uber.data.kafka.connect.integration;
+
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.Admin;
+import org.apache.kafka.clients.admin.ListTopicsOptions;
+import org.apache.kafka.clients.admin.NewPartitions;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.producer.Producer;
+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.utils.Utils;
+import org.apache.kafka.connect.converters.ByteArrayConverter;
+import org.apache.kafka.connect.runtime.isolation.PluginDiscoveryMode;
+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.ConnectorOffset;
+import org.apache.kafka.connect.runtime.rest.entities.ConnectorOffsets;
+import org.apache.kafka.connect.runtime.rest.entities.CreateConnectorRequest;
+import org.apache.kafka.connect.runtime.rest.errors.ConnectRestException;
+import org.apache.kafka.connect.util.clusters.EmbeddedConnectCluster;
+import org.apache.kafka.connect.util.clusters.EmbeddedKafkaCluster;
+import org.apache.kafka.connect.util.clusters.WorkerHandle;
+
+import com.fasterxml.jackson.annotation.JsonAutoDetect.Visibility;
+import com.fasterxml.jackson.annotation.PropertyAccessor;
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
+import com.uber.data.kafka.connect.distributed.assignment.UberClusterAssignor;
+import com.uber.data.kafka.connect.rest.UberConnectRestExtension;
+import com.uber.data.kafka.connect.rest.entities.AssignedPartitionEntity;
+import com.uber.data.kafka.connect.rest.entities.PartitionWorkloadEntity;
+import com.uber.data.kafka.connect.rest.entities.TaskAssignmentEntity;
+import com.uber.data.kafka.connect.rest.entities.WorkerWorkloadEntity;
+import com.uber.data.kafka.connect.rest.entities.WorkloadEntity;
+import com.uber.data.kafka.connect.ureplicator3.UReplicator3Connector;
+import com.uber.data.kafka.connect.ureplicator3.UReplicator3Task;
+import com.uber.data.kafka.connect.ureplicator3.coordination.AssignedPartition;
+import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStoreConfig;
+import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStores;
+import com.uber.data.kafka.connect.workload.common.Workload;
+import com.uber.data.kafka.connect.workload.kafka.KafkaBasedWorkloadProvider;
+
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+import java.net.URI;
+import java.net.http.HttpClient;
+import java.net.http.HttpRequest;
+import java.net.http.HttpRequest.BodyPublishers;
+import java.net.http.HttpResponse;
+import java.net.http.HttpResponse.BodyHandlers;
+import java.nio.charset.StandardCharsets;
+import java.util.AbstractMap;
+import java.util.Arrays;
+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.ConcurrentHashMap;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
+
+import jakarta.ws.rs.core.Response;
+import jakarta.ws.rs.core.Response.Status;
+
+import static com.uber.data.kafka.connect.distributed.assignment.UberClusterAssignorConfig.UBER_WORKLOAD_PROVIDER_CONFIG;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Config.TARGET_CLIENTS_PREFIX;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Config.TARGET_PRODUCER_PREFIX;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.REFRESH_INTERVAL_MS;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.SOURCE_ADMIN_PREFIX;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.SOURCE_CLUSTER_CONFIG;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.SOURCE_CONSUMER_PREFIX;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.TARGET_ADMIN_PREFIX;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.TARGET_CLUSTER_CONFIG;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.TOPICS_CONFIG;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.TOPIC_RENAME_PREFIX;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Task.SOURCE_OFFSET_KAFKA_OFFSET;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3TaskConfig.WORKLOAD_REPORT_INTERVAL_MS_CONFIG;
+import static com.uber.data.kafka.connect.workload.kafka.KafkaBasedWorkloadProviderConfig.STORAGE_TOPIC_CONFIG;
+import static com.uber.data.kafka.connect.workload.kafka.KafkaBasedWorkloadProviderConfig.STORAGE_TOPIC_REPLICATION_FACTOR_CONFIG;
+import static jakarta.ws.rs.core.HttpHeaders.CONTENT_TYPE;
+import static jakarta.ws.rs.core.MediaType.APPLICATION_FORM_URLENCODED;
+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.NAME_CONFIG;
+import static org.apache.kafka.connect.runtime.ConnectorConfig.TASKS_MAX_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.CONNECTOR_CLIENT_POLICY_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.KEY_CONVERTER_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.PLUGIN_DISCOVERY_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.UBER_CLUSTER_ASSIGNOR_CONFIG;
+import static org.apache.kafka.connect.runtime.rest.RestServerConfig.REST_EXTENSION_CLASSES_CONFIG;
+import static org.apache.kafka.test.TestUtils.waitForCondition;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+@Tag("integration")
+public class ReplicationIntegrationTest {
... 1295 lines suppressed ...
(kafka) 02/02: [DRAFT] Add pluggable worker assignor, remove Uber module
Posted by ce...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
C0urante pushed a commit to branch metamorphomover-prototype
in repository https://gitbox.apache.org/repos/asf/kafka.git
commit 08d656e584d507971532f337f679375b7fa57545
Author: chrise <ch...@uber.com>
AuthorDate: Wed Apr 29 05:38:03 2026 +0000
[DRAFT] Add pluggable worker assignor, remove Uber module
---
build-prototype.sh | 70 ----------------------
build.gradle | 23 ++-----
.../distributed/ClusterConfigStateImpl.java | 16 +++++
.../runtime/distributed/DistributedConfig.java | 3 +-
.../runtime/distributed/DistributedHerder.java | 3 +-
.../connect/runtime/distributed/UberAssignor.java | 5 +-
.../runtime/distributed/WorkerCoordinator.java | 3 +-
.../runtime/distributed/WorkerGroupMember.java | 3 +-
.../runtime/health/ConnectClusterStateImpl.java | 1 -
.../runtime/isolation/PluginScanResult.java | 3 +-
.../connect/runtime/isolation/PluginType.java | 2 +-
publish-connect.sh | 51 ++++++++++++++++
12 files changed, 87 insertions(+), 96 deletions(-)
diff --git a/build-prototype.sh b/build-prototype.sh
deleted file mode 100755
index 1ab302b0ea8..00000000000
--- a/build-prototype.sh
+++ /dev/null
@@ -1,70 +0,0 @@
-#! /usr/bin/env bash
-
-if [[ "$#" -ne 1 ]]; then
- echo "Usage: $0 <tb_path>"
- exit 1
-fi
-
-TB_PATH="$1"
-
-upload() {
- echo "Uploading $1 to TerraBlob"
- if ! [ -z ${FORCE+x} ]; then
- tb-cli delete "$TB_PATH"/"$1" || true
- fi
- tb-cli put --timeout 90s --multipart "$1" "$TB_PATH"/"$1"
-}
-
-set -e
-
-SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd)
-
-if tb-cli ls "$TB_PATH" > /dev/null 2> /dev/null; then
- if [ -z ${FORCE+x} ] && [[ $(tb-cli ls "$TB_PATH" | wc -l) != 0 ]]; then
- echo "Path $TB_PATH already exists and is non-empty; cannot use"
- exit 1
- elif ! [ -z ${FORCE+x} ]; then
- echo "Path $TB_PATH is non-empty but FORCE is set; proceeding"
- else
- echo "Path $TB_PATH already exists and is empty; can use safely"
- fi
-else
- tb-cli mkdir -p "$TB_PATH"
-fi
-
-
-if [ -z ${NOBUILD+x} ]; then
-
- if [ -z ${NOCLEAN+x} ]; then
- "$SCRIPT_DIR"/gradlew clean
- else
- echo "Skipping clean"
- fi
-
- "$SCRIPT_DIR"/gradlew :connect:uber:{spotlessApply,build} releaseTarGz -x test
-
-else
- echo "Skipping build"
-fi
-
-if ! [ -z ${NOPUSH+x} ]; then
- echo "Skipping upload"
- exit 0
-fi
-
-
-pushd "$SCRIPT_DIR"/core/build/distributions > /dev/null
-upload kafka_2.13-4.3.0-SNAPSHOT.tgz
-popd > /dev/null
-
-pushd "$SCRIPT_DIR"/connect/uber/build > /dev/null
-TEMP_DIR="$(mktemp -d)"
-mkdir "$TEMP_DIR"/connect-uber-plugins
-cp libs/* dependant-libs/* "$TEMP_DIR"/connect-uber-plugins
-pushd "$TEMP_DIR" > /dev/null
-tar czf connect-uber-plugins.tgz connect-uber-plugins
-upload connect-uber-plugins.tgz
-popd > /dev/null
-popd > /dev/null
-
-echo "Finished uploading artifacts to $TB_PATH"
diff --git a/build.gradle b/build.gradle
index 404e1a0ded6..abc18ef9e92 100644
--- a/build.gradle
+++ b/build.gradle
@@ -430,6 +430,7 @@ subprojects {
if (shouldSign) {
signing {
+ useGpgCmd()
sign publishing.publications.mavenJava
}
}
@@ -3644,10 +3645,7 @@ project(':connect:runtime') {
implementation libs.classgraph
implementation libs.mavenArtifact
implementation libs.swaggerAnnotations
- // TODO: We should not be hacking the :connect:runtime module with our own dependencies; instead,
- // we should be adding them to the classpath
- // This is also causing embedded integration tests to fail ☹️
- implementation ('com.uber.kafka:kafka-upki-provider:1.0.1-217732-a4c6b554904e-java-monorepo')
+ implementation 'com.uber.kafka:kafka-upki-provider-shaded:1.20.0-chrise-5'
compileOnly libs.bndlib
compileOnly libs.spotbugs
@@ -3677,8 +3675,7 @@ project(':connect:runtime') {
testImplementation libs.junitJupiter
testImplementation libs.mockitoCore
testImplementation libs.mockitoJunitJupiter
- // TODO: Conflicts with Kafka UPKI provider
-// testImplementation libs.httpclient
+ testImplementation libs.httpclient
testImplementation testLog4j2Libs
testCompileOnly libs.bndlib
@@ -3990,14 +3987,6 @@ project(':connect:uber') {
archivesName = "connect-uber"
}
- task genUReplicator3Config(type: JavaExec) {
- classpath sourceSets.main.runtimeClasspath
- classpath sourceSets.main.compileClasspath
- mainClass = 'com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig'
- if( !generatedDocsDir.exists() ) { generatedDocsDir.mkdirs() }
- standardOutput = new File(generatedDocsDir, "ureplicator3_config.html").newOutputStream()
- }
-
dependencies {
implementation 'com.google.guava:guava:23.0'
implementation 'com.uber.m3:tally-core:0.11.1'
@@ -4153,9 +4142,9 @@ def updateVersionTask = tasks.register('updateVersion') {
def rawVersion
if (project.hasProperty('newVersion')) {
rawVersion = newVersion
- if (!rawVersion.matches(/^[0-9]+\.[0-9]+\.[0-9]+(-SNAPSHOT)?$/)) {
- throw new GradleException("Invalid version format: '${rawVersion}'. Expected format: X.Y.Z or X.Y.Z-SNAPSHOT")
- }
+// if (!rawVersion.matches(/^[0-9]+\.[0-9]+\.[0-9]+(-SNAPSHOT)?$/)) {
+// throw new GradleException("Invalid version format: '${rawVersion}'. Expected format: X.Y.Z or X.Y.Z-SNAPSHOT")
+// }
// Update gradle.properties
def gradlePropsFile = file("${project.rootDir}/gradle.properties")
if (gradlePropsFile.exists()) {
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ClusterConfigStateImpl.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ClusterConfigStateImpl.java
index def0f344d56..9321a113843 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ClusterConfigStateImpl.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ClusterConfigStateImpl.java
@@ -1,3 +1,19 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY 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 com.uber.data.kafka.connect.distributed.ClusterConfigState;
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 391597f664c..1026581a167 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
@@ -16,7 +16,6 @@
*/
package org.apache.kafka.connect.runtime.distributed;
-import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.clients.MetadataRecoveryStrategy;
import org.apache.kafka.clients.producer.ProducerConfig;
@@ -27,6 +26,8 @@ import org.apache.kafka.common.security.auth.SecurityProtocol;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.connect.runtime.WorkerConfig;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
+
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
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 3b85f561e16..381db371aac 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,7 +16,6 @@
*/
package org.apache.kafka.connect.runtime.distributed;
-import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.common.KafkaFuture;
import org.apache.kafka.common.config.ConfigDef;
@@ -80,6 +79,8 @@ import org.apache.kafka.connect.util.SinkUtils;
import org.apache.kafka.connect.util.Stage;
import org.apache.kafka.connect.util.TemporaryStage;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
+
import org.slf4j.Logger;
import java.util.ArrayList;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberAssignor.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberAssignor.java
index 4a466575ad3..1166199afec 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberAssignor.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberAssignor.java
@@ -16,11 +16,12 @@
*/
package org.apache.kafka.connect.runtime.distributed;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.connect.storage.ClusterConfigState;
+
import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import com.uber.data.kafka.connect.distributed.ConnectorTaskId;
import com.uber.data.kafka.connect.distributed.ConnectorsAndTasks;
-import org.apache.kafka.common.utils.LogContext;
-import org.apache.kafka.connect.storage.ClusterConfigState;
import org.slf4j.Logger;
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 192f035c6b7..1f98f92c628 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
@@ -16,7 +16,6 @@
*/
package org.apache.kafka.connect.runtime.distributed;
-import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import org.apache.kafka.clients.GroupRebalanceConfig;
import org.apache.kafka.clients.consumer.CloseOptions;
import org.apache.kafka.clients.consumer.internals.AbstractCoordinator;
@@ -31,6 +30,8 @@ import org.apache.kafka.connect.storage.ClusterConfigState;
import org.apache.kafka.connect.storage.ConfigBackingStore;
import org.apache.kafka.connect.util.ConnectorTaskId;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
+
import org.slf4j.Logger;
import java.io.Closeable;
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 c3dc2acdd3f..f3d0e664110 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
@@ -16,7 +16,6 @@
*/
package org.apache.kafka.connect.runtime.distributed;
-import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import org.apache.kafka.clients.ApiVersions;
import org.apache.kafka.clients.ClientUtils;
import org.apache.kafka.clients.CommonClientConfigs;
@@ -43,6 +42,8 @@ import org.apache.kafka.connect.runtime.WorkerConfig;
import org.apache.kafka.connect.storage.ConfigBackingStore;
import org.apache.kafka.connect.util.ConnectorTaskId;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
+
import org.slf4j.Logger;
import java.net.InetSocketAddress;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImpl.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImpl.java
index 3f2ce96766b..4c82e8f3a57 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImpl.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImpl.java
@@ -29,7 +29,6 @@ import org.apache.kafka.connect.health.TaskState;
import org.apache.kafka.connect.runtime.Herder;
import org.apache.kafka.connect.runtime.distributed.RebalanceNeededException;
import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo;
-import org.apache.kafka.connect.runtime.rest.entities.TaskInfo;
import org.apache.kafka.connect.util.FutureCallback;
import org.slf4j.Logger;
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 3c6a193a4bf..5ade6a6ef2e 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,7 +16,6 @@
*/
package org.apache.kafka.connect.runtime.isolation;
-import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import org.apache.kafka.common.config.provider.ConfigProvider;
import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
import org.apache.kafka.connect.rest.ConnectRestExtension;
@@ -27,6 +26,8 @@ import org.apache.kafka.connect.storage.HeaderConverter;
import org.apache.kafka.connect.transforms.Transformation;
import org.apache.kafka.connect.transforms.predicates.Predicate;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
+
import java.util.List;
import java.util.SortedSet;
import java.util.TreeSet;
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 05a4766df71..0b686ac68d4 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,6 @@
*/
package org.apache.kafka.connect.runtime.isolation;
-import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import org.apache.kafka.common.config.provider.ConfigProvider;
import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
import org.apache.kafka.connect.rest.ConnectRestExtension;
@@ -28,6 +27,7 @@ import org.apache.kafka.connect.transforms.Transformation;
import org.apache.kafka.connect.transforms.predicates.Predicate;
import com.fasterxml.jackson.annotation.JsonValue;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import java.util.Locale;
diff --git a/publish-connect.sh b/publish-connect.sh
new file mode 100755
index 00000000000..0bb8d7211de
--- /dev/null
+++ b/publish-connect.sh
@@ -0,0 +1,51 @@
+#! /usr/bin/env bash
+
+if [[ "$#" -ne 1 ]]; then
+ echo "Usage: $0 <version>"
+ exit 1
+fi
+
+set -e
+
+VERSION=4.3.0-chrise-$1-uber
+
+./gradlew clean
+./gradlew updateVersion -PnewVersion=$VERSION
+./gradlew build -x test -x rat -x spotbugsMain -x spotbugsTest -x checkstyleMain -x checkstyleTest
+./gradlew releaseTarGz
+./gradlew publish -x test -x rat \
+ -PmavenUrl=https://artifacts.uberinternal.com/artifactory/libs-release-local \
+ -PmavenPassword="$(usso -ussh artifacts.uberinternal.com -print)" \
+ -PmaxParallelForks=8 --max-workers=8 \
+ -x :streams:publish \
+ -x :streams:examples:publish \
+ -x :streams:test-utils:publish \
+ -x :streams:streams-scala:publish \
+ -x :streams:integration-tests:publish \
+ -x :streams:upgrade-system-tests-0110:publish \
+ -x :streams:upgrade-system-tests-10:publish \
+ -x :streams:upgrade-system-tests-11:publish \
+ -x :streams:upgrade-system-tests-20:publish \
+ -x :streams:upgrade-system-tests-21:publish \
+ -x :streams:upgrade-system-tests-22:publish \
+ -x :streams:upgrade-system-tests-23:publish \
+ -x :streams:upgrade-system-tests-24:publish \
+ -x :streams:upgrade-system-tests-25:publish \
+ -x :streams:upgrade-system-tests-26:publish \
+ -x :streams:upgrade-system-tests-27:publish \
+ -x :streams:upgrade-system-tests-28:publish \
+ -x :streams:upgrade-system-tests-30:publish \
+ -x :streams:upgrade-system-tests-31:publish \
+ -x :streams:upgrade-system-tests-32:publish \
+ -x :streams:upgrade-system-tests-33:publish \
+ -x :streams:upgrade-system-tests-34:publish \
+ -x :streams:upgrade-system-tests-35:publish \
+ -x :streams:upgrade-system-tests-36:publish \
+ -x :streams:upgrade-system-tests-37:publish \
+ -x :streams:upgrade-system-tests-38:publish \
+ -x :streams:upgrade-system-tests-39:publish \
+ -x :streams:upgrade-system-tests-40:publish \
+ -x :streams:upgrade-system-tests-41:publish
+curl -H "Authorization: Bearer $(usso -ussh artifacts -print)" \
+ https://artifacts.uberinternal.com/artifactory/libs-release-local/org/apache/kafka/kafka_2.13/$VERSION/kafka_2.13-$VERSION.tgz \
+ -T core/build/distributions/kafka_2.13-$VERSION.tgz