Java Code Examples for org.apache.flink.core.fs.CloseableRegistry#registerCloseable()

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

@VisibleForTesting
SubtaskCheckpointCoordinatorImpl(
		CheckpointStorageWorkerView checkpointStorage,
		String taskName,
		StreamTaskActionExecutor actionExecutor,
		CloseableRegistry closeableRegistry,
		ExecutorService executorService,
		Environment env,
		AsyncExceptionHandler asyncExceptionHandler,
		BiFunctionWithException<ChannelStateWriter, Long, CompletableFuture<Void>, IOException> prepareInputSnapshot,
		int maxRecordAbortedCheckpoints,
		ChannelStateWriter channelStateWriter) throws IOException {
	this.checkpointStorage = new CachingCheckpointStorageWorkerView(checkNotNull(checkpointStorage));
	this.taskName = checkNotNull(taskName);
	this.checkpoints = new HashMap<>();
	this.lock = new Object();
	this.executorService = checkNotNull(executorService);
	this.env = checkNotNull(env);
	this.asyncExceptionHandler = checkNotNull(asyncExceptionHandler);
	this.actionExecutor = checkNotNull(actionExecutor);
	this.channelStateWriter = checkNotNull(channelStateWriter);
	this.prepareInputSnapshot = prepareInputSnapshot;
	this.abortedCheckpointIds = createAbortedCheckpointSetWithLimitSize(maxRecordAbortedCheckpoints);
	this.lastCheckpointId = -1L;
	closeableRegistry.registerCloseable(this);
	this.closed = false;
}
 

/**
 * Copies the file from a single state handle to the given path.
 */
private void downloadDataForStateHandle(
	Path restoreFilePath,
	StreamStateHandle remoteFileHandle,
	CloseableRegistry closeableRegistry) throws IOException {

	FSDataInputStream inputStream = null;
	FSDataOutputStream outputStream = null;

	try {
		FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
		inputStream = remoteFileHandle.openInputStream();
		closeableRegistry.registerCloseable(inputStream);

		outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
		closeableRegistry.registerCloseable(outputStream);

		byte[] buffer = new byte[8 * 1024];
		while (true) {
			int numBytes = inputStream.read(buffer);
			if (numBytes == -1) {
				break;
			}

			outputStream.write(buffer, 0, numBytes);
		}
	} finally {
		if (closeableRegistry.unregisterCloseable(inputStream)) {
			inputStream.close();
		}

		if (closeableRegistry.unregisterCloseable(outputStream)) {
			outputStream.close();
		}
	}
}
 

/**
 * Copies the file from a single state handle to the given path.
 */
private void downloadDataForStateHandle(
	Path restoreFilePath,
	StreamStateHandle remoteFileHandle,
	CloseableRegistry closeableRegistry) throws IOException {

	FSDataInputStream inputStream = null;
	OutputStream outputStream = null;

	try {
		inputStream = remoteFileHandle.openInputStream();
		closeableRegistry.registerCloseable(inputStream);

		Files.createDirectories(restoreFilePath.getParent());
		outputStream = Files.newOutputStream(restoreFilePath);
		closeableRegistry.registerCloseable(outputStream);

		byte[] buffer = new byte[8 * 1024];
		while (true) {
			int numBytes = inputStream.read(buffer);
			if (numBytes == -1) {
				break;
			}

			outputStream.write(buffer, 0, numBytes);
		}
	} finally {
		if (closeableRegistry.unregisterCloseable(inputStream)) {
			inputStream.close();
		}

		if (closeableRegistry.unregisterCloseable(outputStream)) {
			outputStream.close();
		}
	}
}
 

/**
 * Copies the file from a single state handle to the given path.
 */
private void downloadDataForStateHandle(
	Path restoreFilePath,
	StreamStateHandle remoteFileHandle,
	CloseableRegistry closeableRegistry) throws IOException {

	FSDataInputStream inputStream = null;
	FSDataOutputStream outputStream = null;

	try {
		FileSystem restoreFileSystem = restoreFilePath.getFileSystem();
		inputStream = remoteFileHandle.openInputStream();
		closeableRegistry.registerCloseable(inputStream);

		outputStream = restoreFileSystem.create(restoreFilePath, FileSystem.WriteMode.OVERWRITE);
		closeableRegistry.registerCloseable(outputStream);

		byte[] buffer = new byte[8 * 1024];
		while (true) {
			int numBytes = inputStream.read(buffer);
			if (numBytes == -1) {
				break;
			}

			outputStream.write(buffer, 0, numBytes);
		}
	} finally {
		if (closeableRegistry.unregisterCloseable(inputStream)) {
			inputStream.close();
		}

		if (closeableRegistry.unregisterCloseable(outputStream)) {
			outputStream.close();
		}
	}
}
 

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

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

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

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

private AsyncSnapshotTask(@Nonnull CloseableRegistry taskRegistry) throws IOException {
	super(AsyncSnapshotCallable.this);
	this.cancelOnClose = () -> cancel(true);
	this.taskRegistry = taskRegistry;
	taskRegistry.registerCloseable(cancelOnClose);
}
 

private StreamStateHandle uploadLocalFileToCheckpointFs(
	Path filePath,
	CheckpointStreamFactory checkpointStreamFactory,
	CloseableRegistry closeableRegistry) throws IOException {

	InputStream inputStream = null;
	CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;

	try {
		final byte[] buffer = new byte[READ_BUFFER_SIZE];

		inputStream = Files.newInputStream(filePath);
		closeableRegistry.registerCloseable(inputStream);

		outputStream = checkpointStreamFactory
			.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
		closeableRegistry.registerCloseable(outputStream);

		while (true) {
			int numBytes = inputStream.read(buffer);

			if (numBytes == -1) {
				break;
			}

			outputStream.write(buffer, 0, numBytes);
		}

		StreamStateHandle result = null;
		if (closeableRegistry.unregisterCloseable(outputStream)) {
			result = outputStream.closeAndGetHandle();
			outputStream = null;
		}
		return result;

	} finally {

		if (closeableRegistry.unregisterCloseable(inputStream)) {
			IOUtils.closeQuietly(inputStream);
		}

		if (closeableRegistry.unregisterCloseable(outputStream)) {
			IOUtils.closeQuietly(outputStream);
		}
	}
}
 

/**
 * Tests that a failing snapshot method call to the keyed state backend will trigger the closing
 * of the StateSnapshotContextSynchronousImpl and the cancellation of the
 * OperatorSnapshotResult. The latter is supposed to also cancel all assigned futures.
 */
@Test
public void testFailingBackendSnapshotMethod() throws Exception {
	final long checkpointId = 42L;
	final long timestamp = 1L;

	final Exception failingException = new Exception("Test exception");

	final CloseableRegistry closeableRegistry = new CloseableRegistry();

	RunnableFuture<SnapshotResult<KeyedStateHandle>> futureKeyedStateHandle = mock(RunnableFuture.class);
	RunnableFuture<SnapshotResult<OperatorStateHandle>> futureOperatorStateHandle = mock(RunnableFuture.class);

	StateSnapshotContextSynchronousImpl context = spy(new StateSnapshotContextSynchronousImpl(checkpointId, timestamp));
	when(context.getKeyedStateStreamFuture()).thenReturn(futureKeyedStateHandle);
	when(context.getOperatorStateStreamFuture()).thenReturn(futureOperatorStateHandle);

	OperatorSnapshotFutures operatorSnapshotResult = spy(new OperatorSnapshotFutures());

	whenNew(StateSnapshotContextSynchronousImpl.class)
		.withArguments(
			anyLong(),
			anyLong(),
			any(CheckpointStreamFactory.class),
			nullable(KeyGroupRange.class),
			any(CloseableRegistry.class))
		.thenReturn(context);
	whenNew(OperatorSnapshotFutures.class).withAnyArguments().thenReturn(operatorSnapshotResult);

	StreamTask<Void, AbstractStreamOperator<Void>> containingTask = mock(StreamTask.class);
	when(containingTask.getCancelables()).thenReturn(closeableRegistry);

	AbstractStreamOperator<Void> operator = mock(AbstractStreamOperator.class);
	when(operator.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class))).thenCallRealMethod();

	doCallRealMethod().when(operator).close();
	doCallRealMethod().when(operator).dispose();

	doReturn(containingTask).when(operator).getContainingTask();

	RunnableFuture<SnapshotResult<OperatorStateHandle>> futureManagedOperatorStateHandle = mock(RunnableFuture.class);

	OperatorStateBackend operatorStateBackend = mock(OperatorStateBackend.class);
	when(operatorStateBackend.snapshot(
		eq(checkpointId),
		eq(timestamp),
		any(CheckpointStreamFactory.class),
		any(CheckpointOptions.class))).thenReturn(futureManagedOperatorStateHandle);

	AbstractKeyedStateBackend<?> keyedStateBackend = mock(AbstractKeyedStateBackend.class);
	when(keyedStateBackend.snapshot(
		eq(checkpointId),
		eq(timestamp),
		any(CheckpointStreamFactory.class),
		eq(CheckpointOptions.forCheckpointWithDefaultLocation()))).thenThrow(failingException);

	closeableRegistry.registerCloseable(operatorStateBackend);
	closeableRegistry.registerCloseable(keyedStateBackend);

	Whitebox.setInternalState(operator, "operatorStateBackend", operatorStateBackend);
	Whitebox.setInternalState(operator, "keyedStateBackend", keyedStateBackend);

	try {
		operator.snapshotState(
				checkpointId,
				timestamp,
				CheckpointOptions.forCheckpointWithDefaultLocation(),
				new MemCheckpointStreamFactory(Integer.MAX_VALUE));
		fail("Exception expected.");
	} catch (Exception e) {
		assertEquals(failingException, e.getCause());
	}

	// verify that the context has been closed, the operator snapshot result has been cancelled
	// and that all futures have been cancelled.
	verify(operatorSnapshotResult).cancel();

	verify(futureKeyedStateHandle).cancel(anyBoolean());
	verify(futureOperatorStateHandle).cancel(anyBoolean());
	verify(futureKeyedStateHandle).cancel(anyBoolean());

	operator.close();

	operator.dispose();

	verify(operatorStateBackend).close();
	verify(keyedStateBackend).close();
	verify(operatorStateBackend).dispose();
	verify(keyedStateBackend).dispose();
}
 

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

private AsyncSnapshotTask(@Nonnull CloseableRegistry taskRegistry) throws IOException {
	super(AsyncSnapshotCallable.this);
	this.cancelOnClose = () -> cancel(true);
	this.taskRegistry = taskRegistry;
	taskRegistry.registerCloseable(cancelOnClose);
}
 

private StreamStateHandle uploadLocalFileToCheckpointFs(
	Path filePath,
	CheckpointStreamFactory checkpointStreamFactory,
	CloseableRegistry closeableRegistry) throws IOException {
	FSDataInputStream inputStream = null;
	CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;

	try {
		final byte[] buffer = new byte[READ_BUFFER_SIZE];

		FileSystem backupFileSystem = filePath.getFileSystem();
		inputStream = backupFileSystem.open(filePath);
		closeableRegistry.registerCloseable(inputStream);

		outputStream = checkpointStreamFactory
			.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
		closeableRegistry.registerCloseable(outputStream);

		while (true) {
			int numBytes = inputStream.read(buffer);

			if (numBytes == -1) {
				break;
			}

			outputStream.write(buffer, 0, numBytes);
		}

		StreamStateHandle result = null;
		if (closeableRegistry.unregisterCloseable(outputStream)) {
			result = outputStream.closeAndGetHandle();
			outputStream = null;
		}
		return result;

	} finally {

		if (closeableRegistry.unregisterCloseable(inputStream)) {
			IOUtils.closeQuietly(inputStream);
		}

		if (closeableRegistry.unregisterCloseable(outputStream)) {
			IOUtils.closeQuietly(outputStream);
		}
	}
}
 

/**
 * Tests that a failing snapshot method call to the keyed state backend will trigger the closing
 * of the StateSnapshotContextSynchronousImpl and the cancellation of the
 * OperatorSnapshotResult. The latter is supposed to also cancel all assigned futures.
 */
@Test
public void testFailingBackendSnapshotMethod() throws Exception {
	final long checkpointId = 42L;
	final long timestamp = 1L;

	final Exception failingException = new Exception("Test exception");

	final CloseableRegistry closeableRegistry = new CloseableRegistry();

	RunnableFuture<SnapshotResult<KeyedStateHandle>> futureKeyedStateHandle = mock(RunnableFuture.class);
	RunnableFuture<SnapshotResult<OperatorStateHandle>> futureOperatorStateHandle = mock(RunnableFuture.class);

	StateSnapshotContextSynchronousImpl context = spy(new StateSnapshotContextSynchronousImpl(checkpointId, timestamp));
	when(context.getKeyedStateStreamFuture()).thenReturn(futureKeyedStateHandle);
	when(context.getOperatorStateStreamFuture()).thenReturn(futureOperatorStateHandle);

	OperatorSnapshotFutures operatorSnapshotResult = spy(new OperatorSnapshotFutures());

	whenNew(StateSnapshotContextSynchronousImpl.class)
		.withArguments(
			anyLong(),
			anyLong(),
			any(CheckpointStreamFactory.class),
			nullable(KeyGroupRange.class),
			any(CloseableRegistry.class))
		.thenReturn(context);
	whenNew(OperatorSnapshotFutures.class).withAnyArguments().thenReturn(operatorSnapshotResult);

	StreamTask<Void, AbstractStreamOperator<Void>> containingTask = mock(StreamTask.class);
	when(containingTask.getCancelables()).thenReturn(closeableRegistry);

	AbstractStreamOperator<Void> operator = mock(AbstractStreamOperator.class);
	when(operator.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class), any(CheckpointStreamFactory.class))).thenCallRealMethod();

	doCallRealMethod().when(operator).close();
	doCallRealMethod().when(operator).dispose();

	doReturn(containingTask).when(operator).getContainingTask();

	RunnableFuture<SnapshotResult<OperatorStateHandle>> futureManagedOperatorStateHandle = mock(RunnableFuture.class);

	OperatorStateBackend operatorStateBackend = mock(OperatorStateBackend.class);
	when(operatorStateBackend.snapshot(
		eq(checkpointId),
		eq(timestamp),
		any(CheckpointStreamFactory.class),
		any(CheckpointOptions.class))).thenReturn(futureManagedOperatorStateHandle);

	AbstractKeyedStateBackend<?> keyedStateBackend = mock(AbstractKeyedStateBackend.class);
	when(keyedStateBackend.snapshot(
		eq(checkpointId),
		eq(timestamp),
		any(CheckpointStreamFactory.class),
		eq(CheckpointOptions.forCheckpointWithDefaultLocation()))).thenThrow(failingException);

	closeableRegistry.registerCloseable(operatorStateBackend);
	closeableRegistry.registerCloseable(keyedStateBackend);

	Whitebox.setInternalState(operator, "operatorStateBackend", operatorStateBackend);
	Whitebox.setInternalState(operator, "keyedStateBackend", keyedStateBackend);

	try {
		operator.snapshotState(
				checkpointId,
				timestamp,
				CheckpointOptions.forCheckpointWithDefaultLocation(),
				new MemCheckpointStreamFactory(Integer.MAX_VALUE));
		fail("Exception expected.");
	} catch (Exception e) {
		assertEquals(failingException, e.getCause());
	}

	// verify that the context has been closed, the operator snapshot result has been cancelled
	// and that all futures have been cancelled.
	verify(context).close();
	verify(operatorSnapshotResult).cancel();

	verify(futureKeyedStateHandle).cancel(anyBoolean());
	verify(futureOperatorStateHandle).cancel(anyBoolean());
	verify(futureKeyedStateHandle).cancel(anyBoolean());

	operator.close();

	operator.dispose();

	verify(operatorStateBackend).close();
	verify(keyedStateBackend).close();
	verify(operatorStateBackend).dispose();
	verify(keyedStateBackend).dispose();
}
 

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

private AsyncSnapshotTask(@Nonnull CloseableRegistry taskRegistry) throws IOException {
	super(AsyncSnapshotCallable.this);
	this.cancelOnClose = () -> cancel(true);
	this.taskRegistry = taskRegistry;
	taskRegistry.registerCloseable(cancelOnClose);
}
 

private StreamStateHandle uploadLocalFileToCheckpointFs(
	Path filePath,
	CheckpointStreamFactory checkpointStreamFactory,
	CloseableRegistry closeableRegistry) throws IOException {
	FSDataInputStream inputStream = null;
	CheckpointStreamFactory.CheckpointStateOutputStream outputStream = null;

	try {
		final byte[] buffer = new byte[READ_BUFFER_SIZE];

		FileSystem backupFileSystem = filePath.getFileSystem();
		inputStream = backupFileSystem.open(filePath);
		closeableRegistry.registerCloseable(inputStream);

		outputStream = checkpointStreamFactory
			.createCheckpointStateOutputStream(CheckpointedStateScope.SHARED);
		closeableRegistry.registerCloseable(outputStream);

		while (true) {
			int numBytes = inputStream.read(buffer);

			if (numBytes == -1) {
				break;
			}

			outputStream.write(buffer, 0, numBytes);
		}

		StreamStateHandle result = null;
		if (closeableRegistry.unregisterCloseable(outputStream)) {
			result = outputStream.closeAndGetHandle();
			outputStream = null;
		}
		return result;

	} finally {

		if (closeableRegistry.unregisterCloseable(inputStream)) {
			IOUtils.closeQuietly(inputStream);
		}

		if (closeableRegistry.unregisterCloseable(outputStream)) {
			IOUtils.closeQuietly(outputStream);
		}
	}
}