@Test
public void testTimerAssignmentToKeyGroups() {
int totalNoOfTimers = 100;
int totalNoOfKeyGroups = 100;
int startKeyGroupIdx = 0;
int endKeyGroupIdx = totalNoOfKeyGroups - 1; // we have 0 to 99
@SuppressWarnings("unchecked")
Set<InternalTimer<Integer, String>>[] expectedNonEmptyTimerSets =
new HashSet[totalNoOfKeyGroups];
TestKeyContext keyContext = new TestKeyContext();
HeapInternalTimerService<Integer, String> timerService =
new HeapInternalTimerService<>(
totalNoOfKeyGroups,
new KeyGroupRange(startKeyGroupIdx, endKeyGroupIdx),
keyContext,
new TestProcessingTimeService());
timerService.startTimerService(
IntSerializer.INSTANCE, StringSerializer.INSTANCE, mock(Triggerable.class));
for (int i = 0; i < totalNoOfTimers; i++) {
// create the timer to be registered
InternalTimer<Integer, String> timer = new InternalTimer<>(10 + i, i, "hello_world_" + i);
int keyGroupIdx =
KeyGroupRangeAssignment.assignToKeyGroup(timer.getKey(), totalNoOfKeyGroups);
// add it in the adequate expected set of timers per keygroup
Set<InternalTimer<Integer, String>> timerSet = expectedNonEmptyTimerSets[keyGroupIdx];
if (timerSet == null) {
timerSet = new HashSet<>();
expectedNonEmptyTimerSets[keyGroupIdx] = timerSet;
}
timerSet.add(timer);
// register the timer as both processing and event time one
keyContext.setCurrentKey(timer.getKey());
timerService.registerEventTimeTimer(timer.getNamespace(), timer.getTimestamp());
timerService.registerProcessingTimeTimer(timer.getNamespace(), timer.getTimestamp());
}
Set<InternalTimer<Integer, String>>[] eventTimeTimers =
timerService.getEventTimeTimersPerKeyGroup();
Set<InternalTimer<Integer, String>>[] processingTimeTimers =
timerService.getProcessingTimeTimersPerKeyGroup();
// finally verify that the actual timers per key group sets are the expected ones.
for (int i = 0; i < expectedNonEmptyTimerSets.length; i++) {
Set<InternalTimer<Integer, String>> expected = expectedNonEmptyTimerSets[i];
Set<InternalTimer<Integer, String>> actualEvent = eventTimeTimers[i];
Set<InternalTimer<Integer, String>> actualProcessing = processingTimeTimers[i];
if (expected == null) {
Assert.assertNull(actualEvent);
Assert.assertNull(actualProcessing);
} else {
Assert.assertArrayEquals(expected.toArray(), actualEvent.toArray());
Assert.assertArrayEquals(expected.toArray(), actualProcessing.toArray());
}
}
}
