Java Code Examples for org.apache.flink.streaming.runtime.tasks.TestProcessingTimeService#setCurrentTime()

@Test
public void testCurrentProcessingTime() throws Exception {

	@SuppressWarnings("unchecked")
	Triggerable<Integer, String> mockTriggerable = mock(Triggerable.class);

	TestKeyContext keyContext = new TestKeyContext();
	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
	InternalTimerServiceImpl<Integer, String> timerService =
			createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());

	processingTimeService.setCurrentTime(17L);
	assertEquals(17, timerService.currentProcessingTime());

	processingTimeService.setCurrentTime(42);
	assertEquals(42, timerService.currentProcessingTime());
}
 

@Test
public void testCurrentProcessingTime() throws Exception {

	@SuppressWarnings("unchecked")
	Triggerable<Integer, String> mockTriggerable = mock(Triggerable.class);

	TestKeyContext keyContext = new TestKeyContext();
	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
	InternalTimerServiceImpl<Integer, String> timerService =
			createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());

	processingTimeService.setCurrentTime(17L);
	assertEquals(17, timerService.currentProcessingTime());

	processingTimeService.setCurrentTime(42);
	assertEquals(42, timerService.currentProcessingTime());
}
 

@Test
public void testAutomaticWatermarkContext() throws Exception {

	// regular stream source operator
	final StoppableStreamSource<String, InfiniteSource<String>> operator =
		new StoppableStreamSource<>(new InfiniteSource<String>());

	long watermarkInterval = 10;
	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
	processingTimeService.setCurrentTime(0);

	setupSourceOperator(operator, TimeCharacteristic.IngestionTime, watermarkInterval, processingTimeService);

	final List<StreamElement> output = new ArrayList<>();

	StreamSourceContexts.getSourceContext(TimeCharacteristic.IngestionTime,
		operator.getContainingTask().getProcessingTimeService(),
		operator.getContainingTask().getCheckpointLock(),
		operator.getContainingTask().getStreamStatusMaintainer(),
		new CollectorOutput<String>(output),
		operator.getExecutionConfig().getAutoWatermarkInterval(),
		-1);

	// periodically emit the watermarks
	// even though we start from 1 the watermark are still
	// going to be aligned with the watermark interval.

	for (long i = 1; i < 100; i += watermarkInterval)  {
		processingTimeService.setCurrentTime(i);
	}

	assertTrue(output.size() == 9);

	long nextWatermark = 0;
	for (StreamElement el : output) {
		nextWatermark += watermarkInterval;
		Watermark wm = (Watermark) el;
		assertTrue(wm.getTimestamp() == nextWatermark);
	}
}
 

private AbstractStreamOperatorTestHarness(
		StreamOperator<OUT> operator,
		MockEnvironment env,
		boolean environmentIsInternal,
		OperatorID operatorID) throws Exception {
	this.operator = operator;
	this.outputList = new ConcurrentLinkedQueue<>();
	this.sideOutputLists = new HashMap<>();

	Configuration underlyingConfig = env.getTaskConfiguration();
	this.config = new StreamConfig(underlyingConfig);
	this.config.setCheckpointingEnabled(true);
	this.config.setOperatorID(operatorID);
	this.executionConfig = env.getExecutionConfig();
	this.closableRegistry = new CloseableRegistry();
	this.checkpointLock = new Object();

	this.environment = Preconditions.checkNotNull(env);

	this.taskStateManager = (TestTaskStateManager) env.getTaskStateManager();
	this.internalEnvironment = environmentIsInternal ? Optional.of(environment) : Optional.empty();

	processingTimeService = new TestProcessingTimeService();
	processingTimeService.setCurrentTime(0);

	this.streamTaskStateInitializer = createStreamTaskStateManager(environment, stateBackend, processingTimeService);

	BiConsumer<String, Throwable> handleAsyncException = (message, t) -> {
		wasFailedExternally = true;
	};

	mockTask = new MockStreamTaskBuilder(env)
		.setCheckpointLock(checkpointLock)
		.setConfig(config)
		.setExecutionConfig(executionConfig)
		.setStreamTaskStateInitializer(streamTaskStateInitializer)
		.setClosableRegistry(closableRegistry)
		.setCheckpointStorage(checkpointStorage)
		.setProcessingTimeService(processingTimeService)
		.setHandleAsyncException(handleAsyncException)
		.build();
}
 

/**
 * This also verifies that we don't have leakage between keys/namespaces.
 */
@Test
public void testSetAndFireProcessingTimeTimers() throws Exception {
	@SuppressWarnings("unchecked")
	Triggerable<Integer, String> mockTriggerable = mock(Triggerable.class);

	TestKeyContext keyContext = new TestKeyContext();
	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
	InternalTimerServiceImpl<Integer, String> timerService =
			createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());

	// get two different keys
	int key1 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	int key2 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	while (key2 == key1) {
		key2 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	}

	keyContext.setCurrentKey(key1);

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

	keyContext.setCurrentKey(key2);

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

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

	processingTimeService.setCurrentTime(10);

	verify(mockTriggerable, times(4)).onProcessingTime(anyInternalTimer());
	verify(mockTriggerable, times(1)).onProcessingTime(eq(new TimerHeapInternalTimer<>(10, key1, "ciao")));
	verify(mockTriggerable, times(1)).onProcessingTime(eq(new TimerHeapInternalTimer<>(10, key1, "hello")));
	verify(mockTriggerable, times(1)).onProcessingTime(eq(new TimerHeapInternalTimer<>(10, key2, "ciao")));
	verify(mockTriggerable, times(1)).onProcessingTime(eq(new TimerHeapInternalTimer<>(10, key2, "hello")));

	assertEquals(0, timerService.numProcessingTimeTimers());
}
 

private void testLatencyMarkEmission(int numberLatencyMarkers, OperatorSetupOperation operatorSetup) throws Exception {
	final List<StreamElement> output = new ArrayList<>();

	final TestProcessingTimeService testProcessingTimeService = new TestProcessingTimeService();
	testProcessingTimeService.setCurrentTime(0L);
	final List<Long> processingTimes = Arrays.asList(1L, 10L, 11L, 21L, maxProcessingTime);

	// regular stream source operator
	final StreamSource<Long, ProcessingTimeServiceSource> operator =
		new StreamSource<>(new ProcessingTimeServiceSource(testProcessingTimeService, processingTimes));

	operatorSetup.setupSourceOperator(operator, testProcessingTimeService);

	// run and wait to be stopped
	OperatorChain<?, ?> operatorChain = new OperatorChain<>(
		operator.getContainingTask(),
		StreamTask.createRecordWriterDelegate(operator.getOperatorConfig(), new MockEnvironmentBuilder().build()));
	try {
		operator.run(new Object(), mock(StreamStatusMaintainer.class), new CollectorOutput<>(output), operatorChain);
		operator.close();
	} finally {
		operatorChain.releaseOutputs();
	}

	assertEquals(
		numberLatencyMarkers + 1, // + 1 is the final watermark element
		output.size());

	long timestamp = 0L;
	int expectedLatencyIndex = 0;

	int i = 0;
	// verify that its only latency markers + a final watermark
	for (; i < numberLatencyMarkers; i++) {
		StreamElement se = output.get(i);
		Assert.assertTrue(se.isLatencyMarker());
		Assert.assertEquals(operator.getOperatorID(), se.asLatencyMarker().getOperatorId());
		Assert.assertEquals(0, se.asLatencyMarker().getSubtaskIndex());

		// determines the next latency mark that should've been emitted
		// latency marks are emitted once per latencyMarkInterval,
		// as a result of which we never emit both 10 and 11
		while (timestamp > processingTimes.get(expectedLatencyIndex)) {
			expectedLatencyIndex++;
		}
		Assert.assertEquals(processingTimes.get(expectedLatencyIndex).longValue(), se.asLatencyMarker().getMarkedTime());

		timestamp += latencyMarkInterval;
	}

	Assert.assertTrue(output.get(i).isWatermark());
}
 

/**
 * 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 testPeriodicWatermarks() throws Exception {
	final String testTopic = "test topic name";
	Map<KafkaTopicPartition, Long> originalPartitions = new HashMap<>();
	originalPartitions.put(
			new KafkaTopicPartition(testTopic, 7),
			KafkaTopicPartitionStateSentinel.LATEST_OFFSET);
	originalPartitions.put(
			new KafkaTopicPartition(testTopic, 13),
			KafkaTopicPartitionStateSentinel.LATEST_OFFSET);
	originalPartitions.put(
			new KafkaTopicPartition(testTopic, 21),
			KafkaTopicPartitionStateSentinel.LATEST_OFFSET);

	TestSourceContext<Long> sourceContext = new TestSourceContext<>();

	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();

	TestFetcher<Long> fetcher = new TestFetcher<>(
			sourceContext,
			originalPartitions,
			new SerializedValue<>(testWmStrategy),
			processingTimeService,
			10);

	final KafkaTopicPartitionState<Long, Object> part1 =
			fetcher.subscribedPartitionStates().get(0);
	final KafkaTopicPartitionState<Long, Object> part2 =
			fetcher.subscribedPartitionStates().get(1);
	final KafkaTopicPartitionState<Long, Object> part3 =
			fetcher.subscribedPartitionStates().get(2);

	// elements generate a watermark if the timestamp is a multiple of three

	// elements for partition 1
	emitRecord(fetcher, 1L, part1, 1L);
	emitRecord(fetcher, 1L, part1, 1L);
	emitRecord(fetcher, 2L, part1, 2L);
	emitRecord(fetcher, 3L, part1, 3L);
	assertEquals(3L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(3L, sourceContext.getLatestElement().getTimestamp());

	// elements for partition 2
	emitRecord(fetcher, 12L, part2, 1L);
	assertEquals(12L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(12L, sourceContext.getLatestElement().getTimestamp());

	// elements for partition 3
	emitRecord(fetcher, 101L, part3, 1L);
	emitRecord(fetcher, 102L, part3, 2L);
	assertEquals(102L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(102L, sourceContext.getLatestElement().getTimestamp());

	processingTimeService.setCurrentTime(10);

	// now, we should have a watermark (this blocks until the periodic thread emitted the watermark)
	assertEquals(3L, sourceContext.getLatestWatermark().getTimestamp());

	// advance partition 3
	emitRecord(fetcher, 1003L, part3, 3L);
	emitRecord(fetcher, 1004L, part3, 4L);
	emitRecord(fetcher, 1005L, part3, 5L);
	assertEquals(1005L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(1005L, sourceContext.getLatestElement().getTimestamp());

	// advance partition 1 beyond partition 2 - this bumps the watermark
	emitRecord(fetcher, 30L, part1, 4L);
	assertEquals(30L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(30L, sourceContext.getLatestElement().getTimestamp());

	processingTimeService.setCurrentTime(20);

	// this blocks until the periodic thread emitted the watermark
	assertEquals(12L, sourceContext.getLatestWatermark().getTimestamp());

	// advance partition 2 again - this bumps the watermark
	emitRecord(fetcher, 13L, part2, 2L);
	emitRecord(fetcher, 14L, part2, 3L);
	emitRecord(fetcher, 15L, part2, 3L);

	processingTimeService.setCurrentTime(30);
	// this blocks until the periodic thread emitted the watermark
	long watermarkTs = sourceContext.getLatestWatermark().getTimestamp();
	assertTrue(watermarkTs >= 13L && watermarkTs <= 15L);
}
 

@Test
public void testPeriodicWatermarks() throws Exception {
	final String testTopic = "test topic name";
	Map<KafkaTopicPartition, Long> originalPartitions = new HashMap<>();
	originalPartitions.put(new KafkaTopicPartition(testTopic, 7), KafkaTopicPartitionStateSentinel.LATEST_OFFSET);
	originalPartitions.put(new KafkaTopicPartition(testTopic, 13), KafkaTopicPartitionStateSentinel.LATEST_OFFSET);
	originalPartitions.put(new KafkaTopicPartition(testTopic, 21), KafkaTopicPartitionStateSentinel.LATEST_OFFSET);

	TestSourceContext<Long> sourceContext = new TestSourceContext<>();

	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();

	TestFetcher<Long> fetcher = new TestFetcher<>(
			sourceContext,
			originalPartitions,
			new SerializedValue<AssignerWithPeriodicWatermarks<Long>>(new PeriodicTestExtractor()),
			null, /* punctuated watermarks assigner*/
			processingTimeService,
			10);

	final KafkaTopicPartitionState<Object> part1 = fetcher.subscribedPartitionStates().get(0);
	final KafkaTopicPartitionState<Object> part2 = fetcher.subscribedPartitionStates().get(1);
	final KafkaTopicPartitionState<Object> part3 = fetcher.subscribedPartitionStates().get(2);

	// elements generate a watermark if the timestamp is a multiple of three

	// elements for partition 1
	fetcher.emitRecord(1L, part1, 1L);
	fetcher.emitRecord(2L, part1, 2L);
	fetcher.emitRecord(3L, part1, 3L);
	assertEquals(3L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(3L, sourceContext.getLatestElement().getTimestamp());

	// elements for partition 2
	fetcher.emitRecord(12L, part2, 1L);
	assertEquals(12L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(12L, sourceContext.getLatestElement().getTimestamp());

	// elements for partition 3
	fetcher.emitRecord(101L, part3, 1L);
	fetcher.emitRecord(102L, part3, 2L);
	assertEquals(102L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(102L, sourceContext.getLatestElement().getTimestamp());

	processingTimeService.setCurrentTime(10);

	// now, we should have a watermark (this blocks until the periodic thread emitted the watermark)
	assertEquals(3L, sourceContext.getLatestWatermark().getTimestamp());

	// advance partition 3
	fetcher.emitRecord(1003L, part3, 3L);
	fetcher.emitRecord(1004L, part3, 4L);
	fetcher.emitRecord(1005L, part3, 5L);
	assertEquals(1005L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(1005L, sourceContext.getLatestElement().getTimestamp());

	// advance partition 1 beyond partition 2 - this bumps the watermark
	fetcher.emitRecord(30L, part1, 4L);
	assertEquals(30L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(30L, sourceContext.getLatestElement().getTimestamp());

	processingTimeService.setCurrentTime(20);

	// this blocks until the periodic thread emitted the watermark
	assertEquals(12L, sourceContext.getLatestWatermark().getTimestamp());

	// advance partition 2 again - this bumps the watermark
	fetcher.emitRecord(13L, part2, 2L);
	fetcher.emitRecord(14L, part2, 3L);
	fetcher.emitRecord(15L, part2, 3L);

	processingTimeService.setCurrentTime(30);
	// this blocks until the periodic thread emitted the watermark
	long watermarkTs = sourceContext.getLatestWatermark().getTimestamp();
	assertTrue(watermarkTs >= 13L && watermarkTs <= 15L);
}
 

@Test
public void testSkipCorruptedRecordWithPeriodicWatermarks() throws Exception {
	final String testTopic = "test topic name";
	Map<KafkaTopicPartition, Long> originalPartitions = new HashMap<>();
	originalPartitions.put(new KafkaTopicPartition(testTopic, 1), KafkaTopicPartitionStateSentinel.LATEST_OFFSET);

	TestSourceContext<Long> sourceContext = new TestSourceContext<>();

	TestProcessingTimeService processingTimeProvider = new TestProcessingTimeService();

	TestFetcher<Long> fetcher = new TestFetcher<>(
		sourceContext,
		originalPartitions,
		new SerializedValue<AssignerWithPeriodicWatermarks<Long>>(new PeriodicTestExtractor()), /* periodic watermark assigner */
		null, /* punctuated watermark assigner */
		processingTimeProvider,
		10);

	final KafkaTopicPartitionState<Object> partitionStateHolder = fetcher.subscribedPartitionStates().get(0);

	// elements generate a watermark if the timestamp is a multiple of three
	fetcher.emitRecord(1L, partitionStateHolder, 1L);
	fetcher.emitRecord(2L, partitionStateHolder, 2L);
	fetcher.emitRecord(3L, partitionStateHolder, 3L);
	assertEquals(3L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(3L, sourceContext.getLatestElement().getTimestamp());
	assertEquals(3L, partitionStateHolder.getOffset());

	// advance timer for watermark emitting
	processingTimeProvider.setCurrentTime(10L);
	assertTrue(sourceContext.hasWatermark());
	assertEquals(3L, sourceContext.getLatestWatermark().getTimestamp());

	// emit null record
	fetcher.emitRecord(null, partitionStateHolder, 4L);

	// no elements should have been collected
	assertEquals(3L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(3L, sourceContext.getLatestElement().getTimestamp());
	// the offset in state still should have advanced
	assertEquals(4L, partitionStateHolder.getOffset());

	// no watermarks should be collected
	processingTimeProvider.setCurrentTime(20L);
	assertFalse(sourceContext.hasWatermark());
}
 

@Test
public void testPeriodicWatermarks() throws Exception {
    String testTopic = "tp";
    Map<String, MessageId> offset = new HashMap<>();
    offset.put(topicName(testTopic, 1), MessageId.latest);
    offset.put(topicName(testTopic, 2), MessageId.latest);
    offset.put(topicName(testTopic, 3), MessageId.latest);

    TestSourceContext<Long> sourceContext = new TestSourceContext<Long>();
    TestProcessingTimeService processingTimeProvider = new TestProcessingTimeService();
    TestFetcher<Long> fetcher = new TestFetcher<>(
            sourceContext,
            offset,
            new SerializedValue<>(new PeriodicTestExtractor()), /* periodic watermark assigner */
            null,
            processingTimeProvider,
            10,
            null,
            null);

    PulsarTopicState part1 = fetcher.getSubscribedTopicStates().get(0);
    PulsarTopicState part2 = fetcher.getSubscribedTopicStates().get(1);
    PulsarTopicState part3 = fetcher.getSubscribedTopicStates().get(2);

    // elements for partition 1
    fetcher.emitRecord(1L, part1, dummyMessageId(1));
    fetcher.emitRecord(2L, part1, dummyMessageId(2));
    fetcher.emitRecord(3L, part1, dummyMessageId(3));
    assertEquals(3L, sourceContext.getLatestElement().getValue().longValue());
    assertEquals(3L, sourceContext.getLatestElement().getTimestamp());

    fetcher.emitRecord(12L, part2, dummyMessageId(1));
    assertEquals(12L, sourceContext.getLatestElement().getValue().longValue());
    assertEquals(12L, sourceContext.getLatestElement().getTimestamp());

    // element for partition 3
    fetcher.emitRecord(101L, part3, dummyMessageId(1));
    fetcher.emitRecord(102L, part3, dummyMessageId(2));
    assertEquals(102L, sourceContext.getLatestElement().getValue().longValue());
    assertEquals(102L, sourceContext.getLatestElement().getTimestamp());

    processingTimeProvider.setCurrentTime(10);

    // now, we should have a watermark (this blocks until the periodic thread emitted the watermark)
    assertEquals(3L, sourceContext.getLatestWatermark().getTimestamp());

    // advance partition 3
    fetcher.emitRecord(1003L, part3, dummyMessageId(3));
    fetcher.emitRecord(1004L, part3, dummyMessageId(4));
    fetcher.emitRecord(1005L, part3, dummyMessageId(5));
    assertEquals(1005L, sourceContext.getLatestElement().getValue().longValue());
    assertEquals(1005L, sourceContext.getLatestElement().getTimestamp());

    // advance partition 1 beyond partition 2 - this bumps the watermark
    fetcher.emitRecord(30L, part1, dummyMessageId(4));
    assertEquals(30L, sourceContext.getLatestElement().getValue().longValue());
    assertEquals(30L, sourceContext.getLatestElement().getTimestamp());

    processingTimeProvider.setCurrentTime(20);

    // this blocks until the periodic thread emitted the watermark
    assertEquals(12L, sourceContext.getLatestWatermark().getTimestamp());

    // advance partition 2 again - this bumps the watermark
    fetcher.emitRecord(13L, part2, dummyMessageId(2));
    fetcher.emitRecord(14L, part2, dummyMessageId(3));
    fetcher.emitRecord(15L, part2, dummyMessageId(4));

    processingTimeProvider.setCurrentTime(30);
    long watermarkTs = sourceContext.getLatestWatermark().getTimestamp();
    assertTrue(watermarkTs >= 13L && watermarkTs <= 15L);

}
 

@Test
public void testSkipCorruptedRecordWithPeriodicWatermarks() throws Exception {
	final String testTopic = "test topic name";
	Map<KafkaTopicPartition, Long> originalPartitions = new HashMap<>();
	originalPartitions.put(
			new KafkaTopicPartition(testTopic, 1),
			KafkaTopicPartitionStateSentinel.LATEST_OFFSET);

	TestSourceContext<Long> sourceContext = new TestSourceContext<>();

	TestProcessingTimeService processingTimeProvider = new TestProcessingTimeService();

	TestFetcher<Long> fetcher = new TestFetcher<>(
			sourceContext,
			originalPartitions,
			new SerializedValue<>(testWmStrategy),
			processingTimeProvider,
			10);

	final KafkaTopicPartitionState<Long, Object> partitionStateHolder =
			fetcher.subscribedPartitionStates().get(0);

	// elements generate a watermark if the timestamp is a multiple of three
	emitRecord(fetcher, 1L, partitionStateHolder, 1L);
	emitRecord(fetcher, 2L, partitionStateHolder, 2L);
	emitRecord(fetcher, 3L, partitionStateHolder, 3L);
	assertEquals(3L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(3L, sourceContext.getLatestElement().getTimestamp());
	assertEquals(3L, partitionStateHolder.getOffset());

	// advance timer for watermark emitting
	processingTimeProvider.setCurrentTime(10L);
	assertTrue(sourceContext.hasWatermark());
	assertEquals(3L, sourceContext.getLatestWatermark().getTimestamp());

	// emit no records
	fetcher.emitRecordsWithTimestamps(
			emptyQueue(),
			partitionStateHolder,
			4L,
			Long.MIN_VALUE);

	// no elements should have been collected
	assertEquals(3L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(3L, sourceContext.getLatestElement().getTimestamp());
	// the offset in state still should have advanced
	assertEquals(4L, partitionStateHolder.getOffset());

	// no watermarks should be collected
	processingTimeProvider.setCurrentTime(20L);
	assertFalse(sourceContext.hasWatermark());
}
 

@Test
public void testSkipCorruptedRecordWithPeriodicWatermarks() throws Exception {
	final String testTopic = "test topic name";
	Map<KafkaTopicPartition, Long> originalPartitions = new HashMap<>();
	originalPartitions.put(new KafkaTopicPartition(testTopic, 1), KafkaTopicPartitionStateSentinel.LATEST_OFFSET);

	TestSourceContext<Long> sourceContext = new TestSourceContext<>();

	TestProcessingTimeService processingTimeProvider = new TestProcessingTimeService();

	TestFetcher<Long> fetcher = new TestFetcher<>(
		sourceContext,
		originalPartitions,
		new SerializedValue<AssignerWithPeriodicWatermarks<Long>>(new PeriodicTestExtractor()), /* periodic watermark assigner */
		null, /* punctuated watermark assigner */
		processingTimeProvider,
		10);

	final KafkaTopicPartitionState<Object> partitionStateHolder = fetcher.subscribedPartitionStates().get(0);

	// elements generate a watermark if the timestamp is a multiple of three
	fetcher.emitRecord(1L, partitionStateHolder, 1L);
	fetcher.emitRecord(2L, partitionStateHolder, 2L);
	fetcher.emitRecord(3L, partitionStateHolder, 3L);
	assertEquals(3L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(3L, sourceContext.getLatestElement().getTimestamp());
	assertEquals(3L, partitionStateHolder.getOffset());

	// advance timer for watermark emitting
	processingTimeProvider.setCurrentTime(10L);
	assertTrue(sourceContext.hasWatermark());
	assertEquals(3L, sourceContext.getLatestWatermark().getTimestamp());

	// emit null record
	fetcher.emitRecord(null, partitionStateHolder, 4L);

	// no elements should have been collected
	assertEquals(3L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(3L, sourceContext.getLatestElement().getTimestamp());
	// the offset in state still should have advanced
	assertEquals(4L, partitionStateHolder.getOffset());

	// no watermarks should be collected
	processingTimeProvider.setCurrentTime(20L);
	assertFalse(sourceContext.hasWatermark());
}
 

private void testLatencyMarkEmission(int numberLatencyMarkers, OperatorSetupOperation operatorSetup) throws Exception {
	final List<StreamElement> output = new ArrayList<>();

	final TestProcessingTimeService testProcessingTimeService = new TestProcessingTimeService();
	testProcessingTimeService.setCurrentTime(0L);
	final List<Long> processingTimes = Arrays.asList(1L, 10L, 11L, 21L, maxProcessingTime);

	// regular stream source operator
	final StreamSource<Long, ProcessingTimeServiceSource> operator =
		new StreamSource<>(new ProcessingTimeServiceSource(testProcessingTimeService, processingTimes));

	operatorSetup.setupSourceOperator(operator, testProcessingTimeService);

	// run and wait to be stopped
	operator.run(new Object(), mock(StreamStatusMaintainer.class), new CollectorOutput<Long>(output));

	assertEquals(
		numberLatencyMarkers + 1, // + 1 is the final watermark element
		output.size());

	long timestamp = 0L;
	int expectedLatencyIndex = 0;

	int i = 0;
	// verify that its only latency markers + a final watermark
	for (; i < numberLatencyMarkers; i++) {
		StreamElement se = output.get(i);
		Assert.assertTrue(se.isLatencyMarker());
		Assert.assertEquals(operator.getOperatorID(), se.asLatencyMarker().getOperatorId());
		Assert.assertEquals(0, se.asLatencyMarker().getSubtaskIndex());

		// determines the next latency mark that should've been emitted
		// latency marks are emitted once per latencyMarkInterval,
		// as a result of which we never emit both 10 and 11
		while (timestamp > processingTimes.get(expectedLatencyIndex)) {
			expectedLatencyIndex++;
		}
		Assert.assertEquals(processingTimes.get(expectedLatencyIndex).longValue(), se.asLatencyMarker().getMarkedTime());

		timestamp += latencyMarkInterval;
	}

	Assert.assertTrue(output.get(i).isWatermark());
}
 

/**
 * This also verifies that we don't have leakage between keys/namespaces.
 *
 * <p>This also verifies that deleted timers don't fire.
 */
@Test
public void testDeleteProcessingTimeTimers() throws Exception {
	@SuppressWarnings("unchecked")
	Triggerable<Integer, String> mockTriggerable = mock(Triggerable.class);

	TestKeyContext keyContext = new TestKeyContext();
	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
	InternalTimerServiceImpl<Integer, String> timerService =
			createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());

	// get two different keys
	int key1 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	int key2 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	while (key2 == key1) {
		key2 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	}

	keyContext.setCurrentKey(key1);

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

	keyContext.setCurrentKey(key2);

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

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

	keyContext.setCurrentKey(key1);
	timerService.deleteProcessingTimeTimer("hello", 10);

	keyContext.setCurrentKey(key2);
	timerService.deleteProcessingTimeTimer("ciao", 10);

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

	processingTimeService.setCurrentTime(10);

	verify(mockTriggerable, times(2)).onProcessingTime(anyInternalTimer());
	verify(mockTriggerable, times(1)).onProcessingTime(eq(new TimerHeapInternalTimer<>(10, key1, "ciao")));
	verify(mockTriggerable, times(0)).onProcessingTime(eq(new TimerHeapInternalTimer<>(10, key1, "hello")));
	verify(mockTriggerable, times(0)).onProcessingTime(eq(new TimerHeapInternalTimer<>(10, key2, "ciao")));
	verify(mockTriggerable, times(1)).onProcessingTime(eq(new TimerHeapInternalTimer<>(10, key2, "hello")));

	assertEquals(0, timerService.numEventTimeTimers());
}
 

/**
 * Test scenario (idleTimeout = 100):
 * (1) Start from 0 as initial time.
 * (2) As soon as time reaches 100, status should have been toggled to IDLE.
 * (3) After some arbitrary time (until 300), the status should remain IDLE.
 * (4) Emit a record at 310. Status should become ACTIVE. This should fire a idleness detection at 410.
 * (5) Emit another record at 320 (which is before the next check). This should make the idleness check pass.
 * (6) Advance time to 410 and trigger idleness detection.
 *     The status should still be ACTIVE due to step (5). Another idleness detection should be fired at 510.
 * (7) Advance time to 510 and trigger idleness detection. Since no records were collected in-between the two
 *     idleness detections, status should have been toggle back to IDLE.
 *
 * <p>Inline comments will refer to the corresponding tested steps in the scenario.
 */
@Test
public void testManualWatermarkContext() throws Exception {
	long idleTimeout = 100;

	long initialTime = 0;
	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
	processingTimeService.setCurrentTime(initialTime);

	final List<StreamElement> output = new ArrayList<>();

	MockStreamStatusMaintainer mockStreamStatusMaintainer = new MockStreamStatusMaintainer();

	SourceFunction.SourceContext<String> context = StreamSourceContexts.getSourceContext(
		TimeCharacteristic.EventTime,
		processingTimeService,
		new Object(),
		mockStreamStatusMaintainer,
		new CollectorOutput<String>(output),
		0,
		idleTimeout);

	// -------------------------- begin test scenario --------------------------

	// corresponds to step (2) of scenario (please see method-level Javadoc comment)
	processingTimeService.setCurrentTime(initialTime + idleTimeout);
	assertTrue(mockStreamStatusMaintainer.getStreamStatus().isIdle());

	// corresponds to step (3) of scenario (please see method-level Javadoc comment)
	processingTimeService.setCurrentTime(initialTime + 2 * idleTimeout);
	processingTimeService.setCurrentTime(initialTime + 3 * idleTimeout);
	assertTrue(mockStreamStatusMaintainer.getStreamStatus().isIdle());

	// corresponds to step (4) of scenario (please see method-level Javadoc comment)
	processingTimeService.setCurrentTime(initialTime + 3 * idleTimeout + idleTimeout / 10);
	switch (testMethod) {
		case COLLECT:
			context.collect("msg");
			break;
		case COLLECT_WITH_TIMESTAMP:
			context.collectWithTimestamp("msg", processingTimeService.getCurrentProcessingTime());
			break;
		case EMIT_WATERMARK:
			context.emitWatermark(new Watermark(processingTimeService.getCurrentProcessingTime()));
			break;
	}
	assertTrue(mockStreamStatusMaintainer.getStreamStatus().isActive());

	// corresponds to step (5) of scenario (please see method-level Javadoc comment)
	processingTimeService.setCurrentTime(initialTime + 3 * idleTimeout + 2 * idleTimeout / 10);
	switch (testMethod) {
		case COLLECT:
			context.collect("msg");
			break;
		case COLLECT_WITH_TIMESTAMP:
			context.collectWithTimestamp("msg", processingTimeService.getCurrentProcessingTime());
			break;
		case EMIT_WATERMARK:
			context.emitWatermark(new Watermark(processingTimeService.getCurrentProcessingTime()));
			break;
	}
	assertTrue(mockStreamStatusMaintainer.getStreamStatus().isActive());

	// corresponds to step (6) of scenario (please see method-level Javadoc comment)
	processingTimeService.setCurrentTime(initialTime + 4 * idleTimeout + idleTimeout / 10);
	assertTrue(mockStreamStatusMaintainer.getStreamStatus().isActive());

	// corresponds to step (7) of scenario (please see method-level Javadoc comment)
	processingTimeService.setCurrentTime(initialTime + 5 * idleTimeout + idleTimeout / 10);
	assertTrue(mockStreamStatusMaintainer.getStreamStatus().isIdle());
}
 

/**
 * 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());
}
 

/**
 * This also verifies that we don't have leakage between keys/namespaces.
 */
@Test
public void testSetAndFireProcessingTimeTimers() throws Exception {
	@SuppressWarnings("unchecked")
	Triggerable<Integer, String> mockTriggerable = mock(Triggerable.class);

	TestKeyContext keyContext = new TestKeyContext();
	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();
	InternalTimerServiceImpl<Integer, String> timerService =
			createAndStartInternalTimerService(mockTriggerable, keyContext, processingTimeService, testKeyGroupRange, createQueueFactory());

	// get two different keys
	int key1 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	int key2 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	while (key2 == key1) {
		key2 = getKeyInKeyGroupRange(testKeyGroupRange, maxParallelism);
	}

	keyContext.setCurrentKey(key1);

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

	keyContext.setCurrentKey(key2);

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

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

	processingTimeService.setCurrentTime(10);

	verify(mockTriggerable, times(4)).onProcessingTime(anyInternalTimer());
	verify(mockTriggerable, times(1)).onProcessingTime(eq(new TimerHeapInternalTimer<>(10, key1, "ciao")));
	verify(mockTriggerable, times(1)).onProcessingTime(eq(new TimerHeapInternalTimer<>(10, key1, "hello")));
	verify(mockTriggerable, times(1)).onProcessingTime(eq(new TimerHeapInternalTimer<>(10, key2, "ciao")));
	verify(mockTriggerable, times(1)).onProcessingTime(eq(new TimerHeapInternalTimer<>(10, key2, "hello")));

	assertEquals(0, timerService.numProcessingTimeTimers());
}
 

@Test
public void testPeriodicWatermarks() throws Exception {
	final String testTopic = "test topic name";
	Map<KafkaTopicPartition, Long> originalPartitions = new HashMap<>();
	originalPartitions.put(new KafkaTopicPartition(testTopic, 7), KafkaTopicPartitionStateSentinel.LATEST_OFFSET);
	originalPartitions.put(new KafkaTopicPartition(testTopic, 13), KafkaTopicPartitionStateSentinel.LATEST_OFFSET);
	originalPartitions.put(new KafkaTopicPartition(testTopic, 21), KafkaTopicPartitionStateSentinel.LATEST_OFFSET);

	TestSourceContext<Long> sourceContext = new TestSourceContext<>();

	TestProcessingTimeService processingTimeService = new TestProcessingTimeService();

	TestFetcher<Long> fetcher = new TestFetcher<>(
			sourceContext,
			originalPartitions,
			new SerializedValue<AssignerWithPeriodicWatermarks<Long>>(new PeriodicTestExtractor()),
			null, /* punctuated watermarks assigner*/
			processingTimeService,
			10);

	final KafkaTopicPartitionState<Object> part1 = fetcher.subscribedPartitionStates().get(0);
	final KafkaTopicPartitionState<Object> part2 = fetcher.subscribedPartitionStates().get(1);
	final KafkaTopicPartitionState<Object> part3 = fetcher.subscribedPartitionStates().get(2);

	// elements generate a watermark if the timestamp is a multiple of three

	// elements for partition 1
	fetcher.emitRecord(1L, part1, 1L);
	fetcher.emitRecord(2L, part1, 2L);
	fetcher.emitRecord(3L, part1, 3L);
	assertEquals(3L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(3L, sourceContext.getLatestElement().getTimestamp());

	// elements for partition 2
	fetcher.emitRecord(12L, part2, 1L);
	assertEquals(12L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(12L, sourceContext.getLatestElement().getTimestamp());

	// elements for partition 3
	fetcher.emitRecord(101L, part3, 1L);
	fetcher.emitRecord(102L, part3, 2L);
	assertEquals(102L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(102L, sourceContext.getLatestElement().getTimestamp());

	processingTimeService.setCurrentTime(10);

	// now, we should have a watermark (this blocks until the periodic thread emitted the watermark)
	assertEquals(3L, sourceContext.getLatestWatermark().getTimestamp());

	// advance partition 3
	fetcher.emitRecord(1003L, part3, 3L);
	fetcher.emitRecord(1004L, part3, 4L);
	fetcher.emitRecord(1005L, part3, 5L);
	assertEquals(1005L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(1005L, sourceContext.getLatestElement().getTimestamp());

	// advance partition 1 beyond partition 2 - this bumps the watermark
	fetcher.emitRecord(30L, part1, 4L);
	assertEquals(30L, sourceContext.getLatestElement().getValue().longValue());
	assertEquals(30L, sourceContext.getLatestElement().getTimestamp());

	processingTimeService.setCurrentTime(20);

	// this blocks until the periodic thread emitted the watermark
	assertEquals(12L, sourceContext.getLatestWatermark().getTimestamp());

	// advance partition 2 again - this bumps the watermark
	fetcher.emitRecord(13L, part2, 2L);
	fetcher.emitRecord(14L, part2, 3L);
	fetcher.emitRecord(15L, part2, 3L);

	processingTimeService.setCurrentTime(30);
	// this blocks until the periodic thread emitted the watermark
	long watermarkTs = sourceContext.getLatestWatermark().getTimestamp();
	assertTrue(watermarkTs >= 13L && watermarkTs <= 15L);
}
 

@Test
public void testCustomTimeServiceProvider() throws Throwable {
	final TestProcessingTimeService tp = new TestProcessingTimeService();

	final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<>(
			(env) -> new OneInputStreamTask<>(env, tp),
			BasicTypeInfo.STRING_TYPE_INFO,
			BasicTypeInfo.STRING_TYPE_INFO);

	testHarness.setupOutputForSingletonOperatorChain();

	StreamConfig streamConfig = testHarness.getStreamConfig();

	StreamMap<String, String> mapOperator = new StreamMap<>(new StreamTaskTimerTest.DummyMapFunction<>());
	streamConfig.setStreamOperator(mapOperator);
	streamConfig.setOperatorID(new OperatorID());

	testHarness.invoke();
	testHarness.waitForTaskRunning();

	ProcessingTimeService processingTimeService = ((StreamMap<?, ?>) testHarness.getHeadOperator()).getProcessingTimeService();

	assertEquals(Long.MIN_VALUE, processingTimeService.getCurrentProcessingTime());

	tp.setCurrentTime(11);
	assertEquals(processingTimeService.getCurrentProcessingTime(), 11);

	tp.setCurrentTime(15);
	tp.setCurrentTime(16);
	assertEquals(processingTimeService.getCurrentProcessingTime(), 16);

	// register 2 tasks
	processingTimeService.registerTimer(30, timestamp -> {});

	processingTimeService.registerTimer(40, timestamp -> {});

	assertEquals(2, tp.getNumActiveTimers());

	tp.setCurrentTime(35);
	assertEquals(1, tp.getNumActiveTimers());

	tp.setCurrentTime(40);
	assertEquals(0, tp.getNumActiveTimers());

	tp.shutdownService();
}