Skip to content

Commit 6217fb1

Browse files
SCNiehSCNieh
authored
feat(core): reuse unregistered node when requesting for next node id (#2200)
* feat(core): reuse unregistered node when requesting for next node id (#2183) Signed-off-by: Shichao Nie <[email protected]> * fix(core): fix getting duplicated node id Signed-off-by: Shichao Nie <[email protected]> --------- Signed-off-by: Shichao Nie <[email protected]>
1 parent d6641d7 commit 6217fb1

File tree

4 files changed

+161
-12
lines changed

4 files changed

+161
-12
lines changed

‎metadata/src/main/java/org/apache/kafka/controller/ClusterControlManager.java

+98-11
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@
2727
import org.apache.kafka.common.errors.UnsupportedVersionException;
2828
import org.apache.kafka.common.message.BrokerRegistrationRequestData;
2929
import org.apache.kafka.common.message.ControllerRegistrationRequestData;
30+
import org.apache.kafka.common.message.GetKVsRequestData;
31+
import org.apache.kafka.common.message.PutKVsRequestData;
3032
import org.apache.kafka.common.metadata.BrokerRegistrationChangeRecord;
3133
import org.apache.kafka.common.metadata.FenceBrokerRecord;
3234
import org.apache.kafka.common.metadata.RegisterBrokerRecord;
@@ -39,6 +41,7 @@
3941
import org.apache.kafka.common.protocol.ApiMessage;
4042
import org.apache.kafka.common.utils.LogContext;
4143
import org.apache.kafka.common.utils.Time;
44+
import org.apache.kafka.controller.stream.KVControlManager;
4245
import org.apache.kafka.metadata.BrokerRegistration;
4346
import org.apache.kafka.metadata.BrokerRegistrationFencingChange;
4447
import org.apache.kafka.metadata.BrokerRegistrationInControlledShutdownChange;
@@ -59,8 +62,10 @@
5962

6063
import org.slf4j.Logger;
6164

65+
import java.nio.ByteBuffer;
6266
import java.util.AbstractMap;
6367
import java.util.ArrayList;
68+
import java.util.Collections;
6469
import java.util.HashMap;
6570
import java.util.HashSet;
6671
import java.util.Iterator;
@@ -100,11 +105,17 @@ static class Builder {
100105

101106
// AutoMQ for Kafka inject start
102107
private List<String> quorumVoters;
108+
private KVControlManager kvControlManager;
103109

104110
Builder setQuorumVoters(List<String> quorumVoters) {
105111
this.quorumVoters = quorumVoters;
106112
return this;
107113
}
114+
115+
Builder setKVControlManager(KVControlManager kvControlManager) {
116+
this.kvControlManager = kvControlManager;
117+
return this;
118+
}
108119
// AutoMQ for Kafka inject end
109120

110121
Builder setLogContext(LogContext logContext) {
@@ -180,7 +191,10 @@ ClusterControlManager build() {
180191
featureControl,
181192
zkMigrationEnabled,
182193
brokerUncleanShutdownHandler,
183-
quorumVoters
194+
// AutoMQ inject start
195+
quorumVoters,
196+
kvControlManager
197+
// AutoMQ inject end
184198
);
185199
}
186200
}
@@ -292,6 +306,12 @@ boolean check() {
292306
* The real next available node id is generally one greater than this value.
293307
*/
294308
private AtomicInteger nextNodeId = new AtomicInteger(-1);
309+
310+
/**
311+
* A set of node IDs that have been unregistered and can be reused for new node assignments.
312+
*/
313+
private final KVControlManager kvControlManager;
314+
private static final String REUSABLE_NODE_IDS_KEY = "__automq_reusable_node_ids/";
295315
// AutoMQ for Kafka inject end
296316

297317
private ClusterControlManager(
@@ -304,7 +324,10 @@ private ClusterControlManager(
304324
FeatureControlManager featureControl,
305325
boolean zkMigrationEnabled,
306326
BrokerUncleanShutdownHandler brokerUncleanShutdownHandler,
307-
List<String> quorumVoters
327+
// AutoMQ inject start
328+
List<String> quorumVoters,
329+
KVControlManager kvControlManager
330+
// AutoMQ inject end
308331
) {
309332
this.logContext = logContext;
310333
this.clusterId = clusterId;
@@ -323,6 +346,7 @@ private ClusterControlManager(
323346
this.brokerUncleanShutdownHandler = brokerUncleanShutdownHandler;
324347
// AutoMQ for Kafka inject start
325348
this.maxControllerId = QuorumConfig.parseVoterConnections(quorumVoters).keySet().stream().max(Integer::compareTo).orElse(0);
349+
this.kvControlManager = kvControlManager;
326350
// AutoMQ for Kafka inject end
327351
}
328352

@@ -369,16 +393,73 @@ boolean zkRegistrationAllowed() {
369393

370394
// AutoMQ for Kafka inject start
371395
public ControllerResult<Integer> getNextNodeId() {
372-
int maxBrokerId = brokerRegistrations.keySet().stream().max(Integer::compareTo).orElse(0);
373-
int maxNodeId = Math.max(maxBrokerId, maxControllerId);
374-
int nextId = this.nextNodeId.accumulateAndGet(maxNodeId, (x, y) -> Math.max(x, y) + 1);
375-
// Let the broker's nodeId start from 1000 to easily distinguish broker and controller.
376-
nextId = Math.max(nextId, 1000);
377-
UpdateNextNodeIdRecord record = new UpdateNextNodeIdRecord().setNodeId(nextId);
396+
int nextId;
397+
Set<Integer> reusableNodeIds = getReusableNodeIds();
398+
if (!reusableNodeIds.isEmpty()) {
399+
Iterator<Integer> iterator = reusableNodeIds.iterator();
400+
nextId = iterator.next();
401+
// we simply remove the id from reusable id set because we're unable to determine if the id
402+
// will finally be used.
403+
iterator.remove();
404+
return ControllerResult.atomicOf(putReusableNodeIds(reusableNodeIds), nextId);
405+
} else {
406+
int maxBrokerId = brokerRegistrations.keySet().stream().max(Integer::compareTo).orElse(0);
407+
int maxNodeId = Math.max(maxBrokerId, maxControllerId);
408+
nextId = this.nextNodeId.accumulateAndGet(maxNodeId, (x, y) -> Math.max(x, y) + 1);
409+
// Let the broker's nodeId start from 1000 to easily distinguish broker and controller.
410+
nextId = Math.max(nextId, 1000);
411+
UpdateNextNodeIdRecord record = new UpdateNextNodeIdRecord().setNodeId(nextId);
378412

379-
List<ApiMessageAndVersion> records = new ArrayList<>();
380-
records.add(new ApiMessageAndVersion(record, (short) 0));
381-
return ControllerResult.atomicOf(records, nextId);
413+
List<ApiMessageAndVersion> records = new ArrayList<>();
414+
records.add(new ApiMessageAndVersion(record, (short) 0));
415+
return ControllerResult.atomicOf(records, nextId);
416+
}
417+
}
418+
419+
Set<Integer> getReusableNodeIds() {
420+
return deserializeReusableNodeIds(kvControlManager.getKV(
421+
new GetKVsRequestData.GetKVRequest().setKey(REUSABLE_NODE_IDS_KEY)).value());
422+
}
423+
424+
List<ApiMessageAndVersion> putReusableNodeIds(Set<Integer> reusableNodeIds) {
425+
return kvControlManager.putKV(new PutKVsRequestData.PutKVRequest()
426+
.setKey(REUSABLE_NODE_IDS_KEY)
427+
.setValue(serializeReusableNodeIds(reusableNodeIds))
428+
.setOverwrite(true))
429+
.records();
430+
}
431+
432+
private Set<Integer> deserializeReusableNodeIds(byte[] value) {
433+
if (value == null) {
434+
return new HashSet<>();
435+
}
436+
ByteBuffer buffer = ByteBuffer.wrap(value);
437+
Set<Integer> reusableNodeIds = new HashSet<>();
438+
while (buffer.hasRemaining()) {
439+
reusableNodeIds.add(buffer.getInt());
440+
}
441+
return reusableNodeIds;
442+
}
443+
444+
private byte[] serializeReusableNodeIds(Set<Integer> reusableNodeIds) {
445+
ByteBuffer buffer = ByteBuffer.allocate(reusableNodeIds.size() * Integer.BYTES);
446+
reusableNodeIds.forEach(buffer::putInt);
447+
return buffer.array();
448+
}
449+
450+
public List<ApiMessageAndVersion> registerBrokerRecords(int brokerId) {
451+
Set<Integer> reusableNodeIds = getReusableNodeIds();
452+
if (reusableNodeIds.contains(brokerId)) {
453+
reusableNodeIds.remove(brokerId);
454+
return putReusableNodeIds(reusableNodeIds);
455+
}
456+
return Collections.emptyList();
457+
}
458+
459+
public List<ApiMessageAndVersion> unRegisterBrokerRecords(int brokerId) {
460+
Set<Integer> reusableNodeIds = getReusableNodeIds();
461+
reusableNodeIds.add(brokerId);
462+
return putReusableNodeIds(reusableNodeIds);
382463
}
383464
// AutoMQ for Kafka inject end
384465

@@ -496,6 +577,10 @@ public ControllerResult<BrokerRegistrationReply> registerBroker(
496577
}
497578
heartbeatManager.register(brokerId, record.fenced());
498579

580+
// AutoMQ for Kafka inject start
581+
records.addAll(registerBrokerRecords(brokerId));
582+
// AutoMQ for Kafka inject end
583+
499584
return ControllerResult.atomicOf(records, new BrokerRegistrationReply(record.brokerEpoch()));
500585
}
501586

@@ -583,6 +668,7 @@ public void replay(RegisterBrokerRecord record, long offset) {
583668
if (prevRegistration != null) heartbeatManager.remove(brokerId);
584669
heartbeatManager.register(brokerId, record.fenced());
585670
}
671+
586672
if (prevRegistration == null) {
587673
log.info("Replayed initial RegisterBrokerRecord for broker {}: {}", record.brokerId(), record);
588674
} else if (prevRegistration.incarnationId().equals(record.incarnationId())) {
@@ -608,6 +694,7 @@ public void replay(UnregisterBrokerRecord record) {
608694
if (heartbeatManager != null) heartbeatManager.remove(brokerId);
609695
updateDirectories(brokerId, registration.directories(), null);
610696
brokerRegistrations.remove(brokerId);
697+
// AutoMQ injection end
611698
log.info("Replayed {}", record);
612699
}
613700
}

‎metadata/src/main/java/org/apache/kafka/controller/QuorumController.java

+4-1
Original file line numberDiff line numberDiff line change
@@ -2047,6 +2047,9 @@ private QuorumController(
20472047
this.time = time;
20482048
this.controllerMetrics = controllerMetrics;
20492049
this.snapshotRegistry = new SnapshotRegistry(logContext);
2050+
// AutoMQ for Kafka inject start
2051+
this.kvControlManager = new KVControlManager(snapshotRegistry, logContext);
2052+
// AutoMQ for Kafka inject end
20502053
this.deferredEventQueue = new DeferredEventQueue(logContext);
20512054
this.deferredUnstableEventQueue = new DeferredEventQueue(logContext);
20522055
this.offsetControl = new OffsetControlManager.Builder().
@@ -2094,6 +2097,7 @@ private QuorumController(
20942097
setZkMigrationEnabled(zkMigrationEnabled).
20952098
// AutoMQ for Kafka inject start
20962099
setQuorumVoters(quorumVoters).
2100+
setKVControlManager(kvControlManager).
20972101
// AutoMQ for Kafka inject end
20982102
setBrokerUncleanShutdownHandler(this::handleUncleanBrokerShutdown).
20992103
build();
@@ -2156,7 +2160,6 @@ private QuorumController(
21562160
featureControl::autoMQVersion, time);
21572161
this.streamControlManager = new StreamControlManager(this, snapshotRegistry, logContext,
21582162
this.s3ObjectControlManager, clusterControl, featureControl, replicationControl);
2159-
this.kvControlManager = new KVControlManager(snapshotRegistry, logContext);
21602163
this.topicDeletionManager = new TopicDeletionManager(snapshotRegistry, this, streamControlManager, kvControlManager);
21612164
this.nodeControlManager = new NodeControlManager(snapshotRegistry, new DefaultNodeRuntimeInfoGetter(clusterControl, streamControlManager));
21622165
this.extension = extension.apply(this);

‎metadata/src/main/java/org/apache/kafka/controller/ReplicationControlManager.java

+1
Original file line numberDiff line numberDiff line change
@@ -1486,6 +1486,7 @@ void handleBrokerUnregistered(int brokerId, long brokerEpoch,
14861486
(short) 0));
14871487
// AutoMQ for Kafka inject start
14881488
records.add(nodeControlManager.unregisterNodeRecord(brokerId));
1489+
records.addAll(clusterControl.unRegisterBrokerRecords(brokerId));
14891490
// AutoMQ for Kafka inject end
14901491
}
14911492

‎metadata/src/test/java/org/apache/kafka/controller/ClusterControlManagerTest.java

+58
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import org.apache.kafka.common.message.ControllerRegistrationRequestData;
2929
import org.apache.kafka.common.metadata.BrokerRegistrationChangeRecord;
3030
import org.apache.kafka.common.metadata.FenceBrokerRecord;
31+
import org.apache.kafka.common.metadata.KVRecord;
3132
import org.apache.kafka.common.metadata.RegisterBrokerRecord;
3233
import org.apache.kafka.common.metadata.RegisterBrokerRecord.BrokerEndpoint;
3334
import org.apache.kafka.common.metadata.RegisterBrokerRecord.BrokerEndpointCollection;
@@ -36,6 +37,7 @@
3637
import org.apache.kafka.common.security.auth.SecurityProtocol;
3738
import org.apache.kafka.common.utils.LogContext;
3839
import org.apache.kafka.common.utils.MockTime;
40+
import org.apache.kafka.controller.stream.KVControlManager;
3941
import org.apache.kafka.image.writer.ImageWriterOptions;
4042
import org.apache.kafka.metadata.BrokerRegistration;
4143
import org.apache.kafka.metadata.BrokerRegistrationFencingChange;
@@ -52,6 +54,7 @@
5254
import org.apache.kafka.server.common.MetadataVersion;
5355
import org.apache.kafka.timeline.SnapshotRegistry;
5456

57+
import org.junit.jupiter.api.Assertions;
5558
import org.junit.jupiter.api.Test;
5659
import org.junit.jupiter.api.Timeout;
5760
import org.junit.jupiter.params.ParameterizedTest;
@@ -60,12 +63,14 @@
6063
import org.junit.jupiter.params.provider.MethodSource;
6164
import org.junit.jupiter.params.provider.ValueSource;
6265

66+
import java.util.ArrayList;
6367
import java.util.Arrays;
6468
import java.util.Collections;
6569
import java.util.HashSet;
6670
import java.util.Iterator;
6771
import java.util.List;
6872
import java.util.Optional;
73+
import java.util.Set;
6974
import java.util.stream.Stream;
7075

7176
import static java.util.Arrays.asList;
@@ -718,4 +723,57 @@ public void testReRegistrationAndBrokerEpoch(boolean newIncarnationId) {
718723
clusterControl.brokerRegistrations().get(1).epoch());
719724
}
720725
}
726+
727+
@Test
728+
public void testReusableNodeIds() {
729+
MockTime time = new MockTime(0, 0, 0);
730+
SnapshotRegistry snapshotRegistry = new SnapshotRegistry(new LogContext());
731+
KVControlManager kvControl = new KVControlManager(snapshotRegistry, new LogContext());
732+
FeatureControlManager featureControl = new FeatureControlManager.Builder().
733+
setSnapshotRegistry(snapshotRegistry).
734+
setQuorumFeatures(new QuorumFeatures(0,
735+
QuorumFeatures.defaultFeatureMap(true),
736+
Collections.singletonList(0))).
737+
setMetadataVersion(MetadataVersion.IBP_3_9_IV0).
738+
build();
739+
ClusterControlManager clusterControl = new ClusterControlManager.Builder().
740+
setTime(time).
741+
setSnapshotRegistry(snapshotRegistry).
742+
setSessionTimeoutNs(1000).
743+
setFeatureControlManager(featureControl).
744+
setBrokerUncleanShutdownHandler((brokerId, records) -> { }).
745+
setQuorumVoters(new ArrayList<>()).
746+
setKVControlManager(kvControl).
747+
build();
748+
clusterControl.activate();
749+
Set<Integer> nodeIds = clusterControl.getReusableNodeIds();
750+
Assertions.assertTrue(nodeIds.isEmpty());
751+
clusterControl.putReusableNodeIds(Set.of(1, 2, 3)).forEach(r -> {
752+
kvControl.replay((KVRecord) r.message());
753+
});
754+
nodeIds = clusterControl.getReusableNodeIds();
755+
Assertions.assertEquals(Set.of(1, 2, 3), nodeIds);
756+
757+
clusterControl.unRegisterBrokerRecords(4).forEach(r -> {
758+
kvControl.replay((KVRecord) r.message());
759+
});
760+
nodeIds = clusterControl.getReusableNodeIds();
761+
Assertions.assertEquals(Set.of(1, 2, 3, 4), nodeIds);
762+
763+
clusterControl.registerBrokerRecords(2).forEach(r -> {
764+
kvControl.replay((KVRecord) r.message());
765+
});
766+
nodeIds = clusterControl.getReusableNodeIds();
767+
Assertions.assertEquals(Set.of(1, 3, 4), nodeIds);
768+
769+
ControllerResult<Integer> result = clusterControl.getNextNodeId();
770+
result.records().forEach(r -> {
771+
if (r.message() instanceof KVRecord) {
772+
kvControl.replay((KVRecord) r.message());
773+
}
774+
});
775+
Set<Integer> remainIds = new HashSet<>(Set.of(1, 3, 4));
776+
remainIds.remove(result.response());
777+
Assertions.assertEquals(remainIds, clusterControl.getReusableNodeIds());
778+
}
721779
}

0 commit comments

Comments
 (0)