Skip to content

MINOR: Support co-resident mode in KRaft TestKit #12157

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
dengziming wants to merge 1 commit into from
Closed

MINOR: Support co-resident mode in KRaft TestKit #12157

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions core/src/test/java/kafka/test/ClusterTestExtensionsTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,11 @@ public void testClusterTemplate() {
@ClusterConfigProperty(key = "foo", value = "baz"),
@ClusterConfigProperty(key = "spam", value = "eggz")
})
,
@ClusterTest(name = "cluster-tests-3", clusterType = Type.CO_KRAFT, serverProperties = {
@ClusterConfigProperty(key = "foo", value = "baz"),
@ClusterConfigProperty(key = "spam", value = "eggz")
})
})
public void testClusterTests() {
if (clusterInstance.clusterType().equals(ClusterInstance.ClusterType.ZK)) {
Expand Down
18 changes: 16 additions & 2 deletions core/src/test/java/kafka/test/annotation/Type.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,13 @@ public enum Type {
KRAFT {
@Override
public void invocationContexts(ClusterConfig config, Consumer<TestTemplateInvocationContext> invocationConsumer) {
invocationConsumer.accept(new RaftClusterInvocationContext(config.copyOf()));
invocationConsumer.accept(new RaftClusterInvocationContext(config.copyOf(), false));
}
},
CO_KRAFT {
@Override
public void invocationContexts(ClusterConfig config, Consumer<TestTemplateInvocationContext> invocationConsumer) {
invocationConsumer.accept(new RaftClusterInvocationContext(config.copyOf(), true));
}
},
ZK {
Expand All @@ -43,7 +49,15 @@ public void invocationContexts(ClusterConfig config, Consumer<TestTemplateInvoca
BOTH {
@Override
public void invocationContexts(ClusterConfig config, Consumer<TestTemplateInvocationContext> invocationConsumer) {
invocationConsumer.accept(new RaftClusterInvocationContext(config.copyOf()));
invocationConsumer.accept(new RaftClusterInvocationContext(config.copyOf(), false));
invocationConsumer.accept(new ZkClusterInvocationContext(config.copyOf()));
}
},
ALL {
@Override
public void invocationContexts(ClusterConfig config, Consumer<TestTemplateInvocationContext> invocationConsumer) {
invocationConsumer.accept(new RaftClusterInvocationContext(config.copyOf(), false));
invocationConsumer.accept(new RaftClusterInvocationContext(config.copyOf(), true));
invocationConsumer.accept(new ZkClusterInvocationContext(config.copyOf()));
}
},
Expand Down
14 changes: 9 additions & 5 deletions core/src/test/java/kafka/test/junit/RaftClusterInvocationContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,18 +66,20 @@ public class RaftClusterInvocationContext implements TestTemplateInvocationConte

private final ClusterConfig clusterConfig;
private final AtomicReference<KafkaClusterTestKit> clusterReference;
private final boolean isCoResident;

public RaftClusterInvocationContext(ClusterConfig clusterConfig) {
public RaftClusterInvocationContext(ClusterConfig clusterConfig, boolean isCoResident) {
this.clusterConfig = clusterConfig;
this.clusterReference = new AtomicReference<>();
this.isCoResident = isCoResident;
}

@Override
public String getDisplayName(int invocationIndex) {
String clusterDesc = clusterConfig.nameTags().entrySet().stream()
.map(Object::toString)
.collect(Collectors.joining(", "));
return String.format("[%d] Type=Raft, %s", invocationIndex, clusterDesc);
.map(Object::toString)
.collect(Collectors.joining(", "));
return String.format("[%d] Type=Raft-%s, %s", invocationIndex, isCoResident ? "CoReside" : "Distributed", clusterDesc);
}

@Override
Expand All @@ -87,8 +89,10 @@ public List<Extension> getAdditionalExtensions() {
(BeforeTestExecutionCallback) context -> {
TestKitNodes nodes = new TestKitNodes.Builder().
setBootstrapMetadataVersion(clusterConfig.metadataVersion().orElse(MetadataVersion.latest())).
setCoResident(isCoResident).
setNumBrokerNodes(clusterConfig.numBrokers()).
setNumControllerNodes(clusterConfig.numControllers()).build();
setNumControllerNodes(clusterConfig.numControllers())
.build();
nodes.brokerNodes().forEach((brokerId, brokerNode) -> {
clusterConfig.brokerServerProperties(brokerId).forEach(
(key, value) -> brokerNode.propertyOverrides().put(key.toString(), value.toString()));
Expand Down
59 changes: 50 additions & 9 deletions core/src/test/java/kafka/testkit/KafkaClusterTestKit.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -150,15 +150,16 @@ public KafkaClusterTestKit build() throws Exception {
ThreadUtils.createThreadFactory("KafkaClusterTestKit%d", false));
for (ControllerNode node : nodes.controllerNodes().values()) {
Map<String, String> props = new HashMap<>(configProps);
props.put(KafkaConfig$.MODULE$.ProcessRolesProp(), "controller");
props.put(KafkaConfig$.MODULE$.ProcessRolesProp(), roles(node.id()));

@hachikuji hachikuji May 25, 2022 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Something that was never quite clear to me is why we bother working with the lower-level ControllerServer and BrokerServer instances in this class. It seems like it would be simpler to use KafkaRaftServer which is already doing the dirty work of building KafkaRaftManager and the server objects. Not saying you have to do that here, but it would be useful to know if there is an actual benefit.

@dengziming dengziming May 26, 2022

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I think that's a good point. I will try to improve this.

props.put(KafkaConfig$.MODULE$.NodeIdProp(),
Integer.toString(node.id()));
props.put(KafkaConfig$.MODULE$.MetadataLogDirProp(),
node.metadataDirectory());
props.put(KafkaConfig$.MODULE$.ListenerSecurityProtocolMapProp(),
"CONTROLLER:PLAINTEXT");
props.put(KafkaConfig$.MODULE$.ListenersProp(),
"CONTROLLER://localhost:0");
"EXTERNAL:PLAINTEXT,CONTROLLER:PLAINTEXT");
props.put(KafkaConfig$.MODULE$.ListenersProp(), listeners(node.id()));
props.put(KafkaConfig$.MODULE$.InterBrokerListenerNameProp(),
nodes.interBrokerListenerName().value());
props.put(KafkaConfig$.MODULE$.ControllerListenerNamesProp(),
"CONTROLLER");
// Note: we can't accurately set controller.quorum.voters yet, since we don't
Expand Down Expand Up @@ -203,7 +204,7 @@ metaProperties, config, new MetadataRecordSerde(), metadataPartition, KafkaRaftS
}
for (BrokerNode node : nodes.brokerNodes().values()) {
Map<String, String> props = new HashMap<>(configProps);
props.put(KafkaConfig$.MODULE$.ProcessRolesProp(), "broker");
props.put(KafkaConfig$.MODULE$.ProcessRolesProp(), roles(node.id()));
props.put(KafkaConfig$.MODULE$.BrokerIdProp(),
Integer.toString(node.id()));
props.put(KafkaConfig$.MODULE$.MetadataLogDirProp(),
Expand All @@ -212,8 +213,7 @@ metaProperties, config, new MetadataRecordSerde(), metadataPartition, KafkaRaftS
String.join(",", node.logDataDirectories()));
props.put(KafkaConfig$.MODULE$.ListenerSecurityProtocolMapProp(),
"EXTERNAL:PLAINTEXT,CONTROLLER:PLAINTEXT");
props.put(KafkaConfig$.MODULE$.ListenersProp(),
"EXTERNAL://localhost:0");
props.put(KafkaConfig$.MODULE$.ListenersProp(), listeners(node.id()));
props.put(KafkaConfig$.MODULE$.InterBrokerListenerNameProp(),
nodes.interBrokerListenerName().value());
props.put(KafkaConfig$.MODULE$.ControllerListenerNamesProp(),
Expand All @@ -231,9 +231,15 @@ metaProperties, config, new MetadataRecordSerde(), metadataPartition, KafkaRaftS
String threadNamePrefix = String.format("broker%d_", node.id());
MetaProperties metaProperties = MetaProperties.apply(nodes.clusterId().toString(), node.id());
TopicPartition metadataPartition = new TopicPartition(KafkaRaftServer.MetadataTopic(), 0);
KafkaRaftManager<ApiMessageAndVersion> raftManager = new KafkaRaftManager<>(
KafkaRaftManager<ApiMessageAndVersion> raftManager;
if (raftManagers.containsKey(node.id())) {
raftManager = raftManagers.get(node.id());
} else {
raftManager = new KafkaRaftManager<>(
metaProperties, config, new MetadataRecordSerde(), metadataPartition, KafkaRaftServer.MetadataTopicId(),
Time.SYSTEM, new Metrics(), Option.apply(threadNamePrefix), connectFutureManager.future);
raftManagers.put(node.id(), raftManager);
}
BrokerServer broker = new BrokerServer(
config,
nodes.brokerProperties(node.id()),
Expand All @@ -245,7 +251,6 @@ metaProperties, config, new MetadataRecordSerde(), metadataPartition, KafkaRaftS
connectFutureManager.future
);
brokers.put(node.id(), broker);
raftManagers.put(node.id(), raftManager);
}
} catch (Exception e) {
if (executorService != null) {
Expand All @@ -271,6 +276,42 @@ metaProperties, config, new MetadataRecordSerde(), metadataPartition, KafkaRaftS
brokers, raftManagers, connectFutureManager, baseDirectory);
}

private String listeners(int node) {
if (nodes.isCoResidentNode(node)) {
return "EXTERNAL://localhost:0,CONTROLLER://localhost:0";
}
if (nodes.controllerNodes().containsKey(node)) {
return "CONTROLLER://localhost:0";
}
return "EXTERNAL://localhost:0";
}

private String roles(int node) {
if (nodes.isCoResidentNode(node)) {
return "broker,controller";
}
if (nodes.controllerNodes().containsKey(node)) {
return "controller";
}
return "broker";
}

public Map<String, String> specificControllerProps(boolean coResident) {
Map<String, String> props = new HashMap<>();
props.put(KafkaConfig$.MODULE$.ProcessRolesProp(), coResident ? "controller,broker" : "controller");
props.put(KafkaConfig$.MODULE$.ListenersProp(),
coResident ? "EXTERNAL://localhost:0,CONTROLLER://localhost:0": "CONTROLLER://localhost:0");
return props;
}

public Map<String, String> specificBrokerProps(boolean coResident) {
Map<String, String> props = new HashMap<>();
props.put(KafkaConfig$.MODULE$.ProcessRolesProp(), coResident ? "controller,broker" : "broker");
props.put(KafkaConfig$.MODULE$.ListenersProp(),
coResident ? "EXTERNAL://localhost:0,CONTROLLER://localhost:0": "EXTERNAL://localhost:0");
return props;
}

static private void setupNodeDirectories(File baseDirectory,
String metadataDirectory,
Collection<String> logDataDirectories) throws Exception {
Expand Down
25 changes: 23 additions & 2 deletions core/src/test/java/kafka/testkit/TestKitNodes.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@

public class TestKitNodes {
public static class Builder {
private boolean coResident = false;
private Uuid clusterId = null;
private MetadataVersion bootstrapMetadataVersion = null;
private final NavigableMap<Integer, ControllerNode> controllerNodes = new TreeMap<>();
Expand All @@ -48,6 +49,11 @@ public Builder setBootstrapMetadataVersion(MetadataVersion metadataVersion) {
return this;
}

public Builder setCoResident(boolean coResident) {
this.coResident = coResident;
return this;
}

public Builder addNodes(TestKitNode[] nodes) {
for (TestKitNode node : nodes) {
addNode(node);
Expand Down Expand Up @@ -78,7 +84,7 @@ public Builder setNumControllerNodes(int numControllerNodes) {
controllerNodes.pollFirstEntry();
}
while (controllerNodes.size() < numControllerNodes) {
int nextId = 3000;
int nextId = startControllerId();
if (!controllerNodes.isEmpty()) {
nextId = controllerNodes.lastKey() + 1;
}
Expand All @@ -96,7 +102,7 @@ public Builder setNumBrokerNodes(int numBrokerNodes) {
brokerNodes.pollFirstEntry();
}
while (brokerNodes.size() < numBrokerNodes) {
int nextId = 0;
int nextId = startBrokerId();
if (!brokerNodes.isEmpty()) {
nextId = brokerNodes.lastKey() + 1;
}
Expand All @@ -115,13 +121,28 @@ public TestKitNodes build() {
}
return new TestKitNodes(clusterId, bootstrapMetadataVersion, controllerNodes, brokerNodes);
}

private int startBrokerId() {
return 0;
}

private int startControllerId() {
if (coResident) {
return startBrokerId();
}
return startBrokerId() + 3000;
}
}

private final Uuid clusterId;
private final MetadataVersion bootstrapMetadataVersion;
private final NavigableMap<Integer, ControllerNode> controllerNodes;
private final NavigableMap<Integer, BrokerNode> brokerNodes;

public boolean isCoResidentNode(int node) {
return controllerNodes.containsKey(node) && brokerNodes.containsKey(node);
}

private TestKitNodes(Uuid clusterId,
MetadataVersion bootstrapMetadataVersion,
NavigableMap<Integer, ControllerNode> controllerNodes,
Expand Down
2 changes: 1 addition & 1 deletion core/src/test/scala/unit/kafka/admin/LeaderElectionCommandTest.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ import org.junit.jupiter.api.extension.ExtendWith
import org.junit.jupiter.api.{BeforeEach, Tag}

@ExtendWith(value = Array(classOf[ClusterTestExtensions]))
@ClusterTestDefaults(clusterType = Type.BOTH, brokers = 3)
@ClusterTestDefaults(clusterType = Type.ALL, brokers = 3)
@Tag("integration")
final class LeaderElectionCommandTest(cluster: ClusterInstance) {
import LeaderElectionCommandTest._
Expand Down
2 changes: 1 addition & 1 deletion core/src/test/scala/unit/kafka/server/BrokerMetricNamesTest.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ import org.junit.jupiter.api.extension.ExtendWith

import scala.jdk.CollectionConverters._

@ClusterTestDefaults(clusterType = Type.BOTH)
@ClusterTestDefaults(clusterType = Type.ALL)
@ExtendWith(value = Array(classOf[ClusterTestExtensions]))
class BrokerMetricNamesTest(cluster: ClusterInstance) {
@AfterEach
Expand Down
2 changes: 1 addition & 1 deletion core/src/test/scala/unit/kafka/server/ClientQuotasRequestTest.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ import org.junit.jupiter.api.extension.ExtendWith

import scala.jdk.CollectionConverters._

@ClusterTestDefaults(clusterType = Type.BOTH)
@ClusterTestDefaults(clusterType = Type.ALL)
@ExtendWith(value = Array(classOf[ClusterTestExtensions]))
@Tag("integration")
class ClientQuotasRequestTest(cluster: ClusterInstance) {
Expand Down