diff --git a/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java b/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java index 873c2a3e5c51b..3c433875ac03c 100644 --- a/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java +++ b/clients/src/main/java/org/apache/kafka/clients/admin/KafkaAdminClient.java @@ -1811,7 +1811,7 @@ static Cluster parseDescribeClusterResponse(DescribeClusterResponseData response List nodes = new ArrayList<>(); Node controllerNode = null; for (DescribeClusterResponseData.DescribeClusterBroker node : response.brokers()) { - Node newNode = new Node(node.brokerId(), node.host(), node.port(), node.rack()); + Node newNode = new Node(node.brokerId(), node.host(), node.port(), node.rack(), node.pod()); nodes.add(newNode); if (node.brokerId() == response.controllerId()) { controllerNode = newNode; diff --git a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/GroupCoordinatorNode.java b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/GroupCoordinatorNode.java index e6d98b6125036..b36833960accf 100644 --- a/clients/src/main/java/org/apache/kafka/clients/consumer/internals/GroupCoordinatorNode.java +++ b/clients/src/main/java/org/apache/kafka/clients/consumer/internals/GroupCoordinatorNode.java @@ -30,7 +30,7 @@ */ public class GroupCoordinatorNode extends Node { public GroupCoordinatorNode(int id, String host, int port) { - super(GroupCoordinatorNode.validateId(id), host, port, null, false, "+" + id); + super(GroupCoordinatorNode.validateId(id), host, port, null, null, false, "+" + id); } private static int validateId(int id) { diff --git a/clients/src/main/java/org/apache/kafka/common/Node.java b/clients/src/main/java/org/apache/kafka/common/Node.java index 013230ae01acc..8962f44169726 100644 --- a/clients/src/main/java/org/apache/kafka/common/Node.java +++ b/clients/src/main/java/org/apache/kafka/common/Node.java @@ -34,29 +34,35 @@ public class Node { private final String host; private final int port; private final String rack; + private final String pod; private final boolean isFenced; // Cache hashCode as it is called in performance sensitive parts of the code (e.g. RecordAccumulator.ready) private Integer hash; public Node(int id, String host, int port) { - this(id, host, port, null, false); + this(id, host, port, null, null, false); } public Node(int id, String host, int port, String rack) { - this(id, host, port, rack, false); + this(id, host, port, rack, null, false); } - public Node(int id, String host, int port, String rack, boolean isFenced) { - this(id, host, port, rack, isFenced, Integer.toString(id)); + public Node(int id, String host, int port, String rack, String pod) { + this(id, host, port, rack, pod, false); } - protected Node(int id, String host, int port, String rack, boolean isFenced, String idString) { + public Node(int id, String host, int port, String rack, String pod, boolean isFenced) { + this(id, host, port, rack, pod, isFenced, Integer.toString(id)); + } + + protected Node(int id, String host, int port, String rack, String pod, boolean isFenced, String idString) { this.id = id; this.idString = idString; this.host = host; this.port = port; this.rack = rack; + this.pod = pod; this.isFenced = isFenced; } @@ -127,6 +133,20 @@ public boolean isFenced() { return isFenced; } + /** + * True if this node has a defined pod + */ + public boolean hasPod() { + return pod != null; + } + + /** + * @return the pod of the node + */ + public String pod() { + return pod; + } + @Override public int hashCode() { Integer h = this.hash; @@ -135,6 +155,7 @@ public int hashCode() { result = 31 * result + id; result = 31 * result + port; result = 31 * result + ((rack == null) ? 0 : rack.hashCode()); + result = 31 * result + ((pod == null) ? 0 : pod.hashCode()); result = 31 * result + Objects.hashCode(isFenced); this.hash = result; return result; @@ -154,11 +175,12 @@ public boolean equals(Object obj) { port == other.port && Objects.equals(host, other.host) && Objects.equals(rack, other.rack) && + Objects.equals(pod, other.pod) && Objects.equals(isFenced, other.isFenced); } @Override public String toString() { - return host + ":" + port + " (id: " + idString + " rack: " + rack + " isFenced: " + isFenced + ")"; + return host + ":" + port + " (id: " + idString + " rack: " + rack + " pod: " + pod + " isFenced: " + isFenced + ")"; } } diff --git a/clients/src/main/java/org/apache/kafka/common/requests/DescribeClusterResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/DescribeClusterResponse.java index 0b8724a57602b..e262f7a006714 100644 --- a/clients/src/main/java/org/apache/kafka/common/requests/DescribeClusterResponse.java +++ b/clients/src/main/java/org/apache/kafka/common/requests/DescribeClusterResponse.java @@ -38,7 +38,7 @@ public DescribeClusterResponse(DescribeClusterResponseData data) { public Map nodes() { return data.brokers().valuesList().stream() - .map(b -> new Node(b.brokerId(), b.host(), b.port(), b.rack(), b.isFenced())) + .map(b -> new Node(b.brokerId(), b.host(), b.port(), b.rack(), b.pod(), b.isFenced())) .collect(Collectors.toMap(Node::id, Function.identity())); } diff --git a/clients/src/main/java/org/apache/kafka/common/requests/MetadataResponse.java b/clients/src/main/java/org/apache/kafka/common/requests/MetadataResponse.java index bbcebdd484618..cb045f60a27e1 100644 --- a/clients/src/main/java/org/apache/kafka/common/requests/MetadataResponse.java +++ b/clients/src/main/java/org/apache/kafka/common/requests/MetadataResponse.java @@ -442,7 +442,7 @@ private static class Holder { } private Map createBrokers(MetadataResponseData data) { - return data.brokers().valuesList().stream().map(b -> new Node(b.nodeId(), b.host(), b.port(), b.rack())) + return data.brokers().valuesList().stream().map(b -> new Node(b.nodeId(), b.host(), b.port(), b.rack(), b.pod())) .collect(Collectors.toMap(Node::id, Function.identity())); } @@ -503,7 +503,8 @@ public static MetadataResponse prepareResponse(boolean hasReliableEpoch, .setNodeId(broker.id()) .setHost(broker.host()) .setPort(broker.port()) - .setRack(broker.rack())) + .setRack(broker.rack()) + .setPod(broker.pod())) ); responseData.setClusterId(clusterId); diff --git a/clients/src/main/resources/common/message/BrokerRegistrationRequest.json b/clients/src/main/resources/common/message/BrokerRegistrationRequest.json index 04cd1b65a79ca..5f14572d6072a 100644 --- a/clients/src/main/resources/common/message/BrokerRegistrationRequest.json +++ b/clients/src/main/resources/common/message/BrokerRegistrationRequest.json @@ -60,6 +60,8 @@ { "name": "LogDirs", "type": "[]uuid", "versions": "2+", "about": "Log directories configured in this broker which are available.", "ignorable": true }, { "name": "PreviousBrokerEpoch", "type": "int64", "versions": "3+", "default": "-1", "ignorable": true, - "about": "The epoch before a clean shutdown." } + "about": "The epoch before a clean shutdown." }, + { "name": "Pod", "type": "string", "versions": "0+", "nullableVersions": "0+", "taggedVersions": "0+", "tag": 0, "default": "null", + "about": "The Pod which this broker is in." } ] } diff --git a/clients/src/main/resources/common/message/DescribeClusterResponse.json b/clients/src/main/resources/common/message/DescribeClusterResponse.json index 1911b1ec33d6f..9cbada0f0c7ce 100644 --- a/clients/src/main/resources/common/message/DescribeClusterResponse.json +++ b/clients/src/main/resources/common/message/DescribeClusterResponse.json @@ -48,7 +48,9 @@ { "name": "Rack", "type": "string", "versions": "0+", "nullableVersions": "0+", "default": "null", "about": "The rack of the broker, or null if it has not been assigned to a rack." }, { "name": "IsFenced", "type": "bool", "versions": "2+", - "about": "Whether the broker is fenced" } + "about": "Whether the broker is fenced" }, + { "name": "Pod", "type": "string", "versions": "0+", "nullableVersions": "0+", "default": "null", + "about": "The pod of the broker, or null if it has not been assigned to a pod.", "taggedVersions": "0+", "tag": 0 } ]}, { "name": "ClusterAuthorizedOperations", "type": "int32", "versions": "0+", "default": "-2147483648", "about": "32-bit bitfield to represent authorized operations for this cluster." } diff --git a/clients/src/main/resources/common/message/MetadataResponse.json b/clients/src/main/resources/common/message/MetadataResponse.json index 07ee7010e5e54..f29589cec8435 100644 --- a/clients/src/main/resources/common/message/MetadataResponse.json +++ b/clients/src/main/resources/common/message/MetadataResponse.json @@ -56,7 +56,9 @@ { "name": "Port", "type": "int32", "versions": "0+", "about": "The broker port." }, { "name": "Rack", "type": "string", "versions": "1+", "nullableVersions": "1+", "ignorable": true, "default": "null", - "about": "The rack of the broker, or null if it has not been assigned to a rack." } + "about": "The rack of the broker, or null if it has not been assigned to a rack." }, + { "name": "Pod", "type": "string", "versions": "11+", "nullableVersions": "11+", "ignorable": true, "default": "null", + "about": "The pod of the broker, or null if it has not been assigned to a pod.", "taggedVersions": "11+", "tag": 0 } ]}, { "name": "ClusterId", "type": "string", "nullableVersions": "2+", "versions": "2+", "ignorable": true, "default": "null", "about": "The cluster ID that responding broker belongs to." }, diff --git a/clients/src/test/java/org/apache/kafka/clients/MetadataTest.java b/clients/src/test/java/org/apache/kafka/clients/MetadataTest.java index aff3c3cc473eb..9bdaa92daff94 100644 --- a/clients/src/test/java/org/apache/kafka/clients/MetadataTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/MetadataTest.java @@ -905,7 +905,8 @@ private MetadataResponseBrokerCollection buildBrokerCollection(List nodes) .setNodeId(node.id()) .setHost(node.host()) .setPort(node.port()) - .setRack(node.rack()); + .setRack(node.rack()) + .setPod(node.pod()); brokers.add(broker); } return brokers; diff --git a/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java b/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java index 0162d2faa4962..fcbd725b0d5f9 100644 --- a/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java +++ b/clients/src/test/java/org/apache/kafka/clients/admin/KafkaAdminClientTest.java @@ -3002,7 +3002,8 @@ private static DescribeClusterResponse prepareDescribeClusterResponse( .setHost(broker.host()) .setPort(broker.port()) .setBrokerId(broker.id()) - .setRack(broker.rack()))); + .setRack(broker.rack()) + .setPod(broker.pod()))); return new DescribeClusterResponse(data); } @@ -10536,6 +10537,7 @@ public void testListTransactions() throws Exception { .setNodeId(node.id()) .setPort(node.port()) .setRack(node.rack()) + .setPod(node.pod()) ) ); diff --git a/connect/runtime/src/test/java/org/apache/kafka/connect/util/TopicAdminTest.java b/connect/runtime/src/test/java/org/apache/kafka/connect/util/TopicAdminTest.java index b40683865b454..1ea3a5eba33ca 100644 --- a/connect/runtime/src/test/java/org/apache/kafka/connect/util/TopicAdminTest.java +++ b/connect/runtime/src/test/java/org/apache/kafka/connect/util/TopicAdminTest.java @@ -928,7 +928,8 @@ private DescribeClusterResponse describeClusterResponse(Cluster cluster) { .setHost(broker.host()) .setPort(broker.port()) .setBrokerId(broker.id()) - .setRack(broker.rack()))); + .setRack(broker.rack()) + .setPod(broker.pod()))); return new DescribeClusterResponse(data); } diff --git a/core/src/main/scala/kafka/server/ControllerServer.scala b/core/src/main/scala/kafka/server/ControllerServer.scala index c2906e4c6745e..0065b4228b221 100644 --- a/core/src/main/scala/kafka/server/ControllerServer.scala +++ b/core/src/main/scala/kafka/server/ControllerServer.scala @@ -40,6 +40,7 @@ import org.apache.kafka.image.publisher.{ControllerRegistrationsPublisher, KRaft import org.apache.kafka.metadata.{KafkaConfigSchema, KRaftMetadataCache, ListenerInfo} import org.apache.kafka.metadata.authorizer.ClusterMetadataAuthorizer import org.apache.kafka.metadata.bootstrap.BootstrapMetadata +import org.apache.kafka.metadata.placement.{PodReplicaPlacer, StripedReplicaPlacer} import org.apache.kafka.metadata.publisher.{AclPublisher, DelegationTokenPublisher, DynamicClientQuotaPublisher, DynamicTopicClusterQuotaPublisher, FeaturesPublisher, ScramPublisher} import org.apache.kafka.raft.QuorumConfig import org.apache.kafka.security.{CredentialProvider, DelegationTokenManager} @@ -58,7 +59,7 @@ import org.apache.kafka.server.NodeToControllerChannelManagerImpl import org.apache.kafka.server.RaftControllerNodeProvider import java.util -import java.util.{Optional, OptionalLong} +import java.util.{Optional, OptionalLong, Random} import java.util.concurrent.locks.ReentrantLock import java.util.concurrent.{CompletableFuture, TimeUnit} import scala.jdk.CollectionConverters._ @@ -248,6 +249,7 @@ class ControllerServer( setQuorumFeatures(quorumFeatures). setDefaultReplicationFactor(config.defaultReplicationFactor.toShort). setDefaultNumPartitions(config.numPartitions.intValue()). + setReplicaPlacer(new PodReplicaPlacer(new StripedReplicaPlacer(new Random), config.canarySpec.toMap)). setSessionTimeoutNs(TimeUnit.NANOSECONDS.convert(config.brokerSessionTimeoutMs.longValue(), TimeUnit.MILLISECONDS)). setLeaderImbalanceCheckIntervalNs(leaderImbalanceCheckIntervalNs). diff --git a/core/src/main/scala/kafka/server/KafkaApis.scala b/core/src/main/scala/kafka/server/KafkaApis.scala index d9d5879dac57e..4e33f8847988e 100644 --- a/core/src/main/scala/kafka/server/KafkaApis.scala +++ b/core/src/main/scala/kafka/server/KafkaApis.scala @@ -2537,6 +2537,7 @@ class KafkaApis(val requestChannel: RequestChannel, setHost(node.host). setPort(node.port). setRack(node.rack). + setPod(node.pod). setIsFenced(node.isFenced)) } } diff --git a/core/src/main/scala/kafka/server/KafkaConfig.scala b/core/src/main/scala/kafka/server/KafkaConfig.scala index 03f2c46a89095..ffa6ffd5b03eb 100755 --- a/core/src/main/scala/kafka/server/KafkaConfig.scala +++ b/core/src/main/scala/kafka/server/KafkaConfig.scala @@ -36,6 +36,7 @@ import org.apache.kafka.coordinator.group.Group.GroupType import org.apache.kafka.coordinator.group.modern.share.ShareGroupConfig import org.apache.kafka.coordinator.group.GroupCoordinatorConfig import org.apache.kafka.coordinator.share.ShareCoordinatorConfig +import org.apache.kafka.metadata.placement.CanarySpec import org.apache.kafka.network.SocketServerConfigs import org.apache.kafka.raft.{KRaftConfigs, MetadataLogConfig, QuorumConfig} import org.apache.kafka.security.authorizer.AuthorizerUtils @@ -163,6 +164,7 @@ class KafkaConfig private(doLog: Boolean, val props: util.Map[_, _]) val brokerSessionTimeoutMs: Int = getInt(KRaftConfigs.BROKER_SESSION_TIMEOUT_MS_CONFIG) val controllerPerformanceSamplePeriodMs: Long = getLong(KRaftConfigs.CONTROLLER_PERFORMANCE_SAMPLE_PERIOD_MS) val controllerPerformanceAlwaysLogThresholdMs: Long = getLong(KRaftConfigs.CONTROLLER_PERFORMANCE_ALWAYS_LOG_THRESHOLD_MS) + val canarySpec = new CanarySpec(getString(KRaftConfigs.CANARY_POD_NAME), getInt(KRaftConfigs.CANARY_PARTITION_INTERVAL)) private def parseProcessRoles(): Set[ProcessRole] = { val roles = getList(KRaftConfigs.PROCESS_ROLES_CONFIG).asScala.map { diff --git a/core/src/test/scala/unit/kafka/server/BrokerLifecycleManagerTest.scala b/core/src/test/scala/unit/kafka/server/BrokerLifecycleManagerTest.scala index 51d51967871fe..1dcb62f6dbbc2 100644 --- a/core/src/test/scala/unit/kafka/server/BrokerLifecycleManagerTest.scala +++ b/core/src/test/scala/unit/kafka/server/BrokerLifecycleManagerTest.scala @@ -20,6 +20,7 @@ package kafka.server import java.util import java.util.{Collections, OptionalLong, Properties} import kafka.utils.TestUtils +import org.apache.kafka.clients.MockClient.RequestMatcher import org.apache.kafka.common.Node import org.apache.kafka.common.Uuid import org.apache.kafka.common.message.{BrokerHeartbeatResponseData, BrokerRegistrationResponseData} @@ -27,10 +28,11 @@ import org.apache.kafka.common.protocol.Errors import org.apache.kafka.common.requests.{AbstractRequest, AbstractResponse, BrokerHeartbeatRequest, BrokerHeartbeatResponse, BrokerRegistrationRequest, BrokerRegistrationResponse} import org.apache.kafka.metadata.BrokerState import org.apache.kafka.raft.{KRaftConfigs, QuorumConfig} -import org.apache.kafka.server.config.ServerLogConfigs +import org.apache.kafka.server.config.{ServerLogConfigs, ServerConfigs} import org.apache.kafka.server.BrokerLifecycleManager import org.junit.jupiter.api.Assertions._ -import org.junit.jupiter.api.{AfterEach, Test, Timeout} +import org.junit.jupiter.api.{AfterEach, Assertions, Test, Timeout} + import java.util.concurrent.{CompletableFuture, Future} import scala.jdk.CollectionConverters._ @@ -59,6 +61,7 @@ class BrokerLifecycleManagerTest { properties.setProperty(KRaftConfigs.CONTROLLER_LISTENER_NAMES_CONFIG, "SSL") properties.setProperty(KRaftConfigs.INITIAL_BROKER_REGISTRATION_TIMEOUT_MS_CONFIG, "300000") properties.setProperty(KRaftConfigs.BROKER_HEARTBEAT_INTERVAL_MS_CONFIG, "100") + properties.setProperty(ServerConfigs.BROKER_POD_CONFIG, "pod1") properties } @@ -97,7 +100,12 @@ class BrokerLifecycleManagerTest { assertEquals(1, context.mockChannelManager.unsentQueue.size) assertEquals(10L, context.mockChannelManager.unsentQueue.getFirst.request.build().asInstanceOf[BrokerRegistrationRequest].data().previousBrokerEpoch()) } - context.mockClient.prepareResponseFrom(new BrokerRegistrationResponse( + val podMatches: RequestMatcher = { request => + Assertions.assertNull(request.asInstanceOf[BrokerRegistrationRequest].data().rack()) + Assertions.assertEquals("pod1", request.asInstanceOf[BrokerRegistrationRequest].data().pod()) + true + } + context.mockClient.prepareResponseFrom(podMatches, new BrokerRegistrationResponse( new BrokerRegistrationResponseData().setBrokerEpoch(1000)), controllerNode) TestUtils.retry(10000) { context.poll() diff --git a/core/src/test/scala/unit/kafka/server/DescribeClusterRequestTest.scala b/core/src/test/scala/unit/kafka/server/DescribeClusterRequestTest.scala index 96283d603892d..b374aea2d450d 100644 --- a/core/src/test/scala/unit/kafka/server/DescribeClusterRequestTest.scala +++ b/core/src/test/scala/unit/kafka/server/DescribeClusterRequestTest.scala @@ -62,8 +62,8 @@ class DescribeClusterRequestTest extends BaseRequestTest { .setBrokerId(server.config.brokerId) .setHost("localhost") .setPort(server.socketServer.boundPort(listenerName)) - .setRack(server.config.rack.orElse(null)) - }.toSet + .setRack(server.config.rack().orElse(null)) + .setPod(server.config.pod().orElse(null))}.toSet val expectedClusterId = brokers.last.clusterId diff --git a/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala b/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala index db7dbd8cab619..7dc47afc45e72 100644 --- a/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala +++ b/core/src/test/scala/unit/kafka/server/KafkaApisTest.scala @@ -4418,9 +4418,10 @@ class KafkaApisTest extends Logging { .setSecurityProtocol(SecurityProtocol.PLAINTEXT.id) .setName(plaintextListener.value) ) + MetadataCacheFixtures.updateCache(metadataCache, - util.List.of(new RegisterBrokerRecord().setBrokerId(0).setRack("rack").setFenced(false).setEndPoints(endpoints)) - ) + util.List.of(new RegisterBrokerRecord().setBrokerId(0).setRack("rack").setPod("pod").setFenced(false).setEndPoints(endpoints)) +) // 2. Set up authorizer val authorizer: Authorizer = mock(classOf[Authorizer]) @@ -10177,6 +10178,7 @@ class KafkaApisTest extends Logging { util.List.of(new RegisterBrokerRecord() .setBrokerId(brokerId) .setRack("rack") + .setPod("pod") .setFenced(false) .setEndPoints(endpoints))) @@ -10201,7 +10203,6 @@ class KafkaApisTest extends Logging { private def updateMetadataCacheWithInconsistentListeners(): (ListenerName, ListenerName) = { val plaintextListener = ListenerName.forSecurityProtocol(SecurityProtocol.PLAINTEXT) val anotherListener = new ListenerName("LISTENER2") - val endpoints0 = new BrokerEndpointCollection() endpoints0.add( new BrokerEndpoint() @@ -10229,8 +10230,7 @@ class KafkaApisTest extends Logging { MetadataCacheFixtures.updateCache(metadataCache, util.List.of(new RegisterBrokerRecord().setBrokerId(0).setRack("rack").setFenced(false).setEndPoints(endpoints0), - new RegisterBrokerRecord().setBrokerId(1).setRack("rack").setFenced(false).setEndPoints(endpoints1)) - ) + new RegisterBrokerRecord().setBrokerId(1).setRack("rack").setPod("pod").setFenced(false).setEndPoints(endpoints1))) (plaintextListener, anotherListener) } @@ -10483,6 +10483,7 @@ class KafkaApisTest extends Logging { new RegisterBrokerRecord() .setBrokerId(brokerId) .setRack("rack") + .setPod("pod") .setFenced(false) .setEndPoints(endpoints) .setBrokerEpoch(brokerEpoch) diff --git a/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala b/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala index 34f4863558177..ab213ce48b5e1 100755 --- a/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala +++ b/core/src/test/scala/unit/kafka/server/KafkaConfigTest.scala @@ -19,7 +19,7 @@ package kafka.server import java.net.InetSocketAddress import java.util -import java.util.{Arrays, Collections, Properties} +import java.util.{Arrays, Collections, Optional, Properties} import kafka.utils.TestUtils.assertBadConfigContainingMessage import kafka.utils.TestUtils import org.apache.kafka.common.{Endpoint, Node} @@ -823,6 +823,8 @@ class KafkaConfigTest { case MetadataLogConfig.INTERNAL_METADATA_MAX_BATCH_SIZE_IN_BYTES_CONFIG => // no op case MetadataLogConfig.INTERNAL_METADATA_DELETE_DELAY_MILLIS_CONFIG => // no op case KRaftConfigs.CONTROLLER_LISTENER_NAMES_CONFIG => // ignore string + case KRaftConfigs.CANARY_POD_NAME => assertPropertyInvalid(baseProperties, name, " ") + case KRaftConfigs.CANARY_PARTITION_INTERVAL => assertPropertyInvalid(baseProperties, name, "not_a_number", "-1") case MetadataLogConfig.METADATA_MAX_IDLE_INTERVAL_MS_CONFIG => assertPropertyInvalid(baseProperties, name, "not_a_number") case ServerConfigs.AUTHORIZER_CLASS_NAME_CONFIG => //ignore string @@ -919,6 +921,7 @@ class KafkaConfigTest { case MetricConfigs.METRIC_REPORTER_CLASSES_CONFIG => // ignore string case MetricConfigs.METRIC_RECORDING_LEVEL_CONFIG => // ignore string case ServerConfigs.BROKER_RACK_CONFIG => // ignore string + case ServerConfigs.BROKER_POD_CONFIG => // ignore string case ServerConfigs.COMPRESSION_GZIP_LEVEL_CONFIG => assertPropertyInvalid(baseProperties, name, "not_a_number", "0") case ServerConfigs.COMPRESSION_LZ4_LEVEL_CONFIG => assertPropertyInvalid(baseProperties, name, "not_a_number", "0") @@ -2031,4 +2034,19 @@ class KafkaConfigTest { "Please increase broker.session.timeout.ms or decrease broker.heartbeat.interval.ms.")) } } + + @Test + def testPodAndRackProperties(): Unit = { + val props = new Properties() + props.put(KRaftConfigs.PROCESS_ROLES_CONFIG, "broker") + props.setProperty(QuorumConfig.QUORUM_VOTERS_CONFIG, "2@localhost:9093") + props.setProperty(KRaftConfigs.CONTROLLER_LISTENER_NAMES_CONFIG, "CONTROLLER") + props.put(KRaftConfigs.NODE_ID_CONFIG, "1") + props.put(ServerConfigs.BROKER_RACK_CONFIG, "rack-1") + props.put(ServerConfigs.BROKER_POD_CONFIG, "pod-1") + assertTrue(isValidKafkaConfig(props)) + val config = KafkaConfig.fromProps(props) + assertEquals(Optional.of("rack-1"), config.rack) + assertEquals(Optional.of("pod-1"), config.pod) + } } diff --git a/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala b/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala index 31e836123bffe..cdc85ee7401e6 100644 --- a/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala +++ b/core/src/test/scala/unit/kafka/server/ReplicaManagerTest.scala @@ -1397,8 +1397,8 @@ class ReplicaManagerTest { try { val leaderBrokerId = 0 val followerBrokerId = 1 - val leaderNode = new Node(leaderBrokerId, "host1", 0, "rack-a") - val followerNode = new Node(followerBrokerId, "host2", 1, "rack-b") + val leaderNode = new Node(leaderBrokerId, "host1", 0, "rack-a", "pod-a") + val followerNode = new Node(followerBrokerId, "host2", 1, "rack-b", "pod-b") val brokerList = Seq[Integer](leaderBrokerId, followerBrokerId).asJava val tp0 = new TopicPartition(topic, 0) val tidp0 = new TopicIdPartition(topicId, tp0) @@ -1508,8 +1508,8 @@ class ReplicaManagerTest { tp0, new ListenerName("default") )).thenReturn(util.Map.of( - leaderBrokerId, new Node(leaderBrokerId, "host1", 9092, "rack-a"), - followerBrokerId, new Node(followerBrokerId, "host2", 9092, "rack-b") + leaderBrokerId, new Node(leaderBrokerId, "host1", 9092, "rack-a", "pod-a"), + followerBrokerId, new Node(followerBrokerId, "host2", 9092, "rack-b", "pod-b") )) // Make this replica the leader @@ -2582,8 +2582,8 @@ class ReplicaManagerTest { mockGetAliveBrokerFunctions(metadataCache, aliveBrokers) when(metadataCache.getPartitionReplicaEndpoints( any[TopicPartition], any[ListenerName])). - thenReturn(util.Map.of(leaderBrokerId, new Node(leaderBrokerId, "host1", 9092, "rack-a"), - followerBrokerId, new Node(followerBrokerId, "host2", 9092, "rack-b"))) + thenReturn(util.Map.of(leaderBrokerId, new Node(leaderBrokerId, "host1", 9092, "rack-a", "pod-a"), + followerBrokerId, new Node(followerBrokerId, "host2", 9092, "rack-b", "pod-b"))) when(metadataCache.metadataVersion()).thenReturn(MetadataVersion.MINIMUM_VERSION) when(metadataCache.getAliveBrokerEpoch(leaderBrokerId)).thenReturn(util.Optional.of(brokerEpoch)) val mockProducePurgatory = new DelayedOperationPurgatory[DelayedProduce]( diff --git a/metadata/src/main/java/org/apache/kafka/controller/BrokerHeartbeatManager.java b/metadata/src/main/java/org/apache/kafka/controller/BrokerHeartbeatManager.java index 48ca047508440..0a198da49163b 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/BrokerHeartbeatManager.java +++ b/metadata/src/main/java/org/apache/kafka/controller/BrokerHeartbeatManager.java @@ -329,22 +329,26 @@ void maybeUpdateControlledShutdownOffset(int brokerId, long controlledShutDownOf Iterator usableBrokers( Function> idToRack, + Function> idToPod, Function hasUncordonedDirs ) { - return new UsableBrokerIterator(brokers.values().iterator(), idToRack, hasUncordonedDirs); + return new UsableBrokerIterator(brokers.values().iterator(), idToRack, idToPod, hasUncordonedDirs); } static class UsableBrokerIterator implements Iterator { private final Iterator iterator; private final Function> idToRack; + private final Function> idToPod; private final Function hasUncordonedDirs; private UsableBroker next; UsableBrokerIterator(Iterator iterator, Function> idToRack, + Function> idToPod, Function hasUncordonedDirs) { this.iterator = iterator; this.idToRack = idToRack; + this.idToPod = idToPod; this.hasUncordonedDirs = hasUncordonedDirs; this.next = null; } @@ -362,7 +366,8 @@ public boolean hasNext() { result = iterator.next(); } while (result.shuttingDown() || !hasUncordonedDirs.apply(result.id())); Optional rack = idToRack.apply(result.id()); - next = new UsableBroker(result.id(), rack, result.fenced()); + Optional pod = idToPod.apply(result.id()); + next = new UsableBroker(result.id(), rack, pod, result.fenced()); return true; } diff --git a/metadata/src/main/java/org/apache/kafka/controller/ClusterControlManager.java b/metadata/src/main/java/org/apache/kafka/controller/ClusterControlManager.java index b7665eeb8c069..f5f277708c81b 100644 --- a/metadata/src/main/java/org/apache/kafka/controller/ClusterControlManager.java +++ b/metadata/src/main/java/org/apache/kafka/controller/ClusterControlManager.java @@ -405,6 +405,7 @@ public ControllerResult registerBroker( setIsMigratingZkBroker(request.isMigratingZkBroker()). setIncarnationId(request.incarnationId()). setRack(request.rack()). + setPod(request.pod()). setEndPoints(listenerInfo.toBrokerRegistrationRecord()); // Track which finalized features we have not yet verified are supported by the broker. @@ -570,6 +571,7 @@ public void replay(RegisterBrokerRecord record, long offset) { setListeners(listenerInfo.listeners()). setSupportedFeatures(features). setRack(Optional.ofNullable(record.rack())). + setPod(Optional.ofNullable(record.pod())). setFenced(record.fenced()). setInControlledShutdown(record.inControlledShutdown()). setIsMigratingZkBroker(record.isMigratingZkBroker()). @@ -712,6 +714,7 @@ Iterator usableBrokers() { } return heartbeatManager.usableBrokers( id -> brokerRegistrations.get(id).rack(), + id -> brokerRegistrations.get(id).pod(), id -> brokerRegistrations.get(id).hasUncordonedDirs()); } diff --git a/metadata/src/main/java/org/apache/kafka/metadata/BrokerRegistration.java b/metadata/src/main/java/org/apache/kafka/metadata/BrokerRegistration.java index 0545869901c4d..0ce347c36a4f0 100644 --- a/metadata/src/main/java/org/apache/kafka/metadata/BrokerRegistration.java +++ b/metadata/src/main/java/org/apache/kafka/metadata/BrokerRegistration.java @@ -51,6 +51,7 @@ public static class Builder { private Map listeners; private Map supportedFeatures; private Optional rack; + private Optional pod; private boolean fenced; private boolean inControlledShutdown; private boolean isMigratingZkBroker; @@ -64,6 +65,7 @@ public Builder() { this.listeners = new HashMap<>(); this.supportedFeatures = new HashMap<>(); this.rack = Optional.empty(); + this.pod = Optional.empty(); this.fenced = false; this.inControlledShutdown = false; this.isMigratingZkBroker = false; @@ -112,6 +114,12 @@ public Builder setRack(Optional rack) { return this; } + public Builder setPod(Optional pod) { + Objects.requireNonNull(pod); + this.pod = pod; + return this; + } + public Builder setFenced(boolean fenced) { this.fenced = fenced; return this; @@ -145,6 +153,7 @@ public BrokerRegistration build() { listeners, supportedFeatures, rack, + pod, fenced, inControlledShutdown, isMigratingZkBroker, @@ -159,6 +168,7 @@ public BrokerRegistration build() { private final Map listeners; private final Map supportedFeatures; private final Optional rack; + private final Optional pod; private final boolean fenced; private final boolean inControlledShutdown; private final boolean isMigratingZkBroker; @@ -172,6 +182,7 @@ private BrokerRegistration( Map listeners, Map supportedFeatures, Optional rack, + Optional pod, boolean fenced, boolean inControlledShutdown, boolean isMigratingZkBroker, @@ -192,6 +203,7 @@ private BrokerRegistration( Objects.requireNonNull(supportedFeatures); this.supportedFeatures = new HashMap<>(supportedFeatures); this.rack = rack; + this.pod = pod; this.fenced = fenced; this.inControlledShutdown = inControlledShutdown; this.isMigratingZkBroker = isMigratingZkBroker; @@ -220,6 +232,7 @@ public static BrokerRegistration fromRecord(RegisterBrokerRecord record) { listeners, supportedFeatures, Optional.ofNullable(record.rack()), + Optional.ofNullable(record.pod()), record.fenced(), record.inControlledShutdown(), record.isMigratingZkBroker(), @@ -248,7 +261,7 @@ public Optional node(String listenerName) { if (endpoint == null) { return Optional.empty(); } - return Optional.of(new Node(id, endpoint.host(), endpoint.port(), rack.orElse(null), fenced)); + return Optional.of(new Node(id, endpoint.host(), endpoint.port(), rack.orElse(null), pod.orElse(null), fenced)); } public List nodes() { @@ -263,6 +276,10 @@ public Optional rack() { return rack; } + public Optional pod() { + return pod; + } + public boolean fenced() { return fenced; } @@ -323,6 +340,7 @@ public ApiMessageAndVersion toRecord(ImageWriterOptions options) { RegisterBrokerRecord registrationRecord = new RegisterBrokerRecord(). setBrokerId(id). setRack(rack.orElse(null)). + setPod(pod.orElse(null)). setBrokerEpoch(epoch). setIncarnationId(incarnationId). setFenced(fenced). @@ -371,7 +389,7 @@ public ApiMessageAndVersion toRecord(ImageWriterOptions options) { @Override public int hashCode() { return Objects.hash(id, epoch, incarnationId, listeners, supportedFeatures, - rack, fenced, inControlledShutdown, isMigratingZkBroker, directories, cordonedDirectories); + rack, pod, fenced, inControlledShutdown, isMigratingZkBroker, directories, cordonedDirectories); } @Override @@ -383,6 +401,7 @@ public boolean equals(Object o) { other.listeners.equals(listeners) && other.supportedFeatures.equals(supportedFeatures) && other.rack.equals(rack) && + other.pod.equals(pod) && other.fenced == fenced && other.inControlledShutdown == inControlledShutdown && other.isMigratingZkBroker == isMigratingZkBroker && @@ -405,6 +424,7 @@ public String toString() { collect(Collectors.joining(", ")) + "}" + ", rack=" + rack + + ", pod=" + pod + ", fenced=" + fenced + ", inControlledShutdown=" + inControlledShutdown + ", isMigratingZkBroker=" + isMigratingZkBroker + @@ -437,6 +457,7 @@ public BrokerRegistration cloneWith( listeners, supportedFeatures, rack, + pod, newFenced, newInControlledShutdownChange, isMigratingZkBroker, diff --git a/metadata/src/main/java/org/apache/kafka/metadata/placement/CanarySpec.java b/metadata/src/main/java/org/apache/kafka/metadata/placement/CanarySpec.java new file mode 100644 index 0000000000000..ba63d757af5bf --- /dev/null +++ b/metadata/src/main/java/org/apache/kafka/metadata/placement/CanarySpec.java @@ -0,0 +1,54 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.metadata.placement; + +import java.util.Collections; +import java.util.Map; +import java.util.function.Predicate; + +/** + * The class defined specification of kafka canary + * partition meet defined condition should be placed in canary pod + * The condition is controlled by an interval, which designates every Nth partition as a canary partition. + * for example, + * if interval is 10, partition 9, 19, 29,,, etc will be canary partition + * if interval is 50, partition 49, 99, 149 will be canary partition + *

+ * Because the last partition of each group is selected, a topic with fewer than {@code interval} + * partitions has no canary partition at all. + *

+ */ +public class CanarySpec { + protected static final Predicate ALWAYS_FALSE = i -> false; + + private final String canaryPodName; + private final Predicate predicate; + + public CanarySpec(String canaryPodName, int interval) { + this.canaryPodName = canaryPodName; + if (interval > 0) { + predicate = i -> i % interval == interval - 1; + } else { + predicate = ALWAYS_FALSE; + } + } + + public Map> toMap() { + return Collections.singletonMap(canaryPodName, predicate); + } +} diff --git a/metadata/src/main/java/org/apache/kafka/metadata/placement/PodReplicaPlacer.java b/metadata/src/main/java/org/apache/kafka/metadata/placement/PodReplicaPlacer.java new file mode 100644 index 0000000000000..e027f1f811fe6 --- /dev/null +++ b/metadata/src/main/java/org/apache/kafka/metadata/placement/PodReplicaPlacer.java @@ -0,0 +1,237 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.metadata.placement; + +import org.apache.kafka.common.Uuid; +import org.apache.kafka.common.errors.InvalidReplicationFactorException; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.Iterator; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.function.Predicate; +import java.util.stream.Collectors; + +/** + * The class is a delegator of ReplicaPlacer to achieve pod isolation by placing partitions by broker pod + * {@code podPartitionMap} defined all pods that need isolate certain partition into + * for a partition that matches with any pod defined in {@code podPartitionMap}, the partition will be placed on that specific pod + * otherwise the partition will be placed on all the other pods not defined in podPartitionMap + * + * IMPORTANT: Each partition can match at most ONE isolation rule. If multiple pod predicates match the same + * partition, an IllegalStateException will be thrown during placement. + * + * special cases: + * - A broker with empty pod will be treated as any pod + * - A pod defined in podPartitionMap but not match any broker pod will be ignored + */ +public class PodReplicaPlacer implements ReplicaPlacer { + private final ReplicaPlacer impl; + private final Map> partitionPredicateByPod; + + public PodReplicaPlacer(ReplicaPlacer impl, Map> partitionPredicateByPod) { + this.impl = impl; + this.partitionPredicateByPod = partitionPredicateByPod; + } + + public PodReplicaPlacer(ReplicaPlacer impl) { + this(impl, Collections.emptyMap()); + } + + @Override + public TopicAssignment place(PlacementSpec placement, ClusterDescriber cluster) throws InvalidReplicationFactorException { + List clusterPlacements = calculateClusterPlacements(placement, cluster); + List partitionAssignments = new ArrayList<>(); + for (ClusterPlacement clusterPlacement : clusterPlacements) { + partitionAssignments.addAll(impl.place(clusterPlacement.placement, clusterPlacement.cluster).assignments()); + } + return new TopicAssignment(partitionAssignments); + } + + /** + * calculate placement for each partition by broker pod to achieve partition isolation by pod + * @param placement placement + * @param cluster cluster + * @return list of placement and cluster + */ + private List calculateClusterPlacements(PlacementSpec placement, ClusterDescriber cluster) { + List result = new ArrayList<>(); + Map, ClusterDescriber> podClusterCache = initCache(cluster); + Set allPods = podClusterCache.keySet() + .stream() + .flatMap(Collection::stream) + .collect(Collectors.toSet()); + int lastStartPartition = -1; + ClusterDescriber lastCluster = null; + for (int i = 0; i < placement.numPartitions(); ++i) { + int partition = placement.startPartition() + i; + Set pods = selectPodsByPartition(partition, allPods); + ClusterDescriber podCluster = clusterOfPods(pods, cluster, podClusterCache); + if (lastStartPartition == -1 || lastCluster != podCluster) { + if (lastStartPartition != -1) { + result.add(new ClusterPlacement(lastCluster, new PlacementSpec(lastStartPartition, partition - lastStartPartition, placement.numReplicas()))); + } + lastStartPartition = partition; + lastCluster = podCluster; + } + } + if (lastCluster != null) { + result.add(new ClusterPlacement(lastCluster, new PlacementSpec(lastStartPartition, placement.numPartitions() - lastStartPartition + placement.startPartition(), placement.numReplicas()))); + } + return result; + } + + /** + * evaluates pods with given partition + *

+ * - if the partition matches any isolation rule of a pod, return the pod + * - otherwise, return other pods + *

+ * @param partition the partition to map from + * @param pods pods of the whole cluster + * @return pods of the partition + * @throws IllegalStateException if multiple pod isolation rules match the same partition + */ + private Set selectPodsByPartition(int partition, Set pods) { + List matchingPods = new ArrayList<>(); + + // Check all pods for matches + for (String pod : pods) { + if (partitionPredicateByPod.getOrDefault(pod, x -> false).test(partition)) { + matchingPods.add(pod); + } + } + + // Validate: at most one match allowed + if (matchingPods.size() > 1) { + Collections.sort(matchingPods); // Sort for deterministic error messages + throw new IllegalStateException( + String.format( + "Partition %d matches multiple pod isolation rules: %s. " + + "Each partition must map to at most one isolated pod. " + + "Please review your canary specification.", + partition, matchingPods + ) + ); + } + + // Single match - isolate to this pod + if (matchingPods.size() == 1) { + return Collections.singleton(matchingPods.get(0)); + } + + // No matches - distribute to all non-isolated pods + return pods.stream() + .filter(pod -> !partitionPredicateByPod.containsKey(pod)) + .collect(Collectors.toSet()); + } + + /** + * Gets brokers for given pod set + * @param pods list of pod + * @param defaultCluster the full broker list + * @param cache the initial pods to brokers map + * @return brokers of given pods + */ + private ClusterDescriber clusterOfPods(Set pods, ClusterDescriber defaultCluster, Map, ClusterDescriber> cache) { + if (cache.containsKey(pods)) { + return cache.get(pods); + } + // if pods is empty but there is no broker with empty pod, fallback to full cluster + if (pods.isEmpty()) { + return defaultCluster; + } + + List clusters = pods.stream() + .map(Collections::singleton) + .map(cache::get) + .toList(); + + Set brokers = new LinkedHashSet<>(); + for (ClusterDescriber cluster : clusters) { + Iterator iter = cluster.usableBrokers(); + while (iter.hasNext()) { + brokers.add(iter.next()); + } + } + + ClusterDescriber result = new ClusterDescriber() { + @Override + public Iterator usableBrokers() { + return brokers.iterator(); + } + + @Override + public Uuid defaultDir(int brokerId) { + return defaultCluster.defaultDir(brokerId); + } + }; + + cache.put(pods, result); + return result; + } + + /** + * initialize mapping from pod set to brokers + * @param cluster indicates the full broker set + * @return pod set to brokers map + */ + private Map, ClusterDescriber> initCache(ClusterDescriber cluster) { + Map, Set> podBrokers = new HashMap<>(); + Iterator iter = cluster.usableBrokers(); + while (iter.hasNext()) { + UsableBroker broker = iter.next(); + Optional pod = broker.pod(); + podBrokers.computeIfAbsent(pod.map(Collections::singleton).orElse(Collections.emptySet()), k -> new LinkedHashSet<>()).add(broker); + } + + // broker without pod can be treated as any pod + Set podLessBrokers = podBrokers.get(Collections.emptySet()); + if (podLessBrokers != null) { + for (Map.Entry, Set> entry : podBrokers.entrySet()) { + if (!entry.getKey().isEmpty()) { + entry.getValue().addAll(podLessBrokers); + } + } + } + return podBrokers.entrySet().stream().collect(Collectors.toMap( + Map.Entry::getKey, // Keep the original key + entry -> new ClusterDescriber() { + @Override + public Iterator usableBrokers() { + return entry.getValue().iterator(); + } + + @Override + public Uuid defaultDir(int brokerId) { + return cluster.defaultDir(brokerId); + } + }) + ); + } + + private record ClusterPlacement(ClusterDescriber cluster, PlacementSpec placement) { + } + +} diff --git a/metadata/src/main/java/org/apache/kafka/metadata/placement/UsableBroker.java b/metadata/src/main/java/org/apache/kafka/metadata/placement/UsableBroker.java index 2c24d937f9e16..1d22dbdadc3d8 100644 --- a/metadata/src/main/java/org/apache/kafka/metadata/placement/UsableBroker.java +++ b/metadata/src/main/java/org/apache/kafka/metadata/placement/UsableBroker.java @@ -25,5 +25,5 @@ * A broker where a replica can be placed. */ @InterfaceStability.Unstable -public record UsableBroker(int id, Optional rack, boolean fenced) { +public record UsableBroker(int id, Optional rack, Optional pod, boolean fenced) { } diff --git a/metadata/src/main/resources/common/metadata/RegisterBrokerRecord.json b/metadata/src/main/resources/common/metadata/RegisterBrokerRecord.json index 12af50ce17d4c..51fb4668e0440 100644 --- a/metadata/src/main/resources/common/metadata/RegisterBrokerRecord.json +++ b/metadata/src/main/resources/common/metadata/RegisterBrokerRecord.json @@ -62,6 +62,8 @@ "about": "Log directories configured in this broker which are available." }, { "name": "CordonedLogDirs", "type": "[]uuid", "versions": "4+", "taggedVersions": "4+", "tag": 1, "nullableVersions": "4+", "default": "null", - "about": "Log directories that are cordoned. This is initially null before the broker starts sending the value via its heartbeats." } + "about": "Log directories that are cordoned. This is initially null before the broker starts sending the value via its heartbeats." }, + { "name": "Pod", "type": "string", "versions": "0+", "nullableVersions": "0+", "default": "null", + "taggedVersions": "0+", "tag": 2, "about": "The broker pod." } ] } diff --git a/metadata/src/test/java/org/apache/kafka/controller/BrokerHeartbeatManagerTest.java b/metadata/src/test/java/org/apache/kafka/controller/BrokerHeartbeatManagerTest.java index fe6e39600d9b7..ccad859ad142e 100644 --- a/metadata/src/test/java/org/apache/kafka/controller/BrokerHeartbeatManagerTest.java +++ b/metadata/src/test/java/org/apache/kafka/controller/BrokerHeartbeatManagerTest.java @@ -111,6 +111,7 @@ private static Set usableBrokersToSet(BrokerHeartbeatManager manag for (Iterator iterator = new UsableBrokerIterator( manager.brokers().iterator(), id -> id % 2 == 0 ? Optional.of("rack1") : Optional.of("rack2"), + id -> Optional.empty(), id -> id % 3 != 0); iterator.hasNext(); ) { brokers.add(iterator.next()); @@ -132,9 +133,9 @@ public void testUsableBrokerIterator() { manager.touch(4, true, 100); assertEquals(98L, manager.lowestActiveOffset()); Set expected = new HashSet<>(); - expected.add(new UsableBroker(1, Optional.of("rack2"), false)); - expected.add(new UsableBroker(2, Optional.of("rack1"), false)); - expected.add(new UsableBroker(4, Optional.of("rack1"), true)); + expected.add(new UsableBroker(1, Optional.of("rack2"), Optional.empty(), false)); + expected.add(new UsableBroker(2, Optional.of("rack1"), Optional.empty(), false)); + expected.add(new UsableBroker(4, Optional.of("rack1"), Optional.empty(), true)); assertEquals(expected, usableBrokersToSet(manager)); manager.maybeUpdateControlledShutdownOffset(2, 0); assertEquals(100L, manager.lowestActiveOffset()); @@ -142,8 +143,8 @@ public void testUsableBrokerIterator() { () -> manager.maybeUpdateControlledShutdownOffset(4, 0)); manager.touch(4, false, 100); manager.maybeUpdateControlledShutdownOffset(4, 0); - expected.remove(new UsableBroker(2, Optional.of("rack1"), false)); - expected.remove(new UsableBroker(4, Optional.of("rack1"), true)); + expected.remove(new UsableBroker(2, Optional.of("rack1"), Optional.empty(), false)); + expected.remove(new UsableBroker(4, Optional.of("rack1"), Optional.empty(), true)); assertEquals(expected, usableBrokersToSet(manager)); } diff --git a/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java b/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java index 71dc3c00cfe74..876e1659117ae 100644 --- a/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java +++ b/metadata/src/test/java/org/apache/kafka/controller/ReplicationControlManagerTest.java @@ -3021,11 +3021,11 @@ public void testKRaftClusterDescriber() { HashSet brokers = new HashSet<>(); describer.usableBrokers().forEachRemaining(brokers::add); assertEquals(Set.of( - new UsableBroker(0, Optional.empty(), true), - new UsableBroker(1, Optional.empty(), true), - new UsableBroker(2, Optional.empty(), false), - new UsableBroker(3, Optional.empty(), false), - new UsableBroker(4, Optional.empty(), false)), brokers); + new UsableBroker(0, Optional.empty(), Optional.empty(), true), + new UsableBroker(1, Optional.empty(), Optional.empty(), true), + new UsableBroker(2, Optional.empty(), Optional.empty(), false), + new UsableBroker(3, Optional.empty(), Optional.empty(), false), + new UsableBroker(4, Optional.empty(), Optional.empty(), false)), brokers); assertEquals(DirectoryId.MIGRATING, describer.defaultDir(1)); assertEquals(Uuid.fromString("ozwqsVMFSNiYQUPSJA3j0w"), describer.defaultDir(2)); assertEquals(DirectoryId.UNASSIGNED, describer.defaultDir(3)); diff --git a/metadata/src/test/java/org/apache/kafka/image/node/ClusterImageBrokersNodeTest.java b/metadata/src/test/java/org/apache/kafka/image/node/ClusterImageBrokersNodeTest.java index 1fed0e0dbad98..96bea7eb3c948 100644 --- a/metadata/src/test/java/org/apache/kafka/image/node/ClusterImageBrokersNodeTest.java +++ b/metadata/src/test/java/org/apache/kafka/image/node/ClusterImageBrokersNodeTest.java @@ -66,6 +66,7 @@ public void testNode1Child() { "listeners=[], " + "supportedFeatures={metadata.version: 1-4}, " + "rack=Optional.empty, " + + "pod=Optional.empty, " + "fenced=false, " + "inControlledShutdown=false, " + "isMigratingZkBroker=false, " + diff --git a/metadata/src/test/java/org/apache/kafka/metadata/BrokerRegistrationTest.java b/metadata/src/test/java/org/apache/kafka/metadata/BrokerRegistrationTest.java index c661658c9096a..f08fdacb6b315 100644 --- a/metadata/src/test/java/org/apache/kafka/metadata/BrokerRegistrationTest.java +++ b/metadata/src/test/java/org/apache/kafka/metadata/BrokerRegistrationTest.java @@ -53,6 +53,7 @@ public class BrokerRegistrationTest { setListeners(List.of(new Endpoint("INTERNAL", SecurityProtocol.PLAINTEXT, "localhost", 9090))). setSupportedFeatures(Map.of("foo", VersionRange.of((short) 1, (short) 2))). setRack(Optional.empty()). + setPod(Optional.empty()). setFenced(false). setInControlledShutdown(false).build(), new BrokerRegistration.Builder(). @@ -62,6 +63,7 @@ public class BrokerRegistrationTest { setListeners(List.of(new Endpoint("INTERNAL", SecurityProtocol.PLAINTEXT, "localhost", 9091))). setSupportedFeatures(Map.of("foo", VersionRange.of((short) 1, (short) 2))). setRack(Optional.empty()). + setPod(Optional.empty()). setFenced(true). setInControlledShutdown(false).build(), new BrokerRegistration.Builder(). @@ -73,6 +75,7 @@ public class BrokerRegistrationTest { new SimpleEntry<>("bar", VersionRange.of((short) 1, (short) 4))).collect( Collectors.toMap(SimpleEntry::getKey, SimpleEntry::getValue))). setRack(Optional.of("myrack")). + setPod(Optional.of("pod")). setFenced(false). setInControlledShutdown(true).build(), new BrokerRegistration.Builder(). @@ -83,6 +86,7 @@ public class BrokerRegistrationTest { setSupportedFeatures(Stream.of(new SimpleEntry<>("metadata.version", VersionRange.of((short) 7, (short) 7))) .collect(Collectors.toMap(SimpleEntry::getKey, SimpleEntry::getValue))). setRack(Optional.empty()). + setPod(Optional.empty()). setFenced(false). setInControlledShutdown(true). setIsMigratingZkBroker(true). @@ -134,28 +138,28 @@ public void testToString() { "incarnationId=3MfdxWlNSn2UDYsmDP1pYg, listeners=[Endpoint(" + "listenerName='INTERNAL', securityProtocol=PLAINTEXT, " + "host='localhost', port=9091)], supportedFeatures={foo: 1-2}, " + - "rack=Optional.empty, fenced=true, inControlledShutdown=false, isMigratingZkBroker=false, " + + "rack=Optional.empty, pod=Optional.empty, fenced=true, inControlledShutdown=false, isMigratingZkBroker=false, " + "directories=[], cordonedDirectories=null)", REGISTRATIONS.get(1).toString()); assertEquals("BrokerRegistration(id=2, epoch=0, " + "incarnationId=eY7oaG1RREie5Kk9uy1l6g, listeners=[Endpoint(" + "listenerName='INTERNAL', securityProtocol=PLAINTEXT, " + "host='localhost', port=9092)], supportedFeatures={bar: 1-4, foo: 2-3}, " + - "rack=Optional[myrack], fenced=false, inControlledShutdown=true, isMigratingZkBroker=false, " + + "rack=Optional[myrack], pod=Optional[pod], fenced=false, inControlledShutdown=true, isMigratingZkBroker=false, " + "directories=[], cordonedDirectories=null)", REGISTRATIONS.get(2).toString()); assertEquals("BrokerRegistration(id=3, epoch=0, " + "incarnationId=1t8VyWx2TCSTpUWuqj-FOw, listeners=[Endpoint(" + "listenerName='INTERNAL', securityProtocol=PLAINTEXT, " + "host='localhost', port=9093)], supportedFeatures={metadata.version: 7}, " + - "rack=Optional.empty, fenced=false, inControlledShutdown=true, isMigratingZkBroker=true, " + + "rack=Optional.empty, pod=Optional.empty, fenced=false, inControlledShutdown=true, isMigratingZkBroker=true, " + "directories=[r4HpEsMuST6nQ4rznIEJVA], cordonedDirectories=[r4HpEsMuST6nQ4rznIEJVA])", REGISTRATIONS.get(3).toString()); assertEquals("BrokerRegistration(id=4, epoch=0, " + "incarnationId=Xkq84F5bTsSEwHqceVxcOQ, listeners=[Endpoint(" + "listenerName='INTERNAL', securityProtocol=PLAINTEXT, " + "host='localhost', port=9094)], supportedFeatures={foo: 1-2}, " + - "rack=Optional.empty, fenced=true, inControlledShutdown=false, isMigratingZkBroker=false, " + + "rack=Optional.empty, pod=Optional.empty, fenced=true, inControlledShutdown=false, isMigratingZkBroker=false, " + "directories=[r4HpEsMuST6nQ4rznIEJVA], cordonedDirectories=null)", REGISTRATIONS.get(4).toString()); } @@ -186,9 +190,9 @@ public void testToNode() { assertEquals(Optional.empty(), REGISTRATIONS.get(0).node("NONEXISTENT")); assertEquals(Optional.of(new Node(0, "localhost", 9090, null)), REGISTRATIONS.get(0).node("INTERNAL")); - assertEquals(Optional.of(new Node(1, "localhost", 9091, null, true)), + assertEquals(Optional.of(new Node(1, "localhost", 9091, null, null, true)), REGISTRATIONS.get(1).node("INTERNAL")); - assertEquals(Optional.of(new Node(2, "localhost", 9092, "myrack")), + assertEquals(Optional.of(new Node(2, "localhost", 9092, "myrack", "pod")), REGISTRATIONS.get(2).node("INTERNAL")); assertEquals(Optional.of(new Node(3, "localhost", 9093, null)), REGISTRATIONS.get(3).node("INTERNAL")); @@ -203,6 +207,7 @@ public void testDirectoriesAreSorted() { setListeners(List.of(new Endpoint("INTERNAL", SecurityProtocol.PLAINTEXT, "localhost", 9090))). setSupportedFeatures(Map.of("foo", VersionRange.of((short) 1, (short) 2))). setRack(Optional.empty()). + setPod(Optional.empty()). setFenced(false). setInControlledShutdown(false). setDirectories(List.of( @@ -233,6 +238,7 @@ void testHasOnlineDir() { setListeners(List.of(new Endpoint("INTERNAL", SecurityProtocol.PLAINTEXT, "localhost", 9090))). setSupportedFeatures(Map.of("foo", VersionRange.of((short) 1, (short) 2))). setRack(Optional.empty()). + setPod(Optional.empty()). setFenced(false). setInControlledShutdown(false). setDirectories(List.of( diff --git a/metadata/src/test/java/org/apache/kafka/metadata/MetadataCacheTest.java b/metadata/src/test/java/org/apache/kafka/metadata/MetadataCacheTest.java index c193e48b58a1c..723dcb2bb62c4 100644 --- a/metadata/src/test/java/org/apache/kafka/metadata/MetadataCacheTest.java +++ b/metadata/src/test/java/org/apache/kafka/metadata/MetadataCacheTest.java @@ -132,7 +132,8 @@ public void getTopicMetadata() { records.add(new RegisterBrokerRecord() .setBrokerId(brokerId) .setEndPoints(endpoints) - .setRack("rack1")); + .setRack("rack1") + .setPod("pod")); } records.add(new TopicRecord().setName(topic0).setTopicId(topicIds.get(topic0))); records.add(new TopicRecord().setName(topic1).setTopicId(topicIds.get(topic1))); @@ -398,6 +399,7 @@ public void getTopicMetadataIsrNotAvailable() { new RegisterBrokerRecord() .setBrokerId(0) .setRack("rack1") + .setPod("pod") .setFenced(false) .setEndPoints(endpoints), new TopicRecord().setName(topic).setTopicId(topicId), @@ -456,6 +458,7 @@ public void getTopicMetadataWithNonSupportedSecurityProtocol() { RegisterBrokerRecord broker = new RegisterBrokerRecord() .setBrokerId(0) .setRack("") + .setPod("") .setEndPoints(new BrokerEndpointCollection(List.of( new BrokerEndpoint() .setHost("foo") @@ -516,6 +519,7 @@ private void updateCacheWithBrokers(MetadataCache cache, List brokerIds records.add(new RegisterBrokerRecord() .setBrokerId(brokerId) .setRack("") + .setPod("") .setFenced(false) .setBrokerEpoch(BROKER_EPOCH) .setEndPoints(new BrokerEndpointCollection(List.of( @@ -552,6 +556,7 @@ public void testGetPartitionReplicaEndpoints() { .setBrokerId(brokerId) .setFenced(brokerId == fencedBrokerId) .setRack("rack" + (brokerId % 3)) + .setPod("pod") .setEndPoints(new BrokerEndpointCollection(List.of( new BrokerEndpoint() .setHost("foo" + brokerId) @@ -958,6 +963,7 @@ public void testGetLeaderAndIsr() { .setBrokerId(0) .setBrokerEpoch(BROKER_EPOCH) .setRack("rack1") + .setPod("pod") .setEndPoints(new BrokerEndpointCollection(List.of( new BrokerEndpoint() .setHost("foo") @@ -1064,6 +1070,32 @@ public void testToClusterIncludesFencedReplicasAsOffline() { assertTrue(cluster.nodeById(broker1).isFenced()); } + @Test + public void testGetClusterMetadataWithPod() { + MetadataCache cache = createCache(); + SecurityProtocol securityProtocol = SecurityProtocol.PLAINTEXT; + ListenerName listenerName = ListenerName.forSecurityProtocol(securityProtocol); + List brokers = List.of( + new RegisterBrokerRecord() + .setBrokerId(0) + .setFenced(false) + .setBrokerEpoch(BROKER_EPOCH) + .setRack("rack1") + .setPod("pod") + .setEndPoints(new BrokerEndpointCollection(List.of( + new BrokerEndpoint() + .setHost("foo") + .setPort(9092) + .setSecurityProtocol(securityProtocol.id) + .setName(listenerName.value()) + ))) + ); + updateCache(cache, brokers); + Node aliveBroker = cache.getAliveBrokerNodes(listenerName).iterator().next(); + assertEquals("rack1", aliveBroker.rack()); + assertEquals("pod", aliveBroker.pod()); + } + private record Partition(int id, List replicas, List dirs) { } diff --git a/metadata/src/test/java/org/apache/kafka/metadata/placement/CanarySpecTest.java b/metadata/src/test/java/org/apache/kafka/metadata/placement/CanarySpecTest.java new file mode 100644 index 0000000000000..8b89ad6c1a973 --- /dev/null +++ b/metadata/src/test/java/org/apache/kafka/metadata/placement/CanarySpecTest.java @@ -0,0 +1,91 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.metadata.placement; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +import java.util.function.Predicate; + +public class CanarySpecTest { + @Test + public void testCanaryPodName() { + CanarySpec canarySpec = new CanarySpec("pod1", 50); + Assertions.assertEquals("pod1", canarySpec.toMap().keySet().iterator().next()); + } + + @Test + public void testEmptyCanaryPodName() { + CanarySpec canarySpec = new CanarySpec("", 50); + Assertions.assertTrue(canarySpec.toMap().keySet().iterator().next().isEmpty()); + } + + @Test + public void testZeroCanaryInterval() { + CanarySpec canarySpec = new CanarySpec("pod1", 0); + Assertions.assertEquals(CanarySpec.ALWAYS_FALSE, canarySpec.toMap().values().iterator().next()); + } + + @Test + public void testNegativeCanaryInterval() { + CanarySpec canarySpec = new CanarySpec("pod1", -1); + Assertions.assertEquals(CanarySpec.ALWAYS_FALSE, canarySpec.toMap().values().iterator().next()); + } + + @Test + public void testCanaryIntervalOne() { + CanarySpec canarySpec = new CanarySpec("pod1", 1); + Predicate predicate = canarySpec.toMap().values().iterator().next(); + for (int i = 0; i < 10; ++i) { + Assertions.assertTrue(predicate.test(i)); + } + } + + @Test + public void testCanaryIntervalFifty() { + CanarySpec canarySpec = new CanarySpec("pod1", 50); + Predicate predicate = canarySpec.toMap().values().iterator().next(); + for (int i = 0; i < 49; ++i) { + Assertions.assertFalse(predicate.test(i)); + } + Assertions.assertTrue(predicate.test(49)); + for (int i = 50; i < 99; ++i) { + Assertions.assertFalse(predicate.test(i)); + } + Assertions.assertTrue(predicate.test(99)); + } + + @Test + public void testCanaryIntervalCustom() { + CanarySpec canarySpec = new CanarySpec("pod1", 32); + Predicate predicate = canarySpec.toMap().values().iterator().next(); + Assertions.assertFalse(predicate.test(0)); + Assertions.assertTrue(predicate.test(31)); + Assertions.assertTrue(predicate.test(63)); + Assertions.assertTrue(predicate.test(95)); + } + + @Test + public void testTopicSmallerThanIntervalHasNoCanaryPartition() { + CanarySpec canarySpec = new CanarySpec("pod1", 50); + Predicate predicate = canarySpec.toMap().values().iterator().next(); + for (int i = 0; i < 12; ++i) { + Assertions.assertFalse(predicate.test(i)); + } + } +} diff --git a/metadata/src/test/java/org/apache/kafka/metadata/placement/PodReplicaPlacerTest.java b/metadata/src/test/java/org/apache/kafka/metadata/placement/PodReplicaPlacerTest.java new file mode 100644 index 0000000000000..1aead1b858db0 --- /dev/null +++ b/metadata/src/test/java/org/apache/kafka/metadata/placement/PodReplicaPlacerTest.java @@ -0,0 +1,354 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.kafka.metadata.placement; + +import org.apache.kafka.common.DirectoryId; +import org.apache.kafka.common.Uuid; +import org.apache.kafka.common.errors.InvalidReplicationFactorException; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; + +import java.util.Arrays; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collectors; + +public class PodReplicaPlacerTest { + + @Test + public void testDefaultPlaceWithEmptyCluster() { + PodReplicaPlacer placer = new PodReplicaPlacer(new MockReplicaPlacer()); + //place on all brokers with empty pods + TopicAssignment assignment = place(placer, 0, 4, (short) 1, Collections.EMPTY_LIST); + Assertions.assertEquals(0, assignment.assignments().size()); + } + + @Test + public void testDefaultPlaceWithEmptyPods() { + PodReplicaPlacer placer = new PodReplicaPlacer(new MockReplicaPlacer()); + //place on all brokers with empty pods + TopicAssignment assignment = place(placer, 0, 4, (short) 1, List.of( + new UsableBroker(1, Optional.empty(), Optional.empty(), true), + new UsableBroker(2, Optional.empty(), Optional.empty(), true))); + Assertions.assertEquals(4, assignment.assignments().size()); + Assertions.assertEquals(1, assignment.assignments().get(0).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(1).replicas().get(0)); + Assertions.assertEquals(1, assignment.assignments().get(2).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(3).replicas().get(0)); + } + + @Test + public void testAddPartitionDefaultPlaceWithEmptyPods() { + PodReplicaPlacer placer = new PodReplicaPlacer(new MockReplicaPlacer()); + //place on all brokers with empty pods + TopicAssignment assignment = place(placer, 4, 4, (short) 1, List.of( + new UsableBroker(1, Optional.empty(), Optional.empty(), true), + new UsableBroker(2, Optional.empty(), Optional.empty(), true))); + Assertions.assertEquals(4, assignment.assignments().size()); + Assertions.assertEquals(1, assignment.assignments().get(0).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(1).replicas().get(0)); + Assertions.assertEquals(1, assignment.assignments().get(2).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(3).replicas().get(0)); + } + + @Test + public void testDefaultPlaceWithTwoPods() { + PodReplicaPlacer placer = new PodReplicaPlacer(new MockReplicaPlacer()); + //there is no broker with empty pod, fallback to all brokers + TopicAssignment assignment = place(placer, 0, 4, (short) 1, List.of( + new UsableBroker(1, Optional.empty(), Optional.of("a"), true), + new UsableBroker(2, Optional.empty(), Optional.of("b"), true))); + Assertions.assertEquals(4, assignment.assignments().size()); + Assertions.assertEquals(1, assignment.assignments().get(0).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(1).replicas().get(0)); + Assertions.assertEquals(1, assignment.assignments().get(2).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(3).replicas().get(0)); + } + + @Test + public void testAddPartitionDefaultPlaceWithTwoPods() { + PodReplicaPlacer placer = new PodReplicaPlacer(new MockReplicaPlacer()); + //there is no broker with empty pod, fallback to all brokers + TopicAssignment assignment = place(placer, 4, 4, (short) 1, List.of( + new UsableBroker(1, Optional.empty(), Optional.of("a"), true), + new UsableBroker(2, Optional.empty(), Optional.of("b"), true))); + Assertions.assertEquals(4, assignment.assignments().size()); + Assertions.assertEquals(1, assignment.assignments().get(0).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(1).replicas().get(0)); + Assertions.assertEquals(1, assignment.assignments().get(2).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(3).replicas().get(0)); + } + + @Test + public void testDefaultPlaceWithPartialPods() { + PodReplicaPlacer placer = new PodReplicaPlacer(new MockReplicaPlacer()); + //there is one broker with empty pod, place all replica to the empty pod + TopicAssignment assignment = place(placer, 0, 4, (short) 1, List.of( + new UsableBroker(1, Optional.empty(), Optional.of("a"), true), + new UsableBroker(2, Optional.empty(), Optional.empty(), true))); + Assertions.assertEquals(4, assignment.assignments().size()); + Assertions.assertEquals(1, assignment.assignments().get(0).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(1).replicas().get(0)); + Assertions.assertEquals(1, assignment.assignments().get(2).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(3).replicas().get(0)); + } + + @Test + public void testPodIsolationPlaceWithEmptyPod() { + PodReplicaPlacer placer = new PodReplicaPlacer(new MockReplicaPlacer(), Collections.singletonMap("pod0", i -> i == 0)); + TopicAssignment assignment = place(placer, 0, 4, (short) 1, List.of( + new UsableBroker(1, Optional.empty(), Optional.empty(), true), + new UsableBroker(2, Optional.empty(), Optional.empty(), true))); + Assertions.assertEquals(4, assignment.assignments().size()); + Assertions.assertEquals(1, assignment.assignments().get(0).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(1).replicas().get(0)); + Assertions.assertEquals(1, assignment.assignments().get(2).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(3).replicas().get(0)); + } + + + @Test + public void testAddPartitionPodIsolationPlaceWithEmptyPod() { + PodReplicaPlacer placer = new PodReplicaPlacer(new MockReplicaPlacer(), Collections.singletonMap("pod0", i -> i == 4)); + TopicAssignment assignment = place(placer, 4, 4, (short) 1, List.of( + new UsableBroker(1, Optional.empty(), Optional.empty(), true), + new UsableBroker(2, Optional.empty(), Optional.empty(), true))); + Assertions.assertEquals(4, assignment.assignments().size()); + Assertions.assertEquals(1, assignment.assignments().get(0).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(1).replicas().get(0)); + Assertions.assertEquals(1, assignment.assignments().get(2).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(3).replicas().get(0)); + } + + @Test + public void testPodIsolationPlaceWithPartialPod() { + PodReplicaPlacer placer = new PodReplicaPlacer(new MockReplicaPlacer(), Collections.singletonMap("pod0", i -> i == 0)); + TopicAssignment assignment = place(placer, 0, 4, (short) 1, List.of( + new UsableBroker(1, Optional.empty(), Optional.empty(), true), + new UsableBroker(2, Optional.empty(), Optional.of("pod0"), true))); + Assertions.assertEquals(4, assignment.assignments().size()); + Assertions.assertEquals(2, assignment.assignments().get(0).replicas().get(0)); + Assertions.assertEquals(1, assignment.assignments().get(1).replicas().get(0)); + Assertions.assertEquals(1, assignment.assignments().get(2).replicas().get(0)); + Assertions.assertEquals(1, assignment.assignments().get(3).replicas().get(0)); + } + + @Test + public void testPodIsolationPlaceWithMultiplePods() { + PodReplicaPlacer placer = new PodReplicaPlacer(new MockReplicaPlacer(), Collections.singletonMap("pod0", i -> i == 0)); + TopicAssignment assignment = place(placer, 0, 8, (short) 1, List.of( + new UsableBroker(1, Optional.empty(), Optional.of("pod1"), true), + new UsableBroker(2, Optional.empty(), Optional.of("pod2"), true), + new UsableBroker(3, Optional.empty(), Optional.of("pod3"), true), + new UsableBroker(4, Optional.empty(), Optional.of("pod0"), true))); + Assertions.assertEquals(8, assignment.assignments().size()); + Assertions.assertEquals(4, assignment.assignments().get(0).replicas().get(0)); + Assertions.assertEquals(1, assignment.assignments().get(1).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(2).replicas().get(0)); + Assertions.assertEquals(3, assignment.assignments().get(3).replicas().get(0)); + Assertions.assertEquals(1, assignment.assignments().get(4).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(5).replicas().get(0)); + Assertions.assertEquals(3, assignment.assignments().get(6).replicas().get(0)); + Assertions.assertEquals(1, assignment.assignments().get(7).replicas().get(0)); + } + + @Test + public void testAddPartitionPodIsolationPlaceWithMultiplePods() { + PodReplicaPlacer placer = new PodReplicaPlacer(new MockReplicaPlacer(), Collections.singletonMap("pod0", i -> i == 8)); + TopicAssignment assignment = place(placer, 8, 8, (short) 1, List.of( + new UsableBroker(1, Optional.empty(), Optional.of("pod1"), true), + new UsableBroker(2, Optional.empty(), Optional.of("pod2"), true), + new UsableBroker(3, Optional.empty(), Optional.of("pod3"), true), + new UsableBroker(4, Optional.empty(), Optional.of("pod0"), true))); + Assertions.assertEquals(8, assignment.assignments().size()); + Assertions.assertEquals(4, assignment.assignments().get(0).replicas().get(0)); + Assertions.assertEquals(1, assignment.assignments().get(1).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(2).replicas().get(0)); + Assertions.assertEquals(3, assignment.assignments().get(3).replicas().get(0)); + Assertions.assertEquals(1, assignment.assignments().get(4).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(5).replicas().get(0)); + Assertions.assertEquals(3, assignment.assignments().get(6).replicas().get(0)); + Assertions.assertEquals(1, assignment.assignments().get(7).replicas().get(0)); + } + + @Test + public void testInvalidPodIsolationPlaceWithMultiplePods() { + PodReplicaPlacer placer = new PodReplicaPlacer(new MockReplicaPlacer(), Collections.singletonMap("foo", i -> i == 0)); + TopicAssignment assignment = place(placer, 0, 8, (short) 1, List.of( + new UsableBroker(1, Optional.empty(), Optional.of("pod1"), true), + new UsableBroker(2, Optional.empty(), Optional.of("pod2"), true), + new UsableBroker(3, Optional.empty(), Optional.of("pod3"), true), + new UsableBroker(4, Optional.empty(), Optional.of("pod0"), true))); + Assertions.assertEquals(8, assignment.assignments().size()); + Assertions.assertEquals(1, assignment.assignments().get(0).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(1).replicas().get(0)); + Assertions.assertEquals(3, assignment.assignments().get(2).replicas().get(0)); + Assertions.assertEquals(4, assignment.assignments().get(3).replicas().get(0)); + Assertions.assertEquals(1, assignment.assignments().get(4).replicas().get(0)); + Assertions.assertEquals(2, assignment.assignments().get(5).replicas().get(0)); + Assertions.assertEquals(3, assignment.assignments().get(6).replicas().get(0)); + Assertions.assertEquals(4, assignment.assignments().get(7).replicas().get(0)); + } + + @Test + public void testOverlappingPodIsolationRulesShouldFail() { + Map> rules = new HashMap<>(); + rules.put("pod0", i -> i == 0); + rules.put("pod1", i -> i < 5); // Also matches partition 0 + + PodReplicaPlacer placer = new PodReplicaPlacer(new MockReplicaPlacer(), rules); + + IllegalStateException exception = Assertions.assertThrows( + IllegalStateException.class, + () -> place(placer, 0, 1, (short) 1, Arrays.asList( + new UsableBroker(1, Optional.empty(), Optional.of("pod0"), true), + new UsableBroker(2, Optional.empty(), Optional.of("pod1"), true) + )) + ); + + Assertions.assertTrue(exception.getMessage().contains("Partition 0")); + Assertions.assertTrue(exception.getMessage().contains("pod0")); + Assertions.assertTrue(exception.getMessage().contains("pod1")); + Assertions.assertTrue(exception.getMessage().contains("multiple pod isolation rules")); + } + + @Test + public void testPodIsolationWithMockedReplicaPlacer() { + ReplicaPlacer mockReplicaPlacer = Mockito.mock(ReplicaPlacer.class); + ArgumentCaptor clusterCaptor = ArgumentCaptor.forClass(ClusterDescriber.class); + ArgumentCaptor placementSpecCaptor = ArgumentCaptor.forClass(PlacementSpec.class); + Mockito.when(mockReplicaPlacer.place(Mockito.any(PlacementSpec.class), Mockito.any(ClusterDescriber.class))).thenReturn(new TopicAssignment(Collections.EMPTY_LIST)); + PodReplicaPlacer placer = new PodReplicaPlacer(mockReplicaPlacer, Collections.singletonMap("pod0", i -> i == 0)); + place(placer, 0, 8, (short) 1, List.of( + new UsableBroker(1, Optional.of("rack-1"), Optional.of("pod0"), false), + new UsableBroker(2, Optional.of("rack-2"), Optional.of("pod0"), false), + new UsableBroker(3, Optional.of("rack-3"), Optional.of("pod0"), false), + new UsableBroker(4, Optional.of("rack-1"), Optional.of("pod1"), false), + new UsableBroker(5, Optional.of("rack-2"), Optional.of("pod2"), false), + new UsableBroker(6, Optional.of("rack-3"), Optional.empty(), false))); + Mockito.verify(mockReplicaPlacer, Mockito.times(2)).place(placementSpecCaptor.capture(), clusterCaptor.capture()); + List placementSpecs = placementSpecCaptor.getAllValues(); + List clusters = clusterCaptor.getAllValues(); + Assertions.assertEquals(2, placementSpecs.size()); + Assertions.assertEquals(2, clusters.size()); + Assertions.assertEquals(new PlacementSpec(0, 1, (short) 1), placementSpecs.get(0)); + Assertions.assertEquals(new PlacementSpec(1, 7, (short) 1), placementSpecs.get(1)); + Assertions.assertEquals(Set.of(1, 2, 3, 6), toList(clusters.get(0).usableBrokers()) + .stream() + .map(UsableBroker::id) + .collect(Collectors.toSet())); + Assertions.assertEquals(Set.of(4, 5, 6), toList(clusters.get(1).usableBrokers()) + .stream() + .map(UsableBroker::id) + .collect(Collectors.toSet())); + } + + @Test + public void testAddPartitionPodIsolationWithMockedReplicaPlacer() { + ReplicaPlacer mockReplicaPlacer = Mockito.mock(ReplicaPlacer.class); + ArgumentCaptor clusterCaptor = ArgumentCaptor.forClass(ClusterDescriber.class); + ArgumentCaptor placementSpecCaptor = ArgumentCaptor.forClass(PlacementSpec.class); + Mockito.when(mockReplicaPlacer.place(Mockito.any(PlacementSpec.class), Mockito.any(ClusterDescriber.class))).thenReturn(new TopicAssignment(Collections.EMPTY_LIST)); + PodReplicaPlacer placer = new PodReplicaPlacer(mockReplicaPlacer, Collections.singletonMap("pod0", i -> i == 8)); + place(placer, 8, 8, (short) 1, List.of( + new UsableBroker(1, Optional.of("rack-1"), Optional.of("pod0"), false), + new UsableBroker(2, Optional.of("rack-2"), Optional.of("pod0"), false), + new UsableBroker(3, Optional.of("rack-3"), Optional.of("pod0"), false), + new UsableBroker(4, Optional.of("rack-1"), Optional.of("pod1"), false), + new UsableBroker(5, Optional.of("rack-2"), Optional.of("pod2"), false), + new UsableBroker(6, Optional.of("rack-3"), Optional.empty(), false))); + Mockito.verify(mockReplicaPlacer, Mockito.times(2)).place(placementSpecCaptor.capture(), clusterCaptor.capture()); + List placementSpecs = placementSpecCaptor.getAllValues(); + List clusters = clusterCaptor.getAllValues(); + Assertions.assertEquals(2, placementSpecs.size()); + Assertions.assertEquals(2, clusters.size()); + Assertions.assertEquals(new PlacementSpec(8, 1, (short) 1), placementSpecs.get(0)); + Assertions.assertEquals(new PlacementSpec(9, 7, (short) 1), placementSpecs.get(1)); + Assertions.assertEquals(Set.of(1, 2, 3, 6), toList(clusters.get(0).usableBrokers()) + .stream() + .map(UsableBroker::id) + .collect(Collectors.toSet())); + Assertions.assertEquals(Set.of(4, 5, 6), toList(clusters.get(1).usableBrokers()) + .stream() + .map(UsableBroker::id) + .collect(Collectors.toSet())); + } + + private List toList(Iterator iterator) { + List result = new ArrayList<>(); + while (iterator.hasNext()) { + result.add(iterator.next()); + } + return result; + } + + /** + * Place partition round-robin + */ + private class MockReplicaPlacer implements ReplicaPlacer { + @Override + public TopicAssignment place(PlacementSpec placement, ClusterDescriber cluster) throws InvalidReplicationFactorException { + Iterator iterator = cluster.usableBrokers(); + List partitionAssignments = new ArrayList<>(); + if (!iterator.hasNext()) { + return new TopicAssignment(partitionAssignments); + } + for (int i = 0; i < placement.numPartitions(); i++) { + if (!iterator.hasNext()) { + iterator = cluster.usableBrokers(); + } + List replicas = new ArrayList<>(); + replicas.add(iterator.next().id()); + partitionAssignments.add(new PartitionAssignment(replicas, cluster)); + } + + return new TopicAssignment(partitionAssignments); + } + } + + private TopicAssignment place( + ReplicaPlacer placer, + int startPartition, + int numPartitions, + short replicationFactor, + List brokers + ) { + PlacementSpec placementSpec = new PlacementSpec(startPartition, + numPartitions, + replicationFactor); + return placer.place(placementSpec, new ClusterDescriber() { + @Override + public Iterator usableBrokers() { + return brokers.iterator(); + } + + @Override + public Uuid defaultDir(int brokerId) { + return DirectoryId.MIGRATING; + } + }); + } +} diff --git a/metadata/src/test/java/org/apache/kafka/metadata/placement/StripedReplicaPlacerTest.java b/metadata/src/test/java/org/apache/kafka/metadata/placement/StripedReplicaPlacerTest.java index 0db3a42c3bdc2..afb4b450a32cf 100644 --- a/metadata/src/test/java/org/apache/kafka/metadata/placement/StripedReplicaPlacerTest.java +++ b/metadata/src/test/java/org/apache/kafka/metadata/placement/StripedReplicaPlacerTest.java @@ -70,11 +70,11 @@ public void testBrokerList() { public void testAvoidFencedReplicaIfPossibleOnSingleRack() { MockRandom random = new MockRandom(); RackList rackList = new RackList(random, List.of( - new UsableBroker(3, Optional.empty(), false), - new UsableBroker(1, Optional.empty(), true), - new UsableBroker(0, Optional.empty(), false), - new UsableBroker(4, Optional.empty(), false), - new UsableBroker(2, Optional.empty(), false)).iterator()); + new UsableBroker(3, Optional.empty(), Optional.empty(), false), + new UsableBroker(1, Optional.empty(), Optional.empty(), true), + new UsableBroker(0, Optional.empty(), Optional.empty(), false), + new UsableBroker(4, Optional.empty(), Optional.empty(), false), + new UsableBroker(2, Optional.empty(), Optional.empty(), false)).iterator()); assertEquals(5, rackList.numTotalBrokers()); assertEquals(4, rackList.numUnfencedBrokers()); assertEquals(List.of(Optional.empty()), rackList.rackNames()); @@ -122,8 +122,8 @@ public void testMultiPartitionTopicPlacementOnSingleUnfencedBroker() { partitionAssignment(List.of(0)), partitionAssignment(List.of(0)))), place(placer, 0, 3, (short) 1, List.of( - new UsableBroker(0, Optional.empty(), false), - new UsableBroker(1, Optional.empty(), true)))); + new UsableBroker(0, Optional.empty(), Optional.empty(), false), + new UsableBroker(1, Optional.empty(), Optional.empty(), true)))); } /** @@ -133,9 +133,9 @@ public void testMultiPartitionTopicPlacementOnSingleUnfencedBroker() { public void testPlacementOnFencedReplicaOnSingleRack() { MockRandom random = new MockRandom(); RackList rackList = new RackList(random, List.of( - new UsableBroker(3, Optional.empty(), false), - new UsableBroker(1, Optional.empty(), true), - new UsableBroker(2, Optional.empty(), false)).iterator()); + new UsableBroker(3, Optional.empty(), Optional.empty(), false), + new UsableBroker(1, Optional.empty(), Optional.empty(), true), + new UsableBroker(2, Optional.empty(), Optional.empty(), false)).iterator()); assertEquals(3, rackList.numTotalBrokers()); assertEquals(2, rackList.numUnfencedBrokers()); assertEquals(List.of(Optional.empty()), rackList.rackNames()); @@ -149,12 +149,12 @@ public void testPlacementOnFencedReplicaOnSingleRack() { public void testRackListWithMultipleRacks() { MockRandom random = new MockRandom(); RackList rackList = new RackList(random, List.of( - new UsableBroker(11, Optional.of("1"), false), - new UsableBroker(10, Optional.of("1"), false), - new UsableBroker(30, Optional.of("3"), false), - new UsableBroker(31, Optional.of("3"), false), - new UsableBroker(21, Optional.of("2"), false), - new UsableBroker(20, Optional.of("2"), true)).iterator()); + new UsableBroker(11, Optional.of("1"), Optional.empty(), false), + new UsableBroker(10, Optional.of("1"), Optional.empty(), false), + new UsableBroker(30, Optional.of("3"), Optional.empty(), false), + new UsableBroker(31, Optional.of("3"), Optional.empty(), false), + new UsableBroker(21, Optional.of("2"), Optional.empty(), false), + new UsableBroker(20, Optional.of("2"), Optional.empty(), true)).iterator()); assertEquals(6, rackList.numTotalBrokers()); assertEquals(5, rackList.numUnfencedBrokers()); assertEquals(List.of(Optional.of("1"), Optional.of("2"), Optional.of("3")), rackList.rackNames()); @@ -167,14 +167,14 @@ public void testRackListWithMultipleRacks() { public void testRackListWithInvalidRacks() { MockRandom random = new MockRandom(); RackList rackList = new RackList(random, List.of( - new UsableBroker(11, Optional.of("1"), false), - new UsableBroker(10, Optional.of("1"), false), - new UsableBroker(30, Optional.of("3"), true), - new UsableBroker(31, Optional.of("3"), true), - new UsableBroker(20, Optional.of("2"), true), - new UsableBroker(21, Optional.of("2"), true), - new UsableBroker(41, Optional.of("4"), false), - new UsableBroker(40, Optional.of("4"), true)).iterator()); + new UsableBroker(11, Optional.of("1"), Optional.empty(), false), + new UsableBroker(10, Optional.of("1"), Optional.empty(), false), + new UsableBroker(30, Optional.of("3"), Optional.empty(), true), + new UsableBroker(31, Optional.of("3"), Optional.empty(), true), + new UsableBroker(20, Optional.of("2"), Optional.empty(), true), + new UsableBroker(21, Optional.of("2"), Optional.empty(), true), + new UsableBroker(41, Optional.of("4"), Optional.empty(), false), + new UsableBroker(40, Optional.of("4"), Optional.empty(), true)).iterator()); assertEquals(8, rackList.numTotalBrokers()); assertEquals(3, rackList.numUnfencedBrokers()); assertEquals(List.of(Optional.of("1"), @@ -193,8 +193,8 @@ public void testAllBrokersFenced() { assertEquals("All brokers are currently fenced, or have all their log directories cordoned.", assertThrows(InvalidReplicationFactorException.class, () -> place(placer, 0, 1, (short) 1, List.of( - new UsableBroker(11, Optional.of("1"), true), - new UsableBroker(10, Optional.of("1"), true)))).getMessage()); + new UsableBroker(11, Optional.of("1"), Optional.empty(), true), + new UsableBroker(10, Optional.of("1"), Optional.empty(), true)))).getMessage()); } @Test @@ -205,8 +205,8 @@ public void testNotEnoughBrokers() { "2 broker(s) are registered or some brokers have all their log directories cordoned.", assertThrows(InvalidReplicationFactorException.class, () -> place(placer, 0, 1, (short) 3, List.of( - new UsableBroker(11, Optional.of("1"), false), - new UsableBroker(10, Optional.of("1"), false)))).getMessage()); + new UsableBroker(11, Optional.of("1"), Optional.empty(), false), + new UsableBroker(10, Optional.of("1"), Optional.empty(), false)))).getMessage()); } @Test @@ -216,8 +216,8 @@ public void testNonPositiveReplicationFactor() { assertEquals("Invalid replication factor 0: the replication factor must be positive.", assertThrows(InvalidReplicationFactorException.class, () -> place(placer, 0, 1, (short) 0, List.of( - new UsableBroker(11, Optional.of("1"), false), - new UsableBroker(10, Optional.of("1"), false)))).getMessage()); + new UsableBroker(11, Optional.of("1"), Optional.empty(), false), + new UsableBroker(10, Optional.of("1"), Optional.empty(), false)))).getMessage()); } @Test @@ -230,10 +230,10 @@ public void testSuccessfulPlacement() { partitionAssignment(List.of(1, 2, 3)), partitionAssignment(List.of(1, 0, 2)))), place(placer, 0, 5, (short) 3, List.of( - new UsableBroker(0, Optional.empty(), false), - new UsableBroker(3, Optional.empty(), false), - new UsableBroker(2, Optional.empty(), false), - new UsableBroker(1, Optional.empty(), false)))); + new UsableBroker(0, Optional.empty(), Optional.empty(), false), + new UsableBroker(3, Optional.empty(), Optional.empty(), false), + new UsableBroker(2, Optional.empty(), Optional.empty(), false), + new UsableBroker(1, Optional.empty(), Optional.empty(), false)))); } @Test @@ -241,10 +241,10 @@ public void testEvenDistribution() { MockRandom random = new MockRandom(); StripedReplicaPlacer placer = new StripedReplicaPlacer(random); TopicAssignment topicAssignment = place(placer, 0, 200, (short) 2, List.of( - new UsableBroker(0, Optional.empty(), false), - new UsableBroker(1, Optional.empty(), false), - new UsableBroker(2, Optional.empty(), false), - new UsableBroker(3, Optional.empty(), false))); + new UsableBroker(0, Optional.empty(), Optional.empty(), false), + new UsableBroker(1, Optional.empty(), Optional.empty(), false), + new UsableBroker(2, Optional.empty(), Optional.empty(), false), + new UsableBroker(3, Optional.empty(), Optional.empty(), false))); Map, Integer> counts = new HashMap<>(); for (PartitionAssignment partitionAssignment : topicAssignment.assignments()) { counts.put(partitionAssignment.replicas(), counts.getOrDefault(partitionAssignment.replicas(), 0) + 1); @@ -268,9 +268,9 @@ public void testRackListAllBrokersFenced() { // ensure we can place N replicas on a rack when the rack has less than N brokers MockRandom random = new MockRandom(); RackList rackList = new RackList(random, List.of( - new UsableBroker(0, Optional.empty(), true), - new UsableBroker(1, Optional.empty(), true), - new UsableBroker(2, Optional.empty(), true)).iterator()); + new UsableBroker(0, Optional.empty(), Optional.empty(), true), + new UsableBroker(1, Optional.empty(), Optional.empty(), true), + new UsableBroker(2, Optional.empty(), Optional.empty(), true)).iterator()); assertEquals(3, rackList.numTotalBrokers()); assertEquals(0, rackList.numUnfencedBrokers()); assertEquals(List.of(Optional.empty()), rackList.rackNames()); @@ -283,8 +283,8 @@ public void testRackListAllBrokersFenced() { public void testRackListNotEnoughBrokers() { MockRandom random = new MockRandom(); RackList rackList = new RackList(random, List.of( - new UsableBroker(11, Optional.of("1"), false), - new UsableBroker(10, Optional.of("1"), false)).iterator()); + new UsableBroker(11, Optional.of("1"), Optional.empty(), false), + new UsableBroker(10, Optional.of("1"), Optional.empty(), false)).iterator()); assertEquals("The target replication factor of 3 cannot be reached because only " + "2 broker(s) are registered or some brokers have all their log directories cordoned.", assertThrows(InvalidReplicationFactorException.class, @@ -295,8 +295,8 @@ public void testRackListNotEnoughBrokers() { public void testRackListNonPositiveReplicationFactor() { MockRandom random = new MockRandom(); RackList rackList = new RackList(random, List.of( - new UsableBroker(11, Optional.of("1"), false), - new UsableBroker(10, Optional.of("1"), false)).iterator()); + new UsableBroker(11, Optional.of("1"), Optional.empty(), false), + new UsableBroker(10, Optional.of("1"), Optional.empty(), false)).iterator()); assertEquals("Invalid replication factor -1: the replication factor must be positive.", assertThrows(InvalidReplicationFactorException.class, () -> rackList.place(-1)).getMessage()); diff --git a/raft/src/main/java/org/apache/kafka/raft/KRaftConfigs.java b/raft/src/main/java/org/apache/kafka/raft/KRaftConfigs.java index 58c4965ff8f2b..4562b2451e828 100644 --- a/raft/src/main/java/org/apache/kafka/raft/KRaftConfigs.java +++ b/raft/src/main/java/org/apache/kafka/raft/KRaftConfigs.java @@ -69,6 +69,16 @@ public class KRaftConfigs { public static final long CONTROLLER_PERFORMANCE_ALWAYS_LOG_THRESHOLD_MS_DEFAULT = 2000; public static final String CONTROLLER_PERFORMANCE_ALWAYS_LOG_THRESHOLD_MS_DOC = "We will log an error message about controller events that take longer than this threshold."; + public static final String CANARY_POD_NAME = "canary.pod.name"; + public static final String CANARY_POD_NAME_DEFAULT = "canary-broker"; + public static final String CANARY_POD_NAME_DOC = "The name of canary pod that should place canary partition into"; + + public static final String CANARY_PARTITION_INTERVAL = "canary.partition.interval"; + public static final int CANARY_PARTITION_INTERVAL_DEFAULT = 0; + public static final String CANARY_PARTITION_INTERVAL_DOC = "Every Nth partition of a topic is a canary partition, placed on the canary pod. " + + "For example, 10 designates partitions 9, 19, 29 and so on, so a topic with fewer than 10 partitions has none. " + + "0 means disabled."; + public static final ConfigDef CONFIG_DEF = new ConfigDef() .define(PROCESS_ROLES_CONFIG, LIST, ConfigDef.NO_DEFAULT_VALUE, ConfigDef.ValidList.in(false, "broker", "controller"), HIGH, PROCESS_ROLES_DOC) .define(NODE_ID_CONFIG, INT, ConfigDef.NO_DEFAULT_VALUE, atLeast(0), HIGH, NODE_ID_DOC) @@ -77,6 +87,8 @@ public class KRaftConfigs { .define(BROKER_SESSION_TIMEOUT_MS_CONFIG, INT, BROKER_SESSION_TIMEOUT_MS_DEFAULT, null, MEDIUM, BROKER_SESSION_TIMEOUT_MS_DOC) .define(CONTROLLER_LISTENER_NAMES_CONFIG, LIST, ConfigDef.NO_DEFAULT_VALUE, ConfigDef.ValidList.anyNonDuplicateValues(false, false), HIGH, CONTROLLER_LISTENER_NAMES_DOC) .define(SASL_MECHANISM_CONTROLLER_PROTOCOL_CONFIG, STRING, SaslConfigs.DEFAULT_SASL_MECHANISM, null, HIGH, SASL_MECHANISM_CONTROLLER_PROTOCOL_DOC) + .define(CANARY_POD_NAME, STRING, CANARY_POD_NAME_DEFAULT, new ConfigDef.NonEmptyString(), MEDIUM, CANARY_POD_NAME_DOC) + .define(CANARY_PARTITION_INTERVAL, INT, CANARY_PARTITION_INTERVAL_DEFAULT, atLeast(0), MEDIUM, CANARY_PARTITION_INTERVAL_DOC) .defineInternal(CONTROLLER_PERFORMANCE_SAMPLE_PERIOD_MS, LONG, CONTROLLER_PERFORMANCE_SAMPLE_PERIOD_MS_DEFAULT, atLeast(100), MEDIUM, CONTROLLER_PERFORMANCE_SAMPLE_PERIOD_MS_DOC) .defineInternal(CONTROLLER_PERFORMANCE_ALWAYS_LOG_THRESHOLD_MS, LONG, CONTROLLER_PERFORMANCE_ALWAYS_LOG_THRESHOLD_MS_DEFAULT, atLeast(0), MEDIUM, CONTROLLER_PERFORMANCE_ALWAYS_LOG_THRESHOLD_MS_DOC) .defineInternal(SERVER_MAX_STARTUP_TIME_MS_CONFIG, LONG, SERVER_MAX_STARTUP_TIME_MS_DEFAULT, atLeast(0), MEDIUM, SERVER_MAX_STARTUP_TIME_MS_DOC); diff --git a/server-common/src/main/java/org/apache/kafka/server/config/ServerConfigs.java b/server-common/src/main/java/org/apache/kafka/server/config/ServerConfigs.java index 04a2353591fb7..c29197489bd7e 100644 --- a/server-common/src/main/java/org/apache/kafka/server/config/ServerConfigs.java +++ b/server-common/src/main/java/org/apache/kafka/server/config/ServerConfigs.java @@ -88,6 +88,11 @@ public class ServerConfigs { public static final String COMPRESSION_ZSTD_LEVEL_CONFIG = ServerTopicConfigSynonyms.serverSynonym(TopicConfig.COMPRESSION_ZSTD_LEVEL_CONFIG); public static final String COMPRESSION_ZSTD_LEVEL_DOC = "The compression level to use if " + COMPRESSION_TYPE_CONFIG + " is set to 'zstd'."; + /** ********* Pod Configuration ***********/ + public static final String BROKER_POD_CONFIG = "broker.pod"; + public static final String BROKER_POD_DEFAULT = "default"; + public static final String BROKER_POD_DOC = "The pod for the server, if unset, default pod id will be used"; + /***************** rack configuration *************/ public static final String BROKER_RACK_CONFIG = "broker.rack"; public static final String BROKER_RACK_DOC = "Rack of the broker. This will be used in rack aware replication assignment for fault tolerance. Examples: RACK1, us-east-1d"; @@ -135,6 +140,7 @@ public class ServerConfigs { .define(SOCKET_CONNECTION_SETUP_TIMEOUT_MS_CONFIG, LONG, DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MS, MEDIUM, SOCKET_CONNECTION_SETUP_TIMEOUT_MS_DOC) .define(SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_CONFIG, LONG, SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS, MEDIUM, SOCKET_CONNECTION_SETUP_TIMEOUT_MAX_MS_DOC) .define(CONFIG_PROVIDERS_CONFIG, ConfigDef.Type.LIST, List.of(), ConfigDef.ValidList.anyNonDuplicateValues(true, false), ConfigDef.Importance.LOW, CONFIG_PROVIDERS_DOC) + .define(BROKER_POD_CONFIG, STRING, null, MEDIUM, BROKER_POD_DOC) /************* Authorizer Configuration ***********/ .define(AUTHORIZER_CLASS_NAME_CONFIG, STRING, AUTHORIZER_CLASS_NAME_DEFAULT, new ConfigDef.NonNullValidator(), LOW, AUTHORIZER_CLASS_NAME_DOC) .define(EARLY_START_LISTENERS_CONFIG, LIST, null, ConfigDef.ValidList.anyNonDuplicateValues(true, true), HIGH, EARLY_START_LISTENERS_DOC) diff --git a/server/src/main/java/org/apache/kafka/server/BrokerLifecycleManager.java b/server/src/main/java/org/apache/kafka/server/BrokerLifecycleManager.java index 3c8bc8807cd12..ecad46fa337cf 100644 --- a/server/src/main/java/org/apache/kafka/server/BrokerLifecycleManager.java +++ b/server/src/main/java/org/apache/kafka/server/BrokerLifecycleManager.java @@ -92,6 +92,11 @@ public class BrokerLifecycleManager { */ private final Optional rack; + /** + * The broker pod, or null if there is no configured pod. + */ + private final Optional pod; + /** * How long to wait for registration to succeed before failing the startup process. */ @@ -239,6 +244,7 @@ public BrokerLifecycleManager( this.logger = logContext.logger(BrokerLifecycleManager.class); this.nodeId = config.nodeId(); this.rack = config.rack(); + this.pod = config.pod(); this.initialTimeoutNs = MILLISECONDS.toNanos(config.initialRegistrationTimeoutMs()); this.eventQueue = new KafkaEventQueue( time, @@ -490,6 +496,7 @@ private void sendBrokerRegistration() { .setIncarnationId(incarnationId) .setListeners(advertisedListeners) .setRack(rack.orElse(null)) + .setPod(pod.orElse(null)) .setPreviousBrokerEpoch(previousBrokerEpoch.orElse(-1L)) .setLogDirs(sortedLogDirs); if (logger.isDebugEnabled()) { diff --git a/server/src/main/java/org/apache/kafka/server/config/AbstractKafkaConfig.java b/server/src/main/java/org/apache/kafka/server/config/AbstractKafkaConfig.java index 6220651ed03d6..8195bbf0c48af 100644 --- a/server/src/main/java/org/apache/kafka/server/config/AbstractKafkaConfig.java +++ b/server/src/main/java/org/apache/kafka/server/config/AbstractKafkaConfig.java @@ -170,6 +170,10 @@ public Optional rack() { return Optional.ofNullable(getString(ServerConfigs.BROKER_RACK_CONFIG)); } + public Optional pod() { + return Optional.ofNullable(getString(ServerConfigs.BROKER_POD_CONFIG)); + } + public int nodeId() { return getInt(KRaftConfigs.NODE_ID_CONFIG); } diff --git a/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java b/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java index 89595b63750ef..3eeae6bff8e67 100644 --- a/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java +++ b/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommand.java @@ -37,8 +37,10 @@ import org.apache.kafka.common.utils.Time; import org.apache.kafka.common.utils.Utils; import org.apache.kafka.common.utils.internals.Exit; +import org.apache.kafka.metadata.placement.CanarySpec; import org.apache.kafka.metadata.placement.ClusterDescriber; import org.apache.kafka.metadata.placement.PlacementSpec; +import org.apache.kafka.metadata.placement.PodReplicaPlacer; import org.apache.kafka.metadata.placement.ReplicaPlacer; import org.apache.kafka.metadata.placement.StripedReplicaPlacer; import org.apache.kafka.metadata.placement.TopicAssignment; @@ -172,7 +174,9 @@ private static void handleAction(Admin adminClient, ReassignPartitionsCommandOpt generateAssignment(adminClient, Utils.readFileAsString(opts.options.valueOf(opts.topicsToMoveJsonFileOpt)), opts.options.valueOf(opts.brokerListOpt), - !opts.options.has(opts.disableRackAware)); + !opts.options.has(opts.disableRackAware), + opts.options.valueOf(opts.canaryNameOpt), + opts.options.valueOf(opts.canaryIntervalOpt)); } else if (opts.options.has(opts.executeOpt)) { executeAssignment(adminClient, opts.options.has(opts.additionalOpt), @@ -558,7 +562,9 @@ private static void clearTopicLevelThrottles(Admin adminClient, Set topi public static Entry>, Map>> generateAssignment(Admin adminClient, String reassignmentJson, String brokerListString, - boolean enableRackAwareness + boolean enableRackAwareness, + String canaryPod, + int canaryInterval ) throws ExecutionException, InterruptedException, JsonProcessingException { Entry, List> t0 = parseGenerateAssignmentArgs(reassignmentJson, brokerListString); @@ -569,7 +575,7 @@ public static Entry>, Map currentReplicaLogDirs = getReplicaToLogDir(adminClient, currentAssignments); List usableBrokers = getBrokerMetadata(adminClient, brokersToReassign, enableRackAwareness); Map> currentParts = toReplicaIds(currentAssignments); - Map> proposedAssignments = calculateAssignment(currentParts, usableBrokers); + Map> proposedAssignments = calculateAssignment(currentParts, usableBrokers, canaryPod, canaryInterval); System.out.printf("Current partition replica assignment%n%s%n%n", formatAsReassignmentJson(currentParts, currentReplicaLogDirs)); System.out.printf("Proposed partition reassignment configuration%n%s%n", @@ -586,7 +592,9 @@ public static Entry>, Map> calculateAssignment(Map> currentAssignment, - List usableBrokers) { + List usableBrokers, + String canaryPod, + int canaryInterval) { Map>>> groupedByTopic = new HashMap<>(); for (Entry> e : currentAssignment.entrySet()) groupedByTopic.computeIfAbsent(e.getKey().topic(), k -> new ArrayList<>()).add(e); @@ -595,7 +603,7 @@ private static Map> calculateAssignment(Map replicas = assignment.get(0).getValue(); int partitionNum = assignment.size(); // generate topic assignments - TopicAssignment topicAssignment = REPLICA_PLACER.place( + TopicAssignment topicAssignment = new PodReplicaPlacer(REPLICA_PLACER, new CanarySpec(canaryPod, canaryInterval).toMap()).place( new PlacementSpec(0, partitionNum, (short) replicas.size()), new ClusterDescriber() { @Override @@ -712,8 +720,8 @@ static List getBrokerMetadata(Admin adminClient, List bro List results = adminClient.describeCluster().nodes().get().stream() .filter(node -> brokerSet.contains(node.id())) .map(node -> (enableRackAwareness && node.rack() != null) - ? new UsableBroker(node.id(), Optional.of(node.rack()), false) - : new UsableBroker(node.id(), Optional.empty(), false) + ? new UsableBroker(node.id(), Optional.of(node.rack()), Optional.ofNullable(node.pod()), false) + : new UsableBroker(node.id(), Optional.empty(), Optional.ofNullable(node.pod()), false) ).collect(Collectors.toList()); long numRackless = results.stream().filter(m -> m.rack().isEmpty()).count(); diff --git a/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandOptions.java b/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandOptions.java index 3954128871222..455891221e694 100644 --- a/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandOptions.java +++ b/tools/src/main/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandOptions.java @@ -43,6 +43,8 @@ public class ReassignPartitionsCommandOptions extends CommandDefaultOptions { final OptionSpec additionalOpt; final OptionSpec preserveThrottlesOpt; final OptionSpec disallowReplicationFactorChangeOpt; + final OptionSpec canaryNameOpt; + final OptionSpec canaryIntervalOpt; public ReassignPartitionsCommandOptions(String[] args) { super(args); @@ -117,7 +119,16 @@ public ReassignPartitionsCommandOptions(String[] args) { "other ongoing ones. This option can also be used to change the throttle of an ongoing reassignment."); preserveThrottlesOpt = parser.accepts("preserve-throttles", "Do not modify broker or topic throttles."); disallowReplicationFactorChangeOpt = parser.accepts("disallow-replication-factor-change", "Denies the ability to change a partition's replication factor as part of this reassignment through adding validation against it."); - + canaryNameOpt = parser.accepts("canary-name", "Specify the name of canary pod, canary partition will be placed on canary pod according to name defined") + .withRequiredArg() + .describedAs("canary pod name") + .ofType(String.class) + .defaultsTo("canary-broker"); + canaryIntervalOpt = parser.accepts("canary-interval", "Specify that every Nth partition is a canary partition; 0 disables canary placement") + .withRequiredArg() + .describedAs("canary partition interval") + .ofType(Integer.class) + .defaultsTo(0); options = parser.parse(args); } } diff --git a/tools/src/test/java/org/apache/kafka/tools/BrokerApiVersionsCommandTest.java b/tools/src/test/java/org/apache/kafka/tools/BrokerApiVersionsCommandTest.java index cfc9b686644a7..fb63bdc08d479 100644 --- a/tools/src/test/java/org/apache/kafka/tools/BrokerApiVersionsCommandTest.java +++ b/tools/src/test/java/org/apache/kafka/tools/BrokerApiVersionsCommandTest.java @@ -50,7 +50,7 @@ public void testBrokerApiVersionsCommandOutput(ClusterInstance clusterInstance) BrokerApiVersionsCommand.mainNoExit("--bootstrap-server", clusterInstance.bootstrapServers())); Iterator lineIter = Arrays.stream(output.split("\n")).iterator(); assertTrue(lineIter.hasNext()); - assertEquals(clusterInstance.bootstrapServers() + " (id: 0 rack: null isFenced: false) -> (", lineIter.next()); + assertEquals(clusterInstance.bootstrapServers() + " (id: 0 rack: null pod: null isFenced: false) -> (", lineIter.next()); ApiMessageType.ListenerType listenerType = ApiMessageType.ListenerType.BROKER; diff --git a/tools/src/test/java/org/apache/kafka/tools/ClusterToolTest.java b/tools/src/test/java/org/apache/kafka/tools/ClusterToolTest.java index bbd5f2d1c2300..d37bc517f40e0 100644 --- a/tools/src/test/java/org/apache/kafka/tools/ClusterToolTest.java +++ b/tools/src/test/java/org/apache/kafka/tools/ClusterToolTest.java @@ -227,9 +227,9 @@ private void testApiVersionsCommandOutput(ClusterInstance clusterInstance, boole if (usingBootstrapController) { int id = clusterInstance.type() == Type.CO_KRAFT ? 0 : 3000; - assertEquals(clusterInstance.bootstrapControllers() + " (id: " + id + " rack: null isFenced: false) -> (", lineIter.next()); + assertEquals(clusterInstance.bootstrapControllers() + " (id: " + id + " rack: null pod: null isFenced: false) -> (", lineIter.next()); } else { - assertEquals(clusterInstance.bootstrapServers() + " (id: 0 rack: null isFenced: false) -> (", lineIter.next()); + assertEquals(clusterInstance.bootstrapServers() + " (id: 0 rack: null pod: null isFenced: false) -> (", lineIter.next()); } EnumSet apiKeys = EnumSet.copyOf(ApiKeys.clientApis()); diff --git a/tools/src/test/java/org/apache/kafka/tools/other/ReplicationQuotasTestRig.java b/tools/src/test/java/org/apache/kafka/tools/other/ReplicationQuotasTestRig.java index 5df994ec1aac8..f4a40f260691a 100644 --- a/tools/src/test/java/org/apache/kafka/tools/other/ReplicationQuotasTestRig.java +++ b/tools/src/test/java/org/apache/kafka/tools/other/ReplicationQuotasTestRig.java @@ -240,7 +240,7 @@ public void run(ExperimentDef config, Journal journal, boolean displayChartsOnSc cluster.brokers().values().stream() .map(server -> String.valueOf(server.replicaManager().localBrokerId())) .collect(Collectors.joining(",")), - true + true, "pod1", 0 ).getKey(); System.out.println("Starting Reassignment"); diff --git a/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java b/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java index a71e4779a3e4e..86a6908e151e1 100644 --- a/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java +++ b/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsCommandTest.java @@ -161,7 +161,7 @@ public void testGenerateAssignmentWithBootstrapServer() throws Exception { "version": 1 } """; - var assignment = generateAssignment(admin, topicsToMoveJson, "1,2,3", false); + var assignment = generateAssignment(admin, topicsToMoveJson, "1,2,3", false, "pod1", 0); Map> proposedAssignments = assignment.getKey(); String assignmentJson = String.format(""" { @@ -523,7 +523,7 @@ public void testGenerateAssignmentWithOneBootstrapServerShutdownWontTimeout() th () -> clusterInstance.aliveBrokers().size() == 4, "Waiting for broker to shutdown failed" ); - generateAssignment(admin, topicsToMoveJson, "1,2,3", false); + generateAssignment(admin, topicsToMoveJson, "1,2,3", false, "pod1", 0); } } diff --git a/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsUnitTest.java b/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsUnitTest.java index 6a59a1513fc39..58da062f7bee3 100644 --- a/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsUnitTest.java +++ b/tools/src/test/java/org/apache/kafka/tools/reassign/ReassignPartitionsUnitTest.java @@ -322,25 +322,25 @@ public void testGetReplicaAssignments() throws Exception { @Test public void testGetBrokerRackInformation() throws Exception { try (MockAdminClient adminClient = new MockAdminClient.Builder(). - brokers(List.of(new Node(0, "localhost", 9092, "rack0"), - new Node(1, "localhost", 9093, "rack1"), - new Node(2, "localhost", 9094, null))). + brokers(List.of(new Node(0, "localhost", 9092, "rack0", "pod0"), + new Node(1, "localhost", 9093, "rack1", "pod1"), + new Node(2, "localhost", 9094, null, null))). build()) { assertEquals(List.of( - new UsableBroker(0, Optional.of("rack0"), false), - new UsableBroker(1, Optional.of("rack1"), false) + new UsableBroker(0, Optional.of("rack0"), Optional.of("pod0"), false), + new UsableBroker(1, Optional.of("rack1"), Optional.of("pod1"), false) ), getBrokerMetadata(adminClient, List.of(0, 1), true)); assertEquals(List.of( - new UsableBroker(0, Optional.empty(), false), - new UsableBroker(1, Optional.empty(), false) + new UsableBroker(0, Optional.empty(), Optional.of("pod0"), false), + new UsableBroker(1, Optional.empty(), Optional.of("pod1"), false) ), getBrokerMetadata(adminClient, List.of(0, 1), false)); assertStartsWith("Not all brokers have rack information", assertThrows(AdminOperationException.class, () -> getBrokerMetadata(adminClient, List.of(1, 2), true)).getMessage()); assertEquals(List.of( - new UsableBroker(1, Optional.empty(), false), - new UsableBroker(2, Optional.empty(), false) + new UsableBroker(1, Optional.empty(), Optional.of("pod1"), false), + new UsableBroker(2, Optional.empty(), Optional.empty(), false) ), getBrokerMetadata(adminClient, List.of(1, 2), false)); } } @@ -372,7 +372,7 @@ public void testGenerateAssignmentFailsWithoutEnoughReplicas() { addTopics(adminClient); assertStartsWith("The target replication factor of 3 cannot be reached because only 2 broker(s) are registered", assertThrows(InvalidReplicationFactorException.class, - () -> generateAssignment(adminClient, "{\"topics\":[{\"topic\":\"foo\"},{\"topic\":\"bar\"}]}", "0,1", false), + () -> generateAssignment(adminClient, "{\"topics\":[{\"topic\":\"foo\"},{\"topic\":\"bar\"}]}", "0,1", false, "pod1", 0), "Expected generateAssignment to fail").getMessage()); } } @@ -383,7 +383,7 @@ public void testGenerateAssignmentWithInvalidPartitionsFails() { addTopics(adminClient); assertStartsWith("Topic quux not found", assertThrows(ExecutionException.class, - () -> generateAssignment(adminClient, "{\"topics\":[{\"topic\":\"foo\"},{\"topic\":\"quux\"}]}", "0,1", false), + () -> generateAssignment(adminClient, "{\"topics\":[{\"topic\":\"foo\"},{\"topic\":\"quux\"}]}", "0,1", false, "pod1", 0), "Expected generateAssignment to fail").getCause().getMessage()); } } @@ -393,21 +393,21 @@ public void testGenerateAssignmentWithInconsistentRacks() throws Exception { try (MockAdminClient adminClient = new MockAdminClient.Builder(). brokers(List.of( new Node(0, "localhost", 9092, "rack0"), - new Node(1, "localhost", 9093, "rack0"), - new Node(2, "localhost", 9094, null), - new Node(3, "localhost", 9095, "rack1"), - new Node(4, "localhost", 9096, "rack1"), - new Node(5, "localhost", 9097, "rack2"))). + new Node(1, "localhost", 9093, "rack0", "pod0"), + new Node(2, "localhost", 9094, null, null), + new Node(3, "localhost", 9095, "rack1", "pod1"), + new Node(4, "localhost", 9096, "rack1", "pod1"), + new Node(5, "localhost", 9097, "rack2", "pod2"))). build()) { addTopics(adminClient); assertStartsWith("Not all brokers have rack information.", assertThrows(AdminOperationException.class, - () -> generateAssignment(adminClient, "{\"topics\":[{\"topic\":\"foo\"}]}", "0,1,2,3", true), + () -> generateAssignment(adminClient, "{\"topics\":[{\"topic\":\"foo\"}]}", "0,1,2,3", true, "pod1", 0), "Expected generateAssignment to fail").getMessage()); // It should succeed when --disable-rack-aware is used. Entry>, Map>> - proposedCurrent = generateAssignment(adminClient, "{\"topics\":[{\"topic\":\"foo\"}]}", "0,1,2,3", false); + proposedCurrent = generateAssignment(adminClient, "{\"topics\":[{\"topic\":\"foo\"}]}", "0,1,2,3", false, "pod1", 0); Map> expCurrent = new HashMap<>(); @@ -418,6 +418,36 @@ public void testGenerateAssignmentWithInconsistentRacks() throws Exception { } } + @Test + public void testGenerateAssignmentWithCanaryIsolation() throws Exception { + try (MockAdminClient adminClient = new MockAdminClient.Builder(). + brokers(List.of( + new Node(0, "localhost", 9092, "rack0", "pod0"), + new Node(1, "localhost", 9093, "rack0", "pod0"), + new Node(2, "localhost", 9094, null, null), + new Node(3, "localhost", 9095, "rack1", "pod1"), + new Node(4, "localhost", 9096, "rack1", "pod1"), + new Node(5, "localhost", 9097, "rack2", "pod2"))). + build()) { + addTopics(adminClient); + + Entry>, Map>> + assignment = generateAssignment(adminClient, "{\"topics\":[{\"topic\":\"foo\"}]}", "0,1,2,3,4", false, "pod1", 2); + + Map> expCurrent = new HashMap<>(); + expCurrent.put(new TopicPartition("foo", 0), List.of(0, 1, 2)); + expCurrent.put(new TopicPartition("foo", 1), List.of(1, 2, 3)); + assertEquals(expCurrent, assignment.getValue()); + + Map> expProposal = new HashMap<>(); + expProposal.put(new TopicPartition("foo", 0), Set.of(0, 1, 2)); + expProposal.put(new TopicPartition("foo", 1), Set.of(2, 3, 4)); + + assertEquals(expProposal, assignment.getKey().entrySet().stream() + .collect(Collectors.toMap(Map.Entry::getKey, entry -> Set.copyOf(entry.getValue())))); + } + } + @Test public void testGenerateAssignmentWithFewerBrokers() throws Exception { try (MockAdminClient adminClient = new MockAdminClient.Builder().numBrokers(4).build()) { @@ -427,7 +457,7 @@ public void testGenerateAssignmentWithFewerBrokers() throws Exception { Entry>, Map>> proposedCurrent = generateAssignment(adminClient, "{\"topics\":[{\"topic\":\"foo\"},{\"topic\":\"bar\"}]}", - goalBrokers.stream().map(Object::toString).collect(Collectors.joining(",")), false); + goalBrokers.stream().map(Object::toString).collect(Collectors.joining(",")), false, "pod1", 0); Map> expCurrent = new HashMap<>();