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

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

@Override
@SuppressWarnings("unchecked")
public void open(KeyGroupRangeInputSplit split) throws IOException {
	registry = new CloseableRegistry();

	final Environment environment = new SavepointEnvironment
		.Builder(getRuntimeContext(), split.getNumKeyGroups())
		.setSubtaskIndex(split.getSplitNumber())
		.setPrioritizedOperatorSubtaskState(split.getPrioritizedOperatorSubtaskState())
		.build();

	final StreamOperatorStateContext context = getStreamOperatorStateContext(environment);

	keyedStateBackend = (AbstractKeyedStateBackend<K>) context.keyedStateBackend();

	final DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, getRuntimeContext().getExecutionConfig());
	SavepointRuntimeContext ctx = new SavepointRuntimeContext(getRuntimeContext(), keyedStateStore);
	FunctionUtils.setFunctionRuntimeContext(userFunction, ctx);

	keys = getKeyIterator(ctx);

	final InternalTimerService<VoidNamespace> timerService = restoreTimerService(context);
	try {
		this.ctx = new Context(keyedStateBackend, timerService);
	} catch (Exception e) {
		throw new IOException("Failed to restore timer state", e);
	}
}
 

@Override
@SuppressWarnings("unchecked")
public void open(KeyGroupRangeInputSplit split) throws IOException {
	registry = new CloseableRegistry();

	final Environment environment = new SavepointEnvironment
		.Builder(getRuntimeContext(), split.getNumKeyGroups())
		.setConfiguration(configuration)
		.setSubtaskIndex(split.getSplitNumber())
		.setPrioritizedOperatorSubtaskState(split.getPrioritizedOperatorSubtaskState())
		.build();

	final StreamOperatorStateContext context = getStreamOperatorStateContext(environment);

	AbstractKeyedStateBackend<K> keyedStateBackend = (AbstractKeyedStateBackend<K>) context.keyedStateBackend();

	final DefaultKeyedStateStore keyedStateStore = new DefaultKeyedStateStore(keyedStateBackend, getRuntimeContext().getExecutionConfig());
	SavepointRuntimeContext ctx = new SavepointRuntimeContext(getRuntimeContext(), keyedStateStore);

	InternalTimeServiceManager<K> timeServiceManager = (InternalTimeServiceManager<K>) context.internalTimerServiceManager();
	try {
		operator.setup(getRuntimeContext().getExecutionConfig(), keyedStateBackend, timeServiceManager, ctx);
		operator.open();
		keysAndNamespaces = operator.getKeysAndNamespaces(ctx);
	} catch (Exception e) {
		throw new IOException("Failed to restore timer state", e);
	}
}
 

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

@SuppressWarnings("unchecked")
private InternalTimerService<VoidNamespace> restoreTimerService(StreamOperatorStateContext context) {
	InternalTimeServiceManager<K> timeServiceManager = (InternalTimeServiceManager<K>) context.internalTimerServiceManager();
	TimerSerializer<K, VoidNamespace> timerSerializer = new TimerSerializer<>(keySerializer, VoidNamespaceSerializer.INSTANCE);
	return timeServiceManager.getInternalTimerService(USER_TIMERS_NAME, timerSerializer, VoidTriggerable.instance());
}
 

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

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