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/03/12 22:02:54 UTC
(kafka) 01/01: [DRAFT] Prototype workload-based load balancing, Kafka-based workload tracking, and UReplicator3 connector
This is an automated email from the ASF dual-hosted git repository.
cegerton pushed a commit to branch chrise-ureplicator3-prototype
in repository https://gitbox.apache.org/repos/asf/kafka.git
commit 132f9d3c07872a23bb565c03982afb493dceba05
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.gradle | 66 ++
.../connect/uber/distributed/WorkloadProvider.java | 18 +
.../kafka/connect/mirror/MirrorSourceConfig.java | 2 +-
.../runtime/distributed/ClusterAssignment.java | 109 +++
.../runtime/distributed/DistributedConfig.java | 53 +-
.../runtime/distributed/DistributedHerder.java | 11 +-
.../IncrementalCooperativeAssignor.java | 98 ---
.../connect/runtime/distributed/UberAssignor.java | 532 ++++++++++++++
.../runtime/distributed/WorkerCoordinator.java | 38 +-
.../runtime/distributed/WorkerGroupMember.java | 10 +-
.../runtime/isolation/PluginScanResult.java | 15 +-
.../connect/runtime/isolation/PluginType.java | 4 +-
.../runtime/isolation/ReflectionScanner.java | 3 +-
.../runtime/isolation/ServiceLoaderScanner.java | 3 +-
.../IncrementalCooperativeAssignorTest.java | 1 -
.../WorkerCoordinatorIncrementalTest.java | 6 +-
.../runtime/distributed/WorkerCoordinatorTest.java | 5 +-
.../runtime/distributed/WorkerGroupMemberTest.java | 4 +-
.../isolation/DelegatingClassLoaderTest.java | 1 +
.../connect/runtime/isolation/PluginUtilsTest.java | 4 +
.../connect/runtime/isolation/PluginsTest.java | 1 +
.../connect/ureplicator3/UReplicator3Config.java | 261 +++++++
.../ureplicator3/UReplicator3Connector.java | 186 +++++
.../ureplicator3/UReplicator3ConnectorConfig.java | 282 ++++++++
.../connect/ureplicator3/UReplicator3Task.java | 211 ++++++
.../ureplicator3/UReplicator3TaskConfig.java | 115 +++
.../data/kafka/connect/utils/ExecutorUtils.java | 88 +++
.../kafka/connect/utils/M3ReporterProvider.java | 43 ++
.../data/kafka/connect/utils/MetricsUtils.java | 37 +
.../data/kafka/connect/utils/UberConnectUtils.java | 134 ++++
.../workload/common/ReplicatedTopicPartition.java | 67 ++
.../kafka/connect/workload/common/Workload.java | 60 ++
.../connect/workload/common/WorkloadCache.java | 231 ++++++
.../workload/kafka/KafkaBasedWorkloadProvider.java | 170 +++++
.../kafka/KafkaBasedWorkloadProviderConfig.java | 126 ++++
.../workload/kafka/KafkaBasedWorkloadStore.java | 316 ++++++++
.../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 | 61 ++
.../WorkloadTrackingTransformationConfig.java | 62 ++
...org.apache.kafka.connect.source.SourceConnector | 16 +
...kafka.connect.uber.distributed.WorkloadProvider | 16 +
.../integration/ReplicationIntegrationTest.java | 800 +++++++++++++++++++++
settings.gradle | 1 +
47 files changed, 4686 insertions(+), 119 deletions(-)
diff --git a/build.gradle b/build.gradle
index bb11d6eb3bc..f77716aaeae 100644
--- a/build.gradle
+++ b/build.gradle
@@ -168,6 +168,14 @@ allprojects {
repositories {
mavenCentral()
+ maven {
+ url 'http://artifactory.uber.internal:4587/artifactory/libs-release-local/'
+ allowInsecureProtocol = true
+ metadataSources {
+ mavenPom()
+ artifact()
+ }
+ }
}
dependencyUpdates {
@@ -3969,6 +3977,64 @@ project(':connect:mirror-client') {
}
}
+project(':connect:uber') {
+ base {
+ archivesName = "connect-uber"
+ }
+
+ dependencies {
+ implementation project(':connect:api')
+ implementation project(':connect:runtime')
+ implementation project(':connect:mirror')
+ implementation project(':clients')
+ 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 libs.slf4jApi
+
+ testImplementation testLog4j2Libs
+ testImplementation libs.junitJupiter
+ testImplementation libs.bndlib
+ testImplementation libs.mockitoCore
+ testImplementation project(':clients').sourceSets.test.output
+ testImplementation project(':connect:runtime').sourceSets.test.output
+ 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/connect/api/src/main/java/org/apache/kafka/connect/uber/distributed/WorkloadProvider.java b/connect/api/src/main/java/org/apache/kafka/connect/uber/distributed/WorkloadProvider.java
new file mode 100644
index 00000000000..2e6f2cc45a5
--- /dev/null
+++ b/connect/api/src/main/java/org/apache/kafka/connect/uber/distributed/WorkloadProvider.java
@@ -0,0 +1,18 @@
+package org.apache.kafka.connect.uber.distributed;
+
+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 per-second byte rate of the given topic partition,
+ * if it is known
+ * @return the byte rate for the topic partition, or null if none is known
+ */
+ Double byteRate(String connector, int taskId, Map<String, String> connectorConfig, Map<String, String> taskConfig);
+
+}
diff --git a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorSourceConfig.java b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorSourceConfig.java
index 2ec663ad2fc..1c18e27ef65 100644
--- a/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorSourceConfig.java
+++ b/connect/mirror/src/main/java/org/apache/kafka/connect/mirror/MirrorSourceConfig.java
@@ -52,7 +52,7 @@ public class MirrorSourceConfig extends MirrorConnectorConfig {
public static final String OFFSET_SYNCS_TOPIC_REPLICATION_FACTOR_DOC = "Replication factor for offset-syncs topic.";
public static final short OFFSET_SYNCS_TOPIC_REPLICATION_FACTOR_DEFAULT = 3;
- static final String TASK_TOPIC_PARTITIONS = "task.assigned.partitions";
+ public static final String TASK_TOPIC_PARTITIONS = "task.assigned.partitions";
public static final String CONSUMER_POLL_TIMEOUT_MILLIS = "consumer.poll.timeout.ms";
private static final String CONSUMER_POLL_TIMEOUT_MILLIS_DOC = "Timeout when polling source cluster.";
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ClusterAssignment.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ClusterAssignment.java
new file mode 100644
index 00000000000..a34ff67d925
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ClusterAssignment.java
@@ -0,0 +1,109 @@
+package org.apache.kafka.connect.runtime.distributed;
+
+import org.apache.kafka.connect.util.ConnectorTaskId;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static org.apache.kafka.connect.util.ConnectUtils.combineCollections;
+
+class ClusterAssignment {
+
+ private final Map<String, Collection<String>> newlyAssignedConnectors;
+ private final Map<String, Collection<ConnectorTaskId>> newlyAssignedTasks;
+ private final Map<String, Collection<String>> newlyRevokedConnectors;
+ private final Map<String, Collection<ConnectorTaskId>> newlyRevokedTasks;
+ private final Map<String, Collection<String>> allAssignedConnectors;
+ private final Map<String, Collection<ConnectorTaskId>> allAssignedTasks;
+ private final Set<String> allWorkers;
+
+ public static final ClusterAssignment EMPTY = new ClusterAssignment(
+ Map.of(),
+ Map.of(),
+ Map.of(),
+ Map.of(),
+ Map.of(),
+ Map.of()
+ );
+
+ public ClusterAssignment(
+ Map<String, Collection<String>> newlyAssignedConnectors,
+ Map<String, Collection<ConnectorTaskId>> newlyAssignedTasks,
+ Map<String, Collection<String>> newlyRevokedConnectors,
+ Map<String, Collection<ConnectorTaskId>> newlyRevokedTasks,
+ Map<String, Collection<String>> allAssignedConnectors,
+ Map<String, Collection<ConnectorTaskId>> allAssignedTasks
+ ) {
+ this.newlyAssignedConnectors = newlyAssignedConnectors;
+ this.newlyAssignedTasks = newlyAssignedTasks;
+ this.newlyRevokedConnectors = newlyRevokedConnectors;
+ this.newlyRevokedTasks = newlyRevokedTasks;
+ this.allAssignedConnectors = allAssignedConnectors;
+ this.allAssignedTasks = allAssignedTasks;
+ this.allWorkers = combineCollections(
+ List.of(newlyAssignedConnectors, newlyAssignedTasks, newlyRevokedConnectors, newlyRevokedTasks, allAssignedConnectors, allAssignedTasks),
+ Map::keySet,
+ Collectors.toSet()
+ );
+ }
+
+ public Map<String, Collection<String>> newlyAssignedConnectors() {
+ return newlyAssignedConnectors;
+ }
+
+ public Collection<String> newlyAssignedConnectors(String worker) {
+ return newlyAssignedConnectors.getOrDefault(worker, Set.of());
+ }
+
+ public Map<String, Collection<ConnectorTaskId>> newlyAssignedTasks() {
+ return newlyAssignedTasks;
+ }
+
+ public Collection<ConnectorTaskId> newlyAssignedTasks(String worker) {
+ return newlyAssignedTasks.getOrDefault(worker, Set.of());
+ }
+
+ public Map<String, Collection<String>> newlyRevokedConnectors() {
+ return newlyRevokedConnectors;
+ }
+
+ public Collection<String> newlyRevokedConnectors(String worker) {
+ return newlyRevokedConnectors.getOrDefault(worker, Set.of());
+ }
+
+ public Map<String, Collection<ConnectorTaskId>> newlyRevokedTasks() {
+ return newlyRevokedTasks;
+ }
+
+ public Collection<ConnectorTaskId> newlyRevokedTasks(String worker) {
+ return newlyRevokedTasks.getOrDefault(worker, Set.of());
+ }
+
+ public Map<String, Collection<String>> allAssignedConnectors() {
+ return allAssignedConnectors;
+ }
+
+ public Map<String, Collection<ConnectorTaskId>> allAssignedTasks() {
+ return allAssignedTasks;
+ }
+
+ public Set<String> allWorkers() {
+ return allWorkers;
+ }
+
+ @Override
+ public String toString() {
+ return "ClusterAssignment{"
+ + "newlyAssignedConnectors=" + newlyAssignedConnectors
+ + ", newlyAssignedTasks=" + newlyAssignedTasks
+ + ", newlyRevokedConnectors=" + newlyRevokedConnectors
+ + ", newlyRevokedTasks=" + newlyRevokedTasks
+ + ", allAssignedConnectors=" + allAssignedConnectors
+ + ", allAssignedTasks=" + allAssignedTasks
+ + ", allWorkers=" + allWorkers
+ + '}';
+ }
+}
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..8793ec67908 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
@@ -20,11 +20,13 @@ import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.clients.MetadataRecoveryStrategy;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.config.ConfigDef.Importance;
import org.apache.kafka.common.config.ConfigException;
import org.apache.kafka.common.config.TopicConfig;
import org.apache.kafka.common.security.auth.SecurityProtocol;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.connect.runtime.WorkerConfig;
+import org.apache.kafka.connect.uber.distributed.WorkloadProvider;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -255,6 +257,22 @@ 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_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";
+ public static final String UBER_WORKER_BYTE_RATE_THRESHOLD_CONFIG = "uber.worker.byte.rate.threshold";
+ public static final double UBER_WORKER_BYTE_RATE_THRESHOLD_DEFAULT = 70 * 1024 * 1024;
+ public static final String UBER_WORKER_BYTE_RATE_THRESHOLD_DOC = "The maximum total byte rate "
+ + "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";
+ public static final String UBER_DEFAULT_PARTITION_BYTE_RATE_CONFIG = "uber.default.partition.byte.rate";
+ public static final double UBER_DEFAULT_PARTITION_BYTE_RATE_DEFAULT = 8 * 1024 * 1024;
+ public static final String UBER_DEFAULT_PARTITION_BYTE_RATE_DOC = "The default byte rate to assume "
+ + "for partitions that do not have a known byte rate";
+
private static Object defaultKeyGenerationAlgorithm(Crypto crypto) {
try {
validateKeyAlgorithm(crypto, INTER_WORKER_KEY_GENERATION_ALGORITHM_CONFIG, INTER_WORKER_KEY_GENERATION_ALGORITHM_DEFAULT);
@@ -309,7 +327,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 +605,18 @@ public final class DistributedConfig extends WorkerConfig {
return getString(GROUP_ID_CONFIG);
}
+ public String uberWorkloadProvider() {
+ return getString(UBER_WORKLOAD_PROVIDER_CONFIG);
+ }
+
+ public double uberWorkerByteRateThreshold() {
+ return getDouble(UBER_WORKER_BYTE_RATE_THRESHOLD_CONFIG);
+ }
+
+ public double uberDefaultPartitionByteRate() {
+ return getDouble(UBER_DEFAULT_PARTITION_BYTE_RATE_CONFIG);
+ }
+
@Override
protected Map<String, Object> postProcessParsedConfig(final Map<String, Object> parsedValues) {
CommonClientConfigs.warnDisablingExponentialBackoff(this);
@@ -722,4 +752,25 @@ public final class DistributedConfig extends WorkerConfig {
return result;
}
+ private static ConfigDef addUberExtensions(ConfigDef baseConfigDef) {
+ return baseConfigDef
+ .define(UBER_WORKLOAD_PROVIDER_CONFIG,
+ ConfigDef.Type.STRING,
+ 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_BYTE_RATE_THRESHOLD_CONFIG,
+ ConfigDef.Type.DOUBLE,
+ UBER_WORKER_BYTE_RATE_THRESHOLD_DEFAULT,
+ atLeast(0),
+ Importance.MEDIUM,
+ UBER_WORKER_BYTE_RATE_THRESHOLD_DOC)
+ .define(UBER_DEFAULT_PARTITION_BYTE_RATE_CONFIG,
+ ConfigDef.Type.DOUBLE,
+ UBER_DEFAULT_PARTITION_BYTE_RATE_DEFAULT,
+ atLeast(0),
+ Importance.MEDIUM,
+ UBER_DEFAULT_PARTITION_BYTE_RATE_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..fa7eb29da21 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
@@ -71,6 +71,7 @@ import org.apache.kafka.connect.storage.ClusterConfigState;
import org.apache.kafka.connect.storage.ConfigBackingStore;
import org.apache.kafka.connect.storage.PrivilegedWriteException;
import org.apache.kafka.connect.storage.StatusBackingStore;
+import org.apache.kafka.connect.uber.distributed.WorkloadProvider;
import org.apache.kafka.connect.util.Callback;
import org.apache.kafka.connect.util.ConnectUtils;
import org.apache.kafka.connect.util.ConnectorTaskId;
@@ -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 workloadProviderClass = config.uberWorkloadProvider();
+ final WorkloadProvider workloadProvider;
+ if (workloadProviderClass != null) {
+ workloadProvider = plugins().newPlugin(workloadProviderClass, config, WorkloadProvider.class);
+ } else {
+ workloadProvider = null;
+ }
+
this.member = member != null
? member
: new WorkerGroupMember(config, restUrl, this.configBackingStore,
- new RebalanceListener(time), time, clientId, logContext);
+ new RebalanceListener(time), workloadProvider, time, clientId, logContext);
this.herderExecutor = new ThreadPoolExecutor(1, 1, 0L,
TimeUnit.MILLISECONDS,
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeAssignor.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeAssignor.java
index 2b8f87c81c7..d731b88fd04 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeAssignor.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeAssignor.java
@@ -874,102 +874,4 @@ public class IncrementalCooperativeAssignor implements ConnectAssignor {
});
}
- static class ClusterAssignment {
-
- private final Map<String, Collection<String>> newlyAssignedConnectors;
- private final Map<String, Collection<ConnectorTaskId>> newlyAssignedTasks;
- private final Map<String, Collection<String>> newlyRevokedConnectors;
- private final Map<String, Collection<ConnectorTaskId>> newlyRevokedTasks;
- private final Map<String, Collection<String>> allAssignedConnectors;
- private final Map<String, Collection<ConnectorTaskId>> allAssignedTasks;
- private final Set<String> allWorkers;
-
- public static final ClusterAssignment EMPTY = new ClusterAssignment(
- Map.of(),
- Map.of(),
- Map.of(),
- Map.of(),
- Map.of(),
- Map.of()
- );
-
- public ClusterAssignment(
- Map<String, Collection<String>> newlyAssignedConnectors,
- Map<String, Collection<ConnectorTaskId>> newlyAssignedTasks,
- Map<String, Collection<String>> newlyRevokedConnectors,
- Map<String, Collection<ConnectorTaskId>> newlyRevokedTasks,
- Map<String, Collection<String>> allAssignedConnectors,
- Map<String, Collection<ConnectorTaskId>> allAssignedTasks
- ) {
- this.newlyAssignedConnectors = newlyAssignedConnectors;
- this.newlyAssignedTasks = newlyAssignedTasks;
- this.newlyRevokedConnectors = newlyRevokedConnectors;
- this.newlyRevokedTasks = newlyRevokedTasks;
- this.allAssignedConnectors = allAssignedConnectors;
- this.allAssignedTasks = allAssignedTasks;
- this.allWorkers = combineCollections(
- List.of(newlyAssignedConnectors, newlyAssignedTasks, newlyRevokedConnectors, newlyRevokedTasks, allAssignedConnectors, allAssignedTasks),
- Map::keySet,
- Collectors.toSet()
- );
- }
-
- public Map<String, Collection<String>> newlyAssignedConnectors() {
- return newlyAssignedConnectors;
- }
-
- public Collection<String> newlyAssignedConnectors(String worker) {
- return newlyAssignedConnectors.getOrDefault(worker, Set.of());
- }
-
- public Map<String, Collection<ConnectorTaskId>> newlyAssignedTasks() {
- return newlyAssignedTasks;
- }
-
- public Collection<ConnectorTaskId> newlyAssignedTasks(String worker) {
- return newlyAssignedTasks.getOrDefault(worker, Set.of());
- }
-
- public Map<String, Collection<String>> newlyRevokedConnectors() {
- return newlyRevokedConnectors;
- }
-
- public Collection<String> newlyRevokedConnectors(String worker) {
- return newlyRevokedConnectors.getOrDefault(worker, Set.of());
- }
-
- public Map<String, Collection<ConnectorTaskId>> newlyRevokedTasks() {
- return newlyRevokedTasks;
- }
-
- public Collection<ConnectorTaskId> newlyRevokedTasks(String worker) {
- return newlyRevokedTasks.getOrDefault(worker, Set.of());
- }
-
- public Map<String, Collection<String>> allAssignedConnectors() {
- return allAssignedConnectors;
- }
-
- public Map<String, Collection<ConnectorTaskId>> allAssignedTasks() {
- return allAssignedTasks;
- }
-
- public Set<String> allWorkers() {
- return allWorkers;
- }
-
- @Override
- public String toString() {
- return "ClusterAssignment{"
- + "newlyAssignedConnectors=" + newlyAssignedConnectors
- + ", newlyAssignedTasks=" + newlyAssignedTasks
- + ", newlyRevokedConnectors=" + newlyRevokedConnectors
- + ", newlyRevokedTasks=" + newlyRevokedTasks
- + ", allAssignedConnectors=" + allAssignedConnectors
- + ", allAssignedTasks=" + allAssignedTasks
- + ", allWorkers=" + allWorkers
- + '}';
- }
- }
-
}
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..c00f8bf0599
--- /dev/null
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberAssignor.java
@@ -0,0 +1,532 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.kafka.connect.runtime.distributed;
+
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.connect.runtime.distributed.WorkerCoordinator.ConnectorsAndTasks;
+import org.apache.kafka.connect.runtime.distributed.WorkerCoordinator.WorkerLoad;
+import org.apache.kafka.connect.storage.ClusterConfigState;
+import org.apache.kafka.connect.uber.distributed.WorkloadProvider;
+import org.apache.kafka.connect.util.ConnectUtils;
+import org.apache.kafka.connect.util.ConnectorTaskId;
+
+import org.slf4j.Logger;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Comparator;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.PriorityQueue;
+import java.util.Set;
+import java.util.TreeSet;
+import java.util.function.Function;
+import java.util.function.Predicate;
+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.IncrementalCooperativeConnectProtocol.CONNECT_PROTOCOL_V2;
+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 computes a distribution of connectors and tasks that takes task workload
+ * into account
+ */
+public class UberAssignor implements ConnectAssignor {
+
+ private final Logger log;
+ private final WorkloadProvider workloadProvider;
+ private final double workerByteRateThreshold;
+ private final double defaultPartitionByteRate;
+
+ public UberAssignor(
+ LogContext logContext,
+ WorkloadProvider workloadProvider,
+ double workerByteRateThreshold,
+ double defaultPartitionByteRate
+ ) {
+ this.log = logContext.logger(UberAssignor.class);
+ this.workloadProvider = workloadProvider;
+ this.workerByteRateThreshold = workerByteRateThreshold;
+ this.defaultPartitionByteRate = defaultPartitionByteRate;
+ }
+
+ @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, protocolVersion);
+ }
+ 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();
+ }
+
+ /**
+ * Performs task assignment based on the incremental cooperative connect protocol.
+ * Read more on the design and implementation in:
+ * <a href="https://cwiki.apache.org/confluence/display/KAFKA/KIP-415%3A+Incremental+Cooperative+Rebalancing+in+Kafka+Connect">
+ * KIP-415</a>
+ *
+ * @param leaderId the ID of the group leader
+ * @param configSnapshot the latest known snapshot of the configuration topic
+ * @param memberConfigs the metadata of all the members of the group as gather in the current
+ * round of rebalancing
+ * @param coordinator the worker coordinator instance that provide the configuration snapshot
+ * and get assigned the leader state during this assignment
+ * @param protocolVersion the Connect subprotocol version
+ * @return the serialized assignment of tasks to the whole group, including assigned or
+ * revoked tasks
+ */
+ protected Map<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()
+ )
+ );
+ 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, protocolVersion);
+ }
+
+ // Visible for testing
+ ClusterAssignment performTaskAssignment(
+ ClusterConfigState configSnapshot,
+ Map<String, ConnectorsAndTasks> currentAssignments
+ ) {
+ Set<String> configuredConnectors = new TreeSet<>(configSnapshot.connectors());
+ Set<ConnectorTaskId> configuredTasks = combineCollections(configuredConnectors, configSnapshot::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);
+ log.debug("Configured assignments: {}", configured);
+
+ Map<String, ConnectorsAndTasks> nextAssignments = computeNextAssignment(configSnapshot, configured, currentAssignments);
+
+ 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,
+ addedTasks,
+ revokedConnectors,
+ revokedTasks,
+ nextConnectorAssignments,
+ nextTaskAssignments
+ );
+ }
+
+ private Map<String, ConnectorsAndTasks> computeNextAssignment(
+ ClusterConfigState configSnapshot,
+ ConnectorsAndTasks configured,
+ Map<String, ConnectorsAndTasks> currentAssignment
+ ) {
+ // 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 = new ConnectorsAndTasks.Builder()
+ .addAll(deleted)
+ .addAll(duplicated)
+ .build();
+
+ Collection<WorkerLoad> workerLoads = workerLoads(currentAssignment);
+
+ revokeAll(workerLoads, toRemove);
+ assignAll(configSnapshot, workerLoads, toAssign);
+
+ applyWorkloadThreshold(configSnapshot, workerLoads);
+
+ Map<String, ConnectorsAndTasks.Builder> resultBuilder = new HashMap<>();
+ workerLoads.forEach(worker ->
+ resultBuilder.computeIfAbsent(worker.worker(), w -> new ConnectorsAndTasks.Builder())
+ .addConnectors(worker.connectors())
+ .addTasks(worker.tasks())
+ );
+ return transformValues(resultBuilder, ConnectorsAndTasks.Builder::build);
+ }
+
+ 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 <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());
+ }
+
+ 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<ConnectorTaskId> tasksToStart = clusterAssignment.newlyAssignedTasks(member);
+ Collection<String> connectorsToStop = clusterAssignment.newlyRevokedConnectors(member);
+ Collection<ConnectorTaskId> tasksToStop = clusterAssignment.newlyRevokedTasks(member);
+ ExtendedAssignment assignment =
+ new ExtendedAssignment(protocolVersion, error, leaderId, leaderUrl, maxOffset,
+ connectorsToStart, tasksToStart, connectorsToStop, tasksToStop, 0);
+ 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, short protocolVersion) {
+ boolean sessioned = protocolVersion >= CONNECT_PROTOCOL_V2;
+ return assignments.entrySet()
+ .stream()
+ .collect(Collectors.toMap(
+ Map.Entry::getKey,
+ e -> IncrementalCooperativeConnectProtocol.serializeAssignment(e.getValue(), sessioned)));
+ }
+
+ private static ConnectorsAndTasks diff(ConnectorsAndTasks base,
+ ConnectorsAndTasks... toSubtract) {
+ Collection<String> connectors = new TreeSet<>(base.connectors());
+ Collection<ConnectorTaskId> tasks = new TreeSet<>(base.tasks());
+ for (ConnectorsAndTasks sub : toSubtract) {
+ connectors.removeAll(sub.connectors());
+ tasks.removeAll(sub.tasks());
+ }
+ return ConnectorsAndTasks.of(connectors, tasks);
+ }
+
+ 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 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)
+ );
+ }
+
+ /**
+ * 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<WorkerLoad> workerLoads, ConnectorsAndTasks toRevoke) {
+ workerLoads.forEach(w -> {
+ w.connectors().removeAll(toRevoke.connectors());
+ w.tasks().removeAll(toRevoke.tasks());
+ });
+ }
+
+ /**
+ * Assign all of the specified connectors and tasks across the workers
+ * @param configSnapshot the latest snapshot of the config topic
+ * @param workers the workers in the cluster; may not be null or empty
+ * @param toAssign the connectors and tasks to assign; may be empty, but not null
+ */
+ private void assignAll(ClusterConfigState configSnapshot, Collection<WorkerLoad> workers, ConnectorsAndTasks toAssign) {
+ assert !workers.isEmpty();
+
+ // Assign largest jobs first
+ List<ConnectorTaskId> sortedTasks = new ArrayList<>(toAssign.tasks());
+ sortedTasks.sort(taskWorkloadComparator(configSnapshot));
+
+ // Attempt to assign to least-loaded workers first
+ PriorityQueue<WorkerLoad> sortedWorkers = new PriorityQueue<>(workerTaskWorkloadComparator(configSnapshot));
+ sortedWorkers.addAll(workers);
+ for (ConnectorTaskId task : sortedTasks) {
+ WorkerLoad worker = sortedWorkers.poll();
+ worker.assign(task);
+ assert sortedWorkers.add(worker);
+ }
+
+ // Assign connectors without taking workload into account; just assign them to the workers with the
+ // fewest connectors first
+ Comparator<WorkerLoad> workerConnectorComparator = Comparator.comparing(WorkerLoad::connectorsSize);
+ sortedWorkers = new PriorityQueue<>(workerConnectorComparator);
+ sortedWorkers.addAll(workers);
+ for (String connector : toAssign.connectors()) {
+ WorkerLoad worker = sortedWorkers.poll();
+ worker.assign(connector);
+ assert sortedWorkers.add(worker);
+ }
+ }
+
+ /**
+ * Attempt to ensure that no worker is assigned a workload that causes it to exceed the configured
+ * thresholds for byte and message rate, by reassigning tasks across workers if necessary (and if that
+ * reassignment would not cause the new worker to become overloaded itself)
+ * @param workers the workers in the cluster; may not be empty or null
+ */
+ private void applyWorkloadThreshold(ClusterConfigState configSnapshot, Collection<WorkerLoad> workers) {
+ assert !workers.isEmpty();
+
+ /*
+ - Load-balancing reassignment process:
+ - Sort workers by load
+ - Iterating from most-loaded to least-loaded worker:
+ - While worker is overloaded:
+ - Iterating from least-heavy to most-heavy job on worker:
+ - Attempt to reassign job to least-loaded worker that can accept job
+ without exceeding threshold
+ - If job cannot be reassigned without exceeding threshold, stop trying
+ to take load off of worker (and move on to next overloaded worker,
+ if there is one)
+ */
+ List<WorkerLoad> overloadedWorkers = workers.stream()
+ .filter(w -> isOverloaded(configSnapshot, w))
+ .sorted(workerTaskWorkloadComparator(configSnapshot))
+ .toList();
+
+ PriorityQueue<WorkerLoad> nonOverloadedWorkers = new PriorityQueue<>(workerTaskWorkloadComparator(configSnapshot).reversed());
+ workers.stream().filter(Predicate.not(w -> isOverloaded(configSnapshot, w))).forEach(nonOverloadedWorkers::add);
+
+ for (WorkerLoad overloaded : overloadedWorkers) {
+ List<ConnectorTaskId> sortedTasks = new ArrayList<>(overloaded.tasks());
+ sortedTasks.sort(taskWorkloadComparator(configSnapshot));
+
+ for (ConnectorTaskId task : sortedTasks) {
+ WorkerLoad leastLoaded = nonOverloadedWorkers.peek();
+ if (leastLoaded == null) {
+ return; // Nothing else we can do, every worker is overloaded
+ }
+
+ if (!canAccept(configSnapshot, leastLoaded, task)) {
+ // Least-heavy job on this worker cannot be assigned to the least-loaded
+ // worker in the cluster; nothing left we can do
+ break;
+ }
+
+ // Reassign from the overloaded worker to the non-overloaded one
+ overloaded.tasks().remove(task);
+ leastLoaded.tasks().add(task);
+
+ // Re-insert into the queue to reflect new workload
+ nonOverloadedWorkers.poll();
+ nonOverloadedWorkers.add(leastLoaded);
+
+ // Check to see if the worker is still overloaded
+ if (!isOverloaded(configSnapshot, overloaded)) {
+ // If the worker is no longer overloaded, we can stop trying to transfer
+ // tasks off of it, and add it to the non-overloaded queue
+ // so that it can potentially receive new work from other overloaded workers
+ nonOverloadedWorkers.add(overloaded);
+ break;
+ }
+ }
+ }
+ }
+
+ private Comparator<WorkerLoad> workerTaskWorkloadComparator(ClusterConfigState configSnapshot) {
+ return Comparator.comparing(worker -> workload(configSnapshot, worker));
+ }
+
+ private Comparator<ConnectorTaskId> taskWorkloadComparator(ClusterConfigState configSnapshot) {
+ return Comparator.comparing(task -> workload(configSnapshot, task));
+ }
+
+ private boolean canAccept(ClusterConfigState configSnapshot, WorkerLoad worker, ConnectorTaskId task) {
+ return workload(configSnapshot, worker) + workload(configSnapshot, task) < workerByteRateThreshold;
+ }
+
+ private boolean isOverloaded(ClusterConfigState configSnapshot, WorkerLoad worker) {
+ return workload(configSnapshot, worker) >= workerByteRateThreshold;
+ }
+
+ private double workload(ClusterConfigState configSnapshot, WorkerLoad worker) {
+ return worker.tasks().stream().mapToDouble(w -> workload(configSnapshot, w)).sum();
+ }
+
+ private double workload(ClusterConfigState configSnapshot, ConnectorTaskId task) {
+ Map<String, String> taskConfig = configSnapshot.taskConfig(task);
+ if (taskConfig == null) {
+ log.warn("Unable to find config for task {}; falling back on default workload of {}", task, defaultPartitionByteRate);
+ return defaultPartitionByteRate;
+ }
+
+ String connector = task.connector();
+ Map<String, String> connectorConfig = configSnapshot.connectorConfig(connector);
+ if (connectorConfig == null) {
+ log.warn("Unable to find config for connector {}; falling back on default workload of {}", connector, defaultPartitionByteRate);
+ return defaultPartitionByteRate;
+ }
+
+ Double byteRate;
+ try {
+ byteRate = workloadProvider.byteRate(connector, task.task(), connectorConfig, taskConfig);
+ } catch (Exception e) {
+ log.error("Failed to retrieve workload for {}; falling back on default workload of {}", task, defaultPartitionByteRate, e);
+ return defaultPartitionByteRate;
+ }
+
+ if (byteRate == null) {
+ log.debug("No byte rate available for {}; assuming default of {}", task, defaultPartitionByteRate);
+ return defaultPartitionByteRate;
+ }
+
+ if (byteRate < 0) {
+ log.warn("Negative byte rate provided for {}; will fall back on default of {}", task, defaultPartitionByteRate);
+ return defaultPartitionByteRate;
+ }
+
+ log.trace("Found byte rate of {} for task {}", byteRate, task);
+
+ return byteRate;
+ }
+
+ private static List<WorkerLoad> workerLoads(Map<String, ConnectorsAndTasks> memberAssignments) {
+ return memberAssignments.entrySet().stream()
+ .map(e -> new WorkerLoad.Builder(e.getKey()).withCopies(e.getValue().connectors(), e.getValue().tasks()).build())
+ .collect(Collectors.toList());
+ }
+
+}
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..41117899589 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
@@ -28,6 +28,7 @@ import org.apache.kafka.common.utils.Time;
import org.apache.kafka.common.utils.Timer;
import org.apache.kafka.connect.storage.ClusterConfigState;
import org.apache.kafka.connect.storage.ConfigBackingStore;
+import org.apache.kafka.connect.uber.distributed.WorkloadProvider;
import org.apache.kafka.connect.util.ConnectorTaskId;
import org.slf4j.Logger;
@@ -70,12 +71,14 @@ 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;
/**
* Initialize the coordination manager.
*/
- public WorkerCoordinator(GroupRebalanceConfig config,
+ public WorkerCoordinator(DistributedConfig connectConfig,
+ GroupRebalanceConfig groupConfig,
LogContext logContext,
ConsumerNetworkClient client,
Metrics metrics,
@@ -84,9 +87,10 @@ public class WorkerCoordinator extends AbstractCoordinator implements Closeable
String restUrl,
ConfigBackingStore configStorage,
WorkerRebalanceListener listener,
+ WorkloadProvider workloadProvider,
ConnectProtocolCompatibility protocolCompatibility,
int maxDelay) {
- super(config,
+ super(groupConfig,
logContext,
client,
metrics,
@@ -100,10 +104,20 @@ public class WorkerCoordinator extends AbstractCoordinator implements Closeable
this.listener = listener;
this.rejoinRequested = false;
this.protocolCompatibility = protocolCompatibility;
+
+ if (workloadProvider != null) {
+ double workerByteRateThreshold = connectConfig.uberWorkerByteRateThreshold();
+ double defaultPartitionByteRate = 8 * 1024 * 1024;
+ this.uberAssignor = new UberAssignor(logContext, workloadProvider, workerByteRateThreshold, defaultPartitionByteRate);
+ } else {
+ log.warn("No workload provider given; workload-based load balancing will be disabled");
+ this.uberAssignor = null;
+ }
+
this.incrementalAssignor = new IncrementalCooperativeAssignor(logContext, time, maxDelay);
this.eagerAssignor = new EagerAssignor(logContext);
this.currentConnectProtocol = protocolCompatibility;
- this.coordinatorDiscoveryTimeoutMs = config.heartbeatIntervalMs;
+ this.coordinatorDiscoveryTimeoutMs = groupConfig.heartbeatIntervalMs;
this.lastCompletedGenerationId = Generation.NO_GENERATION.generationId;
}
@@ -229,9 +243,17 @@ 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 {
+ // TODO: Configurable (possibly via feature flag)?
+ boolean useUberAssignor = true;
+ assignor = useUberAssignor && uberAssignor != null ? uberAssignor : incrementalAssignor;
+ }
+
+ return assignor.performAssignment(leaderId, protocolCompatibility, allMemberMetadata, this);
}
@Override
@@ -469,6 +491,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..bccb0fd3bc6 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
@@ -20,6 +20,7 @@ import org.apache.kafka.clients.ApiVersions;
import org.apache.kafka.clients.ClientUtils;
import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.clients.GroupRebalanceConfig;
+import org.apache.kafka.clients.GroupRebalanceConfig.ProtocolType;
import org.apache.kafka.clients.Metadata;
import org.apache.kafka.clients.MetadataRecoveryStrategy;
import org.apache.kafka.clients.NetworkClient;
@@ -40,6 +41,7 @@ import org.apache.kafka.common.utils.Time;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.connect.runtime.WorkerConfig;
import org.apache.kafka.connect.storage.ConfigBackingStore;
+import org.apache.kafka.connect.uber.distributed.WorkloadProvider;
import org.apache.kafka.connect.util.ConnectorTaskId;
import org.slf4j.Logger;
@@ -68,6 +70,7 @@ public class WorkerGroupMember {
private final String clientId;
private final ConsumerNetworkClient client;
private final Metrics metrics;
+ private final WorkloadProvider workloadProvider;
private final WorkerCoordinator coordinator;
private boolean stopped = false;
@@ -76,6 +79,7 @@ public class WorkerGroupMember {
String restUrl,
ConfigBackingStore configStorage,
WorkerRebalanceListener listener,
+ WorkloadProvider workloadProvider,
Time time,
String clientId,
LogContext logContext) {
@@ -133,8 +137,10 @@ public class WorkerGroupMember {
retryBackoffMs,
config.getInt(CommonClientConfigs.REQUEST_TIMEOUT_MS_CONFIG),
Integer.MAX_VALUE);
+ this.workloadProvider = workloadProvider;
this.coordinator = new WorkerCoordinator(
- new GroupRebalanceConfig(config, GroupRebalanceConfig.ProtocolType.CONNECT),
+ config,
+ new GroupRebalanceConfig(config, ProtocolType.CONNECT),
logContext,
this.client,
metrics,
@@ -143,6 +149,7 @@ public class WorkerGroupMember {
restUrl,
configStorage,
listener,
+ workloadProvider,
ConnectProtocolCompatibility.compatibility(config.getString(DistributedConfig.CONNECT_PROTOCOL_CONFIG)),
config.getInt(DistributedConfig.SCHEDULED_REBALANCE_MAX_DELAY_MS_CONFIG));
@@ -230,6 +237,7 @@ public class WorkerGroupMember {
Utils.closeQuietly(coordinator, "coordinator", firstException);
Utils.closeQuietly(metrics, "consumer metrics", firstException);
Utils.closeQuietly(client, "consumer network client", firstException);
+ Utils.closeQuietly(workloadProvider, "workload provider", 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/isolation/PluginScanResult.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginScanResult.java
index 04be35f2d26..6bdcb81f686 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
@@ -25,6 +25,7 @@ import org.apache.kafka.connect.storage.Converter;
import org.apache.kafka.connect.storage.HeaderConverter;
import org.apache.kafka.connect.transforms.Transformation;
import org.apache.kafka.connect.transforms.predicates.Predicate;
+import org.apache.kafka.connect.uber.distributed.WorkloadProvider;
import java.util.List;
import java.util.SortedSet;
@@ -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<WorkloadProvider>> workloadProviders;
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<WorkloadProvider>> workloadProviders
) {
this.sinkConnectors = sinkConnectors;
this.sourceConnectors = sourceConnectors;
@@ -65,9 +68,10 @@ public class PluginScanResult {
this.configProviders = configProviders;
this.restExtensions = restExtensions;
this.connectorClientConfigPolicies = connectorClientConfigPolicies;
+ this.workloadProviders = workloadProviders;
this.allPlugins =
List.of(sinkConnectors, sourceConnectors, converters, headerConverters, transformations, predicates,
- configProviders, restExtensions, connectorClientConfigPolicies);
+ configProviders, restExtensions, connectorClientConfigPolicies, workloadProviders);
}
/**
@@ -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::workloadProviders)
);
}
@@ -131,6 +136,10 @@ public class PluginScanResult {
return connectorClientConfigPolicies;
}
+ public SortedSet<PluginDesc<WorkloadProvider>> workloadProviders() {
+ return workloadProviders;
+ }
+
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..74b07295cd0 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
@@ -25,6 +25,7 @@ import org.apache.kafka.connect.storage.Converter;
import org.apache.kafka.connect.storage.HeaderConverter;
import org.apache.kafka.connect.transforms.Transformation;
import org.apache.kafka.connect.transforms.predicates.Predicate;
+import org.apache.kafka.connect.uber.distributed.WorkloadProvider;
import com.fasterxml.jackson.annotation.JsonValue;
@@ -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),
+ WORKLOAD_PROVIDER(WorkloadProvider.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..e0c68efb6bc 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.WORKLOAD_PROVIDER, 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..c24de2b01d3 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.WORKLOAD_PROVIDER, source)
);
}
diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeAssignorTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeAssignorTest.java
index 84bb8b145ad..5d685f2a033 100644
--- a/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeAssignorTest.java
+++ b/connect/runtime/src/test/java/org/apache/kafka/connect/runtime/distributed/IncrementalCooperativeAssignorTest.java
@@ -47,7 +47,6 @@ import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
-import static org.apache.kafka.connect.runtime.distributed.IncrementalCooperativeAssignor.ClusterAssignment;
import static org.apache.kafka.connect.runtime.distributed.WorkerCoordinator.WorkerLoad;
import static org.apache.kafka.connect.util.ConnectUtils.transformValues;
import static org.junit.jupiter.api.Assertions.assertEquals;
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..cb19c88c6c2 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
@@ -91,6 +91,8 @@ public class WorkerCoordinatorIncrementalTest {
private MockRebalanceListener rebalanceListener;
@Mock
private KafkaConfigBackingStore configStorage;
+ @Mock
+ private DistributedConfig connectConfig;
private GroupRebalanceConfig rebalanceConfig;
private WorkerCoordinator coordinator;
private int rebalanceDelay = DistributedConfig.SCHEDULED_REBALANCE_MAX_DELAY_MS_DEFAULT;
@@ -150,7 +152,8 @@ public class WorkerCoordinatorIncrementalTest {
null,
retryBackoffMs,
retryBackoffMaxMs);
- this.coordinator = new WorkerCoordinator(rebalanceConfig,
+ this.coordinator = new WorkerCoordinator(connectConfig,
+ rebalanceConfig,
loggerFactory,
consumerClient,
metrics,
@@ -159,6 +162,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..0656219c4b6 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
@@ -104,6 +104,7 @@ public class WorkerCoordinatorTest {
private ConsumerNetworkClient consumerClient;
private MockRebalanceListener rebalanceListener;
@Mock private KafkaConfigBackingStore configStorage;
+ @Mock private DistributedConfig connectConfig;
private GroupRebalanceConfig rebalanceConfig;
private WorkerCoordinator coordinator;
@@ -141,7 +142,8 @@ public class WorkerCoordinatorTest {
null,
retryBackoffMs,
retryBackoffMaxMs);
- this.coordinator = new WorkerCoordinator(rebalanceConfig,
+ this.coordinator = new WorkerCoordinator(connectConfig,
+ rebalanceConfig,
logContext,
consumerClient,
metrics,
@@ -150,6 +152,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/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..499bed57324
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Config.java
@@ -0,0 +1,261 @@
+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.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 org.apache.kafka.common.serialization.ByteArrayDeserializer;
+import org.apache.kafka.connect.runtime.ConnectorConfig;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.uber.data.kafka.common.UberKafkaCluster;
+import com.uber.data.kafka.exception.ResolverException;
+import com.uber.data.kafka.resolvers.HostPortResolver;
+import com.uber.data.kafka.resolvers.brokers.BrokerResolver;
+import com.uber.data.kafka.resolvers.common.Source;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Properties;
+
+import static com.uber.data.kafka.connect.utils.UberConnectUtils.addError;
+import static com.uber.data.kafka.connect.utils.UberConnectUtils.createErrorMessage;
+
+public abstract class UReplicator3Config extends AbstractConfig {
+
+ private static final Logger log = LoggerFactory.getLogger(UReplicator3Config.class);
+
+ private static final HostPortResolver BROKER_RESOLVER;
+ static {
+ HostPortResolver brokerResolver;
+ try {
+ brokerResolver = new BrokerResolver.Builder(new Properties())
+ .addResolver(Source.YAML)
+ .addResolver(Source.UNS)
+ .build();
+ } catch (Throwable t) {
+ log.warn("Failed to initialize broker resolver; this should only happen in local testing environments", t);
+ brokerResolver = null;
+ }
+ BROKER_RESOLVER = brokerResolver;
+ }
+
+ 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 TOPIC_RENAME_PREFIX = "topic.rename.";
+
+ 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.";
+ public static final String SOURCE_ADMIN_PREFIX = "source.admin.";
+ public static final String TARGET_ADMIN_PREFIX = "target.admin.";
+
+ 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;
+
+ protected final Map<String, String> topicRenames;
+
+ protected UReplicator3Config(ConfigDef config, Map<String, ?> props) {
+ super(config, props);
+ this.topicRenames = Collections.unmodifiableMap(parseTopicRenames(originalsStrings(), null));
+ }
+
+ 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
+ ).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 String targetTopic(String sourceTopic) {
+ return topicRenames.getOrDefault(sourceTopic, sourceTopic);
+ }
+
+ public String connectorName() {
+ return getString(ConnectorConfig.NAME_CONFIG);
+ }
+
+ public Admin sourceAdmin() {
+ return admin(
+ sourceBootstrap(),
+ clientId(true, "admin"),
+ SOURCE_CLIENTS_PREFIX,
+ SOURCE_ADMIN_PREFIX
+ );
+ }
+
+ public Admin targetAdmin() {
+ return admin(
+ targetBootstrap(),
+ clientId(false, "admin"),
+ TARGET_CLIENTS_PREFIX,
+ TARGET_ADMIN_PREFIX
+ );
+ }
+
+ public Consumer<byte[], byte[]> sourceConsumer() {
+ return consumer(
+ sourceBootstrap(),
+ clientId(true, "consumer"),
+ 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);
+ }
+
+ private String clientId(boolean source, String clientType) {
+ return "ureplicator3-%s-%s-%s-%s".formatted(
+ role(),
+ connectorName(),
+ 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);
+ }
+
+ protected 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;
+ }
+
+ @VisibleForTesting
+ static String resolveBootstrap(String prop, String cluster) {
+ if (BROKER_RESOLVER == null)
+ 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"
+ );
+
+ @SuppressWarnings("deprecation")
+ UberKafkaCluster kafkaCluster = new UberKafkaCluster(cluster);
+ try {
+ return kafkaCluster.getBrokerConnectionString(BROKER_RESOLVER);
+ } 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..51e279aa236
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Connector.java
@@ -0,0 +1,186 @@
+package com.uber.data.kafka.connect.ureplicator3;
+
+import com.uber.data.kafka.connect.utils.ExecutorUtils;
+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.connect.connector.Task;
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.source.SourceConnector;
+import org.apache.kafka.connect.util.ConnectorUtils;
+
+import com.uber.data.kafka.connect.utils.UberConnectUtils;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+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;
+
+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 ScheduledExecutorService executor;
+
+ // The values for these may be changed at any time, but the contents of any map will never be mutated
+ private volatile Map<String, Integer> sourcePartitionCounts;
+ private volatile Map<String, Integer> targetPartitionCounts;
+
+ @Override
+ public ConfigDef config() {
+ return UReplicator3ConnectorConfig.config();
+ }
+
+ @Override
+ public String version() {
+ return "0.1.0";
+ }
+
+ @Override
+ public Class<? extends Task> taskClass() {
+ return UReplicator3Task.class;
+ }
+
+ @Override
+ public void start(Map<String, String> props) {
+ config = new UReplicator3ConnectorConfig(props);
+
+ // Synchronously collect partition counts the first time around
+ // to make sure they're available before we generate task configs
+ fetchPartitionCounts();
+
+ long refreshIntervalMs = config.partitionCountRefreshIntervalMs();
+ executor = ExecutorUtils.singleThreadedScheduled("executor-" + config.connectorName());
+ executor.scheduleWithFixedDelay(this::refreshPartitionCounts, refreshIntervalMs, refreshIntervalMs, TimeUnit.MILLISECONDS);
+ }
+
+ @Override
+ public List<Map<String, String>> taskConfigs(int maxTasks) {
+ Map<String, Integer> sourcePartitionCountsSnapshot = sourcePartitionCounts;
+ Map<String, Integer> targetPartitionCountsSnapshot = targetPartitionCounts;
+
+ List<TopicPartition> topicPartitions = 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;
+ }
+ }).flatMap(e -> IntStream.range(0, e.getValue())
+ .mapToObj(partition -> new TopicPartition(e.getKey(), partition))
+ ).toList();
+
+ int numTasks = Math.min(maxTasks, topicPartitions.size());
+
+ // We don't do workload-based load balancing here; let the rebalancing protocol handle that for now
+ List<List<TopicPartition>> taskPartitions = ConnectorUtils.groupPartitions(topicPartitions, numTasks);
+
+ return IntStream.range(0, taskPartitions.size())
+ .mapToObj(i -> config.taskConfig(sourcePartitionCountsSnapshot, targetPartitionCountsSnapshot, taskPartitions.get(i), i))
+ .collect(Collectors.toList());
+ }
+
+ @Override
+ public void stop() {
+ }
+
+ @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;
+ }
+
+ private void refreshPartitionCounts() {
+ try {
+ fetchPartitionCounts();
+ } catch (Exception e) {
+ log.error("Failed to refresh partition counts", e);
+ }
+ }
+
+ private void fetchPartitionCounts() {
+ Map<String, Integer> oldSourceCounts = sourcePartitionCounts;
+ Map<String, Integer> newSourceCounts;
+ try (Admin sourceAdmin = config.sourceAdmin()) {
+ newSourceCounts = partitionCounts(sourceAdmin, config.sourceTopics(), "source");
+ }
+
+ Map<String, Integer> oldTargetCounts = targetPartitionCounts;
+ Map<String, Integer> newTargetCounts;
+ try (Admin targetAdmin = config.targetAdmin()) {
+ newTargetCounts = partitionCounts(targetAdmin, config.targetTopics(), "target");
+ }
+
+ // Only update these if we're able to load both successfully
+ sourcePartitionCounts = Collections.unmodifiableMap(newSourceCounts);
+ targetPartitionCounts = Collections.unmodifiableMap(newTargetCounts);
+
+ // Old counts will be null if this is the first time we've loaded partitions, in which case
+ // we don't need to log anything or explicitly request a task reconfiguration
+ if (oldSourceCounts != null && !oldSourceCounts.equals(sourcePartitionCounts)) {
+ log.info("Detected change in source partitions; will generate new task configs");
+ context.requestTaskReconfiguration();
+ } else if (oldTargetCounts != null && !oldTargetCounts.equals(targetPartitionCounts)) {
+ log.info("Detected change in target partitions; will generate new task configs");
+ context.requestTaskReconfiguration();
+ }
+ }
+
+ 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);
+ }
+ }
+}
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..5257c590a2e
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3ConnectorConfig.java
@@ -0,0 +1,282 @@
+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.TopicPartition;
+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 org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.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.TARGET_PARTITIONS_OVERRIDE_PREFIX;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3TaskConfig.TASK_ID_CONFIG;
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3TaskConfig.TOPIC_PARTITIONS_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.encodeTopicPartitions;
+import static com.uber.data.kafka.connect.utils.UberConnectUtils.listTopics;
+import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
+
+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 PARTITION_COUNT_REFRESH_INTERVAL_MS_CONFIG = "partition.count.refresh.interval.ms";
+ public static final long PARTITION_COUNT_REFRESH_INTERVAL_MS_DEFAULT = TimeUnit.MINUTES.toMillis(1);
+ private static final String PARTITION_COUNT_REFRESH_INTERVAL_MS_DOC = "How often to re-fetch the "
+ + "partition counts of replicated topics on the source and target clusters";
+
+ private final Set<String> sourceTopics;
+ private final Set<String> targetTopics;
+
+ public UReplicator3ConnectorConfig(Map<String, ?> props) {
+ super(config(), props);
+ 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(
+ PARTITION_COUNT_REFRESH_INTERVAL_MS_CONFIG,
+ Type.LONG,
+ PARTITION_COUNT_REFRESH_INTERVAL_MS_DEFAULT,
+ atLeast(1),
+ Importance.LOW,
+ PARTITION_COUNT_REFRESH_INTERVAL_MS_DOC
+ );
+ }
+
+ @Override
+ protected String role() {
+ return "connector";
+ }
+
+ 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);
+ }
+
+ public Set<String> sourceTopics() {
+ return sourceTopics;
+ }
+
+ public Set<String> targetTopics() {
+ return targetTopics;
+ }
+
+ public long partitionCountRefreshIntervalMs() {
+ return getLong(PARTITION_COUNT_REFRESH_INTERVAL_MS_CONFIG);
+ }
+
+ public Map<String, String> taskConfig(
+ Map<String, Integer> sourcePartitionCounts,
+ Map<String, Integer> targetPartitionCounts,
+ Collection<TopicPartition> sourcePartitions,
+ int taskId
+ ) {
+ Map<String, String> result = new HashMap<>(originalsStrings());
+
+ String encodedPartitions = encodeTopicPartitions(sourcePartitions);
+ result.put(TOPIC_PARTITIONS_CONFIG, encodedPartitions);
+
+ result.put(TASK_ID_CONFIG, Integer.toString(taskId));
+
+ sourcePartitions.stream().map(TopicPartition::topic).distinct().forEach(sourceTopic -> {
+ Integer sourcePartitionCount = sourcePartitionCounts.get(sourceTopic);
+ if (sourcePartitionCount == null) {
+ log.warn("Could not find partition count for topic {} on source cluster; this should never happen", sourceTopic);
+ return;
+ }
+
+ String targetTopic = targetTopic(sourceTopic);
+ Integer targetPartitionCount = targetPartitionCounts.get(targetTopic);
+ if (targetPartitionCount == null) {
+ log.warn("Could not find partition count for topic {} on target cluster; this should never happen", targetTopic);
+ return;
+ }
+
+ if (!sourcePartitionCount.equals(targetPartitionCount)) {
+ log.debug(
+ "Storing partition override for source topic {} ({} partitions) and target topic {} ({} partitions)",
+ sourceTopic, sourcePartitionCount, targetTopic, targetPartitionCount
+ );
+ result.put(TARGET_PARTITIONS_OVERRIDE_PREFIX + targetTopic, Integer.toString(targetPartitionCount));
+ }
+ });
+
+ // Only include information on renames for topics actually assigned to the task
+ Set<String> sourceTopics = sourcePartitions.stream()
+ .map(TopicPartition::topic)
+ .collect(Collectors.toSet());
+ result.keySet().removeIf(key -> {
+ if (!key.startsWith(TOPIC_RENAME_PREFIX))
+ return false;
+ String sourceTopic = key.substring(TOPIC_RENAME_PREFIX.length());
+ return !sourceTopics.contains(sourceTopic);
+ });
+
+ 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 static Set<String> computeTargetTopics(Set<String> sourceTopics, Map<String, String> topicRenames) {
+ return sourceTopics.stream()
+ .map(t -> topicRenames.getOrDefault(t, t))
+ .collect(Collectors.toSet());
+ }
+
+}
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..339da756e44
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Task.java
@@ -0,0 +1,211 @@
+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.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.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.List;
+import java.util.Map;
+
+public class UReplicator3Task extends SourceTask {
+
+ private static final Logger log = LoggerFactory.getLogger(UReplicator3Task.class);
+
+ static final String SOURCE_PARTITION_SOURCE_CLUSTER = "source.cluster";
+ static final String SOURCE_PARTITION_TARGET_CLUSTER = "target.cluster";
+ static final String SOURCE_PARTITION_KAFKA_TOPIC = "topic";
+ static final String SOURCE_PARTITION_KAFKA_PARTITION = "partition";
+ static final String SOURCE_OFFSET_KAFKA_OFFSET = "offset";
+
+ private UReplicator3TaskConfig config;
+ private Consumer<byte[], byte[]> sourceConsumer;
+ private WorkloadTrackingProcessor workloadTracker;
+
+ @Override
+ public void start(Map<String, String> props) {
+ config = new UReplicator3TaskConfig(props);
+ initializeSourceConsumer();
+ initializeWorkloadTracker();
+ }
+
+ @Override
+ public List<SourceRecord> poll() {
+ ConsumerRecords<byte[], byte[]> records = sourceConsumer.poll(Duration.ofSeconds(3));
+
+ 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);
+ }
+
+ return result;
+ }
+
+ @Override
+ public void stop() {
+ Utils.closeQuietly(sourceConsumer, "source consumer");
+ Utils.closeQuietly(workloadTracker, "workload tracker");
+ }
+
+ @Override
+ public String version() {
+ return "0.1.0";
+ }
+
+ @Override
+ public void commitRecord(SourceRecord record, RecordMetadata metadata) {
+ // TODO: Report to aakafkaoffsetmgmt
+ }
+
+ private void initializeSourceConsumer() {
+ sourceConsumer = config.sourceConsumer();
+ sourceConsumer.assign(config.sourceTopicPartitions());
+
+ List<Map<String, Object>> sourcePartitions = config.sourceTopicPartitions().stream()
+ .map(tp -> encodeSourcePartition(tp.topic(), tp.partition()))
+ .toList();
+ context.offsetStorageReader().offsets(sourcePartitions).forEach((sourcePartition, sourceOffset) -> {
+ TopicPartition topicPartition = decodeSourcePartition(sourcePartition);
+ if (topicPartition == null) {
+ return;
+ }
+
+ Long offset = decodeSourceOffset(topicPartition, sourceOffset);
+ if (offset == null) {
+ return;
+ }
+
+ sourceConsumer.seek(topicPartition, offset);
+ });
+ }
+
+ private void initializeWorkloadTracker() {
+ // TODO: Enable workload tracking to be performed even if the workload provider isn't set up on this worker
+ Duration workloadReportInterval = config.workloadReportInterval();
+ if (workloadReportInterval != null) {
+ workloadTracker = new WorkloadTrackingProcessor(
+ config.sourceCluster(),
+ config.targetCluster(),
+ KafkaBasedWorkloadProvider::publishWorkload,
+ config::sourceTopicPartitions,
+ workloadReportInterval
+ );
+ workloadTracker.start();
+ }
+ }
+
+ private SourceRecord convertRecord(ConsumerRecord<byte[], byte[]> record) {
+ String sourceTopic = record.topic();
+ String targetTopic = config.targetTopic(sourceTopic);
+ int sourcePartition = record.partition();
+ int targetPartition = convertPartition(targetTopic, sourcePartition);
+
+ return new SourceRecord(
+ encodeSourcePartition(sourceTopic, sourcePartition),
+ encodeSourceOffset(record),
+ targetTopic,
+ targetPartition,
+ 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 Map.of(
+ SOURCE_PARTITION_SOURCE_CLUSTER, config.sourceCluster(),
+ SOURCE_PARTITION_TARGET_CLUSTER, config.targetCluster(),
+ SOURCE_PARTITION_KAFKA_TOPIC, topic,
+ SOURCE_PARTITION_KAFKA_PARTITION, partition
+ );
+ }
+
+ private Map<String, Object> encodeSourceOffset(ConsumerRecord<?, ?> record) {
+ return Map.of(
+ SOURCE_OFFSET_KAFKA_OFFSET, record.offset()
+ );
+ }
+
+ private TopicPartition decodeSourcePartition(Map<String, Object> 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, Object> 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("Ignorning invalid source offset {}; field {} is not a number", sourceOffset, SOURCE_OFFSET_KAFKA_OFFSET);
+ return null;
+ }
+
+ return ((Number) rawOffset).longValue();
+ }
+
+ private int convertPartition(String targetTopic, int sourcePartition) {
+ Integer targetPartitionsOverride = config.targetPartitionsOverride(targetTopic);
+ return targetPartitionsOverride != null ? sourcePartition % targetPartitionsOverride : sourcePartition;
+ }
+
+ 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;
+ }
+
+}
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..f6c4a334951
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3TaskConfig.java
@@ -0,0 +1,115 @@
+package com.uber.data.kafka.connect.ureplicator3;
+
+import org.apache.kafka.common.TopicPartition;
+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.utils.UberConnectUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.time.Duration;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+public class UReplicator3TaskConfig extends UReplicator3Config {
+
+ private static final Logger log = LoggerFactory.getLogger(UReplicator3TaskConfig.class);
+
+ public static final String TOPIC_PARTITIONS_CONFIG = "topic.partitions";
+
+ public static final String TASK_ID_CONFIG = "task.id";
+
+ public static final String WORKLOAD_REPORT_INTERVAL_MS_CONFIG = "workload.report.interval.ms";
+
+ public static final String TARGET_PARTITIONS_OVERRIDE_PREFIX = "target.partitions.override.";
+
+ private final Map<String, Integer> targetPartitionsOverrides;
+
+ public UReplicator3TaskConfig(Map<String, ?> props) {
+ super(config(), props);
+ this.targetPartitionsOverrides = Collections.unmodifiableMap(parseTargetPartitionsOverrides(originalsStrings()));
+ }
+
+ protected static ConfigDef config() {
+ return baseConfig()
+ .defineInternal(
+ TOPIC_PARTITIONS_CONFIG,
+ Type.LIST,
+ ConfigDef.NO_DEFAULT_VALUE,
+ Importance.HIGH
+ ).defineInternal(
+ TASK_ID_CONFIG,
+ Type.INT,
+ ConfigDef.NO_DEFAULT_VALUE,
+ Importance.HIGH
+ ).defineInternal(
+ WORKLOAD_REPORT_INTERVAL_MS_CONFIG,
+ Type.LONG,
+ Duration.ofMinutes(10).toMillis(),
+ Importance.LOW
+ );
+ }
+
+ @Override
+ protected String role() {
+ return "task-" + getInt(TASK_ID_CONFIG);
+ }
+
+ public Set<TopicPartition> sourceTopicPartitions() {
+ return getList(TOPIC_PARTITIONS_CONFIG).stream()
+ .map(UberConnectUtils::parseTopicPartition)
+ .collect(Collectors.toSet());
+ }
+
+ public Duration workloadReportInterval() {
+ Long durationMs = getLong(WORKLOAD_REPORT_INTERVAL_MS_CONFIG);
+ return durationMs == null || durationMs <= 0
+ ? null
+ : Duration.ofMillis(durationMs);
+ }
+
+ public Integer targetPartitionsOverride(String targetTopic) {
+ return targetPartitionsOverrides.get(targetTopic);
+ }
+
+ private static Map<String, Integer> parseTargetPartitionsOverrides(Map<String, String> originals) {
+ Map<String, Integer> result = new HashMap<>();
+ originals.forEach((key, value) -> {
+ if (!key.startsWith(TARGET_PARTITIONS_OVERRIDE_PREFIX))
+ return;
+
+ String targetTopic = key.substring(TARGET_PARTITIONS_OVERRIDE_PREFIX.length());
+
+ if (targetTopic.isBlank()) {
+ log.warn("Ignoring invalid property {} with value {}", key, value);
+ }
+
+ if (value == null || value.isBlank()) {
+ log.warn("Ignoring property {} with invalid value {}", key, value);
+ return;
+ }
+
+ int numPartitions;
+ try {
+ numPartitions = Integer.parseInt(value, 10);
+ } catch (NumberFormatException e) {
+ log.warn("Ignoring property {} with invalid value {}", key, value, e);
+ return;
+ }
+
+ if (numPartitions <= 0) {
+ log.warn("Ignoring property {} with invalid value {}", key, value);
+ return;
+ }
+
+ result.putIfAbsent(targetTopic, numPartitions);
+ });
+
+ return result;
+ }
+}
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..0fc2c5994f0
--- /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 LOGGER = 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) {
+ LOGGER.warn("Could not shut down {} in time", executorDescription);
+ } else {
+ LOGGER.info("Gracefully shut down {}", executorDescription);
+ }
+ } catch (InterruptedException e) {
+ LOGGER.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/UberConnectUtils.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/UberConnectUtils.java
new file mode 100644
index 00000000000..340326aa308
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/UberConnectUtils.java
@@ -0,0 +1,134 @@
+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.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Objects;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
+
+public interface UberConnectUtils {
+
+ Logger log = LoggerFactory.getLogger(UberConnectUtils.class);
+
+ @SafeVarargs
+ static <E> Set<E> diff(Set<E> base, Set<E>... toSubtract) {
+ Set<E> result = new HashSet<>(base);
+ for (Set<E> subtracted : toSubtract) {
+ result.removeAll(subtracted);
+ }
+ return result;
+ }
+
+ 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());
+ }
+
+ 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;
+ }
+
+ static String encodeTopicPartitions(Collection<TopicPartition> partitions) {
+ return partitions.stream()
+ .map(Object::toString)
+ .collect(Collectors.joining(","));
+ }
+
+ static TopicPartition parseTopicPartition(String topicPartitionString) {
+ int sep = topicPartitionString.lastIndexOf('-');
+
+ String topic = topicPartitionString.substring(0, sep);
+ String partitionString = topicPartitionString.substring(sep + 1);
+ int partition = Integer.parseInt(partitionString);
+
+ return new TopicPartition(topic, partition);
+ }
+
+ /**
+ * 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
+ */
+ 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);
+ }
+
+ 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;
+ }
+
+ static void addError(Config configValidationResult, String prop, Object value, String baseErrorMessage, Throwable error) {
+ addError(configValidationResult, prop, value, createErrorMessage(baseErrorMessage, error));
+ }
+
+ static void addError(Config configValidationResult, String prop, Object value, String errorMessage) {
+ findOrCreate(configValidationResult, prop, value)
+ .addErrorMessage(errorMessage);
+ }
+
+ 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;
+ }
+}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/ReplicatedTopicPartition.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/ReplicatedTopicPartition.java
new file mode 100644
index 00000000000..8b4696ab499
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/ReplicatedTopicPartition.java
@@ -0,0 +1,67 @@
+package com.uber.data.kafka.connect.workload.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);
+ }
+
+ @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/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..84771a3720f
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/Workload.java
@@ -0,0 +1,60 @@
+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;
+ }
+}
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..e45e6b94477
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/WorkloadCache.java
@@ -0,0 +1,231 @@
+package com.uber.data.kafka.connect.workload.common;
+
+import org.apache.kafka.common.utils.Time;
+
+import com.google.common.annotations.VisibleForTesting;
+
+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 LOGGER = 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())) {
+ 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) {
+ LOGGER.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;
+ }
+
+ /**
+ * 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);
+ }
+ }
+
+ private 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/kafka/KafkaBasedWorkloadProvider.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadProvider.java
new file mode 100644
index 00000000000..886be62b6ed
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadProvider.java
@@ -0,0 +1,170 @@
+package com.uber.data.kafka.connect.workload.kafka;
+
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.utils.Utils;
+import org.apache.kafka.connect.mirror.MirrorConnectorConfig;
+import org.apache.kafka.connect.mirror.MirrorSourceConfig;
+import org.apache.kafka.connect.uber.distributed.WorkloadProvider;
+
+import com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig;
+import com.uber.data.kafka.connect.ureplicator3.UReplicator3TaskConfig;
+import com.uber.data.kafka.connect.utils.UberConnectUtils;
+import com.uber.data.kafka.connect.workload.common.ReplicatedTopicPartition;
+import com.uber.data.kafka.connect.workload.common.Workload;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import static org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLASS_CONFIG;
+
+public class KafkaBasedWorkloadProvider implements WorkloadProvider {
+
+ private static final Logger log = LoggerFactory.getLogger(KafkaBasedWorkloadProvider.class);
+
+ private static KafkaBasedWorkloadStore workloadStore;
+
+ private Workload defaultPartitionWorkload;
+
+ @Override
+ public void configure(Map<String, ?> props) {
+ log.info("Starting provider");
+ initializeStore(props);
+ defaultPartitionWorkload = Workload.EMPTY; // TODO: Configurable
+ log.info("Finished Starting provider");
+ }
+
+ @Override
+ public Double byteRate(String connector, int taskId, Map<String, String> connectorConfig, Map<String, String> taskConfig) {
+ Set<ReplicatedTopicPartition> replicatedTopicPartitions = assignedPartitions(connector, taskId, connectorConfig, taskConfig);
+ if (replicatedTopicPartitions == null) {
+ return null;
+ }
+
+ return replicatedTopicPartitions.stream()
+ .map(rtp -> workloadStore.getWorkload(rtp.srcKafka(), rtp.dstKafka(), rtp.topicPartition()))
+ .map(w -> w != null ? w : defaultPartitionWorkload)
+ .mapToDouble(Workload::getBytesPerSecond)
+ .sum();
+ }
+
+ @Override
+ public void close() {
+ log.info("Shutting down provider");
+ destroyStore();
+ log.info("Finished shutting down provider");
+ }
+
+ @Override
+ public String version() {
+ return "0.1.0";
+ }
+
+ public static void publishWorkload(String srcKafka, String dstKafka, TopicPartition topicPartition, Workload workload) {
+ assertStoreInitialized();
+
+ workloadStore.putWorkload(srcKafka, dstKafka, topicPartition, workload);
+ }
+
+ public static void refreshWorkloadCache() {
+ assertStoreInitialized();
+
+ workloadStore.refresh();
+ }
+
+ public static Map<ReplicatedTopicPartition, Workload> getAllWorkloads() {
+ assertStoreInitialized();
+
+ return workloadStore.getAllWorkloads();
+ }
+
+ private static void assertStoreInitialized() {
+ if (workloadStore == null)
+ throw new IllegalStateException("Workload store has not been initialized yet, or has already been shut down");
+ }
+
+ private static synchronized void initializeStore(Map<String, ?> props) {
+ if (workloadStore != null) {
+ log.warn("Provider appears to have been configured multiple times on the same JVM. This should only happen during testing");
+ return;
+ }
+ KafkaBasedWorkloadProviderConfig config = new KafkaBasedWorkloadProviderConfig(props);
+ workloadStore = new KafkaBasedWorkloadStore(config);
+ workloadStore.start();
+ }
+
+ private static synchronized void destroyStore() {
+ Utils.closeQuietly(workloadStore, "workload store");
+ workloadStore = null;
+ }
+
+ private static Set<ReplicatedTopicPartition> assignedPartitions(String connector, int taskId, Map<String, String> connectorConfig, Map<String, String> taskConfig) {
+ String connectorClass = connectorConfig.get(CONNECTOR_CLASS_CONFIG);
+ if (connectorClass == null) {
+ log.warn("No connector class found for connector {}", connector);
+ return null;
+ }
+
+ String srcKafka;
+ String dstKafka;
+ String assignedPartitionsProp;
+ if (isURep3(connectorClass)) {
+ srcKafka = connectorConfig.get(UReplicator3ConnectorConfig.SOURCE_CLUSTER_CONFIG);
+ dstKafka = connectorConfig.get(UReplicator3ConnectorConfig.TARGET_CLUSTER_CONFIG);
+ assignedPartitionsProp = UReplicator3TaskConfig.TOPIC_PARTITIONS_CONFIG;
+ } else if (isMM2(connectorClass)) {
+ srcKafka = connectorConfig.get(MirrorConnectorConfig.SOURCE_CLUSTER_ALIAS);
+ dstKafka = connectorConfig.get(MirrorConnectorConfig.TARGET_CLUSTER_ALIAS);
+ assignedPartitionsProp = MirrorSourceConfig.TASK_TOPIC_PARTITIONS;
+ } else {
+ log.warn("Ignoring task {}-{} with unsupported connector type {}", connector, taskId, connectorClass);
+ return null;
+ }
+
+ Collection<TopicPartition> partitions = assignedPartitions(assignedPartitionsProp, taskConfig);
+ if (partitions == null) {
+ log.warn("Unable to parse assigned partitions for task {}-{}", connector, taskId);
+ return null;
+ }
+
+ return partitions.stream()
+ .map(tp -> new ReplicatedTopicPartition(srcKafka, dstKafka, tp))
+ .collect(Collectors.toSet());
+ }
+
+ private static boolean isMM2(String connectorClass) {
+ return Set.of(
+ "org.apache.kafka.connect.mirror.MirrorSourceConnector",
+ "MirrorSourceConnector",
+ "MirrorSource"
+ ).contains(connectorClass);
+ }
+
+ private static boolean isURep3(String connectorClass) {
+ return Set.of(
+ "com.uber.data.kafka.connect.UReplicator3Connector",
+ "UReplicator3Connector",
+ "UReplicator3"
+ ).contains(connectorClass);
+ }
+
+ private static Set<TopicPartition> assignedPartitions(String prop, Map<String, String> taskConfig) {
+ String rawPartitions = taskConfig.get(prop);
+
+ @SuppressWarnings("unchecked")
+ List<String> partitionsList = (List<String>) ConfigDef.parseType(prop, rawPartitions, ConfigDef.Type.LIST);
+ if (partitionsList == null) {
+ return null;
+ }
+
+ return partitionsList.stream()
+ .map(UberConnectUtils::parseTopicPartition)
+ .collect(Collectors.toSet());
+ }
+}
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..16174fb48fe
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadProviderConfig.java
@@ -0,0 +1,126 @@
+package com.uber.data.kafka.connect.workload.kafka;
+
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.common.config.AbstractConfig;
+import org.apache.kafka.common.config.ConfigDef;
+import org.apache.kafka.common.config.ConfigDef.Importance;
+import org.apache.kafka.common.config.ConfigDef.Type;
+import org.apache.kafka.common.config.TopicConfig;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+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.";
+
+ 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 KafkaBasedWorkloadProviderConfig(Map<String, ?> props) {
+ super(config(), props);
+ }
+
+ private static ConfigDef config() {
+ return new ConfigDef()
+ .define(
+ CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG,
+ Type.STRING,
+ Importance.HIGH,
+ "Bootstrap servers for the Kafka cluster hosting the workload topic"
+ ).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.MEDIUM,
+ STORAGE_TOPIC_PARTITIONS_DOC
+ ).define(
+ STORAGE_TOPIC_REPLICATION_FACTOR_CONFIG,
+ Type.SHORT,
+ STORAGE_TOPIC_REPLICATION_FACTOR_DEFAULT,
+ atLeast(1),
+ Importance.MEDIUM,
+ 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 String bootstrapServers() {
+ return getString(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG);
+ }
+
+ 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;
+ }
+}
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..289fefdb80b
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadStore.java
@@ -0,0 +1,316 @@
+package com.uber.data.kafka.connect.workload.kafka;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.uber.data.kafka.connect.utils.UberConnectUtils;
+import com.uber.data.kafka.connect.workload.common.WorkloadCache;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.clients.admin.NewTopic;
+import org.apache.kafka.clients.consumer.ConsumerConfig;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.common.KafkaException;
+import org.apache.kafka.common.TopicPartition;
+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.workload.common.ReplicatedTopicPartition;
+import com.uber.data.kafka.connect.workload.common.Workload;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+
+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 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();
+ Map<String, Object> baseClientProps = Map.of(CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, config.bootstrapServers());
+ this.producerProps = producerProps(baseClientProps);
+ this.consumerProps = consumerProps(baseClientProps);
+ this.adminProps = baseClientProps;
+ 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(mutableCopy(adminProps));
+ this.kafkaLog = new KafkaBasedLog<>(
+ topic,
+ mutableCopy(producerProps),
+ mutableCopy(consumerProps),
+ () -> topicAdmin,
+ this::onConsume,
+ Time.SYSTEM,
+ this::createAndValidateTopic
+ );
+ kafkaLog.start();
+ }
+
+ private static Map<String, Object> producerProps(Map<String, Object> baseProps) {
+ Map<String, Object> result = mutableCopy(baseProps);
+ result.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
+ result.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
+ result.put(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG, Integer.MAX_VALUE);
+ return result;
+ }
+
+ private static Map<String, Object> consumerProps(Map<String, Object> baseProps) {
+ Map<String, Object> result = mutableCopy(baseProps);
+ result.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
+ result.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
+ return result;
+ }
+
+ private static <K, V> Map<K, V> mutableCopy(Map<K, V> map) {
+ return new HashMap<>(map);
+ }
+
+ 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 void 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;
+ }
+
+ 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;
+ }
+
+ // TODO: Error handling (probably just log and emit a metric)
+ kafkaLog.send(key, value);
+ }
+
+ public Workload getWorkload(String srcKafka, String dstKafka, TopicPartition topicPartition) {
+ ReplicatedTopicPartition replicatedTopicPartition = new ReplicatedTopicPartition(srcKafka, dstKafka, topicPartition);
+ synchronized (this) {
+ return workloadsCache.get(replicatedTopicPartition).orElse(null);
+ }
+ }
+
+ public synchronized Map<ReplicatedTopicPartition, Workload> getAllWorkloads() {
+ return workloadsCache.getAll();
+ }
+
+ public void refresh() {
+ try {
+ kafkaLog.readToEnd().get();
+ } catch (InterruptedException | ExecutionException e) {
+ throw new KafkaException("Failed to read to end of workload log", e);
+ }
+ }
+
+ 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 {
+ Map<String, Object> map = Map.of(
+ SRC_KAFKA_FIELD, replicatedTopicPartition.srcKafka(),
+ DST_KAFKA_FIELD, replicatedTopicPartition.dstKafka(),
+ TOPIC_FIELD, replicatedTopicPartition.topicPartition().topic(),
+ PARTITION_FIELD, replicatedTopicPartition.topicPartition().partition()
+ );
+ return OBJECT_MAPPER.writeValueAsBytes(map);
+ }
+
+ 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<>() { });
+
+ String srcKafka = parseStringField(true, map, SRC_KAFKA_FIELD);
+ String dstKafka = parseStringField(true, map, DST_KAFKA_FIELD);
+ String topic = parseStringField(true, map, TOPIC_FIELD);
+ Number partition = parseNumberField(true, map, 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());
+ }
+
+ private static String parseStringField(boolean isKey, Map<String, Object> map, String field) throws IOException {
+ Object fieldValue = map.get(field);
+ if (fieldValue == null)
+ throw new MissingFieldException(isKey, field);
+
+ if (!(fieldValue instanceof String))
+ throw new InvalidFieldTypeException(isKey, field, "string", fieldValue);
+
+ return (String) fieldValue;
+ }
+
+ private static Number parseNumberField(boolean isKey, Map<String, Object> map, String field) throws IOException {
+ Object fieldValue = map.get(field);
+ if (fieldValue == null)
+ throw new MissingFieldException(isKey, field);
+
+ if (!(fieldValue instanceof Number))
+ throw new InvalidFieldTypeException(isKey, field, "string", fieldValue);
+
+ return (Number) fieldValue;
+ }
+
+ private static class MissingFieldException extends IOException {
+ public MissingFieldException(boolean isKey, String field) {
+ super(
+ "record has invalid " + (isKey ? "key" : "value")
+ + ": missing '" + field + "' field"
+ );
+ }
+ }
+
+ private 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/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..6e26dd80876
--- /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 LOGGER = 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) {
+ LOGGER.warn("Ignoring invalid workload with negative number of bytes: {}", bytes);
+ return Optional.empty();
+ }
+ if (messages < 0) {
+ LOGGER.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) {
+ LOGGER.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..9806bb6ccf2
--- /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 LOGGER = 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 {
+ LOGGER.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..14d5f57fc68
--- /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 com.uber.data.kafka.connect.workload.transformation.WorkloadTrackingTransformation;
+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.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..df62f40cfef
--- /dev/null
+++ b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/transformation/WorkloadTrackingTransformation.java
@@ -0,0 +1,61 @@
+package com.uber.data.kafka.connect.workload.transformation;
+
+import com.uber.data.kafka.connect.workload.tracking.WorkloadStore;
+import com.uber.data.kafka.connect.workload.tracking.WorkloadTrackingProcessor;
+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 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/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..418531696c5
--- /dev/null
+++ b/connect/uber/src/main/resources/META-INF/services/org.apache.kafka.connect.source.SourceConnector
@@ -0,0 +1,16 @@
+ # Licensed to the Apache Software Foundation (ASF) under one or more
+ # contributor license agreements. See the NOTICE file distributed with
+ # this work for additional information regarding copyright ownership.
+ # The ASF licenses this file to You under the Apache License, Version 2.0
+ # (the "License"); you may not use this file except in compliance with
+ # the License. You may obtain a copy of the License at
+ #
+ # http://www.apache.org/licenses/LICENSE-2.0
+ #
+ # Unless required by applicable law or agreed to in writing, software
+ # distributed under the License is distributed on an "AS IS" BASIS,
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ # See the License for the specific language governing permissions and
+ # limitations under the License.
+
+com.uber.data.kafka.connect.ureplicator3.UReplicator3Connector
\ No newline at end of file
diff --git a/connect/uber/src/main/resources/META-INF/services/org.apache.kafka.connect.uber.distributed.WorkloadProvider b/connect/uber/src/main/resources/META-INF/services/org.apache.kafka.connect.uber.distributed.WorkloadProvider
new file mode 100644
index 00000000000..ec6d132c42c
--- /dev/null
+++ b/connect/uber/src/main/resources/META-INF/services/org.apache.kafka.connect.uber.distributed.WorkloadProvider
@@ -0,0 +1,16 @@
+ # Licensed to the Apache Software Foundation (ASF) under one or more
+ # contributor license agreements. See the NOTICE file distributed with
+ # this work for additional information regarding copyright ownership.
+ # The ASF licenses this file to You under the Apache License, Version 2.0
+ # (the "License"); you may not use this file except in compliance with
+ # the License. You may obtain a copy of the License at
+ #
+ # http://www.apache.org/licenses/LICENSE-2.0
+ #
+ # Unless required by applicable law or agreed to in writing, software
+ # distributed under the License is distributed on an "AS IS" BASIS,
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ # See the License for the specific language governing permissions and
+ # limitations under the License.
+
+com.uber.data.kafka.connect.workload.kafka.KafkaBasedWorkloadProvider
\ 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..9eaf35c4bea
--- /dev/null
+++ b/connect/uber/src/test/java/com/uber/data/kafka/connect/integration/ReplicationIntegrationTest.java
@@ -0,0 +1,800 @@
+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.NewPartitions;
+import org.apache.kafka.clients.consumer.ConsumerRecords;
+import org.apache.kafka.clients.producer.Producer;
+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.mirror.IdentityReplicationPolicy;
+import org.apache.kafka.connect.mirror.MirrorSourceConnector;
+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.transforms.RegexRouter;
+import org.apache.kafka.connect.util.clusters.EmbeddedConnectCluster;
+import org.apache.kafka.connect.util.clusters.EmbeddedKafkaCluster;
+
+import com.uber.data.kafka.connect.ureplicator3.UReplicator3Connector;
+import com.uber.data.kafka.connect.workload.common.ReplicatedTopicPartition;
+import com.uber.data.kafka.connect.workload.common.Workload;
+import com.uber.data.kafka.connect.workload.kafka.KafkaBasedWorkloadProvider;
+import com.uber.data.kafka.connect.workload.kafka.KafkaBasedWorkloadTrackingTransformation;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Tag;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+import java.nio.charset.StandardCharsets;
+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.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import java.util.stream.Stream;
+
+import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.PARTITION_COUNT_REFRESH_INTERVAL_MS_CONFIG;
+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.UReplicator3TaskConfig.TARGET_PARTITIONS_OVERRIDE_PREFIX;
+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 com.uber.data.kafka.connect.workload.transformation.WorkloadTrackingTransformationConfig.DST_CLUSTER_NAME_CONFIG;
+import static com.uber.data.kafka.connect.workload.transformation.WorkloadTrackingTransformationConfig.REPORT_INTERVAL_MS_CONFIG;
+import static com.uber.data.kafka.connect.workload.transformation.WorkloadTrackingTransformationConfig.SRC_CLUSTER_NAME_CONFIG;
+import static org.apache.kafka.connect.mirror.MirrorConnectorConfig.EMIT_OFFSET_SYNCS_ENABLED;
+import static org.apache.kafka.connect.mirror.MirrorConnectorConfig.REPLICATION_POLICY_CLASS;
+import static org.apache.kafka.connect.mirror.MirrorConnectorConfig.SOURCE_CLUSTER_ALIAS;
+import static org.apache.kafka.connect.mirror.MirrorConnectorConfig.TARGET_CLUSTER_ALIAS;
+import static org.apache.kafka.connect.mirror.MirrorSourceConfig.SYNC_TOPIC_ACLS_ENABLED;
+import static org.apache.kafka.connect.mirror.MirrorSourceConfig.SYNC_TOPIC_CONFIGS_ENABLED;
+import static org.apache.kafka.connect.mirror.MirrorSourceConfig.TOPICS;
+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.ConnectorConfig.TRANSFORMS_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.CONNECTOR_CLIENT_POLICY_CLASS_CONFIG;
+import static org.apache.kafka.connect.runtime.WorkerConfig.PLUGIN_DISCOVERY_CONFIG;
+import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.UBER_WORKLOAD_PROVIDER_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.assertTrue;
+import static org.junit.jupiter.api.Assertions.fail;
+
+@Tag("integration")
+public class ReplicationIntegrationTest {
+
+ private static final String WORKLOAD_TOPIC = "uber-workloads";
+
+ private EmbeddedConnectCluster connect;
+ private Map<String, EmbeddedKafkaCluster> extraKafkaClusters;
+
+ public enum ConnectorType {
+ MM2,
+ UREP3
+ }
+
+ @BeforeEach
+ public void setup() throws Exception {
+ Map<String, String> workerProps = new HashMap<>();
+ workerProps.put(UBER_WORKLOAD_PROVIDER_CONFIG, KafkaBasedWorkloadProvider.class.getName());
+ workerProps.put(STORAGE_TOPIC_CONFIG, WORKLOAD_TOPIC);
+ workerProps.put(STORAGE_TOPIC_REPLICATION_FACTOR_CONFIG, "1"); // We run a single embedded broker
+ workerProps.put(PLUGIN_DISCOVERY_CONFIG, PluginDiscoveryMode.SERVICE_LOAD.name()); // Fast startup time ⚡️
+ workerProps.put(CONNECTOR_CLIENT_POLICY_CLASS_CONFIG, "All");
+
+ Properties brokerProps = new Properties();
+
+ int numWorkers = 3;
+
+ extraKafkaClusters = new HashMap<>();
+
+ // build a Connect cluster backed by Kafka
+ connect = new EmbeddedConnectCluster.Builder()
+ .numWorkers(numWorkers)
+ .numBrokers(1)
+ .workerProps(workerProps)
+ .brokerProps(brokerProps)
+ .build();
+
+ connect.start();
+ connect.assertions().assertExactlyNumWorkersAreUp(numWorkers, "Cluster did not start in time");
+ }
+
+ @AfterEach
+ public void close() {
+ AtomicReference<Throwable> shutdownFailure = new AtomicReference<>();
+
+ // stop the Connect cluster and its backing Kafka cluster.
+ if (connect != null) {
+ Utils.closeQuietly(connect::stop, "Embedded Connect cluster", shutdownFailure);
+ }
+
+ // Stop every extra Kafka cluster
+ extraKafkaClusters.forEach((name, kafkaCluster) ->
+ Utils.closeQuietly(kafkaCluster::stop, "Embedded Kafka cluster '" + name + "'", shutdownFailure)
+ );
+
+ if (shutdownFailure.get() != null) {
+ throw new AssertionError("Failed to cleanly shut down testing resources", shutdownFailure.get());
+ }
+ }
+
+ @Test
+ public void testClusterStartupAndManualPublish() {
+ assertEquals(
+ Map.of(),
+ KafkaBasedWorkloadProvider.getAllWorkloads()
+ );
+
+ ReplicatedTopicPartition tp1 = new ReplicatedTopicPartition("src1", "dst1", "t", 0);
+ Workload w1 = new Workload(3.0, 4.0);
+ KafkaBasedWorkloadProvider.publishWorkload(tp1.srcKafka(), tp1.dstKafka(), tp1.topicPartition(), w1);
+
+ KafkaBasedWorkloadProvider.refreshWorkloadCache();
+
+ assertEquals(
+ Map.of(tp1, w1),
+ KafkaBasedWorkloadProvider.getAllWorkloads()
+ );
+
+ ReplicatedTopicPartition tp2 = new ReplicatedTopicPartition("src2", "dst2", "t", 0);
+ Workload w2 = new Workload(35.0, 14.0);
+ KafkaBasedWorkloadProvider.publishWorkload(tp2.srcKafka(), tp2.dstKafka(), tp2.topicPartition(), w2);
+
+ KafkaBasedWorkloadProvider.refreshWorkloadCache();
+
+ assertEquals(
+ Map.of(tp1, w1, tp2, w2),
+ KafkaBasedWorkloadProvider.getAllWorkloads()
+ );
+
+ KafkaBasedWorkloadProvider.publishWorkload(tp1.srcKafka(), tp1.dstKafka(), tp1.topicPartition(), null);
+
+ KafkaBasedWorkloadProvider.refreshWorkloadCache();
+
+ assertEquals(
+ Map.of(tp2, w2),
+ KafkaBasedWorkloadProvider.getAllWorkloads()
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(ConnectorType.class)
+ public void testNormalConnectors(ConnectorType connectorType) throws Exception {
+ String srcKafkaName = "src-kafka";
+ String dstKafkaName = "dst-kafka";
+
+ EmbeddedKafkaCluster srcCluster = startExtraKafkaCluster(srcKafkaName, 1, new Properties());
+ EmbeddedKafkaCluster dstCluster = startExtraKafkaCluster(dstKafkaName, 1, new Properties());
+
+ int numProduced = 0;
+ try (Producer<byte[], byte[]> srcProducer = srcCluster.createProducer(Map.of())) {
+ for (int i = 1; i <= 6; i++) {
+ String topic = "t" + i;
+ srcCluster.createTopic(topic, i);
+ dstCluster.createTopic(topic, i);
+
+ for (int partition = 0; partition < i; partition++) {
+ produceTestRecord(srcProducer, topic, partition, 0);
+ numProduced++;
+ }
+ }
+
+ srcProducer.flush();
+ }
+
+ String conn1 = "connector-one";
+ String conn2 = "connector-two";
+
+ startConnector(connectorType, conn1, srcKafkaName, dstKafkaName, 3, Set.of("t1", "t2", "t3"), null, false);
+ startConnector(connectorType, conn2, srcKafkaName, dstKafkaName, 1, Set.of("t4", "t5", "t6"), null, false);
+
+ dstCluster.consume(numProduced, TimeUnit.MINUTES.toMillis(1), "t1", "t2", "t3", "t4", "t5", "t6");
+ }
+
+ @ParameterizedTest
+ @EnumSource(ConnectorType.class)
+ public void testTopicRenaming(ConnectorType connectorType) throws Exception {
+ String srcKafkaName = "src-kafka";
+ String dstKafkaName = "dst-kafka";
+
+ EmbeddedKafkaCluster srcCluster = startExtraKafkaCluster(srcKafkaName, 1, new Properties());
+ EmbeddedKafkaCluster dstCluster = startExtraKafkaCluster(dstKafkaName, 1, new Properties());
+
+ int numProduced = 0;
+ try (Producer<byte[], byte[]> srcProducer = srcCluster.createProducer(Map.of())) {
+ for (int i = 1; i <= 6; i++) {
+ String srcTopic = "s" + i;
+ srcCluster.createTopic(srcTopic, i);
+ String dstTopic = "t" + i;
+ dstCluster.createTopic(dstTopic, i);
+
+ for (int partition = 0; partition < i; partition++) {
+ produceTestRecord(srcProducer, srcTopic, partition, 0);
+ numProduced++;
+ }
+ }
+
+ srcProducer.flush();
+ }
+
+ String conn1 = "connector-one";
+ String conn2 = "connector-two";
+
+ startConnector(connectorType, conn1, srcKafkaName, dstKafkaName, 3, Set.of("s1", "s2", "s3"), Map.of("s1", "t1", "s2", "t2", "s3", "t3"), false);
+ startConnector(connectorType, conn2, srcKafkaName, dstKafkaName, 1, Set.of("s4", "s5", "s6"), Map.of("s4", "t4", "s5", "t5", "s6", "t6"), false);
+
+ dstCluster.consume(numProduced, TimeUnit.MINUTES.toMillis(1), "t1", "t2", "t3", "t4", "t5", "t6");
+ }
+
+ @Test // MM2 doesn't handle this case, we only test URep3
+ public void testPartitionCountTracking() throws Exception {
+ String srcKafkaName = "src-kafka";
+ String dstKafkaName = "dst-kafka";
+
+ EmbeddedKafkaCluster srcCluster = startExtraKafkaCluster(srcKafkaName, 1, new Properties());
+ EmbeddedKafkaCluster dstCluster = startExtraKafkaCluster(dstKafkaName, 1, new Properties());
+
+ int numProduced = 0;
+ try (Producer<byte[], byte[]> srcProducer = srcCluster.createProducer(Map.of())) {
+ for (int i = 1; i <= 6; i++) {
+ String topic = "t" + i;
+ srcCluster.createTopic(topic, i);
+ dstCluster.createTopic(topic, 7 - i);
+
+ for (int partition = 0; partition < i; partition++) {
+ produceTestRecord(srcProducer, topic, partition, 0);
+ numProduced++;
+ }
+ }
+
+ srcProducer.flush();
+ }
+
+ String conn1 = "connector-one";
+ String conn2 = "connector-two";
+
+ startURep3Connector(
+ conn1,
+ srcKafkaName,
+ dstKafkaName,
+ 3,
+ Set.of("t1", "t2", "t3"),
+ null,
+ PARTITION_COUNT_REFRESH_INTERVAL_MS_CONFIG,
+ "1000"
+ );
+ startURep3Connector(
+ conn2,
+ srcKafkaName,
+ dstKafkaName,
+ 1,
+ Set.of("t4", "t5", "t6"),
+ null,
+ PARTITION_COUNT_REFRESH_INTERVAL_MS_CONFIG,
+ "1000"
+ );
+
+ ConsumerRecords<byte[], byte[]> records = dstCluster.consume(numProduced, TimeUnit.MINUTES.toMillis(1), "t1", "t2", "t3", "t4", "t5", "t6");
+ // Ensure that every possible target partition was written
+ Set<TopicPartition> expectedPartitions = IntStream.rangeClosed(1, 6).boxed().flatMap(i -> {
+ String topic = "t" + i;
+ int targetPartitions = 7 - i;
+ return IntStream.range(0, Math.min(i, targetPartitions))
+ .mapToObj(targetPartition -> new TopicPartition(topic, targetPartition));
+ }).collect(Collectors.toSet());
+ Set<TopicPartition> actualPartitions = records.partitions();
+ assertEquals(expectedPartitions, actualPartitions);
+
+ // Change the partition counts of the source cluster and produce more test traffic
+ int newNumSrcPartitions = 10;
+ Map<String, NewPartitions> newPartitions = IntStream.rangeClosed(1, 6)
+ .mapToObj(i -> "t" + i)
+ .collect(Collectors.toMap(
+ Function.identity(),
+ topic -> NewPartitions.increaseTo(newNumSrcPartitions)
+ ));
+ try (Admin srcAdmin = srcCluster.createAdminClient()) {
+ srcAdmin.createPartitions(newPartitions);
+ }
+ try (Producer<byte[], byte[]> srcProducer = srcCluster.createProducer(Map.of())) {
+ for (int i = 1; i <= 6; i++) {
+ String topic = "t" + i;
+ for (int partition = 0; partition < newNumSrcPartitions; partition++) {
+ produceTestRecord(srcProducer, topic, partition, 0);
+ numProduced++;
+ }
+ }
+
+ srcProducer.flush();
+ }
+ // Make sure all the records we produced (including ones to the newly-created partitions) were replicated
+ dstCluster.consume(numProduced, TimeUnit.MINUTES.toMillis(1), "t1", "t2", "t3", "t4", "t5", "t6");
+
+ // Change the partition counts--of the target cluster this time
+ int newNumDstPartitions = 20;
+ newPartitions = IntStream.rangeClosed(1, 6)
+ .mapToObj(i -> "t" + i)
+ .collect(Collectors.toMap(
+ Function.identity(),
+ topic -> NewPartitions.increaseTo(newNumDstPartitions)
+ ));
+ try (Admin dstAdmin = dstCluster.createAdminClient()) {
+ dstAdmin.createPartitions(newPartitions);
+ }
+
+ // Wait for the connectors to pick up on the partition count changes
+ waitForCondition(() -> {
+ Map<String, String> allTaskConfigs = new HashMap<>();
+ connect.taskConfigs(conn1).forEach(taskInfo -> allTaskConfigs.putAll(taskInfo.config()));
+ connect.taskConfigs(conn2).forEach(taskInfo -> allTaskConfigs.putAll(taskInfo.config()));
+ Map<String, String> expectedPartitionProps = IntStream.rangeClosed(1, 6)
+ .mapToObj(i -> "t" + i)
+ .collect(Collectors.toMap(
+ topic -> TARGET_PARTITIONS_OVERRIDE_PREFIX + topic,
+ topic -> Integer.toString(newNumDstPartitions)
+ ));
+ allTaskConfigs.keySet().retainAll(expectedPartitionProps.keySet());
+ return expectedPartitionProps.equals(allTaskConfigs);
+ }, TimeUnit.MINUTES.toMillis(1), "Connector did not generate new tasks in response to target partition count change in time");
+
+ // Produce even more test traffic
+ try (Producer<byte[], byte[]> srcProducer = srcCluster.createProducer(Map.of())) {
+ for (int i = 1; i <= 6; i++) {
+ String topic = "t" + i;
+ for (int partition = 0; partition < newNumSrcPartitions; partition++) {
+ produceTestRecord(srcProducer, topic, partition, 0);
+ numProduced++;
+ }
+ }
+
+ srcProducer.flush();
+ }
+
+ records = dstCluster.consume(numProduced, TimeUnit.MINUTES.toMillis(1), "t1", "t2", "t3", "t4", "t5", "t6");
+ // Ensure that every possible target partition was written
+ expectedPartitions = IntStream.rangeClosed(1, 6).boxed().flatMap(i -> {
+ String topic = "t" + i;
+ return IntStream.range(0, newNumSrcPartitions)
+ .mapToObj(targetPartition -> new TopicPartition(topic, targetPartition));
+ }).collect(Collectors.toSet());
+ actualPartitions = records.partitions();
+ assertEquals(expectedPartitions, actualPartitions);
+ }
+
+ @Test // MM2 doesn't validate topics on src/dst clusters, we only test URep3
+ public void testPreflightValidation() {
+ String srcKafkaName = "src-kafka";
+ String dstKafkaName = "dst-kafka";
+ String srcTopic = "src-topic";
+ String dstTopic = "dst-topic";
+
+ Map<String, String> connectorProps = baseURep3Config(srcKafkaName, dstKafkaName, 1, Set.of(srcTopic));
+ connectorProps.put(NAME_CONFIG, "connector-validation-test");
+ connectorProps.put(SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG, "bad-host-src:9092");
+ connectorProps.put(TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG, "bad-host-dst:9092");
+ validateURep3Config(connectorProps,
+ // Neither of the clusters has been started yet
+ SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG,
+ TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG
+ );
+
+ EmbeddedKafkaCluster srcCluster = startExtraKafkaCluster(srcKafkaName, 1, new Properties());
+ connectorProps.put(SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG, srcCluster.bootstrapServers());
+ validateURep3Config(connectorProps,
+ // Only the source cluster has been started; target cluster is still down
+ TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG,
+ // Topic hasn't been created on the source cluster yet
+ TOPICS_CONFIG
+ );
+
+ srcCluster.createTopic(srcTopic);
+ validateURep3Config(connectorProps,
+ // Only the source cluster has been started; target cluster is still down
+ TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG
+ );
+
+ EmbeddedKafkaCluster dstCluster = startExtraKafkaCluster(dstKafkaName, 1, new Properties());
+ connectorProps.put(TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG, dstCluster.bootstrapServers());
+ validateURep3Config(connectorProps,
+ // Topic hasn't been created on the target cluster yet
+ TOPICS_CONFIG
+ );
+
+ dstCluster.createTopic(srcTopic);
+ validateURep3Config(connectorProps);
+
+ String renamedTopicProp = TOPIC_RENAME_PREFIX + srcTopic;
+ connectorProps.put(renamedTopicProp, dstTopic);
+ validateURep3Config(connectorProps,
+ // Renamed topic doesn't exist on target cluster yet
+ renamedTopicProp
+ );
+
+ dstCluster.createTopic(dstTopic);
+ validateURep3Config(connectorProps); // Everything should work now
+
+ // Can't override individual client bootstrap servers
+ String srcAdminBootstrapProp = SOURCE_ADMIN_PREFIX + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG;
+ connectorProps.put(srcAdminBootstrapProp, srcCluster.bootstrapServers());
+ String dstAdminBootstrapProp = TARGET_ADMIN_PREFIX + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG;
+ connectorProps.put(dstAdminBootstrapProp, dstCluster.bootstrapServers());
+ String srcConsumerBootstrapProp = SOURCE_CONSUMER_PREFIX + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG;
+ connectorProps.put(srcConsumerBootstrapProp, srcCluster.bootstrapServers());
+ validateURep3Config(connectorProps,
+ srcAdminBootstrapProp,
+ dstAdminBootstrapProp,
+ srcConsumerBootstrapProp);
+
+ connectorProps.keySet().removeAll(Set.of(srcAdminBootstrapProp, dstAdminBootstrapProp, srcConsumerBootstrapProp));
+ String emptyRenamedSrcTopicProp = TOPIC_RENAME_PREFIX;
+ connectorProps.put(emptyRenamedSrcTopicProp, "some-dst-topic");
+ String blankRenamedSrcTopicProp = TOPIC_RENAME_PREFIX + " ";
+ connectorProps.put(blankRenamedSrcTopicProp, "some-dst-topic");
+ String emptyRenamedDstTopicProp = TOPIC_RENAME_PREFIX + "some-src-topic-1";
+ connectorProps.put(emptyRenamedDstTopicProp, "");
+ String blankRenamedDstTopicProp = TOPIC_RENAME_PREFIX + "some-src-topic-2";
+ connectorProps.put(blankRenamedDstTopicProp, " ");
+ validateURep3Config(connectorProps,
+ emptyRenamedSrcTopicProp,
+ blankRenamedSrcTopicProp,
+ emptyRenamedDstTopicProp,
+ blankRenamedDstTopicProp
+ );
+ }
+
+ @ParameterizedTest
+ @EnumSource(ConnectorType.class)
+ public void testWorkloadTracking(ConnectorType connectorType) throws Exception {
+ String srcKafka1Name = "src-kafka-1";
+ String srcKafka2Name = "src-kafka-2";
+ String dstKafka1Name = "dst-kafka-1";
+ String dstKafka2Name = "dst-kafka-2";
+
+ EmbeddedKafkaCluster srcCluster1 = startExtraKafkaCluster(srcKafka1Name, 1, new Properties());
+ EmbeddedKafkaCluster srcCluster2 = startExtraKafkaCluster(srcKafka2Name, 1, new Properties());
+ EmbeddedKafkaCluster dstCluster1 = startExtraKafkaCluster(dstKafka1Name, 1, new Properties());
+ EmbeddedKafkaCluster dstCluster2 = startExtraKafkaCluster(dstKafka2Name, 1, new Properties());
+
+ IntStream.rangeClosed(1, 7).forEach(i -> {
+ String topic = "t" + i;
+ // i = number of partitions in each topic
+ Stream.of(srcCluster1, srcCluster2, dstCluster1, dstCluster2).forEach(
+ cluster -> cluster.createTopic(topic, i)
+ );
+ });
+
+ String conn1 = "connector-one";
+ String conn2 = "connector-two";
+ String conn3 = "connector-three";
+ String conn4 = "connector-four";
+
+ startConnector(connectorType, conn1, srcKafka1Name, dstKafka1Name, 7, Set.of("t1", "t2", "t3", "t4", "t5", "t6", "t7"), null, true);
+ startConnector(connectorType, conn2, srcKafka1Name, dstKafka2Name, 7, Set.of("t1", "t2", "t3", "t4", "t5", "t6", "t7"), null, true);
+ // Only replicate a subset of topics from src 2
+ startConnector(connectorType, conn3, srcKafka2Name, dstKafka1Name, 7, Set.of("t1", "t3", "t4", "t6", "t7"), null, true);
+ startConnector(connectorType, conn4, srcKafka2Name, dstKafka2Name, 7, Set.of("t1", "t3", "t4", "t6", "t7"), null, true);
+
+ int numProduced = 0;
+ try (Producer<byte[], byte[]> src1Producer = srcCluster1.createProducer(Map.of())) {
+ for (int t = 1; t <= 3; t++) {
+ String topic = "t" + t;
+ for (int i = 0; i < 1000; i++) {
+ int partition = i % t;
+ produceTestRecord(src1Producer, topic, partition, i);
+ numProduced++;
+ }
+ }
+ for (int t = 4; t <= 6; t++) {
+ String topic = "t" + t;
+ for (int i = 0; i < 100; i++) {
+ int partition = i % t;
+ produceTestRecord(src1Producer, topic, partition, i);
+ numProduced++;
+ }
+ }
+ // Intentionally refrain from producing to t7 to verify that empty partitions are still tracked
+ src1Producer.flush();
+ }
+
+ try (Producer<byte[], byte[]> src2Producer = srcCluster2.createProducer(Map.of())) {
+ for (int t : Set.of(1, 3)) {
+ String topic = "t" + t;
+ for (int i = 0; i < 100; i++) {
+ int partition = i % t;
+ produceTestRecord(src2Producer, topic, partition, i);
+ numProduced++;
+ }
+ }
+ for (int t : Set.of(4, 6)) {
+ String topic = "t" + t;
+ for (int i = 0; i < 1000; i++) {
+ int partition = i % t;
+ produceTestRecord(src2Producer, topic, partition, i);
+ numProduced++;
+ }
+ }
+ // Intentionally refrain from producing to t7 to verify that empty partitions are still tracked
+ src2Producer.flush();
+ }
+
+ // Ensure that the same number of records produced to both of the source clusters is available on both of the destination clusters
+ dstCluster1.consume(numProduced, TimeUnit.MINUTES.toMillis(1), "t1", "t2", "t3", "t4", "t5", "t6", "t7");
+ dstCluster2.consume(numProduced, TimeUnit.MINUTES.toMillis(1), "t1", "t2", "t3", "t4", "t5", "t6", "t7");
+
+ // Have to potentially retry here since tasks aren't guaranteed to immediately publish workloads
+ waitForCondition(() -> {
+ KafkaBasedWorkloadProvider.refreshWorkloadCache();
+ Map<ReplicatedTopicPartition, Workload> trackedWorkloads = KafkaBasedWorkloadProvider.getAllWorkloads();
+
+ Set<ReplicatedTopicPartition> expectedReplicatedTopicPartitions = new HashSet<>();
+ for (int i = 1; i <= 7; i++) {
+ if (ConnectorType.MM2.equals(connectorType) && i == 7)
+ continue; // Empty workloads aren't supported by the SMT
+
+ String topic = "t" + i;
+ for (int partition = 0; partition < i; partition++) {
+ ReplicatedTopicPartition tpSrc1Dst1 = new ReplicatedTopicPartition(srcKafka1Name, dstKafka1Name, topic, partition);
+ expectedReplicatedTopicPartitions.add(tpSrc1Dst1);
+ ReplicatedTopicPartition tpSrc1Dst2 = new ReplicatedTopicPartition(srcKafka1Name, dstKafka2Name, topic, partition);
+ expectedReplicatedTopicPartitions.add(tpSrc1Dst2);
+
+ if (i != 2 && i != 5) {
+ ReplicatedTopicPartition tpSrc2Dst1 = new ReplicatedTopicPartition(srcKafka2Name, dstKafka1Name, topic, partition);
+ expectedReplicatedTopicPartitions.add(tpSrc2Dst1);
+ ReplicatedTopicPartition tpSrc2Dst2 = new ReplicatedTopicPartition(srcKafka2Name, dstKafka2Name, topic, partition);
+ expectedReplicatedTopicPartitions.add(tpSrc2Dst2);
+ }
+ }
+ }
+ assertEquals(expectedReplicatedTopicPartitions, trackedWorkloads.keySet());
+
+ trackedWorkloads.forEach((tp, workload) -> {
+ if ("t7".equals(tp.topic())) {
+ assertEquals(0, workload.getMessagesPerSecond());
+ assertEquals(0, workload.getBytesPerSecond());
+ } else {
+ assertNotEquals(0, workload.getMessagesPerSecond(), "empty message rate found for " + tp);
+
+ double expectedBytesPerSecond = workload.getMessagesPerSecond() * 10; // Values are always 7 bytes, keys are always 3
+ double delta = 0.1;
+ assertEquals(expectedBytesPerSecond, workload.getBytesPerSecond(), delta);
+ }
+ });
+
+ Map<ReplicatedTopicPartition, Workload> heavyWorkloads = trackedWorkloads.entrySet().stream()
+ .filter(e -> {
+ String topic = e.getKey().topic();
+ String srcKafka = e.getKey().srcKafka();
+ return switch (topic) {
+ case "t1", "t2", "t3" -> srcKafka1Name.equals(srcKafka);
+ case "t4", "t6" -> srcKafka2Name.equals(srcKafka);
+ default -> false;
+ };
+ })
+ .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+
+ Map<ReplicatedTopicPartition, Workload> lightWorkloads = trackedWorkloads.entrySet().stream()
+ .filter(e -> !heavyWorkloads.containsKey(e.getKey()))
+ .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+
+ heavyWorkloads.forEach((heavyTp, heavyLoad) ->
+ lightWorkloads.forEach((lightTp, lightLoad) -> {
+ assertTrue(heavyLoad.getBytesPerSecond() > lightLoad.getBytesPerSecond(), "workload for " + heavyTp + " should be larger than for " + lightTp);
+ assertTrue(heavyLoad.getMessagesPerSecond() > lightLoad.getMessagesPerSecond(), "workload for " + heavyTp + " should be larger than for " + lightTp);
+ })
+ );
+ return true;
+ }, TimeUnit.SECONDS.toMillis(10), "");
+ }
+
+ private EmbeddedKafkaCluster startExtraKafkaCluster(String name, int numBrokers, Properties brokerProperties) {
+ if (extraKafkaClusters.containsKey(name))
+ throw new IllegalStateException("Cannot register multiple Kafka clusters with the same name");
+
+ EmbeddedKafkaCluster result = new EmbeddedKafkaCluster(numBrokers, brokerProperties);
+ extraKafkaClusters.put(name, result);
+
+ result.start();
+
+ return result;
+ }
+
+ private void startConnector(
+ ConnectorType connectorType,
+ String connectorName,
+ String srcKafkaName,
+ String dstKafkaName,
+ int tasksMax,
+ Set<String> topics,
+ Map<String, String> topicRenames,
+ boolean workloadTracking
+ ) throws InterruptedException {
+ switch (connectorType) {
+ case MM2:
+ if (workloadTracking)
+ startMM2ConnectorWithWorkloadTracking(connectorName, srcKafkaName, dstKafkaName, tasksMax, topics, topicRenames);
+ else
+ startMM2Connector(connectorName, srcKafkaName, dstKafkaName, tasksMax, topics, topicRenames);
+ break;
+ case UREP3:
+ if (workloadTracking)
+ startURep3ConnectorWithWorkloadTracking(connectorName, srcKafkaName, dstKafkaName, tasksMax, topics, topicRenames);
+ else
+ startURep3Connector(connectorName, srcKafkaName, dstKafkaName, tasksMax, topics, topicRenames);
+ break;
+ default:
+ fail("Unrecognized connector type: " + connectorType);
+ }
+
+ connect.assertions().assertConnectorAndExactlyNumTasksAreRunning(connectorName, tasksMax, "Connector and/or tasks did not start in time");
+ }
+
+ private void startMM2ConnectorWithWorkloadTracking(String connectorName, String srcKafkaName, String dstKafkaName, int tasksMax, Set<String> topics, Map<String, String> topicRenames) {
+ String transformName = "TrackWorkloads";
+ String transformPrefix = TRANSFORMS_CONFIG + "." + transformName + ".";
+ startMM2Connector(
+ connectorName,
+ srcKafkaName,
+ dstKafkaName,
+ tasksMax,
+ topics,
+ topicRenames,
+ TRANSFORMS_CONFIG, transformName,
+ transformPrefix + "type", KafkaBasedWorkloadTrackingTransformation.class.getName(),
+ transformPrefix + SRC_CLUSTER_NAME_CONFIG, srcKafkaName,
+ transformPrefix + DST_CLUSTER_NAME_CONFIG, dstKafkaName,
+ transformPrefix + REPORT_INTERVAL_MS_CONFIG, "1000" // Force rapid reporting so that we don't have to wait forever to verify workloads
+ );
+ }
+
+ private void startMM2Connector(String connectorName, String srcKafkaName, String dstKafkaName, int tasksMax, Set<String> topics, Map<String, String> topicRenames, String... extraProps) {
+ assertEquals(0, extraProps.length % 2, "Extra props should be given as a series of key and value pairs");
+ EmbeddedKafkaCluster srcCluster = extraKafkaClusters.get(srcKafkaName);
+ assertNotNull(srcCluster, "No Kafka cluster found with name '" + srcKafkaName + "'");
+
+ EmbeddedKafkaCluster dstCluster = extraKafkaClusters.get(dstKafkaName);
+ assertNotNull(dstCluster, "No Kafka cluster found with name '" + dstKafkaName + "'");
+
+ String topicsString = String.join(",", topics);
+
+ Map<String, String> connectorProps = new HashMap<>();
+ connectorProps.put(CONNECTOR_CLASS_CONFIG, MirrorSourceConnector.class.getName());
+ connectorProps.put(TASKS_MAX_CONFIG, Integer.toString(tasksMax));
+ connectorProps.put(SOURCE_CLUSTER_ALIAS, srcKafkaName);
+ connectorProps.put(TARGET_CLUSTER_ALIAS, dstKafkaName);
+ connectorProps.put("source.cluster." + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, srcCluster.bootstrapServers());
+ connectorProps.put("target.cluster." + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, dstCluster.bootstrapServers());
+ connectorProps.put(CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, dstCluster.bootstrapServers());
+ connectorProps.put(TOPICS, topicsString);
+ connectorProps.put(SYNC_TOPIC_CONFIGS_ENABLED, "false");
+ connectorProps.put(SYNC_TOPIC_ACLS_ENABLED, "false");
+ connectorProps.put(EMIT_OFFSET_SYNCS_ENABLED, "false");
+ connectorProps.put(REPLICATION_POLICY_CLASS, IdentityReplicationPolicy.class.getName());
+ if (topicRenames != null) {
+ Set<String> routingTransforms = new HashSet<>();
+ topicRenames.forEach((srcTopic, dstTopic) -> {
+ String transformName = "Redirect" + srcTopic;
+ routingTransforms.add(transformName);
+
+ String transformPrefix = TRANSFORMS_CONFIG + "." + transformName + ".";
+ connectorProps.put(transformPrefix + "type", RegexRouter.class.getName());
+ connectorProps.put(transformPrefix + "regex", Pattern.quote(srcTopic));
+ connectorProps.put(transformPrefix + "replacement", dstTopic);
+ });
+
+ String routingTransformsString = String.join(",", routingTransforms);
+ connectorProps.compute(
+ TRANSFORMS_CONFIG,
+ (key, value) -> {
+ if (value == null) {
+ return routingTransformsString;
+ } else {
+ return value + "," + routingTransformsString;
+ }
+ }
+ );
+ }
+
+ for (int i = 0; i < extraProps.length; i += 2)
+ connectorProps.put(extraProps[i], extraProps[i + 1]);
+
+ connect.configureConnector(connectorName, connectorProps);
+ }
+
+ private void startURep3ConnectorWithWorkloadTracking(String connectorName, String srcKafkaName, String dstKafkaName, int tasksMax, Set<String> topics, Map<String, String> topicRenames) {
+ startURep3Connector(
+ connectorName,
+ srcKafkaName,
+ dstKafkaName,
+ tasksMax,
+ topics,
+ topicRenames,
+ WORKLOAD_REPORT_INTERVAL_MS_CONFIG, "1000" // Force rapid reporting so that we don't have to wait forever to verify workloads
+ );
+ }
+
+ private void startURep3Connector(String connectorName, String srcKafkaName, String dstKafkaName, int tasksMax, Set<String> topics, Map<String, String> topicRenames, String... extraProps) {
+ assertEquals(0, extraProps.length % 2, "Extra props should be given as a series of key and value pairs");
+ EmbeddedKafkaCluster srcCluster = extraKafkaClusters.get(srcKafkaName);
+ assertNotNull(srcCluster, "No Kafka cluster found with name '" + srcKafkaName + "'");
+
+ EmbeddedKafkaCluster dstCluster = extraKafkaClusters.get(dstKafkaName);
+ assertNotNull(dstCluster, "No Kafka cluster found with name '" + dstKafkaName + "'");
+
+ Map<String, String> connectorProps = baseURep3Config(srcKafkaName, dstKafkaName, tasksMax, topics);
+ connectorProps.put(SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG, srcCluster.bootstrapServers());
+ connectorProps.put(TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG, dstCluster.bootstrapServers());
+ connectorProps.put(CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, dstCluster.bootstrapServers());
+
+ if (topicRenames != null) {
+ topicRenames.forEach((srcTopic, dstTopic) ->
+ connectorProps.put(TOPIC_RENAME_PREFIX + srcTopic, dstTopic)
+ );
+ }
+
+ for (int i = 0; i < extraProps.length; i += 2)
+ connectorProps.put(extraProps[i], extraProps[i + 1]);
+
+ connect.configureConnector(connectorName, connectorProps);
+ }
+
+ private Map<String, String> baseURep3Config(String srcKafkaName, String dstKafkaName, int tasksMax, Set<String> topics) {
+ Map<String, String> result = new HashMap<>();
+ result.put(CONNECTOR_CLASS_CONFIG, UReplicator3Connector.class.getName());
+ result.put(TASKS_MAX_CONFIG, Integer.toString(tasksMax));
+ result.put(SOURCE_CLUSTER_CONFIG, srcKafkaName);
+ result.put(TARGET_CLUSTER_CONFIG, dstKafkaName);
+ result.put(TOPICS_CONFIG, String.join(",", topics));
+ result.put(WORKLOAD_REPORT_INTERVAL_MS_CONFIG, "-1"); // Disable workload publishing by default
+ return result;
+ }
+
+ private void validateURep3Config(Map<String, String> config, String... expectedErrorProps) {
+ Set<String> actualErrorProps = validateConnectorConfig(UReplicator3Connector.class.getCanonicalName(), config).keySet();
+ assertEquals(Set.of(expectedErrorProps), actualErrorProps);
+ }
+
+ private Map<String, List<String>> validateConnectorConfig(String connectorClass, Map<String, String> config) {
+ ConfigInfos validationResult = connect.validateConnectorConfig(connectorClass, config);
+ return validationResult.configs().stream()
+ .map(ConfigInfo::configValue)
+ .filter(configValue -> !configValue.errors().isEmpty())
+ .collect(Collectors.toMap(
+ ConfigValueInfo::name,
+ ConfigValueInfo::errors
+ ));
+ }
+
+ private void produceTestRecord(Producer<byte[], byte[]> producer, String topic, int partition, int iteration) {
+ byte[] keyBytes = "key".getBytes(StandardCharsets.UTF_8);
+ byte[] valueBytes = ("value " + (iteration % 10)).getBytes(StandardCharsets.UTF_8);
+ ProducerRecord<byte[], byte[]> record = new ProducerRecord<>(topic, partition, keyBytes, valueBytes);
+ producer.send(record);
+ }
+}
diff --git a/settings.gradle b/settings.gradle
index 7c37a046838..20cccf47aa5 100644
--- a/settings.gradle
+++ b/settings.gradle
@@ -74,6 +74,7 @@ include 'clients',
'connect:runtime',
'connect:test-plugins',
'connect:transforms',
+ 'connect:uber',
'coordinator-common',
'core',
'examples',