org.apache.flink.runtime.checkpoint.PrioritizedOperatorSubtaskState Java Examples

private SavepointEnvironment(RuntimeContext ctx, Configuration configuration, int maxParallelism, int indexOfSubtask, PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskState) {
	this.jobID = new JobID();
	this.vertexID = new JobVertexID();
	this.attemptID = new ExecutionAttemptID();
	this.ctx = Preconditions.checkNotNull(ctx);
	this.configuration = Preconditions.checkNotNull(configuration);

	Preconditions.checkArgument(maxParallelism > 0 && indexOfSubtask < maxParallelism);
	this.maxParallelism = maxParallelism;
	this.indexOfSubtask = indexOfSubtask;

	this.registry = new KvStateRegistry().createTaskRegistry(jobID, vertexID);
	this.taskStateManager = new SavepointTaskStateManager(prioritizedOperatorSubtaskState);
	this.ioManager = new IOManagerAsync(ConfigurationUtils.parseTempDirectories(configuration));
	this.memoryManager = MemoryManager.forDefaultPageSize(64 * 1024 * 1024);
	this.accumulatorRegistry = new AccumulatorRegistry(jobID, attemptID);
}
 

private SavepointEnvironment(RuntimeContext ctx, Configuration configuration, int maxParallelism, int indexOfSubtask, PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskState) {
	this.jobID = new JobID();
	this.vertexID = new JobVertexID();
	this.attemptID = new ExecutionAttemptID();
	this.ctx = Preconditions.checkNotNull(ctx);
	this.configuration = Preconditions.checkNotNull(configuration);

	Preconditions.checkArgument(maxParallelism > 0 && indexOfSubtask < maxParallelism);
	this.maxParallelism = maxParallelism;
	this.indexOfSubtask = indexOfSubtask;

	this.registry = new KvStateRegistry().createTaskRegistry(jobID, vertexID);
	this.taskStateManager = new SavepointTaskStateManager(prioritizedOperatorSubtaskState);
	this.ioManager = new IOManagerAsync();
	this.accumulatorRegistry = new AccumulatorRegistry(jobID, attemptID);
}
 

@Nonnull
@Override
public PrioritizedOperatorSubtaskState prioritizedOperatorState(OperatorID operatorID) {
	TaskStateSnapshot jmTaskStateSnapshot = getLastJobManagerTaskStateSnapshot();
	TaskStateSnapshot tmTaskStateSnapshot = getLastTaskManagerTaskStateSnapshot();

	if (jmTaskStateSnapshot == null) {

		return PrioritizedOperatorSubtaskState.emptyNotRestored();
	} else {

		OperatorSubtaskState jmOpState = jmTaskStateSnapshot.getSubtaskStateByOperatorID(operatorID);

		if (jmOpState == null) {

			return PrioritizedOperatorSubtaskState.emptyNotRestored();
		} else {

			List<OperatorSubtaskState> tmStateCollection = Collections.emptyList();

			if (tmTaskStateSnapshot != null) {
				OperatorSubtaskState tmOpState = tmTaskStateSnapshot.getSubtaskStateByOperatorID(operatorID);
				if (tmOpState != null) {
					tmStateCollection = Collections.singletonList(tmOpState);
				}
			}
			PrioritizedOperatorSubtaskState.Builder builder =
				new PrioritizedOperatorSubtaskState.Builder(jmOpState, tmStateCollection);
			return builder.build();
		}
	}
}
 

protected OperatorStateBackend operatorStateBackend(
	String operatorIdentifierText,
	PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates,
	CloseableRegistry backendCloseableRegistry) throws Exception {

	String logDescription = "operator state backend for " + operatorIdentifierText;

	// Now restore processing is included in backend building/constructing process, so we need to make sure
	// each stream constructed in restore could also be closed in case of task cancel, for example the data
	// input stream opened for serDe during restore.
	CloseableRegistry cancelStreamRegistryForRestore = new CloseableRegistry();
	backendCloseableRegistry.registerCloseable(cancelStreamRegistryForRestore);
	BackendRestorerProcedure<OperatorStateBackend, OperatorStateHandle> backendRestorer =
		new BackendRestorerProcedure<>(
			(stateHandles) -> stateBackend.createOperatorStateBackend(
				environment,
				operatorIdentifierText,
				stateHandles,
				cancelStreamRegistryForRestore),
			backendCloseableRegistry,
			logDescription);

	try {
		return backendRestorer.createAndRestore(
			prioritizedOperatorSubtaskStates.getPrioritizedManagedOperatorState());
	} finally {
		if (backendCloseableRegistry.unregisterCloseable(cancelStreamRegistryForRestore)) {
			IOUtils.closeQuietly(cancelStreamRegistryForRestore);
		}
	}
}
 

protected OperatorStateBackend operatorStateBackend(
	String operatorIdentifierText,
	PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates,
	CloseableRegistry backendCloseableRegistry) throws Exception {

	String logDescription = "operator state backend for " + operatorIdentifierText;

	// Now restore processing is included in backend building/constructing process, so we need to make sure
	// each stream constructed in restore could also be closed in case of task cancel, for example the data
	// input stream opened for serDe during restore.
	CloseableRegistry cancelStreamRegistryForRestore = new CloseableRegistry();
	backendCloseableRegistry.registerCloseable(cancelStreamRegistryForRestore);
	BackendRestorerProcedure<OperatorStateBackend, OperatorStateHandle> backendRestorer =
		new BackendRestorerProcedure<>(
			(stateHandles) -> stateBackend.createOperatorStateBackend(
				environment,
				operatorIdentifierText,
				stateHandles,
				cancelStreamRegistryForRestore),
			backendCloseableRegistry,
			logDescription);

	try {
		return backendRestorer.createAndRestore(
			prioritizedOperatorSubtaskStates.getPrioritizedManagedOperatorState());
	} finally {
		if (backendCloseableRegistry.unregisterCloseable(cancelStreamRegistryForRestore)) {
			IOUtils.closeQuietly(cancelStreamRegistryForRestore);
		}
	}
}
 

@Nullable
static StreamStateHandle readSnapshot(TaskStateManager stateManager, OperatorID operatorId) {
	final PrioritizedOperatorSubtaskState poss = stateManager.prioritizedOperatorState(operatorId);
	if (!poss.isRestored()) {
		return null;
	}

	final StateObjectCollection<OperatorStateHandle> opState =
		stateManager.prioritizedOperatorState(operatorId).getPrioritizedManagedOperatorState().get(0);
	final OperatorStateHandle handle = Iterators.getOnlyElement(opState.iterator());
	return handle.getDelegateStateHandle();
}
 

@Nonnull
@Override
public PrioritizedOperatorSubtaskState prioritizedOperatorState(OperatorID operatorID) {
	TaskStateSnapshot jmTaskStateSnapshot = getLastJobManagerTaskStateSnapshot();
	TaskStateSnapshot tmTaskStateSnapshot = getLastTaskManagerTaskStateSnapshot();

	if (jmTaskStateSnapshot == null) {

		return PrioritizedOperatorSubtaskState.emptyNotRestored();
	} else {

		OperatorSubtaskState jmOpState = jmTaskStateSnapshot.getSubtaskStateByOperatorID(operatorID);

		if (jmOpState == null) {

			return PrioritizedOperatorSubtaskState.emptyNotRestored();
		} else {

			List<OperatorSubtaskState> tmStateCollection = Collections.emptyList();

			if (tmTaskStateSnapshot != null) {
				OperatorSubtaskState tmOpState = tmTaskStateSnapshot.getSubtaskStateByOperatorID(operatorID);
				if (tmOpState != null) {
					tmStateCollection = Collections.singletonList(tmOpState);
				}
			}
			PrioritizedOperatorSubtaskState.Builder builder =
				new PrioritizedOperatorSubtaskState.Builder(jmOpState, tmStateCollection);
			return builder.build();
		}
	}
}
 

public Builder(RuntimeContext ctx, int maxParallelism) {
	this.ctx = Preconditions.checkNotNull(ctx);

	Preconditions.checkArgument(maxParallelism > 0);
	this.maxParallelism = maxParallelism;

	this.prioritizedOperatorSubtaskState = PrioritizedOperatorSubtaskState.emptyNotRestored();
	this.configuration = new Configuration();
	this.indexOfSubtask = ctx.getIndexOfThisSubtask();
}
 

protected OperatorStateBackend operatorStateBackend(
	String operatorIdentifierText,
	PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates,
	CloseableRegistry backendCloseableRegistry) throws Exception {

	String logDescription = "operator state backend for " + operatorIdentifierText;

	// Now restore processing is included in backend building/constructing process, so we need to make sure
	// each stream constructed in restore could also be closed in case of task cancel, for example the data
	// input stream opened for serDe during restore.
	CloseableRegistry cancelStreamRegistryForRestore = new CloseableRegistry();
	backendCloseableRegistry.registerCloseable(cancelStreamRegistryForRestore);
	BackendRestorerProcedure<OperatorStateBackend, OperatorStateHandle> backendRestorer =
		new BackendRestorerProcedure<>(
			(stateHandles) -> stateBackend.createOperatorStateBackend(
				environment,
				operatorIdentifierText,
				stateHandles,
				cancelStreamRegistryForRestore),
			backendCloseableRegistry,
			logDescription);

	try {
		return backendRestorer.createAndRestore(
			prioritizedOperatorSubtaskStates.getPrioritizedManagedOperatorState());
	} finally {
		if (backendCloseableRegistry.unregisterCloseable(cancelStreamRegistryForRestore)) {
			IOUtils.closeQuietly(cancelStreamRegistryForRestore);
		}
	}
}
 

public PrioritizedOperatorSubtaskState getPrioritizedOperatorSubtaskState() {
	return new PrioritizedOperatorSubtaskState.Builder(
		new OperatorSubtaskState(
			StateObjectCollection.empty(),
			StateObjectCollection.empty(),
			new StateObjectCollection<>(managedKeyedState),
			new StateObjectCollection<>(rawKeyedState)
		),
		Collections.emptyList()
	).build();
}
 

public PrioritizedOperatorSubtaskState getPrioritizedOperatorSubtaskState() {
	return new PrioritizedOperatorSubtaskState.Builder(
		new OperatorSubtaskState(
			StateObjectCollection.empty(),
			StateObjectCollection.empty(),
			new StateObjectCollection<>(managedKeyedState),
			new StateObjectCollection<>(rawKeyedState),
			StateObjectCollection.empty(),
			StateObjectCollection.empty()
		),
		Collections.emptyList()
	).build();
}
 

public Builder(RuntimeContext ctx, int maxParallelism) {
	this.ctx = Preconditions.checkNotNull(ctx);

	Preconditions.checkArgument(maxParallelism > 0);
	this.maxParallelism = maxParallelism;

	this.prioritizedOperatorSubtaskState = PrioritizedOperatorSubtaskState.emptyNotRestored();
	this.configuration = new Configuration();
	this.indexOfSubtask = ctx.getIndexOfThisSubtask();
}
 

@Nonnull
@Override
public PrioritizedOperatorSubtaskState prioritizedOperatorState(OperatorID operatorID) {
	TaskStateSnapshot jmTaskStateSnapshot = getLastJobManagerTaskStateSnapshot();
	TaskStateSnapshot tmTaskStateSnapshot = getLastTaskManagerTaskStateSnapshot();

	if (jmTaskStateSnapshot == null) {

		return PrioritizedOperatorSubtaskState.emptyNotRestored();
	} else {

		OperatorSubtaskState jmOpState = jmTaskStateSnapshot.getSubtaskStateByOperatorID(operatorID);

		if (jmOpState == null) {

			return PrioritizedOperatorSubtaskState.emptyNotRestored();
		} else {

			List<OperatorSubtaskState> tmStateCollection = Collections.emptyList();

			if (tmTaskStateSnapshot != null) {
				OperatorSubtaskState tmOpState = tmTaskStateSnapshot.getSubtaskStateByOperatorID(operatorID);
				if (tmOpState != null) {
					tmStateCollection = Collections.singletonList(tmOpState);
				}
			}
			PrioritizedOperatorSubtaskState.Builder builder =
				new PrioritizedOperatorSubtaskState.Builder(jmOpState, tmStateCollection);
			return builder.build();
		}
	}
}
 

@Nonnull
@Override
public PrioritizedOperatorSubtaskState prioritizedOperatorState(OperatorID operatorID) {
	return prioritizedOperatorSubtaskState;
}
 

@Override
public StreamOperatorStateContext streamOperatorStateContext(
	@Nonnull OperatorID operatorID,
	@Nonnull String operatorClassName,
	@Nonnull KeyContext keyContext,
	@Nullable TypeSerializer<?> keySerializer,
	@Nonnull CloseableRegistry streamTaskCloseableRegistry,
	@Nonnull MetricGroup metricGroup) throws Exception {

	TaskInfo taskInfo = environment.getTaskInfo();
	OperatorSubtaskDescriptionText operatorSubtaskDescription =
		new OperatorSubtaskDescriptionText(
			operatorID,
			operatorClassName,
			taskInfo.getIndexOfThisSubtask(),
			taskInfo.getNumberOfParallelSubtasks());

	final String operatorIdentifierText = operatorSubtaskDescription.toString();

	final PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates =
		taskStateManager.prioritizedOperatorState(operatorID);

	AbstractKeyedStateBackend<?> keyedStatedBackend = null;
	OperatorStateBackend operatorStateBackend = null;
	CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs = null;
	CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs = null;
	InternalTimeServiceManager<?> timeServiceManager;

	try {

		// -------------- Keyed State Backend --------------
		keyedStatedBackend = keyedStatedBackend(
			keySerializer,
			operatorIdentifierText,
			prioritizedOperatorSubtaskStates,
			streamTaskCloseableRegistry,
			metricGroup);

		// -------------- Operator State Backend --------------
		operatorStateBackend = operatorStateBackend(
			operatorIdentifierText,
			prioritizedOperatorSubtaskStates,
			streamTaskCloseableRegistry);

		// -------------- Raw State Streams --------------
		rawKeyedStateInputs = rawKeyedStateInputs(
			prioritizedOperatorSubtaskStates.getPrioritizedRawKeyedState().iterator());
		streamTaskCloseableRegistry.registerCloseable(rawKeyedStateInputs);

		rawOperatorStateInputs = rawOperatorStateInputs(
			prioritizedOperatorSubtaskStates.getPrioritizedRawOperatorState().iterator());
		streamTaskCloseableRegistry.registerCloseable(rawOperatorStateInputs);

		// -------------- Internal Timer Service Manager --------------
		timeServiceManager = internalTimeServiceManager(keyedStatedBackend, keyContext, rawKeyedStateInputs);

		// -------------- Preparing return value --------------

		return new StreamOperatorStateContextImpl(
			prioritizedOperatorSubtaskStates.isRestored(),
			operatorStateBackend,
			keyedStatedBackend,
			timeServiceManager,
			rawOperatorStateInputs,
			rawKeyedStateInputs);
	} catch (Exception ex) {

		// cleanup if something went wrong before results got published.
		if (keyedStatedBackend != null) {
			if (streamTaskCloseableRegistry.unregisterCloseable(keyedStatedBackend)) {
				IOUtils.closeQuietly(keyedStatedBackend);
			}
			// release resource (e.g native resource)
			keyedStatedBackend.dispose();
		}

		if (operatorStateBackend != null) {
			if (streamTaskCloseableRegistry.unregisterCloseable(operatorStateBackend)) {
				IOUtils.closeQuietly(operatorStateBackend);
			}
			operatorStateBackend.dispose();
		}

		if (streamTaskCloseableRegistry.unregisterCloseable(rawKeyedStateInputs)) {
			IOUtils.closeQuietly(rawKeyedStateInputs);
		}

		if (streamTaskCloseableRegistry.unregisterCloseable(rawOperatorStateInputs)) {
			IOUtils.closeQuietly(rawOperatorStateInputs);
		}

		throw new Exception("Exception while creating StreamOperatorStateContext.", ex);
	}
}
 

public Builder setPrioritizedOperatorSubtaskState(PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskState) {
	this.prioritizedOperatorSubtaskState = prioritizedOperatorSubtaskState;
	return this;
}
 

@Nonnull
@Override
public PrioritizedOperatorSubtaskState prioritizedOperatorState(OperatorID operatorID) {

	if (jobManagerTaskRestore == null) {
		return PrioritizedOperatorSubtaskState.emptyNotRestored();
	}

	TaskStateSnapshot jobManagerStateSnapshot =
		jobManagerTaskRestore.getTaskStateSnapshot();

	OperatorSubtaskState jobManagerSubtaskState =
		jobManagerStateSnapshot.getSubtaskStateByOperatorID(operatorID);

	if (jobManagerSubtaskState == null) {
		return PrioritizedOperatorSubtaskState.emptyNotRestored();
	}

	long restoreCheckpointId = jobManagerTaskRestore.getRestoreCheckpointId();

	TaskStateSnapshot localStateSnapshot =
		localStateStore.retrieveLocalState(restoreCheckpointId);

	localStateStore.pruneMatchingCheckpoints((long checkpointId) -> checkpointId != restoreCheckpointId);

	List<OperatorSubtaskState> alternativesByPriority = Collections.emptyList();

	if (localStateSnapshot != null) {
		OperatorSubtaskState localSubtaskState = localStateSnapshot.getSubtaskStateByOperatorID(operatorID);

		if (localSubtaskState != null) {
			alternativesByPriority = Collections.singletonList(localSubtaskState);
		}
	}

	LOG.debug("Operator {} has remote state {} from job manager and local state alternatives {} from local " +
			"state store {}.", operatorID, jobManagerSubtaskState, alternativesByPriority, localStateStore);

	PrioritizedOperatorSubtaskState.Builder builder = new PrioritizedOperatorSubtaskState.Builder(
		jobManagerSubtaskState,
		alternativesByPriority,
		true);

	return builder.build();
}
 

@Override
public StreamOperatorStateContext streamOperatorStateContext(
	@Nonnull OperatorID operatorID,
	@Nonnull String operatorClassName,
	@Nonnull ProcessingTimeService processingTimeService,
	@Nonnull KeyContext keyContext,
	@Nullable TypeSerializer<?> keySerializer,
	@Nonnull CloseableRegistry streamTaskCloseableRegistry,
	@Nonnull MetricGroup metricGroup) throws Exception {

	TaskInfo taskInfo = environment.getTaskInfo();
	OperatorSubtaskDescriptionText operatorSubtaskDescription =
		new OperatorSubtaskDescriptionText(
			operatorID,
			operatorClassName,
			taskInfo.getIndexOfThisSubtask(),
			taskInfo.getNumberOfParallelSubtasks());

	final String operatorIdentifierText = operatorSubtaskDescription.toString();

	final PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates =
		taskStateManager.prioritizedOperatorState(operatorID);

	AbstractKeyedStateBackend<?> keyedStatedBackend = null;
	OperatorStateBackend operatorStateBackend = null;
	CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs = null;
	CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs = null;
	InternalTimeServiceManager<?> timeServiceManager;

	try {

		// -------------- Keyed State Backend --------------
		keyedStatedBackend = keyedStatedBackend(
			keySerializer,
			operatorIdentifierText,
			prioritizedOperatorSubtaskStates,
			streamTaskCloseableRegistry,
			metricGroup);

		// -------------- Operator State Backend --------------
		operatorStateBackend = operatorStateBackend(
			operatorIdentifierText,
			prioritizedOperatorSubtaskStates,
			streamTaskCloseableRegistry);

		// -------------- Raw State Streams --------------
		rawKeyedStateInputs = rawKeyedStateInputs(
			prioritizedOperatorSubtaskStates.getPrioritizedRawKeyedState().iterator());
		streamTaskCloseableRegistry.registerCloseable(rawKeyedStateInputs);

		rawOperatorStateInputs = rawOperatorStateInputs(
			prioritizedOperatorSubtaskStates.getPrioritizedRawOperatorState().iterator());
		streamTaskCloseableRegistry.registerCloseable(rawOperatorStateInputs);

		// -------------- Internal Timer Service Manager --------------
		timeServiceManager = internalTimeServiceManager(keyedStatedBackend, keyContext, processingTimeService, rawKeyedStateInputs);

		// -------------- Preparing return value --------------

		return new StreamOperatorStateContextImpl(
			prioritizedOperatorSubtaskStates.isRestored(),
			operatorStateBackend,
			keyedStatedBackend,
			timeServiceManager,
			rawOperatorStateInputs,
			rawKeyedStateInputs);
	} catch (Exception ex) {

		// cleanup if something went wrong before results got published.
		if (keyedStatedBackend != null) {
			if (streamTaskCloseableRegistry.unregisterCloseable(keyedStatedBackend)) {
				IOUtils.closeQuietly(keyedStatedBackend);
			}
			// release resource (e.g native resource)
			keyedStatedBackend.dispose();
		}

		if (operatorStateBackend != null) {
			if (streamTaskCloseableRegistry.unregisterCloseable(operatorStateBackend)) {
				IOUtils.closeQuietly(operatorStateBackend);
			}
			operatorStateBackend.dispose();
		}

		if (streamTaskCloseableRegistry.unregisterCloseable(rawKeyedStateInputs)) {
			IOUtils.closeQuietly(rawKeyedStateInputs);
		}

		if (streamTaskCloseableRegistry.unregisterCloseable(rawOperatorStateInputs)) {
			IOUtils.closeQuietly(rawOperatorStateInputs);
		}

		throw new Exception("Exception while creating StreamOperatorStateContext.", ex);
	}
}
 

protected <K> AbstractKeyedStateBackend<K> keyedStatedBackend(
	TypeSerializer<K> keySerializer,
	String operatorIdentifierText,
	PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates,
	CloseableRegistry backendCloseableRegistry,
	MetricGroup metricGroup) throws Exception {

	if (keySerializer == null) {
		return null;
	}

	String logDescription = "keyed state backend for " + operatorIdentifierText;

	TaskInfo taskInfo = environment.getTaskInfo();

	final KeyGroupRange keyGroupRange = KeyGroupRangeAssignment.computeKeyGroupRangeForOperatorIndex(
		taskInfo.getMaxNumberOfParallelSubtasks(),
		taskInfo.getNumberOfParallelSubtasks(),
		taskInfo.getIndexOfThisSubtask());

	// Now restore processing is included in backend building/constructing process, so we need to make sure
	// each stream constructed in restore could also be closed in case of task cancel, for example the data
	// input stream opened for serDe during restore.
	CloseableRegistry cancelStreamRegistryForRestore = new CloseableRegistry();
	backendCloseableRegistry.registerCloseable(cancelStreamRegistryForRestore);
	BackendRestorerProcedure<AbstractKeyedStateBackend<K>, KeyedStateHandle> backendRestorer =
		new BackendRestorerProcedure<>(
			(stateHandles) -> stateBackend.createKeyedStateBackend(
				environment,
				environment.getJobID(),
				operatorIdentifierText,
				keySerializer,
				taskInfo.getMaxNumberOfParallelSubtasks(),
				keyGroupRange,
				environment.getTaskKvStateRegistry(),
				ttlTimeProvider,
				metricGroup,
				stateHandles,
				cancelStreamRegistryForRestore),
			backendCloseableRegistry,
			logDescription);

	try {
		return backendRestorer.createAndRestore(
			prioritizedOperatorSubtaskStates.getPrioritizedManagedKeyedState());
	} finally {
		if (backendCloseableRegistry.unregisterCloseable(cancelStreamRegistryForRestore)) {
			IOUtils.closeQuietly(cancelStreamRegistryForRestore);
		}
	}
}
 

SavepointTaskStateManager(PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskState) {
	Preconditions.checkNotNull(prioritizedOperatorSubtaskState, "Operator subtask state must not be null");
	this.prioritizedOperatorSubtaskState = prioritizedOperatorSubtaskState;
}
 

protected <K> AbstractKeyedStateBackend<K> keyedStatedBackend(
	TypeSerializer<K> keySerializer,
	String operatorIdentifierText,
	PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates,
	CloseableRegistry backendCloseableRegistry,
	MetricGroup metricGroup) throws Exception {

	if (keySerializer == null) {
		return null;
	}

	String logDescription = "keyed state backend for " + operatorIdentifierText;

	TaskInfo taskInfo = environment.getTaskInfo();

	final KeyGroupRange keyGroupRange = KeyGroupRangeAssignment.computeKeyGroupRangeForOperatorIndex(
		taskInfo.getMaxNumberOfParallelSubtasks(),
		taskInfo.getNumberOfParallelSubtasks(),
		taskInfo.getIndexOfThisSubtask());

	// Now restore processing is included in backend building/constructing process, so we need to make sure
	// each stream constructed in restore could also be closed in case of task cancel, for example the data
	// input stream opened for serDe during restore.
	CloseableRegistry cancelStreamRegistryForRestore = new CloseableRegistry();
	backendCloseableRegistry.registerCloseable(cancelStreamRegistryForRestore);
	BackendRestorerProcedure<AbstractKeyedStateBackend<K>, KeyedStateHandle> backendRestorer =
		new BackendRestorerProcedure<>(
			(stateHandles) -> stateBackend.createKeyedStateBackend(
				environment,
				environment.getJobID(),
				operatorIdentifierText,
				keySerializer,
				taskInfo.getMaxNumberOfParallelSubtasks(),
				keyGroupRange,
				environment.getTaskKvStateRegistry(),
				TtlTimeProvider.DEFAULT,
				metricGroup,
				stateHandles,
				cancelStreamRegistryForRestore),
			backendCloseableRegistry,
			logDescription);

	try {
		return backendRestorer.createAndRestore(
			prioritizedOperatorSubtaskStates.getPrioritizedManagedKeyedState());
	} finally {
		if (backendCloseableRegistry.unregisterCloseable(cancelStreamRegistryForRestore)) {
			IOUtils.closeQuietly(cancelStreamRegistryForRestore);
		}
	}
}
 

@Nonnull
@Override
public PrioritizedOperatorSubtaskState prioritizedOperatorState(OperatorID operatorID) {

	if (jobManagerTaskRestore == null) {
		return PrioritizedOperatorSubtaskState.emptyNotRestored();
	}

	TaskStateSnapshot jobManagerStateSnapshot =
		jobManagerTaskRestore.getTaskStateSnapshot();

	OperatorSubtaskState jobManagerSubtaskState =
		jobManagerStateSnapshot.getSubtaskStateByOperatorID(operatorID);

	if (jobManagerSubtaskState == null) {
		return PrioritizedOperatorSubtaskState.emptyNotRestored();
	}

	long restoreCheckpointId = jobManagerTaskRestore.getRestoreCheckpointId();

	TaskStateSnapshot localStateSnapshot =
		localStateStore.retrieveLocalState(restoreCheckpointId);

	localStateStore.pruneMatchingCheckpoints((long checkpointId) -> checkpointId != restoreCheckpointId);

	List<OperatorSubtaskState> alternativesByPriority = Collections.emptyList();

	if (localStateSnapshot != null) {
		OperatorSubtaskState localSubtaskState = localStateSnapshot.getSubtaskStateByOperatorID(operatorID);

		if (localSubtaskState != null) {
			alternativesByPriority = Collections.singletonList(localSubtaskState);
		}
	}

	LOG.debug("Operator {} has remote state {} from job manager and local state alternatives {} from local " +
			"state store {}.", operatorID, jobManagerSubtaskState, alternativesByPriority, localStateStore);

	PrioritizedOperatorSubtaskState.Builder builder = new PrioritizedOperatorSubtaskState.Builder(
		jobManagerSubtaskState,
		alternativesByPriority,
		true);

	return builder.build();
}
 

public Builder setPrioritizedOperatorSubtaskState(PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskState) {
	this.prioritizedOperatorSubtaskState = prioritizedOperatorSubtaskState;
	return this;
}
 

@Nonnull
@Override
public PrioritizedOperatorSubtaskState prioritizedOperatorState(OperatorID operatorID) {
	return prioritizedOperatorSubtaskState;
}
 

SavepointTaskStateManager(PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskState) {
	Preconditions.checkNotNull(prioritizedOperatorSubtaskState, "Operator subtask state must not be null");
	this.prioritizedOperatorSubtaskState = prioritizedOperatorSubtaskState;
}
 

protected <K> AbstractKeyedStateBackend<K> keyedStatedBackend(
	TypeSerializer<K> keySerializer,
	String operatorIdentifierText,
	PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates,
	CloseableRegistry backendCloseableRegistry,
	MetricGroup metricGroup) throws Exception {

	if (keySerializer == null) {
		return null;
	}

	String logDescription = "keyed state backend for " + operatorIdentifierText;

	TaskInfo taskInfo = environment.getTaskInfo();

	final KeyGroupRange keyGroupRange = KeyGroupRangeAssignment.computeKeyGroupRangeForOperatorIndex(
		taskInfo.getMaxNumberOfParallelSubtasks(),
		taskInfo.getNumberOfParallelSubtasks(),
		taskInfo.getIndexOfThisSubtask());

	// Now restore processing is included in backend building/constructing process, so we need to make sure
	// each stream constructed in restore could also be closed in case of task cancel, for example the data
	// input stream opened for serDe during restore.
	CloseableRegistry cancelStreamRegistryForRestore = new CloseableRegistry();
	backendCloseableRegistry.registerCloseable(cancelStreamRegistryForRestore);
	BackendRestorerProcedure<AbstractKeyedStateBackend<K>, KeyedStateHandle> backendRestorer =
		new BackendRestorerProcedure<>(
			(stateHandles) -> stateBackend.createKeyedStateBackend(
				environment,
				environment.getJobID(),
				operatorIdentifierText,
				keySerializer,
				taskInfo.getMaxNumberOfParallelSubtasks(),
				keyGroupRange,
				environment.getTaskKvStateRegistry(),
				TtlTimeProvider.DEFAULT,
				metricGroup,
				stateHandles,
				cancelStreamRegistryForRestore),
			backendCloseableRegistry,
			logDescription);

	try {
		return backendRestorer.createAndRestore(
			prioritizedOperatorSubtaskStates.getPrioritizedManagedKeyedState());
	} finally {
		if (backendCloseableRegistry.unregisterCloseable(cancelStreamRegistryForRestore)) {
			IOUtils.closeQuietly(cancelStreamRegistryForRestore);
		}
	}
}
 

@Override
public StreamOperatorStateContext streamOperatorStateContext(
	@Nonnull OperatorID operatorID,
	@Nonnull String operatorClassName,
	@Nonnull KeyContext keyContext,
	@Nullable TypeSerializer<?> keySerializer,
	@Nonnull CloseableRegistry streamTaskCloseableRegistry,
	@Nonnull MetricGroup metricGroup) throws Exception {

	TaskInfo taskInfo = environment.getTaskInfo();
	OperatorSubtaskDescriptionText operatorSubtaskDescription =
		new OperatorSubtaskDescriptionText(
			operatorID,
			operatorClassName,
			taskInfo.getIndexOfThisSubtask(),
			taskInfo.getNumberOfParallelSubtasks());

	final String operatorIdentifierText = operatorSubtaskDescription.toString();

	final PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates =
		taskStateManager.prioritizedOperatorState(operatorID);

	AbstractKeyedStateBackend<?> keyedStatedBackend = null;
	OperatorStateBackend operatorStateBackend = null;
	CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs = null;
	CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs = null;
	InternalTimeServiceManager<?> timeServiceManager;

	try {

		// -------------- Keyed State Backend --------------
		keyedStatedBackend = keyedStatedBackend(
			keySerializer,
			operatorIdentifierText,
			prioritizedOperatorSubtaskStates,
			streamTaskCloseableRegistry,
			metricGroup);

		// -------------- Operator State Backend --------------
		operatorStateBackend = operatorStateBackend(
			operatorIdentifierText,
			prioritizedOperatorSubtaskStates,
			streamTaskCloseableRegistry);

		// -------------- Raw State Streams --------------
		rawKeyedStateInputs = rawKeyedStateInputs(
			prioritizedOperatorSubtaskStates.getPrioritizedRawKeyedState().iterator());
		streamTaskCloseableRegistry.registerCloseable(rawKeyedStateInputs);

		rawOperatorStateInputs = rawOperatorStateInputs(
			prioritizedOperatorSubtaskStates.getPrioritizedRawOperatorState().iterator());
		streamTaskCloseableRegistry.registerCloseable(rawOperatorStateInputs);

		// -------------- Internal Timer Service Manager --------------
		timeServiceManager = internalTimeServiceManager(keyedStatedBackend, keyContext, rawKeyedStateInputs);

		// -------------- Preparing return value --------------

		return new StreamOperatorStateContextImpl(
			prioritizedOperatorSubtaskStates.isRestored(),
			operatorStateBackend,
			keyedStatedBackend,
			timeServiceManager,
			rawOperatorStateInputs,
			rawKeyedStateInputs);
	} catch (Exception ex) {

		// cleanup if something went wrong before results got published.
		if (keyedStatedBackend != null) {
			if (streamTaskCloseableRegistry.unregisterCloseable(keyedStatedBackend)) {
				IOUtils.closeQuietly(keyedStatedBackend);
			}
			// release resource (e.g native resource)
			keyedStatedBackend.dispose();
		}

		if (operatorStateBackend != null) {
			if (streamTaskCloseableRegistry.unregisterCloseable(operatorStateBackend)) {
				IOUtils.closeQuietly(operatorStateBackend);
			}
			operatorStateBackend.dispose();
		}

		if (streamTaskCloseableRegistry.unregisterCloseable(rawKeyedStateInputs)) {
			IOUtils.closeQuietly(rawKeyedStateInputs);
		}

		if (streamTaskCloseableRegistry.unregisterCloseable(rawOperatorStateInputs)) {
			IOUtils.closeQuietly(rawOperatorStateInputs);
		}

		throw new Exception("Exception while creating StreamOperatorStateContext.", ex);
	}
}
 

@Nonnull
@Override
public PrioritizedOperatorSubtaskState prioritizedOperatorState(OperatorID operatorID) {

	if (jobManagerTaskRestore == null) {
		return PrioritizedOperatorSubtaskState.emptyNotRestored();
	}

	TaskStateSnapshot jobManagerStateSnapshot =
		jobManagerTaskRestore.getTaskStateSnapshot();

	OperatorSubtaskState jobManagerSubtaskState =
		jobManagerStateSnapshot.getSubtaskStateByOperatorID(operatorID);

	if (jobManagerSubtaskState == null) {
		return PrioritizedOperatorSubtaskState.emptyNotRestored();
	}

	long restoreCheckpointId = jobManagerTaskRestore.getRestoreCheckpointId();

	TaskStateSnapshot localStateSnapshot =
		localStateStore.retrieveLocalState(restoreCheckpointId);

	localStateStore.pruneMatchingCheckpoints((long checkpointId) -> checkpointId != restoreCheckpointId);

	List<OperatorSubtaskState> alternativesByPriority = Collections.emptyList();

	if (localStateSnapshot != null) {
		OperatorSubtaskState localSubtaskState = localStateSnapshot.getSubtaskStateByOperatorID(operatorID);

		if (localSubtaskState != null) {
			alternativesByPriority = Collections.singletonList(localSubtaskState);
		}
	}

	LOG.debug("Operator {} has remote state {} from job manager and local state alternatives {} from local " +
			"state store {}.", operatorID, jobManagerSubtaskState, alternativesByPriority, localStateStore);

	PrioritizedOperatorSubtaskState.Builder builder = new PrioritizedOperatorSubtaskState.Builder(
		jobManagerSubtaskState,
		alternativesByPriority,
		true);

	return builder.build();
}
 

/**
 * Returns means to restore previously reported state of an operator running in the owning task.
 *
 * @param operatorID the id of the operator for which we request state.
 * @return Previous state for the operator. The previous state can be empty if the operator had no previous state.
 */
@Nonnull
PrioritizedOperatorSubtaskState prioritizedOperatorState(OperatorID operatorID);
 

/**
 * Returns means to restore previously reported state of an operator running in the owning task.
 *
 * @param operatorID the id of the operator for which we request state.
 * @return Previous state for the operator. The previous state can be empty if the operator had no previous state.
 */
@Nonnull
PrioritizedOperatorSubtaskState prioritizedOperatorState(OperatorID operatorID);