org.apache.flink.streaming.util.MockStreamTaskBuilder Java Examples

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

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

@Test
public void testProcessWithUnAvailableOutput() throws Exception {
	try (final MockEnvironment environment = setupEnvironment(new boolean[] {true, false})) {
		final int numberOfProcessCalls = 10;
		final AvailabilityTestInputProcessor inputProcessor = new AvailabilityTestInputProcessor(numberOfProcessCalls);
		final StreamTask task = new MockStreamTaskBuilder(environment)
			.setStreamInputProcessor(inputProcessor)
			.build();
		final MailboxExecutor executor = task.mailboxProcessor.getMainMailboxExecutor();

		final RunnableWithException completeFutureTask = () -> {
			assertEquals(1, inputProcessor.currentNumProcessCalls);
			assertTrue(task.mailboxProcessor.isDefaultActionUnavailable());
			environment.getWriter(1).getAvailableFuture().complete(null);
		};

		executor.submit(() -> {
			executor.submit(completeFutureTask, "This task will complete the future to resume process input action."); },
			"This task will submit another task to execute after processing input once.");

		task.invoke();
		assertEquals(numberOfProcessCalls, inputProcessor.currentNumProcessCalls);
	}
}
 

@SuppressWarnings("unchecked")
private static <T> void setupSourceOperator(StreamSource<T, ?> operator,
											TimeCharacteristic timeChar,
											long watermarkInterval,
											final ProcessingTimeService timeProvider) throws Exception {

	ExecutionConfig executionConfig = new ExecutionConfig();
	executionConfig.setAutoWatermarkInterval(watermarkInterval);

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

	cfg.setTimeCharacteristic(timeChar);
	cfg.setOperatorID(new OperatorID());

	Environment env = new DummyEnvironment("MockTwoInputTask", 1, 0);

	StreamStatusMaintainer streamStatusMaintainer = mock(StreamStatusMaintainer.class);
	when(streamStatusMaintainer.getStreamStatus()).thenReturn(StreamStatus.ACTIVE);

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

	operator.setup(mockTask, cfg, (Output<StreamRecord<T>>) mock(Output.class));
}
 

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")
private static <T> MockStreamTask setupSourceOperator(
		StreamSource<T, ?> operator,
		TimeCharacteristic timeChar,
		long watermarkInterval,
		final TimerService timeProvider) throws Exception {

	ExecutionConfig executionConfig = new ExecutionConfig();
	executionConfig.setAutoWatermarkInterval(watermarkInterval);

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

	cfg.setTimeCharacteristic(timeChar);
	cfg.setOperatorID(new OperatorID());

	Environment env = new DummyEnvironment("MockTwoInputTask", 1, 0);

	StreamStatusMaintainer streamStatusMaintainer = mock(StreamStatusMaintainer.class);
	when(streamStatusMaintainer.getStreamStatus()).thenReturn(StreamStatus.ACTIVE);

	MockStreamTask mockTask = new MockStreamTaskBuilder(env)
		.setConfig(cfg)
		.setExecutionConfig(executionConfig)
		.setStreamStatusMaintainer(streamStatusMaintainer)
		.setTimerService(timeProvider)
		.build();

	operator.setup(mockTask, cfg, (Output<StreamRecord<T>>) mock(Output.class));
	return mockTask;
}
 

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

@Test
public void testProcessWithAvailableOutput() throws Exception {
	try (final MockEnvironment environment = setupEnvironment(new boolean[] {true, true})) {
		final int numberOfProcessCalls = 10;
		final AvailabilityTestInputProcessor inputProcessor = new AvailabilityTestInputProcessor(numberOfProcessCalls);
		final StreamTask task = new MockStreamTaskBuilder(environment)
			.setStreamInputProcessor(inputProcessor)
			.build();

		task.invoke();
		assertEquals(numberOfProcessCalls, inputProcessor.currentNumProcessCalls);
	}
}
 

@Test
public void testBeforeInvokeWithoutChannelStates() throws Exception {
	int numWriters = 2;
	int numGates = 2;
	RecoveryResultPartition[] partitions = new RecoveryResultPartition[numWriters];
	for (int i = 0; i < numWriters; i++) {
		partitions[i] = new RecoveryResultPartition();
	}
	RecoveryInputGate[] gates = new RecoveryInputGate[numGates];
	for (int i = 0; i < numGates; i++) {
		gates[i] = new RecoveryInputGate(partitions);
	}

	MockEnvironment mockEnvironment = new MockEnvironmentBuilder().build();
	mockEnvironment.addOutputs(asList(partitions));
	mockEnvironment.addInputs(asList(gates));
	StreamTask task = new MockStreamTaskBuilder(mockEnvironment).build();
	try {
		verifyResults(gates, partitions, false, false);

		task.beforeInvoke();

		verifyResults(gates, partitions, false, true);
	} finally {
		task.cleanUpInvoke();
	}
}
 

@Before
public void setup() throws Exception {
	this.operatorWrappers = new ArrayList<>();
	this.output = new ConcurrentLinkedQueue<>();

	try (MockEnvironment env = MockEnvironment.builder().build()) {
		this.containingTask = new MockStreamTaskBuilder(env).build();

		// initialize operator wrappers
		for (int i = 0; i < numOperators; i++) {
			MailboxExecutor mailboxExecutor = containingTask.getMailboxExecutorFactory().createExecutor(i);

			TimerMailController timerMailController = new TimerMailController(containingTask, mailboxExecutor);
			ProcessingTimeServiceImpl processingTimeService = new ProcessingTimeServiceImpl(
				timerService,
				timerMailController::wrapCallback);

			TestOneInputStreamOperator streamOperator = new TestOneInputStreamOperator(
				"Operator" + i, output, processingTimeService, mailboxExecutor, timerMailController);
			streamOperator.setProcessingTimeService(processingTimeService);

			StreamOperatorWrapper<?, ?> operatorWrapper = new StreamOperatorWrapper<>(
				streamOperator,
				Optional.ofNullable(streamOperator.getProcessingTimeService()),
				mailboxExecutor);
			operatorWrappers.add(operatorWrapper);
		}

		StreamOperatorWrapper<?, ?> previous = null;
		for (StreamOperatorWrapper<?, ?> current : operatorWrappers) {
			if (previous != null) {
				previous.setNext(current);
			}
			current.setPrevious(previous);
			previous = current;
		}
	}
}
 

@Test
public void testSavepointNotResultingInPriorityEvents() throws Exception {
	MockEnvironment mockEnvironment = MockEnvironment.builder().build();

	SubtaskCheckpointCoordinator coordinator = new MockSubtaskCheckpointCoordinatorBuilder()
			.setUnalignedCheckpointEnabled(true)
			.setEnvironment(mockEnvironment)
			.build();

	AtomicReference<Boolean> broadcastedPriorityEvent = new AtomicReference<>(null);
	final OperatorChain<?, ?> operatorChain = new OperatorChain(
			new MockStreamTaskBuilder(mockEnvironment).build(),
			new NonRecordWriter<>()) {
		@Override
		public void broadcastEvent(AbstractEvent event, boolean isPriorityEvent) throws IOException {
			super.broadcastEvent(event, isPriorityEvent);
			broadcastedPriorityEvent.set(isPriorityEvent);
		}
	};

	coordinator.checkpointState(
			new CheckpointMetaData(0, 0),
			new CheckpointOptions(SAVEPOINT, CheckpointStorageLocationReference.getDefault()),
			new CheckpointMetrics(),
			operatorChain,
			() -> false);

	assertEquals(false, broadcastedPriorityEvent.get());
}
 

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")
private static <T> void setupSourceOperator(StreamSource<T, ?> operator,
											TimeCharacteristic timeChar,
											long watermarkInterval,
											final ProcessingTimeService timeProvider) throws Exception {

	ExecutionConfig executionConfig = new ExecutionConfig();
	executionConfig.setAutoWatermarkInterval(watermarkInterval);

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

	cfg.setTimeCharacteristic(timeChar);
	cfg.setOperatorID(new OperatorID());

	Environment env = new DummyEnvironment("MockTwoInputTask", 1, 0);

	StreamStatusMaintainer streamStatusMaintainer = mock(StreamStatusMaintainer.class);
	when(streamStatusMaintainer.getStreamStatus()).thenReturn(StreamStatus.ACTIVE);

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

	operator.setup(mockTask, cfg, (Output<StreamRecord<T>>) mock(Output.class));
}
 

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

private OperatorChain<?, ?> getOperatorChain(MockEnvironment mockEnvironment) throws Exception {
	return new OperatorChain<>(
		new MockStreamTaskBuilder(mockEnvironment).build(),
		new NonRecordWriter<>());
}
 

@Test
public void testBeforeInvokeWithChannelStates() throws Exception {
	int numWriters = 2;
	int numGates = 2;
	RecoveryResultPartition[] partitions = new RecoveryResultPartition[numWriters];
	for (int i = 0; i < numWriters; i++) {
		partitions[i] = new RecoveryResultPartition();
	}
	RecoveryInputGate[] gates = new RecoveryInputGate[numGates];
	for (int i = 0; i < numGates; i++) {
		gates[i] = new RecoveryInputGate(partitions);
	}

	ChannelStateReader reader = new ResultPartitionTest.FiniteChannelStateReader(1, new int[] {0});
	TaskStateManager taskStateManager = new TaskStateManagerImpl(
		new JobID(),
		new ExecutionAttemptID(),
		new TestTaskLocalStateStore(),
		null,
		NoOpCheckpointResponder.INSTANCE,
		reader);
	MockEnvironment mockEnvironment = new MockEnvironmentBuilder().setTaskStateManager(taskStateManager).build();
	mockEnvironment.addOutputs(asList(partitions));
	mockEnvironment.addInputs(asList(gates));
	StreamTask task = new MockStreamTaskBuilder(mockEnvironment).build();
	try {
		verifyResults(gates, partitions, false, false);

		task.beforeInvoke();

		verifyResults(gates, partitions, true, false);

		// execute the partition request mail inserted after input recovery completes
		task.mailboxProcessor.drain();

		for (RecoveryInputGate inputGate : gates) {
			assertTrue(inputGate.isPartitionRequested());
		}
	} finally {
		task.cleanUpInvoke();
	}
}
 

@SafeVarargs
public static <T, OP extends StreamOperator<T>> OperatorChain<T, OP> setupOperatorChain(
		OneInputStreamOperator<T, T>... operators) throws Exception {

	checkNotNull(operators);
	checkArgument(operators.length > 0);

	try (MockEnvironment env = MockEnvironment.builder().build()) {
		final StreamTask<?, ?> containingTask = new MockStreamTaskBuilder(env).build();

		final StreamStatusProvider statusProvider = mock(StreamStatusProvider.class);
		final StreamConfig cfg = new StreamConfig(new Configuration());

		final List<StreamOperatorWrapper<?, ?>> operatorWrappers = new ArrayList<>();

		// initial output goes to nowhere
		@SuppressWarnings({"unchecked", "rawtypes"})
		WatermarkGaugeExposingOutput<StreamRecord<T>> lastWriter = new BroadcastingOutputCollector<>(
				new Output[0], statusProvider);

		// build the reverse operators array
		for (int i = 0; i < operators.length; i++) {
			OneInputStreamOperator<T, T> op = operators[operators.length - i - 1];
			if (op instanceof SetupableStreamOperator) {
				((SetupableStreamOperator) op).setup(containingTask, cfg, lastWriter);
			}
			lastWriter = new ChainingOutput<>(op, statusProvider, null);

			ProcessingTimeService processingTimeService = null;
			if (op instanceof AbstractStreamOperator) {
				processingTimeService = ((AbstractStreamOperator) op).getProcessingTimeService();
			}
			operatorWrappers.add(new StreamOperatorWrapper<>(
				op,
				Optional.ofNullable(processingTimeService),
				containingTask.getMailboxExecutorFactory().createExecutor(i)));
		}

		@SuppressWarnings("unchecked")
		final StreamOperatorWrapper<T, OP> headOperatorWrapper = (StreamOperatorWrapper<T, OP>) operatorWrappers.get(operatorWrappers.size() - 1);

		return new OperatorChain<>(
			operatorWrappers,
			new RecordWriterOutput<?>[0],
			lastWriter,
			headOperatorWrapper);
	}
}
 

@Test
public void testCreateCheckpointedMultipleInputGate() throws Exception {
	try (CloseableRegistry registry = new CloseableRegistry()) {
		MockEnvironment environment = new MockEnvironmentBuilder().build();
		MockStreamTask streamTask = new MockStreamTaskBuilder(environment).build();
		StreamConfig streamConfig = new StreamConfig(environment.getJobConfiguration());
		streamConfig.setCheckpointMode(CheckpointingMode.EXACTLY_ONCE);
		streamConfig.setUnalignedCheckpointsEnabled(true);

		// First input gate has index larger than the second
		List<IndexedInputGate>[] inputGates = new List[] {
			Collections.singletonList(new MockIndexedInputGate(1, 4)),
			Collections.singletonList(new MockIndexedInputGate(0, 2)),
		};

		CheckpointedInputGate[] checkpointedMultipleInputGate = InputProcessorUtil.createCheckpointedMultipleInputGate(
			streamTask,
			streamConfig,
			new TestSubtaskCheckpointCoordinator(new MockChannelStateWriter()),
			environment.getMetricGroup().getIOMetricGroup(),
			streamTask.getName(),
			inputGates);
		for (CheckpointedInputGate checkpointedInputGate : checkpointedMultipleInputGate) {
			registry.registerCloseable(checkpointedInputGate);
		}

		CheckpointBarrierHandler barrierHandler = checkpointedMultipleInputGate[0].getCheckpointBarrierHandler();
		assertTrue(barrierHandler.getBufferReceivedListener().isPresent());
		BufferReceivedListener bufferReceivedListener = barrierHandler.getBufferReceivedListener().get();

		List<IndexedInputGate> allInputGates = Arrays.stream(inputGates).flatMap(gates -> gates.stream()).collect(Collectors.toList());
		for (IndexedInputGate inputGate : allInputGates) {
			for (int channelId = 0; channelId < inputGate.getNumberOfInputChannels(); channelId++) {
				bufferReceivedListener.notifyBarrierReceived(
					new CheckpointBarrier(1, 42, CheckpointOptions.forCheckpointWithDefaultLocation(true, true)),
					new InputChannelInfo(inputGate.getGateIndex(), channelId));
			}
		}
		assertTrue(barrierHandler.getAllBarriersReceivedFuture(1).isDone());
	}
}