org.apache.flink.runtime.state.PriorityQueueSetFactory Java Examples

private static <K, N> InternalTimerServiceImpl<K, N> createInternalTimerService(
	KeyGroupRange keyGroupsList,
	KeyContext keyContext,
	ProcessingTimeService processingTimeService,
	TypeSerializer<K> keySerializer,
	TypeSerializer<N> namespaceSerializer,
	PriorityQueueSetFactory priorityQueueSetFactory) {

	TimerSerializer<K, N> timerSerializer = new TimerSerializer<>(keySerializer, namespaceSerializer);

	return new InternalTimerServiceImpl<>(
		keyGroupsList,
		keyContext,
		processingTimeService,
		createTimerQueue("__test_processing_timers", timerSerializer, priorityQueueSetFactory),
		createTimerQueue("__test_event_timers", timerSerializer, priorityQueueSetFactory));
}
 

private static InternalTimerServiceImpl<Integer, String> createAndStartInternalTimerService(
		Triggerable<Integer, String> triggerable,
		KeyContext keyContext,
		ProcessingTimeService processingTimeService,
		KeyGroupRange keyGroupList,
		PriorityQueueSetFactory priorityQueueSetFactory) {
	InternalTimerServiceImpl<Integer, String> service = createInternalTimerService(
		keyGroupList,
		keyContext,
		processingTimeService,
		IntSerializer.INSTANCE,
		StringSerializer.INSTANCE,
		priorityQueueSetFactory);

	service.startTimerService(IntSerializer.INSTANCE, StringSerializer.INSTANCE, triggerable);
	return service;
}
 

private static InternalTimerServiceImpl<Integer, String> createAndStartInternalTimerService(
		Triggerable<Integer, String> triggerable,
		KeyContext keyContext,
		ProcessingTimeService processingTimeService,
		KeyGroupRange keyGroupList,
		PriorityQueueSetFactory priorityQueueSetFactory) {
	InternalTimerServiceImpl<Integer, String> service = createInternalTimerService(
		keyGroupList,
		keyContext,
		processingTimeService,
		IntSerializer.INSTANCE,
		StringSerializer.INSTANCE,
		priorityQueueSetFactory);

	service.startTimerService(IntSerializer.INSTANCE, StringSerializer.INSTANCE, triggerable);
	return service;
}
 

private static <K, N> InternalTimerServiceImpl<K, N> createInternalTimerService(
	KeyGroupRange keyGroupsList,
	KeyContext keyContext,
	ProcessingTimeService processingTimeService,
	TypeSerializer<K> keySerializer,
	TypeSerializer<N> namespaceSerializer,
	PriorityQueueSetFactory priorityQueueSetFactory) {

	TimerSerializer<K, N> timerSerializer = new TimerSerializer<>(keySerializer, namespaceSerializer);

	return new InternalTimerServiceImpl<>(
		keyGroupsList,
		keyContext,
		processingTimeService,
		createTimerQueue("__test_processing_timers", timerSerializer, priorityQueueSetFactory),
		createTimerQueue("__test_event_timers", timerSerializer, priorityQueueSetFactory));
}
 

private static <K, N> InternalTimerServiceImpl<K, N> createInternalTimerService(
	KeyGroupRange keyGroupsList,
	KeyContext keyContext,
	ProcessingTimeService processingTimeService,
	TypeSerializer<K> keySerializer,
	TypeSerializer<N> namespaceSerializer,
	PriorityQueueSetFactory priorityQueueSetFactory) {

	TimerSerializer<K, N> timerSerializer = new TimerSerializer<>(keySerializer, namespaceSerializer);

	return new InternalTimerServiceImpl<>(
		keyGroupsList,
		keyContext,
		processingTimeService,
		createTimerQueue("__test_processing_timers", timerSerializer, priorityQueueSetFactory),
		createTimerQueue("__test_event_timers", timerSerializer, priorityQueueSetFactory));
}
 

private static InternalTimerServiceImpl<Integer, String> createAndStartInternalTimerService(
		Triggerable<Integer, String> triggerable,
		KeyContext keyContext,
		ProcessingTimeService processingTimeService,
		KeyGroupRange keyGroupList,
		PriorityQueueSetFactory priorityQueueSetFactory) {
	InternalTimerServiceImpl<Integer, String> service = createInternalTimerService(
		keyGroupList,
		keyContext,
		processingTimeService,
		IntSerializer.INSTANCE,
		StringSerializer.INSTANCE,
		priorityQueueSetFactory);

	service.startTimerService(IntSerializer.INSTANCE, StringSerializer.INSTANCE, triggerable);
	return service;
}
 

private static <K, N> KeyGroupedInternalPriorityQueue<TimerHeapInternalTimer<K, N>> createTimerQueue(
	String name,
	TimerSerializer<K, N> timerSerializer,
	PriorityQueueSetFactory priorityQueueSetFactory) {
	return priorityQueueSetFactory.create(
		name,
		timerSerializer);
}
 

private static InternalTimerServiceImpl<Integer, String> restoreTimerService(
		Map<Integer, byte[]> state,
		int snapshotVersion,
		Triggerable<Integer, String> triggerable,
		KeyContext keyContext,
		ProcessingTimeService processingTimeService,
		KeyGroupRange keyGroupsList,
		PriorityQueueSetFactory priorityQueueSetFactory) throws Exception {

	// create an empty service
	InternalTimerServiceImpl<Integer, String> service = createInternalTimerService(
		keyGroupsList,
		keyContext,
		processingTimeService,
		IntSerializer.INSTANCE,
		StringSerializer.INSTANCE,
		priorityQueueSetFactory);

	// restore the timers
	for (Integer keyGroupIndex : keyGroupsList) {
		if (state.containsKey(keyGroupIndex)) {
			try (ByteArrayInputStream inputStream = new ByteArrayInputStream(state.get(keyGroupIndex))) {
				InternalTimersSnapshot<?, ?> restoredTimersSnapshot =
					InternalTimersSnapshotReaderWriters
						.getReaderForVersion(snapshotVersion, InternalTimerServiceImplTest.class.getClassLoader())
						.readTimersSnapshot(new DataInputViewStreamWrapper(inputStream));

				service.restoreTimersForKeyGroup(restoredTimersSnapshot, keyGroupIndex);
			}
		}
	}

	// initialize the service
	service.startTimerService(IntSerializer.INSTANCE, StringSerializer.INSTANCE, triggerable);
	return service;
}
 

InternalTimeServiceManager(
	KeyGroupRange localKeyGroupRange,
	KeyContext keyContext,
	PriorityQueueSetFactory priorityQueueSetFactory,
	ProcessingTimeService processingTimeService, boolean useLegacySynchronousSnapshots) {

	this.localKeyGroupRange = Preconditions.checkNotNull(localKeyGroupRange);
	this.priorityQueueSetFactory = Preconditions.checkNotNull(priorityQueueSetFactory);
	this.keyContext = Preconditions.checkNotNull(keyContext);
	this.processingTimeService = Preconditions.checkNotNull(processingTimeService);
	this.useLegacySynchronousSnapshots = useLegacySynchronousSnapshots;

	this.timerServices = new HashMap<>();
}
 

private static <K, N> KeyGroupedInternalPriorityQueue<TimerHeapInternalTimer<K, N>> createTimerQueue(
	String name,
	TimerSerializer<K, N> timerSerializer,
	PriorityQueueSetFactory priorityQueueSetFactory) {
	return priorityQueueSetFactory.create(
		name,
		timerSerializer);
}
 

private static InternalTimerServiceImpl<Integer, String> restoreTimerService(
		Map<Integer, byte[]> state,
		int snapshotVersion,
		Triggerable<Integer, String> triggerable,
		KeyContext keyContext,
		ProcessingTimeService processingTimeService,
		KeyGroupRange keyGroupsList,
		PriorityQueueSetFactory priorityQueueSetFactory) throws Exception {

	// create an empty service
	InternalTimerServiceImpl<Integer, String> service = createInternalTimerService(
		keyGroupsList,
		keyContext,
		processingTimeService,
		IntSerializer.INSTANCE,
		StringSerializer.INSTANCE,
		priorityQueueSetFactory);

	// restore the timers
	for (Integer keyGroupIndex : keyGroupsList) {
		if (state.containsKey(keyGroupIndex)) {
			try (ByteArrayInputStream inputStream = new ByteArrayInputStream(state.get(keyGroupIndex))) {
				InternalTimersSnapshot<?, ?> restoredTimersSnapshot =
					InternalTimersSnapshotReaderWriters
						.getReaderForVersion(snapshotVersion, InternalTimerServiceImplTest.class.getClassLoader())
						.readTimersSnapshot(new DataInputViewStreamWrapper(inputStream));

				service.restoreTimersForKeyGroup(restoredTimersSnapshot, keyGroupIndex);
			}
		}
	}

	// initialize the service
	service.startTimerService(IntSerializer.INSTANCE, StringSerializer.INSTANCE, triggerable);
	return service;
}
 

InternalTimeServiceManager(
	KeyGroupRange localKeyGroupRange,
	KeyContext keyContext,
	PriorityQueueSetFactory priorityQueueSetFactory,
	ProcessingTimeService processingTimeService, boolean useLegacySynchronousSnapshots) {

	this.localKeyGroupRange = Preconditions.checkNotNull(localKeyGroupRange);
	this.priorityQueueSetFactory = Preconditions.checkNotNull(priorityQueueSetFactory);
	this.keyContext = Preconditions.checkNotNull(keyContext);
	this.processingTimeService = Preconditions.checkNotNull(processingTimeService);
	this.useLegacySynchronousSnapshots = useLegacySynchronousSnapshots;

	this.timerServices = new HashMap<>();
}
 

private static <K, N> KeyGroupedInternalPriorityQueue<TimerHeapInternalTimer<K, N>> createTimerQueue(
	String name,
	TimerSerializer<K, N> timerSerializer,
	PriorityQueueSetFactory priorityQueueSetFactory) {
	return priorityQueueSetFactory.create(
		name,
		timerSerializer);
}
 

InternalTimeServiceManager(
	KeyGroupRange localKeyGroupRange,
	KeyContext keyContext,
	PriorityQueueSetFactory priorityQueueSetFactory,
	ProcessingTimeService processingTimeService, boolean useLegacySynchronousSnapshots) {

	this.localKeyGroupRange = Preconditions.checkNotNull(localKeyGroupRange);
	this.priorityQueueSetFactory = Preconditions.checkNotNull(priorityQueueSetFactory);
	this.keyContext = Preconditions.checkNotNull(keyContext);
	this.processingTimeService = Preconditions.checkNotNull(processingTimeService);
	this.useLegacySynchronousSnapshots = useLegacySynchronousSnapshots;

	this.timerServices = new HashMap<>();
}
 

private static InternalTimerServiceImpl<Integer, String> restoreTimerService(
		Map<Integer, byte[]> state,
		int snapshotVersion,
		Triggerable<Integer, String> triggerable,
		KeyContext keyContext,
		ProcessingTimeService processingTimeService,
		KeyGroupRange keyGroupsList,
		PriorityQueueSetFactory priorityQueueSetFactory) throws Exception {

	// create an empty service
	InternalTimerServiceImpl<Integer, String> service = createInternalTimerService(
		keyGroupsList,
		keyContext,
		processingTimeService,
		IntSerializer.INSTANCE,
		StringSerializer.INSTANCE,
		priorityQueueSetFactory);

	// restore the timers
	for (Integer keyGroupIndex : keyGroupsList) {
		if (state.containsKey(keyGroupIndex)) {
			try (ByteArrayInputStream inputStream = new ByteArrayInputStream(state.get(keyGroupIndex))) {
				InternalTimersSnapshot<?, ?> restoredTimersSnapshot =
					InternalTimersSnapshotReaderWriters
						.getReaderForVersion(snapshotVersion, InternalTimerServiceImplTest.class.getClassLoader())
						.readTimersSnapshot(new DataInputViewStreamWrapper(inputStream));

				service.restoreTimersForKeyGroup(restoredTimersSnapshot, keyGroupIndex);
			}
		}
	}

	// initialize the service
	service.startTimerService(IntSerializer.INSTANCE, StringSerializer.INSTANCE, triggerable);
	return service;
}
 

public RocksDBKeyedStateBackend(
	ClassLoader userCodeClassLoader,
	File instanceBasePath,
	RocksDBResourceContainer optionsContainer,
	Function<String, ColumnFamilyOptions> columnFamilyOptionsFactory,
	TaskKvStateRegistry kvStateRegistry,
	TypeSerializer<K> keySerializer,
	ExecutionConfig executionConfig,
	TtlTimeProvider ttlTimeProvider,
	RocksDB db,
	LinkedHashMap<String, RocksDbKvStateInfo> kvStateInformation,
	int keyGroupPrefixBytes,
	CloseableRegistry cancelStreamRegistry,
	StreamCompressionDecorator keyGroupCompressionDecorator,
	ResourceGuard rocksDBResourceGuard,
	RocksDBSnapshotStrategyBase<K> checkpointSnapshotStrategy,
	RocksDBSnapshotStrategyBase<K> savepointSnapshotStrategy,
	RocksDBWriteBatchWrapper writeBatchWrapper,
	ColumnFamilyHandle defaultColumnFamilyHandle,
	RocksDBNativeMetricMonitor nativeMetricMonitor,
	RocksDBSerializedCompositeKeyBuilder<K> sharedRocksKeyBuilder,
	PriorityQueueSetFactory priorityQueueFactory,
	RocksDbTtlCompactFiltersManager ttlCompactFiltersManager,
	InternalKeyContext<K> keyContext,
	@Nonnegative long writeBatchSize) {

	super(
		kvStateRegistry,
		keySerializer,
		userCodeClassLoader,
		executionConfig,
		ttlTimeProvider,
		cancelStreamRegistry,
		keyGroupCompressionDecorator,
		keyContext);

	this.ttlCompactFiltersManager = ttlCompactFiltersManager;

	// ensure that we use the right merge operator, because other code relies on this
	this.columnFamilyOptionsFactory = Preconditions.checkNotNull(columnFamilyOptionsFactory);

	this.optionsContainer = Preconditions.checkNotNull(optionsContainer);

	this.instanceBasePath = Preconditions.checkNotNull(instanceBasePath);

	this.keyGroupPrefixBytes = keyGroupPrefixBytes;
	this.kvStateInformation = kvStateInformation;

	this.writeOptions = optionsContainer.getWriteOptions();
	this.readOptions = optionsContainer.getReadOptions();
	checkArgument(writeBatchSize >= 0, "Write batch size have to be no negative value.");
	this.writeBatchSize = writeBatchSize;
	this.db = db;
	this.rocksDBResourceGuard = rocksDBResourceGuard;
	this.checkpointSnapshotStrategy = checkpointSnapshotStrategy;
	this.savepointSnapshotStrategy = savepointSnapshotStrategy;
	this.writeBatchWrapper = writeBatchWrapper;
	this.defaultColumnFamily = defaultColumnFamilyHandle;
	this.nativeMetricMonitor = nativeMetricMonitor;
	this.sharedRocksKeyBuilder = sharedRocksKeyBuilder;
	this.priorityQueueFactory = priorityQueueFactory;
}
 

public RocksDBKeyedStateBackend(
	ClassLoader userCodeClassLoader,
	File instanceBasePath,
	DBOptions dbOptions,
	Function<String, ColumnFamilyOptions> columnFamilyOptionsFactory,
	TaskKvStateRegistry kvStateRegistry,
	StateSerializerProvider<K> keySerializerProvider,
	int numberOfKeyGroups,
	KeyGroupRange keyGroupRange,
	ExecutionConfig executionConfig,
	TtlTimeProvider ttlTimeProvider,
	RocksDB db,
	LinkedHashMap<String, RocksDbKvStateInfo> kvStateInformation,
	int keyGroupPrefixBytes,
	CloseableRegistry cancelStreamRegistry,
	StreamCompressionDecorator keyGroupCompressionDecorator,
	ResourceGuard rocksDBResourceGuard,
	RocksDBSnapshotStrategyBase<K> checkpointSnapshotStrategy,
	RocksDBSnapshotStrategyBase<K> savepointSnapshotStrategy,
	RocksDBWriteBatchWrapper writeBatchWrapper,
	ColumnFamilyHandle defaultColumnFamilyHandle,
	RocksDBNativeMetricMonitor nativeMetricMonitor,
	RocksDBSerializedCompositeKeyBuilder<K> sharedRocksKeyBuilder,
	PriorityQueueSetFactory priorityQueueFactory,
	RocksDbTtlCompactFiltersManager ttlCompactFiltersManager) {

	super(kvStateRegistry, keySerializerProvider, userCodeClassLoader, numberOfKeyGroups,
		keyGroupRange, executionConfig, ttlTimeProvider, cancelStreamRegistry, keyGroupCompressionDecorator);

	this.ttlCompactFiltersManager = ttlCompactFiltersManager;

	// ensure that we use the right merge operator, because other code relies on this
	this.columnFamilyOptionsFactory = Preconditions.checkNotNull(columnFamilyOptionsFactory);

	this.dbOptions = Preconditions.checkNotNull(dbOptions);

	this.instanceBasePath = Preconditions.checkNotNull(instanceBasePath);

	this.keyGroupPrefixBytes = keyGroupPrefixBytes;
	this.kvStateInformation = kvStateInformation;

	this.writeOptions = new WriteOptions().setDisableWAL(true);
	this.db = db;
	this.rocksDBResourceGuard = rocksDBResourceGuard;
	this.checkpointSnapshotStrategy = checkpointSnapshotStrategy;
	this.savepointSnapshotStrategy = savepointSnapshotStrategy;
	this.writeBatchWrapper = writeBatchWrapper;
	this.defaultColumnFamily = defaultColumnFamilyHandle;
	this.nativeMetricMonitor = nativeMetricMonitor;
	this.sharedRocksKeyBuilder = sharedRocksKeyBuilder;
	this.priorityQueueFactory = priorityQueueFactory;
}
 

@VisibleForTesting
PriorityQueueSetFactory getPriorityQueueFactory() {
	return priorityQueueFactory;
}
 

protected PriorityQueueSetFactory createQueueFactory(KeyGroupRange keyGroupRange, int numKeyGroups) {
	return new HeapPriorityQueueSetFactory(keyGroupRange, numKeyGroups, 128);
}
 

private PriorityQueueSetFactory createQueueFactory() {
	return createQueueFactory(testKeyGroupRange, maxParallelism);
}
 

@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<TimerHeapInternalTimer<Integer, String>>[] expectedNonEmptyTimerSets = new HashSet[totalNoOfKeyGroups];
	TestKeyContext keyContext = new TestKeyContext();

	final KeyGroupRange keyGroupRange = new KeyGroupRange(startKeyGroupIdx, endKeyGroupIdx);

	final PriorityQueueSetFactory priorityQueueSetFactory =
		createQueueFactory(keyGroupRange, totalNoOfKeyGroups);

	InternalTimerServiceImpl<Integer, String> timerService = createInternalTimerService(
		keyGroupRange,
		keyContext,
		new TestProcessingTimeService(),
		IntSerializer.INSTANCE,
		StringSerializer.INSTANCE,
		priorityQueueSetFactory);

	timerService.startTimerService(IntSerializer.INSTANCE, StringSerializer.INSTANCE, mock(Triggerable.class));

	for (int i = 0; i < totalNoOfTimers; i++) {

		// create the timer to be registered
		TimerHeapInternalTimer<Integer, String> timer = new TimerHeapInternalTimer<>(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<TimerHeapInternalTimer<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());
	}

	List<Set<TimerHeapInternalTimer<Integer, String>>> eventTimeTimers =
		timerService.getEventTimeTimersPerKeyGroup();
	List<Set<TimerHeapInternalTimer<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<TimerHeapInternalTimer<Integer, String>> expected = expectedNonEmptyTimerSets[i];
		Set<TimerHeapInternalTimer<Integer, String>> actualEvent = eventTimeTimers.get(i);
		Set<TimerHeapInternalTimer<Integer, String>> actualProcessing = processingTimeTimers.get(i);

		if (expected == null) {
			Assert.assertTrue(actualEvent.isEmpty());
			Assert.assertTrue(actualProcessing.isEmpty());
		} else {
			Assert.assertEquals(expected, actualEvent);
			Assert.assertEquals(expected, actualProcessing);
		}
	}
}
 

/**
 * Verify that we only ever have one processing-time task registered at the
 * {@link ProcessingTimeService}.
 */
@Test
public void testOnlySetsOnePhysicalProcessingTimeTimer() throws Exception {
	@SuppressWarnings("unchecked")
	Triggerable<Integer, String> mockTriggerable = mock(Triggerable.class);

	TestKeyContext keyContext = new TestKeyContext();

	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
	PriorityQueueSetFactory priorityQueueSetFactory =
		new HeapPriorityQueueSetFactory(testKeyGroupRange, maxParallelism, 128);
	InternalTimerServiceImpl<Integer, String> timerService =
			createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, priorityQueueSetFactory);

	int key = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	keyContext.setCurrentKey(key);

	timerService.registerProcessingTimeTimer("ciao", 10);
	timerService.registerProcessingTimeTimer("ciao", 20);
	timerService.registerProcessingTimeTimer("ciao", 30);
	timerService.registerProcessingTimeTimer("hello", 10);
	timerService.registerProcessingTimeTimer("hello", 20);

	assertEquals(5, timerService.numProcessingTimeTimers());
	assertEquals(2, timerService.numProcessingTimeTimers("hello"));
	assertEquals(3, timerService.numProcessingTimeTimers("ciao"));

	assertEquals(1, processingTimeService.getNumActiveTimers());
	assertThat(processingTimeService.getActiveTimerTimestamps(), containsInAnyOrder(10L));

	processingTimeService.setCurrentTime(10);

	assertEquals(3, timerService.numProcessingTimeTimers());
	assertEquals(1, timerService.numProcessingTimeTimers("hello"));
	assertEquals(2, timerService.numProcessingTimeTimers("ciao"));

	assertEquals(1, processingTimeService.getNumActiveTimers());
	assertThat(processingTimeService.getActiveTimerTimestamps(), containsInAnyOrder(20L));

	processingTimeService.setCurrentTime(20);

	assertEquals(1, timerService.numProcessingTimeTimers());
	assertEquals(0, timerService.numProcessingTimeTimers("hello"));
	assertEquals(1, timerService.numProcessingTimeTimers("ciao"));

	assertEquals(1, processingTimeService.getNumActiveTimers());
	assertThat(processingTimeService.getActiveTimerTimestamps(), containsInAnyOrder(30L));

	processingTimeService.setCurrentTime(30);

	assertEquals(0, timerService.numProcessingTimeTimers());

	assertEquals(0, processingTimeService.getNumActiveTimers());

	timerService.registerProcessingTimeTimer("ciao", 40);

	assertEquals(1, processingTimeService.getNumActiveTimers());
}
 

@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<TimerHeapInternalTimer<Integer, String>>[] expectedNonEmptyTimerSets = new HashSet[totalNoOfKeyGroups];
	TestKeyContext keyContext = new TestKeyContext();

	final KeyGroupRange keyGroupRange = new KeyGroupRange(startKeyGroupIdx, endKeyGroupIdx);

	final PriorityQueueSetFactory priorityQueueSetFactory =
		createQueueFactory(keyGroupRange, totalNoOfKeyGroups);

	InternalTimerServiceImpl<Integer, String> timerService = createInternalTimerService(
		keyGroupRange,
		keyContext,
		new TestProcessingTimeService(),
		IntSerializer.INSTANCE,
		StringSerializer.INSTANCE,
		priorityQueueSetFactory);

	timerService.startTimerService(IntSerializer.INSTANCE, StringSerializer.INSTANCE, mock(Triggerable.class));

	for (int i = 0; i < totalNoOfTimers; i++) {

		// create the timer to be registered
		TimerHeapInternalTimer<Integer, String> timer = new TimerHeapInternalTimer<>(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<TimerHeapInternalTimer<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());
	}

	List<Set<TimerHeapInternalTimer<Integer, String>>> eventTimeTimers =
		timerService.getEventTimeTimersPerKeyGroup();
	List<Set<TimerHeapInternalTimer<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<TimerHeapInternalTimer<Integer, String>> expected = expectedNonEmptyTimerSets[i];
		Set<TimerHeapInternalTimer<Integer, String>> actualEvent = eventTimeTimers.get(i);
		Set<TimerHeapInternalTimer<Integer, String>> actualProcessing = processingTimeTimers.get(i);

		if (expected == null) {
			Assert.assertTrue(actualEvent.isEmpty());
			Assert.assertTrue(actualProcessing.isEmpty());
		} else {
			Assert.assertEquals(expected, actualEvent);
			Assert.assertEquals(expected, actualProcessing);
		}
	}
}
 

/**
 * Verify that we only ever have one processing-time task registered at the
 * {@link ProcessingTimeService}.
 */
@Test
public void testOnlySetsOnePhysicalProcessingTimeTimer() throws Exception {
	@SuppressWarnings("unchecked")
	Triggerable<Integer, String> mockTriggerable = mock(Triggerable.class);

	TestKeyContext keyContext = new TestKeyContext();

	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
	PriorityQueueSetFactory priorityQueueSetFactory =
		new HeapPriorityQueueSetFactory(testKeyGroupRange, maxParallelism, 128);
	InternalTimerServiceImpl<Integer, String> timerService =
			createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, priorityQueueSetFactory);

	int key = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	keyContext.setCurrentKey(key);

	timerService.registerProcessingTimeTimer("ciao", 10);
	timerService.registerProcessingTimeTimer("ciao", 20);
	timerService.registerProcessingTimeTimer("ciao", 30);
	timerService.registerProcessingTimeTimer("hello", 10);
	timerService.registerProcessingTimeTimer("hello", 20);

	assertEquals(5, timerService.numProcessingTimeTimers());
	assertEquals(2, timerService.numProcessingTimeTimers("hello"));
	assertEquals(3, timerService.numProcessingTimeTimers("ciao"));

	assertEquals(1, processingTimeService.getNumActiveTimers());
	assertThat(processingTimeService.getActiveTimerTimestamps(), containsInAnyOrder(10L));

	processingTimeService.setCurrentTime(10);

	assertEquals(3, timerService.numProcessingTimeTimers());
	assertEquals(1, timerService.numProcessingTimeTimers("hello"));
	assertEquals(2, timerService.numProcessingTimeTimers("ciao"));

	assertEquals(1, processingTimeService.getNumActiveTimers());
	assertThat(processingTimeService.getActiveTimerTimestamps(), containsInAnyOrder(20L));

	processingTimeService.setCurrentTime(20);

	assertEquals(1, timerService.numProcessingTimeTimers());
	assertEquals(0, timerService.numProcessingTimeTimers("hello"));
	assertEquals(1, timerService.numProcessingTimeTimers("ciao"));

	assertEquals(1, processingTimeService.getNumActiveTimers());
	assertThat(processingTimeService.getActiveTimerTimestamps(), containsInAnyOrder(30L));

	processingTimeService.setCurrentTime(30);

	assertEquals(0, timerService.numProcessingTimeTimers());

	assertEquals(0, processingTimeService.getNumActiveTimers());

	timerService.registerProcessingTimeTimer("ciao", 40);

	assertEquals(1, processingTimeService.getNumActiveTimers());
}
 

@VisibleForTesting
PriorityQueueSetFactory getPriorityQueueFactory() {
	return priorityQueueFactory;
}
 

@VisibleForTesting
PriorityQueueSetFactory getPriorityQueueFactory() {
	return priorityQueueFactory;
}
 

private PriorityQueueSetFactory createQueueFactory() {
	return createQueueFactory(testKeyGroupRange, maxParallelism);
}
 

protected PriorityQueueSetFactory createQueueFactory(KeyGroupRange keyGroupRange, int numKeyGroups) {
	return new HeapPriorityQueueSetFactory(keyGroupRange, numKeyGroups, 128);
}
 

private PriorityQueueSetFactory createQueueFactory() {
	return createQueueFactory(testKeyGroupRange, maxParallelism);
}
 

protected PriorityQueueSetFactory createQueueFactory(KeyGroupRange keyGroupRange, int numKeyGroups) {
	return new HeapPriorityQueueSetFactory(keyGroupRange, numKeyGroups, 128);
}