You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@kafka.apache.org by ce...@apache.org on 2026/04/30 18:51:34 UTC
(kafka) branch metamorphomover-prototype updated (6f7ecc729a6 -> 9e14706043b)
This is an automated email from the ASF dual-hosted git repository.
C0urante pushed a change to branch metamorphomover-prototype
in repository https://gitbox.apache.org/repos/asf/kafka.git
discard 6f7ecc729a6 [DRAFT] Add pluggable worker assignor, remove Uber module
new 9e14706043b [DRAFT] Add pluggable worker assignor, remove Uber module
This update added new revisions after undoing existing revisions.
That is to say, some revisions that were in the old version of the
branch are not in the new version. This situation occurs
when a user --force pushes a change and generates a repository
containing something like this:
* -- * -- B -- O -- O -- O (6f7ecc729a6)
\
N -- N -- N refs/heads/metamorphomover-prototype (9e14706043b)
You should already have received notification emails for all of the O
revisions, and so the following emails describe only the N revisions
from the common base, B.
Any revisions marked "omit" are not gone; other references still
refer to them. Any revisions marked "discard" are gone forever.
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
.buildkite/publish-to-maven.sh | 4 +++-
committer-tools/kafka-merge-pr.py | 2 +-
gradle.properties | 2 +-
streams/quickstart/java/pom.xml | 2 +-
.../quickstart/java/src/main/resources/archetype-resources/pom.xml | 2 +-
streams/quickstart/pom.xml | 2 +-
tests/kafkatest/__init__.py | 2 +-
tests/kafkatest/version.py | 2 +-
8 files changed, 10 insertions(+), 8 deletions(-)
(kafka) 01/01: [DRAFT] Add pluggable worker assignor, remove Uber module
Posted by ce...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
C0urante pushed a commit to branch metamorphomover-prototype
in repository https://gitbox.apache.org/repos/asf/kafka.git
commit 9e14706043be25b2caaa4a4e7e58fb17455188da
Author: chrise <ch...@uber.com>
AuthorDate: Wed Apr 29 05:38:03 2026 +0000
[DRAFT] Add pluggable worker assignor, remove Uber module
---
.buildkite/publish-to-maven.sh | 60 +
.buildkite/udj/publish-prod.yaml | 26 +
.buildkite/udj/publish-snapshot.yaml | 26 +
build-prototype.sh | 70 --
build.gradle | 102 +-
committer-tools/kafka-merge-pr.py | 2 +-
.../distributed/ClusterConfigStateImpl.java | 16 +
.../runtime/distributed/DistributedConfig.java | 3 +-
.../runtime/distributed/DistributedHerder.java | 3 +-
.../connect/runtime/distributed/UberAssignor.java | 5 +-
.../runtime/distributed/WorkerCoordinator.java | 3 +-
.../runtime/distributed/WorkerGroupMember.java | 3 +-
.../runtime/health/ConnectClusterStateImpl.java | 1 -
.../runtime/isolation/PluginScanResult.java | 3 +-
.../connect/runtime/isolation/PluginType.java | 2 +-
.../aakafkaoffsetmgmt/OffsetSnapshotClient.java | 217 ----
.../aakafkaoffsetmgmt/OffsetSnapshotReporter.java | 207 ----
.../connect/aakafkaoffsetmgmt/OffsetTracker.java | 32 -
.../aakafkaoffsetmgmt/PartitionOffsetMap.java | 28 -
.../connect/common/ReplicatedTopicPartition.java | 72 --
.../distributed/assignment/AbstractAssignor.java | 375 ------
.../assignment/UberClusterAssignor.java | 195 ---
.../assignment/UberClusterAssignorConfig.java | 98 --
.../distributed/assignment/UberWorkerLoad.java | 70 --
.../assignment/WorkerConnectorAssignor.java | 48 -
.../distributed/assignment/WorkerTaskAssignor.java | 81 --
.../distributed/assignment/WorkloadProvider.java | 39 -
.../connect/rest/UberConnectRestExtension.java | 60 -
.../rest/entities/AssignedPartitionEntity.java | 29 -
.../rest/entities/CachedWorkloadEntity.java | 15 -
.../rest/entities/PartitionWorkloadEntity.java | 59 -
.../entities/ReplicatedTopicPartitionEntity.java | 22 -
.../rest/entities/TaskAssignmentEntity.java | 26 -
.../connect/rest/entities/TaskWorkloadEntity.java | 10 -
.../rest/entities/TopicPartitionEntity.java | 20 -
.../rest/entities/WorkerWorkloadEntity.java | 13 -
.../connect/rest/entities/WorkloadEntity.java | 18 -
.../rest/filter/ContentTypeHeaderFilter.java | 56 -
.../interceptor/ConnectorConfigInterceptor.java | 235 ----
.../interceptor/ConnectorOffsetsInterceptor.java | 147 ---
.../rest/resources/AssignmentsResource.java | 85 --
.../connect/rest/resources/WorkloadsResource.java | 316 -----
.../ureplicator3/TaskPartitionAssignor.java | 112 --
.../kafka/connect/ureplicator3/TaskPartitions.java | 25 -
.../connect/ureplicator3/UReplicator3Config.java | 251 ----
.../ureplicator3/UReplicator3Connector.java | 449 -------
.../ureplicator3/UReplicator3ConnectorConfig.java | 385 ------
.../connect/ureplicator3/UReplicator3Task.java | 367 ------
.../ureplicator3/UReplicator3TaskConfig.java | 42 -
.../coordination/AssignedPartition.java | 27 -
.../coordination/AssignmentListener.java | 11 -
.../ureplicator3/coordination/AssignmentStore.java | 368 ------
.../coordination/AssignmentStoreConfig.java | 145 ---
.../coordination/AssignmentStores.java | 39 -
.../data/kafka/connect/utils/BrokerResolver.java | 39 -
.../data/kafka/connect/utils/ExecutorUtils.java | 88 --
.../kafka/connect/utils/M3ReporterProvider.java | 43 -
.../data/kafka/connect/utils/MetricsUtils.java | 37 -
.../kafka/connect/utils/SerializationUtils.java | 78 --
.../data/kafka/connect/utils/UberConnectUtils.java | 205 ----
.../connect/workload/common/CachedWorkload.java | 28 -
.../kafka/connect/workload/common/Workload.java | 70 --
.../connect/workload/common/WorkloadCache.java | 221 ----
.../connect/workload/common/WorkloadProvider.java | 39 -
.../workload/kafka/KafkaBasedWorkloadProvider.java | 132 ---
.../kafka/KafkaBasedWorkloadProviderConfig.java | 225 ----
.../workload/kafka/KafkaBasedWorkloadStore.java | 286 -----
.../KafkaBasedWorkloadTrackingTransformation.java | 15 -
.../workload/tracking/WorkloadCollector.java | 103 --
.../connect/workload/tracking/WorkloadStore.java | 12 -
.../connect/workload/tracking/WorkloadTracker.java | 132 ---
.../tracking/WorkloadTrackingProcessor.java | 276 -----
.../WorkloadTrackingTransformation.java | 62 -
.../WorkloadTrackingTransformationConfig.java | 62 -
....data.kafka.connect.distributed.ClusterAssignor | 1 -
....apache.kafka.connect.rest.ConnectRestExtension | 1 -
...org.apache.kafka.connect.source.SourceConnector | 1 -
.../integration/ReplicationIntegrationTest.java | 1243 --------------------
.../ureplicator3/TaskPartitionAssignorTest.java | 50 -
.../kafka/connect/utils/BrokerResolverTest.java | 43 -
connect/uber/src/test/resources/log4j2.yaml | 47 -
settings.gradle | 1 -
streams/quickstart/java/pom.xml | 2 +-
.../src/main/resources/archetype-resources/pom.xml | 2 +-
streams/quickstart/pom.xml | 2 +-
tests/kafkatest/__init__.py | 2 +-
tests/kafkatest/version.py | 2 +-
87 files changed, 155 insertions(+), 8514 deletions(-)
diff --git a/.buildkite/publish-to-maven.sh b/.buildkite/publish-to-maven.sh
new file mode 100755
index 00000000000..43a2a29b0e5
--- /dev/null
+++ b/.buildkite/publish-to-maven.sh
@@ -0,0 +1,60 @@
+#! /usr/bin/env bash
+
+set -e
+
+if [[ $# -eq 2 ]]; then
+ MAVEN_PASSWORD="$ARTIFACTS_USSO_TOKEN"
+elif [[ $# -eq 3 ]]; then
+ # If running locally, can use `usso -ussh artifacts.uberinternal.com -print` to generate
+ MAVEN_PASSWORD="$3"
+else
+ echo "Usage: $0 <version> <maven_repo_url> [<maven_repo_password=\$ARTIFACTS_USSO_TOKEN>]"
+ exit 1
+fi
+
+VERSION="$1"
+MAVEN_URL="$2"
+
+./gradlew clean
+./gradlew updateVersion -PnewVersion=$VERSION
+./gradlew build -x test -x rat -x spotbugsMain -x spotbugsTest -x checkstyleMain -x checkstyleTest
+./gradlew releaseTarGz
+./gradlew publish -x test -x rat \
+ -PskipSigning=true \
+ -PmavenUrl="$MAVEN_URL" \
+ -PmavenPassword="$MAVEN_PASSWORD" \
+ -x :examples:publish \
+ -x :trogdor:publish \
+ -x :streams:publish \
+ -x :streams:examples:publish \
+ -x :streams:test-utils:publish \
+ -x :streams:streams-scala:publish \
+ -x :streams:integration-tests:publish \
+ -x :streams:upgrade-system-tests-0110:publish \
+ -x :streams:upgrade-system-tests-10:publish \
+ -x :streams:upgrade-system-tests-11:publish \
+ -x :streams:upgrade-system-tests-20:publish \
+ -x :streams:upgrade-system-tests-21:publish \
+ -x :streams:upgrade-system-tests-22:publish \
+ -x :streams:upgrade-system-tests-23:publish \
+ -x :streams:upgrade-system-tests-24:publish \
+ -x :streams:upgrade-system-tests-25:publish \
+ -x :streams:upgrade-system-tests-26:publish \
+ -x :streams:upgrade-system-tests-27:publish \
+ -x :streams:upgrade-system-tests-28:publish \
+ -x :streams:upgrade-system-tests-30:publish \
+ -x :streams:upgrade-system-tests-31:publish \
+ -x :streams:upgrade-system-tests-32:publish \
+ -x :streams:upgrade-system-tests-33:publish \
+ -x :streams:upgrade-system-tests-34:publish \
+ -x :streams:upgrade-system-tests-35:publish \
+ -x :streams:upgrade-system-tests-36:publish \
+ -x :streams:upgrade-system-tests-37:publish \
+ -x :streams:upgrade-system-tests-38:publish \
+ -x :streams:upgrade-system-tests-39:publish \
+ -x :streams:upgrade-system-tests-40:publish \
+ -x :streams:upgrade-system-tests-41:publish
+
+curl -H "Authorization: Bearer $MAVEN_PASSWORD" \
+ "$MAVEN_URL"/org/apache/kafka/kafka_2.13/$VERSION/kafka_2.13-$VERSION.tgz \
+ -T core/build/distributions/kafka_2.13-$VERSION.tgz
diff --git a/.buildkite/udj/publish-prod.yaml b/.buildkite/udj/publish-prod.yaml
new file mode 100644
index 00000000000..bdd4d929043
--- /dev/null
+++ b/.buildkite/udj/publish-prod.yaml
@@ -0,0 +1,26 @@
+#
+# Copyright (C) 2015-2026 Uber Technologies, Inc. (streaming-data@uber.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+version: 1.0
+
+uown: 5afc78ba-5e4a-43b7-92a5-5c2c040a3c83 # kafka-connect
+
+command: ".buildkite/publish-to-maven.sh \"$VERSION\" https://artifacts.uberinternal.com/artifactory/libs-release-local"
+
+docker:
+ base-image: 127.0.0.1:5055/uber-usi/generic-udj-buildkite-cmd-image:bkt1-produ-1759722358-0fadd
+ shell: true
+ with-host-docker: true
\ No newline at end of file
diff --git a/.buildkite/udj/publish-snapshot.yaml b/.buildkite/udj/publish-snapshot.yaml
new file mode 100644
index 00000000000..e47ef27faf9
--- /dev/null
+++ b/.buildkite/udj/publish-snapshot.yaml
@@ -0,0 +1,26 @@
+#
+# Copyright (C) 2015-2026 Uber Technologies, Inc. (streaming-data@uber.com)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+version: 1.0
+
+uown: 5afc78ba-5e4a-43b7-92a5-5c2c040a3c83 # kafka-connect
+
+command: ".buildkite/publish-to-maven.sh \"$VERSION\" https://artifacts.uberinternal.com/artifactory/libs-snapshot-local"
+
+docker:
+ base-image: 127.0.0.1:5055/uber-usi/generic-udj-buildkite-cmd-image:bkt1-produ-1759722358-0fadd
+ shell: true
+ with-host-docker: true
\ No newline at end of file
diff --git a/build-prototype.sh b/build-prototype.sh
deleted file mode 100755
index 1ab302b0ea8..00000000000
--- a/build-prototype.sh
+++ /dev/null
@@ -1,70 +0,0 @@
-#! /usr/bin/env bash
-
-if [[ "$#" -ne 1 ]]; then
- echo "Usage: $0 <tb_path>"
- exit 1
-fi
-
-TB_PATH="$1"
-
-upload() {
- echo "Uploading $1 to TerraBlob"
- if ! [ -z ${FORCE+x} ]; then
- tb-cli delete "$TB_PATH"/"$1" || true
- fi
- tb-cli put --timeout 90s --multipart "$1" "$TB_PATH"/"$1"
-}
-
-set -e
-
-SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &> /dev/null && pwd)
-
-if tb-cli ls "$TB_PATH" > /dev/null 2> /dev/null; then
- if [ -z ${FORCE+x} ] && [[ $(tb-cli ls "$TB_PATH" | wc -l) != 0 ]]; then
- echo "Path $TB_PATH already exists and is non-empty; cannot use"
- exit 1
- elif ! [ -z ${FORCE+x} ]; then
- echo "Path $TB_PATH is non-empty but FORCE is set; proceeding"
- else
- echo "Path $TB_PATH already exists and is empty; can use safely"
- fi
-else
- tb-cli mkdir -p "$TB_PATH"
-fi
-
-
-if [ -z ${NOBUILD+x} ]; then
-
- if [ -z ${NOCLEAN+x} ]; then
- "$SCRIPT_DIR"/gradlew clean
- else
- echo "Skipping clean"
- fi
-
- "$SCRIPT_DIR"/gradlew :connect:uber:{spotlessApply,build} releaseTarGz -x test
-
-else
- echo "Skipping build"
-fi
-
-if ! [ -z ${NOPUSH+x} ]; then
- echo "Skipping upload"
- exit 0
-fi
-
-
-pushd "$SCRIPT_DIR"/core/build/distributions > /dev/null
-upload kafka_2.13-4.3.0-SNAPSHOT.tgz
-popd > /dev/null
-
-pushd "$SCRIPT_DIR"/connect/uber/build > /dev/null
-TEMP_DIR="$(mktemp -d)"
-mkdir "$TEMP_DIR"/connect-uber-plugins
-cp libs/* dependant-libs/* "$TEMP_DIR"/connect-uber-plugins
-pushd "$TEMP_DIR" > /dev/null
-tar czf connect-uber-plugins.tgz connect-uber-plugins
-upload connect-uber-plugins.tgz
-popd > /dev/null
-popd > /dev/null
-
-echo "Finished uploading artifacts to $TB_PATH"
diff --git a/build.gradle b/build.gradle
index 404e1a0ded6..cbce3d9f110 100644
--- a/build.gradle
+++ b/build.gradle
@@ -375,6 +375,7 @@ subprojects {
repositories {
// To test locally, invoke gradlew with `-PmavenUrl=file:///some/local/path`
maven {
+ allowInsecureProtocol = true
url = mavenUrl
credentials {
username = mavenUsername
@@ -430,6 +431,7 @@ subprojects {
if (shouldSign) {
signing {
+ useGpgCmd()
sign publishing.publications.mavenJava
}
}
@@ -3644,10 +3646,7 @@ project(':connect:runtime') {
implementation libs.classgraph
implementation libs.mavenArtifact
implementation libs.swaggerAnnotations
- // TODO: We should not be hacking the :connect:runtime module with our own dependencies; instead,
- // we should be adding them to the classpath
- // This is also causing embedded integration tests to fail ☹️
- implementation ('com.uber.kafka:kafka-upki-provider:1.0.1-217732-a4c6b554904e-java-monorepo')
+ implementation 'com.uber.kafka:kafka-upki-provider-shaded:1.20.0-chrise-5'
compileOnly libs.bndlib
compileOnly libs.spotbugs
@@ -3677,8 +3676,7 @@ project(':connect:runtime') {
testImplementation libs.junitJupiter
testImplementation libs.mockitoCore
testImplementation libs.mockitoJunitJupiter
- // TODO: Conflicts with Kafka UPKI provider
-// testImplementation libs.httpclient
+ testImplementation libs.httpclient
testImplementation testLog4j2Libs
testCompileOnly libs.bndlib
@@ -3985,92 +3983,6 @@ project(':connect:mirror-client') {
}
}
-project(':connect:uber') {
- base {
- archivesName = "connect-uber"
- }
-
- task genUReplicator3Config(type: JavaExec) {
- classpath sourceSets.main.runtimeClasspath
- classpath sourceSets.main.compileClasspath
- mainClass = 'com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig'
- if( !generatedDocsDir.exists() ) { generatedDocsDir.mkdirs() }
- standardOutput = new File(generatedDocsDir, "ureplicator3_config.html").newOutputStream()
- }
-
- dependencies {
- implementation 'com.google.guava:guava:23.0'
- implementation 'com.uber.m3:tally-core:0.11.1'
- implementation 'com.uber.m3:tally-m3:0.11.1'
- implementation('com.uber.data:streaming-common:1.13.1-230620-a54bfb9c5127-java-monorepo') {
- exclude group: 'jakarta.annotation', module: 'jakarta.annotation-api'
- exclude group: 'jakarta.ws.rs', module: 'jakarta.ws.rs-api'
- exclude group: 'jakarta.servlet', module: 'jakarta.servlet-api'
- exclude group: 'jakarta.validation', module: 'jakarta.validation-api'
- }
- implementation 'com.alibaba:fastjson:1.1.41'
- implementation 'com.google.guava:guava:23.0'
-
- implementation('io.opentracing.contrib:opentracing-apache-httpclient:0.3.3') {
- /*
- opentracing client is using 0.33.0 of opentracing-api and conflict with 0.32 used in javamonorepo, the only way is to downgrade it.
- Monorepo won't ugprade opentracing to 0.33 due to https://t3.uberinternal.com/browse/JAG-2163
- More details: https://github.com/opentracing/opentracing-java/blob/master/CHANGELOG.md#v0330-2019-05-06
- */
- exclude group: 'io.opentracing', module: 'opentracing-api'
- exclude group: 'io.opentracing', module: 'opentracing-util'
- }
- implementation libs.slf4jApi
- implementation libs.swaggerAnnotations
-
- compileOnly libs.jakartaRsApi
- compileOnly libs.jakartaServletApi
- compileOnly libs.jaxAnnotationApi
- compileOnly libs.jerseyContainerServlet
- compileOnly project(':connect:api')
- compileOnly project(':clients')
- compileOnly project(':connect:runtime')
- compileOnly project(':connect:mirror')
-
- testImplementation testLog4j2Libs
- testImplementation libs.jakartaRsApi
- testImplementation libs.junitJupiter
- testImplementation libs.bndlib
- testImplementation libs.mockitoCore
- testImplementation project(':clients').sourceSets.test.output
- testImplementation project(':connect:runtime')
- testImplementation project(':connect:runtime').sourceSets.test.output
- testImplementation project(':connect:mirror')
- testImplementation project(':connect:mirror-client')
- testImplementation project(':core')
- testImplementation project(':test-common:test-common-runtime')
- testImplementation project(':raft')
- testImplementation project(':server')
- testImplementation project(':server-common')
-
- testRuntimeOnly project(':connect:runtime')
- testRuntimeOnly libs.bcpkix
- testRuntimeOnly runtimeTestLibs
- }
-
- javadoc {
- enabled = false
- }
-
- tasks.create(name: "copyDependantLibs", type: Copy) {
- from (configurations.runtimeClasspath) {
- exclude('kafka-clients*')
- exclude('connect-*')
- }
- into "${layout.buildDirectory.get().asFile.path}/dependant-libs"
- duplicatesStrategy 'exclude'
- }
-
- jar {
- dependsOn copyDependantLibs
- }
-}
-
project(':connect:test-plugins') {
base {
archivesName = "connect-test-plugins"
@@ -4153,9 +4065,9 @@ def updateVersionTask = tasks.register('updateVersion') {
def rawVersion
if (project.hasProperty('newVersion')) {
rawVersion = newVersion
- if (!rawVersion.matches(/^[0-9]+\.[0-9]+\.[0-9]+(-SNAPSHOT)?$/)) {
- throw new GradleException("Invalid version format: '${rawVersion}'. Expected format: X.Y.Z or X.Y.Z-SNAPSHOT")
- }
+// if (!rawVersion.matches(/^[0-9]+\.[0-9]+\.[0-9]+(-SNAPSHOT)?$/)) {
+// throw new GradleException("Invalid version format: '${rawVersion}'. Expected format: X.Y.Z or X.Y.Z-SNAPSHOT")
+// }
// Update gradle.properties
def gradlePropsFile = file("${project.rootDir}/gradle.properties")
if (gradlePropsFile.exists()) {
diff --git a/committer-tools/kafka-merge-pr.py b/committer-tools/kafka-merge-pr.py
index b4498ad6afe..dfdd511d6da 100755
--- a/committer-tools/kafka-merge-pr.py
+++ b/committer-tools/kafka-merge-pr.py
@@ -70,7 +70,7 @@ TEMP_BRANCH_PREFIX = "PR_TOOL"
DEV_BRANCH_NAME = "trunk"
-DEFAULT_FIX_VERSION = os.environ.get("DEFAULT_FIX_VERSION", "4.3.0")
+DEFAULT_FIX_VERSION = os.environ.get("DEFAULT_FIX_VERSION", "4.3.0-chrise-33-uber")
ORIGINAL_HEAD = ""
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ClusterConfigStateImpl.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ClusterConfigStateImpl.java
index def0f344d56..9321a113843 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ClusterConfigStateImpl.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/ClusterConfigStateImpl.java
@@ -1,3 +1,19 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
package org.apache.kafka.connect.runtime.distributed;
import com.uber.data.kafka.connect.distributed.ClusterConfigState;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedConfig.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedConfig.java
index 391597f664c..1026581a167 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedConfig.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedConfig.java
@@ -16,7 +16,6 @@
*/
package org.apache.kafka.connect.runtime.distributed;
-import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.clients.MetadataRecoveryStrategy;
import org.apache.kafka.clients.producer.ProducerConfig;
@@ -27,6 +26,8 @@ import org.apache.kafka.common.security.auth.SecurityProtocol;
import org.apache.kafka.common.utils.Utils;
import org.apache.kafka.connect.runtime.WorkerConfig;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
+
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedHerder.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedHerder.java
index 3b85f561e16..381db371aac 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedHerder.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/DistributedHerder.java
@@ -16,7 +16,6 @@
*/
package org.apache.kafka.connect.runtime.distributed;
-import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.common.KafkaFuture;
import org.apache.kafka.common.config.ConfigDef;
@@ -80,6 +79,8 @@ import org.apache.kafka.connect.util.SinkUtils;
import org.apache.kafka.connect.util.Stage;
import org.apache.kafka.connect.util.TemporaryStage;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
+
import org.slf4j.Logger;
import java.util.ArrayList;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberAssignor.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberAssignor.java
index 4a466575ad3..1166199afec 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberAssignor.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/UberAssignor.java
@@ -16,11 +16,12 @@
*/
package org.apache.kafka.connect.runtime.distributed;
+import org.apache.kafka.common.utils.LogContext;
+import org.apache.kafka.connect.storage.ClusterConfigState;
+
import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import com.uber.data.kafka.connect.distributed.ConnectorTaskId;
import com.uber.data.kafka.connect.distributed.ConnectorsAndTasks;
-import org.apache.kafka.common.utils.LogContext;
-import org.apache.kafka.connect.storage.ClusterConfigState;
import org.slf4j.Logger;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinator.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinator.java
index 192f035c6b7..1f98f92c628 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinator.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerCoordinator.java
@@ -16,7 +16,6 @@
*/
package org.apache.kafka.connect.runtime.distributed;
-import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import org.apache.kafka.clients.GroupRebalanceConfig;
import org.apache.kafka.clients.consumer.CloseOptions;
import org.apache.kafka.clients.consumer.internals.AbstractCoordinator;
@@ -31,6 +30,8 @@ import org.apache.kafka.connect.storage.ClusterConfigState;
import org.apache.kafka.connect.storage.ConfigBackingStore;
import org.apache.kafka.connect.util.ConnectorTaskId;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
+
import org.slf4j.Logger;
import java.io.Closeable;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerGroupMember.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerGroupMember.java
index c3dc2acdd3f..f3d0e664110 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerGroupMember.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/distributed/WorkerGroupMember.java
@@ -16,7 +16,6 @@
*/
package org.apache.kafka.connect.runtime.distributed;
-import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import org.apache.kafka.clients.ApiVersions;
import org.apache.kafka.clients.ClientUtils;
import org.apache.kafka.clients.CommonClientConfigs;
@@ -43,6 +42,8 @@ import org.apache.kafka.connect.runtime.WorkerConfig;
import org.apache.kafka.connect.storage.ConfigBackingStore;
import org.apache.kafka.connect.util.ConnectorTaskId;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
+
import org.slf4j.Logger;
import java.net.InetSocketAddress;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImpl.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImpl.java
index 3f2ce96766b..4c82e8f3a57 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImpl.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/health/ConnectClusterStateImpl.java
@@ -29,7 +29,6 @@ import org.apache.kafka.connect.health.TaskState;
import org.apache.kafka.connect.runtime.Herder;
import org.apache.kafka.connect.runtime.distributed.RebalanceNeededException;
import org.apache.kafka.connect.runtime.rest.entities.ConnectorStateInfo;
-import org.apache.kafka.connect.runtime.rest.entities.TaskInfo;
import org.apache.kafka.connect.util.FutureCallback;
import org.slf4j.Logger;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginScanResult.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginScanResult.java
index 3c6a193a4bf..5ade6a6ef2e 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginScanResult.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginScanResult.java
@@ -16,7 +16,6 @@
*/
package org.apache.kafka.connect.runtime.isolation;
-import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import org.apache.kafka.common.config.provider.ConfigProvider;
import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
import org.apache.kafka.connect.rest.ConnectRestExtension;
@@ -27,6 +26,8 @@ import org.apache.kafka.connect.storage.HeaderConverter;
import org.apache.kafka.connect.transforms.Transformation;
import org.apache.kafka.connect.transforms.predicates.Predicate;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
+
import java.util.List;
import java.util.SortedSet;
import java.util.TreeSet;
diff --git a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginType.java b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginType.java
index 05a4766df71..0b686ac68d4 100644
--- a/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginType.java
+++ b/connect/runtime/src/main/java/org/apache/kafka/connect/runtime/isolation/PluginType.java
@@ -16,7 +16,6 @@
*/
package org.apache.kafka.connect.runtime.isolation;
-import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import org.apache.kafka.common.config.provider.ConfigProvider;
import org.apache.kafka.connect.connector.policy.ConnectorClientConfigOverridePolicy;
import org.apache.kafka.connect.rest.ConnectRestExtension;
@@ -28,6 +27,7 @@ import org.apache.kafka.connect.transforms.Transformation;
import org.apache.kafka.connect.transforms.predicates.Predicate;
import com.fasterxml.jackson.annotation.JsonValue;
+import com.uber.data.kafka.connect.distributed.ClusterAssignor;
import java.util.Locale;
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/OffsetSnapshotClient.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/OffsetSnapshotClient.java
deleted file mode 100644
index 0a9b3111a4c..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/OffsetSnapshotClient.java
+++ /dev/null
@@ -1,217 +0,0 @@
-package com.uber.data.kafka.connect.aakafkaoffsetmgmt;
-
-import org.apache.kafka.common.TopicPartition;
-
-import com.alibaba.fastjson.JSON;
-import com.alibaba.fastjson.JSONArray;
-import com.alibaba.fastjson.JSONObject;
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.Charsets;
-
-import org.apache.commons.io.IOUtils;
-import org.apache.commons.lang3.tuple.Pair;
-import org.apache.http.client.ResponseHandler;
-import org.apache.http.client.config.RequestConfig;
-import org.apache.http.client.methods.CloseableHttpResponse;
-import org.apache.http.client.methods.HttpGet;
-import org.apache.http.client.methods.HttpPost;
-import org.apache.http.conn.HttpClientConnectionManager;
-import org.apache.http.entity.ByteArrayEntity;
-import org.apache.http.impl.client.CloseableHttpClient;
-import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.Closeable;
-import java.io.IOException;
-import java.net.SocketTimeoutException;
-import java.util.Map;
-import java.util.Optional;
-
-import io.opentracing.contrib.apache.http.client.TracingHttpClientBuilder;
-import io.opentracing.util.GlobalTracer;
-
-import static java.nio.charset.StandardCharsets.UTF_8;
-
-public class OffsetSnapshotClient implements Closeable {
-
- private static final Logger log = LoggerFactory.getLogger(OffsetSnapshotClient.class);
-
- private static final String getUrlFormatString = "http://127.0.0.1:18668/snapshot/%s/%s/%s/%d/%d?srcdst=src";
- private static final int DEFAULT_MAX_CONNECTION_PER_ROUTE = 8;
- private static final String DEFAULT_UBER_SOURCE_VALUE = "ureplicator2-worker";
- private static final String X_UBER_SOURCE_HEADER = "x-uber-source";
-
- private final String snapshotPostUrl;
- private final CloseableHttpClient httpClient;
- private final RequestConfig requestConfig;
-// protected final Meter snapshotReportCount = new Meter();
-// protected final Meter snapshotReportFailure = new Meter();
-// protected final Timer snapshotReportResponse = new Timer();
-// protected final Meter snapshotFetchCount = new Meter();
-// protected final Meter snapshotFetchFailure = new Meter();
-// protected final Timer snapshotFetchResponse = new Timer();
-
- public OffsetSnapshotClient(String snapshotPostUrl, int requestTimeOutInMs) {
- this(snapshotPostUrl, requestTimeOutInMs, buildDefaultConnectionManager());
- }
-
- private static HttpClientConnectionManager buildDefaultConnectionManager() {
- PoolingHttpClientConnectionManager result = new PoolingHttpClientConnectionManager();
- result.setDefaultMaxPerRoute(DEFAULT_MAX_CONNECTION_PER_ROUTE);
- return result;
- }
-
- @VisibleForTesting
- OffsetSnapshotClient(
- String snapshotPostUrl,
- int requestTimeOutInMs,
- HttpClientConnectionManager connMgrForHttpClient
- ) {
- this.snapshotPostUrl = snapshotPostUrl;
- // requestConfig is immutable. These three timeouts are for
- // 1. getting connection from connection manager;
- // 2. establishing connection with server;
- // 3. getting next data snippet from server.
- this.requestConfig = RequestConfig.custom()
- .setConnectionRequestTimeout(requestTimeOutInMs)
- .setConnectTimeout(requestTimeOutInMs)
- .setSocketTimeout(requestTimeOutInMs)
- .build();
-
- TracingHttpClientBuilder builder = TracingHttpClientBuilder.create()
- .withTracer(GlobalTracer.get());
- this.httpClient = builder.setConnectionManager(connMgrForHttpClient).build();
-// KafkaUReplicatorMetricsReporter.get().registerMetric("snapshot.report-count", snapshotReportCount);
-// KafkaUReplicatorMetricsReporter.get().registerMetric("snapshot.report-failure", snapshotReportFailure);
-// KafkaUReplicatorMetricsReporter.get().registerMetric("snapshot.report-time", snapshotReportResponse);
-// KafkaUReplicatorMetricsReporter.get().registerMetric("snapshot.fetch-count", snapshotFetchCount);
-// KafkaUReplicatorMetricsReporter.get().registerMetric("snapshot.fetch-failure", snapshotFetchFailure);
-// KafkaUReplicatorMetricsReporter.get().registerMetric("snapshot.fetch-time", snapshotFetchResponse);
- }
-
- public Optional<Long> convertOffsetBySourceCluster(
- TopicPartition topicPartition,
- String srcCluster,
- String dstCluster,
- long offsetInSrcCluster
- ) {
-// snapshotFetchCount.mark();
- // Rule is that the last 3 params are all for the same cluster
- String urlString = String.format(getUrlFormatString, topicPartition.topic(),
- srcCluster, dstCluster, topicPartition.partition(), offsetInSrcCluster);
- log.info("Converting offset via url {}", urlString);
- HttpGet httpGet = new HttpGet(urlString);
- httpGet.setConfig(requestConfig);
- httpGet.addHeader(X_UBER_SOURCE_HEADER, DEFAULT_UBER_SOURCE_VALUE);
-// return snapshotFetchResponse.timeSupplier(() -> {
- try {
- CloseableHttpResponse resp = httpClient.execute(httpGet);
- String retVal = IOUtils.toString(resp.getEntity().getContent(), UTF_8);
- resp.close();
- JSONObject parsed = JSON.parseObject(retVal);
- JSONArray offsets = parsed.getJSONArray(String.valueOf(topicPartition.partition()));
- if (offsets != null && !offsets.isEmpty()) {
- Long offset = offsets.getLong(0);
- return Optional.of(offset);
- } else {
- log.warn("Failed to fetch offset from response {}", retVal);
- }
- }catch(IOException ex) {
-// snapshotReportFailure.mark();
- log.error("Failed to GET via {}", urlString, ex);
- }
- return Optional.empty();
-// });
- }
-
- private static boolean isSuccessStatus(final int status) {
- // 2xx indicates success
- return status >= 200 && status < 300;
- }
-
- public void reportSnapshot(
- Map<String, PartitionOffsetMap> topicPartitionOffsetMap,
- String srcClusterName,
- String dstClusterName
- ) {
- long startInMs = System.currentTimeMillis();
- log.info("Start to report offset snapshot at={}", startInMs);
- HttpPost httpPost = createHttpPost(topicPartitionOffsetMap, srcClusterName, dstClusterName);
- if (httpPost == null) {
- log.info("Skipped posting offset because there is no data");
- } else {
-// snapshotReportCount.mark();
-// snapshotReportResponse.time(() -> {
- try {
- int respStatus = httpClient.execute(httpPost, createResponseCodeExtractor());
- if (isSuccessStatus(respStatus)) {
- long endInMs = System.currentTimeMillis();
- log.info("Report offset snapshot done at={} tookMs={}", endInMs, endInMs - startInMs);
- } else {
-// snapshotReportFailure.mark();
- log.info("Got error respStatus={} to post offset snapshot", respStatus);
- }
- } catch (SocketTimeoutException e) {
-// snapshotReportFailure.mark();
- log.warn("Got timeout exception to post offset snapshot, {}", e.getMessage());
- } catch (Throwable t) {
-// snapshotReportFailure.mark();
- log.warn("Got unknown exception to post offset snapshot", t);
- }
-// });
- }
- }
-
- public void markSnapshotReportFailure() {
-// snapshotReportFailure.mark();
- }
-
- private HttpPost createHttpPost(
- Map<String, PartitionOffsetMap> topicPartitionOffsetMap,
- String srcClusterName,
- String dstClusterName
- ) {
- if (topicPartitionOffsetMap.isEmpty()) {
- return null;
- }
- JSONArray postBody = new JSONArray();
- for (Map.Entry<String, PartitionOffsetMap> entry : topicPartitionOffsetMap.entrySet()) {
- String topicName = entry.getKey();
- Map<Pair<Integer, Integer>, Pair<Long, Long>> partitionOffsetMap = entry.getValue().getPartitionOffsetMap();
- for (Map.Entry<Pair<Integer, Integer>, Pair<Long, Long>> mapping : partitionOffsetMap.entrySet()) {
- JSONObject mappingInJson = new JSONObject();
- mappingInJson.put("topicName", topicName);
- mappingInJson.put("srcClusterName", srcClusterName);
- mappingInJson.put("dstClusterName", dstClusterName);
- mappingInJson.put("dstPartitionId", mapping.getKey().getLeft());
- mappingInJson.put("dstOffset", mapping.getValue().getLeft());
-
- JSONObject srcOffsetInJson = new JSONObject();
- srcOffsetInJson.put("" + mapping.getKey().getRight(), mapping.getValue().getRight());
- mappingInJson.put("srcPartitionOffsets", srcOffsetInJson);
- postBody.add(mappingInJson);
- }
- }
-
- // there is no strict synchronization, so few offset mapping might be dropped
- // this is negligible, since we overwrite offset mappings during post interval anyway
- topicPartitionOffsetMap.clear();
-
- HttpPost httpPost = new HttpPost(snapshotPostUrl);
- // since the post interval is pretty long, thus close conn to release resource
- httpPost.addHeader("Connection", "close");
- httpPost.addHeader("Content-Type", "application/json; charset=UTF-8");
- httpPost.setEntity(new ByteArrayEntity(postBody.toString().getBytes(Charsets.UTF_8)));
- httpPost.setConfig(requestConfig);
- return httpPost;
- }
-
- private static ResponseHandler<Integer> createResponseCodeExtractor() {
- return response -> response.getStatusLine().getStatusCode();
- }
-
- public void close() throws IOException {
- httpClient.close();
- }
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/OffsetSnapshotReporter.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/OffsetSnapshotReporter.java
deleted file mode 100644
index 3f69529fb98..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/OffsetSnapshotReporter.java
+++ /dev/null
@@ -1,207 +0,0 @@
-package com.uber.data.kafka.connect.aakafkaoffsetmgmt;
-
-import org.apache.kafka.common.TopicPartition;
-
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.collect.ImmutableMap;
-import com.google.common.util.concurrent.ThreadFactoryBuilder;
-import com.uber.data.kafka.connect.utils.ExecutorUtils;
-import com.uber.jfx.core.UberEnvironment;
-import com.uber.m3.tally.RootScopeBuilder;
-import com.uber.m3.tally.Scope;
-import com.uber.m3.util.Duration;
-
-import org.apache.commons.lang3.tuple.Pair;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.IOException;
-import java.util.Map;
-import java.util.Optional;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.Executors;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.ThreadLocalRandom;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicBoolean;
-
-/**
- * Report offset snapshot to AAKafkaOffsetMgmt backend to persist the offset mapping information
- */
-public class OffsetSnapshotReporter {
-
- private static final Logger log = LoggerFactory.getLogger(OffsetSnapshotReporter.class);
-
- private static final String PELOTON_INSTANCE_ID_KEY = "PELOTON_INSTANCE_ID";
-
- private final String srcClusterName;
- private final String dstClusterName;
- private final Map<String, PartitionOffsetMap> topicPartitionOffsetMap = new ConcurrentHashMap<>();
- private final ScheduledExecutorService reportExecutor;
- private final AtomicBoolean isStarted = new AtomicBoolean(false);
- private final int postIntervalInMs;
- private final Map<TopicPartition, Long> lastCommittedSrcOffset = new ConcurrentHashMap<>();
- private final OffsetTracker offsetTracker;
- private final OffsetSnapshotClient snapshotClient;
- private final Scope scope;
- private final String routeId;
-
- public OffsetSnapshotReporter(
- String srcClusterName,
- String dstClusterName,
- int postIntervalInMs,
- OffsetSnapshotClient snapshotClient,
- String routeId
- ) {
- this(srcClusterName, dstClusterName, postIntervalInMs, snapshotClient, routeId, new OffsetTracker(), buildDefaultScope());
- }
-
- private static Scope buildDefaultScope() {
- return new RootScopeBuilder().reportEvery(Duration.ofSeconds(10));
- }
-
- @VisibleForTesting
- OffsetSnapshotReporter(
- String srcClusterName,
- String dstClusterName,
- int postIntervalInMs,
- OffsetSnapshotClient snapshotClient,
- String routeId,
- OffsetTracker offsetTracker,
- Scope scope
- ) {
- this.srcClusterName = srcClusterName;
- this.dstClusterName = dstClusterName;
- this.postIntervalInMs = postIntervalInMs;
- this.offsetTracker = offsetTracker;
- this.snapshotClient = snapshotClient;
-
- this.reportExecutor = Executors.newSingleThreadScheduledExecutor(
- new ThreadFactoryBuilder().setNameFormat("offset-snapshot-reporter-%d").build());
- Map<String, String> envVars = System.getenv();
-
- @SuppressWarnings("deprecation")
- String deployment = envVars.getOrDefault(UberEnvironment.DEPLOYMENT_KEY, "unknown");
- String pelotonInstanceId = envVars.getOrDefault(PELOTON_INSTANCE_ID_KEY, "unknown");
- this.scope = scope.tagged(ImmutableMap.of(
- "deployment", deployment,
- "peloton_instance_id", pelotonInstanceId));
- this.routeId = routeId;
- }
-
- /**
- * Starts to report offset
- */
- public void start() {
- int delayInMs = ThreadLocalRandom.current().nextInt();
- log.info("Start Offset Snapshot Reporter with srcClusterName={}, dstClusterName={}, "
- + "delayInMs={} and postIntervalInMs={}", srcClusterName, dstClusterName, delayInMs,
- postIntervalInMs);
- reportExecutor.scheduleWithFixedDelay(createReportSnapshotTask(), delayInMs, postIntervalInMs, TimeUnit.MILLISECONDS);
- isStarted.set(true);
- }
-
- public void record(String srcTopic, int srcPartition, long srcOffset, int dstPartition, long dstOffset) {
- PartitionOffsetMap partitionOffsetMap = topicPartitionOffsetMap.computeIfAbsent(srcTopic, k -> new PartitionOffsetMap());
- partitionOffsetMap.put(dstPartition, srcPartition, dstOffset, srcOffset);
- }
-
- void updateTopicPartitionOffset(TopicPartition topicPartition) {
- long lastCommitted = lastCommittedSrcOffset.getOrDefault(topicPartition, -1L);
- long lastOffsetCopied = offsetTracker.getHighestCopied().getOrDefault(topicPartition, -1L);
- long lastSkipped = offsetTracker.getHighestSkipped().getOrDefault(topicPartition, -1L);
- if (lastOffsetCopied == -1 && lastCommitted == -1L) {
- log.info(
- "TopicPartition {} has skipped range of {}->{}, last committed {}. Fetching from AA..",
- topicPartition, lastOffsetCopied, lastSkipped, lastCommitted);
-
- Optional<Long> reversedLookup = fetchSnapshotFromAA(topicPartition, lastSkipped);
- if (reversedLookup.isPresent()) {
- PartitionOffsetMap offsetMap = topicPartitionOffsetMap
- .computeIfAbsent(topicPartition.topic(), k -> new PartitionOffsetMap());
- log.info(
- "Fetched from AA. Snapshot offset {}, {} for {}->{}. Faking snapshot of {}->{} ",
- reversedLookup.get(), topicPartition, dstClusterName, srcClusterName, lastSkipped,
- reversedLookup.get());
- offsetMap.put(topicPartition.partition(), topicPartition.partition(), reversedLookup.get(),
- lastSkipped);
- } else {
- log.info("Failed to fetch snapshot from AA {}", topicPartition);
- }
- }
- }
-
- private Runnable createReportSnapshotTask() {
- return () -> {
- try{
- reportSnapshot();
- scope.tagged(ImmutableMap.of("routeid", routeId))
- .counter("computemapping")
- .inc(1);
- } catch (Exception e) {
- log.error("Failed to report snapshot", e);
- snapshotClient.markSnapshotReportFailure();
- }
- clearAfterReport();
- };
- }
-
- @VisibleForTesting
- void reportSnapshot() {
- mergeToLastCommitedMap();
- offsetTracker.getHighestSkipped().keySet().forEach(this::updateTopicPartitionOffset);
- snapshotClient.reportSnapshot(topicPartitionOffsetMap, srcClusterName, dstClusterName);
- }
-
- @VisibleForTesting
- void clearAfterReport() {
- lastCommittedSrcOffset.clear();
- topicPartitionOffsetMap.clear();
- offsetTracker.clear();
- }
-
- @VisibleForTesting
- Optional<Long> fetchSnapshotFromAA(TopicPartition topicPartition, long lastSkipped) {
- return snapshotClient.convertOffsetBySourceCluster(topicPartition, dstClusterName, srcClusterName, lastSkipped);
- }
-
- @VisibleForTesting
- void mergeToLastCommitedMap() {
- topicPartitionOffsetMap.keySet().forEach( topic -> {
- PartitionOffsetMap offsetMap = topicPartitionOffsetMap.get(topic);
- Map<Pair<Integer, Integer>, Pair<Long, Long>> partitionOffsetMap = offsetMap.getPartitionOffsetMap();
- for (Map.Entry<Pair<Integer, Integer>, Pair<Long, Long>> mapping : partitionOffsetMap.entrySet()) {
- int paritionId = mapping.getKey().getRight();
- long offset = mapping.getValue().getRight();
- TopicPartition topicPartition = new TopicPartition(topic, paritionId);
- lastCommittedSrcOffset.put(topicPartition, offset);
- }
- });
- }
-
- @VisibleForTesting
- Optional<Long> getLastCommittedOffset(TopicPartition tp) {
- return Optional.ofNullable(lastCommittedSrcOffset.get(tp));
- }
-
- /**
- * Shuts down the reporter
- */
- public void shutdown() {
- log.info("Shutdown Offset Snapshot Reporter");
- if (isStarted.compareAndSet(true, false)) {
- ExecutorUtils.shutdown(reportExecutor, "offset snapshot reporter thread", 1, TimeUnit.SECONDS);
- try {
- snapshotClient.close();
- } catch (IOException ex) {
- log.error("Failed to close snapshot client");
- }
- topicPartitionOffsetMap.clear();
- }
- }
-
- public boolean isStarted() {
- return isStarted.get();
- }
-
-}
\ No newline at end of file
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/OffsetTracker.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/OffsetTracker.java
deleted file mode 100644
index aa25a1d562d..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/OffsetTracker.java
+++ /dev/null
@@ -1,32 +0,0 @@
-package com.uber.data.kafka.connect.aakafkaoffsetmgmt;
-
-import org.apache.kafka.common.TopicPartition;
-
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
-
-public class OffsetTracker {
- private final Map<TopicPartition, Long> highestOffsetCopied = new ConcurrentHashMap<>();
- private final Map<TopicPartition, Long> highestOffsetSkipped = new ConcurrentHashMap<>();
-
- public void trackSkipped(TopicPartition tp, long offset) {
- highestOffsetSkipped.put(tp, offset);
- }
-
- public void trackCopied(TopicPartition tp, long offset) {
- highestOffsetCopied.put(tp, offset);
- }
-
- public Map<TopicPartition, Long> getHighestSkipped() {
- return highestOffsetSkipped;
- }
-
- public Map<TopicPartition, Long> getHighestCopied() {
- return highestOffsetCopied;
- }
-
- public void clear() {
- highestOffsetCopied.clear();
- highestOffsetSkipped.clear();
- }
-}
\ No newline at end of file
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/PartitionOffsetMap.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/PartitionOffsetMap.java
deleted file mode 100644
index 674d4ab0024..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/aakafkaoffsetmgmt/PartitionOffsetMap.java
+++ /dev/null
@@ -1,28 +0,0 @@
-package com.uber.data.kafka.connect.aakafkaoffsetmgmt;
-
-import com.google.common.collect.ImmutableMap;
-
-import org.apache.commons.lang3.tuple.Pair;
-
-import java.util.Map;
-import java.util.concurrent.ConcurrentHashMap;
-
-// helper class to keep track of offset mapping for topic partitions
-public class PartitionOffsetMap {
-
- // first Pair is <dst partition, src partition>, second Pair is <dst offset, source offset>
- private final Map<Pair<Integer, Integer>, Pair<Long, Long>> partitionOffsetMap = new ConcurrentHashMap<>();
-
- public void put(int dstPartition, int srcPartition, long dstOffset, long srcOffset) {
- partitionOffsetMap.put(Pair.of(dstPartition, srcPartition), Pair.of(dstOffset, srcOffset));
- }
-
- public void removePartition(int partition) {
- partitionOffsetMap.remove(Pair.of(partition, partition));
- }
-
- public Map<Pair<Integer, Integer>, Pair<Long, Long>> getPartitionOffsetMap() {
- return ImmutableMap.copyOf(partitionOffsetMap);
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/common/ReplicatedTopicPartition.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/common/ReplicatedTopicPartition.java
deleted file mode 100644
index 5913296a216..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/common/ReplicatedTopicPartition.java
+++ /dev/null
@@ -1,72 +0,0 @@
-package com.uber.data.kafka.connect.common;
-
-import org.apache.kafka.common.TopicPartition;
-
-import java.util.Objects;
-
-/**
- * POJO containing a topic partition and a source and destination Kafka cluster
- */
-public class ReplicatedTopicPartition {
-
- private final String srcKafka;
- private final String dstKafka;
- private final String topic;
- private final int partition;
-
- public ReplicatedTopicPartition(String srcKafka, String dstKafka, String topic, int partition) {
- this.srcKafka = srcKafka;
- this.dstKafka = dstKafka;
- this.topic = topic;
- this.partition = partition;
- }
-
- public ReplicatedTopicPartition(String srcKafka, String dstKafka, TopicPartition topicPartition) {
- this(srcKafka, dstKafka, topicPartition.topic(), topicPartition.partition());
- }
-
- public String srcKafka() {
- return srcKafka;
- }
-
- public String dstKafka() {
- return dstKafka;
- }
-
- public String topic() {
- return topic;
- }
-
- public int partition() {
- return partition;
- }
-
- public TopicPartition topicPartition() {
- return new TopicPartition(topic, partition);
- }
-
- public boolean matchesPipeline(String srcKafka, String dstKafka) {
- return Objects.equals(srcKafka, this.srcKafka)
- && Objects.equals(dstKafka, this.dstKafka);
- }
-
- @Override
- public boolean equals(Object o) {
- if (o == null || getClass() != o.getClass()) return false;
- ReplicatedTopicPartition that = (ReplicatedTopicPartition) o;
- return partition == that.partition
- && Objects.equals(srcKafka, that.srcKafka)
- && Objects.equals(dstKafka, that.dstKafka)
- && Objects.equals(topic, that.topic);
- }
-
- @Override
- public int hashCode() {
- return Objects.hash(srcKafka, dstKafka, topic, partition);
- }
-
- @Override
- public String toString() {
- return "@" + srcKafka + "@" + dstKafka + "@" + topicPartition();
- }
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/AbstractAssignor.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/AbstractAssignor.java
deleted file mode 100644
index 134c89ee758..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/AbstractAssignor.java
+++ /dev/null
@@ -1,375 +0,0 @@
-package com.uber.data.kafka.connect.distributed.assignment;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Comparator;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Locale;
-import java.util.Map;
-import java.util.PriorityQueue;
-import java.util.Set;
-import java.util.function.Predicate;
-
-// Hey just FYI this class is absolutely not thread-safe and you're deranged if you try to use it that way
-public abstract class AbstractAssignor<G, J> {
-
- private static final Logger log = LoggerFactory.getLogger(AbstractAssignor.class);
-
- private final Comparator<G> groupComparator;
- private final Comparator<J> jobComparator;
- private final double groupLoadThreshold;
- private final double defaultJobLoad;
- private final int groupJobLimit;
- private final Map<G, Double> cachedGroupLoads;
- private final Map<J, Double> cachedJobLoads;
-
- @SuppressWarnings("this-escape")
- protected AbstractAssignor(double groupLoadThreshold, double defaultJobLoad, int groupJobLimit) {
- this.groupComparator = Comparator.comparing(this::cachedGroupLoad);
- this.jobComparator = Comparator.comparing(this::cachedJobLoad);
- this.groupLoadThreshold = groupLoadThreshold;
- this.defaultJobLoad = defaultJobLoad;
- this.groupJobLimit = groupJobLimit;
- this.cachedGroupLoads = new HashMap<>();
- this.cachedJobLoads = new HashMap<>();
- }
-
- protected abstract Double jobLoad(J job) throws Exception;
-
- protected abstract void assign(G group, J job);
- protected abstract void revoke(G group, J job);
- protected abstract Collection<J> jobs(G group);
-
- protected abstract String groupType();
- protected abstract String jobType();
-
- // Can be overridden by subclasses if the lawless chaos of the English language demands it
- protected String groupTypePlural() {
- return groupType() + "(s)";
- }
-
- // Can be overridden by subclasses if the lawless chaos of the English language demands it
- protected String jobTypePlural() {
- return jobType() + "(s)";
- }
-
- protected String jobName(J job) {
- return job.toString();
- }
-
- protected String groupName(G group) {
- return group.toString();
- }
-
- // toAdd contains only new jobs that aren't already assigned and which should
- // be assigned to a group if possible
- public void assign(Collection<G> groups, Collection<J> toAdd) {
- log.info("Rebalancing {} {} with {} new {} to assign", groups.size(), groupTypePlural(), toAdd.size(), jobTypePlural());
-
- cachedGroupLoads.clear();
- cachedJobLoads.clear();
-
- Collection<J> toAssign = new HashSet<>(toAdd);
- Collection<J> removed = enforceGroupJobLimit(groups);
- toAssign.addAll(removed);
-
- assignAll(groups, toAssign);
-
- applyLoadThreshold(groups);
-
- log.info("Finished rebalancing {} {} with {} new {} to assign", groups.size(), groupTypePlural(), toAdd.size(), jobTypePlural());
- }
-
- private Collection<J> enforceGroupJobLimit(Collection<G> groups) {
- if (groupJobLimit <= 0) {
- log.debug("Skipping limit enforcement for {} per {} since limit is non-positive", jobTypePlural(), groupType());
- return Set.of();
- }
-
- log.debug("Enforcing limit for {} per {} across {} {}", jobTypePlural(), groupType(), groups.size(), groupTypePlural());
-
- Collection<J> result = new ArrayList<>();
- for (G group : groups) {
- List<J> jobs = new ArrayList<>(jobs(group));
- int numJobs = jobs.size();
- int numRevoked = numJobs - groupJobLimit;
- if (numRevoked <= 0) {
- log.trace("{} {} has {} {}, which is within the limit of {}",
- capitalize(groupType()), groupName(group), numJobs, jobTypePlural(), groupJobLimit
- );
- continue;
- }
-
- log.debug("{} {} has {} {}, which is above the limit of {}; will revoke {} {}",
- capitalize(groupType()), groupName(group), numJobs, jobTypePlural(), groupJobLimit, numRevoked, jobTypePlural()
- );
- // Revoke smallest jobs in group
- jobs.sort(jobComparator);
- for (int i = 0; i < numRevoked; i++) {
- J revoked = jobs.get(i);
- log.trace("Revoking {} {} from {} {}", jobType(), jobName(revoked), groupType(), groupName(group));
- revokeAndUpdateLoad(group, revoked);
- result.add(revoked);
- }
- }
-
- log.debug("Finished enforcing limit for {} per {} across {} {}", jobTypePlural(), groupType(), groups.size(), groupTypePlural());
-
- return result;
- }
-
- private void assignAll(Collection<G> groups, Collection<J> toAssign) {
- assert !groups.isEmpty();
-
- if (toAssign.isEmpty()) {
- log.debug("No new {} to assign", jobTypePlural());
- return;
- }
-
- log.debug("Assigning {} new {} across {} {}", toAssign.size(), jobTypePlural(), groups.size(), groupTypePlural());
-
- // Assign largest jobs first
- List<J> sortedJobs = new ArrayList<>(toAssign);
- sortedJobs.sort(jobComparator);
-
- // Attempt to assign to least-loaded groups first, but only to groups
- // that are under the job limit
- PriorityQueue<G> eligibleGroups = new PriorityQueue<>(groupComparator);
- groups.forEach(group -> {
- if (beneathJobLimit(group)) {
- eligibleGroups.add(group);
- }
- });
-
- int assigned = 0;
- int assignedPercent = 0;
-
- for (J job : sortedJobs) {
- G group = eligibleGroups.poll();
-
- if (group == null) {
- if (groupJobLimit > 0) {
- // TODO: Metrics
- log.warn("No {} found to accept {} {}; all {} may be above {} limit of {}",
- groupTypePlural(), jobType(), jobName(job), groupTypePlural(), jobType(), groupJobLimit
- );
- } else {
- log.warn(
- "No {} found to accept {} {}; this should never happen",
- groupTypePlural(), jobType(), jobName(job)
- );
- assert false; // Fail if we're in a testing environment
- }
- continue;
- }
-
- log.trace("Assigning {} {} to {} {}", jobType(), jobName(job), groupType(), groupName(group));
- assignAndUpdateLoad(group, job);
-
- if (beneathJobLimit(group)) {
- boolean added = eligibleGroups.add(group);
-
- if (!added) {
- log.warn(
- "Failed to re-insert {} {} into queue after assigning {} {}; this should never happen",
- groupType(), groupName(group), jobType(), jobName(job)
- );
- assert false; // Fail if we're in a testing environment
- }
- } else {
- log.debug("{} {} has now reached limit of {}; will not assign any more {} to it",
- capitalize(groupType()), groupName(group), groupJobLimit, jobTypePlural()
- );
- }
-
- int newAssignedPercent = ++assigned * 100 / toAssign.size();
- if (newAssignedPercent != assignedPercent) {
- assignedPercent = newAssignedPercent;
- log.trace("Assigned {}% of total {} ({} / {})", assignedPercent, jobTypePlural(), assigned, toAssign.size());
- }
- }
-
- log.debug("Finished assigning {} new {} across {} {}", toAssign.size(), jobTypePlural(), groups.size(), groupTypePlural());
- }
-
- // TODO: Dynamic load threshold based on ratio of least-loaded group to others?
- // For example, if loads are 3, 5, 6, 7, 8, and 9, and the threshold is 2, we might
- // still choose to do some balancing even though all groups are above the threshold;
- // the heuristic we could use is to try to balance when groups have 1.5x, 2x, 3x, 5x, etc.
- // the lowest load in the group; if the factor were 2x, we'd try to balance out the groups
- // with loads 6, 7, 8, and 9; if the factor were 2.5x, we'd try to balance out the groups with
- // loads 8 and 9; if it were 3x, we'd try to balance out the group with load 9, and if it were
- // higher, we wouldn't try to do any balancing
-
- private void applyLoadThreshold(Collection<G> groups) {
- assert !groups.isEmpty();
-
- log.debug("Applying load threshold across {} {}", groups.size(), groupTypePlural());
-
- /*
- - Load-balancing reassignment process:
- - Sort groups by load
- - Iterating from most-loaded to least-loaded group:
- - While group is overloaded:
- - Iterating from least-heavy to most-heavy job in group:
- - Attempt to reassign job to least-loaded group that can accept job
- without exceeding threshold
- - If job cannot be reassigned without exceeding threshold, stop trying
- to take load off of group (and move on to next overloaded group,
- if there is one)
- */
- List<G> overloadedGroups = groups.stream()
- .filter(this::isOverloaded)
- .sorted(groupComparator.reversed())
- .toList();
-
- int numOverloaded = overloadedGroups.size();
- log.debug("Found {} overloaded {}", numOverloaded, groupTypePlural());
-
- if (overloadedGroups.isEmpty())
- return;
-
- PriorityQueue<G> eligibleGroups = new PriorityQueue<>(groupComparator);
- groups.stream()
- .filter(this::beneathJobLimit)
- .filter(Predicate.not(this::isOverloaded))
- .forEach(eligibleGroups::add);
-
- log.debug("Found {} eligible {} to receive new {} for load balancing",
- eligibleGroups.size(), groupTypePlural(), jobTypePlural()
- );
-
- for (G overloaded : overloadedGroups) {
- List<J> sortedJobs = new ArrayList<>(jobs(overloaded));
- sortedJobs.sort(jobComparator);
-
- for (J job : sortedJobs) {
- G leastLoaded = eligibleGroups.peek();
- if (leastLoaded == null) {
- return; // Nothing else we can do, every group is overloaded
- }
-
- if (!canAccept(leastLoaded, job)) {
- log.trace("Cannot reassign {} {} to {} {} (would become overloaded)",
- jobType(), jobName(job), groupType(), groupName(leastLoaded));
- // Least-heavy job in this group cannot be assigned to the least-loaded group;
- // nothing left we can do
- break;
- }
-
- log.trace("Reassigning {} {} from {} {} to {} {}",
- jobType(), jobName(job), groupType(), groupName(overloaded), groupType(), groupName(leastLoaded));
- // Reassign from the overloaded group to the non-overloaded one
- revokeAndUpdateLoad(overloaded, job);
- assignAndUpdateLoad(leastLoaded, job);
-
- if (beneathJobLimit(leastLoaded)) {
- // Re-insert into the queue to reflect new workload
- eligibleGroups.poll();
- eligibleGroups.add(leastLoaded);
- } else {
- log.debug("{} {} is at or above {} limit of {}; will not assign any more {} to it",
- capitalize(groupType()), groupName(leastLoaded), jobType(), groupJobLimit, jobTypePlural()
- );
- }
-
- // Check to see if the group is still overloaded
- if (!isOverloaded(overloaded)) {
- numOverloaded--;
- log.debug("{} {} is no longer overloaded", capitalize(groupType()), groupName(overloaded));
- // If the group is no longer overloaded, we can stop trying to transfer
- // jobs off of it, and add it to the non-overloaded queue
- // so that it can potentially receive new jobs from other overloaded groups
- if (beneathJobLimit(overloaded)) {
- eligibleGroups.add(overloaded);
- } else {
- log.debug("{} {} is at or above {} limit of {}; will not assign any more {} to it",
- capitalize(groupType()), groupName(overloaded), jobType(), groupJobLimit, jobTypePlural()
- );
- }
- break;
- } else {
- log.trace("{} {} is still overloaded", capitalize(groupType()), groupName(overloaded));
- }
- }
- }
-
- log.debug("Finished applying load threshold across {} {}; there are {} overloaded {} remaining",
- groups.size(), groupTypePlural(), numOverloaded, groupTypePlural()
- );
- }
-
- private boolean canAccept(G group, J job) {
- double newLoad = cachedGroupLoad(group) + cachedJobLoad(job);
- return !isOverloaded(newLoad);
- }
-
- private boolean isOverloaded(G group) {
- return isOverloaded(cachedGroupLoad(group));
- }
-
- private double cachedGroupLoad(G group) {
- return cachedGroupLoads.computeIfAbsent(group, g ->
- jobs(g).stream().mapToDouble(this::cachedJobLoad).sum()
- );
- }
-
- private boolean isOverloaded(double load) {
- return load > groupLoadThreshold;
- }
-
- private double cachedJobLoad(J job) {
- return cachedJobLoads.computeIfAbsent(job, j -> {
- Double load;
- try {
- load = jobLoad(job);
- } catch (Exception e) {
- log.error("Failed to retrieve load for {} {}; falling back on default load of {}", jobType(), jobName(job), defaultJobLoad);
- return defaultJobLoad;
- }
-
- if (load == null) {
- log.debug("No load available for {} {}; falling back on default load of {}", jobType(), jobName(job), defaultJobLoad);
- return defaultJobLoad;
- }
-
- if (load < 0) {
- log.warn("Negative load provided for {} {}; will fall back on default load of {}", jobType(), jobName(job), defaultJobLoad);
- return defaultJobLoad;
- }
-
- log.trace("Found load of {} for {} {}", load, jobType(), jobName(job));
-
- return load;
- });
- }
-
- private void revokeAndUpdateLoad(G group, J job) {
- revoke(group, job);
- double newLoad = cachedGroupLoad(group) - cachedJobLoad(job);
- cachedGroupLoads.put(group, newLoad);
- }
-
- private void assignAndUpdateLoad(G group, J job) {
- assign(group, job);
- double newLoad = cachedGroupLoad(group) + cachedJobLoad(job);
- cachedGroupLoads.put(group, newLoad);
- }
-
- private int numJobs(G group) {
- return jobs(group).size();
- }
-
- private boolean beneathJobLimit(G group) {
- return groupJobLimit <= 0 || numJobs(group) < groupJobLimit;
- }
-
- private static String capitalize(String s) {
- return (s.substring(0, 1).toUpperCase(Locale.ROOT)) + s.substring(1);
- }
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/UberClusterAssignor.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/UberClusterAssignor.java
deleted file mode 100644
index ed99b72b646..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/UberClusterAssignor.java
+++ /dev/null
@@ -1,195 +0,0 @@
-package com.uber.data.kafka.connect.distributed.assignment;
-
-import org.apache.kafka.common.utils.AppInfoParser;
-import org.apache.kafka.connect.runtime.distributed.UberWorkerLoad;
-import org.apache.kafka.connect.util.ConnectUtils;
-
-import com.uber.data.kafka.connect.distributed.ClusterAssignment;
-import com.uber.data.kafka.connect.distributed.ClusterAssignor;
-import com.uber.data.kafka.connect.distributed.ClusterConfigState;
-import com.uber.data.kafka.connect.distributed.ConnectorTaskId;
-import com.uber.data.kafka.connect.distributed.ConnectorsAndTasks;
-import com.uber.data.kafka.connect.workload.common.WorkloadProvider;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Map.Entry;
-import java.util.Set;
-import java.util.TreeSet;
-import java.util.function.Function;
-import java.util.stream.Collector;
-import java.util.stream.Collectors;
-
-import static org.apache.kafka.common.utils.Utils.closeQuietly;
-
-public class UberClusterAssignor implements ClusterAssignor {
-
- private static final Logger log = LoggerFactory.getLogger(UberClusterAssignor.class);
-
- private WorkloadProvider workloadProvider;
- private double workerLoadThreshold;
- private double defaultTaskLoad;
- private int periodicRebalanceIntervalMs;
- private int workerTaskLimit;
-
- @Override
- public void configure(Map<String, ?> props) {
- UberClusterAssignorConfig config = new UberClusterAssignorConfig(props);
- this.workloadProvider = config.workloadProvider();
- this.workerLoadThreshold = config.uberWorkerLoadThreshold();
- this.defaultTaskLoad = config.uberDefaultTaskLoad();
- this.periodicRebalanceIntervalMs = config.uberPeriodicRebalanceIntervalMs();
- this.workerTaskLimit = 10; // TODO: Don't even bother making this configurable, just put it into Flipr with a sane default
- }
-
- @Override
- public String version() {
- return AppInfoParser.getVersion();
- }
-
- @Override
- public void close() throws Exception {
- closeQuietly(workloadProvider, "Workload provider for Uber cluster assignor");
- }
-
- @Override
- public ClusterAssignment assign(ClusterConfigState clusterConfigState, Map<String, ConnectorsAndTasks> currentAssignment) {
- Set<String> configuredConnectors = new TreeSet<>(clusterConfigState.connectors());
- Set<ConnectorTaskId> configuredTasks = combineCollections(configuredConnectors, clusterConfigState::tasks, Collectors.toSet());
-
- // The connectors and tasks that should be running on the cluster, based on the contents of the
- // config topic (which is the source of truth for the cluster's intended connectors and tasks)
- ConnectorsAndTasks configured = ConnectorsAndTasks.of(configuredConnectors, configuredTasks);
-
- // The connectors and tasks currently running on each member that's participating in this rebalance
- ConnectorsAndTasks activeAssignments = assignment(currentAssignment);
- log.debug("Active assignments: {}", activeAssignments);
-
- // The complete set of connectors and tasks that should be newly-assigned during this round,
- // either because they were just created, or the worker they were previously assigned to
- // is no longer running
- ConnectorsAndTasks toAssign = diff(configured, activeAssignments);
- log.debug("Unassigned: {}", toAssign);
-
- // The set of deleted connectors and tasks that members of the cluster are currently running
- // and which should now be explicitly revoked
- ConnectorsAndTasks deleted = diff(activeAssignments, configured);
- log.debug("Deleted assignments: {}", deleted);
-
- // The connectors and tasks that are currently running on more than one worker each
- ConnectorsAndTasks duplicated = duplicatedAssignments(currentAssignment);
- log.debug("Duplicated assignments: {}", duplicated);
-
- // The connectors and tasks that need to be completely revoked from all workers
- // (A follow-up rebalance is automatically triggered whenever a revocation takes place,
- // during which the duplicated connectors and tasks revoked in this round can be safely
- // reassigned)
- ConnectorsAndTasks toRemove = deleted.toBuilder()
- .add(duplicated)
- .build();
-
- Collection<UberWorkerLoad> workerLoads = workerLoads(currentAssignment);
-
- revokeAll(workerLoads, toRemove);
-
- int connectorsPerWorker = (int) Math.ceil(configured.connectors().size() / (double) workerLoads.size());
- AbstractAssignor<UberWorkerLoad, String> connectorAssignor = new WorkerConnectorAssignor(connectorsPerWorker);
- connectorAssignor.assign(workerLoads, toAssign.connectors());
-
- AbstractAssignor<UberWorkerLoad, ConnectorTaskId> taskAssignor =
- new WorkerTaskAssignor(workloadProvider, clusterConfigState, workerLoadThreshold, defaultTaskLoad, workerTaskLimit);
- taskAssignor.assign(workerLoads, toAssign.tasks());
-
- Map<String, ConnectorsAndTasks.Builder> resultBuilder = new HashMap<>();
- workerLoads.forEach(worker ->
- resultBuilder.computeIfAbsent(worker.worker(), w -> ConnectorsAndTasks.builder())
- .addConnectors(worker.connectors())
- .addTasks(worker.tasks())
- );
-
- Map<String, ConnectorsAndTasks> workerAssignments = transformValues(resultBuilder, ConnectorsAndTasks.Builder::build);
- return new ClusterAssignment(workerAssignments, periodicRebalanceIntervalMs);
- }
-
- private ConnectorsAndTasks assignment(Map<String, ConnectorsAndTasks> memberAssignments) {
- log.debug("Received assignments: {}", memberAssignments);
- return ConnectorsAndTasks.of(
- ConnectUtils.combineCollections(memberAssignments.values(), ConnectorsAndTasks::connectors),
- ConnectUtils.combineCollections(memberAssignments.values(), ConnectorsAndTasks::tasks)
- );
- }
-
- private ConnectorsAndTasks duplicatedAssignments(Map<String, ConnectorsAndTasks> memberAssignments) {
- Set<String> duplicatedConnectors = duplicatedElements(memberAssignments, ConnectorsAndTasks::connectors);
- Set<ConnectorTaskId> duplicatedTasks = duplicatedElements(memberAssignments, ConnectorsAndTasks::tasks);
-
- return ConnectorsAndTasks.of(duplicatedConnectors, duplicatedTasks);
- }
-
- private static <E> Set<E> duplicatedElements(
- Map<String, ConnectorsAndTasks> memberAssignments,
- Function<ConnectorsAndTasks, Collection<E>> assignmentExtractor
- ) {
- Map<E, Long> elementCounts = combineCollections(
- memberAssignments.values(),
- assignmentExtractor,
- Collectors.groupingBy(Function.identity(), Collectors.counting())
- );
- return elementCounts
- .entrySet().stream()
- .filter(entry -> entry.getValue() > 1L)
- .map(Entry::getKey)
- .collect(Collectors.toSet());
- }
-
- /**
- * Revoke all of the specified connectors and tasks from all workers
- * @param workerLoads the workers; may be empty, but not null
- * @param toRevoke the connectors and tasks to revoke; may be empty, but not null
- */
- private void revokeAll(Collection<UberWorkerLoad> workerLoads, ConnectorsAndTasks toRevoke) {
- workerLoads.forEach(w -> {
- w.connectors().removeAll(toRevoke.connectors());
- w.tasks().removeAll(toRevoke.tasks());
- });
- }
-
- private static List<UberWorkerLoad> workerLoads(Map<String, ConnectorsAndTasks> memberAssignments) {
- return memberAssignments.entrySet().stream()
- .map(e -> new UberWorkerLoad(e.getKey(), e.getValue()))
- .collect(Collectors.toList());
- }
-
- private static <K, I, O> Map<K, O> transformValues(Map<K, I> map, Function<I, O> transformation) {
- return map.entrySet().stream().collect(Collectors.toMap(
- Map.Entry::getKey,
- transformation.compose(Map.Entry::getValue)
- ));
- }
-
- private static ConnectorsAndTasks diff(ConnectorsAndTasks base, ConnectorsAndTasks... toSubtract) {
- Set<String> connectors = new TreeSet<>(base.connectors());
- Set<ConnectorTaskId> tasks = new TreeSet<>(base.tasks());
- for (ConnectorsAndTasks sub : toSubtract) {
- connectors.removeAll(sub.connectors());
- tasks.removeAll(sub.tasks());
- }
- return new ConnectorsAndTasks(connectors, tasks);
- }
-
- private static <I, T, C> C combineCollections(
- Collection<I> collection,
- Function<I, Collection<T>> extractCollection,
- Collector<T, ?, C> collector
- ) {
- return collection.stream()
- .map(extractCollection)
- .flatMap(Collection::stream)
- .collect(collector);
- }
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/UberClusterAssignorConfig.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/UberClusterAssignorConfig.java
deleted file mode 100644
index feecfde4dfe..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/UberClusterAssignorConfig.java
+++ /dev/null
@@ -1,98 +0,0 @@
-package com.uber.data.kafka.connect.distributed.assignment;
-
-import org.apache.kafka.common.config.AbstractConfig;
-import org.apache.kafka.common.config.ConfigDef;
-import org.apache.kafka.common.config.ConfigDef.Importance;
-import org.apache.kafka.common.config.ConfigDef.Type;
-
-import com.uber.data.kafka.connect.workload.common.WorkloadProvider;
-
-import java.util.Map;
-import java.util.concurrent.TimeUnit;
-
-import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
-
-public class UberClusterAssignorConfig extends AbstractConfig {
-
- // TODO: Does this still need to be configurable?
- public static final String UBER_WORKLOAD_PROVIDER_CONFIG = "uber.workload.provider";
- public static final String UBER_WORKLOAD_PROVIDER_DOC = "Implementation of the "
- + WorkloadProvider.class.getName() + " interface that will be used to determine the total workload "
- + "of each worker. If left unset, then the default Kafka Connect rebalancing algorithm will be used instead, "
- + "which does not take workload into account and weighs all tasks equally";
-
- // TODO: This can be moved to Flipr
- public static final String UBER_WORKER_LOAD_THRESHOLD_CONFIG = "uber.worker.load.threshold";
- public static final double UBER_WORKER_LOAD_THRESHOLD_DEFAULT = 1;
- public static final String UBER_WORKER_LOAD_THRESHOLD_DOC = "The maximum total load "
- + "that a single worker should be assigned during rebalance. This will be enforced on a best-effort basis; "
- + "in some cases (such as if the cluster is underprovisioned), this threshold will be breached instead of "
- + "leaving tasks unassigned";
-
- // TODO: Flipr
- public static final String UBER_DEFAULT_TASK_LOAD_CONFIG = "uber.default.task.load";
- public static final double UBER_DEFAULT_TASK_LOAD_DEFAULT = 0.1; // Allow up to 10 unknown tasks per worker
- public static final String UBER_DEFAULT_TASK_LOAD_DOC = "The default load to assume "
- + "for tasks that do not have a known load according to the configured workload provider";
-
- // TODO: Flipr
- public static final String UBER_PERIODIC_REBALANCE_INTERVAL_MS_CONFIG = "uber.periodic.rebalance.interval.ms";
- public static final int UBER_PERIODIC_REBALANCE_INTERVAL_MS_DEFAULT = (int) TimeUnit.MINUTES.toMillis(5);
- public static final String UBER_PERIODIC_REBALANCE_INTERVAL_MS_DOC = "How frequently to trigger a "
- + "cluster rebalance, even if there has been no change in the number of workers, connectors, or tasks. "
- + "This can be used to respond to changes in traffic and balance workloads across workers. "
- + "Set to zero or a negative value to disable periodic rebalancing";
-
- private final WorkloadProvider workloadProvider;
-
- @SuppressWarnings("this-escape")
- public UberClusterAssignorConfig(Map<?, ?> props) {
- super(config(), props);
- WorkloadProvider configuredWorkloadProvider = getConfiguredInstance(UBER_WORKLOAD_PROVIDER_CONFIG, WorkloadProvider.class);
- this.workloadProvider = configuredWorkloadProvider != null ? configuredWorkloadProvider : WorkloadProvider.noOp();
- }
-
- public static ConfigDef config() {
- return new ConfigDef()
- .define(UBER_WORKLOAD_PROVIDER_CONFIG,
- ConfigDef.Type.CLASS,
- null, // Null instead of ConfigDef.NO_DEFAULT_VALUE to signify that the property is optional
- ConfigDef.Importance.HIGH,
- UBER_WORKLOAD_PROVIDER_DOC)
- .define(UBER_WORKER_LOAD_THRESHOLD_CONFIG,
- ConfigDef.Type.DOUBLE,
- UBER_WORKER_LOAD_THRESHOLD_DEFAULT,
- atLeast(0),
- Importance.MEDIUM,
- UBER_WORKER_LOAD_THRESHOLD_DOC)
- .define(UBER_DEFAULT_TASK_LOAD_CONFIG,
- ConfigDef.Type.DOUBLE,
- UBER_DEFAULT_TASK_LOAD_DEFAULT,
- atLeast(0),
- Importance.MEDIUM,
- UBER_DEFAULT_TASK_LOAD_DOC
- ).define(UBER_PERIODIC_REBALANCE_INTERVAL_MS_CONFIG,
- Type.INT,
- UBER_PERIODIC_REBALANCE_INTERVAL_MS_DEFAULT,
- Importance.MEDIUM,
- UBER_PERIODIC_REBALANCE_INTERVAL_MS_DOC
- );
- }
-
- public WorkloadProvider workloadProvider() {
- return workloadProvider;
- }
-
- public double uberWorkerLoadThreshold() {
- return getDouble(UBER_WORKER_LOAD_THRESHOLD_CONFIG);
- }
-
- public double uberDefaultTaskLoad() {
- return getDouble(UBER_DEFAULT_TASK_LOAD_CONFIG);
- }
-
- public int uberPeriodicRebalanceIntervalMs() {
- return getInt(UBER_PERIODIC_REBALANCE_INTERVAL_MS_CONFIG);
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/UberWorkerLoad.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/UberWorkerLoad.java
deleted file mode 100644
index fc081e64d55..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/UberWorkerLoad.java
+++ /dev/null
@@ -1,70 +0,0 @@
-package com.uber.data.kafka.connect.distributed.assignment;
-
-import com.uber.data.kafka.connect.distributed.ConnectorTaskId;
-import com.uber.data.kafka.connect.distributed.ConnectorsAndTasks;
-
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Objects;
-
-public class UberWorkerLoad {
-
- private final String worker;
- private final Collection<String> connectors;
- private final Collection<ConnectorTaskId> tasks;
-
- public UberWorkerLoad(String worker, ConnectorsAndTasks connectorsAndTasks) {
- this(worker, connectorsAndTasks.connectors(), connectorsAndTasks.tasks());
- }
-
- public UberWorkerLoad(
- String worker,
- Collection<String> connectors,
- Collection<ConnectorTaskId> tasks
- ) {
- this.worker = worker;
- this.connectors = connectors != null ? new ArrayList<>(connectors) : new ArrayList<>();
- this.tasks = tasks != null ? new ArrayList<>(tasks) : new ArrayList<>();
- }
-
- public String worker() {
- return worker;
- }
-
- public Collection<String> connectors() {
- return connectors;
- }
-
- public Collection<ConnectorTaskId> tasks() {
- return tasks;
- }
-
- public void assign(String connector) {
- connectors.add(connector);
- }
-
- public void assign(ConnectorTaskId task) {
- tasks.add(task);
- }
-
- @Override
- public String toString() {
- return "UberWorkerLoad{ worker=" + worker + ", connectorIds=" + connectors + ", taskIds=" + tasks + '}';
- }
-
- @Override
- public boolean equals(Object o) {
- if (this == o) {
- return true;
- }
- if (!(o instanceof UberWorkerLoad that)) {
- return false;
- }
- return Objects.equals(this.worker, that.worker);
- }
-
- @Override
- public int hashCode() {
- return Objects.hash(worker);
- }
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/WorkerConnectorAssignor.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/WorkerConnectorAssignor.java
deleted file mode 100644
index 47b66658936..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/WorkerConnectorAssignor.java
+++ /dev/null
@@ -1,48 +0,0 @@
-package com.uber.data.kafka.connect.distributed.assignment;
-
-import org.apache.kafka.connect.runtime.distributed.UberWorkerLoad;
-
-import java.util.Collection;
-
-public class WorkerConnectorAssignor extends AbstractAssignor<UberWorkerLoad, String> {
-
- public WorkerConnectorAssignor(int connectorWorkerThreshold) {
- // Try to have at most one connector per worker, and permit at most 10 at a time
- super(connectorWorkerThreshold, 1, 10);
- }
-
- @Override
- protected Double jobLoad(String job) {
- return 1.0;
- }
-
- @Override
- protected void assign(UberWorkerLoad worker, String connector) {
- worker.assign(connector);
- }
-
- @Override
- protected void revoke(UberWorkerLoad worker, String connector) {
- worker.connectors().remove(connector);
- }
-
- @Override
- protected Collection<String> jobs(UberWorkerLoad worker) {
- return worker.connectors();
- }
-
- @Override
- protected String groupType() {
- return "worker";
- }
-
- @Override
- protected String jobType() {
- return "connector";
- }
-
- @Override
- protected String groupName(UberWorkerLoad worker) {
- return worker.worker();
- }
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/WorkerTaskAssignor.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/WorkerTaskAssignor.java
deleted file mode 100644
index ad886b59894..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/WorkerTaskAssignor.java
+++ /dev/null
@@ -1,81 +0,0 @@
-package com.uber.data.kafka.connect.distributed.assignment;
-
-import org.apache.kafka.connect.runtime.distributed.UberWorkerLoad;
-
-import com.uber.data.kafka.connect.distributed.ClusterConfigState;
-import com.uber.data.kafka.connect.distributed.ConnectorTaskId;
-import com.uber.data.kafka.connect.workload.common.WorkloadProvider;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.Collection;
-import java.util.Map;
-
-public class WorkerTaskAssignor extends AbstractAssignor<UberWorkerLoad, ConnectorTaskId> {
-
- private static final Logger log = LoggerFactory.getLogger(WorkerTaskAssignor.class);
-
- private final WorkloadProvider workloadProvider;
- private final ClusterConfigState clusterConfigState;
-
- public WorkerTaskAssignor(
- WorkloadProvider workloadProvider,
- ClusterConfigState clusterConfigState,
- double workerLoadThreshold,
- double defaultTaskLoad,
- int workerTaskLimit
- ) {
- super(workerLoadThreshold, defaultTaskLoad, workerTaskLimit);
- this.workloadProvider = workloadProvider;
- this.clusterConfigState = clusterConfigState;
- }
-
- @Override
- protected Double jobLoad(ConnectorTaskId task) {
- Map<String, String> taskConfig = clusterConfigState.taskConfig(task);
- if (taskConfig == null) {
- log.warn("Unable to find config for task {}", task);
- return null;
- }
-
- String connector = task.connector();
- Map<String, String> connectorConfig = clusterConfigState.connectorConfig(connector);
- if (connectorConfig == null) {
- log.warn("Unable to find config for connector {}", connector);
- return null;
- }
-
- return workloadProvider.taskLoad(connector, task.task(), connectorConfig, taskConfig);
- }
-
- @Override
- protected void assign(UberWorkerLoad worker, ConnectorTaskId task) {
- worker.assign(task);
- }
-
- @Override
- protected void revoke(UberWorkerLoad worker, ConnectorTaskId task) {
- worker.tasks().remove(task);
- }
-
- @Override
- protected Collection<ConnectorTaskId> jobs(UberWorkerLoad worker) {
- return worker.tasks();
- }
-
- @Override
- protected String groupType() {
- return "worker";
- }
-
- @Override
- protected String jobType() {
- return "task";
- }
-
- @Override
- protected String groupName(UberWorkerLoad worker) {
- return worker.worker();
- }
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/WorkloadProvider.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/WorkloadProvider.java
deleted file mode 100644
index 8d22d41f2c9..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/distributed/assignment/WorkloadProvider.java
+++ /dev/null
@@ -1,39 +0,0 @@
-package com.uber.data.kafka.connect.distributed.assignment;
-
-import org.apache.kafka.common.Configurable;
-import org.apache.kafka.connect.components.Versioned;
-
-import java.io.Closeable;
-import java.util.Map;
-
-public interface WorkloadProvider extends Versioned, Configurable, Closeable {
-
- /**
- * Determine and return the load for the given task, if it is known
- * @return the load for the task, or null if none is known
- */
- Double taskLoad(String connector, int taskId, Map<String, String> connectorConfig, Map<String, String> taskConfig);
-
- static com.uber.data.kafka.connect.workload.common.WorkloadProvider noOp() {
- return new com.uber.data.kafka.connect.workload.common.WorkloadProvider() {
- @Override
- public Double taskLoad(String connector, int taskId, Map<String, String> connectorConfig, Map<String, String> taskConfig) {
- return null;
- }
-
- @Override
- public void close() {
- }
-
- @Override
- public void configure(Map<String, ?> configs) {
- }
-
- @Override
- public String version() {
- return "no-op";
- }
- };
- }
-
-}
\ No newline at end of file
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/UberConnectRestExtension.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/UberConnectRestExtension.java
deleted file mode 100644
index 77abd774dda..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/UberConnectRestExtension.java
+++ /dev/null
@@ -1,60 +0,0 @@
-package com.uber.data.kafka.connect.rest;
-
-import org.apache.kafka.common.utils.AppInfoParser;
-import org.apache.kafka.connect.health.ConnectClusterState;
-import org.apache.kafka.connect.rest.ConnectRestExtension;
-import org.apache.kafka.connect.rest.ConnectRestExtensionContext;
-
-import com.uber.data.kafka.connect.rest.filter.ContentTypeHeaderFilter;
-import com.uber.data.kafka.connect.rest.interceptor.ConnectorConfigInterceptor;
-import com.uber.data.kafka.connect.rest.interceptor.ConnectorOffsetsInterceptor;
-import com.uber.data.kafka.connect.rest.resources.AssignmentsResource;
-import com.uber.data.kafka.connect.rest.resources.WorkloadsResource;
-
-import org.glassfish.jersey.internal.inject.AbstractBinder;
-
-import java.util.Map;
-
-import jakarta.ws.rs.core.Configurable;
-
-public class UberConnectRestExtension implements ConnectRestExtension {
-
- @Override
- public void configure(Map<String, ?> configs) {
- }
-
- @Override
- public void register(ConnectRestExtensionContext restPluginContext) {
- Configurable<?> configurable = restPluginContext.configurable();
- configurable.register(ContentTypeHeaderFilter.class);
- configurable.register(new Binder(restPluginContext.clusterState()));
- configurable.register(ConnectorConfigInterceptor.class);
- configurable.register(ConnectorOffsetsInterceptor.class);
- configurable.register(WorkloadsResource.class);
- configurable.register(AssignmentsResource.class);
- }
-
- @Override
- public void close() {
- }
-
- @Override
- public String version() {
- return AppInfoParser.getVersion();
- }
-
- private static class Binder extends AbstractBinder {
- private final ConnectClusterState clusterState;
-
- public Binder(ConnectClusterState clusterState) {
- this.clusterState = clusterState;
- }
-
- @Override
- protected void configure() {
- bind(clusterState).to(ConnectClusterState.class);
- }
- }
-
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/AssignedPartitionEntity.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/AssignedPartitionEntity.java
deleted file mode 100644
index 859b6523b6a..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/AssignedPartitionEntity.java
+++ /dev/null
@@ -1,29 +0,0 @@
-package com.uber.data.kafka.connect.rest.entities;
-
-import com.fasterxml.jackson.annotation.JsonInclude;
-import com.fasterxml.jackson.annotation.JsonInclude.Include;
-import com.fasterxml.jackson.annotation.JsonProperty;
-import com.uber.data.kafka.connect.ureplicator3.coordination.AssignedPartition;
-
-@JsonInclude(Include.NON_NULL)
-public record AssignedPartitionEntity(
- @JsonProperty String topic,
- @JsonProperty int partition,
- @JsonProperty String dstTopic,
- @JsonProperty Integer dstPartition
-) {
-
- public static AssignedPartitionEntity of(AssignedPartition assignedPartition) {
- return new AssignedPartitionEntity(
- assignedPartition.srcTopic(),
- assignedPartition.srcPartition(),
- assignedPartition.dstTopic(),
- assignedPartition.dstPartition()
- );
- }
-
- public AssignedPartition assignedPartition() {
- return new AssignedPartition(topic, partition, dstTopic, dstPartition);
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/CachedWorkloadEntity.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/CachedWorkloadEntity.java
deleted file mode 100644
index 9bb16ca8d34..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/CachedWorkloadEntity.java
+++ /dev/null
@@ -1,15 +0,0 @@
-package com.uber.data.kafka.connect.rest.entities;
-
-import com.fasterxml.jackson.annotation.JsonProperty;
-import com.uber.data.kafka.connect.workload.common.CachedWorkload;
-
-public record CachedWorkloadEntity(
- @JsonProperty long lastUpdated,
- @JsonProperty WorkloadEntity workload
-) {
-
- public static CachedWorkloadEntity of(CachedWorkload cachedWorkload) {
- WorkloadEntity workload = WorkloadEntity.of(cachedWorkload.workload());
- return new CachedWorkloadEntity(cachedWorkload.lastUpdated(), workload);
- }
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/PartitionWorkloadEntity.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/PartitionWorkloadEntity.java
deleted file mode 100644
index 6dada99a8c2..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/PartitionWorkloadEntity.java
+++ /dev/null
@@ -1,59 +0,0 @@
-package com.uber.data.kafka.connect.rest.entities;
-
-import com.fasterxml.jackson.annotation.JsonProperty;
-import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
-import com.uber.data.kafka.connect.workload.common.Workload;
-
-import java.util.Comparator;
-
-// TODO: Property name overrides aren't being used in REST responses for some reason
-public record PartitionWorkloadEntity(
- @JsonProperty("replicated_partition") ReplicatedTopicPartitionEntity replicatedPartition,
- @JsonProperty("bytes_per_second") double bytesPerSecond,
- @JsonProperty("messages_per_second") double messagesPerSecond
-) implements Comparable<PartitionWorkloadEntity> {
-
- public static PartitionWorkloadEntity of(ReplicatedTopicPartition replicatedPartition, Workload workload) {
- return new PartitionWorkloadEntity(
- ReplicatedTopicPartitionEntity.of(replicatedPartition),
- workload.getBytesPerSecond(),
- workload.getMessagesPerSecond()
- );
- }
-
- public static Comparator<PartitionWorkloadEntity> lexicographicComparator() {
- return Comparator.<PartitionWorkloadEntity, String>comparing(w -> w.replicatedPartition().sourceKafka())
- .thenComparing(w -> w.replicatedPartition().targetKafka())
- .thenComparing(w -> w.replicatedPartition().topic())
- .thenComparing(w -> w.replicatedPartition().partition())
- .thenComparing(Comparator.naturalOrder());
- }
-
- public Workload workload() {
- return new Workload(bytesPerSecond, messagesPerSecond);
- }
-
- public ReplicatedTopicPartition replicatedTopicPartition() {
- return new ReplicatedTopicPartition(
- replicatedPartition.sourceKafka(),
- replicatedPartition.targetKafka(),
- replicatedPartition.topic(),
- replicatedPartition.partition()
- );
- }
-
- @Override
- public int compareTo(PartitionWorkloadEntity that) {
- if (that == null)
- return 1;
-
- if (this == that)
- return 0;
-
- int result = Double.compare(this.bytesPerSecond, that.bytesPerSecond);
- if (result != 0)
- return result;
-
- return Double.compare(this.messagesPerSecond, that.messagesPerSecond);
- }
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/ReplicatedTopicPartitionEntity.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/ReplicatedTopicPartitionEntity.java
deleted file mode 100644
index d79ef6ca9de..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/ReplicatedTopicPartitionEntity.java
+++ /dev/null
@@ -1,22 +0,0 @@
-package com.uber.data.kafka.connect.rest.entities;
-
-import com.fasterxml.jackson.annotation.JsonProperty;
-import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
-
-public record ReplicatedTopicPartitionEntity(
- @JsonProperty("src_kafka") String sourceKafka,
- @JsonProperty("dst_kafka") String targetKafka,
- @JsonProperty("topic") String topic,
- @JsonProperty("partition") int partition
-) {
-
- public static ReplicatedTopicPartitionEntity of(ReplicatedTopicPartition replicatedTopicPartition) {
- return new ReplicatedTopicPartitionEntity(
- replicatedTopicPartition.srcKafka(),
- replicatedTopicPartition.dstKafka(),
- replicatedTopicPartition.topic(),
- replicatedTopicPartition.partition()
- );
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/TaskAssignmentEntity.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/TaskAssignmentEntity.java
deleted file mode 100644
index e3d4288f1aa..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/TaskAssignmentEntity.java
+++ /dev/null
@@ -1,26 +0,0 @@
-package com.uber.data.kafka.connect.rest.entities;
-
-
-import com.fasterxml.jackson.annotation.JsonProperty;
-import com.uber.data.kafka.connect.ureplicator3.coordination.AssignedPartition;
-
-import java.util.Collection;
-import java.util.List;
-
-import static com.uber.data.kafka.connect.utils.UberConnectUtils.compareLexicographically;
-
-public record TaskAssignmentEntity(
- @JsonProperty String connector,
- @JsonProperty int task,
- @JsonProperty List<AssignedPartitionEntity> partitions
-) {
-
- public static TaskAssignmentEntity of(String connector, int task, Collection<AssignedPartition> partitions) {
- List<AssignedPartitionEntity> sortedPartitions = partitions.stream()
- .sorted((ap1, ap2) -> compareLexicographically(ap1.source(), ap2.source()))
- .map(AssignedPartitionEntity::of)
- .toList();
-
- return new TaskAssignmentEntity(connector, task, sortedPartitions);
- }
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/TaskWorkloadEntity.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/TaskWorkloadEntity.java
deleted file mode 100644
index 4a0d53b527f..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/TaskWorkloadEntity.java
+++ /dev/null
@@ -1,10 +0,0 @@
-package com.uber.data.kafka.connect.rest.entities;
-
-import com.fasterxml.jackson.annotation.JsonProperty;
-
-public record TaskWorkloadEntity(
- @JsonProperty String connector,
- @JsonProperty int taskId,
- @JsonProperty WorkloadEntity workload
-) {
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/TopicPartitionEntity.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/TopicPartitionEntity.java
deleted file mode 100644
index de42622d360..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/TopicPartitionEntity.java
+++ /dev/null
@@ -1,20 +0,0 @@
-package com.uber.data.kafka.connect.rest.entities;
-
-import org.apache.kafka.common.TopicPartition;
-
-import com.fasterxml.jackson.annotation.JsonProperty;
-
-public record TopicPartitionEntity(
- @JsonProperty String topic,
- @JsonProperty int partition
-) {
-
- public static TopicPartitionEntity of(TopicPartition topicPartition) {
- return new TopicPartitionEntity(topicPartition.topic(), topicPartition.partition());
- }
-
- public TopicPartition topicPartition() {
- return new TopicPartition(topic, partition);
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/WorkerWorkloadEntity.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/WorkerWorkloadEntity.java
deleted file mode 100644
index 6aa4d0ce366..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/WorkerWorkloadEntity.java
+++ /dev/null
@@ -1,13 +0,0 @@
-package com.uber.data.kafka.connect.rest.entities;
-
-import com.fasterxml.jackson.annotation.JsonProperty;
-
-import java.util.List;
-
-
-public record WorkerWorkloadEntity(
- @JsonProperty String worker,
- @JsonProperty("total_workload") WorkloadEntity totalWorkload,
- @JsonProperty("task_workloads") List<TaskWorkloadEntity> taskWorkloads
-) {
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/WorkloadEntity.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/WorkloadEntity.java
deleted file mode 100644
index b2396e187c6..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/entities/WorkloadEntity.java
+++ /dev/null
@@ -1,18 +0,0 @@
-package com.uber.data.kafka.connect.rest.entities;
-
-import com.fasterxml.jackson.annotation.JsonProperty;
-import com.uber.data.kafka.connect.workload.common.Workload;
-
-public record WorkloadEntity(
- @JsonProperty("bytes_per_second") double bytesPerSecond,
- @JsonProperty("messages_per_second") double messagesPerSecond
-) {
-
- public static WorkloadEntity of(Workload workload) {
- return new WorkloadEntity(workload.getBytesPerSecond(), workload.getMessagesPerSecond());
- }
-
- public Workload workload() {
- return new Workload(bytesPerSecond, messagesPerSecond);
- }
-}
\ No newline at end of file
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/filter/ContentTypeHeaderFilter.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/filter/ContentTypeHeaderFilter.java
deleted file mode 100644
index e299c23f3ad..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/filter/ContentTypeHeaderFilter.java
+++ /dev/null
@@ -1,56 +0,0 @@
-package com.uber.data.kafka.connect.rest.filter;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.Set;
-
-import jakarta.ws.rs.container.ContainerRequestContext;
-import jakarta.ws.rs.container.ContainerRequestFilter;
-import jakarta.ws.rs.container.PreMatching;
-import jakarta.ws.rs.core.MediaType;
-
-import static jakarta.ws.rs.core.HttpHeaders.CONTENT_TYPE;
-import static jakarta.ws.rs.core.MediaType.APPLICATION_FORM_URLENCODED_TYPE;
-import static jakarta.ws.rs.core.MediaType.APPLICATION_JSON;
-
-@PreMatching
-public class ContentTypeHeaderFilter implements ContainerRequestFilter {
-
- private static final Logger log = LoggerFactory.getLogger(ContentTypeHeaderFilter.class);
-
- private static final String SKIP_PARAM = "keepContentType";
-
- // Media types that we override unless explicitly instructed not to via the skip URL param
- // TODO: We could technically make this configurable but it's probably not worth the effort
- private static final Set<MediaType> STUPID_MEDIA_TYPES = Set.of(
- APPLICATION_FORM_URLENCODED_TYPE // Automatically injected by curl (this fucking shit took way too long to debug)
- );
-
- @Override
- public void filter(ContainerRequestContext containerRequestContext) {
- if (containerRequestContext.getUriInfo().getQueryParameters().containsKey(SKIP_PARAM)) {
- log.trace("Skipping {} header injection since {} URL param is specified", CONTENT_TYPE, SKIP_PARAM);
- return;
- }
-
- if (!containerRequestContext.hasEntity()) {
- log.trace("Skipping {} header injection since request has no body", CONTENT_TYPE);
- return;
- }
-
- MediaType requestMediaType = containerRequestContext.getMediaType();
- if (requestMediaType == null) {
- log.trace("Request does not have media type; will inject {} header", CONTENT_TYPE);
- } else if (STUPID_MEDIA_TYPES.contains(requestMediaType)) {
- log.trace("Request has stupid media type {}; will override by injecting {} header", requestMediaType, CONTENT_TYPE);
- } else {
- log.trace("Skipping {} header injection since request already has non-stupid media type {}", CONTENT_TYPE, requestMediaType);
- return;
- }
-
- log.debug("Adding {} header with value {} to request", CONTENT_TYPE, APPLICATION_JSON);
- containerRequestContext.getHeaders().putSingle(CONTENT_TYPE, APPLICATION_JSON);
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/interceptor/ConnectorConfigInterceptor.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/interceptor/ConnectorConfigInterceptor.java
deleted file mode 100644
index eca79cfa970..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/interceptor/ConnectorConfigInterceptor.java
+++ /dev/null
@@ -1,235 +0,0 @@
-package com.uber.data.kafka.connect.rest.interceptor;
-
-import org.apache.kafka.clients.CommonClientConfigs;
-import org.apache.kafka.connect.health.ConnectClusterState;
-import org.apache.kafka.connect.runtime.rest.entities.CreateConnectorRequest;
-import org.apache.kafka.connect.runtime.rest.errors.ConnectRestException;
-
-import com.uber.data.kafka.connect.utils.BrokerResolver;
-import com.uber.data.kafka.exception.ResolverException;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.IOException;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.function.Function;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
-
-import jakarta.inject.Inject;
-import jakarta.servlet.http.HttpServletRequest;
-import jakarta.ws.rs.HttpMethod;
-import jakarta.ws.rs.WebApplicationException;
-import jakarta.ws.rs.container.ResourceInfo;
-import jakarta.ws.rs.core.Context;
-import jakarta.ws.rs.core.Response.Status;
-import jakarta.ws.rs.core.UriInfo;
-import jakarta.ws.rs.ext.ReaderInterceptor;
-import jakarta.ws.rs.ext.ReaderInterceptorContext;
-
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Config.CLIENTS_PREFIX;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Config.TARGET_CLIENTS_PREFIX;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Config.TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Config.TARGET_CLUSTER_CONFIG;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Config.TARGET_PRODUCER_PREFIX;
-import static com.uber.data.kafka.connect.utils.UberConnectUtils.isURep3;
-
-public class ConnectorConfigInterceptor implements ReaderInterceptor {
-
- public static final String SKIP_QUERY_PARAM = "uberSkipInterception";
-
- private static final Logger log = LoggerFactory.getLogger(ConnectorConfigInterceptor.class);
-
- private static final String PRODUCER_OVERRIDE_PREFIX = "producer.override.";
- private static final String PRODUCER_OVERRIDE_BOOTSTRAP_CONFIG =
- PRODUCER_OVERRIDE_PREFIX + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG;
-
- private final ConnectClusterState clusterState;
-
- @Inject
- public ConnectorConfigInterceptor(ConnectClusterState clusterState) {
- this.clusterState = clusterState;
- }
-
- @Context
- private UriInfo uriInfo;
-
- @Context
- private ResourceInfo resourceInfo;
-
- @Context
- private HttpServletRequest httpRequest;
-
- @Override
- public Object aroundReadFrom(ReaderInterceptorContext readerInterceptorContext) throws IOException, WebApplicationException {
- Object entity = readerInterceptorContext.proceed();
- if (shouldSkip()) {
- return entity;
- } else if (entity == null) {
- log.debug("Request has no body; skipping interception");
- return null;
- }
-
- if (isConnectorPostRequest()) {
- if (entity instanceof CreateConnectorRequest createRequest) {
- return interceptConnectorConfig(
- createRequest.config(),
- createRequest,
- CreateConnectorRequest::config,
- intercepted -> new CreateConnectorRequest(createRequest.name(), intercepted, createRequest.initialState())
- );
- } else {
- log.error("Unable to intercept connector create request");
- return entity;
- }
- } else if (isConnectorConfigPutRequest()) {
- if (entity instanceof Map<?, ?> rawConfig) {
- @SuppressWarnings("unchecked")
- Map<String, String> originalConfig = (Map<String, String>) rawConfig;
- return interceptConnectorConfig(
- originalConfig,
- originalConfig,
- Function.identity(),
- Function.identity()
- );
- } else {
- log.error("Unable to intercept connector config put request");
- return entity;
- }
- } else if (isConnectorConfigPatchRequest()) {
- String path = uriInfo.getPath();
- Matcher m = Pattern.compile("/*connectors/+([^/]+)/+config").matcher(path);
- if (!m.matches()) {
- log.error("Unable to intercept connector config patch request: unrecognized URL path");
- return entity;
- }
- String connectorName = m.group(1);
- if (connectorName == null) {
- log.error("Unable to intercept connector config patch request: could not deduce connector name from URL path");
- return entity;
- }
- // This is technically prone to race conditions since the connector config may be
- // modified in between us reading it and when we allow the patch request to go through
- // Changing a connector's type is highly unusual and doing it with the kind of frequency
- // necessary to make that kind of race condition actually happen is not a usage pattern
- // worth accommodating right now
- Map<String, String> connectorConfig = clusterState.connectorConfig(connectorName);
- if (connectorConfig == null) {
- return entity;
- }
- if (entity instanceof Map<?, ?> rawPatch) {
- @SuppressWarnings("unchecked")
- Map<String, String> originalPatch = (Map<String, String>) rawPatch;
- return interceptConnectorConfig(
- connectorConfig,
- originalPatch,
- Function.identity(),
- Function.identity()
- );
- } else {
- log.error("Unable to intercept connector config patch request: unrecognized request body type {}", entity);
- return entity;
- }
- } else {
- return entity;
- }
- }
-
- private boolean shouldSkip() {
- if(httpRequest.getParameter(SKIP_QUERY_PARAM) != null) {
- log.debug("Skipping interception (requested by {} URL query parameter)", SKIP_QUERY_PARAM);
- return true;
- }
- return false;
- }
-
- private boolean isConnectorPostRequest() {
- return isConnectorConfigRequest("createConnector", HttpMethod.POST);
- }
-
- private boolean isConnectorConfigPutRequest() {
- return isConnectorConfigRequest("putConnectorConfig", HttpMethod.PUT);
- }
-
- private boolean isConnectorConfigPatchRequest() {
- return isConnectorConfigRequest("patchConnectorConfig", HttpMethod.PATCH);
- }
-
- private boolean isConnectorConfigRequest(String resourceMethod, String httpMethod) {
- return isConnectorRequest()
- && httpMethod.equals(httpRequest.getMethod())
- && resourceMethod.equals(resourceInfo.getResourceMethod().getName());
- }
-
- private boolean isConnectorRequest() {
- return resourceInfo.getResourceClass().getSimpleName().equals("ConnectorsResource");
- }
-
- private <T> T interceptConnectorConfig(
- Map<String, String> completeCurrentConnectorConfig,
- T originalEntity,
- Function<T, Map<String, String>> configExtractor,
- Function<Map<String, String>, T> entityReconstructor
- ) {
- if (!isURep3(completeCurrentConnectorConfig)) {
- log.debug("Ignoring config for non-uReplicator3 connector");
- return originalEntity;
- }
-
- Map<String, String> originalConfig = configExtractor.apply(originalEntity);
- Map<String, String> interceptedConfig = interceptConnectorConfig(originalConfig);
- return entityReconstructor.apply(interceptedConfig);
- }
-
- private Map<String, String> interceptConnectorConfig(Map<String, String> connectorConfig) {
- Map<String, String> result = new HashMap<>(connectorConfig);
-
- String targetCluster = connectorConfig.get(TARGET_CLUSTER_CONFIG);
- if (targetCluster != null) {
- if (!connectorConfig.containsKey(PRODUCER_OVERRIDE_BOOTSTRAP_CONFIG)) {
- String targetBootstrap = connectorConfig.get(TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG);
- if (targetBootstrap == null)
- targetBootstrap = resolveBootstrap(targetCluster);
- result.put(PRODUCER_OVERRIDE_BOOTSTRAP_CONFIG, targetBootstrap);
- }
- }
-
- translatePrefix(connectorConfig, result, CLIENTS_PREFIX, PRODUCER_OVERRIDE_PREFIX);
- translatePrefix(connectorConfig, result, TARGET_CLIENTS_PREFIX, PRODUCER_OVERRIDE_PREFIX);
- translatePrefix(connectorConfig, result, TARGET_PRODUCER_PREFIX, PRODUCER_OVERRIDE_PREFIX);
-
- // TODO: Auto-inject default partition byte and message rates,
- // and possibly max per-task byte and message rates
-
- return result;
- }
-
- private void translatePrefix(Map<String, String> from, Map<String, String> to, String fromPrefix, String toPrefix) {
- from.forEach((prefixedProp, value) -> {
- if (prefixedProp.startsWith(fromPrefix)) {
- String strippedProp = prefixedProp.substring(fromPrefix.length());
- if (strippedProp.isBlank()) {
- log.warn("Ignoring property {} with value {}", prefixedProp, value);
- return;
- }
-
- to.put(toPrefix + strippedProp, value);
- }
- });
- }
-
- private String resolveBootstrap(String cluster) {
- try {
- return BrokerResolver.resolveBootstrap(cluster);
- } catch (ResolverException e) {
- throw new ConnectRestException(
- Status.BAD_REQUEST,
- "Failed to resolve bootstrap servers for cluster " + cluster,
- e
- );
- }
- }
-
-}
\ No newline at end of file
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/interceptor/ConnectorOffsetsInterceptor.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/interceptor/ConnectorOffsetsInterceptor.java
deleted file mode 100644
index 721e38cf95b..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/interceptor/ConnectorOffsetsInterceptor.java
+++ /dev/null
@@ -1,147 +0,0 @@
-package com.uber.data.kafka.connect.rest.interceptor;
-
-import org.apache.kafka.common.config.ConfigDef;
-import org.apache.kafka.common.config.ConfigDef.Type;
-import org.apache.kafka.connect.health.ConnectClusterState;
-import org.apache.kafka.connect.runtime.rest.entities.ConnectorOffsets;
-import org.apache.kafka.connect.runtime.rest.errors.BadRequestException;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.IOException;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.TreeSet;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
-
-import jakarta.inject.Inject;
-import jakarta.servlet.http.HttpServletRequest;
-import jakarta.ws.rs.HttpMethod;
-import jakarta.ws.rs.WebApplicationException;
-import jakarta.ws.rs.container.ResourceInfo;
-import jakarta.ws.rs.core.Context;
-import jakarta.ws.rs.core.UriInfo;
-import jakarta.ws.rs.ext.ReaderInterceptor;
-import jakarta.ws.rs.ext.ReaderInterceptorContext;
-
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.TOPICS_CONFIG;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Task.SOURCE_PARTITION_KAFKA_TOPIC;
-import static com.uber.data.kafka.connect.utils.UberConnectUtils.isURep3;
-import static org.apache.kafka.connect.runtime.rest.resources.ConnectorsResource.UBER_FORCE_OFFSET_PATCH;
-
-public class ConnectorOffsetsInterceptor implements ReaderInterceptor {
-
- public static final String SKIP_QUERY_PARAM = "uberSkipInterception";
-
- private static final Logger log = LoggerFactory.getLogger(ConnectorOffsetsInterceptor.class);
-
- private final ConnectClusterState clusterState;
-
- @Inject
- public ConnectorOffsetsInterceptor(ConnectClusterState clusterState) {
- this.clusterState = clusterState;
- }
-
- @Context
- private UriInfo uriInfo;
-
- @Context
- private ResourceInfo resourceInfo;
-
- @Context
- private HttpServletRequest httpRequest;
-
- @Override
- public Object aroundReadFrom(ReaderInterceptorContext readerInterceptorContext) throws IOException, WebApplicationException {
- Object entity = readerInterceptorContext.proceed();
- if (shouldSkip()) {
- return entity;
- } else if (entity == null) {
- log.debug("Request has no body; skipping interception");
- return null;
- } else if (!Boolean.parseBoolean(uriInfo.getQueryParameters().getFirst(UBER_FORCE_OFFSET_PATCH))) {
- log.debug("Request will not be forced if connector is not stopped; skipping interception since no safeguard is necessary");
- return entity;
- }
-
- if (!isConnectorOffsetPatchRequest()) {
- return entity;
- }
-
- String path = uriInfo.getPath();
- Matcher m = Pattern.compile("/*connectors/+([^/]+)/+offsets").matcher(path);
- if (!m.matches()) {
- log.error("Unable to intercept connector offset patch request: unrecognized URL path");
- return entity;
- }
- String connectorName = m.group(1);
- if (connectorName == null) {
- log.error("Unable to intercept connector offset patch request: could not deduce connector name from URL path");
- return entity;
- }
-
- Map<String, String> connectorConfig = clusterState.connectorConfig(connectorName);
- if (connectorConfig == null) {
- log.warn("Unable to locate config for connector {}; will not be able to safeguard offset patch request", connectorName);
- return entity;
- }
-
- if (!isURep3(connectorConfig)) {
- log.warn("Skipping interception for connector {} since it is not uReplicator3", connectorName);
- return entity;
- }
-
- if (!(entity instanceof ConnectorOffsets offsets)) {
- log.warn("Unrecognized entity type {} for connector {}; will not be able to safeguard offset patch request", entity.getClass(), connectorName);
- return entity;
- }
-
- String topicsValue = connectorConfig.get(TOPICS_CONFIG);
- @SuppressWarnings("unchecked")
- List<String> topicsList = (List<String>) ConfigDef.parseType(TOPICS_CONFIG, topicsValue, Type.LIST);
- if (topicsList == null) {
- log.warn("Unable to parse topics list in config for connector {}; will not be able to safeguard offset patch request", connectorName);
- return entity;
- }
-
- Set<String> invalidTopics = new TreeSet<>();
- Set<String> topicsSet = new HashSet<>(topicsList); // We could also potentially fetch this from the assignment store, which might be more accurate
- offsets.offsets().forEach(offset -> {
- Object rawTopic = offset.partition().get(SOURCE_PARTITION_KAFKA_TOPIC);
- // TODO: Technically this is too broad; we should also narrow down by src/dst Kafka cluster as well
- // I'm too lazy to implement that right now and am currently betting it will not come back to bite me
- // (if you are reading this I was probably wrong, sorry 🥲)
- if (rawTopic instanceof String topic && topicsSet.contains(topic))
- invalidTopics.add(topic);
- });
-
- if (!invalidTopics.isEmpty()) {
- String message = "Cannot forcibly alter offsets for connector " + connectorName + " since it is still configured "
- + "to replicate these topics: " + invalidTopics
- + "\nEither stop the connector before modifying its offsets, or reconfigure the connector to no longer "
- + "replicate from these topics before trying again";
- throw new BadRequestException(message);
- }
-
- return entity;
- }
-
- private boolean shouldSkip() {
- if(httpRequest.getParameter(SKIP_QUERY_PARAM) != null) {
- log.debug("Skipping interception (requested by {} URL query parameter)", SKIP_QUERY_PARAM);
- return true;
- }
- return false;
- }
-
- private boolean isConnectorOffsetPatchRequest() {
- return "ConnectorsResource".equals(resourceInfo.getResourceClass().getSimpleName())
- && HttpMethod.PATCH.equals(httpRequest.getMethod())
- && "alterConnectorOffsets".equals(resourceInfo.getResourceMethod().getName());
- }
-
-}
\ No newline at end of file
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/resources/AssignmentsResource.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/resources/AssignmentsResource.java
deleted file mode 100644
index c263daa3933..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/resources/AssignmentsResource.java
+++ /dev/null
@@ -1,85 +0,0 @@
-package com.uber.data.kafka.connect.rest.resources;
-
-import org.apache.kafka.connect.health.ConnectClusterState;
-
-import com.uber.data.kafka.connect.rest.entities.TaskAssignmentEntity;
-import com.uber.data.kafka.connect.ureplicator3.coordination.AssignedPartition;
-import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStores;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.Collection;
-import java.util.Collections;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.stream.IntStream;
-
-import jakarta.inject.Inject;
-import jakarta.ws.rs.GET;
-import jakarta.ws.rs.Path;
-import jakarta.ws.rs.PathParam;
-import jakarta.ws.rs.Produces;
-import jakarta.ws.rs.core.MediaType;
-import jakarta.ws.rs.core.Response;
-import jakarta.ws.rs.core.Response.Status;
-
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Config.SOURCE_CLUSTER_CONFIG;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Config.TARGET_CLUSTER_CONFIG;
-
-@Path("/uber/assignments")
-@Produces(MediaType.APPLICATION_JSON)
-public class AssignmentsResource {
-
- private static final Logger log = LoggerFactory.getLogger(AssignmentsResource.class);
-
- private final ConnectClusterState clusterState;
-
- @Inject
- public AssignmentsResource(ConnectClusterState clusterState) {
- this.clusterState = clusterState;
- }
-
- @Path("/{connector}/all")
- @GET
- public Response getConnectorAssignments(@PathParam("connector") String connector) {
- Map<String, String> connectorConfig = clusterState.connectorConfig(connector);
- if (connectorConfig == null) {
- return Response.status(Status.NOT_FOUND)
- .entity("Connector '" + connector + " does not appear to exist")
- .build();
- }
-
- String srcKafka = connectorConfig.get(SOURCE_CLUSTER_CONFIG);
- if (srcKafka == null) {
- return Response.serverError()
- .entity("Could not determine source cluster for connector '" + connector + "'")
- .build();
- }
-
- String dstKafka = connectorConfig.get(TARGET_CLUSTER_CONFIG);
- if (dstKafka == null) {
- return Response.serverError()
- .entity("Could not determine target cluster for connector '" + connector + "'")
- .build();
- }
-
- Map<Integer, ? extends Collection<AssignedPartition>> taskAssignments =
- AssignmentStores.get().getAssignments(connector, srcKafka, dstKafka);
-
- int highestTask = Collections.max(taskAssignments.keySet());
-
- List<TaskAssignmentEntity> body = IntStream.rangeClosed(0, highestTask)
- .mapToObj(i -> {
- Collection<AssignedPartition> assignment = taskAssignments.get(i);
- if (assignment == null) {
- assignment = Set.of();
- }
- return TaskAssignmentEntity.of(connector, i, assignment);
- }).toList();
-
- return Response.ok(body).build();
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/resources/WorkloadsResource.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/resources/WorkloadsResource.java
deleted file mode 100644
index 28b1fd28123..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/rest/resources/WorkloadsResource.java
+++ /dev/null
@@ -1,316 +0,0 @@
-package com.uber.data.kafka.connect.rest.resources;
-
-import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.connect.health.ConnectClusterState;
-import org.apache.kafka.connect.health.ConnectorHealth;
-import org.apache.kafka.connect.health.TaskState;
-import org.apache.kafka.connect.util.ConnectorTaskId;
-
-import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
-import com.uber.data.kafka.connect.rest.entities.CachedWorkloadEntity;
-import com.uber.data.kafka.connect.rest.entities.PartitionWorkloadEntity;
-import com.uber.data.kafka.connect.rest.entities.TaskWorkloadEntity;
-import com.uber.data.kafka.connect.rest.entities.WorkerWorkloadEntity;
-import com.uber.data.kafka.connect.rest.entities.WorkloadEntity;
-import com.uber.data.kafka.connect.workload.common.CachedWorkload;
-import com.uber.data.kafka.connect.workload.common.Workload;
-import com.uber.data.kafka.connect.workload.kafka.KafkaBasedWorkloadProvider;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.time.Duration;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Comparator;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.TreeMap;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.Future;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
-import java.util.stream.Collectors;
-import java.util.stream.IntStream;
-
-import jakarta.inject.Inject;
-import jakarta.ws.rs.DELETE;
-import jakarta.ws.rs.DefaultValue;
-import jakarta.ws.rs.GET;
-import jakarta.ws.rs.POST;
-import jakarta.ws.rs.Path;
-import jakarta.ws.rs.PathParam;
-import jakarta.ws.rs.Produces;
-import jakarta.ws.rs.QueryParam;
-import jakarta.ws.rs.core.MediaType;
-import jakarta.ws.rs.core.Response;
-import jakarta.ws.rs.core.Response.Status;
-
-import static com.uber.data.kafka.connect.utils.UberConnectUtils.assignedPartitions;
-import static com.uber.data.kafka.connect.utils.UberConnectUtils.awaitAll;
-import static com.uber.data.kafka.connect.utils.UberConnectUtils.createErrorMessage;
-
-@Path("/uber/workloads")
-@Produces(MediaType.APPLICATION_JSON)
-public class WorkloadsResource {
-
- private static final Logger log = LoggerFactory.getLogger(WorkloadsResource.class);
-
- private final ConnectClusterState clusterState;
-
- @Inject
- public WorkloadsResource(ConnectClusterState clusterState) {
- this.clusterState = clusterState;
- }
-
- @GET
- @Path("/topics/all")
- public List<PartitionWorkloadEntity> getAllTopicWorkloads(
- @QueryParam("src") String srcKafka,
- @QueryParam("dst") String dstKafka,
- @QueryParam("topic") String topic,
- @QueryParam("refresh") @DefaultValue("false") boolean refresh
- ) {
- if (refresh)
- KafkaBasedWorkloadProvider.refreshWorkloadCache();
-
- return KafkaBasedWorkloadProvider.getAllWorkloads().entrySet().stream()
- .filter(e -> matches(e.getKey(), srcKafka, dstKafka, topic))
- .map(e -> PartitionWorkloadEntity.of(e.getKey(), e.getValue()))
- .sorted(PartitionWorkloadEntity.lexicographicComparator())
- .collect(Collectors.toList());
- }
-
- @GET
- @Path("/topics/historical/{src}/{dst}/{topic}/{partition}")
- public Response getHistoricalWorkloads(
- @PathParam("src") String srcKafka,
- @PathParam("dst") String dstKafka,
- @PathParam("topic") String topic,
- @PathParam("partition") int partition,
- @QueryParam("refresh") @DefaultValue("false") boolean refresh
- ) {
- if (refresh)
- KafkaBasedWorkloadProvider.refreshWorkloadCache();
-
- TopicPartition topicPartition = new TopicPartition(topic, partition);
- List<CachedWorkload> cachedWorkloads = KafkaBasedWorkloadProvider.getHistoricalWorkloads(srcKafka, dstKafka, topicPartition);
- if (cachedWorkloads == null)
- return Response.status(Status.NOT_FOUND).entity("no workloads found for the requested topic").build();
-
- List<CachedWorkloadEntity> body = cachedWorkloads.stream()
- .map(CachedWorkloadEntity::of)
- .toList();
- return Response.ok(body).build();
- }
-
- @POST
- @Path("/topics/overrides/{src}/{dst}/{topic}/{partition}")
- public Response publishTopicWorkload(
- @QueryParam("awaitAck") @DefaultValue("true") boolean awaitAck,
- @PathParam("src") String srcKafka,
- @PathParam("dst") String dstKafka,
- @PathParam("topic") String topic,
- @PathParam("partition") int partition,
- WorkloadEntity workload
- ) {
- return publishWorkload(srcKafka, dstKafka, topic, partition, workload.workload(), awaitAck);
- }
-
- @DELETE
- @Path("/topics/overrides/{src}/{dst}/{topic}")
- public Response deleteTopicWorkloads(
- @QueryParam("awaitAck") @DefaultValue("true") boolean awaitAck,
- @QueryParam("partitions") int partitions,
- @PathParam("src") String srcKafka,
- @PathParam("dst") String dstKafka,
- @PathParam("topic") String topic
- ) {
- List<? extends Future<?>> produceFutures = IntStream.range(0, partitions)
- .mapToObj(partition -> KafkaBasedWorkloadProvider.publishWorkload(
- srcKafka,
- dstKafka,
- new TopicPartition(topic, partition),
- null
- )).toList();
-
- if (awaitAck) {
- try {
- awaitAll(produceFutures, Duration.ofMinutes(1));
- } catch (InterruptedException e) {
- return overrideErrorResponse(e);
- } catch (TimeoutException e) {
- return overrideErrorResponse(e);
- } catch (ExecutionException e) {
- return overrideErrorResponse(e);
- }
- } else {
- log.debug("Not awaiting producer ack for workload override on {} partitions of topic {} in pipeline @{}@{}",
- partitions, topic, srcKafka, dstKafka
- );
- }
-
- return Response.noContent().build();
- }
-
- @DELETE
- @Path("/topics/overrides/{src}/{dst}/{topic}/{partition}")
- public Response deleteTopicWorkload(
- @QueryParam("awaitAck") @DefaultValue("true") boolean awaitAck,
- @PathParam("src") String srcKafka,
- @PathParam("dst") String dstKafka,
- @PathParam("topic") String topic,
- @PathParam("partition") int partition
- ) {
- return publishWorkload(srcKafka, dstKafka, topic, partition, null, awaitAck);
- }
-
- private Response publishWorkload(
- String srcKafka,
- String dstKafka,
- String topic,
- int partition,
- Workload workload,
- boolean awaitAck
- ) {
- Future<?> produceFuture = KafkaBasedWorkloadProvider.publishWorkload(
- srcKafka,
- dstKafka,
- new TopicPartition(topic, partition),
- workload
- );
- if (awaitAck) {
- try {
- produceFuture.get(1, TimeUnit.MINUTES);
- } catch (InterruptedException e) {
- return overrideErrorResponse(e);
- } catch (TimeoutException e) {
- return overrideErrorResponse(e);
- } catch (ExecutionException e) {
- return overrideErrorResponse(e);
- }
- } else {
- ReplicatedTopicPartition loggedPartition = new ReplicatedTopicPartition(srcKafka, dstKafka, topic, partition);
- log.debug("Not awaiting producer ack for workload override on {}", loggedPartition);
- }
-
- return Response.noContent().build();
- }
-
- @GET
- @Path("/workers/all")
- public List<WorkerWorkloadEntity> getAllWorkerWorkloads(@QueryParam("refresh") @DefaultValue("false") boolean refresh) {
- Collection<String> connectors = clusterState.connectors();
- Map<String, Set<ConnectorTaskId>> workerTasks = new HashMap<>();
-
- for (String connector : connectors) {
- ConnectorHealth connectorHealth = clusterState.connectorHealth(connector);
- if (connectorHealth == null) {
- // The connector may have just been deleted
- continue;
- }
- for (TaskState taskState : connectorHealth.tasksState().values()) {
- workerTasks.computeIfAbsent(taskState.workerId(), k -> new HashSet<>())
- .add(new ConnectorTaskId(connector, taskState.taskId()));
- }
- }
-
- if (refresh)
- KafkaBasedWorkloadProvider.refreshWorkloadCache();
- Map<ReplicatedTopicPartition, Workload> workloads = KafkaBasedWorkloadProvider.getAllWorkloads();
-
- Map<String, Map<ConnectorTaskId, Workload>> workerTaskWorkloads = new HashMap<>();
- Map<String, Map<String, String>> connectorConfigs = new HashMap<>();
- workerTasks.forEach((worker, tasks) ->
- tasks.forEach(task -> {
- Map<String, String> connectorConfig = connectorConfigs.computeIfAbsent(task.connector(), clusterState::connectorConfig);
- if (connectorConfig == null) {
- // The connector may have just been deleted
- return;
- }
-
- Set<ReplicatedTopicPartition> taskPartitions = assignedPartitions(task.connector(), task.task(), connectorConfig);
-
- if (taskPartitions == null) {
- return;
- }
-
- for (ReplicatedTopicPartition replicatedTopicPartition : taskPartitions) {
- Workload workload = workloads.get(replicatedTopicPartition);
- if (workload == null) {
- continue;
- }
- workerTaskWorkloads.computeIfAbsent(
- worker,
- k -> new TreeMap<>(Comparator.comparing(ConnectorTaskId::connector).thenComparing(ConnectorTaskId::task))
- ).compute(
- task,
- (k, existingWorkload) -> workload.plus(existingWorkload)
- );
- }
- })
- );
-
- return workerTaskWorkloads.entrySet().stream()
- .sorted(Map.Entry.comparingByKey())
- .map(e -> {
- String worker = e.getKey();
- Map<ConnectorTaskId, Workload> taskWorkloads = e.getValue();
-
- List<TaskWorkloadEntity> taskWorkloadEntities = new ArrayList<>();
- double totalByteRate = 0;
- double totalMessageRate = 0;
-
- // TODO: We might want to order this somehow (probably lexicographically by connector name and then task ID)
- for (Map.Entry<ConnectorTaskId, Workload> taskWorkloadEntry : taskWorkloads.entrySet()) {
- totalByteRate += taskWorkloadEntry.getValue().getBytesPerSecond();
- totalMessageRate += taskWorkloadEntry.getValue().getMessagesPerSecond();;
-
- TaskWorkloadEntity taskWorkloadEntity = new TaskWorkloadEntity(
- taskWorkloadEntry.getKey().connector(),
- taskWorkloadEntry.getKey().task(),
- WorkloadEntity.of(taskWorkloadEntry.getValue())
- );
- taskWorkloadEntities.add(taskWorkloadEntity);
- }
-
- WorkloadEntity totalWorkload = new WorkloadEntity(totalByteRate, totalMessageRate);
- return new WorkerWorkloadEntity(worker, totalWorkload, taskWorkloadEntities);
- }).toList();
- }
-
- private boolean matches(ReplicatedTopicPartition rtp, String srcKafka, String dstKafka, String topic) {
- if (srcKafka != null && !srcKafka.equals(rtp.srcKafka()))
- return false;
-
- if (dstKafka != null && !dstKafka.equals(rtp.dstKafka()))
- return false;
-
- if (topic != null && !topic.equals(rtp.topic()))
- return false;
-
- return true;
- }
-
- private Response overrideErrorResponse(InterruptedException e) {
- return Response.serverError()
- .entity(createErrorMessage("Interrupted while waiting for producer ack", e))
- .build();
- }
-
- private Response overrideErrorResponse(TimeoutException e) {
- return Response.serverError()
- .entity(createErrorMessage("Timed out waiting for producer ack", e))
- .build();
- }
-
- private Response overrideErrorResponse(ExecutionException e) {
- return Response.serverError()
- .entity(createErrorMessage("Failed to publish workload override", e))
- .build();
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/TaskPartitionAssignor.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/TaskPartitionAssignor.java
deleted file mode 100644
index d2a51bc60b8..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/TaskPartitionAssignor.java
+++ /dev/null
@@ -1,112 +0,0 @@
-package com.uber.data.kafka.connect.ureplicator3;
-
-import org.apache.kafka.common.TopicPartition;
-
-import com.google.common.annotations.VisibleForTesting;
-import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
-import com.uber.data.kafka.connect.distributed.assignment.AbstractAssignor;
-import com.uber.data.kafka.connect.workload.common.Workload;
-import com.uber.data.kafka.connect.workload.kafka.KafkaBasedWorkloadProvider;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.Collection;
-import java.util.Map;
-import java.util.stream.Collectors;
-
-public class TaskPartitionAssignor extends AbstractAssignor<TaskPartitions, TopicPartition> {
-
- private static final Logger log = LoggerFactory.getLogger(TaskPartitionAssignor.class);
-
- private final Map<TopicPartition, Double> partitionLoads;
-
- public TaskPartitionAssignor(
- String srcKafka,
- String dstKafka,
- double defaultPartitionLoad,
- double taskByteRateThreshold,
- double taskMessageRateThreshold
- ) {
- // TODO: Should we have a limit of partitions per task?
- this(
- defaultPartitionLoad,
- retrievePartitionLoads(srcKafka, dstKafka, taskByteRateThreshold, taskMessageRateThreshold)
- );
- }
-
- @VisibleForTesting
- TaskPartitionAssignor(
- double defaultPartitionLoad,
- Map<TopicPartition, Double> partitionLoads
- ) {
- super(1, defaultPartitionLoad, 0);
- this.partitionLoads = partitionLoads;
- }
-
- @Override
- protected Double jobLoad(TopicPartition topicPartition) {
- return partitionLoads.get(topicPartition);
- }
-
- @Override
- protected void assign(TaskPartitions task, TopicPartition topicPartition) {
- task.partitions().add(topicPartition);
- }
-
- @Override
- protected void revoke(TaskPartitions task, TopicPartition topicPartition) {
- task.partitions().remove(topicPartition);
- }
-
- @Override
- protected Collection<TopicPartition> jobs(TaskPartitions task) {
- return task.partitions();
- }
-
- @Override
- protected String groupType() {
- return "task";
- }
-
- @Override
- protected String jobType() {
- return "partition";
- }
-
- private static Map<TopicPartition, Double> retrievePartitionLoads(
- String srcKafka,
- String dstKafka,
- double byteRateThreshold,
- double messageRateThreshold
- ) {
- log.debug("Retrieving partition loads in preparation for task partition assignment");
-
- Map<ReplicatedTopicPartition, Workload> allWorkloads;
- try {
- allWorkloads = KafkaBasedWorkloadProvider.getAllWorkloads();
- } catch (IllegalStateException e) {
- log.warn(
- "Worker does not appear to be configured with Kafka-based workload provider; "
- + "will not be able to balance task assignment based on partition workloads"
- );
- return Map.of();
- }
-
- Map<TopicPartition, Double> result = allWorkloads.entrySet().stream()
- .filter(e -> srcKafka.equals(e.getKey().srcKafka()))
- .filter(e -> dstKafka.equals(e.getKey().dstKafka()))
- .collect(Collectors.toMap(
- e -> e.getKey().topicPartition(),
- e -> {
- Workload workload = e.getValue();
- double byteRateScale = workload.getBytesPerSecond() / byteRateThreshold;
- double messageRateScale = workload.getMessagesPerSecond() / messageRateThreshold;
- return Math.max(byteRateScale, messageRateScale);
- }
- ));
-
- log.debug("Finished retrieving partition loads in preparation for task partition assignment");
- return result;
- }
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/TaskPartitions.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/TaskPartitions.java
deleted file mode 100644
index 4935a907615..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/TaskPartitions.java
+++ /dev/null
@@ -1,25 +0,0 @@
-package com.uber.data.kafka.connect.ureplicator3;
-
-import org.apache.kafka.common.TopicPartition;
-
-import java.util.HashSet;
-import java.util.Objects;
-import java.util.Set;
-
-public record TaskPartitions(String connectorName, int taskId, Set<TopicPartition> partitions) {
-
- @Override
- public String toString() {
- return connectorName + "-" + taskId;
- }
-
- @Override
- public int hashCode() {
- return Objects.hash(connectorName, taskId);
- }
-
- public TaskPartitions copy() {
- return new TaskPartitions(connectorName, taskId, new HashSet<>(partitions));
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Config.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Config.java
deleted file mode 100644
index 14a415e68c7..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Config.java
+++ /dev/null
@@ -1,251 +0,0 @@
-package com.uber.data.kafka.connect.ureplicator3;
-
-import org.apache.kafka.clients.CommonClientConfigs;
-import org.apache.kafka.clients.admin.Admin;
-import org.apache.kafka.clients.consumer.Consumer;
-import org.apache.kafka.clients.consumer.ConsumerConfig;
-import org.apache.kafka.clients.consumer.KafkaConsumer;
-import org.apache.kafka.common.config.AbstractConfig;
-import org.apache.kafka.common.config.ConfigDef;
-import org.apache.kafka.common.config.ConfigDef.Importance;
-import org.apache.kafka.common.config.ConfigDef.Type;
-import org.apache.kafka.common.config.ConfigException;
-import org.apache.kafka.common.serialization.ByteArrayDeserializer;
-import org.apache.kafka.connect.runtime.ConnectorConfig;
-
-import com.google.common.annotations.VisibleForTesting;
-import com.uber.data.kafka.connect.utils.BrokerResolver;
-import com.uber.data.kafka.exception.ResolverException;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.time.Duration;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.concurrent.TimeUnit;
-
-import static com.uber.data.kafka.connect.utils.UberConnectUtils.createErrorMessage;
-import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
-
-public abstract class UReplicator3Config extends AbstractConfig {
-
- private static final Logger log = LoggerFactory.getLogger(UReplicator3Config.class);
-
- public static final String SOURCE_CLUSTER_CONFIG = "source.cluster";
- private static final String SOURCE_CLUSTER_DOC = "Name of source cluster to consume from (e.g., phx2-dev1)";
-
- public static final String TARGET_CLUSTER_CONFIG = "target.cluster";
- private static final String TARGET_CLUSTER_DOC = "Name of target cluster to produce to (e.g., dca1-dev1)";
-
- public static final String CONSUMER_POLL_TIMEOUT_MS_CONFIG = "consumer.poll.timeout.ms";
- public static final long CONSUMER_POLL_TIMEOUT_MS_DEFAULT = 1000;
- private static final String CONSUMER_POLL_TIMEOUT_MS_DOC = "The timeout to use when invoking Consumer::poll";
-
- public static final String CLIENTS_PREFIX = "all.kafka.clients.";
- public static final String SOURCE_CLIENTS_PREFIX = "source.kafka.clients.";
- public static final String TARGET_CLIENTS_PREFIX = "target.kafka.clients.";
- public static final String SOURCE_CONSUMER_PREFIX = "source.consumer.";
- // This prefix controls "phantom" properties: ones not actually used by our connector code
- // but instead, picked up by our REST extension and translated into properties recognized by
- // the Connect runtime (i.e., "producer.override.<property>")
- public static final String TARGET_PRODUCER_PREFIX = "target.producer.";
- public static final String SOURCE_ADMIN_PREFIX = "source.admin.";
- public static final String TARGET_ADMIN_PREFIX = "target.admin.";
-
- public static final String WORKLOAD_REPORT_INTERVAL_MS_CONFIG = "workload.report.interval.ms";
- private static final long WORKLOAD_REPORT_INTERVAL_MS_DEFAULT = TimeUnit.MINUTES.toMillis(10);
-
- public static final String OFFSET_REPORTER_ENABLED_CONFIG = "offset.reporter.enabled";
- private static final boolean OFFSET_REPORTER_ENABLED_DEFAULT = true;
-
- public static final String SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG = SOURCE_CLIENTS_PREFIX + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG;
- public static final String TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG = TARGET_CLIENTS_PREFIX + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG;
-
- @SuppressWarnings("this-escape")
- protected UReplicator3Config(ConfigDef config, Map<String, ?> props) {
- super(config, props);
- }
-
- protected static ConfigDef baseConfig() {
- return new ConfigDef()
- .define(
- SOURCE_CLUSTER_CONFIG,
- Type.STRING,
- ConfigDef.NO_DEFAULT_VALUE,
- Importance.HIGH,
- SOURCE_CLUSTER_DOC
- ).define(
- TARGET_CLUSTER_CONFIG,
- Type.STRING,
- ConfigDef.NO_DEFAULT_VALUE,
- Importance.HIGH,
- TARGET_CLUSTER_DOC
- ).define(
- CONSUMER_POLL_TIMEOUT_MS_CONFIG,
- Type.LONG,
- CONSUMER_POLL_TIMEOUT_MS_DEFAULT,
- atLeast(1),
- Importance.MEDIUM,
- CONSUMER_POLL_TIMEOUT_MS_DOC
- ).defineInternal(
- WORKLOAD_REPORT_INTERVAL_MS_CONFIG,
- Type.LONG,
- WORKLOAD_REPORT_INTERVAL_MS_DEFAULT,
- Importance.LOW
- ).defineInternal(
- OFFSET_REPORTER_ENABLED_CONFIG,
- Type.BOOLEAN,
- OFFSET_REPORTER_ENABLED_DEFAULT,
- Importance.LOW
- ).defineInternal(
- ConnectorConfig.NAME_CONFIG,
- Type.STRING,
- ConfigDef.NO_DEFAULT_VALUE,
- Importance.LOW
- );
- }
-
- protected abstract String role();
-
- public String sourceCluster() {
- return getString(SOURCE_CLUSTER_CONFIG);
- }
-
- public String targetCluster() {
- return getString(TARGET_CLUSTER_CONFIG);
- }
-
- public Duration consumerPollTimeout() {
- long timeoutMs = getLong(CONSUMER_POLL_TIMEOUT_MS_CONFIG);
- return Duration.ofMillis(timeoutMs);
- }
-
- public String connectorName() {
- return getString(ConnectorConfig.NAME_CONFIG);
- }
-
- public Duration workloadReportInterval() {
- Long durationMs = getLong(WORKLOAD_REPORT_INTERVAL_MS_CONFIG);
- return durationMs == null || durationMs <= 0
- ? null
- : Duration.ofMillis(durationMs);
- }
-
- public boolean offsetReporterEnabled() {
- return getBoolean(OFFSET_REPORTER_ENABLED_CONFIG);
- }
-
- public Admin sourceAdmin() {
- return admin(
- sourceBootstrap(),
- clientId(true, "admin"),
- CLIENTS_PREFIX,
- SOURCE_CLIENTS_PREFIX,
- SOURCE_ADMIN_PREFIX
- );
- }
-
- public Admin targetAdmin() {
- return admin(
- targetBootstrap(),
- clientId(false, "admin"),
- CLIENTS_PREFIX,
- TARGET_CLIENTS_PREFIX,
- TARGET_ADMIN_PREFIX
- );
- }
-
- public Consumer<byte[], byte[]> sourceConsumer() {
- return consumer(
- sourceBootstrap(),
- clientId(true, "consumer"),
- CLIENTS_PREFIX,
- SOURCE_CLIENTS_PREFIX,
- SOURCE_CONSUMER_PREFIX
- );
- }
-
- private String sourceBootstrap() {
- String override = sourceBootstrapOverride();
- if (override != null) {
- log.info("Using overridden bootstrap servers for source cluster: {}", override);
- return override;
- }
-
- return resolveBootstrap(SOURCE_CLUSTER_CONFIG, sourceCluster());
- }
-
- protected String sourceBootstrapOverride() {
- return originalsStrings().get(SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG);
- }
-
- private String targetBootstrap() {
- String override = targetBootstrapOverride();
- if (override != null) {
- log.info("Using overridden bootstrap servers for target cluster: {}", override);
- return override;
- }
-
- return resolveBootstrap(TARGET_CLUSTER_CONFIG, targetCluster());
- }
-
- protected String targetBootstrapOverride() {
- return originalsStrings().get(TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG);
- }
-
- // TODO: Maybe align this with the IDs assigned by the Connect runtime to the
- // clients it instantiates on behalf of connectors and tasks?
- private String clientId(boolean source, String clientType) {
- return "ureplicator3-%s-%s-%s".formatted(
- role(),
- source ? "source" : "target",
- clientType
- );
- }
-
- private Admin admin(String bootstrapServers, String clientId, String... overridesPrefixes) {
- Map<String, Object> adminProps = new HashMap<>(Map.of(
- CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers,
- CommonClientConfigs.CLIENT_ID_CONFIG, clientId
- ));
-
- for (String overridesPrefix : overridesPrefixes)
- adminProps.putAll(originalsWithPrefix(overridesPrefix));
-
- return Admin.create(adminProps);
- }
-
- private Consumer<byte[], byte[]> consumer(String bootstrapServers, String clientId, String... overridesPrefixes) {
- Map<String, Object> consumerProps = new HashMap<>(Map.of(
- CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, bootstrapServers,
- CommonClientConfigs.CLIENT_ID_CONFIG, clientId,
- ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, "earliest",
- ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class,
- ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class
- ));
-
- for (String overridesPrefix : overridesPrefixes)
- consumerProps.putAll(originalsWithPrefix(overridesPrefix));
-
- return new KafkaConsumer<>(consumerProps);
- }
-
- @VisibleForTesting
- static String resolveBootstrap(String prop, String cluster) {
- try {
- return BrokerResolver.resolveBootstrap(cluster);
- } catch (IllegalStateException e) {
- throw new IllegalStateException(
- "Failed to initialize broker resolver. If testing locally, manually specify the bootstrap servers "
- + "for the source and target cluster using the " + SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG + " "
- + "and " + TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG + " properties"
- );
- } catch (ResolverException e) {
- log.error("Failed to resolve bootstrap servers for {}", cluster, e);
- String errorMessage = createErrorMessage("failed to resolve bootstrap servers for cluster", e);
- throw new ConfigException(prop, cluster, errorMessage);
- }
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Connector.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Connector.java
deleted file mode 100644
index a5aba091d08..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Connector.java
+++ /dev/null
@@ -1,449 +0,0 @@
-package com.uber.data.kafka.connect.ureplicator3;
-
-import org.apache.kafka.clients.admin.Admin;
-import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.config.Config;
-import org.apache.kafka.common.config.ConfigDef;
-import org.apache.kafka.common.utils.AppInfoParser;
-import org.apache.kafka.connect.connector.Task;
-import org.apache.kafka.connect.errors.ConnectException;
-import org.apache.kafka.connect.source.SourceConnector;
-
-import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
-import com.uber.data.kafka.connect.distributed.assignment.AbstractAssignor;
-import com.uber.data.kafka.connect.ureplicator3.coordination.AssignedPartition;
-import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStore;
-import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStores;
-import com.uber.data.kafka.connect.utils.ExecutorUtils;
-import com.uber.data.kafka.connect.utils.UberConnectUtils;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.time.Duration;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Collections;
-import java.util.Comparator;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.LinkedHashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Objects;
-import java.util.Set;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.Future;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
-import java.util.stream.Collectors;
-import java.util.stream.IntStream;
-
-import static com.uber.data.kafka.connect.utils.UberConnectUtils.awaitAll;
-import static com.uber.data.kafka.connect.utils.UberConnectUtils.diff;
-
-public class UReplicator3Connector extends SourceConnector {
-
- private static final Logger log = LoggerFactory.getLogger(UReplicator3Connector.class);
-
- public static final String SKIP_CONFIG_VALIDATION = "skip.config.validation";
-
- private UReplicator3ConnectorConfig config;
- private AbstractAssignor<TaskPartitions, TopicPartition> assignor;
- private ScheduledExecutorService executor;
- private AssignmentStore assignmentStore;
-
- // The values for these maps may be changed at any time, but their contents of any single map will never be mutated
- private volatile Map<String, Integer> sourcePartitionCounts;
- private volatile Map<String, Integer> targetPartitionCounts;
- private volatile int clusterSize;
-
- @Override
- public ConfigDef config() {
- return UReplicator3ConnectorConfig.config();
- }
-
- @Override
- public String version() {
- return AppInfoParser.getVersion();
- }
-
- @Override
- public Class<? extends Task> taskClass() {
- return UReplicator3Task.class;
- }
-
- @Override
- public void start(Map<String, String> props) {
- log.info("Starting connector");
-
- config = new UReplicator3ConnectorConfig(props);
-
- assignmentStore = AssignmentStores.get();
-
- assignor = new TaskPartitionAssignor(
- config.sourceCluster(),
- config.targetCluster(),
- 0.1, // This should never really be used so we hardcode it for now
- config.taskByteRateThreshold(),
- config.taskMessageRateThreshold()
- );
-
- // Synchronously collect partition counts the first time around
- // to make sure they're available before we generate task configs
- doRefreshClusterMetadata();
-
- long refreshIntervalMs = config.refreshIntervalMs();
- executor = ExecutorUtils.singleThreadedScheduled("executor-" + config.connectorName());
- executor.scheduleWithFixedDelay(this::refreshClusterMetadata, refreshIntervalMs, refreshIntervalMs, TimeUnit.MILLISECONDS);
- // TODO: Should be two separate intervals for these operations
- executor.scheduleWithFixedDelay(this::refreshAssignments, refreshIntervalMs, refreshIntervalMs, TimeUnit.MILLISECONDS);
-
- log.info("Finished starting connector");
- }
-
- @Override
- public List<Map<String, String>> taskConfigs(int maxTasks) {
- log.info("Generating task configs");
-
- Map<String, Integer> sourcePartitionCountsSnapshot = sourcePartitionCounts;
- Map<String, Integer> targetPartitionCountsSnapshot = targetPartitionCounts;
-
- List<TopicPartition> topicPartitions = replicatedPartitions(sourcePartitionCountsSnapshot, targetPartitionCountsSnapshot);
-
- clusterSize = context.uberClusterSize();
- int numTasks = numTasks(maxTasks, topicPartitions.size(), clusterSize);
-
- List<Map<String, String>> result = IntStream.range(0, numTasks)
- .mapToObj(i -> config.taskConfig(i))
- .collect(Collectors.toList());
-
- // Make sure partition assignments are updated in line with the expected number of tasks
- // Run on a separate thread since we really, really don't want to block too long in this method
- // (it's invoked from the Connect worker's herder thread, which is basically the central nervous
- // system of the control plane)
- executor.execute(() -> refreshAssignments(numTasks));
-
- log.info("Finished generating task configs");
-
- return result;
- }
-
- @Override
- public void stop() {
- log.info("Stopping connector");
- ExecutorUtils.shutdown(executor, "connector executor", 2, TimeUnit.SECONDS);
- log.info("Finished stopping connector");
- }
-
- @Override
- public Config validate(Map<String, String> props) {
- // Breakglass logic: allow the connector config to be updated even if it's failing
- // Should only be used in emergencies
- String skipValidationValue = props.get(SKIP_CONFIG_VALIDATION);
- if (Boolean.parseBoolean(skipValidationValue)) {
- log.warn("Skipping config validation since {} is set to {}", SKIP_CONFIG_VALIDATION, skipValidationValue);
- return new Config(List.of());
- }
-
- Config result = super.validate(props);
- UReplicator3ConnectorConfig config;
-
- try {
- config = new UReplicator3ConnectorConfig(props);
- } catch (Exception e) {
- boolean alreadyHasErrors = result.configValues().stream()
- .anyMatch(configValue -> !configValue.errorMessages().isEmpty());
-
- if (alreadyHasErrors) {
- // We can just log the error; preflight validation will fail without any intervention on our part
- log.debug("Unable to fully validate connector config", e);
- return result;
- } else {
- // This shouldn't really happen, make sure to log at error level if it does so we can investigate
- log.error("Unable to fully validate connector config", e);
- // Force preflight validation to fail
- throw new ConnectException("Unable to fully validate connector config", e);
- }
- }
-
- config.validate(result);
-
- return result;
- }
-
- @Override
- public boolean alterOffsets(Map<String, String> connectorConfig, Map<Map<String, ?>, Map<String, ?>> offsets) {
- return true;
- }
-
- private void refreshClusterMetadata() {
- log.info("Refreshing cluster metadata");
- try {
- doRefreshClusterMetadata();
- log.info("Finished refreshing cluster metadata");
- } catch (Exception e) {
- // TODO: Emit metric
- log.error("Failed to refresh cluster metadata", e);
- }
- }
-
- private void doRefreshClusterMetadata() {
- Map<String, Integer> oldSourcePartitionCounts = sourcePartitionCounts;
- try (Admin sourceAdmin = config.sourceAdmin()) {
- sourcePartitionCounts = partitionCounts(sourceAdmin, config.sourceTopics(), "source");
- }
-
- Map<String, Integer> oldTargetPartitionCounts = targetPartitionCounts;
- try (Admin targetAdmin = config.targetAdmin()) {
- targetPartitionCounts = partitionCounts(targetAdmin, config.targetTopics(), "target");
- }
-
- if (oldSourcePartitionCounts != null && !oldSourcePartitionCounts.equals(sourcePartitionCounts)) {
- context.requestTaskReconfiguration();
- }
- if (oldTargetPartitionCounts != null && !oldTargetPartitionCounts.equals(targetPartitionCounts)) {
- context.requestTaskReconfiguration();
- }
-
- int currentClusterSize = context.uberClusterSize();
- if (currentClusterSize > 0 && clusterSize > 0) {
- double clusterSizeRatio = 100.0 * Math.abs(clusterSize - currentClusterSize) / clusterSize;
- log.trace("Current cluster size: {}; last-used cluster size: {}", currentClusterSize, clusterSize);
- double reconfigurationThreshold = config.clusterSizeReconfigurationThreshold();
- if (clusterSizeRatio >= reconfigurationThreshold) {
- log.debug(
- "Ratio of current cluster size to last-used cluster size exceeds reconfiguration threshold ({}); generating new task configs",
- reconfigurationThreshold
- );
- context.requestTaskReconfiguration();
- } else if (currentClusterSize == clusterSize) {
- log.trace("No change in cluster size");
- } else {
- log.trace("Ratio of current cluster size to last-used cluster size is within reconfiguration threshold; will not generate new task configs");
- }
- }
- }
-
- private static Map<String, Integer> partitionCounts(Admin admin, Set<String> topics, String clusterType) {
- try {
- return UberConnectUtils.listPartitions(topics, admin, 1, TimeUnit.MINUTES);
- } catch (Exception e) {
- throw new ConnectException("Failed to list partition counts for topics on " + clusterType + " cluster", e);
- }
- }
-
- private List<TopicPartition> replicatedPartitions(
- Map<String, Integer> sourcePartitionCountsSnapshot,
- Map<String, Integer> targetPartitionCountsSnapshot
- ) {
- Comparator<Map.Entry<String, Integer>> topicPartitionCountEntryComparator =
- Map.Entry.<String, Integer>comparingByValue()
- .reversed()
- .thenComparing(
- Map.Entry.comparingByKey()
- );
-
- return sourcePartitionCountsSnapshot.entrySet().stream()
- .filter(e -> {
- String sourceTopic = e.getKey();
- String targetTopic = config.targetTopic(sourceTopic);
-
- if (targetPartitionCountsSnapshot.containsKey(targetTopic)) {
- return true;
- } else {
- log.warn(
- "Not assigning topic {} to tasks since its target partition count is unknown "
- + "(topic may have been deleted)",
- sourceTopic);
- return false;
- }
- })
- .sorted(topicPartitionCountEntryComparator)
- .flatMap(e -> IntStream.range(0, e.getValue())
- .mapToObj(partition -> new TopicPartition(e.getKey(), partition))
- )
- .toList();
- }
-
- private int numTasks(int maxTasks, int numTopicPartitions, int clusterSize) {
- int numTasks = Math.min(maxTasks, numTopicPartitions);
- if (clusterSize > 0) {
- int workerTaskThreshold = config.workerTaskLimit();
- if (workerTaskThreshold > 0) {
- int maxWorkerTasks = workerTaskThreshold * clusterSize;
- log.debug("Total cluster-based task limit: {} (cluster size: {}, tasks per worker: {})",
- maxWorkerTasks, clusterSize, workerTaskThreshold);
- numTasks = Math.min(maxWorkerTasks, numTasks);
- } else {
- log.debug("Ignoring cluster size since worker task limit is disabled");
- }
- } else {
- log.warn("Cluster size not available; will generate at most {} tasks", numTasks);
- }
- return numTasks;
- }
-
- private void refreshAssignments() {
- refreshAssignments(context.uberTaskConfigs().size());
- }
-
- private void refreshAssignments(int numTasks) {
- log.info("Refreshing task partition assignments");
- try {
- doRefreshAssignments(numTasks);
- log.info("Finished refreshing task partition assignments");
- } catch (Exception e) {
- // TODO: Emit metric
- log.error("Failed to refresh task partition assignments", e);
- }
- }
-
- private void doRefreshAssignments(int numTasks) {
- Map<String, Integer> sourcePartitionCountsSnapshot = sourcePartitionCounts;
- Map<String, Integer> targetPartitionCountsSnapshot = targetPartitionCounts;
-
- List<TopicPartition> topicPartitions = replicatedPartitions(sourcePartitionCountsSnapshot, targetPartitionCountsSnapshot);
-
- Map<Integer, ? extends Collection<AssignedPartition>> existingStoredAssignments = assignmentStore.getAssignments(
- config.connectorName(),
- config.sourceCluster(),
- config.targetCluster()
- );
-
- List<TaskPartitions> existingAssignments = existingAssignments(existingStoredAssignments);
- Collection<TopicPartition> unassignedPartitions = unassignedPartitions(numTasks, existingAssignments, topicPartitions);
-
- List<TaskPartitions> newAssignments = IntStream.range(0, numTasks)
- .mapToObj(i -> i < existingAssignments.size()
- ? existingAssignments.get(i).copy()
- : new TaskPartitions(config.connectorName(), i, new HashSet<>())
- ).toList();
- assignor.assign(newAssignments, unassignedPartitions);
-
- Map<Integer, Collection<AssignedPartition>> newStoredAssignments = newStoredAssignments(
- targetPartitionCountsSnapshot,
- newAssignments
- );
-
- Collection<Future<?>> storeFutures = new ArrayList<>();
-
- Set<TopicPartition> toExplicitlyRevoke = existingAssignments.stream()
- .map(TaskPartitions::partitions)
- .flatMap(Collection::stream)
- .collect(Collectors.toCollection(HashSet::new));
- toExplicitlyRevoke.removeAll(new HashSet<>(topicPartitions));
-
- if (!toExplicitlyRevoke.isEmpty()) {
- log.debug("Found {} partitions to explicitly revoke", toExplicitlyRevoke.size());
- } else {
- log.trace("No partitions to explicitly revoke");
- }
-
- toExplicitlyRevoke.stream()
- .map(this::replicatedTopicPartition)
- .map(rtp -> assignmentStore.revokePartition(config.connectorName(), rtp))
- .forEach(storeFutures::add);
-
- int tasksWithNewPartitions = 0;
- for (int i = 0; i < newAssignments.size(); i++) {
- final int taskId = i;
- Collection<AssignedPartition> existingPartitions = existingStoredAssignments.get(taskId);
- Collection<AssignedPartition> newPartitions = newStoredAssignments.get(taskId);
-
- Set<AssignedPartition> added = newPartitions != null ? diff(newPartitions, existingPartitions) : Set.of();
-
- if (!added.isEmpty()) {
- log.debug("Found {} partitions to add for task {}-{}", added.size(), config.connectorName(), taskId);
- tasksWithNewPartitions++;
- } else {
- log.trace("No partitions to add for task {}-{}", config.connectorName(), taskId);
- }
-
- added.stream()
- .map(ap -> {
- ReplicatedTopicPartition rtp = replicatedTopicPartition(ap.source());
- return assignmentStore.assignPartition(config.connectorName(), rtp, taskId, ap.dstTopic(), ap.dstPartition());
- })
- .forEach(storeFutures::add);
- }
-
- if (tasksWithNewPartitions > 0) {
- log.debug("Assigned new partitions to {} total tasks", tasksWithNewPartitions);
- } else {
- log.debug("Assigned no new partitions to any tasks");
- }
-
- // We may want to make this duration configurable in the future
- // TODO: Should also add metrics on write failures
- try {
- awaitAll(storeFutures, Duration.ofMinutes(1));
- } catch (InterruptedException | TimeoutException | ExecutionException e) {
- log.error("Failed to publish task partition assignment updates", e);
- }
- }
-
- private List<TaskPartitions> existingAssignments(Map<Integer, ? extends Collection<AssignedPartition>> storedAssignments) {
- if (storedAssignments.isEmpty())
- return List.of();
-
- int highestTask = Collections.max(storedAssignments.keySet());
- return IntStream.rangeClosed(0, highestTask) // rangeClosed because tasks are zero-indexed
- .mapToObj(i -> {
- Collection<AssignedPartition> partitions = storedAssignments.get(i);
- Set<TopicPartition> partitionsSet = new HashSet<>();
- if (partitions != null) {
- partitions.stream()
- .map(AssignedPartition::source)
- .forEach(partitionsSet::add);
- }
- return new TaskPartitions(config.connectorName(), i, partitionsSet);
- }).toList();
- }
-
- private Set<TopicPartition> unassignedPartitions(int numTasks, List<TaskPartitions> taskPartitions, Collection<TopicPartition> topicPartitions) {
- Set<TopicPartition> result = new LinkedHashSet<>(topicPartitions);
-
- for (int i = 0; i < numTasks && i < taskPartitions.size(); i++) {
- result.removeAll(taskPartitions.get(i).partitions());
- }
-
- return result;
- }
-
- private ReplicatedTopicPartition replicatedTopicPartition(TopicPartition tp) {
- return new ReplicatedTopicPartition(config.sourceCluster(), config.targetCluster(), tp);
- }
-
- private Map<Integer, Collection<AssignedPartition>> newStoredAssignments(
- Map<String, Integer> targetPartitionCounts,
- List<TaskPartitions> newAssignments
- ) {
- Map<Integer, Collection<AssignedPartition>> result = new HashMap<>();
- for (int i = 0; i < newAssignments.size(); i++) {
- TaskPartitions newAssignment = newAssignments.get(i);
- Collection<AssignedPartition> assignment = new HashSet<>();
- for (TopicPartition source : newAssignment.partitions()) {
- String dstTopic = config.targetTopic(source.topic());
- Integer targetPartitionCount = targetPartitionCounts.get(dstTopic);
- if (targetPartitionCount == null) {
- log.warn("Could not find partition count for topic {} on target cluster; will not replicate topic partition {} (this should never happen)",
- dstTopic, source);
- continue;
- }
- int dstPartition = source.partition() % targetPartitionCount;
- AssignedPartition assignedPartition = new AssignedPartition(
- source.topic(),
- source.partition(),
- !Objects.equals(dstTopic, source.topic()) ? dstTopic : null,
- dstPartition != source.partition() ? dstPartition : null
- );
- assignment.add(assignedPartition);
- }
- result.put(i, assignment);
- }
- return result;
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3ConnectorConfig.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3ConnectorConfig.java
deleted file mode 100644
index 3f30cec9e9e..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3ConnectorConfig.java
+++ /dev/null
@@ -1,385 +0,0 @@
-package com.uber.data.kafka.connect.ureplicator3;
-
-import org.apache.kafka.clients.CommonClientConfigs;
-import org.apache.kafka.clients.admin.Admin;
-import org.apache.kafka.clients.consumer.Consumer;
-import org.apache.kafka.common.config.Config;
-import org.apache.kafka.common.config.ConfigDef;
-import org.apache.kafka.common.config.ConfigDef.Importance;
-import org.apache.kafka.common.config.ConfigDef.Type;
-import org.apache.kafka.common.config.ConfigException;
-
-import com.google.common.annotations.VisibleForTesting;
-import com.uber.data.kafka.connect.workload.kafka.KafkaBasedWorkloadProvider;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Map;
-import java.util.Set;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
-import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.stream.Collectors;
-import java.util.stream.Stream;
-
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3TaskConfig.TASK_ID_CONFIG;
-import static com.uber.data.kafka.connect.utils.UberConnectUtils.addError;
-import static com.uber.data.kafka.connect.utils.UberConnectUtils.diff;
-import static com.uber.data.kafka.connect.utils.UberConnectUtils.listTopics;
-import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
-import static org.apache.kafka.connect.runtime.ConnectorConfig.TASKS_MAX_CONFIG;
-
-public class UReplicator3ConnectorConfig extends UReplicator3Config {
-
- private static final Logger log = LoggerFactory.getLogger(UReplicator3ConnectorConfig.class);
-
- public static final String TOPICS_CONFIG = "topics";
- private static final String TOPICS_DOC = "Comma-separated list of topics to replicate";
-
- public static final String WORKER_TASK_LIMIT_CONFIG = "worker.task.limit";
- public static final int WORKER_TASK_LIMIT_DEFAULT = 2;
- private static final String WORKER_TASK_LIMIT_DOC = "Maximum number of tasks to generate per worker. "
- + "To allow any number of tasks per worker, set to a non-positive value";
-
- public static final String TASK_BYTE_RATE_THRESHOLD_CONFIG = "task.byte.rate.threshold";
- public static final double TASK_BYTE_RATE_THRESHOLD_DEFAULT = 70 * 1024 * 1024;
- private static final String TASK_BYTE_RATE_THRESHOLD_DOC = "Maximum total byte rate for partitions assigned "
- + "to each task before attempting to revoke partitions and redistribute to other tasks";
-
- public static final String TASK_MESSAGE_RATE_THRESHOLD_CONFIG = "task.message.rate.threshold";
- public static final double TASK_MESSAGE_RATE_THRESHOLD_DEFAULT = 70_000;
- private static final String TASK_MESSAGE_RATE_THRESHOLD_DOC = "Maximum total message rate for partitions assigned "
- + "to each task before attempting to revoke partitions and redistribute to other workers";
-
- public static final String REFRESH_INTERVAL_MS = "refresh.interval.ms";
- public static final long REFRESH_INTERVAL_MS_DEFAULT = TimeUnit.MINUTES.toMillis(1);
- private static final String REFRESH_INTERVAL_MS_DOC = "How often to check for changes in the partition counts of source and target topics, "
- + "the size of the Kafka Connect cluster, and imbalances in the partitions assigned to each task, "
- + "possibly generating new task configs and/or reassigning partitions in response";
-
- public static final String CLUSTER_SIZE_RECONFIGURATION_THRESHOLD_CONFIG = "cluster.size.reconfiguration.threshold";
- public static final double CLUSTER_SIZE_RECONFIGURATION_THRESHOLD_DEFAULT = 10;
- private static final String CLUSTER_SIZE_RECONFIGURATION_THRESHOLD_DOC = "Minimum percentage difference "
- + "in cluster size (i.e., number of workers) before new task configs are generated";
-
- public static final String TOPIC_RENAME_PREFIX = "topic.rename.";
-
- private final Map<String, String> topicRenames;
- private final Set<String> sourceTopics;
- private final Set<String> targetTopics;
-
- @SuppressWarnings("this-escape")
- public UReplicator3ConnectorConfig(Map<String, ?> props) {
- super(config(), props);
- this.topicRenames = Collections.unmodifiableMap(parseTopicRenames(originalsStrings(), null));
- this.sourceTopics = Set.copyOf(getList(TOPICS_CONFIG));
- this.targetTopics = Collections.unmodifiableSet(computeTargetTopics(sourceTopics, topicRenames));
- }
-
- protected static ConfigDef config() {
- return baseConfig()
- .define(
- TOPICS_CONFIG,
- Type.LIST,
- ConfigDef.NO_DEFAULT_VALUE,
- ConfigDef.ValidList.anyNonDuplicateValues(false, false),
- Importance.HIGH,
- TOPICS_DOC
- ).define(
- WORKER_TASK_LIMIT_CONFIG,
- Type.INT,
- WORKER_TASK_LIMIT_DEFAULT,
- Importance.MEDIUM,
- WORKER_TASK_LIMIT_DOC
- ).define(
- TASK_BYTE_RATE_THRESHOLD_CONFIG,
- Type.DOUBLE,
- TASK_BYTE_RATE_THRESHOLD_DEFAULT,
- atLeast(1),
- Importance.MEDIUM,
- TASK_BYTE_RATE_THRESHOLD_DOC
- ).define(
- TASK_MESSAGE_RATE_THRESHOLD_CONFIG,
- Type.DOUBLE,
- TASK_MESSAGE_RATE_THRESHOLD_DEFAULT,
- atLeast(1),
- Importance.MEDIUM,
- TASK_MESSAGE_RATE_THRESHOLD_DOC
- ).define(
- REFRESH_INTERVAL_MS,
- Type.LONG,
- REFRESH_INTERVAL_MS_DEFAULT,
- atLeast(1),
- Importance.MEDIUM,
- REFRESH_INTERVAL_MS_DOC
- ).define(
- CLUSTER_SIZE_RECONFIGURATION_THRESHOLD_CONFIG,
- Type.DOUBLE,
- CLUSTER_SIZE_RECONFIGURATION_THRESHOLD_DEFAULT,
- atLeast(1),
- Importance.LOW,
- CLUSTER_SIZE_RECONFIGURATION_THRESHOLD_DOC
- );
- }
-
- @Override
- protected String role() {
- return "connector-" + connectorName();
- }
-
- public void validate(Config configValidationResult) {
- parseTopicRenames(originalsStrings(), configValidationResult);
-
- if (!validateBootstrapServersOverrides(configValidationResult)) {
- // Don't try to validate more since we might end up contacting the wrong Kafka clusters
- return;
- }
-
- validateCluster(configValidationResult, true);
- validateCluster(configValidationResult, false);
-
- validateWorkloadPublishing(configValidationResult);
- }
-
- public Set<String> sourceTopics() {
- return sourceTopics;
- }
-
- public Set<String> targetTopics() {
- return targetTopics;
- }
-
- public String targetTopic(String sourceTopic) {
- return topicRenames.getOrDefault(sourceTopic, sourceTopic);
- }
-
- public int workerTaskLimit() {
- return getInt(WORKER_TASK_LIMIT_CONFIG);
- }
-
- public double taskByteRateThreshold() {
- return getDouble(TASK_BYTE_RATE_THRESHOLD_CONFIG);
- }
-
- public double taskMessageRateThreshold() {
- return getDouble(TASK_MESSAGE_RATE_THRESHOLD_CONFIG);
- }
-
- public long refreshIntervalMs() {
- return getLong(REFRESH_INTERVAL_MS);
- }
-
- public double clusterSizeReconfigurationThreshold() {
- return getDouble(CLUSTER_SIZE_RECONFIGURATION_THRESHOLD_CONFIG);
- }
-
- public Map<String, String> taskConfig(int taskId) {
- Map<String, String> result = new HashMap<>(originalsStrings());
-
- // If we don't do this, every time a topic is added or removed, a task reconfiguration will be triggered
- result.remove(TOPICS_CONFIG);
- // Same for topic renames
- result.keySet().removeIf(key -> key.startsWith(TOPIC_RENAME_PREFIX));
- // Same for horizontal scaling properties that are only used by the connector and not its tasks
- result.remove(TASKS_MAX_CONFIG);
- result.remove(WORKER_TASK_LIMIT_CONFIG);
- result.remove(CLUSTER_SIZE_RECONFIGURATION_THRESHOLD_CONFIG);
- // And other properties that are also not used by tasks
- result.remove(REFRESH_INTERVAL_MS);
- result.remove(TASK_BYTE_RATE_THRESHOLD_CONFIG);
- result.remove(TASK_MESSAGE_RATE_THRESHOLD_CONFIG);
-
- result.put(TASK_ID_CONFIG, Integer.toString(taskId));
-
- return result;
- }
-
- private boolean validateBootstrapServersOverrides(Config configValidationResult) {
- AtomicBoolean result = new AtomicBoolean(true);
- Stream.of(SOURCE_CONSUMER_PREFIX, SOURCE_ADMIN_PREFIX, TARGET_ADMIN_PREFIX)
- .map(prefix -> prefix + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG)
- .forEach(prop -> {
- Object forbiddenOverride = originals().get(prop);
- if (forbiddenOverride != null) {
- String errorMessage = "cannot override bootstrap servers for individual clients; instead, use the "
- + SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG + " and/or "
- + TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG + " properties";
- addError(configValidationResult, prop, forbiddenOverride, errorMessage);
- result.set(false);
- }
- });
- return result.get();
- }
-
- private void validateCluster(Config configValidationResult, boolean source) {
- String sourceTarget = source ? "source" : "target";
- String bootstrapOverride = source ? sourceBootstrapOverride() : targetBootstrapOverride();
-
- String prop;
- String value;
- if (bootstrapOverride != null) {
- prop = source ? SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG : TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG;
- value = bootstrapOverride;
- } else {
- String clusterName = source ? sourceCluster() : targetCluster();
- prop = source ? SOURCE_CLUSTER_CONFIG : TARGET_CLUSTER_CONFIG;
- try {
- resolveBootstrap(prop, clusterName);
- } catch (ConfigException e) {
- addError(configValidationResult, prop, clusterName, "failed to resolve bootstrap servers for " + sourceTarget + " cluster", e);
- return;
- }
- value = clusterName;
- }
-
- boolean clusterReachable;
- Set<String> topics = source ? sourceTopics() : targetTopics();
- try (Admin admin = source ? sourceAdmin() : targetAdmin()) {
- clusterReachable = true;
- validateTopics(configValidationResult, topics, source, admin);
- } catch (Exception e) {
- clusterReachable = false;
- log.error("Failed to create admin client for {} cluster", sourceTarget, e);
- addError(configValidationResult, prop, value, "failed to create admin client for " + sourceTarget + " cluster", e);
- }
-
- // Only bother validating the source consumer if the Kafka cluster can actually be reached
- if (clusterReachable && source) {
- validateSourceConsumer(configValidationResult);
- }
- }
-
- private void validateTopics(Config configValidationResult, Set<String> configuredTopics, boolean source, Admin admin) {
- String sourceTarget = source ? "source" : "target";
-
- Set<String> presentTopics;
- try {
- presentTopics = listTopics(admin, 1, TimeUnit.MINUTES);
- } catch (InterruptedException | TimeoutException | ExecutionException e) {
- log.error("Failed to list topics on {} cluster", sourceTarget, e);
-
- String clusterProp = source ? SOURCE_CLUSTER_CONFIG : TARGET_CLUSTER_CONFIG;
- addError(configValidationResult, clusterProp, null, "failed to list topics on " + sourceTarget + " cluster", e);
-
- return;
- }
-
- Set<String> missingTopics = diff(configuredTopics, presentTopics);
-
- if (!missingTopics.isEmpty()) {
- if (source) {
- String errorMessage = "the following topics are missing on the source cluster: " + String.join(", ", missingTopics);
- addError(configValidationResult, TOPICS_CONFIG, null, errorMessage);
- } else {
- Map<String, Set<String>> reversedTopicRenames = topicRenames.entrySet().stream()
- .collect(Collectors.groupingBy(
- Map.Entry::getValue,
- Collectors.mapping(Map.Entry::getKey, Collectors.toSet()
- )));
- Set<String> missingNonRenamedTopics = new HashSet<>();
- for (String targetTopic : missingTopics) {
- Set<String> renamedSourceTopics = reversedTopicRenames.get(targetTopic);
- if (renamedSourceTopics != null) {
- renamedSourceTopics.forEach(sourceTopic -> {
- // If a missing target topic is specified via a topic rename, we should raise that to the user by
- // attaching an error message to the topic.rename.<source_topic> property
- addError(
- configValidationResult,
- TOPIC_RENAME_PREFIX + sourceTopic,
- targetTopic,
- "renamed topic is missing on the target cluster"
- );
- });
- }
-
- if (sourceTopics.contains(targetTopic) && !topicRenames.containsKey(targetTopic))
- missingNonRenamedTopics.add(targetTopic);
- }
- if (!missingNonRenamedTopics.isEmpty()) {
- String errorMessage = "the following topics are missing on the target cluster: " + String.join(", ", missingNonRenamedTopics);
- addError(configValidationResult, TOPICS_CONFIG, null, errorMessage);
- }
- }
- }
- }
-
- private void validateSourceConsumer(Config configValidationResult) {
- try (Consumer<?, ?> consumer = sourceConsumer()) {
- } catch (Exception e) {
- // This error could be caused by any number of custom consumer properties and we don't really
- // have a way of knowing which one(s) caused the issue, so we just report the error with a detailed
- // message on the source cluster property and hope that there's enough info there for the user to
- // pinpoint the problem
- addError(configValidationResult, SOURCE_CLUSTER_CONFIG, sourceCluster(), "failed to create consumer for source cluster", e);
- }
- }
-
- private void validateWorkloadPublishing(Config configValidationResult) {
- if (workloadReportInterval() == null) {
- log.debug("Workload publishing is disabled; skipping validation");
- return;
- }
-
- try {
- KafkaBasedWorkloadProvider.assertStoreInitialized();
- } catch (IllegalStateException e) {
- log.debug("Kafka workload store does not appear to be initialized", e);
- String errorMessage = "This worker does not appear to be configured to use the Kafka-based workload provider, "
- + "which means that uReplicator3's workload publishing feature must be disabled. Set the "
- + WORKLOAD_REPORT_INTERVAL_MS_CONFIG + " property to zero or a negative value to disable, or reconfigure "
- + "the workers in your Kafka Connect cluster to use the Kafka-based workload provider";
-
- addError(configValidationResult, WORKLOAD_REPORT_INTERVAL_MS_CONFIG, null, errorMessage);
- }
- }
-
- @VisibleForTesting
- static Map<String, String> parseTopicRenames(Map<String, String> originals, Config configValidationResult) {
- Map<String, String> result = new HashMap<>();
- originals.forEach((key, value) -> {
- if (!key.startsWith(TOPIC_RENAME_PREFIX))
- return;
-
- String sourceTopic = key.substring(TOPIC_RENAME_PREFIX.length());
-
- if (sourceTopic.isBlank()) {
- if (configValidationResult != null) {
- addError(configValidationResult, key, value, "source topic is empty or blank");
- } else {
- log.warn("Ignoring invalid property {} with value {}", key, value);
- }
- return;
- }
-
- if (value == null || value.isBlank()) {
- if (configValidationResult != null) {
- addError(configValidationResult, key, value, "target topic is null, empty, or blank");
- } else {
- log.warn("Ignoring property {} with invalid value {}", key, value);
- }
- return;
- }
-
- result.putIfAbsent(sourceTopic, value);
- });
-
- return result;
- }
-
- private static Set<String> computeTargetTopics(Set<String> sourceTopics, Map<String, String> topicRenames) {
- return sourceTopics.stream()
- .map(t -> topicRenames.getOrDefault(t, t))
- .collect(Collectors.toSet());
- }
-
- // Run with `./gradlew :connect:uber:genUReplicator3Config`
- public static void main(String[] args) {
- System.out.println(config().toHtml());
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Task.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Task.java
deleted file mode 100644
index df70b0c1442..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3Task.java
+++ /dev/null
@@ -1,367 +0,0 @@
-package com.uber.data.kafka.connect.ureplicator3;
-
-import org.apache.kafka.clients.consumer.Consumer;
-import org.apache.kafka.clients.consumer.ConsumerRecord;
-import org.apache.kafka.clients.consumer.ConsumerRecords;
-import org.apache.kafka.clients.producer.RecordMetadata;
-import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.header.Header;
-import org.apache.kafka.common.utils.AppInfoParser;
-import org.apache.kafka.common.utils.Time;
-import org.apache.kafka.common.utils.Utils;
-import org.apache.kafka.connect.data.Schema;
-import org.apache.kafka.connect.header.ConnectHeaders;
-import org.apache.kafka.connect.header.Headers;
-import org.apache.kafka.connect.source.SourceRecord;
-import org.apache.kafka.connect.source.SourceTask;
-
-import com.uber.data.kafka.connect.aakafkaoffsetmgmt.OffsetSnapshotClient;
-import com.uber.data.kafka.connect.aakafkaoffsetmgmt.OffsetSnapshotReporter;
-import com.uber.data.kafka.connect.ureplicator3.coordination.AssignedPartition;
-import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStore;
-import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStores;
-import com.uber.data.kafka.connect.workload.kafka.KafkaBasedWorkloadProvider;
-import com.uber.data.kafka.connect.workload.tracking.WorkloadTrackingProcessor;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.time.Duration;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.LinkedHashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.function.Predicate;
-import java.util.stream.Collectors;
-
-public class UReplicator3Task extends SourceTask {
-
- private static final Logger log = LoggerFactory.getLogger(UReplicator3Task.class);
-
- public static final String SOURCE_PARTITION_SOURCE_CLUSTER = "source.cluster";
- public static final String SOURCE_PARTITION_TARGET_CLUSTER = "target.cluster";
- public static final String SOURCE_PARTITION_KAFKA_TOPIC = "topic";
- public static final String SOURCE_PARTITION_KAFKA_PARTITION = "partition";
- public static final String SOURCE_OFFSET_KAFKA_OFFSET = "offset";
-
- private Time time;
- private UReplicator3TaskConfig config;
- private Consumer<byte[], byte[]> sourceConsumer;
- private WorkloadTrackingProcessor workloadTracker;
- private AssignmentStore partitionAssignmentStore;
- private volatile boolean assignmentChanged;
- private Map<TopicPartition, TopicPartition> overrides;
- private OffsetSnapshotReporter offsetSnapshotReporter;
-
- @Override
- public void start(Map<String, String> props) {
- log.info("Starting task");
-
- time = Time.SYSTEM;
-
- config = new UReplicator3TaskConfig(props);
-
- overrides = Map.of();
-
- partitionAssignmentStore = AssignmentStores.get();
- partitionAssignmentStore.addListener(config.connectorName(), config.taskId(), this::onAssignmentUpdate);
-
- sourceConsumer = config.sourceConsumer();
-
- // Force us to update the consumer assignment before polling it for the first time,
- // since right now it's not assigned anything
- assignmentChanged = true;
-
- initializeWorkloadTracker();
- initializeOffsetSnapshotReporter();
-
- log.info("Finished starting task");
- }
-
- @Override
- public List<SourceRecord> poll() {
- log.trace("Starting to poll for records"); // This gets called a lot; can't put it at INFO like the others
-
- // Consumer isn't thread-safe, so we make sure to update its partition assignment on the same thread
- // that we use to poll it for records
- if (assignmentChanged) {
- // Reset the flag before updating the assignment; that way, any subsequent assignment updates
- // (which will be picked up on a separate thread and may be observed concurrently with our consumer update)
- // won't be skipped
- assignmentChanged = false;
- updateAssignment();
- }
-
- if (sourceConsumer.assignment().isEmpty()) {
- log.trace("Task is not assigned any partitions; will idle and return empty list of records");
- // Idle here so that we don't enter a tight loop where no (real) work is done. We can't block indefinitely though,
- // since the task may be scheduled for shutdown and the same thread that calls this method is also responsible
- // for calling Task::stop
- // We may want to make this configurable in the future but for now this should be fine
- time.sleep(100);
- return List.of();
- }
-
- ConsumerRecords<byte[], byte[]> records = sourceConsumer.poll(config.consumerPollTimeout());
-
- List<SourceRecord> result = new ArrayList<>(records.count());
- for (ConsumerRecord<byte[], byte[]> record : records) {
- SourceRecord sourceRecord = convertRecord(record);
-
- if (workloadTracker != null)
- workloadTracker.track(sourceRecord);
-
- result.add(sourceRecord);
- }
-
- log.trace("Finished polling for records");
- return result;
- }
-
- @Override
- public void stop() {
- log.info("Stopping task");
- Utils.closeQuietly(sourceConsumer, "source consumer");
- Utils.closeQuietly(workloadTracker, "workload tracker");
- if (offsetSnapshotReporter != null)
- Utils.closeQuietly(offsetSnapshotReporter::shutdown, "offset snapshot reporter");
- log.info("Finished stopping task");
- }
-
- @Override
- public String version() {
- return AppInfoParser.getVersion();
- }
-
- @Override
- public void commitRecord(SourceRecord record, RecordMetadata metadata) {
- if (offsetSnapshotReporter != null) {
- TopicPartition dstTopicPartition = new TopicPartition(metadata.topic(), metadata.partition());
- if (!metadata.hasOffset()) {
- log.warn("Record from topic partition {} does not have offset; cannot report offsets for it", dstTopicPartition);
- return;
- }
- long dstOffset = metadata.offset();
-
- TopicPartition srcTopicPartition = decodeSourcePartition(record.sourcePartition());
- if (srcTopicPartition == null) {
- log.warn("Could not decode source partition for record sent to topic partition {}; cannot report offsets for it", dstTopicPartition);
- return;
- }
-
- Long srcOffset = decodeSourceOffset(srcTopicPartition, record.sourceOffset());
- if (srcOffset == null) {
- log.warn("Could not decode source offset for record sent to topic partition {}; cannot report offsets for it", dstTopicPartition);
- return;
- }
-
- offsetSnapshotReporter.record(srcTopicPartition.topic(), srcTopicPartition.partition(), srcOffset, dstTopicPartition.partition(), dstOffset);
- }
- }
-
- private Collection<AssignedPartition> assignedPartitions() {
- return partitionAssignmentStore.getAssignment(
- config.connectorName(),
- config.taskId(),
- config.sourceCluster(),
- config.targetCluster()
- );
- }
-
- private static Set<TopicPartition> sourcePartitions(Collection<AssignedPartition> partitions) {
- return partitions.stream().map(AssignedPartition::source).collect(Collectors.toSet());
- }
-
- private void updateAssignment() {
- log.debug("Updating consumer assignment");
-
- Set<TopicPartition> existingAssignment = sourceConsumer.assignment();
- Collection<AssignedPartition> newAssignment = assignedPartitions();
- Set<TopicPartition> sourceAssignment = sourcePartitions(newAssignment);
- sourceConsumer.assign(sourceAssignment);
-
- List<TopicPartition> newPartitions = sourceAssignment.stream()
- .filter(Predicate.not(existingAssignment::contains))
- .toList();
- List<Map<String, Object>> newSourcePartitions = newPartitions.stream()
- .map(tp -> encodeSourcePartition(tp.topic(), tp.partition()))
- .toList();
-
- if (!newPartitions.isEmpty()) {
- log.debug("Assigned new partition(s): {}", newPartitions);
- } else {
- log.debug("Assigned no new partitions; revocation or topic/partition override may have taken place");
- }
-
- context.offsetStorageReader().offsets(newSourcePartitions).forEach((sourcePartition, sourceOffset) -> {
- TopicPartition topicPartition = decodeSourcePartition(sourcePartition);
- if (topicPartition == null) {
- return;
- }
-
- Long offset = decodeSourceOffset(topicPartition, sourceOffset);
- if (offset == null) {
- log.trace("Found no committed offset for newly-assigned partition {}", topicPartition);
- return;
- }
-
- log.trace("Found committed offset {} for newly-assigned partition {}", offset, topicPartition);
- // Stored offset is what we were last able to replicate; we should seek to one beyond that
- sourceConsumer.seek(topicPartition, offset + 1);
- });
-
- overrides = newAssignment.stream()
- .filter(AssignedPartition::hasOverride)
- .collect(Collectors.toMap(
- AssignedPartition::source,
- AssignedPartition::target
- ));
- }
-
- private void initializeWorkloadTracker() {
- Duration workloadReportInterval = config.workloadReportInterval();
- if (workloadReportInterval == null) {
- log.warn("Workload tracking is disabled");
- return;
- }
-
- log.debug("Initializing workload tracker");
- workloadTracker = new WorkloadTrackingProcessor(
- config.sourceCluster(),
- config.targetCluster(),
- // TODO: Enable workload tracking to be performed even if the workload provider isn't set up on this worker
- KafkaBasedWorkloadProvider::publishWorkload,
- () -> sourcePartitions(assignedPartitions()),
- workloadReportInterval
- );
- workloadTracker.start();
- log.debug("Finished initializing workload tracker");
- }
-
- private void initializeOffsetSnapshotReporter() {
- if (!config.offsetReporterEnabled()) {
- log.warn("Offset reporter is disabled");
- return;
- }
-
- log.debug("Initializing offset reporter");
- // We may want to make some of these values configurable in the future, but for now
- // this is fine since they're the values used across the entire uReplicator2 fleet
- int offsetReportIntervalMs = 600_000;
- String offsetSnapshotPostUrl = "http://127.0.0.1:18668/snapshot";
- int offsetSnapshotRequestTimeoutMs = 10_000;
- OffsetSnapshotClient offsetSnapshotClient = new OffsetSnapshotClient(offsetSnapshotPostUrl, offsetSnapshotRequestTimeoutMs);
-
- // 🤷 Seems as reasonable as anything I guess
- String routeId = config.connectorName() + "-" + config.sourceCluster() + "-" + config.targetCluster();
- offsetSnapshotReporter = new OffsetSnapshotReporter(
- config.sourceCluster(),
- config.targetCluster(),
- offsetReportIntervalMs,
- offsetSnapshotClient,
- routeId
- );
- offsetSnapshotReporter.start();
- log.debug("Initialized offset reporter");
- }
-
- private SourceRecord convertRecord(ConsumerRecord<byte[], byte[]> record) {
- TopicPartition source = new TopicPartition(record.topic(), record.partition());
- TopicPartition target = overrides.getOrDefault(source, source);
-
- return new SourceRecord(
- encodeSourcePartition(source.topic(), source.partition()),
- encodeSourceOffset(record),
- target.topic(),
- target.partition(),
- Schema.OPTIONAL_BYTES_SCHEMA,
- record.key(),
- Schema.OPTIONAL_BYTES_SCHEMA,
- record.value(),
- record.timestamp(),
- convertHeaders(record.headers())
- );
- }
-
- private Map<String, Object> encodeSourcePartition(String topic, int partition) {
- return encodeSourcePartition(config.sourceCluster(), config.targetCluster(), topic, partition);
- }
-
- public static Map<String, Object> encodeSourcePartition(String srcKafka, String dstKafka, String topic, int partition) {
- // LinkedHashMap to guarantee order of key/value pairs across serialization
- Map<String, Object> result = new LinkedHashMap<>();
- result.put(SOURCE_PARTITION_SOURCE_CLUSTER, srcKafka);
- result.put(SOURCE_PARTITION_TARGET_CLUSTER, dstKafka);
- result.put(SOURCE_PARTITION_KAFKA_TOPIC, topic);
- result.put(SOURCE_PARTITION_KAFKA_PARTITION, partition);
- return result;
- }
-
- private Map<String, Object> encodeSourceOffset(ConsumerRecord<?, ?> record) {
- return encodeSourceOffset(record.offset());
- }
-
- public static Map<String, Object> encodeSourceOffset(long offset) {
- return Map.of(SOURCE_OFFSET_KAFKA_OFFSET, offset);
- }
-
- private TopicPartition decodeSourcePartition(Map<String, ?> sourcePartition) {
- Object rawTopic = sourcePartition.get(SOURCE_PARTITION_KAFKA_TOPIC);
- if (rawTopic == null) {
- log.warn("Ignoring invalid source partition {}; missing {} field", sourcePartition, SOURCE_PARTITION_KAFKA_TOPIC);
- return null;
- }
- if (!(rawTopic instanceof String)) {
- log.warn("Ignoring invalid source partition {}; field {} is not a string", sourcePartition, SOURCE_PARTITION_KAFKA_TOPIC);
- return null;
- }
-
- Object rawPartition = sourcePartition.get(SOURCE_PARTITION_KAFKA_PARTITION);
- if (rawPartition == null) {
- log.warn("Ignoring invalid source partition {}; missing {} field", sourcePartition, SOURCE_PARTITION_KAFKA_PARTITION);
- return null;
- }
- if (!(rawPartition instanceof Number)) {
- log.warn("Ignoring invalid source partition {}; field {} is not a number", sourcePartition, SOURCE_PARTITION_KAFKA_PARTITION);
- return null;
- }
-
- String topic = (String) rawTopic;
- int partition = ((Number) rawPartition).intValue();
- return new TopicPartition(topic, partition);
- }
-
- private Long decodeSourceOffset(TopicPartition topicPartition, Map<String, ?> sourceOffset) {
- if (sourceOffset == null) {
- log.debug("No source offset found for {}; this partition is probably being replicated for the first time", topicPartition);
- return null;
- }
-
- Object rawOffset = sourceOffset.get(SOURCE_OFFSET_KAFKA_OFFSET);
- if (rawOffset == null) {
- log.warn("Ignoring invalid source offset {}; missing {} field", sourceOffset, SOURCE_OFFSET_KAFKA_OFFSET);
- return null;
- }
- if (!(rawOffset instanceof Number)) {
- log.warn("Ignoring invalid source offset {}; field {} is not a number", sourceOffset, SOURCE_OFFSET_KAFKA_OFFSET);
- return null;
- }
-
- return ((Number) rawOffset).longValue();
- }
-
- private static Headers convertHeaders(org.apache.kafka.common.header.Headers headers) {
- ConnectHeaders result = new ConnectHeaders();
- for (Header header : headers) {
- result.addBytes(header.key(), header.value());
- }
- return result;
- }
-
- private void onAssignmentUpdate() {
- this.assignmentChanged = true;
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3TaskConfig.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3TaskConfig.java
deleted file mode 100644
index 749a5a2e1d9..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/UReplicator3TaskConfig.java
+++ /dev/null
@@ -1,42 +0,0 @@
-package com.uber.data.kafka.connect.ureplicator3;
-
-import org.apache.kafka.common.config.ConfigDef;
-import org.apache.kafka.common.config.ConfigDef.Importance;
-import org.apache.kafka.common.config.ConfigDef.Type;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.Map;
-
-public class UReplicator3TaskConfig extends UReplicator3Config {
-
- private static final Logger log = LoggerFactory.getLogger(UReplicator3TaskConfig.class);
-
- public static final String TASK_ID_CONFIG = "task.id";
-
- @SuppressWarnings("this-escape")
- public UReplicator3TaskConfig(Map<String, ?> props) {
- super(config(), props);
- }
-
- protected static ConfigDef config() {
- return baseConfig()
- .defineInternal(
- TASK_ID_CONFIG,
- Type.INT,
- ConfigDef.NO_DEFAULT_VALUE,
- Importance.HIGH
- );
- }
-
- @Override
- protected String role() {
- return "task-" + connectorName() + "-" + taskId();
- }
-
- public int taskId() {
- return getInt(TASK_ID_CONFIG);
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignedPartition.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignedPartition.java
deleted file mode 100644
index 5cc6c77604b..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignedPartition.java
+++ /dev/null
@@ -1,27 +0,0 @@
-package com.uber.data.kafka.connect.ureplicator3.coordination;
-
-import org.apache.kafka.common.TopicPartition;
-
-public record AssignedPartition(
- String srcTopic,
- int srcPartition,
- String dstTopic,
- Integer dstPartition
-) {
-
- public TopicPartition source() {
- return new TopicPartition(srcTopic, srcPartition);
- }
-
- public TopicPartition target() {
- return new TopicPartition(
- dstTopic != null ? dstTopic : srcTopic,
- dstPartition != null ? dstPartition : srcPartition
- );
- }
-
- public boolean hasOverride() {
- return !source().equals(target());
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentListener.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentListener.java
deleted file mode 100644
index 3058d3d45ad..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentListener.java
+++ /dev/null
@@ -1,11 +0,0 @@
-package com.uber.data.kafka.connect.ureplicator3.coordination;
-
-public interface AssignmentListener {
-
- /**
- * Invoked when the assignment for the task has been updated. Long-running or blocking
- * operations should <strong>not</strong> be performed inside this method
- */
- void onUpdate();
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentStore.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentStore.java
deleted file mode 100644
index 2e2c413a537..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentStore.java
+++ /dev/null
@@ -1,368 +0,0 @@
-package com.uber.data.kafka.connect.ureplicator3.coordination;
-
-import org.apache.kafka.clients.admin.NewTopic;
-import org.apache.kafka.clients.consumer.ConsumerConfig;
-import org.apache.kafka.clients.consumer.ConsumerRecord;
-import org.apache.kafka.clients.producer.ProducerConfig;
-import org.apache.kafka.common.KafkaException;
-import org.apache.kafka.common.internals.KafkaCompletableFuture;
-import org.apache.kafka.common.serialization.ByteArrayDeserializer;
-import org.apache.kafka.common.serialization.ByteArraySerializer;
-import org.apache.kafka.common.utils.Time;
-import org.apache.kafka.common.utils.Utils;
-import org.apache.kafka.connect.util.KafkaBasedLog;
-import org.apache.kafka.connect.util.TopicAdmin;
-
-import com.fasterxml.jackson.core.JsonProcessingException;
-import com.fasterxml.jackson.core.type.TypeReference;
-import com.fasterxml.jackson.databind.ObjectMapper;
-import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.Closeable;
-import java.io.IOException;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.LinkedHashMap;
-import java.util.Map;
-import java.util.Objects;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.Future;
-import java.util.stream.Collectors;
-
-import static com.uber.data.kafka.connect.utils.SerializationUtils.parseMapField;
-import static com.uber.data.kafka.connect.utils.SerializationUtils.parseNumberField;
-import static com.uber.data.kafka.connect.utils.SerializationUtils.parseOptionalNumberField;
-import static com.uber.data.kafka.connect.utils.SerializationUtils.parseOptionalStringField;
-import static com.uber.data.kafka.connect.utils.SerializationUtils.parseStringField;
-
-// This class is thread-safe
-public class AssignmentStore implements Closeable {
-
- private static final Logger log = LoggerFactory.getLogger(AssignmentStore.class);
- private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
-
- private static final String ASSIGNMENT_FIELD = "assignment-v1";
- private static final String CONNECTOR_FIELD = "connector";
- private static final String SRC_KAFKA_FIELD = "src";
- private static final String DST_KAFKA_FIELD = "dst";
- private static final String TOPIC_FIELD = "topic";
- private static final String PARTITION_FIELD = "partition";
-
- private static final String TASK_FIELD = "task";
- private static final String DST_TOPIC_FIELD = "dst-topic";
- private static final String DST_PARTITION_FIELD = "dst-partition";
-
- private final String topic;
- private final int topicPartitions;
- private final short topicReplicationFactor;
- private final Map<String, Object> topicSettings;
- private final Map<String, Object> producerProps;
- private final Map<String, Object> consumerProps;
- private final Map<String, Object> adminProps;
- private final Map<String, Map<ReplicatedTopicPartition, AssignmentValue>> taskPartitionAssignments;
- private final Map<String, Map<Integer, AssignmentListener>> assignmentListeners;
-
- private TopicAdmin topicAdmin;
- private KafkaBasedLog<byte[], byte[]> kafkaLog;
-
- public AssignmentStore(AssignmentStoreConfig config) {
- this.topic = config.storageTopic();
- this.topicPartitions = config.storageTopicPartitions();
- this.topicReplicationFactor = config.storageTopicReplicationFactor();
- this.topicSettings = config.storageTopicSettings();
- this.producerProps = producerProps(config.producerConfig());
- this.consumerProps = consumerProps(config.consumerConfig());
- this.adminProps = config.adminConfig();
- this.taskPartitionAssignments = new HashMap<>();
- this.assignmentListeners = new HashMap<>();
- }
-
- public void start() {
- this.topicAdmin = new TopicAdmin(adminProps);
- this.kafkaLog = new KafkaBasedLog<>(
- topic,
- producerProps,
- consumerProps,
- () -> topicAdmin,
- this::onConsume,
- Time.SYSTEM,
- this::createAndValidateTopic
- );
- kafkaLog.start();
- }
-
- private static Map<String, Object> producerProps(Map<String, Object> baseProps) {
- baseProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
- baseProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
- baseProps.put(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG, Integer.MAX_VALUE);
- return baseProps;
- }
-
- private static Map<String, Object> consumerProps(Map<String, Object> baseProps) {
- baseProps.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
- baseProps.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
- return baseProps;
- }
-
- private void createAndValidateTopic(TopicAdmin topicAdmin) {
- NewTopic newTopic = TopicAdmin.defineTopic(topic)
- .config(topicSettings)
- .compacted()
- .partitions(topicPartitions)
- .replicationFactor(topicReplicationFactor)
- .build();
- boolean created = topicAdmin.createTopic(newTopic);
- if (!created) {
- // Ensure existing topic is compacted
- boolean verified = topicAdmin.verifyTopicCleanupPolicyOnlyCompact(
- topic,
- AssignmentStoreConfig.STORAGE_TOPIC_CONFIG,
- "task partition assignments"
- );
- // Should never, NEVER happen, but just in case
- if (!verified) {
- throw new KafkaException("Failed to verify that task partition assignment storage topic " + topic + " is compacted");
- }
- }
- }
-
- @Override
- public void close() {
- if (kafkaLog != null) {
- Utils.closeQuietly(kafkaLog::stop, "Kafka-based log");
- }
- Utils.closeQuietly(topicAdmin, "topic admin");
- }
-
- public Future<?> assignPartition(String connector, ReplicatedTopicPartition replicatedTopicPartition, int taskId, String dstTopic, Integer dstPartition) {
- return publishAssignment(connector, replicatedTopicPartition, dstTopic, dstPartition, taskId);
- }
-
- public Future<?> revokePartition(String connector, ReplicatedTopicPartition replicatedTopicPartition) {
- return publishAssignment(connector, replicatedTopicPartition, null, null, null);
- }
-
- public Map<Integer, ? extends Collection<AssignedPartition>> getAssignments(String connector, String srcKafka, String dstKafka) {
- Objects.requireNonNull(connector);
- Objects.requireNonNull(srcKafka);
- Objects.requireNonNull(dstKafka);
-
- Map<ReplicatedTopicPartition, AssignmentValue> connectorSnapshot;
- synchronized (this) {
- connectorSnapshot = new HashMap<>(taskPartitionAssignments.getOrDefault(connector, Map.of()));
- }
-
- return connectorSnapshot.entrySet().stream()
- .filter(e -> e.getKey().matchesPipeline(srcKafka, dstKafka))
- .collect(Collectors.groupingBy(
- e -> e.getValue().taskId(),
- Collectors.mapping(AssignmentStore::assignedPartition, Collectors.toSet())
- ));
- }
-
- public Collection<AssignedPartition> getAssignment(String connector, int taskId, String srcKafka, String dstKafka) {
- synchronized (this) {
- return taskPartitionAssignments.getOrDefault(connector, Map.of()).entrySet().stream()
- .filter(e -> e.getKey().matchesPipeline(srcKafka, dstKafka))
- .filter(e -> taskId == e.getValue().taskId())
- .map(AssignmentStore::assignedPartition)
- .collect(Collectors.toSet());
- }
- }
-
- // TODO: Add an API to remove listeners as well (have to be careful not to overwrite listeners on slow task shutdown, though)
- public void addListener(String connector, int taskId, AssignmentListener listener) {
- synchronized (this) {
- assignmentListeners.computeIfAbsent(connector, k -> new HashMap<>())
- .put(taskId, listener);
- }
- }
-
- public void refresh() {
- try {
- kafkaLog.readToEnd().get();
- } catch (InterruptedException | ExecutionException e) {
- throw new KafkaException("Failed to read to end of task partition assignment log", e);
- }
- }
-
- // For testing only
- public void clear() {
- taskPartitionAssignments.clear();
- }
-
- private Future<?> publishAssignment(String connector, ReplicatedTopicPartition replicatedTopicPartition, String dstTopic, Integer dstPartition, Integer taskId) {
- if (connector == null) {
- log.error("Failed to serialize assignment key with null connector name");
- return KafkaCompletableFuture.failedFuture(new NullPointerException("Connector may not be null"));
- }
- if (replicatedTopicPartition == null) {
- log.error("Failed to serialize assignment key with null replicated topic partition");
- return KafkaCompletableFuture.failedFuture(new NullPointerException("Replicated topic partition may not be null"));
- }
-
- byte[] key;
- try {
- key = serializeAssignmentKey(connector, replicatedTopicPartition);
- } catch (JsonProcessingException e) {
- log.error("Failed to serialize assignment key for connector {} and replicated topic partition {} as JSON", connector, replicatedTopicPartition, e);
- return KafkaCompletableFuture.failedFuture(e);
- }
-
- byte[] value;
- if (taskId == null) {
- value = null;
- } else {
- try {
- value = serializeAssignmentValue(taskId, dstTopic, dstPartition);
- } catch (JsonProcessingException e) {
- log.error(
- "Failed to serialize assignment value (task={}, dstTopic={}, dstPartition={}) for connector {} and replicated topic partition {} as JSON",
- taskId, dstTopic, dstPartition, connector, replicatedTopicPartition, e);
- return KafkaCompletableFuture.failedFuture(e);
- }
- }
-
- // TODO: Error handling (probably just log and emit a metric)
- return kafkaLog.sendWithReceipt(key, value);
- }
-
- private void onConsume(Throwable error, ConsumerRecord<byte[], byte[]> record) {
- if (error != null) {
- log.error("Error occurred while consuming from task partition assignment topic", error);
- return;
- }
-
- int recordPartition = record.partition();
- long recordOffset = record.offset();
-
- if (record.key() == null) {
- log.error("Ignoring record with null key {} in partition {} at offset {}", record.key(), recordPartition, recordOffset);
- return;
- }
-
- AssignmentKey assignmentKey;
- try {
- assignmentKey = deserializeAssignmentKey(record.key());
- } catch (Exception e) {
- log.error("Failed to deserialize task partition assignment key in partition {} at offset {} as JSON", recordPartition, recordOffset, e);
- return;
- }
-
- AssignmentValue assignmentValue;
- try {
- assignmentValue = deserializeAssignmentValue(record.value());
- } catch (IOException e) {
- log.error("Failed to deserialize task partition assignment value in partition {} at offset {} for assignment key {} as JSON", recordPartition, recordOffset, assignmentKey, e);
- return;
- }
-
- // TODO: Notify listeners
- String connector = assignmentKey.connector();
- ReplicatedTopicPartition replicatedTopicPartition = assignmentKey.replicatedTopicPartition();
- if (assignmentValue == null) {
- log.trace("Picked up revocation of {} for connector {}", replicatedTopicPartition, connector);
- synchronized (this) {
- Map<ReplicatedTopicPartition, AssignmentValue> connectorAssignments = taskPartitionAssignments.get(connector);
- if (connectorAssignments == null) {
- log.debug("No cached assignments for connector {}", connector);
- } else {
- AssignmentValue priorAssignment = connectorAssignments.remove(replicatedTopicPartition);
- if (priorAssignment != null) {
- notifyListener(connector, priorAssignment.taskId());
- }
- }
- }
- } else {
- log.trace("Picked up assignment of {} with destination topic {} and destination partition {} to task {} for connector {}",
- replicatedTopicPartition, assignmentValue.dstTopic(), assignmentValue.dstPartition(), assignmentValue.taskId(), connector);
- synchronized (this) {
- Map<ReplicatedTopicPartition, AssignmentValue> connectorAssignments = taskPartitionAssignments.computeIfAbsent(connector, k -> new HashMap<>());
- AssignmentValue priorAssignment = connectorAssignments.put(replicatedTopicPartition, assignmentValue);
- if (priorAssignment != null) {
- notifyListener(connector, priorAssignment.taskId());
- }
- notifyListener(connector, assignmentValue.taskId());
- }
- }
- }
-
- private byte[] serializeAssignmentKey(String connector, ReplicatedTopicPartition replicatedTopicPartition) throws JsonProcessingException {
- // LinkedHashMap preserves order and allows us to deterministically generate a serialized key from a replicated topic partition
- Map<String, Object> innerMap = new LinkedHashMap<>();
- innerMap.put(CONNECTOR_FIELD, connector);
- innerMap.put(SRC_KAFKA_FIELD, replicatedTopicPartition.srcKafka());
- innerMap.put(DST_KAFKA_FIELD, replicatedTopicPartition.dstKafka());
- innerMap.put(TOPIC_FIELD, replicatedTopicPartition.topicPartition().topic());
- innerMap.put(PARTITION_FIELD, replicatedTopicPartition.topicPartition().partition());
- Map<String, Object> wrapped = Map.of(ASSIGNMENT_FIELD, innerMap);
- return OBJECT_MAPPER.writeValueAsBytes(wrapped);
- }
-
- private byte[] serializeAssignmentValue(int taskId, String dstTopic, Integer dstPartition) throws JsonProcessingException {
- Map<String, Object> map = new LinkedHashMap<>();
- map.put(TASK_FIELD, taskId);
- if (dstTopic != null)
- map.put(DST_TOPIC_FIELD, dstTopic);
- if (dstPartition != null)
- map.put(DST_PARTITION_FIELD, dstPartition);
- return OBJECT_MAPPER.writeValueAsBytes(map);
- }
-
- private AssignmentKey deserializeAssignmentKey(byte[] key) throws IOException {
- if (key == null)
- throw new IOException("Key should not be null");
-
- Map<String, Object> map = OBJECT_MAPPER.readValue(key, new TypeReference<>() { });
- Map<?, ?> wrapped = parseMapField(true, map, ASSIGNMENT_FIELD);
-
- String connector = parseStringField(true, wrapped, CONNECTOR_FIELD);
- String srcKafka = parseStringField(true, wrapped, SRC_KAFKA_FIELD);
- String dstKafka = parseStringField(true, wrapped, DST_KAFKA_FIELD);
- String topic = parseStringField(true, wrapped, TOPIC_FIELD);
- Number partition = parseNumberField(true, wrapped, PARTITION_FIELD);
- ReplicatedTopicPartition replicatedTopicPartition = new ReplicatedTopicPartition(srcKafka, dstKafka, topic, partition.intValue());
-
- return new AssignmentKey(connector, replicatedTopicPartition);
- }
-
- private AssignmentValue deserializeAssignmentValue(byte[] value) throws IOException {
- // Permit null values to signify revocation of the partition
- if (value == null)
- return null;
-
- Map<String, Object> map = OBJECT_MAPPER.readValue(value, new TypeReference<>() { });
- int taskId = parseNumberField(false, map, TASK_FIELD).intValue();
- String dstTopic = parseOptionalStringField(false, map, DST_TOPIC_FIELD);
- Number rawDstPartition = parseOptionalNumberField(false, map, DST_PARTITION_FIELD);
- Integer dstPartition = rawDstPartition != null ? rawDstPartition.intValue() : null;
- return new AssignmentValue(taskId, dstTopic, dstPartition);
- }
-
- private void notifyListener(String connector, int taskId) {
- AssignmentListener listener = assignmentListeners.getOrDefault(connector, Map.of())
- .get(taskId);
-
- if (listener == null)
- return;
-
- try {
- listener.onUpdate();
- } catch (Throwable t) {
- log.error("Listener for {}-{} threw an exception", connector, taskId, t);
- }
- }
-
- private static AssignedPartition assignedPartition(Map.Entry<ReplicatedTopicPartition, AssignmentValue> entry) {
- return new AssignedPartition(entry.getKey().topic(), entry.getKey().partition(), entry.getValue().dstTopic(), entry.getValue().dstPartition);
- }
-
- private record AssignmentKey(String connector, ReplicatedTopicPartition replicatedTopicPartition) {
- }
-
- private record AssignmentValue(int taskId, String dstTopic, Integer dstPartition) {
- }
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentStoreConfig.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentStoreConfig.java
deleted file mode 100644
index 27165db2860..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentStoreConfig.java
+++ /dev/null
@@ -1,145 +0,0 @@
-package com.uber.data.kafka.connect.ureplicator3.coordination;
-
-import org.apache.kafka.common.config.AbstractConfig;
-import org.apache.kafka.common.config.ConfigDef;
-import org.apache.kafka.common.config.ConfigDef.Importance;
-import org.apache.kafka.common.config.ConfigDef.Type;
-import org.apache.kafka.common.config.TopicConfig;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.HashMap;
-import java.util.Map;
-
-import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
-
-public class AssignmentStoreConfig extends AbstractConfig {
-
- private static final Logger log = LoggerFactory.getLogger(AssignmentStoreConfig.class);
-
- private static final String CONFIG_PREFIX = "uber.ureplicator3.assignments.";
-
- public static final String STORAGE_TOPIC_CONFIG = CONFIG_PREFIX + "storage.topic";
- public static final String STORAGE_TOPIC_PREFIX = STORAGE_TOPIC_CONFIG + ".";
- private static final String STORAGE_TOPIC_DOC = "Topic used to store assignments. "
- + "This topic must be compacted (i.e., have the 'compacted' cleanup policy) and can have "
- + "any number of partitions, but it's recommended to use a low number (e.g., 5 or 10). "
- + "If the topic does not exist, it will be automatically created on startup. "
- + "Additional properties for the topic can be specified using the " + STORAGE_TOPIC_PREFIX + " "
- + "prefix";
-
- private static final String PARTITIONS_SUFFIX = "partitions";
- public static final String STORAGE_TOPIC_PARTITIONS_CONFIG = STORAGE_TOPIC_PREFIX + PARTITIONS_SUFFIX;
- private static final int STORAGE_TOPIC_PARTITIONS_DEFAULT = 10;
- private static final String STORAGE_TOPIC_PARTITIONS_DOC = "The number of partitions to create "
- + "the assignments topic with, if it does not already exist. If the topic already exists this property "
- + "will be ignored, even if the topic has a different number of partitions";
-
- private static final String REPLICATION_FACTOR_SUFFIX = "replication.factor";
- public static final String STORAGE_TOPIC_REPLICATION_FACTOR_CONFIG = STORAGE_TOPIC_PREFIX + REPLICATION_FACTOR_SUFFIX;
- private static final short STORAGE_TOPIC_REPLICATION_FACTOR_DEFAULT = (short) 3;
- private static final String STORAGE_TOPIC_REPLICATION_FACTOR_DOC = "The replication factor to create "
- + "the assignments topic with, if it does not already exist. If the topic already exists this property "
- + "will be ignored, even if the topic has a different replication factor";
-
- private static final String SEGMENT_BYTES_SUFFIX = TopicConfig.SEGMENT_BYTES_CONFIG;
- public static final String STORAGE_TOPIC_SEGMENT_BYTES_CONFIG = STORAGE_TOPIC_PREFIX + SEGMENT_BYTES_SUFFIX;
- private static final int STORAGE_TOPIC_SEGMENT_BYTES_DEFAULT = 1024 * 1024 * 3; // Lower default segment size in order to enable more-frequent compaction
- private static final String STORAGE_TOPIC_SEGMENT_BYTES_DOC = "The per-partition segment size "
- + "for the assignments topic. This should be relatively low in order to enable topic compaction to take place, which "
- + "will help keep the topic from growing too large and causing slow startup time for workers. Like other topic "
- + "properties, if the topic already exists, this property will be ignored, even if the topic has a different segment size";
-
- public static final String CLIENTS_PREFIX = CONFIG_PREFIX + "kafka.clients";
- public static final String CONSUMER_PREFIX = CONFIG_PREFIX + "consumer";
- public static final String PRODUCER_PREFIX = CONFIG_PREFIX + "producer";
- public static final String ADMIN_PREFIX = CONFIG_PREFIX + "admin";
-
- public AssignmentStoreConfig(Map<String, ?> props) {
- super(config(), props);
- }
-
- public static ConfigDef config() {
- return addToConfig(new ConfigDef());
- }
-
- public static ConfigDef addToConfig(ConfigDef base) {
- return base
- .define(
- STORAGE_TOPIC_CONFIG,
- Type.STRING,
- Importance.HIGH,
- STORAGE_TOPIC_DOC
- ).define(
- STORAGE_TOPIC_PARTITIONS_CONFIG,
- Type.INT,
- STORAGE_TOPIC_PARTITIONS_DEFAULT,
- atLeast(1),
- Importance.LOW,
- STORAGE_TOPIC_PARTITIONS_DOC
- ).define(
- STORAGE_TOPIC_REPLICATION_FACTOR_CONFIG,
- Type.SHORT,
- STORAGE_TOPIC_REPLICATION_FACTOR_DEFAULT,
- atLeast(1),
- Importance.LOW,
- STORAGE_TOPIC_REPLICATION_FACTOR_DOC
- ).define(
- STORAGE_TOPIC_SEGMENT_BYTES_CONFIG,
- Type.INT,
- STORAGE_TOPIC_SEGMENT_BYTES_DEFAULT,
- atLeast(1),
- Importance.LOW,
- STORAGE_TOPIC_SEGMENT_BYTES_DOC
- );
- }
-
- public Map<String, Object> adminConfig() {
- return clientConfig(ADMIN_PREFIX);
- }
-
- public Map<String, Object> producerConfig() {
- return clientConfig(PRODUCER_PREFIX);
- }
-
- public Map<String, Object> consumerConfig() {
- return clientConfig(CONSUMER_PREFIX);
- }
-
- public String storageTopic() {
- return getString(STORAGE_TOPIC_CONFIG);
- }
-
- public int storageTopicPartitions() {
- return getInt(STORAGE_TOPIC_PARTITIONS_CONFIG);
- }
-
- public short storageTopicReplicationFactor() {
- return getShort(STORAGE_TOPIC_REPLICATION_FACTOR_CONFIG);
- }
-
- public Map<String, Object> storageTopicSettings() {
- Map<String, Object> result = originalsWithPrefix(STORAGE_TOPIC_PREFIX);
-
- // AbstractConfig::originalsWithPrefix doesn't include default values, we apply that manually here
- result.computeIfAbsent(TopicConfig.SEGMENT_BYTES_CONFIG, k -> STORAGE_TOPIC_SEGMENT_BYTES_DEFAULT);
-
- Object removedPolicy = result.remove(TopicConfig.CLEANUP_POLICY_CONFIG);
- if (removedPolicy != null) {
- log.warn("Ignoring '{}cleanup.policy={}' setting, since compaction is always used", STORAGE_TOPIC_PREFIX, removedPolicy);
- }
-
- result.remove(REPLICATION_FACTOR_SUFFIX);
- result.remove(PARTITIONS_SUFFIX);
-
- return result;
- }
-
- private Map<String, Object> clientConfig(String specialPrefix) {
- Map<String, Object> result = new HashMap<>(originals());
- result.putAll(originalsWithPrefix(CLIENTS_PREFIX));
- result.putAll(originalsWithPrefix(specialPrefix));
- return result;
- }
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentStores.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentStores.java
deleted file mode 100644
index f58ebc830c9..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/ureplicator3/coordination/AssignmentStores.java
+++ /dev/null
@@ -1,39 +0,0 @@
-package com.uber.data.kafka.connect.ureplicator3.coordination;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-public class AssignmentStores {
-
- private static final Logger log = LoggerFactory.getLogger(AssignmentStores.class);
-
- private static AssignmentStore store;
-
- public static synchronized void initialize(AssignmentStoreConfig config) {
- if (store == null)
- store = initializeStore(config);
- else
- log.warn("Store is already initialized; this should never happen outside an embedded testing environment");
- }
-
- // TODO: Instead of exposing the whole store, we could expose a safe set of operations
- // like assign, revoke, and addListener here (this would prevent callers from doing
- // something stupid like trying to close the store)
- public static AssignmentStore get() {
- if (store == null) {
- throw new IllegalStateException("Store has not been initialized yet");
- }
- return store;
- }
-
- // TODO: private void destroyStore()
-
- private static AssignmentStore initializeStore(AssignmentStoreConfig config) {
- log.info("Initializing assignment store");
- AssignmentStore result = new AssignmentStore(config);
- result.start();
- log.info("Finished initializing assignment store");
- return result;
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/BrokerResolver.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/BrokerResolver.java
deleted file mode 100644
index 0933ebbe7e9..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/BrokerResolver.java
+++ /dev/null
@@ -1,39 +0,0 @@
-package com.uber.data.kafka.connect.utils;
-
-
-import com.uber.data.kafka.common.ClusterNameHelper;
-import com.uber.data.kafka.common.ClusterNameHelper.ClusterNameParseResult;
-import com.uber.data.kafka.exception.ResolverException;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-
-// TODO: Singleton pattern and fault injection in dependent classes
-public final class BrokerResolver {
-
- private static final Logger log = LoggerFactory.getLogger(BrokerResolver.class);
-
- // TODO: We should expose normalization logic so that different aliases for the same Kafka
- // cluster get treated the same (e.g., when storing offsets or logging replication pipelines)
-
- public static String resolveBootstrap(String cluster) throws ResolverException {
- if (cluster == null) {
- throw new ResolverException("Cannot parse null cluster name");
- }
-
- ClusterNameParseResult parseResult;
- try {
- parseResult = ClusterNameHelper.parseClusterName(cluster);
- } catch (IllegalArgumentException e) {
- throw new ResolverException("Failed to parse cluster name '" + cluster + "'", e);
- }
- // TODO: There has to be a better way 💀
- return String.format(
- "kafka-%s.%s.uber.internal:9092",
- parseResult.functionalGroup,
- parseResult.region
- );
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/ExecutorUtils.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/ExecutorUtils.java
deleted file mode 100644
index 74f1517550e..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/ExecutorUtils.java
+++ /dev/null
@@ -1,88 +0,0 @@
-package com.uber.data.kafka.connect.utils;
-
-import com.google.common.util.concurrent.ThreadFactoryBuilder;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.ThreadFactory;
-import java.util.concurrent.TimeUnit;
-
-
-/**
- * This is a simple wrapper around the standard library and Guava's thread utils that can
- * be used to create and work with executors.
- *
- * @see Executors
- * @see ExecutorService
- * @see ScheduledExecutorService
- * @see ThreadFactory
- * @see ThreadFactoryBuilder
- */
-public class ExecutorUtils {
- private static final Logger log = LoggerFactory.getLogger(ExecutorUtils.class);
-
- private ExecutorUtils() {
- }
-
- /**
- * Create a new executor that uses the given thread name format and for which
- * all created threads are daemons
- *
- * @param nameFormat the name format for executor threads
- * @return the executor; never null
- */
- public static ScheduledExecutorService singleThreadedScheduled(String nameFormat) {
- return Executors.newSingleThreadScheduledExecutor(threadFactory(nameFormat));
- }
-
- /**
- * Shut down the executor, waiting for the specified timeout for all running
- * and scheduled tasks to complete
- *
- * @param executor the executor to shut down; may not be null
- * @param executorDescription a description of the executor to use in log messages;
- * for example, "offset monitor" or "topic watcher service"
- * @param timeout the time to wait for all tasks to complete
- * @param unit the unit for the timeout
- */
- public static void shutdown(
- ExecutorService executor,
- String executorDescription,
- long timeout,
- TimeUnit unit
- ) {
- executor.shutdown();
- boolean shutdownSucceeded = false;
- boolean interrupted = false;
- try {
- shutdownSucceeded = executor.awaitTermination(timeout, unit);
- if (!shutdownSucceeded) {
- log.warn("Could not shut down {} in time", executorDescription);
- } else {
- log.info("Gracefully shut down {}", executorDescription);
- }
- } catch (InterruptedException e) {
- log.warn("Interrupted while shutting down {}", executorDescription, e);
- interrupted = true;
- } finally {
- if (!shutdownSucceeded) {
- executor.shutdownNow();
- }
-
- // Preserve the interrupt status when control is returned to the caller
- if (interrupted)
- Thread.currentThread().interrupt();
- }
- }
-
- private static ThreadFactory threadFactory(String nameFormat) {
- return new ThreadFactoryBuilder()
- .setNameFormat(nameFormat)
- .setDaemon(true)
- .build();
- }
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/M3ReporterProvider.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/M3ReporterProvider.java
deleted file mode 100644
index 1538ae9d42b..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/M3ReporterProvider.java
+++ /dev/null
@@ -1,43 +0,0 @@
-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
deleted file mode 100644
index 064e12b43fe..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/MetricsUtils.java
+++ /dev/null
@@ -1,37 +0,0 @@
-package com.uber.data.kafka.connect.utils;
-
-import org.apache.kafka.common.TopicPartition;
-
-import com.google.common.collect.ImmutableMap;
-import com.uber.m3.tally.Scope;
-
-import java.util.Map;
-
-public class MetricsUtils {
-
- public static final MetricsUtils INSTANCE = new MetricsUtils();
-
- private static final String TOPIC = "topic";
- private static final String PARTITION = "partition";
- private static final String SRC_CLUSTER = "srcCluster";
- private static final String DST_CLUSTER = "dstCluster";
-
- private MetricsUtils() {
- }
-
- public Scope routeScope(String srcKafka, String dstKafka) {
- return M3ReporterProvider.SCOPE
- .tagged(ImmutableMap.of(
- SRC_CLUSTER, srcKafka,
- DST_CLUSTER, dstKafka
- ));
- }
-
- public Map<String, String> tags(TopicPartition topicPartition) {
- return ImmutableMap.of(
- TOPIC, topicPartition.topic(),
- PARTITION, Integer.toString(topicPartition.partition())
- );
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/SerializationUtils.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/SerializationUtils.java
deleted file mode 100644
index 60cd5c91418..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/SerializationUtils.java
+++ /dev/null
@@ -1,78 +0,0 @@
-package com.uber.data.kafka.connect.utils;
-
-import java.io.IOException;
-import java.util.Map;
-
-public final class SerializationUtils {
-
- public static Map<?, ?> parseMapField(boolean isKey, Map<?, ?> map, String field) throws IOException {
- Object fieldValue = map.get(field);
- if (fieldValue == null)
- throw new MissingFieldException(isKey, field);
-
- if (!(fieldValue instanceof Map))
- throw new InvalidFieldTypeException(isKey, field, "map", fieldValue);
-
- return (Map<?, ?>) fieldValue;
- }
-
- public static String parseStringField(boolean isKey, Map<?, ?> map, String field) throws IOException {
- String result = parseOptionalStringField(isKey, map, field);
-
- if (result == null)
- throw new MissingFieldException(isKey, field);
-
- return result;
- }
-
- public static String parseOptionalStringField(boolean isKey, Map<?, ?> map, String field) throws IOException {
- Object fieldValue = map.get(field);
- if (fieldValue == null)
- return null;
-
- if (!(fieldValue instanceof String))
- throw new InvalidFieldTypeException(isKey, field, "string", fieldValue);
-
- return (String) fieldValue;
- }
-
- public static Number parseNumberField(boolean isKey, Map<?, ?> map, String field) throws IOException {
- Number result = parseOptionalNumberField(isKey, map, field);
-
- if (result == null)
- throw new MissingFieldException(isKey, field);
-
- return result;
- }
-
- public static Number parseOptionalNumberField(boolean isKey, Map<?, ?> map, String field) throws IOException {
- Object fieldValue = map.get(field);
- if (fieldValue == null)
- return null;
-
- if (!(fieldValue instanceof Number))
- throw new InvalidFieldTypeException(isKey, field, "string", fieldValue);
-
- return (Number) fieldValue;
- }
-
- public static class MissingFieldException extends IOException {
- public MissingFieldException(boolean isKey, String field) {
- super(
- "record has invalid " + (isKey ? "key" : "value")
- + ": missing '" + field + "' field"
- );
- }
- }
-
- public static class InvalidFieldTypeException extends IOException {
- public InvalidFieldTypeException(boolean isKey, String field, String expectedType, Object actualValue) {
- super(
- "record has invalid " + (isKey ? "key" : "value")
- + ": field '" + field + "' should be a(n) " + expectedType
- + "; was " + (actualValue == null ? "null" : actualValue.getClass().getSimpleName())
- );
- }
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/UberConnectUtils.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/UberConnectUtils.java
deleted file mode 100644
index c2e7eb0af22..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/utils/UberConnectUtils.java
+++ /dev/null
@@ -1,205 +0,0 @@
-package com.uber.data.kafka.connect.utils;
-
-import org.apache.kafka.clients.admin.Admin;
-import org.apache.kafka.clients.admin.DescribeTopicsOptions;
-import org.apache.kafka.clients.admin.TopicDescription;
-import org.apache.kafka.clients.admin.TopicListing;
-import org.apache.kafka.clients.consumer.ConsumerRecord;
-import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.config.Config;
-import org.apache.kafka.common.config.ConfigValue;
-import org.apache.kafka.common.record.TimestampType;
-import org.apache.kafka.common.utils.Time;
-import org.apache.kafka.common.utils.Timer;
-
-import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
-import com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig;
-import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStore;
-import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStores;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.time.Duration;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Map;
-import java.util.Objects;
-import java.util.Optional;
-import java.util.Set;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.Future;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
-import java.util.function.Predicate;
-import java.util.stream.Collectors;
-
-import static org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLASS_CONFIG;
-
-public final class UberConnectUtils {
-
- private static final Logger log = LoggerFactory.getLogger(UberConnectUtils.class);
-
-
- @SafeVarargs
- public static <E> Set<E> diff(Collection<E> base, Collection<E>... toSubtract) {
- Set<E> result = new HashSet<>(base);
- for (Collection<E> subtracted : toSubtract) {
- if (subtracted != null)
- result.removeAll(subtracted);
- }
- return result;
- }
-
- public static Set<String> listTopics(
- Admin admin, long timeout, TimeUnit timeUnit
- ) throws InterruptedException, TimeoutException, ExecutionException {
- return admin.listTopics().listings().get(timeout, timeUnit).stream()
- .filter(Predicate.not(TopicListing::isInternal))
- .map(TopicListing::name)
- .collect(Collectors.toSet());
- }
-
- public static Map<String, Integer> listPartitions(Collection<String> topics, Admin admin, long timeout, TimeUnit timeUnit) {
- int timeoutMs = (int) Math.max(Integer.MAX_VALUE, timeUnit.toMillis(timeout));
- DescribeTopicsOptions options = new DescribeTopicsOptions()
- .includeAuthorizedOperations(false) // Don't care about ACLS
- .timeoutMs(timeoutMs);
-
- Map<String, Integer> result = new HashMap<>();
- admin.describeTopics(topics, options).topicNameValues().forEach((topic, topicDescriptionFuture) -> {
- TopicDescription topicDescription;
- try {
- topicDescription = topicDescriptionFuture.get();
- } catch (InterruptedException | ExecutionException e) {
- log.warn("Unable to describe topic {}", topic, e);
- return;
- }
-
- result.put(topic, topicDescription.partitions().size());
- });
-
- return result;
- }
-
- /**
- * Extract and return a record's {@link ConsumerRecord#timestamp() timestamp}, if it has one
- * @param record the record; may not be null
- * @return an {@link Optional} containing the record's timestamp, or {@link Optional#empty()}
- * if the record does not have a timestamp
- */
- public static Optional<Long> extractTimestamp(ConsumerRecord<?, ?> record) {
- return Optional.of(record)
- .filter(r -> !TimestampType.NO_TIMESTAMP_TYPE.equals(record.timestampType()))
- .map(ConsumerRecord::timestamp)
- .filter(ts -> ts >= 0);
- }
-
- public static String createErrorMessage(String base, Throwable t) {
- String separator = ": ";
- String result = base;
- while (t != null) {
- if (t.getMessage() != null)
- result += separator + t.getMessage();
- t = t.getCause();
- }
-
- return result;
- }
-
- public static void addError(Config configValidationResult, String prop, Object value, String baseErrorMessage, Throwable error) {
- addError(configValidationResult, prop, value, createErrorMessage(baseErrorMessage, error));
- }
-
- public static void addError(Config configValidationResult, String prop, Object value, String errorMessage) {
- findOrCreate(configValidationResult, prop, value)
- .addErrorMessage(errorMessage);
- }
-
- public static ConfigValue findOrCreate(Config configValidationResult, String prop, Object value) {
- for (ConfigValue definedValue : configValidationResult.configValues()) {
- if (Objects.equals(prop, definedValue.name())) {
- return definedValue;
- }
- }
-
- ConfigValue newValue = value == null ? new ConfigValue(prop) : new ConfigValue(prop, value, new ArrayList<>(), new ArrayList<>());
- configValidationResult.configValues().add(newValue);
- return newValue;
- }
-
- public static Set<ReplicatedTopicPartition> assignedPartitions(String connector, int taskId, Map<String, String> connectorConfig) {
- String connectorClass = connectorConfig.get(CONNECTOR_CLASS_CONFIG);
- if (connectorClass == null) {
- log.warn("No connector class found for connector {}", connector);
- return null;
- }
-
- if (!isURep3(connectorClass)) {
- log.warn("Ignoring task {}-{} with unsupported connector type {}", connector, taskId, connectorClass);
- return null;
- }
-
- String srcKafka = connectorConfig.get(UReplicator3ConnectorConfig.SOURCE_CLUSTER_CONFIG);
- String dstKafka = connectorConfig.get(UReplicator3ConnectorConfig.TARGET_CLUSTER_CONFIG);
- return uRep3Assignment(connector, taskId, srcKafka, dstKafka);
- }
-
- public static boolean isURep3(Map<String, String> connectorConfig) {
- String connectorClass = connectorConfig.get(CONNECTOR_CLASS_CONFIG);
- if (connectorClass == null) {
- log.warn("No connector class found for connector");
- return false;
- }
-
- return isURep3(connectorClass);
- }
-
- public static boolean isURep3(String connectorClass) {
- return Set.of(
- "com.uber.data.kafka.connect.ureplicator3.UReplicator3Connector",
- "UReplicator3Connector",
- "UReplicator3"
- ).contains(connectorClass);
- }
-
- public static Set<ReplicatedTopicPartition> uRep3Assignment(
- String connector,
- int taskId,
- String srcKafka,
- String dstKafka
- ) {
- AssignmentStore store = AssignmentStores.get();
- return store.getAssignment(connector, taskId, srcKafka, dstKafka)
- .stream()
- .map(ap -> new ReplicatedTopicPartition(srcKafka, dstKafka, ap.source()))
- .collect(Collectors.toSet());
- }
-
- public static void awaitAll(
- Collection<? extends Future<?>> futures, Duration duration
- ) throws InterruptedException, TimeoutException, ExecutionException {
- Timer timer = Time.SYSTEM.timer(duration);
- for (Future<?> future : futures) {
- timer.update();
- long remaining = Math.max(1, timer.remainingMs());
- future.get(remaining, TimeUnit.MILLISECONDS);
- }
- }
-
- public static int compareLexicographically(TopicPartition first, TopicPartition second) {
- if (first == null)
- return second == null ? 0 : -1;
-
- if (second == null)
- return 1;
-
- int result = first.topic().compareTo(second.topic());
- if (result != 0)
- return result;
-
- return Integer.compare(first.partition(), second.partition());
- }
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/CachedWorkload.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/CachedWorkload.java
deleted file mode 100644
index 0a15f453a5d..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/CachedWorkload.java
+++ /dev/null
@@ -1,28 +0,0 @@
-package com.uber.data.kafka.connect.workload.common;
-
-import java.util.Objects;
-
-public record CachedWorkload(long lastUpdated, Workload workload) implements Comparable<CachedWorkload> {
-
- public CachedWorkload(long lastUpdated, Workload workload) {
- this.lastUpdated = lastUpdated;
- this.workload = Objects.requireNonNull(workload, "workload may not be null");
- }
-
- @Override
- public int compareTo(CachedWorkload that) {
- if (this == that)
- return 0;
-
- if (that == null)
- return 1;
-
- int result = this.workload.compareTo(that.workload);
- if (result != 0)
- return result;
-
- return Long.compare(this.lastUpdated, that.lastUpdated);
- }
-
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/Workload.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/Workload.java
deleted file mode 100644
index a6faaac8545..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/Workload.java
+++ /dev/null
@@ -1,70 +0,0 @@
-package com.uber.data.kafka.connect.workload.common;
-
-import java.util.Objects;
-
-/**
- * POJO that tracks two generic workload statistics: bytes per second, and messages per second.
- */
-public class Workload implements Comparable<Workload> {
-
- public static final Workload EMPTY = new Workload(0, 0);
-
- private final double bytesPerSecond;
- private final double messagesPerSecond;
-
- public Workload(double bytesPerSecond, double messagesPerSecond) {
- this.bytesPerSecond = bytesPerSecond;
- this.messagesPerSecond = messagesPerSecond;
- }
-
- public double getBytesPerSecond() {
- return bytesPerSecond;
- }
-
- public double getMessagesPerSecond() {
- return messagesPerSecond;
- }
-
- @Override
- public boolean equals(Object o) {
- if (o == null || getClass() != o.getClass()) return false;
- Workload workload = (Workload) o;
- return Double.compare(bytesPerSecond, workload.bytesPerSecond) == 0 && Double.compare(messagesPerSecond, workload.messagesPerSecond) == 0;
- }
-
- @Override
- public int hashCode() {
- return Objects.hash(bytesPerSecond, messagesPerSecond);
- }
-
- @Override
- public String toString() {
- return "bytesPerSecond=" + bytesPerSecond + ",messagesPerSecond=" + messagesPerSecond;
- }
-
- @Override
- public int compareTo(Workload that) {
- if (this == that)
- return 0;
-
- if (that == null)
- return 1;
-
- int cmp = Double.compare(this.bytesPerSecond, that. bytesPerSecond);
- if (cmp != 0)
- return cmp;
-
- cmp = Double.compare(this.messagesPerSecond, that.messagesPerSecond);
- return cmp;
- }
-
- public Workload plus(Workload that) {
- if (that == null)
- return this;
-
- return new Workload(
- this.bytesPerSecond + that.bytesPerSecond,
- this.messagesPerSecond + that.messagesPerSecond
- );
- }
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/WorkloadCache.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/WorkloadCache.java
deleted file mode 100644
index f1eada2cb82..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/WorkloadCache.java
+++ /dev/null
@@ -1,221 +0,0 @@
-package com.uber.data.kafka.connect.workload.common;
-
-import org.apache.kafka.common.utils.Time;
-
-import com.google.common.annotations.VisibleForTesting;
-import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.LinkedList;
-import java.util.List;
-import java.util.Map;
-import java.util.Objects;
-import java.util.Optional;
-import java.util.Set;
-
-
-/**
- * Cache for topic workloads, with logic for automatic eviction of stale workloads and calculating
- * aggregate workloads when multiple samples from different time windows are available.
- * <p>
- * This class is fully thread-safe: all methods may be safely invoked concurrently.
- */
-public class WorkloadCache {
-
- private static final Logger log = LoggerFactory.getLogger(WorkloadCache.class);
-
- private final Time time;
- private final long maxAgeMs;
- private final long estimationLookbackWindowMs;
- private final Map<ReplicatedTopicPartition, LinkedList<CachedWorkload>> cachedWorkloads;
-
- /**
- * Create a new cache.
- * @param maxAgeMs the maximum age of any workloads before they are expired from the cache
- * @param estimationLookbackWindowMs the preferred age of workloads when computing the
- * aggregate workload for a topic or topic partition
- */
- public WorkloadCache(
- long maxAgeMs,
- long estimationLookbackWindowMs
- ) {
- this(maxAgeMs, estimationLookbackWindowMs, Time.SYSTEM);
- }
-
- /**
- * Create a new cache.
- * @param maxAgeMs the maximum age of any workloads before they are expired from the cache
- * @param estimationLookbackWindowMs the preferred age of workloads when computing the
- * aggregate workload for a topic or topic partition
- * @param time pluggable time interface; useful for mocking; may not be null
- */
- @VisibleForTesting
- WorkloadCache(
- long maxAgeMs,
- long estimationLookbackWindowMs,
- Time time
- ) {
- this.time = time;
- this.maxAgeMs = maxAgeMs;
- this.estimationLookbackWindowMs = estimationLookbackWindowMs;
- this.cachedWorkloads = new HashMap<>();
- }
-
- /**
- * Add a workload for a replicated topic to the cache. Stale workloads for the topic submitted here will
- * also be automatically expired from the cache.
- * @param replicatedTopicPartition the replicated topic partition; may not be null
- * @param workload the topic workload; may not be null
- * @param lastUpdated the time when the workload was recorded
- */
- public void put(ReplicatedTopicPartition replicatedTopicPartition, Workload workload, long lastUpdated) {
- CachedWorkload cachedWorkload = new CachedWorkload(lastUpdated, workload);
- LinkedList<CachedWorkload> workloads;
-
- if (isStale(cachedWorkload)) {
- // Don't bother inserting the workload if it's already stale
- synchronized (this) {
- workloads = cachedWorkloads.get(replicatedTopicPartition);
- }
- } else {
- synchronized (this) {
- workloads = cachedWorkloads.computeIfAbsent(replicatedTopicPartition, ignored -> new LinkedList<>());
- }
- synchronized (workloads) {
- if (workloads.isEmpty() || workloads.getLast().lastUpdated() < cachedWorkload.lastUpdated()) {
- workloads.add(cachedWorkload);
- }
- }
- }
-
- if (workloads == null) {
- return;
- }
-
- synchronized (workloads) {
- // purge the data points out of the valid window
- while (!workloads.isEmpty() && isStale(workloads.getFirst())) {
- long workloadTimestamp = workloads.getFirst().lastUpdated();
- log.trace("Expiring old workload for {} with timestamp of {}", replicatedTopicPartition, workloadTimestamp);
- workloads.removeFirst();
- }
- }
- }
-
- /**
- * Gets the workload for a replicated topic partition
- * @param replicatedTopicPartition the replicated topic partition; may not be null
- * @return the cached and aggregated workload for the replicated topic partition, if one can be found
- */
- public Optional<Workload> get(ReplicatedTopicPartition replicatedTopicPartition) {
- Objects.requireNonNull(replicatedTopicPartition, "replicated topic partition may not be null");
- LinkedList<CachedWorkload> allCachedWorkloads = workloadSnapshot(replicatedTopicPartition);
- if (allCachedWorkloads == null) {
- log.debug("There is no workload for {}.", replicatedTopicPartition);
- return Optional.empty();
- }
-
- // we will return the max workload during the valid window
- long current = time.milliseconds();
- long lookbackWindow = lookbackWindow(current, allCachedWorkloads);
-
- CachedWorkload maxWorkload = null;
- for (CachedWorkload workload : allCachedWorkloads) {
- if (current - workload.lastUpdated() > lookbackWindow) {
- continue;
- }
-
- if (maxWorkload == null || maxWorkload.compareTo(workload) < 0) {
- maxWorkload = workload;
- }
- }
-
- return Optional.ofNullable(maxWorkload).map(CachedWorkload::workload);
- }
-
- /**
- * Gets the workloads of all known replicated topic partitions
- * @return a map from replicated topic partition to workload; may be empty, but never null
- */
- public Map<ReplicatedTopicPartition, Workload> getAll() {
- Set<ReplicatedTopicPartition> allPartitions;
- synchronized (this) {
- allPartitions = new HashSet<>(cachedWorkloads.keySet());
- }
-
- Map<ReplicatedTopicPartition, Workload> result = new HashMap<>();
- for (ReplicatedTopicPartition replicatedTopicPartition : allPartitions) {
- Optional<Workload> workload = get(replicatedTopicPartition);
- workload.ifPresent(w -> result.put(replicatedTopicPartition, w));
- }
- return result;
- }
-
- /**
- * Gets all cached workloads for a replicated topic partition
- * @param replicatedTopicPartition the replicated topic partition; may not be null
- * @return all cached workloads for the replicated topic partition; may be null, and may be empty
- */
- public List<CachedWorkload> getHistorical(ReplicatedTopicPartition replicatedTopicPartition) {
- return workloadSnapshot(replicatedTopicPartition);
- }
-
- /**
- * Removes all cached workloads for a replicated topic partition
- * @param replicatedTopicPartition the replicated topic partition; may not be null
- * @return whether any workloads were removed for the specified partition
- */
- public boolean remove(ReplicatedTopicPartition replicatedTopicPartition) {
- synchronized (this) {
- return cachedWorkloads.remove(replicatedTopicPartition) != null;
- }
- }
-
- /**
- * Clear all entries from the cache
- */
- public void clear() {
- synchronized (this) {
- cachedWorkloads.clear();
- }
- }
-
- private long lookbackWindow(long current, Collection<CachedWorkload> workloads) {
- // if there are workloads inside the estimation lookback window (e.g., 1 hour), then we only consider those
- boolean recentWorkloadsAvailable = workloads.stream()
- .map(CachedWorkload::lastUpdated)
- .anyMatch(lastUpdated -> current - lastUpdated < estimationLookbackWindowMs);
- return recentWorkloadsAvailable ? estimationLookbackWindowMs : maxAgeMs;
- }
-
- private boolean isStale(CachedWorkload workload) {
- return time.milliseconds() - workload.lastUpdated() > maxAgeMs;
- }
-
- /**
- * Take a snapshot of the latest-known workloads for a replicated topic partition. This snapshot will
- * belong solely to the caller and is guaranteed not to be modified by other threads or side effects
- * of other modifications made to the cache.
- * @param replicatedTopicPartition the replicated topic partition; may not be null
- * @return the workload snapshot for the partition; may be empty, may be null
- */
- private LinkedList<CachedWorkload> workloadSnapshot(ReplicatedTopicPartition replicatedTopicPartition) {
- List<CachedWorkload> cachedWorkloads;
- synchronized (this) {
- cachedWorkloads = this.cachedWorkloads.get(replicatedTopicPartition);
- if (cachedWorkloads == null) {
- return null;
- }
- }
-
- synchronized (cachedWorkloads) {
- return new LinkedList<>(cachedWorkloads);
- }
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/WorkloadProvider.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/WorkloadProvider.java
deleted file mode 100644
index bab285169a0..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/common/WorkloadProvider.java
+++ /dev/null
@@ -1,39 +0,0 @@
-package com.uber.data.kafka.connect.workload.common;
-
-import org.apache.kafka.common.Configurable;
-import org.apache.kafka.connect.components.Versioned;
-
-import java.io.Closeable;
-import java.util.Map;
-
-public interface WorkloadProvider extends Versioned, Configurable, Closeable {
-
- /**
- * Determine and return the load for the given task, if it is known
- * @return the load for the task, or null if none is known
- */
- Double taskLoad(String connector, int taskId, Map<String, String> connectorConfig, Map<String, String> taskConfig);
-
- static WorkloadProvider noOp() {
- return new WorkloadProvider() {
- @Override
- public Double taskLoad(String connector, int taskId, Map<String, String> connectorConfig, Map<String, String> taskConfig) {
- return null;
- }
-
- @Override
- public void close() {
- }
-
- @Override
- public void configure(Map<String, ?> configs) {
- }
-
- @Override
- public String version() {
- return "0.1.0";
- }
- };
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadProvider.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadProvider.java
deleted file mode 100644
index e555de970ba..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadProvider.java
+++ /dev/null
@@ -1,132 +0,0 @@
-package com.uber.data.kafka.connect.workload.kafka;
-
-import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.utils.AppInfoParser;
-import org.apache.kafka.common.utils.Utils;
-
-import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
-import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStores;
-import com.uber.data.kafka.connect.workload.common.CachedWorkload;
-import com.uber.data.kafka.connect.workload.common.Workload;
-import com.uber.data.kafka.connect.workload.common.WorkloadProvider;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
-import java.util.concurrent.Future;
-import java.util.concurrent.atomic.AtomicBoolean;
-
-import static com.uber.data.kafka.connect.utils.UberConnectUtils.assignedPartitions;
-
-public class KafkaBasedWorkloadProvider implements WorkloadProvider {
-
- private static final AtomicBoolean DESTROY_ON_CLOSE = new AtomicBoolean(false);
- private static final Logger log = LoggerFactory.getLogger(KafkaBasedWorkloadProvider.class);
-
- private static KafkaBasedWorkloadStore workloadStore;
-
- private KafkaBasedWorkloadProviderConfig config;
-
- @Override
- public void configure(Map<String, ?> props) {
- log.info("Starting provider");
- config = new KafkaBasedWorkloadProviderConfig(props);
- initializeStore(config);
- AssignmentStores.initialize(config.assignmentStoreConfig());
- log.info("Finished Starting provider");
- }
-
- @Override
- public Double taskLoad(String connector, int taskId, Map<String, String> connectorConfig, Map<String, String> taskConfig) {
- Set<ReplicatedTopicPartition> replicatedTopicPartitions = assignedPartitions(connector, taskId, connectorConfig);
- if (replicatedTopicPartitions == null) {
- return null;
- }
-
- return replicatedTopicPartitions.stream()
- .map(rtp -> workloadStore.getWorkload(rtp.srcKafka(), rtp.dstKafka(), rtp.topicPartition()))
- .map(w -> w != null ? w : config.defaultPartitionWorkload())
- .mapToDouble(this::scaleWorkload)
- .sum();
- }
-
- @Override
- public void close() {
- if (!DESTROY_ON_CLOSE.get()) {
- log.debug("Skipping teardown of workload store on close; this should only be done in testing environments");
- return;
- }
- log.info("Shutting down provider");
- destroyStore();
- log.info("Finished shutting down provider");
- }
-
- @Override
- public String version() {
- return AppInfoParser.getVersion();
- }
-
- public static Future<?> publishWorkload(String srcKafka, String dstKafka, TopicPartition topicPartition, Workload workload) {
- assertStoreInitialized();
-
- return workloadStore.putWorkload(srcKafka, dstKafka, topicPartition, workload);
- }
-
- public static void refreshWorkloadCache() {
- assertStoreInitialized();
-
- AssignmentStores.get().refresh();
- workloadStore.refresh();
- }
-
- public static Map<ReplicatedTopicPartition, Workload> getAllWorkloads() {
- assertStoreInitialized();
-
- return workloadStore.getAllWorkloads();
- }
-
- public static List<CachedWorkload> getHistoricalWorkloads(String srcKafka, String dstKafka, TopicPartition topicPartition) {
- assertStoreInitialized();
-
- return workloadStore.getHistoricalWorkloads(srcKafka, dstKafka, topicPartition);
- }
-
- public static void clearCache() {
- assertStoreInitialized();
-
- workloadStore.clear();
- }
-
- public static void assertStoreInitialized() throws IllegalStateException {
- if (workloadStore == null)
- throw new IllegalStateException("Workload store has not been initialized yet, or has already been shut down");
- }
-
- public static void destroyOnClose(boolean destroy) {
- DESTROY_ON_CLOSE.set(destroy);
- }
-
- private static synchronized void initializeStore(KafkaBasedWorkloadProviderConfig config) {
- if (workloadStore != null) {
- log.warn("Provider appears to have been configured multiple times on the same JVM. This should only happen during testing");
- return;
- }
- workloadStore = new KafkaBasedWorkloadStore(config);
- workloadStore.start();
- }
-
- private static synchronized void destroyStore() {
- Utils.closeQuietly(workloadStore, "workload store");
- workloadStore = null;
- }
-
- private double scaleWorkload(Workload taskWorkload) {
- double byteRateLoad = taskWorkload.getBytesPerSecond() / config.maxWorkerWorkload().getBytesPerSecond();
- double messageRateLoad = taskWorkload.getMessagesPerSecond() / config.maxWorkerWorkload().getMessagesPerSecond();
- return Math.max(byteRateLoad, messageRateLoad);
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadProviderConfig.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadProviderConfig.java
deleted file mode 100644
index acbd0d181fb..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadProviderConfig.java
+++ /dev/null
@@ -1,225 +0,0 @@
-package com.uber.data.kafka.connect.workload.kafka;
-
-import org.apache.kafka.common.config.AbstractConfig;
-import org.apache.kafka.common.config.ConfigDef;
-import org.apache.kafka.common.config.ConfigDef.Importance;
-import org.apache.kafka.common.config.ConfigDef.Type;
-import org.apache.kafka.common.config.TopicConfig;
-
-import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStoreConfig;
-import com.uber.data.kafka.connect.workload.common.Workload;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.HashMap;
-import java.util.Map;
-
-import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
-
-public class KafkaBasedWorkloadProviderConfig extends AbstractConfig {
-
- private static final Logger log = LoggerFactory.getLogger(KafkaBasedWorkloadProviderConfig.class);
-
- private static final String CONFIG_PREFIX = "uber.workload.provider.";
-
- // TODO: Flipr
- public static final String MAX_WORKER_BYTE_RATE_CONFIG = CONFIG_PREFIX + "max.worker.byte.rate";
- public static final double MAX_WORKER_BYTE_RATE_DEFAULT = 70 * 1024 * 1024;
- private static final String MAX_WORKER_BYTE_RATE_DOC = "Maximum total byte rate for workloads assigned "
- + "to each worker before attempting to revoke workloads and redistribute to other workers";
-
- // TODO: Flipr
- public static final String MAX_WORKER_MESSAGE_RATE_CONFIG = CONFIG_PREFIX + "max.worker.message.rate";
- public static final double MAX_WORKER_MESSAGE_RATE_DEFAULT = 70_000;
- private static final String MAX_WORKER_MESSAGE_RATE_DOC = "Maximum total message rate for workloads assigned "
- + "to each worker before attempting to revoke workloads and redistribute to other workers";
-
- // TODO: Flipr
- public static final String DEFAULT_PARTITION_BYTE_RATE_CONFIG = CONFIG_PREFIX + "default.partition.byte.rate";
- public static final double DEFAULT_PARTITION_BYTE_RATE_DEFAULT = 1024;
- private static final String DEFAULT_PARTITION_BYTE_RATE_DOC = "The default byte rate to assume for partitions "
- + "that have no workload information available (likely because they have only recently started being replicated)";
-
- // TODO: Flipr
- public static final String DEFAULT_PARTITION_MESSAGE_RATE_CONFIG = CONFIG_PREFIX + "default.partition.message.rate";
- public static final double DEFAULT_PARTITION_MESSAGE_RATE_DEFAULT = 10;
- private static final String DEFAULT_PARTITION_MESSAGE_RATE_DOC = "The default message rate to assume for partitions "
- + "that have no workload information available (likely because they have only recently started being replicated)";
-
- public static final String STORAGE_TOPIC_CONFIG = CONFIG_PREFIX + "storage.topic";
- public static final String STORAGE_TOPIC_PREFIX = STORAGE_TOPIC_CONFIG + ".";
- private static final String STORAGE_TOPIC_DOC = "Topic used to store workload information. "
- + "This topic must be compacted (i.e., have the 'compacted' cleanup policy) and can have "
- + "any number of partitions, but it's recommended to use a high number (e.g., 25 or 50). "
- + "If the topic does not exist, it will be automatically created on startup. "
- + "Additional properties for the topic can be specified using the " + STORAGE_TOPIC_PREFIX + " "
- + "prefix";
-
- private static final String PARTITIONS_SUFFIX = "partitions";
- public static final String STORAGE_TOPIC_PARTITIONS_CONFIG = STORAGE_TOPIC_PREFIX + PARTITIONS_SUFFIX;
- private static final int STORAGE_TOPIC_PARTITIONS_DEFAULT = 50;
- private static final String STORAGE_TOPIC_PARTITIONS_DOC = "The number of partitions to create "
- + "the workload topic with, if it does not already exist. If the topic already exists this property "
- + "will be ignored, even if the topic has a different number of partitions";
-
- private static final String REPLICATION_FACTOR_SUFFIX = "replication.factor";
- public static final String STORAGE_TOPIC_REPLICATION_FACTOR_CONFIG = STORAGE_TOPIC_PREFIX + REPLICATION_FACTOR_SUFFIX;
- private static final short STORAGE_TOPIC_REPLICATION_FACTOR_DEFAULT = (short) 3;
- private static final String STORAGE_TOPIC_REPLICATION_FACTOR_DOC = "The replication factor to create "
- + "the workload topic with, if it does not already exist. If the topic already exists this property "
- + "will be ignored, even if the topic has a different replication factor";
-
- private static final String SEGMENT_BYTES_SUFFIX = TopicConfig.SEGMENT_BYTES_CONFIG;
- public static final String STORAGE_TOPIC_SEGMENT_BYTES_CONFIG = STORAGE_TOPIC_PREFIX + SEGMENT_BYTES_SUFFIX;
- private static final int STORAGE_TOPIC_SEGMENT_BYTES_DEFAULT = 1024 * 1024 * 3; // Lower default segment size in order to enable more-frequent compaction
- private static final String STORAGE_TOPIC_SEGMENT_BYTES_DOC = "The per-partition segment size "
- + "for the workload topic. This should be relatively low in order to enable topic compaction to take place, which "
- + "will help keep the topic from growing too large and causing slow startup time for workers. Like other topic "
- + "properties, if the topic already exists, this property will be ignored, even if the topic has a different segment size";
-
- public static final String CLIENTS_PREFIX = CONFIG_PREFIX + "kafka.clients";
- public static final String CONSUMER_PREFIX = CONFIG_PREFIX + "consumer";
- public static final String PRODUCER_PREFIX = CONFIG_PREFIX + "producer";
- public static final String ADMIN_PREFIX = CONFIG_PREFIX + "admin";
-
- private final Workload maxWorkerWorkload;
- private final Workload defaultPartitionWorkload;
-
- private final AssignmentStoreConfig assignmentStoreConfig;
-
- @SuppressWarnings("this-escape")
- public KafkaBasedWorkloadProviderConfig(Map<String, ?> props) {
- super(config(), props);
- this.maxWorkerWorkload = new Workload(
- getDouble(MAX_WORKER_BYTE_RATE_CONFIG),
- getDouble(MAX_WORKER_MESSAGE_RATE_CONFIG)
- );
- this.defaultPartitionWorkload = new Workload(
- getDouble(DEFAULT_PARTITION_BYTE_RATE_CONFIG),
- getDouble(DEFAULT_PARTITION_MESSAGE_RATE_CONFIG)
- );
- this.assignmentStoreConfig = new AssignmentStoreConfig(props);
- }
-
- private static ConfigDef config() {
- ConfigDef result = new ConfigDef()
- .define(
- STORAGE_TOPIC_CONFIG,
- Type.STRING,
- Importance.HIGH,
- STORAGE_TOPIC_DOC
- ).define(
- MAX_WORKER_BYTE_RATE_CONFIG,
- Type.DOUBLE,
- MAX_WORKER_BYTE_RATE_DEFAULT,
- atLeast(1),
- Importance.MEDIUM,
- MAX_WORKER_BYTE_RATE_DOC
- ).define(
- MAX_WORKER_MESSAGE_RATE_CONFIG,
- Type.DOUBLE,
- MAX_WORKER_MESSAGE_RATE_DEFAULT,
- atLeast(1),
- Importance.MEDIUM,
- MAX_WORKER_MESSAGE_RATE_DOC
- ).define(
- DEFAULT_PARTITION_BYTE_RATE_CONFIG,
- Type.DOUBLE,
- DEFAULT_PARTITION_BYTE_RATE_DEFAULT,
- atLeast(0),
- Importance.MEDIUM,
- DEFAULT_PARTITION_BYTE_RATE_DOC
- ).define(
- DEFAULT_PARTITION_MESSAGE_RATE_CONFIG,
- Type.DOUBLE,
- DEFAULT_PARTITION_MESSAGE_RATE_DEFAULT,
- atLeast(0),
- Importance.MEDIUM,
- DEFAULT_PARTITION_MESSAGE_RATE_DOC
- ).define(
- STORAGE_TOPIC_PARTITIONS_CONFIG,
- Type.INT,
- STORAGE_TOPIC_PARTITIONS_DEFAULT,
- atLeast(1),
- Importance.LOW,
- STORAGE_TOPIC_PARTITIONS_DOC
- ).define(
- STORAGE_TOPIC_REPLICATION_FACTOR_CONFIG,
- Type.SHORT,
- STORAGE_TOPIC_REPLICATION_FACTOR_DEFAULT,
- atLeast(1),
- Importance.LOW,
- STORAGE_TOPIC_REPLICATION_FACTOR_DOC
- ).define(
- STORAGE_TOPIC_SEGMENT_BYTES_CONFIG,
- Type.INT,
- STORAGE_TOPIC_SEGMENT_BYTES_DEFAULT,
- atLeast(1),
- Importance.LOW,
- STORAGE_TOPIC_SEGMENT_BYTES_DOC
- );
- AssignmentStoreConfig.addToConfig(result);
- return result;
- }
-
- public AssignmentStoreConfig assignmentStoreConfig() {
- return assignmentStoreConfig;
- }
-
- public Map<String, Object> adminConfig() {
- return clientConfig(ADMIN_PREFIX);
- }
-
- public Map<String, Object> producerConfig() {
- return clientConfig(PRODUCER_PREFIX);
- }
-
- public Map<String, Object> consumerConfig() {
- return clientConfig(CONSUMER_PREFIX);
- }
-
- public String storageTopic() {
- return getString(STORAGE_TOPIC_CONFIG);
- }
-
- public Workload maxWorkerWorkload() {
- return maxWorkerWorkload;
- }
-
- public Workload defaultPartitionWorkload() {
- return defaultPartitionWorkload;
- }
-
- public int storageTopicPartitions() {
- return getInt(STORAGE_TOPIC_PARTITIONS_CONFIG);
- }
-
- public short storageTopicReplicationFactor() {
- return getShort(STORAGE_TOPIC_REPLICATION_FACTOR_CONFIG);
- }
-
- public Map<String, Object> storageTopicSettings() {
- Map<String, Object> result = originalsWithPrefix(STORAGE_TOPIC_PREFIX);
-
- // AbstractConfig::originalsWithPrefix doesn't include default values, we apply that manually here
- result.computeIfAbsent(TopicConfig.SEGMENT_BYTES_CONFIG, k -> STORAGE_TOPIC_SEGMENT_BYTES_DEFAULT);
-
- Object removedPolicy = result.remove(TopicConfig.CLEANUP_POLICY_CONFIG);
- if (removedPolicy != null) {
- log.warn("Ignoring '{}cleanup.policy={}' setting, since compaction is always used", STORAGE_TOPIC_PREFIX, removedPolicy);
- }
-
- result.remove(REPLICATION_FACTOR_SUFFIX);
- result.remove(PARTITIONS_SUFFIX);
-
- return result;
- }
-
- private Map<String, Object> clientConfig(String specialPrefix) {
- Map<String, Object> result = new HashMap<>(originals());
- result.putAll(originalsWithPrefix(CLIENTS_PREFIX));
- result.putAll(originalsWithPrefix(specialPrefix));
- return result;
- }
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadStore.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadStore.java
deleted file mode 100644
index 92d8246b95f..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadStore.java
+++ /dev/null
@@ -1,286 +0,0 @@
-package com.uber.data.kafka.connect.workload.kafka;
-
-import org.apache.kafka.clients.admin.NewTopic;
-import org.apache.kafka.clients.consumer.ConsumerConfig;
-import org.apache.kafka.clients.consumer.ConsumerRecord;
-import org.apache.kafka.clients.producer.ProducerConfig;
-import org.apache.kafka.common.KafkaException;
-import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.internals.KafkaCompletableFuture;
-import org.apache.kafka.common.serialization.ByteArrayDeserializer;
-import org.apache.kafka.common.serialization.ByteArraySerializer;
-import org.apache.kafka.common.utils.Time;
-import org.apache.kafka.common.utils.Utils;
-import org.apache.kafka.connect.util.KafkaBasedLog;
-import org.apache.kafka.connect.util.TopicAdmin;
-
-import com.fasterxml.jackson.core.JsonProcessingException;
-import com.fasterxml.jackson.core.type.TypeReference;
-import com.fasterxml.jackson.databind.ObjectMapper;
-import com.google.common.annotations.VisibleForTesting;
-import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
-import com.uber.data.kafka.connect.utils.UberConnectUtils;
-import com.uber.data.kafka.connect.workload.common.CachedWorkload;
-import com.uber.data.kafka.connect.workload.common.Workload;
-import com.uber.data.kafka.connect.workload.common.WorkloadCache;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.Closeable;
-import java.io.IOException;
-import java.util.LinkedHashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.Future;
-import java.util.concurrent.TimeUnit;
-
-import static com.uber.data.kafka.connect.utils.SerializationUtils.parseMapField;
-import static com.uber.data.kafka.connect.utils.SerializationUtils.parseNumberField;
-import static com.uber.data.kafka.connect.utils.SerializationUtils.parseStringField;
-
-class KafkaBasedWorkloadStore implements Closeable {
-
- private static final Logger log = LoggerFactory.getLogger(KafkaBasedWorkloadStore.class);
- private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
-
- @VisibleForTesting
- static final long MAX_VALID_TIME_WINDOW_IN_MS = TimeUnit.HOURS.toMillis(25);
- @VisibleForTesting
- static final long ESTIMATION_LOOKBACK_WINDOW_IN_MS = TimeUnit.HOURS.toMillis(1);
-
- private static final String WORKLOAD_FIELD = "workload-v1";
- private static final String SRC_KAFKA_FIELD = "src";
- private static final String DST_KAFKA_FIELD = "dst";
- private static final String TOPIC_FIELD = "topic";
- private static final String PARTITION_FIELD = "partition";
- private static final String BYTE_RATE_FIELD = "bytes";
- private static final String MESSAGE_RATE_FIELD = "messages";
-
- private final String topic;
- private final int topicPartitions;
- private final short topicReplicationFactor;
- private final Map<String, Object> topicSettings;
- private final Map<String, Object> producerProps;
- private final Map<String, Object> consumerProps;
- private final Map<String, Object> adminProps;
- private final WorkloadCache workloadsCache;
- private final Time time;
-
- private TopicAdmin topicAdmin;
- private KafkaBasedLog<byte[], byte[]> kafkaLog;
-
- public KafkaBasedWorkloadStore(KafkaBasedWorkloadProviderConfig config) {
- this(config, Time.SYSTEM);
- }
-
- @VisibleForTesting
- KafkaBasedWorkloadStore(KafkaBasedWorkloadProviderConfig config, Time time) {
- this.topic = config.storageTopic();
- this.topicPartitions = config.storageTopicPartitions();
- this.topicReplicationFactor = config.storageTopicReplicationFactor();
- this.topicSettings = config.storageTopicSettings();
- this.producerProps = producerProps(config.producerConfig());
- this.consumerProps = consumerProps(config.consumerConfig());
- this.adminProps = config.adminConfig();
- this.workloadsCache = new WorkloadCache(MAX_VALID_TIME_WINDOW_IN_MS, ESTIMATION_LOOKBACK_WINDOW_IN_MS);
- this.time = time;
- }
-
- public void start() {
- this.topicAdmin = new TopicAdmin(adminProps);
- this.kafkaLog = new KafkaBasedLog<>(
- topic,
- producerProps,
- consumerProps,
- () -> topicAdmin,
- this::onConsume,
- Time.SYSTEM,
- this::createAndValidateTopic
- );
- kafkaLog.start();
- }
-
- private static Map<String, Object> producerProps(Map<String, Object> baseProps) {
- baseProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
- baseProps.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
- baseProps.put(ProducerConfig.DELIVERY_TIMEOUT_MS_CONFIG, Integer.MAX_VALUE);
- return baseProps;
- }
-
- private static Map<String, Object> consumerProps(Map<String, Object> baseProps) {
- baseProps.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
- baseProps.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class.getName());
- return baseProps;
- }
-
- private void createAndValidateTopic(TopicAdmin topicAdmin) {
- NewTopic newTopic = TopicAdmin.defineTopic(topic)
- .config(topicSettings)
- .compacted()
- .partitions(topicPartitions)
- .replicationFactor(topicReplicationFactor)
- .build();
- boolean created = topicAdmin.createTopic(newTopic);
- if (!created) {
- // Ensure existing topic is compacted
- boolean verified = topicAdmin.verifyTopicCleanupPolicyOnlyCompact(
- topic,
- KafkaBasedWorkloadProviderConfig.STORAGE_TOPIC_CONFIG,
- "workloads"
- );
- // Should never, NEVER happen, but just in case
- if (!verified) {
- throw new KafkaException("Failed to verify that workload storage topic " + topic + " is compacted");
- }
- }
- }
-
- @Override
- public void close() {
- if (kafkaLog != null) {
- Utils.closeQuietly(kafkaLog::stop, "Kafka-based log");
- }
- Utils.closeQuietly(topicAdmin, "topic admin");
- }
-
- public Future<?> putWorkload(String srcKafka, String dstKafka, TopicPartition topicPartition, Workload workload) {
- ReplicatedTopicPartition replicatedTopicPartition = new ReplicatedTopicPartition(srcKafka, dstKafka, topicPartition);
-
- byte[] key;
- try {
- key = serializeWorkloadKey(replicatedTopicPartition);
- } catch (JsonProcessingException e) {
- log.error("Failed to serialize workload key for replicated topic partition {} as JSON", replicatedTopicPartition, e);
- return KafkaCompletableFuture.failedFuture(e);
- }
-
- byte[] value;
- try {
- value = serializeWorkloadValue(workload);
- } catch (JsonProcessingException e) {
- log.error("Failed to serialize workload value {} for replicated topic partition {} as JSON", replicatedTopicPartition, e);
- return KafkaCompletableFuture.failedFuture(e);
- }
-
- // TODO: Error handling (probably just log and emit a metric)
- return kafkaLog.sendWithReceipt(key, value);
- }
-
- public Workload getWorkload(String srcKafka, String dstKafka, TopicPartition topicPartition) {
- ReplicatedTopicPartition replicatedTopicPartition = new ReplicatedTopicPartition(srcKafka, dstKafka, topicPartition);
- return workloadsCache.get(replicatedTopicPartition).orElse(null);
- }
-
- public Map<ReplicatedTopicPartition, Workload> getAllWorkloads() {
- return workloadsCache.getAll();
- }
-
- public List<CachedWorkload> getHistoricalWorkloads(String srcKafka, String dstKafka, TopicPartition topicPartition) {
- ReplicatedTopicPartition replicatedTopicPartition = new ReplicatedTopicPartition(srcKafka, dstKafka, topicPartition);
- return workloadsCache.getHistorical(replicatedTopicPartition);
- }
-
- public void refresh() {
- try {
- kafkaLog.readToEnd().get();
- } catch (InterruptedException | ExecutionException e) {
- throw new KafkaException("Failed to read to end of workload log", e);
- }
- }
-
- public void clear() {
- workloadsCache.clear();
- }
-
- private void onConsume(Throwable error, ConsumerRecord<byte[], byte[]> record) {
- if (error != null) {
- log.error("Error occurred while consuming from workload topic", error);
- return;
- }
-
- int recordPartition = record.partition();
- long recordOffset = record.offset();
-
- ReplicatedTopicPartition replicatedTopicPartition;
- try {
- replicatedTopicPartition = deserializeWorkloadKey(record.key());
- } catch (IOException e) {
- log.error("Failed to deserialize replicated topic partition in partition {} at offset {} as JSON", recordPartition, recordOffset, e);
- return;
- }
-
- Workload workload;
- try {
- workload = deserializeWorkloadValue(record.value());
- } catch (IOException e) {
- log.error("Failed to deserialize workload in partition {} at offset {} for replicated topic partition {} as JSON", recordPartition, recordOffset, replicatedTopicPartition, e);
- return;
- }
-
- if (workload != null) {
- long timestamp = UberConnectUtils.extractTimestamp(record)
- .orElseGet(() -> {
- log.warn("No timestamp found for record in partition {} at offset {}; using current wall-clock time instead", recordPartition, recordOffset);
- return time.milliseconds();
- });
-
- workloadsCache.put(replicatedTopicPartition, workload, timestamp);
- } else {
- workloadsCache.remove(replicatedTopicPartition);
- }
- }
-
- private byte[] serializeWorkloadKey(ReplicatedTopicPartition replicatedTopicPartition) throws JsonProcessingException {
- // LinkedHashMap preserves order and allows us to deterministically generate a serialized key from a replicated topic partition
- Map<String, Object> innerMap = new LinkedHashMap<>();
- innerMap.put(SRC_KAFKA_FIELD, replicatedTopicPartition.srcKafka());
- innerMap.put(DST_KAFKA_FIELD, replicatedTopicPartition.dstKafka());
- innerMap.put(TOPIC_FIELD, replicatedTopicPartition.topicPartition().topic());
- innerMap.put(PARTITION_FIELD, replicatedTopicPartition.topicPartition().partition());
- Map<String, Object> wrapped = Map.of(WORKLOAD_FIELD, innerMap);
- return OBJECT_MAPPER.writeValueAsBytes(wrapped);
- }
-
- private ReplicatedTopicPartition deserializeWorkloadKey(byte[] key) throws IOException {
- if (key == null)
- throw new IOException("record has no key");
-
- Map<String, Object> map = OBJECT_MAPPER.readValue(key, new TypeReference<>() { });
- Map<?, ?> wrapped = parseMapField(true, map, WORKLOAD_FIELD);
-
- String srcKafka = parseStringField(true, wrapped, SRC_KAFKA_FIELD);
- String dstKafka = parseStringField(true, wrapped, DST_KAFKA_FIELD);
- String topic = parseStringField(true, wrapped, TOPIC_FIELD);
- Number partition = parseNumberField(true, wrapped, PARTITION_FIELD);
-
- return new ReplicatedTopicPartition(srcKafka, dstKafka, topic, partition.intValue());
- }
-
- private byte[] serializeWorkloadValue(Workload workload) throws JsonProcessingException {
- // This allows us to write tombstone records for replicated topic partitions
- if (workload == null)
- return null;
-
- Map<String, Object> map = Map.of(
- BYTE_RATE_FIELD, workload.getBytesPerSecond(),
- MESSAGE_RATE_FIELD, workload.getMessagesPerSecond()
- );
- return OBJECT_MAPPER.writeValueAsBytes(map);
- }
-
- private Workload deserializeWorkloadValue(byte[] value) throws IOException {
- // Permit null values (tombstones) that indicate we should clear the cache for this replicated topic partition
- if (value == null)
- return null;
-
- Map<String, Object> map = OBJECT_MAPPER.readValue(value, new TypeReference<>() { });
-
- Number byteRate = parseNumberField(false, map, BYTE_RATE_FIELD);
- Number messageRate = parseNumberField(false, map, MESSAGE_RATE_FIELD);
-
- return new Workload(byteRate.doubleValue(), messageRate.doubleValue());
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadTrackingTransformation.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadTrackingTransformation.java
deleted file mode 100644
index 670a068252a..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/kafka/KafkaBasedWorkloadTrackingTransformation.java
+++ /dev/null
@@ -1,15 +0,0 @@
-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
deleted file mode 100644
index 577ae78d7a4..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadCollector.java
+++ /dev/null
@@ -1,103 +0,0 @@
-package com.uber.data.kafka.connect.workload.tracking;
-
-import org.apache.kafka.common.utils.Time;
-
-import com.google.common.annotations.VisibleForTesting;
-import com.uber.data.kafka.connect.workload.common.Workload;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.time.Duration;
-import java.util.Objects;
-import java.util.Optional;
-
-/**
- * Tracks segments for a single workload. Segments will be considered in-progress until a
- * user-specified duration has elapsed, at which point the segment will be considered complete, and
- * a new active segment will be started.
- */
-class WorkloadCollector {
-
- private static final Logger log = LoggerFactory.getLogger(WorkloadCollector.class);
-
- private final Time time;
- private final Duration maxSegmentAge;
-
- private long startMs;
- private long totalBytes;
- private long totalMessages;
-
- WorkloadCollector(Duration maxSegmentAge, Time time) {
- this.maxSegmentAge = Objects.requireNonNull(maxSegmentAge, "max segment age may not be null");
- this.time = Objects.requireNonNull(time, "time may not be null");
- reset();
- }
-
- /**
- * Record the workload for zero or more messages, and if there is an in-progress segment whose
- * age has exceeded the maximum segment age, return that segment's workload and begin a new
- * segment.
- * <p>
- * Note that if a {@link Optional#isPresent() present Optional} is returned, the
- * workload information contained in it will be dropped by this collector after the method
- * returns. It is the responsibility of the caller to store the workload information for as long as it is needed.
- * @param bytes the total bytes consumed for the messages; may not be negative
- * @param messages the total messages consumed; may not be negative
- * @return the workload for the previous segment, if one was just completed
- */
- public synchronized Optional<Workload> record(long bytes, long messages) {
- if (bytes < 0) {
- log.warn("Ignoring invalid workload with negative number of bytes: {}", bytes);
- return Optional.empty();
- }
- if (messages < 0) {
- log.warn("Ignoring invalid workload with negative number of messages: {}", messages);
- return Optional.empty();
- }
-
- Optional<Workload> result;
- Duration age = age();
-
- if (age.compareTo(maxSegmentAge) >= 0) {
- result = Optional.ofNullable(collect(totalBytes, totalMessages, age));
- reset();
- } else {
- result = Optional.empty();
- }
-
- this.totalBytes += bytes;
- this.totalMessages += messages;
-
- return result;
- }
-
- private Duration age() {
- return Duration.ofMillis(time.milliseconds() - startMs);
- }
-
- @VisibleForTesting
- static Workload collect(long totalBytes, long totalMessages, Duration age) {
- double durationSeconds = Math.ceil(age.toMillis() / 1000.0);
- if (durationSeconds <= 0) {
- log.warn(
- "Tried to collect workload from segment whose age in seconds ({}) is " +
- "non-positive; discarding this segment and starting fresh",
- durationSeconds
- );
- return null;
- }
-
- return new Workload(
- totalBytes / durationSeconds,
- totalMessages / durationSeconds
- );
- }
-
- private void reset() {
- this.startMs = time.milliseconds();
- this.totalBytes = 0;
- this.totalMessages = 0;
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadStore.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadStore.java
deleted file mode 100644
index 816b02867df..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadStore.java
+++ /dev/null
@@ -1,12 +0,0 @@
-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
deleted file mode 100644
index 00edf3581f6..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadTracker.java
+++ /dev/null
@@ -1,132 +0,0 @@
-package com.uber.data.kafka.connect.workload.tracking;
-
-import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.utils.Time;
-
-import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.Ticker;
-import com.google.common.cache.CacheBuilder;
-import com.google.common.cache.CacheLoader;
-import com.google.common.cache.LoadingCache;
-import com.google.common.cache.RemovalCause;
-import com.google.common.cache.RemovalListener;
-import com.google.common.cache.RemovalNotification;
-import com.uber.data.kafka.connect.workload.common.Workload;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.time.Duration;
-import java.util.HashMap;
-import java.util.Map;
-import java.util.Optional;
-import java.util.concurrent.TimeUnit;
-
-/**
- * Tracks {@link Workload workloads} on a per-topic-partition basis, allowing users to periodically
- * {@link #record(TopicPartition, long, long) record} new workload information for partitions and,
- * once workload information for a topic partition has matured, retrieve a complete snapshot of the
- * average per-second workload of that partition.
- */
-class WorkloadTracker {
-
- private static final Logger log = LoggerFactory.getLogger(WorkloadTracker.class);
-
- private final Duration maxSegmentAge;
- private final Time time;
- private final LoadingCache<TopicPartition, WorkloadCollector> activeSegments;
- private final Map<TopicPartition, Workload> staleWorkloads;
-
- public WorkloadTracker(Duration maxSegmentAge) {
- this(
- maxSegmentAge,
- Time.SYSTEM,
- Ticker.systemTicker()
- );
- }
-
- @VisibleForTesting
- WorkloadTracker(
- Duration maxSegmentAge,
- Time time,
- Ticker ticker
- ) {
- this.maxSegmentAge = maxSegmentAge;
- this.time = time;
-
- this.activeSegments = CacheBuilder.newBuilder()
- .ticker(ticker)
- .expireAfterAccess(maxSegmentAge.toMillis(), TimeUnit.MILLISECONDS)
- .removalListener(new ExpiredCollectorListener())
- .build(new WorkloadCollectorCacheLoader());
- this.staleWorkloads = new HashMap<>();
- }
-
- /**
- * Record the throughput for a single record, and if the maximum segment age has been reached
- * for the given topic partition, return the workload for its just-completed segment
- * @param topicPartition the topic partition that the record came from; may not be null
- * @param bytes the size of the record in bytes; may not be negative
- * @param messages the number of records; may not be negative
- * @return the workload for the current segment, if it has reached the maximum segment age
- */
- public Optional<Workload> record(TopicPartition topicPartition, long bytes, long messages) {
- return activeSegments.getUnchecked(topicPartition)
- .record(bytes, messages);
- }
-
- /**
- * Returns and then discards all stale workloads; that is, workloads for segments that have
- * reached or exceeded the maximum segment age, and for which no throughputs have recently been
- * recorded. Should be invoked periodically in order to prevent memory leaks.
- * @return all known stale workloads; may be empty, but never null
- */
- public Map<TopicPartition, Workload> collectStaleWorkloads() {
- activeSegments.cleanUp();
-
- Map<TopicPartition, Workload> result;
- synchronized (staleWorkloads) {
- result = new HashMap<>(staleWorkloads);
- staleWorkloads.clear();
- }
-
- return result;
- }
-
- /**
- * Clear all workloads, both normal and stale.
- */
- public void clear() {
- activeSegments.asMap().clear();
- staleWorkloads.clear();
- }
-
- private class WorkloadCollectorCacheLoader extends CacheLoader<TopicPartition, WorkloadCollector> {
-
- @Override
- public WorkloadCollector load(TopicPartition topicPartition) throws Exception {
- log.debug("Started recording workload for topic partition {}", topicPartition);
- return new WorkloadCollector(maxSegmentAge, time);
- }
- }
-
- private class ExpiredCollectorListener implements RemovalListener<TopicPartition, WorkloadCollector> {
- @Override
- public void onRemoval(RemovalNotification<TopicPartition, WorkloadCollector> removal) {
- if (!RemovalCause.EXPIRED.equals(removal.getCause())) {
- return;
- }
-
- // We just simulate recording a throughput of 0 bytes/0 messages; this forces a check
- // to see if the segment is complete, but doesn't have any effect on the throughputs
- // for the topic partition
- Optional<Workload> finalSegmentWorkload = removal.getValue().record(0, 0);
- finalSegmentWorkload.ifPresent(workload -> {
- synchronized (staleWorkloads) {
- staleWorkloads.put(removal.getKey(), workload);
- }
- });
- }
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadTrackingProcessor.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadTrackingProcessor.java
deleted file mode 100644
index 83c01cf06a2..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/tracking/WorkloadTrackingProcessor.java
+++ /dev/null
@@ -1,276 +0,0 @@
-package com.uber.data.kafka.connect.workload.tracking;
-
-import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.connect.connector.ConnectRecord;
-import org.apache.kafka.connect.header.Header;
-
-import com.google.common.annotations.VisibleForTesting;
-import com.uber.data.kafka.connect.utils.ExecutorUtils;
-import com.uber.data.kafka.connect.utils.MetricsUtils;
-import com.uber.data.kafka.connect.workload.common.Workload;
-import com.uber.data.kafka.connect.workload.transformation.WorkloadTrackingTransformation;
-import com.uber.m3.tally.Scope;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.time.Duration;
-import java.util.Collection;
-import java.util.HashSet;
-import java.util.Map;
-import java.util.Objects;
-import java.util.Set;
-import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.TimeUnit;
-import java.util.function.Supplier;
-
-public class WorkloadTrackingProcessor implements AutoCloseable {
-
- @VisibleForTesting
- static final String WORKLOAD_BYTE_RATE_METRIC = "worker.workloads.byte.rate";
- @VisibleForTesting
- static final String WORKLOAD_MESSAGE_RATE_METRIC = "worker.workloads.message.rate";
- @VisibleForTesting
- static final String WORKLOAD_EMPTY_PARTITIONS_METRIC = "worker.workloads.empty.partitions";
- @VisibleForTesting
- static final String WORKLOAD_STALE_FLUSH_FAILURE_METRIC = "worker.workloads.stale.flush.failures";
- @VisibleForTesting
- static final String WORKLOAD_EMPTY_PUBLISH_FAILURE_METRIC = "worker.workloads.empty.publish.failures";
-
- private static final Logger log = LoggerFactory.getLogger(WorkloadTrackingTransformation.class);
-
- private final WorkloadStore workloadStore;
- private final String srcKafka;
- private final String dstKafka;
- private final Supplier<? extends Collection<TopicPartition>> assignedPartitions;
- private final Set<TopicPartition> emptyPartitions;
- private final WorkloadTracker topicWorkloadTracker;
- private final ScheduledExecutorService executor;
- private final Scope scope;
- private final Duration maxSegmentAge;
-
- public WorkloadTrackingProcessor(
- String srcKafka,
- String dstKafka,
- WorkloadStore workloadStore,
- Supplier<? extends Collection<TopicPartition>> assignedPartitions,
- Duration maxSegmentAge
- ) {
- this(
- srcKafka,
- dstKafka,
- assignedPartitions,
- MetricsUtils.INSTANCE.routeScope(srcKafka, dstKafka),
- workloadStore,
- new WorkloadTracker(maxSegmentAge),
- ExecutorUtils.singleThreadedScheduled("workload-tracking-message-processor-flush"),
- maxSegmentAge
- );
- }
-
- @VisibleForTesting
- WorkloadTrackingProcessor(
- String srcKafka,
- String dstKafka,
- Supplier<? extends Collection<TopicPartition>> assignedPartitions,
- Scope scope,
- WorkloadStore workloadStore,
- WorkloadTracker topicWorkloadTracker,
- ScheduledExecutorService executor,
- Duration maxSegmentAge
- ) {
- this.srcKafka = Objects.requireNonNull(
- srcKafka,
- "source Kafka cluster name may not be null"
- );
- this.dstKafka = Objects.requireNonNull(
- dstKafka,
- "target Kafka cluster name may not be null"
- );
- this.assignedPartitions = Objects.requireNonNull(
- assignedPartitions,
- "assigned topics supplier may not be null"
- );
- this.scope = Objects.requireNonNull(
- scope,
- "metrics scope may not be null"
- );
- this.workloadStore = Objects.requireNonNull(
- workloadStore,
- "workload store may not be nul"
- );
- this.topicWorkloadTracker = Objects.requireNonNull(
- topicWorkloadTracker,
- "topic workload tracker may not be null"
- );
- this.executor = Objects.requireNonNull(
- executor,
- "flush executor may not be null"
- );
- this.maxSegmentAge = Objects.requireNonNull(
- maxSegmentAge,
- "max segment age may not be null"
- );
-
- this.emptyPartitions = new HashSet<>();
- }
-
- public void start() {
- executor.scheduleAtFixedRate(
- this::handleStaleAndEmptyWorkloads,
- maxSegmentAge.toMillis(),
- maxSegmentAge.toMillis(),
- TimeUnit.MILLISECONDS
- );
- }
-
- public void track(ConnectRecord<?> record) {
- track(record.topic(), record.kafkaPartition(), numRecordBytes(record), 1);
- }
-
- @VisibleForTesting
- void track(String topic, int partition, long bytes, long messages) {
- TopicPartition topicPartition = new TopicPartition(topic, partition);
- synchronized (emptyPartitions) {
- emptyPartitions.remove(topicPartition);
- }
- topicWorkloadTracker.record(topicPartition, bytes, messages)
- .ifPresent(workload -> publishWorkload(topicPartition, workload));
- }
-
- @Override
- public void close() {
- ExecutorUtils.shutdown(
- executor,
- "executor for workload tracker",
- 5,
- TimeUnit.SECONDS
- );
-
- // Final attempt to flush complete workloads
- flushStaleWorkloads();
- }
-
- @VisibleForTesting
- void handleStaleAndEmptyWorkloads() {
- try {
- flushStaleWorkloads();
- } catch (Throwable t) {
- log.error("Failed to flush stale workloads", t);
- scope.counter(WORKLOAD_STALE_FLUSH_FAILURE_METRIC).inc(1);
- }
-
- try {
- publishEmptyWorkloads();
- } catch (Throwable t) {
- log.error("Failed to publish empty workloads", t);
- scope.counter(WORKLOAD_EMPTY_PUBLISH_FAILURE_METRIC).inc(1);
- }
- }
-
- /**
- * Check for stale segments (i.e., ones that are old enough to complete, but for which no
- * throughput has been recorded after they became old enough), complete them, and flush them.
- */
- @VisibleForTesting
- void flushStaleWorkloads() {
- Map<TopicPartition, Workload> staleWorkloads = topicWorkloadTracker.collectStaleWorkloads();
- if (staleWorkloads.isEmpty()) {
- return;
- }
-
- synchronized (emptyPartitions) {
- emptyPartitions.removeAll(staleWorkloads.keySet());
- }
-
- log.debug(
- "Flushing {} stale workloads for topic partitions {}",
- staleWorkloads.size(),
- staleWorkloads.keySet()
- );
- staleWorkloads.forEach(this::publishWorkload);
- }
-
- /**
- * Check for empty topic partitions (i.e., ones that haven't received any traffic since the
- * last time we checked), and publish a zero-throughput workload for them.
- * <p>
- * This explicit step is necessary because all other workloads are only published after some
- * traffic is received (via one or more calls to {@link #track(String, int, long, long)});
- * for empty partitions, we never see traffic, and instead have to explicitly check for them
- * and publish workloads.
- */
- @VisibleForTesting
- void publishEmptyWorkloads() {
- Set<TopicPartition> partitions = new HashSet<>(assignedPartitions.get());
-
- // Take a snapshot of the current set of empty partitions so that we don't perform writes
- // to ZooKeeper while holding the lock
- Set<TopicPartition> emptyPartitionsSnapshot;
- synchronized (this.emptyPartitions) {
- emptyPartitionsSnapshot = new HashSet<>(this.emptyPartitions);
-
- // Assume every partition is empty until we see traffic for it
- this.emptyPartitions.clear();
- this.emptyPartitions.addAll(partitions);
- }
-
- // Don't publish empty workloads for partitions that are no longer assigned to us
- emptyPartitionsSnapshot.retainAll(partitions);
-
- scope.gauge(WORKLOAD_EMPTY_PARTITIONS_METRIC)
- .update(emptyPartitionsSnapshot.size());
-
- emptyPartitionsSnapshot.forEach(tp -> publishWorkload(tp, Workload.EMPTY));
- }
-
- private void publishWorkload(TopicPartition topicPartition, Workload workload) {
- log.trace(
- "Publishing workload of {} for topic partition {}",
- workload,
- topicPartition
- );
-
- workloadStore.store(srcKafka, dstKafka, topicPartition, workload);
-
- scope.tagged(MetricsUtils.INSTANCE.tags(topicPartition))
- .gauge(WORKLOAD_BYTE_RATE_METRIC)
- .update(workload.getBytesPerSecond());
- scope.tagged(MetricsUtils.INSTANCE.tags(topicPartition))
- .gauge(WORKLOAD_MESSAGE_RATE_METRIC)
- .update(workload.getMessagesPerSecond());
- }
-
- @VisibleForTesting
- static int numRecordBytes(ConnectRecord<?> record) {
- String topic = record.topic();
- int recordBytes = 0;
-
- if (record.key() instanceof byte[] keyBytes) {
- recordBytes += keyBytes.length;
- } else {
- log.warn("Key for record in topic {} is not byte array; ignoring", topic);
- }
-
- if (record.value() instanceof byte[] valueBytes) {
- recordBytes += valueBytes.length;
- } else {
- log.warn("Value for record in topic {} is not byte array; ignoring", topic);
- }
-
- for (Header header : record.headers()) {
- recordBytes += header.key().length();
-
- if (header.value() instanceof String stringHeader) {
- recordBytes += stringHeader.length();
- } else if (header.value() instanceof byte[] bytesHeader) {
- recordBytes += bytesHeader.length;
- } else {
- log.debug("Value for header {} on record in topic {} is not byte array; ignoring", header.key(), topic);
- }
- }
-
- return recordBytes;
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/transformation/WorkloadTrackingTransformation.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/transformation/WorkloadTrackingTransformation.java
deleted file mode 100644
index 1cb2bc69b0d..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/transformation/WorkloadTrackingTransformation.java
+++ /dev/null
@@ -1,62 +0,0 @@
-package com.uber.data.kafka.connect.workload.transformation;
-
-import org.apache.kafka.common.config.ConfigDef;
-import org.apache.kafka.common.utils.Utils;
-import org.apache.kafka.connect.connector.ConnectRecord;
-import org.apache.kafka.connect.transforms.Transformation;
-
-import com.uber.data.kafka.connect.workload.tracking.WorkloadStore;
-import com.uber.data.kafka.connect.workload.tracking.WorkloadTrackingProcessor;
-
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.Map;
-import java.util.Set;
-
-/**
- * {@link org.apache.kafka.connect.transforms.Transformation} implementation that tracks workloads on a per-topic-partition basis as
- * records are consumed from Kafka, and reports those workloads periodically to a pluggable {@link WorkloadStore}
- */
-public abstract class WorkloadTrackingTransformation<R extends ConnectRecord<R>> implements Transformation<R> {
-
- private static final Logger log = LoggerFactory.getLogger(WorkloadTrackingTransformation.class);
-
- private WorkloadTrackingProcessor processor;
-
- protected abstract WorkloadStore createStore();
-
- @Override
- public ConfigDef config() {
- return WorkloadTrackingTransformationConfig.config();
- }
-
- @Override
- public void configure(Map<String, ?> props) {
- WorkloadTrackingTransformationConfig config = new WorkloadTrackingTransformationConfig(props);
-
- this.processor = new WorkloadTrackingProcessor(
- config.srcKafkaName(),
- config.dstKafkaName(),
- createStore(),
- Set::of, // TODO: Is there any non-insane way to get the set of topic partitions assigned to the task here?
- config.reportInterval()
- );
-
- this.processor.start();
-
- log.info("Topic workload tracking is now enabled");
- }
-
- @Override
- public R apply(R record) {
- processor.track(record);
- return record;
- }
-
- @Override
- public void close() {
- Utils.closeQuietly(processor, "workload processor");
- }
-
-}
diff --git a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/transformation/WorkloadTrackingTransformationConfig.java b/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/transformation/WorkloadTrackingTransformationConfig.java
deleted file mode 100644
index 4c4e83ea745..00000000000
--- a/connect/uber/src/main/java/com/uber/data/kafka/connect/workload/transformation/WorkloadTrackingTransformationConfig.java
+++ /dev/null
@@ -1,62 +0,0 @@
-package com.uber.data.kafka.connect.workload.transformation;
-
-import org.apache.kafka.common.config.AbstractConfig;
-import org.apache.kafka.common.config.ConfigDef;
-import org.apache.kafka.common.config.ConfigDef.Importance;
-import org.apache.kafka.common.config.ConfigDef.Type;
-
-import java.time.Duration;
-import java.util.Map;
-
-import static org.apache.kafka.common.config.ConfigDef.Range.atLeast;
-
-public class WorkloadTrackingTransformationConfig extends AbstractConfig {
-
- public static final String SRC_CLUSTER_NAME_CONFIG = "source.cluster.name";
- private static final String SRC_CLUSTER_NAME_DOC = "Name of the source Kafka cluster";
-
- public static final String DST_CLUSTER_NAME_CONFIG = "target.cluster.name";
- private static final String DST_CLUSTER_NAME_DOC = "Name of the target Kafka cluster";
-
- public static final String REPORT_INTERVAL_MS_CONFIG = "report.interval.ms";
- private static final String REPORT_INTERVAL_MS_DOC = "How long to record traffic for in between reporting recorded values";
-
- public WorkloadTrackingTransformationConfig(Map<String, ?> props) {
- super(config(), props);
- }
-
- static ConfigDef config() {
- return new ConfigDef()
- .define(
- SRC_CLUSTER_NAME_CONFIG,
- Type.STRING,
- Importance.HIGH,
- SRC_CLUSTER_NAME_DOC
- ).define(
- DST_CLUSTER_NAME_CONFIG,
- Type.STRING,
- Importance.HIGH,
- DST_CLUSTER_NAME_DOC
- ).define(
- REPORT_INTERVAL_MS_CONFIG,
- Type.LONG,
- Duration.ofMinutes(10).toMillis(),
- atLeast(Duration.ofSeconds(1).toMillis()),
- Importance.LOW,
- REPORT_INTERVAL_MS_DOC
- );
- }
-
- public String srcKafkaName() {
- return getString(SRC_CLUSTER_NAME_CONFIG);
- }
-
- public String dstKafkaName() {
- return getString(DST_CLUSTER_NAME_CONFIG);
- }
-
- public Duration reportInterval() {
- return Duration.ofMillis(getLong(REPORT_INTERVAL_MS_CONFIG));
- }
-
-}
diff --git a/connect/uber/src/main/resources/META-INF/services/com.uber.data.kafka.connect.distributed.ClusterAssignor b/connect/uber/src/main/resources/META-INF/services/com.uber.data.kafka.connect.distributed.ClusterAssignor
deleted file mode 100644
index 3d72df7440a..00000000000
--- a/connect/uber/src/main/resources/META-INF/services/com.uber.data.kafka.connect.distributed.ClusterAssignor
+++ /dev/null
@@ -1 +0,0 @@
- com.uber.data.kafka.connect.distributed.assignment.UberClusterAssignor
\ No newline at end of file
diff --git a/connect/uber/src/main/resources/META-INF/services/org.apache.kafka.connect.rest.ConnectRestExtension b/connect/uber/src/main/resources/META-INF/services/org.apache.kafka.connect.rest.ConnectRestExtension
deleted file mode 100644
index b12936f3e82..00000000000
--- a/connect/uber/src/main/resources/META-INF/services/org.apache.kafka.connect.rest.ConnectRestExtension
+++ /dev/null
@@ -1 +0,0 @@
-com.uber.data.kafka.connect.rest.UberConnectRestExtension
\ No newline at end of file
diff --git a/connect/uber/src/main/resources/META-INF/services/org.apache.kafka.connect.source.SourceConnector b/connect/uber/src/main/resources/META-INF/services/org.apache.kafka.connect.source.SourceConnector
deleted file mode 100644
index 5822942e2b8..00000000000
--- a/connect/uber/src/main/resources/META-INF/services/org.apache.kafka.connect.source.SourceConnector
+++ /dev/null
@@ -1 +0,0 @@
-com.uber.data.kafka.connect.ureplicator3.UReplicator3Connector
\ No newline at end of file
diff --git a/connect/uber/src/test/java/com/uber/data/kafka/connect/integration/ReplicationIntegrationTest.java b/connect/uber/src/test/java/com/uber/data/kafka/connect/integration/ReplicationIntegrationTest.java
deleted file mode 100644
index 790f082010b..00000000000
--- a/connect/uber/src/test/java/com/uber/data/kafka/connect/integration/ReplicationIntegrationTest.java
+++ /dev/null
@@ -1,1243 +0,0 @@
-package com.uber.data.kafka.connect.integration;
-
-import org.apache.kafka.clients.CommonClientConfigs;
-import org.apache.kafka.clients.admin.Admin;
-import org.apache.kafka.clients.admin.ListTopicsOptions;
-import org.apache.kafka.clients.admin.NewPartitions;
-import org.apache.kafka.clients.consumer.ConsumerRecord;
-import org.apache.kafka.clients.consumer.ConsumerRecords;
-import org.apache.kafka.clients.producer.Producer;
-import org.apache.kafka.clients.producer.ProducerConfig;
-import org.apache.kafka.clients.producer.ProducerRecord;
-import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.common.utils.Utils;
-import org.apache.kafka.connect.converters.ByteArrayConverter;
-import org.apache.kafka.connect.runtime.isolation.PluginDiscoveryMode;
-import org.apache.kafka.connect.runtime.rest.entities.ConfigInfo;
-import org.apache.kafka.connect.runtime.rest.entities.ConfigInfos;
-import org.apache.kafka.connect.runtime.rest.entities.ConfigValueInfo;
-import org.apache.kafka.connect.runtime.rest.entities.ConnectorOffset;
-import org.apache.kafka.connect.runtime.rest.entities.ConnectorOffsets;
-import org.apache.kafka.connect.runtime.rest.entities.CreateConnectorRequest;
-import org.apache.kafka.connect.runtime.rest.errors.ConnectRestException;
-import org.apache.kafka.connect.util.clusters.EmbeddedConnectCluster;
-import org.apache.kafka.connect.util.clusters.EmbeddedKafkaCluster;
-import org.apache.kafka.connect.util.clusters.WorkerHandle;
-
-import com.fasterxml.jackson.annotation.JsonAutoDetect.Visibility;
-import com.fasterxml.jackson.annotation.PropertyAccessor;
-import com.fasterxml.jackson.core.type.TypeReference;
-import com.fasterxml.jackson.databind.ObjectMapper;
-import com.uber.data.kafka.connect.common.ReplicatedTopicPartition;
-import com.uber.data.kafka.connect.distributed.assignment.UberClusterAssignor;
-import com.uber.data.kafka.connect.rest.UberConnectRestExtension;
-import com.uber.data.kafka.connect.rest.entities.AssignedPartitionEntity;
-import com.uber.data.kafka.connect.rest.entities.PartitionWorkloadEntity;
-import com.uber.data.kafka.connect.rest.entities.TaskAssignmentEntity;
-import com.uber.data.kafka.connect.rest.entities.WorkerWorkloadEntity;
-import com.uber.data.kafka.connect.rest.entities.WorkloadEntity;
-import com.uber.data.kafka.connect.ureplicator3.UReplicator3Connector;
-import com.uber.data.kafka.connect.ureplicator3.UReplicator3Task;
-import com.uber.data.kafka.connect.ureplicator3.coordination.AssignedPartition;
-import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStoreConfig;
-import com.uber.data.kafka.connect.ureplicator3.coordination.AssignmentStores;
-import com.uber.data.kafka.connect.workload.common.Workload;
-import com.uber.data.kafka.connect.workload.kafka.KafkaBasedWorkloadProvider;
-
-import org.junit.jupiter.api.AfterAll;
-import org.junit.jupiter.api.AfterEach;
-import org.junit.jupiter.api.BeforeAll;
-import org.junit.jupiter.api.Tag;
-import org.junit.jupiter.api.Test;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.io.IOException;
-import java.net.URI;
-import java.net.http.HttpClient;
-import java.net.http.HttpRequest;
-import java.net.http.HttpRequest.BodyPublishers;
-import java.net.http.HttpResponse;
-import java.net.http.HttpResponse.BodyHandlers;
-import java.nio.charset.StandardCharsets;
-import java.util.AbstractMap;
-import java.util.Arrays;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Map;
-import java.util.Properties;
-import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicReference;
-import java.util.function.Function;
-import java.util.stream.Collectors;
-import java.util.stream.IntStream;
-import java.util.stream.Stream;
-
-import jakarta.ws.rs.core.Response;
-import jakarta.ws.rs.core.Response.Status;
-
-import static com.uber.data.kafka.connect.distributed.assignment.UberClusterAssignorConfig.UBER_WORKLOAD_PROVIDER_CONFIG;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Config.TARGET_CLIENTS_PREFIX;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Config.TARGET_PRODUCER_PREFIX;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.REFRESH_INTERVAL_MS;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.SOURCE_ADMIN_PREFIX;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.SOURCE_CLUSTER_CONFIG;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.SOURCE_CONSUMER_PREFIX;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.TARGET_ADMIN_PREFIX;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.TARGET_CLUSTER_CONFIG;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.TOPICS_CONFIG;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3ConnectorConfig.TOPIC_RENAME_PREFIX;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3Task.SOURCE_OFFSET_KAFKA_OFFSET;
-import static com.uber.data.kafka.connect.ureplicator3.UReplicator3TaskConfig.WORKLOAD_REPORT_INTERVAL_MS_CONFIG;
-import static com.uber.data.kafka.connect.workload.kafka.KafkaBasedWorkloadProviderConfig.STORAGE_TOPIC_CONFIG;
-import static com.uber.data.kafka.connect.workload.kafka.KafkaBasedWorkloadProviderConfig.STORAGE_TOPIC_REPLICATION_FACTOR_CONFIG;
-import static jakarta.ws.rs.core.HttpHeaders.CONTENT_TYPE;
-import static jakarta.ws.rs.core.MediaType.APPLICATION_FORM_URLENCODED;
-import static org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLASS_CONFIG;
-import static org.apache.kafka.connect.runtime.ConnectorConfig.CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX;
-import static org.apache.kafka.connect.runtime.ConnectorConfig.NAME_CONFIG;
-import static org.apache.kafka.connect.runtime.ConnectorConfig.TASKS_MAX_CONFIG;
-import static org.apache.kafka.connect.runtime.WorkerConfig.CONNECTOR_CLIENT_POLICY_CLASS_CONFIG;
-import static org.apache.kafka.connect.runtime.WorkerConfig.KEY_CONVERTER_CLASS_CONFIG;
-import static org.apache.kafka.connect.runtime.WorkerConfig.PLUGIN_DISCOVERY_CONFIG;
-import static org.apache.kafka.connect.runtime.WorkerConfig.VALUE_CONVERTER_CLASS_CONFIG;
-import static org.apache.kafka.connect.runtime.distributed.DistributedConfig.UBER_CLUSTER_ASSIGNOR_CONFIG;
-import static org.apache.kafka.connect.runtime.rest.RestServerConfig.REST_EXTENSION_CLASSES_CONFIG;
-import static org.apache.kafka.test.TestUtils.waitForCondition;
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertNotEquals;
-import static org.junit.jupiter.api.Assertions.assertNotNull;
-import static org.junit.jupiter.api.Assertions.assertThrows;
-import static org.junit.jupiter.api.Assertions.assertTrue;
-
-@Tag("integration")
-public class ReplicationIntegrationTest {
-
- // TODO: These test cases are taking pretty long, even though we share embedded Kafka clusters across cases
- // If we have time we should figure out what's taking so long and try to speed things up
-
- private static final Logger log = LoggerFactory.getLogger(ReplicationIntegrationTest.class);
- private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper()
- .setVisibility(PropertyAccessor.FIELD, Visibility.ANY);
-
- private static final int NUM_WORKERS = 3;
- private static final String WORKLOAD_TOPIC = "uber-workloads";
- private static final String ASSIGNMENTS_TOPIC = "ureplicator3-assignments";
- private static final Map<String, EmbeddedKafkaCluster> EXTRA_KAFKA_CLUSTERS = new ConcurrentHashMap<>();
-
- private static EmbeddedConnectCluster connect;
-
- @BeforeAll
- public static void setup() throws Exception {
- Map<String, String> workerProps = new HashMap<>(); // Has to be mutable to work with the embedded Connect cluster API
- workerProps.put(UBER_CLUSTER_ASSIGNOR_CONFIG, UberClusterAssignor.class.getName());
- 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(AssignmentStoreConfig.STORAGE_TOPIC_CONFIG, ASSIGNMENTS_TOPIC);
- workerProps.put(AssignmentStoreConfig.STORAGE_TOPIC_REPLICATION_FACTOR_CONFIG, "1");
- workerProps.put(PLUGIN_DISCOVERY_CONFIG, PluginDiscoveryMode.SERVICE_LOAD.name()); // Fast startup time ⚡
- workerProps.put(CONNECTOR_CLIENT_POLICY_CLASS_CONFIG, "All");
- workerProps.put(REST_EXTENSION_CLASSES_CONFIG, UberConnectRestExtension.class.getName());
- workerProps.put(KEY_CONVERTER_CLASS_CONFIG, ByteArrayConverter.class.getName());
- workerProps.put(VALUE_CONVERTER_CLASS_CONFIG, ByteArrayConverter.class.getName());
-
- int numBrokers = 1;
- Properties brokerProps = new Properties();
-
- // build a Connect cluster backed by Kafka
- connect = new EmbeddedConnectCluster.Builder()
- .numWorkers(NUM_WORKERS)
- .workerProps(workerProps)
- .numBrokers(numBrokers)
- .brokerProps(brokerProps)
- .build();
-
- connect.start();
- connect.assertions().assertExactlyNumWorkersAreUp(NUM_WORKERS, "Cluster did not start in time");
- }
-
- @AfterEach
- public void close() throws Exception {
- log.info("Cleaning up test case resources");
-
- // Wipe every connector's offsets and then delete it
- for (String connector : connect.connectors()) {
- connect.stopConnector(connector);
- connect.assertions().assertConnectorIsStopped(connector, "connector did not stop in time");
-
- connect.resetConnectorOffsets(connector);
-
- connect.deleteConnector(connector);
- connect.assertions().assertConnectorDoesNotExist(connector, "connector was not deleted in time");
- }
-
- // Delete testing topics from our extra Kafka clusters
- for (EmbeddedKafkaCluster kafkaCluster : EXTRA_KAFKA_CLUSTERS.values()) {
- try (Admin admin = kafkaCluster.createAdminClient()) {
- ListTopicsOptions listOptions = new ListTopicsOptions().listInternal(false);
- Set<String> topics = admin.listTopics(listOptions).names().get();
-
- admin.deleteTopics(topics);
- }
- }
-
- // Wipe out any workloads and assignments generated during this test case so they don't affect other cases
- KafkaBasedWorkloadProvider.clearCache();
- AssignmentStores.get().clear();
-
- log.info("Finished cleaning up test case resources");
- }
-
- @AfterAll
- public static void teardown() {
- 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
- EXTRA_KAFKA_CLUSTERS.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 testContentTypeHeaderFilter() throws Exception {
- String srcKafkaName = "src-kafka";
- String dstKafkaName = "dst-kafka";
-
- EmbeddedKafkaCluster srcCluster = extraKafkaCluster(srcKafkaName, 1, new Properties());
- EmbeddedKafkaCluster dstCluster = extraKafkaCluster(dstKafkaName, 1, new Properties());
-
- String topic = "t1";
- srcCluster.createTopic(topic);
- dstCluster.createTopic(topic);
-
- // If your IDE tells you this should be closed, it's using a more up-to-date JDK version than the project builds with
- HttpClient httpClient = HttpClient.newHttpClient();
- String resource = "connector-plugins/" + UReplicator3Connector.class.getSimpleName() + "/config/validate";
- Map<String, String> connectorConfig = baseURep3Config(srcKafkaName, dstKafkaName, 3, Set.of(topic));
- byte[] body = OBJECT_MAPPER.writeValueAsBytes(connectorConfig);
-
- HttpRequest request = HttpRequest.newBuilder(new URI(connect.endpointForResource(resource)))
- .PUT(BodyPublishers.ofByteArray(body))
- .build();
- HttpResponse<String> response = httpClient.send(request, BodyHandlers.ofString());
- // 415 (unsupported media type) is what we'd see with the header filter disabled,
- // since we didn't specify a content-type header in our request
- assertEquals(200, response.statusCode());
-
- request = HttpRequest.newBuilder(new URI(connect.endpointForResource(resource)))
- .PUT(BodyPublishers.ofByteArray(body))
- .header(CONTENT_TYPE, APPLICATION_FORM_URLENCODED) // Make sure this stupid shit gets overridden too (curl, I'm not mad at you, I'm just disappointed)
- .build();
- response = httpClient.send(request, BodyHandlers.ofString());
- assertEquals(200, response.statusCode());
- }
-
- @Test
- public void testPublishWorkloadsManually() throws Exception {
- assertEquals(
- Map.of(),
- getAllWorkloads()
- );
-
- ReplicatedTopicPartition tp1 = new ReplicatedTopicPartition("src1", "dst1", "t", 0);
- Workload w1 = new Workload(3.0, 4.0);
- manuallyPublishWorkload(tp1, w1);
-
- assertEquals(
- Map.of(tp1, w1),
- getAllWorkloads()
- );
-
- ReplicatedTopicPartition tp2 = new ReplicatedTopicPartition("src2", "dst2", "t", 0);
- Workload w2 = new Workload(35.0, 14.0);
- manuallyPublishWorkload(tp2, w2);
-
- assertEquals(
- Map.of(tp1, w1, tp2, w2),
- getAllWorkloads()
- );
-
- manuallyDeleteWorkload(tp1);
-
- assertEquals(
- Map.of(tp2, w2),
- getAllWorkloads()
- );
- }
-
- @Test
- public void testNormalConnectors() throws Exception {
- String srcKafkaName = "src-kafka";
- String dstKafkaName = "dst-kafka";
-
- EmbeddedKafkaCluster srcCluster = extraKafkaCluster(srcKafkaName, 1, new Properties());
- EmbeddedKafkaCluster dstCluster = extraKafkaCluster(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(conn1, srcKafkaName, dstKafkaName, 3, Set.of("t1", "t2", "t3"), null, false);
- startConnector(conn2, srcKafkaName, dstKafkaName, 1, Set.of("t4", "t5", "t6"), null, false);
-
- dstCluster.consume(numProduced, TimeUnit.MINUTES.toMillis(1), "t1", "t2", "t3", "t4", "t5", "t6");
-
- // Stop the connectors to force offset commits
- connect.stopConnector(conn1);
- connect.stopConnector(conn2);
-
- connect.assertions().assertConnectorIsStopped(conn1, "Connector and/or tasks did not stop in time");
- connect.assertions().assertConnectorIsStopped(conn2, "Connector and/or tasks did not stop in time");
-
- waitForCondition(() -> {
- Map<Map<String, ?>, Map<String, ?>> actualOffsets = connect.connectorOffsets(conn1).toMap();
- Map<Map<String, ?>, Map<String, ?>> expectedOffsets = IntStream.rangeClosed(1, 3)
- .boxed()
- .flatMap(t -> IntStream.range(0, t)
- .mapToObj(p -> {
- Map<String, Object> sourcePartition = UReplicator3Task.encodeSourcePartition(srcKafkaName, dstKafkaName, "t" + t, p);
- // Can't use UReplicator3Task::encodeSourceOffset because it encodes as a long instead of an integer,
- // and the response from the REST API above gets deserialized as an integer instead of a long
- // I lost several hours debugging this fucking bullshit if you want to do better have at it
- Map<String, Object> sourceOffset = Map.of(SOURCE_OFFSET_KAFKA_OFFSET, 0);
- return new AbstractMap.SimpleEntry<>(sourcePartition, sourceOffset);
- })
- ).collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
- assertEquals(actualOffsets, expectedOffsets);
-
- actualOffsets = connect.connectorOffsets(conn2).toMap();
- expectedOffsets = IntStream.rangeClosed(4, 6)
- .boxed()
- .flatMap(t -> IntStream.range(0, t)
- .mapToObj(p -> {
- Map<String, Object> sourcePartition = UReplicator3Task.encodeSourcePartition(srcKafkaName, dstKafkaName, "t" + t, p);
- Map<String, Object> sourceOffset = Map.of(SOURCE_OFFSET_KAFKA_OFFSET, 0);
- return new AbstractMap.SimpleEntry<>(sourcePartition, sourceOffset);
- })
- ).collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
- assertEquals(actualOffsets, expectedOffsets);
-
- return true;
- }, 30_000, "Connector offsets were not updated in time");
-
- // Produce some more 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 < i; partition++) {
- produceTestRecord(srcProducer, topic, partition, 1);
- numProduced++;
- }
- }
-
- srcProducer.flush();
- }
-
- connect.resumeConnector(conn1);
- connect.resumeConnector(conn2);
-
- dstCluster.consume(numProduced, TimeUnit.MINUTES.toMillis(1), "t1", "t2", "t3", "t4", "t5", "t6");
- int numConsumed = dstCluster.consumeAll(60_000, "t1", "t2", "t3", "t4", "t5", "t6").count();
- assertEquals(numProduced, numConsumed);
-
- // Reconfigure to stop replicating from t3
- startConnector(conn1, srcKafkaName, dstKafkaName, 2, Set.of("t1", "t2"), null, false);
-
- ConnectorOffset t1p0Offset = new ConnectorOffset(
- UReplicator3Task.encodeSourcePartition(srcKafkaName, dstKafkaName, "t1", 0),
- UReplicator3Task.encodeSourceOffset(-1)
- );
- ConnectorOffset t3p0Offset = new ConnectorOffset(
- UReplicator3Task.encodeSourcePartition(srcKafkaName, dstKafkaName, "t3", 0),
- UReplicator3Task.encodeSourceOffset(-1)
- );
-
- final ConnectorOffsets invalidOffsets = new ConnectorOffsets(List.of(t1p0Offset, t3p0Offset));
- // Should fail because the connector is still running
- ConnectRestException failure = assertThrows(ConnectRestException.class, () -> connect.alterConnectorOffsets(conn1, invalidOffsets));
- assertEquals(Status.BAD_REQUEST.getStatusCode(), failure.statusCode());
-
- // Should fail because the connector is still assigned one of the altered topics
- failure = assertThrows(ConnectRestException.class, () -> connect.alterConnectorOffsets(conn1, invalidOffsets, true));
- assertEquals(Status.BAD_REQUEST.getStatusCode(), failure.statusCode());
-
- ConnectorOffsets validOffsets = new ConnectorOffsets(List.of(t3p0Offset));
- // Should succeed
- connect.alterConnectorOffsets(conn1, validOffsets, true);
-
- // Reconfigure to resume replicating from t3
- startConnector(conn1, srcKafkaName, dstKafkaName, 1, Set.of("t1", "t2", "t3"), null, false);
-
- dstCluster.consume(4, TimeUnit.MINUTES.toMillis(1), "t3");
- List<ConsumerRecord<byte[], byte[]>> consumed = dstCluster.consumeAll(60_000, "t3")
- .records(new TopicPartition("t3", 0));
- // UReplicator3 currently only provides at-least-once semantics; it's possible there could be some
- // duplication for our reset partition (or any other replicated partition)
- assertTrue(consumed.size() >= 4);
- long replicatedFirstRecords = consumed.stream()
- .filter(record -> Arrays.equals(record.value(), "value 0".getBytes(StandardCharsets.UTF_8)))
- .count();
- assertTrue(replicatedFirstRecords >= 2);
- long replicatedSecondRecords = consumed.stream()
- .filter(record -> Arrays.equals(record.value(), "value 1".getBytes(StandardCharsets.UTF_8)))
- .count();
- assertTrue(replicatedSecondRecords >= 2);
- }
-
- @Test
- public void testConfigTranslation() throws Exception {
- String srcKafkaName = "src-kafka";
- String dstKafkaName = "dst-kafka";
-
- EmbeddedKafkaCluster srcCluster = extraKafkaCluster(srcKafkaName, 1, new Properties());
- EmbeddedKafkaCluster dstCluster = extraKafkaCluster(dstKafkaName, 1, new Properties());
-
- String connector = "connector";
-
- String topic = "t1";
- srcCluster.createTopic(topic);
- dstCluster.createTopic(topic);
-
- String targetClientsReceiveBufferBytes = TARGET_CLIENTS_PREFIX + CommonClientConfigs.RECEIVE_BUFFER_CONFIG;
- String targetClientsSendBufferBytes = TARGET_CLIENTS_PREFIX + CommonClientConfigs.SEND_BUFFER_CONFIG;
- String targetProducerLingerMs = TARGET_PRODUCER_PREFIX + ProducerConfig.LINGER_MS_CONFIG;
- String targetProducerSendBufferBytes = TARGET_PRODUCER_PREFIX + CommonClientConfigs.SEND_BUFFER_CONFIG;
-
- Map<String, String> connectorConfig = baseURep3Config(srcKafkaName, dstKafkaName, 1, Set.of(topic));
- connectorConfig.put(SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG, srcCluster.bootstrapServers());
- connectorConfig.put(TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG, dstCluster.bootstrapServers());
- connectorConfig.put(targetClientsReceiveBufferBytes, "33333");
- connectorConfig.put(targetClientsSendBufferBytes, "11111");
- connectorConfig.put(targetProducerLingerMs, "1000");
- connectorConfig.put(targetProducerSendBufferBytes, "22222");
-
- // Uses the POST endpoint
- // TODO: Flaky (sometimes metadata for the newly-created topic hasn't propagated to every broker)
- // Unclear if we should handle this by adding retry logic to the test, or if we should add retry logic to the connector itself
- connect.configureConnector(new CreateConnectorRequest(connector, connectorConfig, null));
-
- connect.assertions().assertConnectorAndExactlyNumTasksAreRunning(connector, 1, "Connector and/or task did not start in time");
-
- Map<String, String> expectedTranslatedConfig = new HashMap<>(connectorConfig);
- expectedTranslatedConfig.put(
- CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG,
- dstCluster.bootstrapServers()
- );
- expectedTranslatedConfig.put(
- CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + CommonClientConfigs.RECEIVE_BUFFER_CONFIG,
- "33333"
- );
- expectedTranslatedConfig.put(
- CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + ProducerConfig.LINGER_MS_CONFIG,
- "1000"
- );
- expectedTranslatedConfig.put(
- CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + CommonClientConfigs.SEND_BUFFER_CONFIG,
- "22222" // Precedence is given to the target.producer properties over the target.kafka.clients properties
- );
- // Automatically injected by the Connect runtime
- expectedTranslatedConfig.put(NAME_CONFIG, connector);
-
- Map<String, String> actualTranslatedConfig = connect.connectorInfo(connector).config();
- assertEquals(expectedTranslatedConfig, actualTranslatedConfig);
-
- // Swap source and target clusters--inadvisable in production, perfectly fine in a tightly-controlled testing environment
- connectorConfig = baseURep3Config(srcKafkaName, dstKafkaName, 1, Set.of(topic));
- connectorConfig.put(SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG, dstCluster.bootstrapServers());
- connectorConfig.put(TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG, srcCluster.bootstrapServers());
- connectorConfig.put(targetClientsReceiveBufferBytes, "66666");
- connectorConfig.put(targetClientsSendBufferBytes, "22222");
- connectorConfig.put(targetProducerLingerMs, "2000");
- connectorConfig.put(targetProducerSendBufferBytes, "44444");
-
- // Uses the PUT endpoint
- connect.configureConnector(connector, connectorConfig);
-
- expectedTranslatedConfig = new HashMap<>(connectorConfig);
- expectedTranslatedConfig.put(NAME_CONFIG, connector);
- expectedTranslatedConfig.put(
- CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG,
- srcCluster.bootstrapServers()
- );
- expectedTranslatedConfig.put(
- CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + CommonClientConfigs.RECEIVE_BUFFER_CONFIG,
- "66666"
- );
- expectedTranslatedConfig.put(
- CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + ProducerConfig.LINGER_MS_CONFIG,
- "2000"
- );
- expectedTranslatedConfig.put(
- CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + CommonClientConfigs.SEND_BUFFER_CONFIG,
- "44444"
- );
-
- actualTranslatedConfig = connect.connectorInfo(connector).config();
- assertEquals(expectedTranslatedConfig, actualTranslatedConfig);
-
- // Swap source and target back to the correct clusters
- Map<String, String> patch = new HashMap<>();
- patch.put(SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG, srcCluster.bootstrapServers());
- patch.put(TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG, dstCluster.bootstrapServers());
- patch.put(targetClientsReceiveBufferBytes, "99999");
- patch.put(targetClientsSendBufferBytes, "33333");
- patch.put(targetProducerLingerMs, "3000");
- patch.put(targetProducerSendBufferBytes, "66666");
-
- connect.patchConnectorConfig(connector, patch);
-
- actualTranslatedConfig = connect.connectorInfo(connector).config();
- // Only assert that our patch has gone through
- assertEquals(
- dstCluster.bootstrapServers(),
- actualTranslatedConfig.get(CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG)
- );
- assertEquals(
- srcCluster.bootstrapServers(),
- actualTranslatedConfig.get(SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG)
- );
- assertEquals(
- dstCluster.bootstrapServers(),
- actualTranslatedConfig.get(TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG)
- );
- assertEquals(
- "99999",
- actualTranslatedConfig.get(CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + CommonClientConfigs.RECEIVE_BUFFER_CONFIG)
- );
- assertEquals(
- "3000",
- actualTranslatedConfig.get(CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + ProducerConfig.LINGER_MS_CONFIG)
- );
- assertEquals(
- "66666",
- actualTranslatedConfig.get(CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + CommonClientConfigs.SEND_BUFFER_CONFIG)
- );
- }
-
- @Test
- public void testTopicRenaming() throws Exception {
- String srcKafkaName = "src-kafka";
- String dstKafkaName = "dst-kafka";
-
- EmbeddedKafkaCluster srcCluster = extraKafkaCluster(srcKafkaName, 1, new Properties());
- EmbeddedKafkaCluster dstCluster = extraKafkaCluster(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(conn1, srcKafkaName, dstKafkaName, 3, Set.of("s1", "s2", "s3"), Map.of("s1", "t1", "s2", "t2", "s3", "t3"), false);
- startConnector(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
- public void testTasksMax() throws Exception {
- String srcKafkaName = "src-kafka";
- String dstKafkaName = "dst-kafka";
-
- EmbeddedKafkaCluster srcCluster = extraKafkaCluster(srcKafkaName, 1, new Properties());
- EmbeddedKafkaCluster dstCluster = extraKafkaCluster(dstKafkaName, 1, new Properties());
-
- String t1 = "t1";
- srcCluster.createTopic(t1, 1);
- dstCluster.createTopic(t1, 1);
-
- String connector = "connector";
-
- Map<String, String> connectorConfig = baseURep3Config(srcKafkaName, dstKafkaName, 5, Set.of(t1));
- connectorConfig.put(SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG, srcCluster.bootstrapServers());
- connectorConfig.put(TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG, dstCluster.bootstrapServers());
- connectorConfig.put(CONNECTOR_CLIENT_PRODUCER_OVERRIDES_PREFIX + CommonClientConfigs.BOOTSTRAP_SERVERS_CONFIG, dstCluster.bootstrapServers());
- connect.configureConnector(connector, connectorConfig);
- // Only one partition to replicate; expect only one task, even though tasks.max is set to 5
- connect.assertions().assertConnectorAndExactlyNumTasksAreRunning(
- connector, 1, "Connector and/or expected number of tasks did not start in time"
- );
-
- String t2 = "t2";
- srcCluster.createTopic(t2, 2);
- dstCluster.createTopic(t2, 2);
-
- connectorConfig.put(TOPICS_CONFIG, String.join(",", List.of(t1, t2)));
- connect.configureConnector(connector, connectorConfig);
- // Three partitions to replicate; expect three tasks, even though tasks.max is set to 5
- stopAndAssertResume(connector, 3);
-
- String t3 = "t3";
- srcCluster.createTopic(t3, 3);
- dstCluster.createTopic(t3, 3);
-
- connectorConfig.put(TOPICS_CONFIG, String.join(",", List.of(t1, t2, t3)));
- connect.configureConnector(connector, connectorConfig);
- // Six partitions to replicate; expect five tasks, since tasks.max is set to 5
- stopAndAssertResume(connector, 5);
-
- connectorConfig.put(TASKS_MAX_CONFIG, "10");
- connectorConfig.put(TOPICS_CONFIG, String.join(",", List.of(t1, t2, t3)));
- connect.configureConnector(connector, connectorConfig);
- // Six partitions to replicate; expect six tasks now
- stopAndAssertResume(connector, 6);
-
- String t4 = "t4";
- srcCluster.createTopic(t4, 10);
- dstCluster.createTopic(t4, 10);
-
- connectorConfig.put(TOPICS_CONFIG, String.join(",", List.of(t1, t2, t3, t4)));
- connect.configureConnector(connector, connectorConfig);
- // Ten partitions to replicate and tasks.max is 10; however, only three workers in the cluster--should still only be six tasks (two per worker)
- stopAndAssertResume(connector, 6);
-
- WorkerHandle newWorker1 = connect.addWorker();
- connect.assertions().assertExactlyNumWorkersAreUp(4, "Additional worker did not start in time");
- connect.restartConnector(connector);
- // Four workers in the cluster--should now be eight tasks (two per worker)
- stopAndAssertResume(connector, 8);
-
- WorkerHandle newWorker2 = connect.addWorker();
- connect.assertions().assertExactlyNumWorkersAreUp(5, "Additional worker did not start in time");
- connect.restartConnector(connector);
- // Five workers in the cluster--should now be ten tasks (two per worker)
- stopAndAssertResume(connector, 10);
-
- WorkerHandle newWorker3 = connect.addWorker();
- connect.assertions().assertExactlyNumWorkersAreUp(6, "Additional worker did not start in time");
- connect.restartConnector(connector);
- // Six workers in the cluster, but tasks.max is still only 10
- stopAndAssertResume(connector, 10);
-
- // Prevent the shutdown of our extra workers from killing the JVM-global workload store
- KafkaBasedWorkloadProvider.destroyOnClose(false);
-
- newWorker3.stop();
- newWorker2.stop();
- newWorker1.stop();
-
- KafkaBasedWorkloadProvider.destroyOnClose(true);
- }
-
- private void stopAndAssertResume(String connector, int expectedRunningTasks) throws InterruptedException {
- connect.stopConnector(connector);
- connect.assertions().assertConnectorIsStopped(connector, "Connector did not stop in time");
- connect.resumeConnector(connector);
- connect.assertions().assertConnectorAndExactlyNumTasksAreRunning(connector, expectedRunningTasks, "Connector and/or tasks did not start back up in time");
- }
-
- @Test
- public void testPartitionCountTracking() throws Exception {
- String srcKafkaName = "src-kafka";
- String dstKafkaName = "dst-kafka";
-
- EmbeddedKafkaCluster srcCluster = extraKafkaCluster(srcKafkaName, 1, new Properties());
- EmbeddedKafkaCluster dstCluster = extraKafkaCluster(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,
- REFRESH_INTERVAL_MS,
- "1000"
- );
- startURep3Connector(
- conn2,
- srcKafkaName,
- dstKafkaName,
- 1,
- Set.of("t4", "t5", "t6"),
- null,
- REFRESH_INTERVAL_MS,
- "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
- final 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
- final int newNumDstPartitions = 8;
- 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);
- }
-
- Thread.sleep(5_000);
- Set<AssignedPartition> expectedConn1Assignment = IntStream.rangeClosed(1, 3)
- .boxed()
- .flatMap(t -> IntStream.range(0, newNumSrcPartitions)
- .mapToObj(p -> {
- String topic = "t" + t;
- int dstPartition = p % newNumDstPartitions;
- return new AssignedPartition(
- topic,
- p,
- null,
- dstPartition != p ? dstPartition : null
- );
- })
- ).collect(Collectors.toSet());
- Set<AssignedPartition> expectedConn2Assignment = IntStream.rangeClosed(4, 6)
- .boxed()
- .flatMap(t -> IntStream.range(0, newNumSrcPartitions)
- .mapToObj(p -> {
- String topic = "t" + t;
- int dstPartition = p % newNumDstPartitions;
- return new AssignedPartition(
- topic,
- p,
- null,
- dstPartition != p ? dstPartition : null
- );
- })
- ).collect(Collectors.toSet());
-
- // Wait for the connectors to pick up on the partition count changes
- waitForCondition(() -> {
- Set<AssignedPartition> actualConn1Assignment = getAllAssignedPartitions(conn1);
- assertEquals(expectedConn1Assignment, actualConn1Assignment);
-
- Set<AssignedPartition> actualConn2Assignment = getAllAssignedPartitions(conn2);
- assertEquals(expectedConn2Assignment, actualConn2Assignment);
- return true;
- }, TimeUnit.MINUTES.toMillis(1), "Connector did not generate new assignments in response to target partition count change in time");
-
- // Force task restarts to guarantee that the new assignments have been picked up
- connect.restartConnectorAndTasks(conn1, false, true, false);
- connect.restartConnectorAndTasks(conn2, false, true, false);
- connect.assertions().assertConnectorAndExactlyNumTasksAreRunning(conn1, 3, "Connector and/or tasks did not restart in time");
- connect.assertions().assertConnectorAndExactlyNumTasksAreRunning(conn2, 1, "Connector and/or tasks did not restart 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 % newNumDstPartitions));
- }).collect(Collectors.toSet());
- actualPartitions = records.partitions();
- assertEquals(expectedPartitions, actualPartitions);
- }
-
- @Test
- 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");
- // Simulate unreachable src/dst clusters
- 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 is available
- SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG,
- TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG
- );
-
- EmbeddedKafkaCluster srcCluster = extraKafkaCluster(srcKafkaName, 1, new Properties());
- connectorProps.put(SOURCE_CLUSTER_BOOTSTRAP_SERVERS_CONFIG, srcCluster.bootstrapServers());
- validateURep3Config(connectorProps,
- // Only the source cluster is available; 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 is available; target cluster is still down
- TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG
- );
-
- EmbeddedKafkaCluster dstCluster = extraKafkaCluster(dstKafkaName, 1, new Properties());
- connectorProps.put(TARGET_CLUSTER_BOOTSTRAP_SERVERS_CONFIG, dstCluster.bootstrapServers());
- validateURep3Config(connectorProps,
- // Both clusters available now, but the 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
- );
- }
-
- @Test
- public void testWorkloadTracking() throws Exception {
- String srcKafka1Name = "src-kafka";
- String srcKafka2Name = "src-kafka-2";
- String dstKafka1Name = "dst-kafka";
- String dstKafka2Name = "dst-kafka-2";
-
- EmbeddedKafkaCluster srcCluster1 = extraKafkaCluster(srcKafka1Name, 1, new Properties());
- EmbeddedKafkaCluster srcCluster2 = extraKafkaCluster(srcKafka2Name, 1, new Properties());
- EmbeddedKafkaCluster dstCluster1 = extraKafkaCluster(dstKafka1Name, 1, new Properties());
- EmbeddedKafkaCluster dstCluster2 = extraKafkaCluster(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(conn1, srcKafka1Name, dstKafka1Name, 6, Set.of("t1", "t2", "t3", "t4", "t5", "t6", "t7"), null, true);
- startConnector(conn2, srcKafka1Name, dstKafka2Name, 6, Set.of("t1", "t2", "t3", "t4", "t5", "t6", "t7"), null, true);
- // Only replicate a subset of topics from src 2
- startConnector(conn3, srcKafka2Name, dstKafka1Name, 6, Set.of("t1", "t3", "t4", "t6", "t7"), null, true);
- startConnector(conn4, srcKafka2Name, dstKafka2Name, 6, 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(() -> {
- Map<ReplicatedTopicPartition, Workload> trackedWorkloads = getAllWorkloads();
-
- Set<ReplicatedTopicPartition> expectedReplicatedTopicPartitions = new HashSet<>();
- for (int i = 1; i <= 7; i++) {
- 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), "");
-
- Map<String, Workload> workerWorkloads = getWorkerWorkloads();
- assertEquals(NUM_WORKERS, workerWorkloads.size());
- workerWorkloads.forEach((worker, workload) ->
- assertNotEquals(Workload.EMPTY, workload, "Worker " + worker + " has empty workload")
- );
- }
-
- private EmbeddedKafkaCluster extraKafkaCluster(String name, int numBrokers, Properties brokerProperties) {
- return EXTRA_KAFKA_CLUSTERS.computeIfAbsent(name, n -> {
- log.info("Starting embedded Kafka cluster {}", name);
- EmbeddedKafkaCluster result = new EmbeddedKafkaCluster(numBrokers, brokerProperties);
- result.start();
- log.info("Finished starting embedded Kafka cluster {}", name);
- return result;
- });
- }
-
- private void startConnector(
- String connectorName,
- String srcKafkaName,
- String dstKafkaName,
- int tasksMax,
- Set<String> topics,
- Map<String, String> topicRenames,
- boolean workloadTracking
- ) throws InterruptedException {
- if (workloadTracking)
- startURep3ConnectorWithWorkloadTracking(connectorName, srcKafkaName, dstKafkaName, tasksMax, topics, topicRenames);
- else
- startURep3Connector(connectorName, srcKafkaName, dstKafkaName, tasksMax, topics, topicRenames);
- }
-
- private void startURep3ConnectorWithWorkloadTracking(
- String connectorName,
- String srcKafkaName,
- String dstKafkaName,
- int tasksMax,
- Set<String> topics,
- Map<String, String> topicRenames
- ) throws InterruptedException {
- 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
- ) throws InterruptedException {
- assertEquals(0, extraProps.length % 2, "Extra props should be given as a series of key and value pairs");
- EmbeddedKafkaCluster srcCluster = EXTRA_KAFKA_CLUSTERS.get(srcKafkaName);
- assertNotNull(srcCluster, "No Kafka cluster found with name '" + srcKafkaName + "'");
-
- EmbeddedKafkaCluster dstCluster = EXTRA_KAFKA_CLUSTERS.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]);
-
- log.info("Starting connector {}", connectorName);
- connect.configureConnector(connectorName, connectorProps);
- connect.assertions().assertConnectorAndExactlyNumTasksAreRunning(connectorName, tasksMax, "Connector and/or tasks did not start in time");
- log.info("Finished starting connector {}", connectorName);
- }
-
- 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);
- }
-
- private void manuallyPublishWorkload(ReplicatedTopicPartition rtp, Workload w) throws IOException {
- String resource = String.format(
- "uber/workloads/topics/overrides/%s/%s/%s/%d",
- rtp.srcKafka(), rtp.dstKafka(), rtp.topic(), rtp.partition()
- );
- WorkloadEntity entity = WorkloadEntity.of(w);
- String body = OBJECT_MAPPER.writeValueAsString(entity);
- Response response = connect.requestPost(connect.endpointForResource(resource), body, Map.of());
- assertEquals(Response.Status.NO_CONTENT.getStatusCode(), response.getStatus());
- }
-
- private void manuallyDeleteWorkload(ReplicatedTopicPartition rtp) {
- String resource = String.format(
- "uber/workloads/topics/overrides/%s/%s/%s/%d",
- rtp.srcKafka(), rtp.dstKafka(), rtp.topic(), rtp.partition()
- );
- Response response = connect.requestDelete(connect.endpointForResource(resource));
- assertEquals(Response.Status.NO_CONTENT.getStatusCode(), response.getStatus());
- }
-
- private Map<ReplicatedTopicPartition, Workload> getAllWorkloads() throws IOException {
- Response response = connect.requestGet(connect.endpointForResource("uber/workloads/topics/all?refresh=true"));
- String entity = (String) response.getEntity();
-
- return OBJECT_MAPPER
- .readerFor(new TypeReference<List<PartitionWorkloadEntity>>() {})
- .<List<PartitionWorkloadEntity>>readValue(entity)
- .stream()
- .collect(Collectors.toMap(
- PartitionWorkloadEntity::replicatedTopicPartition,
- PartitionWorkloadEntity::workload
- ));
- }
-
- private Map<String, Workload> getWorkerWorkloads() throws IOException {
- Response response = connect.requestGet(connect.endpointForResource("uber/workloads/workers/all?refresh=true"));
- String entity = (String) response.getEntity();
-
- return OBJECT_MAPPER
- .readerFor(new TypeReference<List<WorkerWorkloadEntity>>() {})
- .<List<WorkerWorkloadEntity>>readValue(entity)
- .stream()
- .collect(Collectors.toMap(
- WorkerWorkloadEntity::worker,
- w -> w.totalWorkload().workload()
- ));
- }
-
- private Set<AssignedPartition> getAllAssignedPartitions(String connector) throws IOException {
- Response response = connect.requestGet(connect.endpointForResource("uber/assignments/" + connector + "/all"));
- String entity = (String) response.getEntity();
-
- return OBJECT_MAPPER
- .readerFor(new TypeReference<List<TaskAssignmentEntity>>() {})
- .<List<TaskAssignmentEntity>>readValue(entity)
- .stream()
- .map(TaskAssignmentEntity::partitions)
- .flatMap(List::stream)
- .map(AssignedPartitionEntity::assignedPartition)
- .collect(Collectors.toSet());
- }
-}
diff --git a/connect/uber/src/test/java/com/uber/data/kafka/connect/ureplicator3/TaskPartitionAssignorTest.java b/connect/uber/src/test/java/com/uber/data/kafka/connect/ureplicator3/TaskPartitionAssignorTest.java
deleted file mode 100644
index 26c206094ca..00000000000
--- a/connect/uber/src/test/java/com/uber/data/kafka/connect/ureplicator3/TaskPartitionAssignorTest.java
+++ /dev/null
@@ -1,50 +0,0 @@
-package com.uber.data.kafka.connect.ureplicator3;
-
-import org.apache.kafka.common.TopicPartition;
-
-import org.junit.jupiter.params.ParameterizedTest;
-import org.junit.jupiter.params.provider.CsvSource;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import java.util.Collection;
-import java.util.HashSet;
-import java.util.Map;
-import java.util.Set;
-import java.util.function.Function;
-import java.util.stream.Collectors;
-import java.util.stream.IntStream;
-
-public class TaskPartitionAssignorTest {
-
- private static final Logger log = LoggerFactory.getLogger(TaskPartitionAssignorTest.class);
-
- private static final String CONNECTOR = "connector";
- private static final String TOPIC = "topic";
-
- @ParameterizedTest
- @CsvSource({
- "1,100_000", "2,100_000", "5,100_000", "50,100_000", "500,100_000", "5_000,100_000", "50_000,100_000", "100_000,100_000", "200_000,100_000",
- "1,1_000_000", "2,1_000_000", "5,1_000_000", "50,1_000_000", "500,1_000_000", "5_000,1_000_000", "50_000,1_000_000", "500_000,1_000_000", "1_000_000,1_000_000"
- })
- public void testInitialAssignment(int numTasks, int numPartitions) {
- Collection<TaskPartitions> taskPartitions = IntStream.range(0, numTasks)
- .mapToObj(i -> new TaskPartitions(CONNECTOR, i, new HashSet<>()))
- .toList();
-
- Map<TopicPartition, Double> partitionLoads = IntStream.range(0, numPartitions)
- .mapToObj(i -> new TopicPartition(TOPIC, i))
- .collect(Collectors.toMap(
- Function.identity(),
- tp -> 0.1 * (tp.partition() % 100)
- ));
-
- Set<TopicPartition> toAssign = partitionLoads.keySet();
-
- TaskPartitionAssignor assignor = new TaskPartitionAssignor(0.1, partitionLoads);
-
- log.info("Beginning initial assignment");
- assignor.assign(taskPartitions, toAssign);
- log.info("Finished initial assignment");
- }
-}
diff --git a/connect/uber/src/test/java/com/uber/data/kafka/connect/utils/BrokerResolverTest.java b/connect/uber/src/test/java/com/uber/data/kafka/connect/utils/BrokerResolverTest.java
deleted file mode 100644
index a0d097cc333..00000000000
--- a/connect/uber/src/test/java/com/uber/data/kafka/connect/utils/BrokerResolverTest.java
+++ /dev/null
@@ -1,43 +0,0 @@
-package com.uber.data.kafka.connect.utils;
-
-import com.uber.data.kafka.exception.ResolverException;
-
-import org.junit.jupiter.api.Test;
-
-import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertThrows;
-
-public class BrokerResolverTest {
-
- private static final String SECPOC_DCA = "kafka-secpoc.dca.uber.internal:9092";
-
- private static final String DEV1_PHX = "kafka-dev1.phx.uber.internal:9092";
-
- @Test
- public void testBrokerResolver() {
- assertResolution("dca-secpoc", SECPOC_DCA);
- assertResolution("dca1-secpoc", SECPOC_DCA);
- assertResolution("kloak-dca-secpoc", SECPOC_DCA);
- assertResolution("kloak-dca1-secpoc", SECPOC_DCA);
-
- assertResolution("phx-dev1", DEV1_PHX);
- assertResolution("phx2-dev1", DEV1_PHX);
- assertResolution("kloak-phx-dev1", DEV1_PHX);
- assertResolution("kloak-phx2-dev1", DEV1_PHX);
-
- assertFailedResolution(null);
- assertFailedResolution("");
- assertFailedResolution("idklol");
- }
-
- private void assertResolution(String clusterName, String expected) {
- String actual = assertDoesNotThrow(() -> BrokerResolver.resolveBootstrap(clusterName));
- assertEquals(expected, actual);
- }
-
- private void assertFailedResolution(String clusterName) {
- assertThrows(ResolverException.class, () -> BrokerResolver.resolveBootstrap(clusterName));
- }
-
-}
diff --git a/connect/uber/src/test/resources/log4j2.yaml b/connect/uber/src/test/resources/log4j2.yaml
deleted file mode 100644
index 443e0c12923..00000000000
--- a/connect/uber/src/test/resources/log4j2.yaml
+++ /dev/null
@@ -1,47 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one or more
-# contributor license agreements. See the NOTICE file distributed with
-# this work for additional information regarding copyright ownership.
-# The ASF licenses this file to You under the Apache License, Version 2.0
-# (the "License"); you may not use this file except in compliance with
-# the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-Configuration:
- Properties:
- Property:
- - name: "logPattern"
- value: "[%d] %p %X{connector.context}%m (%c:%L)%n"
-
- Appenders:
- Console:
- name: STDOUT
- PatternLayout:
- pattern: "${logPattern}"
-
- Loggers:
- Root:
- level: INFO
- AppenderRef:
- - ref: STDOUT
- Logger:
- - name: kafka
- level: WARN
-
- - name: state.change.logger
- level: "OFF"
-
- - name: org.apache.kafka.connect
- level: DEBUG
-
- - name: com.uber.data.kafka.connect
- level: DEBUG
-
- - name: com.uber.data.kafka.connect.assignment.AbstractAssignor
- level: DEBUG
\ No newline at end of file
diff --git a/settings.gradle b/settings.gradle
index 20cccf47aa5..7c37a046838 100644
--- a/settings.gradle
+++ b/settings.gradle
@@ -74,7 +74,6 @@ include 'clients',
'connect:runtime',
'connect:test-plugins',
'connect:transforms',
- 'connect:uber',
'coordinator-common',
'core',
'examples',
diff --git a/streams/quickstart/java/pom.xml b/streams/quickstart/java/pom.xml
index 96de39980d7..c139244e9a5 100644
--- a/streams/quickstart/java/pom.xml
+++ b/streams/quickstart/java/pom.xml
@@ -26,7 +26,7 @@
<parent>
<groupId>org.apache.kafka</groupId>
<artifactId>streams-quickstart</artifactId>
- <version>4.3.0-SNAPSHOT</version>
+ <version>4.3.0-chrise-33-uber</version>
<relativePath>..</relativePath>
</parent>
diff --git a/streams/quickstart/java/src/main/resources/archetype-resources/pom.xml b/streams/quickstart/java/src/main/resources/archetype-resources/pom.xml
index 9af8d1a74c9..76f10d45f34 100644
--- a/streams/quickstart/java/src/main/resources/archetype-resources/pom.xml
+++ b/streams/quickstart/java/src/main/resources/archetype-resources/pom.xml
@@ -29,7 +29,7 @@
<properties>
<project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
- <kafka.version>4.3.0-SNAPSHOT</kafka.version>
+ <kafka.version>4.3.0-chrise-33-uber</kafka.version>
<slf4j.version>2.0.16</slf4j.version>
</properties>
diff --git a/streams/quickstart/pom.xml b/streams/quickstart/pom.xml
index f043621022f..a52e149a29a 100644
--- a/streams/quickstart/pom.xml
+++ b/streams/quickstart/pom.xml
@@ -22,7 +22,7 @@
<groupId>org.apache.kafka</groupId>
<artifactId>streams-quickstart</artifactId>
<packaging>pom</packaging>
- <version>4.3.0-SNAPSHOT</version>
+ <version>4.3.0-chrise-33-uber</version>
<name>Kafka Streams :: Quickstart</name>
diff --git a/tests/kafkatest/__init__.py b/tests/kafkatest/__init__.py
index 7df191bd7ce..33562f52f1a 100644
--- a/tests/kafkatest/__init__.py
+++ b/tests/kafkatest/__init__.py
@@ -22,4 +22,4 @@
# Instead, in development branches, the version should have a suffix of the form ".devN"
#
# For example, when Kafka is at version 1.0.0-SNAPSHOT, this should be something like "1.0.0.dev0"
-__version__ = '4.3.0.dev0'
+__version__ = '4.3.0-chrise-33-uber'
diff --git a/tests/kafkatest/version.py b/tests/kafkatest/version.py
index cf384aac835..8c7447ca3db 100644
--- a/tests/kafkatest/version.py
+++ b/tests/kafkatest/version.py
@@ -128,7 +128,7 @@ def get_version(node=None):
return DEV_BRANCH
DEV_BRANCH = KafkaVersion("dev")
-DEV_VERSION = KafkaVersion("4.3.0-SNAPSHOT")
+DEV_VERSION = KafkaVersion("4.3.0-chrise-33-uber")
LATEST_STABLE_TRANSACTION_VERSION = 2
# This should match the LATEST_PRODUCTION version defined in MetadataVersion.java