org.apache.flink.streaming.api.operators.StreamTaskStateInitializer Java Examples

public MockStreamTask(
	Environment environment,
	Object checkpointLock,
	StreamConfig config,
	ExecutionConfig executionConfig,
	StreamTaskStateInitializer streamTaskStateInitializer,
	StreamStatusMaintainer streamStatusMaintainer,
	CheckpointStorageWorkerView checkpointStorage,
	TimerService timerService,
	BiConsumer<String, Throwable> handleAsyncException,
	TaskMailbox taskMailbox,
	StreamTaskActionExecutor.SynchronizedStreamTaskActionExecutor taskActionExecutor,
	StreamInputProcessor inputProcessor) throws Exception {

	super(environment, timerService, FatalExitExceptionHandler.INSTANCE, taskActionExecutor, taskMailbox);
	this.checkpointLock = checkpointLock;
	this.config = config;
	this.executionConfig = executionConfig;
	this.streamTaskStateInitializer = streamTaskStateInitializer;
	this.streamStatusMaintainer = streamStatusMaintainer;
	this.checkpointStorage = checkpointStorage;
	this.processingTimeService = timerService;
	this.handleAsyncException = handleAsyncException;
	this.inputProcessor = inputProcessor;
}
 

private StreamOperatorStateContext getStreamOperatorStateContext(Environment environment) throws IOException {
	StreamTaskStateInitializer initializer = new StreamTaskStateInitializerImpl(
		environment,
		stateBackend,
		new NeverFireProcessingTimeService());

	try {
		return initializer.streamOperatorStateContext(
			operatorState.getOperatorID(),
			operatorState.getOperatorID().toString(),
			this,
			keySerializer,
			registry,
			getRuntimeContext().getMetricGroup());
	} catch (Exception e) {
		throw new IOException("Failed to restore state backend", e);
	}
}
 

private StreamOperatorStateContext getStreamOperatorStateContext(Environment environment) throws IOException {
	StreamTaskStateInitializer initializer = new StreamTaskStateInitializerImpl(
		environment,
		stateBackend);

	try {
		return initializer.streamOperatorStateContext(
			operatorState.getOperatorID(),
			operatorState.getOperatorID().toString(),
			new NeverFireProcessingTimeService(),
			operator,
			operator.getKeyType().createSerializer(environment.getExecutionConfig()),
			registry,
			getRuntimeContext().getMetricGroup());
	} catch (Exception e) {
		throw new IOException("Failed to restore state backend", e);
	}
}
 

protected StreamTaskStateInitializer createStreamTaskStateManager(
	Environment env,
	StateBackend stateBackend,
	TtlTimeProvider ttlTimeProvider) {
	return new StreamTaskStateInitializerImpl(
		env,
		stateBackend,
		ttlTimeProvider);
}
 

public MockStreamTask(
	Environment environment,
	String name,
	Object checkpointLock,
	StreamConfig config,
	ExecutionConfig executionConfig,
	StreamTaskStateInitializer streamTaskStateInitializer,
	CloseableRegistry closableRegistry,
	StreamStatusMaintainer streamStatusMaintainer,
	CheckpointStorage checkpointStorage,
	ProcessingTimeService processingTimeService,
	BiConsumer<String, Throwable> handleAsyncException,
	Map<String, Accumulator<?, ?>> accumulatorMap
) {
	super(environment);
	this.name = name;
	this.checkpointLock = checkpointLock;
	this.config = config;
	this.executionConfig = executionConfig;
	this.streamTaskStateInitializer = streamTaskStateInitializer;
	this.closableRegistry = closableRegistry;
	this.streamStatusMaintainer = streamStatusMaintainer;
	this.checkpointStorage = checkpointStorage;
	this.processingTimeService = processingTimeService;
	this.handleAsyncException = handleAsyncException;
	this.accumulatorMap = accumulatorMap;
}
 

protected StreamTaskStateInitializer createStreamTaskStateManager(
	Environment env,
	StateBackend stateBackend,
	ProcessingTimeService processingTimeService) {
	return new StreamTaskStateInitializerImpl(
		env,
		stateBackend,
		processingTimeService);
}
 

public MockStreamTask(
	Environment environment,
	String name,
	Object checkpointLock,
	StreamConfig config,
	ExecutionConfig executionConfig,
	StreamTaskStateInitializer streamTaskStateInitializer,
	CloseableRegistry closableRegistry,
	StreamStatusMaintainer streamStatusMaintainer,
	CheckpointStorageWorkerView checkpointStorage,
	ProcessingTimeService processingTimeService,
	BiConsumer<String, Throwable> handleAsyncException,
	Map<String, Accumulator<?, ?>> accumulatorMap
) {
	super(environment);
	this.name = name;
	this.checkpointLock = checkpointLock;
	this.config = config;
	this.executionConfig = executionConfig;
	this.streamTaskStateInitializer = streamTaskStateInitializer;
	this.closableRegistry = closableRegistry;
	this.streamStatusMaintainer = streamStatusMaintainer;
	this.checkpointStorage = checkpointStorage;
	this.processingTimeService = processingTimeService;
	this.handleAsyncException = handleAsyncException;
	this.accumulatorMap = accumulatorMap;
}
 

/**
 * Initialize state and open all operators in the chain from <b>tail to head</b>,
 * contrary to {@link StreamOperator#close()} which happens <b>head to tail</b>
 * (see {@link #closeOperators(StreamTaskActionExecutor)}).
 */
protected void initializeStateAndOpenOperators(StreamTaskStateInitializer streamTaskStateInitializer) throws Exception {
	for (StreamOperatorWrapper<?, ?> operatorWrapper : getAllOperators(true)) {
		StreamOperator<?> operator = operatorWrapper.getStreamOperator();
		operator.initializeState(streamTaskStateInitializer);
		operator.open();
	}
}
 

protected StreamTaskStateInitializer createStreamTaskStateManager(
	Environment env,
	StateBackend stateBackend,
	ProcessingTimeService processingTimeService) {
	return new StreamTaskStateInitializerImpl(
		env,
		stateBackend,
		processingTimeService);
}
 

public void setStreamTaskStateInitializer(StreamTaskStateInitializer streamTaskStateInitializer) {
	this.streamTaskStateInitializer = streamTaskStateInitializer;
}
 

public void setStreamTaskStateInitializer(StreamTaskStateInitializer streamTaskStateInitializer) {
	this.streamTaskStateInitializer = streamTaskStateInitializer;
}
 

@Override
public StreamTaskStateInitializer createStreamTaskStateInitializer() {
	return streamTaskStateInitializer;
}
 

public MockStreamTaskBuilder setStreamTaskStateInitializer(StreamTaskStateInitializer streamTaskStateInitializer) {
	this.streamTaskStateInitializer = streamTaskStateInitializer;
	return this;
}
 

@Override
public StreamTaskStateInitializer createStreamTaskStateInitializer() {
	final StreamTaskStateInitializer streamTaskStateManager = super.createStreamTaskStateInitializer();
	return (operatorID, operatorClassName, processingTimeService, keyContext, keySerializer, closeableRegistry, metricGroup) -> {

		final StreamOperatorStateContext controller = streamTaskStateManager.streamOperatorStateContext(
			operatorID,
			operatorClassName,
			processingTimeService,
			keyContext,
			keySerializer,
			closeableRegistry,
			metricGroup);

		return new StreamOperatorStateContext() {
			@Override
			public boolean isRestored() {
				return controller.isRestored();
			}

			@Override
			public OperatorStateBackend operatorStateBackend() {
				return controller.operatorStateBackend();
			}

			@Override
			public AbstractKeyedStateBackend<?> keyedStateBackend() {
				return controller.keyedStateBackend();
			}

			@Override
			public InternalTimeServiceManager<?> internalTimerServiceManager() {
				InternalTimeServiceManager<?> timeServiceManager = controller.internalTimerServiceManager();
				return timeServiceManager != null ? spy(timeServiceManager) : null;
			}

			@Override
			public CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs() {
				return replaceWithSpy(controller.rawOperatorStateInputs());
			}

			@Override
			public CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs() {
				return replaceWithSpy(controller.rawKeyedStateInputs());
			}

			public <T extends Closeable> T replaceWithSpy(T closeable) {
				T spyCloseable = spy(closeable);
				if (closeableRegistry.unregisterCloseable(closeable)) {
					try {
						closeableRegistry.registerCloseable(spyCloseable);
					} catch (IOException e) {
						throw new RuntimeException(e);
					}
				}
				return spyCloseable;
			}
		};
	};
}
 

@Override
public void initializeState(StreamTaskStateInitializer streamTaskStateManager) throws Exception {
}
 

public StreamTaskStateInitializer createStreamTaskStateInitializer() {
	return new StreamTaskStateInitializerImpl(
		getEnvironment(),
		stateBackend);
}
 

@Override
public void initializeState(StreamTaskStateInitializer streamTaskStateManager) throws Exception {
	ACTUAL_ORDER_TRACKING.add("initializeState");
}
 

public StreamTaskStateInitializer createStreamTaskStateInitializer() {
	return new StreamTaskStateInitializerImpl(
		getEnvironment(),
		stateBackend,
		timerService);
}
 

@Override
public StreamTaskStateInitializer createStreamTaskStateInitializer() {
	return streamTaskStateInitializer;
}
 

public MockStreamTaskBuilder setStreamTaskStateInitializer(StreamTaskStateInitializer streamTaskStateInitializer) {
	this.streamTaskStateInitializer = streamTaskStateInitializer;
	return this;
}
 

@Override
public StreamTaskStateInitializer createStreamTaskStateInitializer() {
	final StreamTaskStateInitializer streamTaskStateManager = super.createStreamTaskStateInitializer();
	return (operatorID, operatorClassName, keyContext, keySerializer, closeableRegistry, metricGroup) -> {

		final StreamOperatorStateContext context = streamTaskStateManager.streamOperatorStateContext(
			operatorID,
			operatorClassName,
			keyContext,
			keySerializer,
			closeableRegistry,
			metricGroup);

		return new StreamOperatorStateContext() {
			@Override
			public boolean isRestored() {
				return context.isRestored();
			}

			@Override
			public OperatorStateBackend operatorStateBackend() {
				return context.operatorStateBackend();
			}

			@Override
			public AbstractKeyedStateBackend<?> keyedStateBackend() {
				return context.keyedStateBackend();
			}

			@Override
			public InternalTimeServiceManager<?> internalTimerServiceManager() {
				InternalTimeServiceManager<?> timeServiceManager = context.internalTimerServiceManager();
				return timeServiceManager != null ? spy(timeServiceManager) : null;
			}

			@Override
			public CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs() {
				return replaceWithSpy(context.rawOperatorStateInputs());
			}

			@Override
			public CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs() {
				return replaceWithSpy(context.rawKeyedStateInputs());
			}

			public <T extends Closeable> T replaceWithSpy(T closeable) {
				T spyCloseable = spy(closeable);
				if (closeableRegistry.unregisterCloseable(closeable)) {
					try {
						closeableRegistry.registerCloseable(spyCloseable);
					} catch (IOException e) {
						throw new RuntimeException(e);
					}
				}
				return spyCloseable;
			}
		};
	};
}
 

public StreamTaskStateInitializer createStreamTaskStateInitializer() {
	return new StreamTaskStateInitializerImpl(
		getEnvironment(),
		stateBackend,
		timerService);
}
 

public void setStreamTaskStateInitializer(StreamTaskStateInitializer streamTaskStateInitializer) {
	this.streamTaskStateInitializer = streamTaskStateInitializer;
}
 

@Override
public StreamTaskStateInitializer createStreamTaskStateInitializer() {
	return streamTaskStateInitializer;
}
 

public MockStreamTaskBuilder setStreamTaskStateInitializer(StreamTaskStateInitializer streamTaskStateInitializer) {
	this.streamTaskStateInitializer = streamTaskStateInitializer;
	return this;
}
 

@Override
public StreamTaskStateInitializer createStreamTaskStateInitializer() {
	final StreamTaskStateInitializer streamTaskStateManager = super.createStreamTaskStateInitializer();
	return (operatorID, operatorClassName, keyContext, keySerializer, closeableRegistry, metricGroup) -> {

		final StreamOperatorStateContext context = streamTaskStateManager.streamOperatorStateContext(
			operatorID,
			operatorClassName,
			keyContext,
			keySerializer,
			closeableRegistry,
			metricGroup);

		return new StreamOperatorStateContext() {
			@Override
			public boolean isRestored() {
				return context.isRestored();
			}

			@Override
			public OperatorStateBackend operatorStateBackend() {
				return context.operatorStateBackend();
			}

			@Override
			public AbstractKeyedStateBackend<?> keyedStateBackend() {
				return context.keyedStateBackend();
			}

			@Override
			public InternalTimeServiceManager<?> internalTimerServiceManager() {
				InternalTimeServiceManager<?> timeServiceManager = context.internalTimerServiceManager();
				return timeServiceManager != null ? spy(timeServiceManager) : null;
			}

			@Override
			public CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs() {
				return replaceWithSpy(context.rawOperatorStateInputs());
			}

			@Override
			public CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs() {
				return replaceWithSpy(context.rawKeyedStateInputs());
			}

			public <T extends Closeable> T replaceWithSpy(T closeable) {
				T spyCloseable = spy(closeable);
				if (closeableRegistry.unregisterCloseable(closeable)) {
					try {
						closeableRegistry.registerCloseable(spyCloseable);
					} catch (IOException e) {
						throw new RuntimeException(e);
					}
				}
				return spyCloseable;
			}
		};
	};
}