org.apache.flink.streaming.api.graph.StreamConfig Java Examples

private RecordWriterOutput<OUT> createStreamOutput(
		RecordWriter<SerializationDelegate<StreamRecord<OUT>>> recordWriter,
		StreamEdge edge,
		StreamConfig upStreamConfig,
		Environment taskEnvironment) {
	OutputTag sideOutputTag = edge.getOutputTag(); // OutputTag, return null if not sideOutput

	TypeSerializer outSerializer = null;

	if (edge.getOutputTag() != null) {
		// side output
		outSerializer = upStreamConfig.getTypeSerializerSideOut(
				edge.getOutputTag(), taskEnvironment.getUserClassLoader());
	} else {
		// main output
		outSerializer = upStreamConfig.getTypeSerializerOut(taskEnvironment.getUserClassLoader());
	}

	return new RecordWriterOutput<>(recordWriter, outSerializer, sideOutputTag, this);
}
 

private static <T> void setupSourceOperator(
		StreamSource<T, ?> operator,
		ExecutionConfig executionConfig,
		Environment env,
		ProcessingTimeService timeProvider) {

	StreamConfig cfg = new StreamConfig(new Configuration());
	cfg.setStateBackend(new MemoryStateBackend());

	cfg.setTimeCharacteristic(TimeCharacteristic.EventTime);
	cfg.setOperatorID(new OperatorID());

	try {
		MockStreamTask mockTask = new MockStreamTaskBuilder(env)
			.setConfig(cfg)
			.setExecutionConfig(executionConfig)
			.setProcessingTimeService(timeProvider)
			.build();

		operator.setup(mockTask, cfg, (Output<StreamRecord<T>>) mock(Output.class));
	} catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

/**
 * Tests that some StreamTask methods are called only in the main task's thread.
 * Currently, the main task's thread is the thread that creates the task.
 */
@Test
public void testThreadInvariants() throws Throwable {
	Configuration taskConfiguration = new Configuration();
	StreamConfig streamConfig = new StreamConfig(taskConfiguration);
	streamConfig.setStreamOperator(new StreamMap<>(value -> value));
	streamConfig.setOperatorID(new OperatorID());
	try (MockEnvironment mockEnvironment =
			new MockEnvironmentBuilder()
				.setTaskConfiguration(taskConfiguration)
				.build()) {

		ClassLoader taskClassLoader = new TestUserCodeClassLoader();

		RunningTask<ThreadInspectingTask> runningTask = runTask(() -> {
			Thread.currentThread().setContextClassLoader(taskClassLoader);
			return new ThreadInspectingTask(mockEnvironment);
		});
		runningTask.invocationFuture.get();

		assertThat(runningTask.streamTask.getTaskClassLoader(), is(sameInstance(taskClassLoader)));
	}
}
 

@Test
public void testCanceleablesCanceledOnCancelTaskError() throws Exception {
	syncLatch = new OneShotLatch();

	StreamConfig cfg = new StreamConfig(new Configuration());
	try (NettyShuffleEnvironment shuffleEnvironment = new NettyShuffleEnvironmentBuilder().build()) {

		Task task = createTask(CancelFailingTask.class, shuffleEnvironment, cfg, new Configuration());

		// start the task and wait until it runs
		// execution state RUNNING is not enough, we need to wait until the stream task's run() method
		// is entered
		task.startTaskThread();
		syncLatch.await();

		// cancel the execution - this should lead to smooth shutdown
		task.cancelExecution();
		task.getExecutingThread().join();

		assertEquals(ExecutionState.CANCELED, task.getExecutionState());
	}
}
 

private static CheckpointBarrierHandler createCheckpointBarrierHandler(
		StreamConfig config,
		InputGate[] inputGates,
		SubtaskCheckpointCoordinator checkpointCoordinator,
		String taskName,
		AbstractInvokable toNotifyOnCheckpoint) {
	switch (config.getCheckpointMode()) {
		case EXACTLY_ONCE:
			if (config.isUnalignedCheckpointsEnabled()) {
				return new AlternatingCheckpointBarrierHandler(
					new CheckpointBarrierAligner(taskName, toNotifyOnCheckpoint, inputGates),
					new CheckpointBarrierUnaligner(checkpointCoordinator, taskName, toNotifyOnCheckpoint, inputGates),
					toNotifyOnCheckpoint);
			}
			return new CheckpointBarrierAligner(taskName, toNotifyOnCheckpoint, inputGates);
		case AT_LEAST_ONCE:
			if (config.isUnalignedCheckpointsEnabled()) {
				throw new IllegalStateException("Cannot use unaligned checkpoints with AT_LEAST_ONCE " +
					"checkpointing mode");
			}
			int numInputChannels = Arrays.stream(inputGates).mapToInt(InputGate::getNumberOfInputChannels).sum();
			return new CheckpointBarrierTracker(numInputChannels, toNotifyOnCheckpoint);
		default:
			throw new UnsupportedOperationException("Unrecognized Checkpointing Mode: " + config.getCheckpointMode());
	}
}
 

private static <T> void setupSourceOperator(
		StreamSource<T, ?> operator,
		ExecutionConfig executionConfig,
		Environment env,
		ProcessingTimeService timeProvider) {

	StreamConfig cfg = new StreamConfig(new Configuration());
	cfg.setStateBackend(new MemoryStateBackend());

	cfg.setTimeCharacteristic(TimeCharacteristic.EventTime);
	cfg.setOperatorID(new OperatorID());

	try {
		MockStreamTask mockTask = new MockStreamTaskBuilder(env)
			.setConfig(cfg)
			.setExecutionConfig(executionConfig)
			.setProcessingTimeService(timeProvider)
			.build();

		operator.setup(mockTask, cfg, (Output<StreamRecord<T>>) mock(Output.class));
	} catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

/**
 * If finishing a task doesn't swallow exceptions this test would fail with an exception.
 */
@Test
public void finishingIgnoresExceptions() throws Exception {
	final StreamTaskTestHarness<String> testHarness = new StreamTaskTestHarness<>(
			SourceStreamTask::new,
			BasicTypeInfo.STRING_TYPE_INFO);

	final CompletableFuture<Void> operatorRunningWaitingFuture = new CompletableFuture<>();
	ExceptionThrowingSource.setIsInRunLoopFuture(operatorRunningWaitingFuture);

	testHarness.setupOutputForSingletonOperatorChain();
	StreamConfig streamConfig = testHarness.getStreamConfig();
	streamConfig.setStreamOperator(new StreamSource<>(new ExceptionThrowingSource()));
	streamConfig.setOperatorID(new OperatorID());

	testHarness.invoke();
	operatorRunningWaitingFuture.get();
	testHarness.getTask().finishTask();

	testHarness.waitForTaskCompletion();
}
 

private static <OUT> List<RecordWriter<SerializationDelegate<StreamRecord<OUT>>>> createRecordWriters(
		StreamConfig configuration,
		Environment environment) {
	List<RecordWriter<SerializationDelegate<StreamRecord<OUT>>>> recordWriters = new ArrayList<>();
	List<StreamEdge> outEdgesInOrder = configuration.getOutEdgesInOrder(environment.getUserClassLoader());
	Map<Integer, StreamConfig> chainedConfigs = configuration.getTransitiveChainedTaskConfigsWithSelf(environment.getUserClassLoader());

	for (int i = 0; i < outEdgesInOrder.size(); i++) {
		StreamEdge edge = outEdgesInOrder.get(i);
		recordWriters.add(
			createRecordWriter(
				edge,
				i,
				environment,
				environment.getTaskInfo().getTaskName(),
				chainedConfigs.get(edge.getSourceId()).getBufferTimeout()));
	}
	return recordWriters;
}
 

@Override
public void setup(StreamTask<?, ?> containingTask, StreamConfig config, Output<StreamRecord<OUT>> output) {
	ACTUAL_ORDER_TRACKING.add("OPERATOR::setup");
	super.setup(containingTask, config, output);
	if (simulateCheckpointing) {
		testCheckpointer = new Thread() {
			@Override
			public void run() {
				try {
					runStarted.await();
					if (getContainingTask().isCanceled() || getContainingTask().triggerCheckpoint(
							new CheckpointMetaData(0, System.currentTimeMillis()),
							CheckpointOptions.forCheckpointWithDefaultLocation())) {
						LifecycleTrackingStreamSource.runFinish.trigger();
					}
				} catch (Exception e) {
					e.printStackTrace();
					Assert.fail();
				}
			}
		};
		testCheckpointer.start();
	}
}
 

@Test
public void testLifeCycleCancel() throws Exception {
	ACTUAL_ORDER_TRACKING.clear();

	Configuration taskManagerConfig = new Configuration();
	StreamConfig cfg = new StreamConfig(new Configuration());
	MockSourceFunction srcFun = new MockSourceFunction();
	cfg.setStreamOperator(new LifecycleTrackingStreamSource<>(srcFun, false));
	cfg.setOperatorID(new OperatorID());
	cfg.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	Task task = StreamTaskTest.createTask(SourceStreamTask.class, cfg, taskManagerConfig);

	task.startTaskThread();
	LifecycleTrackingStreamSource.runStarted.await();

	// this should cancel the task even though it is blocked on runFinished
	task.cancelExecution();

	// wait for clean termination
	task.getExecutingThread().join();
	assertEquals(ExecutionState.CANCELED, task.getExecutionState());
	assertEquals(EXPECTED_CALL_ORDER_CANCEL_RUNNING, ACTUAL_ORDER_TRACKING);
}
 

@Override
public void init() throws Exception {
	StreamConfig configuration = getConfiguration();
	int numberOfInputs = configuration.getNumberOfInputs();

	if (numberOfInputs > 0) {
		CheckpointedInputGate inputGate = createCheckpointedInputGate();
		DataOutput<IN> output = createDataOutput();
		StreamTaskInput<IN> input = createTaskInput(inputGate, output);
		inputProcessor = new StreamOneInputProcessor<>(
			input,
			output,
			operatorChain);
	}
	headOperator.getMetricGroup().gauge(MetricNames.IO_CURRENT_INPUT_WATERMARK, this.inputWatermarkGauge);
	// wrap watermark gauge since registered metrics must be unique
	getEnvironment().getMetricGroup().gauge(MetricNames.IO_CURRENT_INPUT_WATERMARK, this.inputWatermarkGauge::getValue);
}
 

@Override
public void setup(StreamTask<?, ?> containingTask, StreamConfig config, Output<StreamRecord<OUT>> output) {
	super.setup(containingTask, config, output);

	this.inStreamElementSerializer = new StreamElementSerializer<>(
		getOperatorConfig().<IN>getTypeSerializerIn1(getUserCodeClassloader()));

	switch (outputMode) {
		case ORDERED:
			queue = new OrderedStreamElementQueue<>(capacity);
			break;
		case UNORDERED:
			queue = new UnorderedStreamElementQueue<>(capacity);
			break;
		default:
			throw new IllegalStateException("Unknown async mode: " + outputMode + '.');
	}

	this.timestampedCollector = new TimestampedCollector<>(output);
}
 

@Test
public void testQuiesceTimerServiceAfterOpClose() throws Exception {

	final OneInputStreamTaskTestHarness<String, String> testHarness = new OneInputStreamTaskTestHarness<>(
			OneInputStreamTask::new,
			2, 2,
			BasicTypeInfo.STRING_TYPE_INFO, BasicTypeInfo.STRING_TYPE_INFO);
	testHarness.setupOutputForSingletonOperatorChain();

	StreamConfig streamConfig = testHarness.getStreamConfig();
	streamConfig.setStreamOperator(new TestOperator());
	streamConfig.setOperatorID(new OperatorID());

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

	SystemProcessingTimeService timeService = (SystemProcessingTimeService)
			testHarness.getTask().getProcessingTimeService();

	// verify that the timer service is running
	Assert.assertTrue(timeService.isAlive());

	testHarness.endInput();
	testHarness.waitForTaskCompletion();
	timeService.shutdownService();
}
 

@SuppressWarnings("unchecked")
public static CheckpointedInputGate createCheckpointedInputGate(
		AbstractInvokable toNotifyOnCheckpoint,
		StreamConfig config,
		SubtaskCheckpointCoordinator checkpointCoordinator,
		IndexedInputGate[] inputGates,
		TaskIOMetricGroup taskIOMetricGroup,
		String taskName) {
	CheckpointedInputGate[] checkpointedInputGates = createCheckpointedMultipleInputGate(
		toNotifyOnCheckpoint,
		config,
		checkpointCoordinator,
		taskIOMetricGroup,
		taskName,
		Arrays.asList(inputGates));
	return Iterables.getOnlyElement(Arrays.asList(checkpointedInputGates));
}
 

@Test
public void testLifeCycleCancel() throws Exception {
	ACTUAL_ORDER_TRACKING.clear();

	Configuration taskManagerConfig = new Configuration();
	StreamConfig cfg = new StreamConfig(new Configuration());
	MockSourceFunction srcFun = new MockSourceFunction();
	cfg.setStreamOperator(new LifecycleTrackingStreamSource<>(srcFun, false));
	cfg.setOperatorID(new OperatorID());
	cfg.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	try (ShuffleEnvironment shuffleEnvironment = new NettyShuffleEnvironmentBuilder().build()) {
		Task task = StreamTaskTest.createTask(SourceStreamTask.class, shuffleEnvironment, cfg, taskManagerConfig);

		task.startTaskThread();
		LifecycleTrackingStreamSource.runStarted.await();

		// this should cancel the task even though it is blocked on runFinished
		task.cancelExecution();

		// wait for clean termination
		task.getExecutingThread().join();
		assertEquals(ExecutionState.CANCELED, task.getExecutionState());
		assertEquals(EXPECTED_CALL_ORDER_CANCEL_RUNNING, ACTUAL_ORDER_TRACKING);
	}
}
 

@Test
public void testCancellationNotBlockedOnLock() throws Exception {
	syncLatch = new OneShotLatch();

	StreamConfig cfg = new StreamConfig(new Configuration());
	Task task = createTask(CancelLockingTask.class, cfg, new Configuration());

	// start the task and wait until it runs
	// execution state RUNNING is not enough, we need to wait until the stream task's run() method
	// is entered
	task.startTaskThread();
	syncLatch.await();

	// cancel the execution - this should lead to smooth shutdown
	task.cancelExecution();
	task.getExecutingThread().join();

	assertEquals(ExecutionState.CANCELED, task.getExecutionState());
}
 

public StreamingRuntimeContext(
		Environment env,
		Map<String, Accumulator<?, ?>> accumulators,
		MetricGroup operatorMetricGroup,
		OperatorID operatorID,
		ProcessingTimeService processingTimeService,
		@Nullable KeyedStateStore keyedStateStore,
		ExternalResourceInfoProvider externalResourceInfoProvider) {
	super(checkNotNull(env).getTaskInfo(),
			env.getUserClassLoader(),
			env.getExecutionConfig(),
			accumulators,
			env.getDistributedCacheEntries(),
			operatorMetricGroup);
	this.taskEnvironment = env;
	this.streamConfig = new StreamConfig(env.getTaskConfiguration());
	this.operatorUniqueID = checkNotNull(operatorID).toString();
	this.processingTimeService = processingTimeService;
	this.keyedStateStore = keyedStateStore;
	this.externalResourceInfoProvider = externalResourceInfoProvider;
}
 

@Override
@SuppressWarnings("unchecked")
public <T extends StreamOperator<E>> T createStreamOperator(
    StreamTask<?, ?> containingTask, StreamConfig config, Output<StreamRecord<E>> output) {
  final TypeSerializer<E> serializer =
      config.getTypeSerializerIn1(containingTask.getUserCodeClassLoader());

  final long totalMemoryUsedForFeedbackCheckpointing =
      config
          .getConfiguration()
          .getInteger(FeedbackConfiguration.TOTAL_MEMORY_USED_FOR_FEEDBACK_CHECKPOINTING);

  FeedbackUnionOperator<E> op =
      new FeedbackUnionOperator<>(
          feedbackKey,
          isBarrierMessage,
          keySelector,
          totalMemoryUsedForFeedbackCheckpointing,
          serializer,
          mailboxExecutor);

  op.setup(containingTask, config, output);

  return (T) op;
}
 

@Test
public void testLifeCycleFull() throws Exception {
	ACTUAL_ORDER_TRACKING.clear();

	Configuration taskManagerConfig = new Configuration();
	StreamConfig cfg = new StreamConfig(new Configuration());
	MockSourceFunction srcFun = new MockSourceFunction();

	cfg.setStreamOperator(new LifecycleTrackingStreamSource<>(srcFun, true));
	cfg.setOperatorID(new OperatorID());
	cfg.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	Task task = StreamTaskTest.createTask(SourceStreamTask.class, cfg, taskManagerConfig);

	task.startTaskThread();

	LifecycleTrackingStreamSource.runStarted.await();

	// wait for clean termination
	task.getExecutingThread().join();
	assertEquals(ExecutionState.FINISHED, task.getExecutionState());
	assertEquals(EXPECTED_CALL_ORDER_FULL, ACTUAL_ORDER_TRACKING);
}
 

@Test
public void testStateBackendClosingOnFailure() throws Exception {
	Configuration taskManagerConfig = new Configuration();
	taskManagerConfig.setString(CheckpointingOptions.STATE_BACKEND, TestMemoryStateBackendFactory.class.getName());

	StreamConfig cfg = new StreamConfig(new Configuration());
	cfg.setStateKeySerializer(mock(TypeSerializer.class));
	cfg.setOperatorID(new OperatorID(4711L, 42L));
	TestStreamSource<Long, MockSourceFunction> streamSource = new TestStreamSource<>(new MockSourceFunction());
	cfg.setStreamOperator(streamSource);
	cfg.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	try (NettyShuffleEnvironment shuffleEnvironment = new NettyShuffleEnvironmentBuilder().build()) {
		Task task = createTask(StateBackendTestSource.class, shuffleEnvironment, cfg, taskManagerConfig);

		StateBackendTestSource.fail = true;
		task.startTaskThread();

		// wait for clean termination
		task.getExecutingThread().join();

		// ensure that the state backends and stream iterables are closed ...
		verify(TestStreamSource.operatorStateBackend).close();
		verify(TestStreamSource.keyedStateBackend).close();
		verify(TestStreamSource.rawOperatorStateInputs).close();
		verify(TestStreamSource.rawKeyedStateInputs).close();
		// ... and disposed
		verify(TestStreamSource.operatorStateBackend).dispose();
		verify(TestStreamSource.keyedStateBackend).dispose();

		assertEquals(ExecutionState.FAILED, task.getExecutionState());
	}
}
 

/**
 * Tests that the StreamTask first closes all of its operators before setting its
 * state to not running (isRunning == false)
 *
 * <p>See FLINK-7430.
 */
@Test
public void testOperatorClosingBeforeStopRunning() throws Throwable {
	BlockingCloseStreamOperator.resetLatches();
	Configuration taskConfiguration = new Configuration();
	StreamConfig streamConfig = new StreamConfig(taskConfiguration);
	streamConfig.setStreamOperator(new BlockingCloseStreamOperator());
	streamConfig.setOperatorID(new OperatorID());

	try (MockEnvironment mockEnvironment =
			new MockEnvironmentBuilder()
				.setTaskName("Test Task")
				.setManagedMemorySize(32L * 1024L)
				.setInputSplitProvider(new MockInputSplitProvider())
				.setBufferSize(1)
				.setTaskConfiguration(taskConfiguration)
				.build()) {

		RunningTask<StreamTask<Void, BlockingCloseStreamOperator>> task = runTask(() -> new NoOpStreamTask<>(mockEnvironment));

		BlockingCloseStreamOperator.inClose.await();

		// check that the StreamTask is not yet in isRunning == false
		assertTrue(task.streamTask.isRunning());

		// let the operator finish its close operation
		BlockingCloseStreamOperator.finishClose.trigger();

		task.waitForTaskCompletion(false);

		// now the StreamTask should no longer be running
		assertFalse(task.streamTask.isRunning());
	}
}
 

@Override
public void init() throws Exception {
	StreamConfig configuration = getConfiguration();

	TypeSerializer<IN> inSerializer = configuration.getTypeSerializerIn1(getUserCodeClassLoader());
	int numberOfInputs = configuration.getNumberOfInputs();

	if (numberOfInputs > 0) {
		InputGate[] inputGates = getEnvironment().getAllInputGates();

		inputProcessor = new StreamOneInputProcessor<>(
				inputGates,
				inSerializer,
				this,
				configuration.getCheckpointMode(),
				getCheckpointLock(),
				getEnvironment().getIOManager(),
				getEnvironment().getTaskManagerInfo().getConfiguration(),
				getStreamStatusMaintainer(),
				this.headOperator,
				getEnvironment().getMetricGroup().getIOMetricGroup(),
				inputWatermarkGauge,
				getTaskNameWithSubtaskAndId(),
				operatorChain);
	}
	headOperator.getMetricGroup().gauge(MetricNames.IO_CURRENT_INPUT_WATERMARK, this.inputWatermarkGauge);
	// wrap watermark gauge since registered metrics must be unique
	getEnvironment().getMetricGroup().gauge(MetricNames.IO_CURRENT_INPUT_WATERMARK, this.inputWatermarkGauge::getValue);
}
 

public StreamingRuntimeContext(AbstractStreamOperator<?> operator,
								Environment env, Map<String, Accumulator<?, ?>> accumulators) {
	super(env.getTaskInfo(),
			env.getUserClassLoader(),
			operator.getExecutionConfig(),
			accumulators,
			env.getDistributedCacheEntries(),
			operator.getMetricGroup());

	this.operator = operator;
	this.taskEnvironment = env;
	this.streamConfig = new StreamConfig(env.getTaskConfiguration());
	this.operatorUniqueID = operator.getOperatorID().toString();
}
 

private <IN, OT extends StreamOperator<IN>> StreamTask<IN, OT> createMockTask(
		StreamConfig streamConfig,
		Environment environment) throws Exception {

	//noinspection unchecked
	return new MockStreamTaskBuilder(environment)
		.setConfig(streamConfig)
		.setExecutionConfig(new ExecutionConfig().enableObjectReuse())
		.build();
}
 

@SuppressWarnings("unchecked")
@Override
public <T extends StreamOperator<OUT>> T createStreamOperator(StreamTask<?, ?> containingTask,
		StreamConfig config, Output<StreamRecord<OUT>> output) {
	return (T) generatedClass.newInstance(containingTask.getUserCodeClassLoader(),
			generatedClass.getReferences(), containingTask, config, output);
}
 

@VisibleForTesting
MapPartitionOperator<T, TaggedOperatorSubtaskState> writeOperatorSubtaskStates(
	OperatorID operatorID,
	StateBackend stateBackend,
	Path savepointPath,
	int localMaxParallelism) {

	DataSet<T> input = dataSet;
	if (originalKeySelector != null) {
		input = dataSet.partitionCustom(new KeyGroupRangePartitioner(localMaxParallelism), hashKeySelector);
	}

	StreamOperator<TaggedOperatorSubtaskState> operator = factory.createOperator(
		System.currentTimeMillis(),
		savepointPath);

	operator = dataSet.clean(operator);

	final StreamConfig config = getConfig(operatorID, stateBackend, operator);

	BoundedOneInputStreamTaskRunner<T> operatorRunner = new BoundedOneInputStreamTaskRunner<>(
		config,
		localMaxParallelism
	);

	MapPartitionOperator<T, TaggedOperatorSubtaskState> subtaskStates = input
		.mapPartition(operatorRunner)
		.name(operatorID.toHexString());

	if (operator instanceof BroadcastStateBootstrapOperator) {
		subtaskStates = subtaskStates.setParallelism(1);
	} else {
		int currentParallelism = getParallelism(subtaskStates);
		if (currentParallelism > localMaxParallelism) {
			subtaskStates.setParallelism(localMaxParallelism);
		}
	}
	return subtaskStates;
}
 

@Override
public void setup(
    StreamTask<?, ?> containingTask,
    StreamConfig config,
    Output<StreamRecord<WindowedValue<OutputT>>> output) {

  // make sure that FileSystems is initialized correctly
  FileSystems.setDefaultPipelineOptions(serializedOptions.get());

  super.setup(containingTask, config, output);
}
 

@Test
public void testStateBackendLoadingAndClosing() throws Exception {
	Configuration taskManagerConfig = new Configuration();
	taskManagerConfig.setString(CheckpointingOptions.STATE_BACKEND, TestMemoryStateBackendFactory.class.getName());

	StreamConfig cfg = new StreamConfig(new Configuration());
	cfg.setStateKeySerializer(mock(TypeSerializer.class));
	cfg.setOperatorID(new OperatorID(4711L, 42L));
	TestStreamSource<Long, MockSourceFunction> streamSource = new TestStreamSource<>(new MockSourceFunction());
	cfg.setStreamOperator(streamSource);
	cfg.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);

	try (ShuffleEnvironment shuffleEnvironment = new NettyShuffleEnvironmentBuilder().build()) {
		Task task = createTask(StateBackendTestSource.class, shuffleEnvironment, cfg, taskManagerConfig);

		StateBackendTestSource.fail = false;
		task.startTaskThread();

		// wait for clean termination
		task.getExecutingThread().join();

		// ensure that the state backends and stream iterables are closed ...
		verify(TestStreamSource.operatorStateBackend).close();
		verify(TestStreamSource.keyedStateBackend).close();
		verify(TestStreamSource.rawOperatorStateInputs).close();
		verify(TestStreamSource.rawKeyedStateInputs).close();
		// ... and disposed
		verify(TestStreamSource.operatorStateBackend).dispose();
		verify(TestStreamSource.keyedStateBackend).dispose();

		assertEquals(ExecutionState.FINISHED, task.getExecutionState());
	}
}
 

@SuppressWarnings("unchecked")
private static <T> void setupSourceOperator(
		StreamSource<T, ?> operator,
		ExecutionConfig executionConfig,
		Environment env,
		TimerService timerService) {

	StreamConfig cfg = new StreamConfig(new Configuration());
	cfg.setStateBackend(new MemoryStateBackend());

	cfg.setTimeCharacteristic(TimeCharacteristic.EventTime);
	cfg.setOperatorID(new OperatorID());

	try {
		MockStreamTask mockTask = new MockStreamTaskBuilder(env)
			.setConfig(cfg)
			.setExecutionConfig(executionConfig)
			.setTimerService(timerService)
			.build();

		operator.setProcessingTimeService(mockTask.getProcessingTimeServiceFactory().createProcessingTimeService(null));
		operator.setup(mockTask, cfg, (Output<StreamRecord<T>>) mock(Output.class));
	} catch (Exception e) {
		e.printStackTrace();
		fail(e.getMessage());
	}
}
 

public static Task createTask(
	Class<? extends AbstractInvokable> invokable,
	StreamConfig taskConfig,
	Configuration taskManagerConfig,
	TaskManagerActions taskManagerActions) throws Exception {
	return createTask(invokable, taskConfig, taskManagerConfig, new TestTaskStateManager(), taskManagerActions);
}