Java Code Examples for org.apache.flink.streaming.util.AbstractStreamOperatorTestHarness#snapshot()

/**
 * Manually run this to write binary snapshot data. Remove @Ignore to run.
 */
@Ignore
@Test
public void writeMonitoringSourceSnapshot() throws Exception {

	File testFolder = tempFolder.newFolder();

	long fileModTime = Long.MIN_VALUE;
	for (int i = 0; i < 1; i++) {
		Tuple2<File, String> file = createFileAndFillWithData(testFolder, "file", i, "This is test line.");
		fileModTime = file.f0.lastModified();
	}

	TextInputFormat format = new TextInputFormat(new Path(testFolder.getAbsolutePath()));

	final ContinuousFileMonitoringFunction<String> monitoringFunction =
		new ContinuousFileMonitoringFunction<>(format, FileProcessingMode.PROCESS_CONTINUOUSLY, 1, INTERVAL);

	StreamSource<TimestampedFileInputSplit, ContinuousFileMonitoringFunction<String>> src =
		new StreamSource<>(monitoringFunction);

	final AbstractStreamOperatorTestHarness<TimestampedFileInputSplit> testHarness =
			new AbstractStreamOperatorTestHarness<>(src, 1, 1, 0);

	testHarness.open();

	final Throwable[] error = new Throwable[1];

	final OneShotLatch latch = new OneShotLatch();

	// run the source asynchronously
	Thread runner = new Thread() {
		@Override
		public void run() {
			try {
				monitoringFunction.run(new DummySourceContext() {
					@Override
					public void collect(TimestampedFileInputSplit element) {
						latch.trigger();
					}

					@Override
					public void markAsTemporarilyIdle() {

					}
				});
			}
			catch (Throwable t) {
				t.printStackTrace();
				error[0] = t;
			}
		}
	};
	runner.start();

	if (!latch.isTriggered()) {
		latch.await();
	}

	final OperatorSubtaskState snapshot;
	synchronized (testHarness.getCheckpointLock()) {
		snapshot = testHarness.snapshot(0L, 0L);
	}

	OperatorSnapshotUtil.writeStateHandle(
			snapshot,
			"src/test/resources/monitoring-function-migration-test-" + fileModTime + "-flink" + flinkGenerateSavepointVersion + "-snapshot");

	monitoringFunction.cancel();
	runner.join();

	testHarness.close();
}
 

@Test
public void testCheckpointAndRestore() {
	try {
		final int numElements = 10000;

		List<Integer> data = new ArrayList<Integer>(numElements);
		List<Integer> result = new ArrayList<Integer>(numElements);

		for (int i = 0; i < numElements; i++) {
			data.add(i);
		}

		final FromElementsFunction<Integer> source = new FromElementsFunction<>(IntSerializer.INSTANCE, data);
		StreamSource<Integer, FromElementsFunction<Integer>> src = new StreamSource<>(source);
		AbstractStreamOperatorTestHarness<Integer> testHarness =
			new AbstractStreamOperatorTestHarness<>(src, 1, 1, 0);
		testHarness.open();

		final SourceFunction.SourceContext<Integer> ctx = new ListSourceContext<Integer>(result, 2L);

		final Throwable[] error = new Throwable[1];

		// run the source asynchronously
		Thread runner = new Thread() {
			@Override
			public void run() {
				try {
					source.run(ctx);
				}
				catch (Throwable t) {
					error[0] = t;
				}
			}
		};
		runner.start();

		// wait for a bit
		Thread.sleep(1000);

		// make a checkpoint
		List<Integer> checkpointData = new ArrayList<>(numElements);
		OperatorSubtaskState handles = null;
		synchronized (ctx.getCheckpointLock()) {
			handles = testHarness.snapshot(566, System.currentTimeMillis());
			checkpointData.addAll(result);
		}

		// cancel the source
		source.cancel();
		runner.join();

		// check for errors
		if (error[0] != null) {
			System.err.println("Error in asynchronous source runner");
			error[0].printStackTrace();
			fail("Error in asynchronous source runner");
		}

		final FromElementsFunction<Integer> sourceCopy = new FromElementsFunction<>(IntSerializer.INSTANCE, data);
		StreamSource<Integer, FromElementsFunction<Integer>> srcCopy = new StreamSource<>(sourceCopy);
		AbstractStreamOperatorTestHarness<Integer> testHarnessCopy =
			new AbstractStreamOperatorTestHarness<>(srcCopy, 1, 1, 0);
		testHarnessCopy.setup();
		testHarnessCopy.initializeState(handles);
		testHarnessCopy.open();

		// recovery run
		SourceFunction.SourceContext<Integer> newCtx = new ListSourceContext<>(checkpointData);

		sourceCopy.run(newCtx);

		assertEquals(data, checkpointData);
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

@Test
public void testFunctionRestore() throws Exception {
	String testBasePath = hdfsURI + "/" + UUID.randomUUID() + "/";

	org.apache.hadoop.fs.Path path = null;
	long fileModTime = Long.MIN_VALUE;
	for (int i = 0; i < 1; i++) {
		Tuple2<org.apache.hadoop.fs.Path, String> file = createFileAndFillWithData(testBasePath, "file", i, "This is test line.");
		path = file.f0;
		fileModTime = hdfs.getFileStatus(file.f0).getModificationTime();
	}

	TextInputFormat format = new TextInputFormat(new Path(testBasePath));

	final ContinuousFileMonitoringFunction<String> monitoringFunction =
		createTestContinuousFileMonitoringFunction(format, FileProcessingMode.PROCESS_CONTINUOUSLY);

	StreamSource<TimestampedFileInputSplit, ContinuousFileMonitoringFunction<String>> src =
		new StreamSource<>(monitoringFunction);

	final AbstractStreamOperatorTestHarness<TimestampedFileInputSplit> testHarness =
		new AbstractStreamOperatorTestHarness<>(src, 1, 1, 0);
	testHarness.open();

	final Throwable[] error = new Throwable[1];

	final OneShotLatch latch = new OneShotLatch();

	final DummySourceContext sourceContext = new DummySourceContext() {
		@Override
		public void collect(TimestampedFileInputSplit element) {
			latch.trigger();
		}
	};

	// run the source asynchronously
	Thread runner = new Thread() {
		@Override
		public void run() {
			try {
				monitoringFunction.run(sourceContext);
			}
			catch (Throwable t) {
				t.printStackTrace();
				error[0] = t;
			}
		}
	};
	runner.start();

	// first condition for the source to have updated its state: emit at least one element
	if (!latch.isTriggered()) {
		latch.await();
	}

	// second condition for the source to have updated its state: it's not on the lock anymore,
	// this means it has processed all the splits and updated its state.
	synchronized (sourceContext.getCheckpointLock()) {}

	OperatorSubtaskState snapshot = testHarness.snapshot(0, 0);
	monitoringFunction.cancel();
	runner.join();

	testHarness.close();

	final ContinuousFileMonitoringFunction<String> monitoringFunctionCopy =
		createTestContinuousFileMonitoringFunction(format, FileProcessingMode.PROCESS_CONTINUOUSLY);

	StreamSource<TimestampedFileInputSplit, ContinuousFileMonitoringFunction<String>> srcCopy =
		new StreamSource<>(monitoringFunctionCopy);

	AbstractStreamOperatorTestHarness<TimestampedFileInputSplit> testHarnessCopy =
		new AbstractStreamOperatorTestHarness<>(srcCopy, 1, 1, 0);
	testHarnessCopy.initializeState(snapshot);
	testHarnessCopy.open();

	Assert.assertNull(error[0]);
	Assert.assertEquals(fileModTime, monitoringFunctionCopy.getGlobalModificationTime());

	hdfs.delete(path, false);
}
 

/**
 * Manually run this to write binary snapshot data. Remove @Ignore to run.
 */
@Ignore
@Test
public void writeMonitoringSourceSnapshot() throws Exception {

	File testFolder = tempFolder.newFolder();

	long fileModTime = Long.MIN_VALUE;
	for (int i = 0; i < 1; i++) {
		Tuple2<File, String> file = createFileAndFillWithData(testFolder, "file", i, "This is test line.");
		fileModTime = file.f0.lastModified();
	}

	TextInputFormat format = new TextInputFormat(new Path(testFolder.getAbsolutePath()));

	final ContinuousFileMonitoringFunction<String> monitoringFunction =
		new ContinuousFileMonitoringFunction<>(format, FileProcessingMode.PROCESS_CONTINUOUSLY, 1, INTERVAL);

	StreamSource<TimestampedFileInputSplit, ContinuousFileMonitoringFunction<String>> src =
		new StreamSource<>(monitoringFunction);

	final AbstractStreamOperatorTestHarness<TimestampedFileInputSplit> testHarness =
			new AbstractStreamOperatorTestHarness<>(src, 1, 1, 0);

	testHarness.open();

	final Throwable[] error = new Throwable[1];

	final OneShotLatch latch = new OneShotLatch();

	// run the source asynchronously
	Thread runner = new Thread() {
		@Override
		public void run() {
			try {
				monitoringFunction.run(new DummySourceContext() {
					@Override
					public void collect(TimestampedFileInputSplit element) {
						latch.trigger();
					}

					@Override
					public void markAsTemporarilyIdle() {

					}
				});
			}
			catch (Throwable t) {
				t.printStackTrace();
				error[0] = t;
			}
		}
	};
	runner.start();

	if (!latch.isTriggered()) {
		latch.await();
	}

	final OperatorSubtaskState snapshot;
	synchronized (testHarness.getCheckpointLock()) {
		snapshot = testHarness.snapshot(0L, 0L);
	}

	OperatorSnapshotUtil.writeStateHandle(
			snapshot,
			"src/test/resources/monitoring-function-migration-test-" + fileModTime + "-flink" + flinkGenerateSavepointVersion + "-snapshot");

	monitoringFunction.cancel();
	runner.join();

	testHarness.close();
}
 

@Test
public void testCheckpointing() throws Exception {
	source.autoAck = false;

	StreamSource<String, RMQSource<String>> src = new StreamSource<>(source);
	AbstractStreamOperatorTestHarness<String> testHarness =
		new AbstractStreamOperatorTestHarness<>(src, 1, 1, 0);
	testHarness.open();

	sourceThread.start();

	Thread.sleep(5);

	final Random random = new Random(System.currentTimeMillis());
	int numSnapshots = 50;
	long previousSnapshotId;
	long lastSnapshotId = 0;

	long totalNumberOfAcks = 0;

	for (int i = 0; i < numSnapshots; i++) {
		long snapshotId = random.nextLong();
		OperatorSubtaskState data;

		synchronized (DummySourceContext.lock) {
			data = testHarness.snapshot(snapshotId, System.currentTimeMillis());
			previousSnapshotId = lastSnapshotId;
			lastSnapshotId = messageId;
		}
		// let some time pass
		Thread.sleep(5);

		// check if the correct number of messages have been snapshotted
		final long numIds = lastSnapshotId - previousSnapshotId;

		RMQTestSource sourceCopy = new RMQTestSource();
		StreamSource<String, RMQTestSource> srcCopy = new StreamSource<>(sourceCopy);
		AbstractStreamOperatorTestHarness<String> testHarnessCopy =
			new AbstractStreamOperatorTestHarness<>(srcCopy, 1, 1, 0);

		testHarnessCopy.setup();
		testHarnessCopy.initializeState(data);
		testHarnessCopy.open();

		ArrayDeque<Tuple2<Long, Set<String>>> deque = sourceCopy.getRestoredState();
		Set<String> messageIds = deque.getLast().f1;

		assertEquals(numIds, messageIds.size());
		if (messageIds.size() > 0) {
			assertTrue(messageIds.contains(Long.toString(lastSnapshotId)));
		}

		// check if the messages are being acknowledged and the transaction committed
		synchronized (DummySourceContext.lock) {
			source.notifyCheckpointComplete(snapshotId);
		}
		totalNumberOfAcks += numIds;

	}

	Mockito.verify(source.channel, Mockito.times((int) totalNumberOfAcks)).basicAck(Mockito.anyLong(), Mockito.eq(false));
	Mockito.verify(source.channel, Mockito.times(numSnapshots)).txCommit();

}
 

@SuppressWarnings("unchecked")
private void writeSnapshot(String path, HashMap<StreamShardMetadata, SequenceNumber> state) throws Exception {
	final List<StreamShardHandle> initialDiscoveryShards = new ArrayList<>(state.size());
	for (StreamShardMetadata shardMetadata : state.keySet()) {
		Shard shard = new Shard();
		shard.setShardId(shardMetadata.getShardId());

		SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
		sequenceNumberRange.withStartingSequenceNumber("1");
		shard.setSequenceNumberRange(sequenceNumberRange);

		initialDiscoveryShards.add(new StreamShardHandle(shardMetadata.getStreamName(), shard));
	}

	final TestFetcher<String> fetcher = new TestFetcher<>(
		Collections.singletonList(TEST_STREAM_NAME),
		new TestSourceContext<>(),
		new TestRuntimeContext(true, 1, 0),
		TestUtils.getStandardProperties(),
		new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()),
		state,
		initialDiscoveryShards);

	final DummyFlinkKinesisConsumer<String> consumer = new DummyFlinkKinesisConsumer<>(
		fetcher, new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()));

	StreamSource<String, DummyFlinkKinesisConsumer<String>> consumerOperator = new StreamSource<>(consumer);

	final AbstractStreamOperatorTestHarness<String> testHarness =
			new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

	testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	testHarness.setup();
	testHarness.open();

	final AtomicReference<Throwable> error = new AtomicReference<>();

	// run the source asynchronously
	Thread runner = new Thread() {
		@Override
		public void run() {
			try {
				consumer.run(new TestSourceContext<>());
			} catch (Throwable t) {
				t.printStackTrace();
				error.set(t);
			}
		}
	};
	runner.start();

	fetcher.waitUntilRun();

	final OperatorSubtaskState snapshot;
	synchronized (testHarness.getCheckpointLock()) {
		snapshot = testHarness.snapshot(0L, 0L);
	}

	OperatorSnapshotUtil.writeStateHandle(snapshot, path);

	consumerOperator.close();
	runner.join();
}
 

private void writeSnapshot(String path, HashMap<KafkaTopicPartition, Long> state) throws Exception {

		final OneShotLatch latch = new OneShotLatch();
		final AbstractFetcher<String, ?> fetcher = mock(AbstractFetcher.class);

		doAnswer(new Answer<Void>() {
			@Override
			public Void answer(InvocationOnMock invocation) throws Throwable {
				latch.trigger();
				return null;
			}
		}).when(fetcher).runFetchLoop();

		when(fetcher.snapshotCurrentState()).thenReturn(state);

		final List<KafkaTopicPartition> partitions = new ArrayList<>(PARTITION_STATE.keySet());

		final DummyFlinkKafkaConsumer<String> consumerFunction =
			new DummyFlinkKafkaConsumer<>(fetcher, TOPICS, partitions, FlinkKafkaConsumerBase.PARTITION_DISCOVERY_DISABLED);

		StreamSource<String, DummyFlinkKafkaConsumer<String>> consumerOperator =
				new StreamSource<>(consumerFunction);

		final AbstractStreamOperatorTestHarness<String> testHarness =
				new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

		testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

		testHarness.setup();
		testHarness.open();

		final Throwable[] error = new Throwable[1];

		// run the source asynchronously
		Thread runner = new Thread() {
			@Override
			public void run() {
				try {
					consumerFunction.run(new DummySourceContext() {
						@Override
						public void collect(String element) {

						}
					});
				}
				catch (Throwable t) {
					t.printStackTrace();
					error[0] = t;
				}
			}
		};
		runner.start();

		if (!latch.isTriggered()) {
			latch.await();
		}

		final OperatorSubtaskState snapshot;
		synchronized (testHarness.getCheckpointLock()) {
			snapshot = testHarness.snapshot(0L, 0L);
		}

		OperatorSnapshotUtil.writeStateHandle(snapshot, path);

		consumerOperator.close();
		runner.join();
	}
 

@Test
public void testNullCheckpoint() throws Exception {
  final int numElements = 20;
  PipelineOptions options = PipelineOptionsFactory.create();

  TestCountingSource source =
      new TestCountingSource(numElements) {
        @Override
        public Coder<CounterMark> getCheckpointMarkCoder() {
          return null;
        }
      };

  UnboundedSourceWrapper<KV<Integer, Integer>, TestCountingSource.CounterMark> flinkWrapper =
      new UnboundedSourceWrapper<>("stepName", options, source, numSplits);

  StreamSource<
          WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>,
          UnboundedSourceWrapper<KV<Integer, Integer>, TestCountingSource.CounterMark>>
      sourceOperator = new StreamSource<>(flinkWrapper);

  AbstractStreamOperatorTestHarness<WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>>
      testHarness =
          new AbstractStreamOperatorTestHarness<>(
              sourceOperator,
              numTasks /* max parallelism */,
              numTasks /* parallelism */,
              0 /* subtask index */);

  testHarness.setTimeCharacteristic(TimeCharacteristic.EventTime);

  testHarness.open();

  OperatorSubtaskState snapshot = testHarness.snapshot(0, 0);

  UnboundedSourceWrapper<KV<Integer, Integer>, TestCountingSource.CounterMark>
      restoredFlinkWrapper =
          new UnboundedSourceWrapper<>(
              "stepName", options, new TestCountingSource(numElements), numSplits);

  StreamSource<
          WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>,
          UnboundedSourceWrapper<KV<Integer, Integer>, TestCountingSource.CounterMark>>
      restoredSourceOperator = new StreamSource<>(restoredFlinkWrapper);

  // set parallelism to 1 to ensure that our testing operator gets all checkpointed state
  AbstractStreamOperatorTestHarness<WindowedValue<ValueWithRecordId<KV<Integer, Integer>>>>
      restoredTestHarness =
          new AbstractStreamOperatorTestHarness<>(
              restoredSourceOperator,
              numTasks /* max parallelism */,
              1 /* parallelism */,
              0 /* subtask index */);

  restoredTestHarness.setup();
  restoredTestHarness.initializeState(snapshot);
  restoredTestHarness.open();

  // when the source checkpointed a null we don't re-initialize the splits, that is we
  // will have no splits.
  assertEquals(0, restoredFlinkWrapper.getLocalSplitSources().size());
}
 

@Test
public void testCheckpointAndRestore() {
	try {
		final int numElements = 10000;

		List<Integer> data = new ArrayList<Integer>(numElements);
		List<Integer> result = new ArrayList<Integer>(numElements);

		for (int i = 0; i < numElements; i++) {
			data.add(i);
		}

		final FromElementsFunction<Integer> source = new FromElementsFunction<>(IntSerializer.INSTANCE, data);
		StreamSource<Integer, FromElementsFunction<Integer>> src = new StreamSource<>(source);
		AbstractStreamOperatorTestHarness<Integer> testHarness =
			new AbstractStreamOperatorTestHarness<>(src, 1, 1, 0);
		testHarness.open();

		final SourceFunction.SourceContext<Integer> ctx = new ListSourceContext<Integer>(result, 2L);

		final Throwable[] error = new Throwable[1];

		// run the source asynchronously
		Thread runner = new Thread() {
			@Override
			public void run() {
				try {
					source.run(ctx);
				}
				catch (Throwable t) {
					error[0] = t;
				}
			}
		};
		runner.start();

		// wait for a bit
		Thread.sleep(1000);

		// make a checkpoint
		List<Integer> checkpointData = new ArrayList<>(numElements);
		OperatorSubtaskState handles = null;
		synchronized (ctx.getCheckpointLock()) {
			handles = testHarness.snapshot(566, System.currentTimeMillis());
			checkpointData.addAll(result);
		}

		// cancel the source
		source.cancel();
		runner.join();

		// check for errors
		if (error[0] != null) {
			System.err.println("Error in asynchronous source runner");
			error[0].printStackTrace();
			fail("Error in asynchronous source runner");
		}

		final FromElementsFunction<Integer> sourceCopy = new FromElementsFunction<>(IntSerializer.INSTANCE, data);
		StreamSource<Integer, FromElementsFunction<Integer>> srcCopy = new StreamSource<>(sourceCopy);
		AbstractStreamOperatorTestHarness<Integer> testHarnessCopy =
			new AbstractStreamOperatorTestHarness<>(srcCopy, 1, 1, 0);
		testHarnessCopy.setup();
		testHarnessCopy.initializeState(handles);
		testHarnessCopy.open();

		// recovery run
		SourceFunction.SourceContext<Integer> newCtx = new ListSourceContext<>(checkpointData);

		sourceCopy.run(newCtx);

		assertEquals(data, checkpointData);
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

@Test
public void testFunctionRestore() throws Exception {
	String testBasePath = hdfsURI + "/" + UUID.randomUUID() + "/";

	org.apache.hadoop.fs.Path path = null;
	long fileModTime = Long.MIN_VALUE;
	for (int i = 0; i < 1; i++) {
		Tuple2<org.apache.hadoop.fs.Path, String> file = createFileAndFillWithData(testBasePath, "file", i, "This is test line.");
		path = file.f0;
		fileModTime = hdfs.getFileStatus(file.f0).getModificationTime();
	}

	TextInputFormat format = new TextInputFormat(new Path(testBasePath));

	final ContinuousFileMonitoringFunction<String> monitoringFunction =
		createTestContinuousFileMonitoringFunction(format, FileProcessingMode.PROCESS_CONTINUOUSLY);

	StreamSource<TimestampedFileInputSplit, ContinuousFileMonitoringFunction<String>> src =
		new StreamSource<>(monitoringFunction);

	final AbstractStreamOperatorTestHarness<TimestampedFileInputSplit> testHarness =
		new AbstractStreamOperatorTestHarness<>(src, 1, 1, 0);
	testHarness.open();

	final Throwable[] error = new Throwable[1];

	final OneShotLatch latch = new OneShotLatch();

	final DummySourceContext sourceContext = new DummySourceContext() {
		@Override
		public void collect(TimestampedFileInputSplit element) {
			latch.trigger();
		}
	};

	// run the source asynchronously
	Thread runner = new Thread() {
		@Override
		public void run() {
			try {
				monitoringFunction.run(sourceContext);
			}
			catch (Throwable t) {
				t.printStackTrace();
				error[0] = t;
			}
		}
	};
	runner.start();

	// first condition for the source to have updated its state: emit at least one element
	if (!latch.isTriggered()) {
		latch.await();
	}

	// second condition for the source to have updated its state: it's not on the lock anymore,
	// this means it has processed all the splits and updated its state.
	synchronized (sourceContext.getCheckpointLock()) {}

	OperatorSubtaskState snapshot = testHarness.snapshot(0, 0);
	monitoringFunction.cancel();
	runner.join();

	testHarness.close();

	final ContinuousFileMonitoringFunction<String> monitoringFunctionCopy =
		createTestContinuousFileMonitoringFunction(format, FileProcessingMode.PROCESS_CONTINUOUSLY);

	StreamSource<TimestampedFileInputSplit, ContinuousFileMonitoringFunction<String>> srcCopy =
		new StreamSource<>(monitoringFunctionCopy);

	AbstractStreamOperatorTestHarness<TimestampedFileInputSplit> testHarnessCopy =
		new AbstractStreamOperatorTestHarness<>(srcCopy, 1, 1, 0);
	testHarnessCopy.initializeState(snapshot);
	testHarnessCopy.open();

	Assert.assertNull(error[0]);
	Assert.assertEquals(fileModTime, monitoringFunctionCopy.getGlobalModificationTime());

	hdfs.delete(path, false);
}
 

private void writeSnapshot(String path, HashMap<KafkaTopicPartition, Long> state) throws Exception {

		final OneShotLatch latch = new OneShotLatch();
		final AbstractFetcher<String, ?> fetcher = mock(AbstractFetcher.class);

		doAnswer(new Answer<Void>() {
			@Override
			public Void answer(InvocationOnMock invocation) throws Throwable {
				latch.trigger();
				return null;
			}
		}).when(fetcher).runFetchLoop();

		when(fetcher.snapshotCurrentState()).thenReturn(state);

		final List<KafkaTopicPartition> partitions = new ArrayList<>(PARTITION_STATE.keySet());

		final DummyFlinkKafkaConsumer<String> consumerFunction =
			new DummyFlinkKafkaConsumer<>(fetcher, TOPICS, partitions, FlinkKafkaConsumerBase.PARTITION_DISCOVERY_DISABLED);

		StreamSource<String, DummyFlinkKafkaConsumer<String>> consumerOperator =
				new StreamSource<>(consumerFunction);

		final AbstractStreamOperatorTestHarness<String> testHarness =
				new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

		testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

		testHarness.setup();
		testHarness.open();

		final Throwable[] error = new Throwable[1];

		// run the source asynchronously
		Thread runner = new Thread() {
			@Override
			public void run() {
				try {
					consumerFunction.run(new DummySourceContext() {
						@Override
						public void collect(String element) {

						}
					});
				}
				catch (Throwable t) {
					t.printStackTrace();
					error[0] = t;
				}
			}
		};
		runner.start();

		if (!latch.isTriggered()) {
			latch.await();
		}

		final OperatorSubtaskState snapshot;
		synchronized (testHarness.getCheckpointLock()) {
			snapshot = testHarness.snapshot(0L, 0L);
		}

		OperatorSnapshotUtil.writeStateHandle(snapshot, path);

		consumerOperator.close();
		runner.join();
	}
 

@Test
public void testCheckpointing() throws Exception {
	source.autoAck = false;

	StreamSource<String, RMQSource<String>> src = new StreamSource<>(source);
	AbstractStreamOperatorTestHarness<String> testHarness =
		new AbstractStreamOperatorTestHarness<>(src, 1, 1, 0);
	testHarness.open();

	sourceThread.start();

	Thread.sleep(5);

	final Random random = new Random(System.currentTimeMillis());
	int numSnapshots = 50;
	long previousSnapshotId;
	long lastSnapshotId = 0;

	long totalNumberOfAcks = 0;

	for (int i = 0; i < numSnapshots; i++) {
		long snapshotId = random.nextLong();
		OperatorSubtaskState data;

		synchronized (DummySourceContext.lock) {
			data = testHarness.snapshot(snapshotId, System.currentTimeMillis());
			previousSnapshotId = lastSnapshotId;
			lastSnapshotId = messageId;
		}
		// let some time pass
		Thread.sleep(5);

		// check if the correct number of messages have been snapshotted
		final long numIds = lastSnapshotId - previousSnapshotId;

		RMQTestSource sourceCopy = new RMQTestSource();
		StreamSource<String, RMQTestSource> srcCopy = new StreamSource<>(sourceCopy);
		AbstractStreamOperatorTestHarness<String> testHarnessCopy =
			new AbstractStreamOperatorTestHarness<>(srcCopy, 1, 1, 0);

		testHarnessCopy.setup();
		testHarnessCopy.initializeState(data);
		testHarnessCopy.open();

		ArrayDeque<Tuple2<Long, Set<String>>> deque = sourceCopy.getRestoredState();
		Set<String> messageIds = deque.getLast().f1;

		assertEquals(numIds, messageIds.size());
		if (messageIds.size() > 0) {
			assertTrue(messageIds.contains(Long.toString(lastSnapshotId)));
		}

		// check if the messages are being acknowledged and the transaction committed
		synchronized (DummySourceContext.lock) {
			source.notifyCheckpointComplete(snapshotId);
		}
		totalNumberOfAcks += numIds;

	}

	Mockito.verify(source.channel, Mockito.times((int) totalNumberOfAcks)).basicAck(Mockito.anyLong(), Mockito.eq(false));
	Mockito.verify(source.channel, Mockito.times(numSnapshots)).txCommit();

}
 

@SuppressWarnings("unchecked")
private void writeSnapshot(String path, HashMap<StreamShardMetadata, SequenceNumber> state) throws Exception {
	final List<StreamShardHandle> initialDiscoveryShards = new ArrayList<>(state.size());
	for (StreamShardMetadata shardMetadata : state.keySet()) {
		Shard shard = new Shard();
		shard.setShardId(shardMetadata.getShardId());

		SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
		sequenceNumberRange.withStartingSequenceNumber("1");
		shard.setSequenceNumberRange(sequenceNumberRange);

		initialDiscoveryShards.add(new StreamShardHandle(shardMetadata.getStreamName(), shard));
	}

	final TestFetcher<String> fetcher = new TestFetcher<>(
		Collections.singletonList(TEST_STREAM_NAME),
		new TestSourceContext<>(),
		new TestRuntimeContext(true, 1, 0),
		TestUtils.getStandardProperties(),
		new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()),
		state,
		initialDiscoveryShards);

	final DummyFlinkKinesisConsumer<String> consumer = new DummyFlinkKinesisConsumer<>(
		fetcher, new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()));

	StreamSource<String, DummyFlinkKinesisConsumer<String>> consumerOperator = new StreamSource<>(consumer);

	final AbstractStreamOperatorTestHarness<String> testHarness =
			new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

	testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	testHarness.setup();
	testHarness.open();

	final AtomicReference<Throwable> error = new AtomicReference<>();

	// run the source asynchronously
	Thread runner = new Thread() {
		@Override
		public void run() {
			try {
				consumer.run(new TestSourceContext<>());
			} catch (Throwable t) {
				t.printStackTrace();
				error.set(t);
			}
		}
	};
	runner.start();

	fetcher.waitUntilRun();

	final OperatorSubtaskState snapshot;
	synchronized (testHarness.getCheckpointLock()) {
		snapshot = testHarness.snapshot(0L, 0L);
	}

	OperatorSnapshotUtil.writeStateHandle(snapshot, path);

	consumerOperator.close();
	runner.join();
}
 

private void writeSnapshot(String path, HashMap<KafkaTopicPartition, Long> state) throws Exception {

		final OneShotLatch latch = new OneShotLatch();
		final AbstractFetcher<String, ?> fetcher = mock(AbstractFetcher.class);

		doAnswer(new Answer<Void>() {
			@Override
			public Void answer(InvocationOnMock invocation) throws Throwable {
				latch.trigger();
				return null;
			}
		}).when(fetcher).runFetchLoop();

		when(fetcher.snapshotCurrentState()).thenReturn(state);

		final List<KafkaTopicPartition> partitions = new ArrayList<>(PARTITION_STATE.keySet());

		final DummyFlinkKafkaConsumer<String> consumerFunction =
			new DummyFlinkKafkaConsumer<>(fetcher, TOPICS, partitions, FlinkKafkaConsumerBase.PARTITION_DISCOVERY_DISABLED);

		StreamSource<String, DummyFlinkKafkaConsumer<String>> consumerOperator =
				new StreamSource<>(consumerFunction);

		final AbstractStreamOperatorTestHarness<String> testHarness =
				new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

		testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

		testHarness.setup();
		testHarness.open();

		final Throwable[] error = new Throwable[1];

		// run the source asynchronously
		Thread runner = new Thread() {
			@Override
			public void run() {
				try {
					consumerFunction.run(new DummySourceContext() {
						@Override
						public void collect(String element) {

						}
					});
				}
				catch (Throwable t) {
					t.printStackTrace();
					error[0] = t;
				}
			}
		};
		runner.start();

		if (!latch.isTriggered()) {
			latch.await();
		}

		final OperatorSubtaskState snapshot;
		synchronized (testHarness.getCheckpointLock()) {
			snapshot = testHarness.snapshot(0L, 0L);
		}

		OperatorSnapshotUtil.writeStateHandle(snapshot, path);

		consumerOperator.close();
		runner.join();
	}
 

@Test
public void testCheckpointAndRestore() {
	try {
		final int numElements = 10000;

		List<Integer> data = new ArrayList<Integer>(numElements);
		List<Integer> result = new ArrayList<Integer>(numElements);

		for (int i = 0; i < numElements; i++) {
			data.add(i);
		}

		final FromElementsFunction<Integer> source = new FromElementsFunction<>(IntSerializer.INSTANCE, data);
		StreamSource<Integer, FromElementsFunction<Integer>> src = new StreamSource<>(source);
		AbstractStreamOperatorTestHarness<Integer> testHarness =
			new AbstractStreamOperatorTestHarness<>(src, 1, 1, 0);
		testHarness.open();

		final SourceFunction.SourceContext<Integer> ctx = new ListSourceContext<Integer>(result, 2L);

		final Throwable[] error = new Throwable[1];

		// run the source asynchronously
		Thread runner = new Thread() {
			@Override
			public void run() {
				try {
					source.run(ctx);
				}
				catch (Throwable t) {
					error[0] = t;
				}
			}
		};
		runner.start();

		// wait for a bit
		Thread.sleep(1000);

		// make a checkpoint
		List<Integer> checkpointData = new ArrayList<>(numElements);
		OperatorSubtaskState handles = null;
		synchronized (ctx.getCheckpointLock()) {
			handles = testHarness.snapshot(566, System.currentTimeMillis());
			checkpointData.addAll(result);
		}

		// cancel the source
		source.cancel();
		runner.join();

		// check for errors
		if (error[0] != null) {
			System.err.println("Error in asynchronous source runner");
			error[0].printStackTrace();
			fail("Error in asynchronous source runner");
		}

		final FromElementsFunction<Integer> sourceCopy = new FromElementsFunction<>(IntSerializer.INSTANCE, data);
		StreamSource<Integer, FromElementsFunction<Integer>> srcCopy = new StreamSource<>(sourceCopy);
		AbstractStreamOperatorTestHarness<Integer> testHarnessCopy =
			new AbstractStreamOperatorTestHarness<>(srcCopy, 1, 1, 0);
		testHarnessCopy.setup();
		testHarnessCopy.initializeState(handles);
		testHarnessCopy.open();

		// recovery run
		SourceFunction.SourceContext<Integer> newCtx = new ListSourceContext<>(checkpointData);

		sourceCopy.run(newCtx);

		assertEquals(data, checkpointData);
	}
	catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

@Test
public void testFunctionRestore() throws Exception {
	String testBasePath = hdfsURI + "/" + UUID.randomUUID() + "/";

	org.apache.hadoop.fs.Path path = null;
	long fileModTime = Long.MIN_VALUE;
	for (int i = 0; i < 1; i++) {
		Tuple2<org.apache.hadoop.fs.Path, String> file = createFileAndFillWithData(testBasePath, "file", i, "This is test line.");
		path = file.f0;
		fileModTime = hdfs.getFileStatus(file.f0).getModificationTime();
	}

	TextInputFormat format = new TextInputFormat(new Path(testBasePath));

	final ContinuousFileMonitoringFunction<String> monitoringFunction =
		createTestContinuousFileMonitoringFunction(format, FileProcessingMode.PROCESS_CONTINUOUSLY);

	StreamSource<TimestampedFileInputSplit, ContinuousFileMonitoringFunction<String>> src =
		new StreamSource<>(monitoringFunction);

	final AbstractStreamOperatorTestHarness<TimestampedFileInputSplit> testHarness =
		new AbstractStreamOperatorTestHarness<>(src, 1, 1, 0);
	testHarness.open();

	final Throwable[] error = new Throwable[1];

	final OneShotLatch latch = new OneShotLatch();

	final DummySourceContext sourceContext = new DummySourceContext() {
		@Override
		public void collect(TimestampedFileInputSplit element) {
			latch.trigger();
		}
	};

	// run the source asynchronously
	Thread runner = new Thread() {
		@Override
		public void run() {
			try {
				monitoringFunction.run(sourceContext);
			}
			catch (Throwable t) {
				t.printStackTrace();
				error[0] = t;
			}
		}
	};
	runner.start();

	// first condition for the source to have updated its state: emit at least one element
	if (!latch.isTriggered()) {
		latch.await();
	}

	// second condition for the source to have updated its state: it's not on the lock anymore,
	// this means it has processed all the splits and updated its state.
	synchronized (sourceContext.getCheckpointLock()) {}

	OperatorSubtaskState snapshot = testHarness.snapshot(0, 0);
	monitoringFunction.cancel();
	runner.join();

	testHarness.close();

	final ContinuousFileMonitoringFunction<String> monitoringFunctionCopy =
		createTestContinuousFileMonitoringFunction(format, FileProcessingMode.PROCESS_CONTINUOUSLY);

	StreamSource<TimestampedFileInputSplit, ContinuousFileMonitoringFunction<String>> srcCopy =
		new StreamSource<>(monitoringFunctionCopy);

	AbstractStreamOperatorTestHarness<TimestampedFileInputSplit> testHarnessCopy =
		new AbstractStreamOperatorTestHarness<>(srcCopy, 1, 1, 0);
	testHarnessCopy.initializeState(snapshot);
	testHarnessCopy.open();

	Assert.assertNull(error[0]);
	Assert.assertEquals(fileModTime, monitoringFunctionCopy.getGlobalModificationTime());

	hdfs.delete(path, false);
}
 

/**
 * Manually run this to write binary snapshot data. Remove @Ignore to run.
 */
@Ignore
@Test
public void writeMonitoringSourceSnapshot() throws Exception {

	File testFolder = tempFolder.newFolder();

	long fileModTime = Long.MIN_VALUE;
	for (int i = 0; i < 1; i++) {
		Tuple2<File, String> file = createFileAndFillWithData(testFolder, "file", i, "This is test line.");
		fileModTime = file.f0.lastModified();
	}

	TextInputFormat format = new TextInputFormat(new Path(testFolder.getAbsolutePath()));

	final ContinuousFileMonitoringFunction<String> monitoringFunction =
		new ContinuousFileMonitoringFunction<>(format, FileProcessingMode.PROCESS_CONTINUOUSLY, 1, INTERVAL);

	StreamSource<TimestampedFileInputSplit, ContinuousFileMonitoringFunction<String>> src =
		new StreamSource<>(monitoringFunction);

	final AbstractStreamOperatorTestHarness<TimestampedFileInputSplit> testHarness =
			new AbstractStreamOperatorTestHarness<>(src, 1, 1, 0);

	testHarness.open();

	final Throwable[] error = new Throwable[1];

	final OneShotLatch latch = new OneShotLatch();

	// run the source asynchronously
	Thread runner = new Thread() {
		@Override
		public void run() {
			try {
				monitoringFunction.run(new DummySourceContext() {
					@Override
					public void collect(TimestampedFileInputSplit element) {
						latch.trigger();
					}

					@Override
					public void markAsTemporarilyIdle() {

					}
				});
			}
			catch (Throwable t) {
				t.printStackTrace();
				error[0] = t;
			}
		}
	};
	runner.start();

	if (!latch.isTriggered()) {
		latch.await();
	}

	final OperatorSubtaskState snapshot;
	synchronized (testHarness.getCheckpointLock()) {
		snapshot = testHarness.snapshot(0L, 0L);
	}

	OperatorSnapshotUtil.writeStateHandle(
			snapshot,
			"src/test/resources/monitoring-function-migration-test-" + fileModTime + "-flink" + flinkGenerateSavepointVersion + "-snapshot");

	monitoringFunction.cancel();
	runner.join();

	testHarness.close();
}
 

@Test
public void testCheckpointing() throws Exception {
	source.autoAck = false;

	StreamSource<String, RMQSource<String>> src = new StreamSource<>(source);
	AbstractStreamOperatorTestHarness<String> testHarness =
		new AbstractStreamOperatorTestHarness<>(src, 1, 1, 0);
	testHarness.open();

	sourceThread.start();

	Thread.sleep(5);

	final Random random = new Random(System.currentTimeMillis());
	int numSnapshots = 50;
	long previousSnapshotId;
	long lastSnapshotId = 0;

	long totalNumberOfAcks = 0;

	for (int i = 0; i < numSnapshots; i++) {
		long snapshotId = random.nextLong();
		OperatorSubtaskState data;

		synchronized (DummySourceContext.lock) {
			data = testHarness.snapshot(snapshotId, System.currentTimeMillis());
			previousSnapshotId = lastSnapshotId;
			lastSnapshotId = messageId;
		}
		// let some time pass
		Thread.sleep(5);

		// check if the correct number of messages have been snapshotted
		final long numIds = lastSnapshotId - previousSnapshotId;

		RMQTestSource sourceCopy = new RMQTestSource();
		StreamSource<String, RMQTestSource> srcCopy = new StreamSource<>(sourceCopy);
		AbstractStreamOperatorTestHarness<String> testHarnessCopy =
			new AbstractStreamOperatorTestHarness<>(srcCopy, 1, 1, 0);

		testHarnessCopy.setup();
		testHarnessCopy.initializeState(data);
		testHarnessCopy.open();

		ArrayDeque<Tuple2<Long, Set<String>>> deque = sourceCopy.getRestoredState();
		Set<String> messageIds = deque.getLast().f1;

		assertEquals(numIds, messageIds.size());
		if (messageIds.size() > 0) {
			assertTrue(messageIds.contains(Long.toString(lastSnapshotId)));
		}

		// check if the messages are being acknowledged and the transaction committed
		synchronized (DummySourceContext.lock) {
			source.notifyCheckpointComplete(snapshotId);
		}
		totalNumberOfAcks += numIds;

	}

	Mockito.verify(source.channel, Mockito.times((int) totalNumberOfAcks)).basicAck(Mockito.anyLong(), Mockito.eq(false));
	Mockito.verify(source.channel, Mockito.times(numSnapshots)).txCommit();

}
 

@SuppressWarnings("unchecked")
private void writeSnapshot(String path, HashMap<StreamShardMetadata, SequenceNumber> state) throws Exception {
	final TestFetcher<String> fetcher = new TestFetcher<>(
		Collections.singletonList(TEST_STREAM_NAME),
		new TestSourceContext<>(),
		new TestRuntimeContext(true, 1, 0),
		new Properties(),
		new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()),
		state,
		null);

	final DummyFlinkKinesisConsumer<String> consumer = new DummyFlinkKinesisConsumer<>(
		fetcher, new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()));

	StreamSource<String, DummyFlinkKinesisConsumer<String>> consumerOperator = new StreamSource<>(consumer);

	final AbstractStreamOperatorTestHarness<String> testHarness =
			new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);

	testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	testHarness.setup();
	testHarness.open();

	final AtomicReference<Throwable> error = new AtomicReference<>();

	// run the source asynchronously
	Thread runner = new Thread() {
		@Override
		public void run() {
			try {
				consumer.run(new TestSourceContext<>());
			} catch (Throwable t) {
				t.printStackTrace();
				error.set(t);
			}
		}
	};
	runner.start();

	fetcher.waitUntilRun();

	final OperatorSubtaskState snapshot;
	synchronized (testHarness.getCheckpointLock()) {
		snapshot = testHarness.snapshot(0L, 0L);
	}

	OperatorSnapshotUtil.writeStateHandle(snapshot, path);

	consumerOperator.close();
	runner.join();
}