Honor lease sync on app bootstrap (#1325)

This commit is contained in:
chenylee-aws 2024-05-02 14:43:03 -07:00 committed by GitHub
parent 35fc72b2c8
commit 69cf5996c5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 25 additions and 25 deletions

View file

@ -476,29 +476,29 @@ public class Scheduler implements Runnable {
final Map<StreamIdentifier, StreamConfig> newStreamConfigMap = streamTracker.streamConfigList()
.stream().collect(Collectors.toMap(StreamConfig::streamIdentifier, Function.identity()));
// This is done to ensure that we clean up the stale streams lingering in the lease table.
if (!leaderSynced.get() || !leasesSyncedOnAppInit) {
// Only sync from lease table again if the currentStreamConfigMap and newStreamConfigMap contain
// different set of streams.
if (!newStreamConfigMap.keySet().equals(currentStreamConfigMap.keySet())) {
log.info("Syncing leases for leader to catch up");
final List<MultiStreamLease> leaseTableLeases = fetchMultiStreamLeases();
syncStreamsFromLeaseTableOnAppInit(leaseTableLeases);
final Set<StreamIdentifier> streamsFromLeaseTable = leaseTableLeases.stream()
.map(lease -> StreamIdentifier.multiStreamInstance(lease.streamIdentifier()))
.collect(Collectors.toSet());
// Remove stream from currentStreamConfigMap if this stream in not in the lease table and newStreamConfigMap.
// This means that the leases have already been deleted by the last leader.
currentStreamConfigMap.keySet().stream()
.filter(streamIdentifier -> !newStreamConfigMap.containsKey(streamIdentifier)
&& !streamsFromLeaseTable.contains(streamIdentifier)).forEach(stream -> {
log.info("Removing stream {} from currentStreamConfigMap due to not being active", stream);
currentStreamConfigMap.remove(stream);
staleStreamDeletionMap.remove(stream);
streamsSynced.add(stream);
});
}
leasesSyncedOnAppInit = true;
// Only sync from lease table again if the currentStreamConfigMap and newStreamConfigMap contain
// different set of streams and Leader has not synced the leases yet
// or this is the first app bootstrap.
if ((!leaderSynced.get() && !newStreamConfigMap.keySet().equals(currentStreamConfigMap.keySet()))
|| !leasesSyncedOnAppInit) {
log.info("Syncing leases for leader to catch up");
final List<MultiStreamLease> leaseTableLeases = fetchMultiStreamLeases();
syncStreamsFromLeaseTableOnAppInit(leaseTableLeases);
final Set<StreamIdentifier> streamsFromLeaseTable = leaseTableLeases.stream()
.map(lease -> StreamIdentifier.multiStreamInstance(lease.streamIdentifier()))
.collect(Collectors.toSet());
// Remove stream from currentStreamConfigMap if this stream in not in the lease table and newStreamConfigMap.
// This means that the leases have already been deleted by the last leader.
currentStreamConfigMap.keySet().stream()
.filter(streamIdentifier -> !newStreamConfigMap.containsKey(streamIdentifier)
&& !streamsFromLeaseTable.contains(streamIdentifier)).forEach(stream -> {
log.info("Removing stream {} from currentStreamConfigMap due to not being active", stream);
currentStreamConfigMap.remove(stream);
staleStreamDeletionMap.remove(stream);
streamsSynced.add(stream);
});
}
leasesSyncedOnAppInit = true;
// For new streams discovered, do a shard sync and update the currentStreamConfigMap
for (StreamIdentifier streamIdentifier : newStreamConfigMap.keySet()) {

View file

@ -995,14 +995,14 @@ public class SchedulerTest {
}
@Test
public void testNoDdbLookUpAsStreamMapContainsAllStreams() throws Exception {
public void testSyncLeaseAsThisIsInitialAppBootstrapEvenThoughStreamMapContainsAllStreams() {
final List<StreamConfig> streamConfigList = createDummyStreamConfigList(1, 6);
when(multiStreamTracker.streamConfigList()).thenReturn(Collections.emptyList());
prepareMultiStreamScheduler(streamConfigList);
// Populate currentStreamConfigMap to simulate that the leader has the latest streams.
multiStreamTracker.streamConfigList().forEach(s -> scheduler.currentStreamConfigMap().put(s.streamIdentifier(), s));
scheduler.checkAndSyncStreamShardsAndLeases();
verify(scheduler, never()).syncStreamsFromLeaseTableOnAppInit(any());
scheduler.runProcessLoop();
verify(scheduler).syncStreamsFromLeaseTableOnAppInit(any());
assertTrue(scheduler.currentStreamConfigMap().size() != 0);
}