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

View file

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