Java Code Examples for java.util.concurrent.RunnableFuture#get()

@Test
public void testDismissingSnapshotNotRunnable() throws Exception {
	setupRocksKeyedStateBackend();
	try {
		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot =
			keyedStateBackend.snapshot(0L, 0L, testStreamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());
		snapshot.cancel(true);
		Thread asyncSnapshotThread = new Thread(snapshot);
		asyncSnapshotThread.start();
		try {
			snapshot.get();
			fail();
		} catch (Exception ignored) {

		}
		asyncSnapshotThread.join();
		verifyRocksObjectsReleased();
	} finally {
		this.keyedStateBackend.dispose();
		this.keyedStateBackend = null;
	}
}
 

public static PsiDirectory writeDirectory(PsiDirectory dir, DirectoryWrapper dirWrapper, Project project) {
    if (dir == null) {
        //todo print error
        return null;
    }

    RunnableFuture<PsiDirectory> runnableFuture = new FutureTask<>(() ->
            ApplicationManager.getApplication().runWriteAction(new Computable<PsiDirectory>() {
                @Override
                public PsiDirectory compute() {
                    return writeDirectoryAction(dir, dirWrapper, project);
                }
            }));

    ApplicationManager.getApplication().invokeLater(runnableFuture);

    try {
        return runnableFuture.get();
    } catch (InterruptedException | ExecutionException e) {
        Logger.log("runnableFuture  " + e.getMessage());
        Logger.printStack(e);
    }

    return null;
}
 

protected KeyedStateHandle runSnapshot(
	RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotRunnableFuture,
	SharedStateRegistry sharedStateRegistry) throws Exception {

	if (!snapshotRunnableFuture.isDone()) {
		snapshotRunnableFuture.run();
	}

	SnapshotResult<KeyedStateHandle> snapshotResult = snapshotRunnableFuture.get();
	KeyedStateHandle jobManagerOwnedSnapshot = snapshotResult.getJobManagerOwnedSnapshot();
	if (jobManagerOwnedSnapshot != null) {
		jobManagerOwnedSnapshot.registerSharedStates(sharedStateRegistry);
	}
	return jobManagerOwnedSnapshot;
}
 

@Test
public void testCancelRunningSnapshot() throws Exception {
	setupRocksKeyedStateBackend();
	try {
		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot =
			keyedStateBackend.snapshot(0L, 0L, testStreamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());
		Thread asyncSnapshotThread = new Thread(snapshot);
		asyncSnapshotThread.start();
		waiter.await(); // wait for snapshot to run
		waiter.reset();
		runStateUpdates();
		snapshot.cancel(true);
		blocker.trigger(); // allow checkpointing to start writing

		for (BlockingCheckpointOutputStream stream : testStreamFactory.getAllCreatedStreams()) {
			assertTrue(stream.isClosed());
		}

		waiter.await(); // wait for snapshot stream writing to run
		try {
			snapshot.get();
			fail();
		} catch (Exception ignored) {
		}

		asyncSnapshotThread.join();
		verifyRocksObjectsReleased();
	} finally {
		this.keyedStateBackend.dispose();
		this.keyedStateBackend = null;
	}
}
 

protected KeyedStateHandle runSnapshot(
	RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotRunnableFuture,
	SharedStateRegistry sharedStateRegistry) throws Exception {

	if (!snapshotRunnableFuture.isDone()) {
		snapshotRunnableFuture.run();
	}

	SnapshotResult<KeyedStateHandle> snapshotResult = snapshotRunnableFuture.get();
	KeyedStateHandle jobManagerOwnedSnapshot = snapshotResult.getJobManagerOwnedSnapshot();
	if (jobManagerOwnedSnapshot != null) {
		jobManagerOwnedSnapshot.registerSharedStates(sharedStateRegistry);
	}
	return jobManagerOwnedSnapshot;
}
 

@Test
public void testCompletingSnapshot() throws Exception {
	setupRocksKeyedStateBackend();
	try {
		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot =
			keyedStateBackend.snapshot(0L, 0L, testStreamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());
		Thread asyncSnapshotThread = new Thread(snapshot);
		asyncSnapshotThread.start();
		waiter.await(); // wait for snapshot to run
		waiter.reset();
		runStateUpdates();
		blocker.trigger(); // allow checkpointing to start writing
		waiter.await(); // wait for snapshot stream writing to run

		SnapshotResult<KeyedStateHandle> snapshotResult = snapshot.get();
		KeyedStateHandle keyedStateHandle = snapshotResult.getJobManagerOwnedSnapshot();
		assertNotNull(keyedStateHandle);
		assertTrue(keyedStateHandle.getStateSize() > 0);
		assertEquals(2, keyedStateHandle.getKeyGroupRange().getNumberOfKeyGroups());

		for (BlockingCheckpointOutputStream stream : testStreamFactory.getAllCreatedStreams()) {
			assertTrue(stream.isClosed());
		}

		asyncSnapshotThread.join();
		verifyRocksObjectsReleased();
	} finally {
		this.keyedStateBackend.dispose();
		this.keyedStateBackend = null;
	}
}
 

void testNonConcurrentSnapshotTransformerAccess() throws Exception {
	List<TestState> testStates = Arrays.asList(
		new TestValueState(),
		new TestListState(),
		new TestMapState()
	);

	for (TestState state : testStates) {
		for (int i = 0; i < 100; i++) {
			backend.setCurrentKey(i);
			state.setToRandomValue();
		}

		CheckpointOptions checkpointOptions = CheckpointOptions.forCheckpointWithDefaultLocation();

		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot1 =
			backend.snapshot(1L, 0L, streamFactory, checkpointOptions);

		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot2 =
			backend.snapshot(2L, 0L, streamFactory, checkpointOptions);

		Thread runner1 = new Thread(snapshot1, "snapshot1");
		runner1.start();
		Thread runner2 = new Thread(snapshot2, "snapshot2");
		runner2.start();

		runner1.join();
		runner2.join();

		snapshot1.get();
		snapshot2.get();
	}
}
 

@Test
public void testCompletingSnapshot() throws Exception {
	setupRocksKeyedStateBackend();
	try {
		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot =
			keyedStateBackend.snapshot(0L, 0L, testStreamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());
		Thread asyncSnapshotThread = new Thread(snapshot);
		asyncSnapshotThread.start();
		waiter.await(); // wait for snapshot to run
		waiter.reset();
		runStateUpdates();
		blocker.trigger(); // allow checkpointing to start writing
		waiter.await(); // wait for snapshot stream writing to run

		SnapshotResult<KeyedStateHandle> snapshotResult = snapshot.get();
		KeyedStateHandle keyedStateHandle = snapshotResult.getJobManagerOwnedSnapshot();
		assertNotNull(keyedStateHandle);
		assertTrue(keyedStateHandle.getStateSize() > 0);
		assertEquals(2, keyedStateHandle.getKeyGroupRange().getNumberOfKeyGroups());

		for (BlockingCheckpointOutputStream stream : testStreamFactory.getAllCreatedStreams()) {
			assertTrue(stream.isClosed());
		}

		asyncSnapshotThread.join();
		verifyRocksObjectsReleased();
	} finally {
		this.keyedStateBackend.dispose();
		this.keyedStateBackend = null;
	}
}
 

private List<Result> evaluateJmcRules(IItemCollection jfrEvents) {
    List<Result> ruleEvaluations = new ArrayList<>();
    for (IRule rule : RuleRegistry.getRules()) {
        RunnableFuture<Result> future = rule.evaluate(jfrEvents, IPreferenceValueProvider.DEFAULT_VALUES);
        future.run();
        Result result;
        try {
            result = future.get();
        } catch (InterruptedException | ExecutionException e) {
            throw new IllegalStateException(e);
        }
        ruleEvaluations.add(result);
    }
    return ruleEvaluations;
}
 

@Test
public void testSnapshotAsyncClose() throws Exception {
	DefaultOperatorStateBackend operatorStateBackend =
		new DefaultOperatorStateBackendBuilder(
			OperatorStateBackendTest.class.getClassLoader(),
			new ExecutionConfig(),
			true,
			emptyStateHandles,
			new CloseableRegistry()).build();

	ListStateDescriptor<MutableType> stateDescriptor1 =
			new ListStateDescriptor<>("test1", new JavaSerializer<MutableType>());

	ListState<MutableType> listState1 = operatorStateBackend.getOperatorState(stateDescriptor1);

	listState1.add(MutableType.of(42));
	listState1.add(MutableType.of(4711));

	MapStateDescriptor<MutableType, MutableType> broadcastStateDescriptor1 =
			new MapStateDescriptor<>("test4", new JavaSerializer<MutableType>(), new JavaSerializer<MutableType>());

	BroadcastState<MutableType, MutableType> broadcastState1 = operatorStateBackend.getBroadcastState(broadcastStateDescriptor1);
	broadcastState1.put(MutableType.of(1), MutableType.of(2));
	broadcastState1.put(MutableType.of(2), MutableType.of(5));

	BlockerCheckpointStreamFactory streamFactory = new BlockerCheckpointStreamFactory(1024 * 1024);

	OneShotLatch waiterLatch = new OneShotLatch();
	OneShotLatch blockerLatch = new OneShotLatch();

	streamFactory.setWaiterLatch(waiterLatch);
	streamFactory.setBlockerLatch(blockerLatch);

	RunnableFuture<SnapshotResult<OperatorStateHandle>> runnableFuture =
			operatorStateBackend.snapshot(1, 1, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());

	ExecutorService executorService = Executors.newFixedThreadPool(1);

	executorService.submit(runnableFuture);

	// wait until the async checkpoint is in the write code, then continue
	waiterLatch.await();

	operatorStateBackend.close();

	blockerLatch.trigger();

	try {
		runnableFuture.get(60, TimeUnit.SECONDS);
		Assert.fail();
	} catch (CancellationException expected) {
	}
}
 

@Test
public void testSnapshotAsyncCancel() throws Exception {
	DefaultOperatorStateBackend operatorStateBackend =
		new DefaultOperatorStateBackendBuilder(
			OperatorStateBackendTest.class.getClassLoader(),
			new ExecutionConfig(),
			true,
			emptyStateHandles,
			new CloseableRegistry()).build();

	ListStateDescriptor<MutableType> stateDescriptor1 =
			new ListStateDescriptor<>("test1", new JavaSerializer<MutableType>());

	ListState<MutableType> listState1 = operatorStateBackend.getOperatorState(stateDescriptor1);

	listState1.add(MutableType.of(42));
	listState1.add(MutableType.of(4711));

	BlockerCheckpointStreamFactory streamFactory = new BlockerCheckpointStreamFactory(1024 * 1024);

	OneShotLatch waiterLatch = new OneShotLatch();
	OneShotLatch blockerLatch = new OneShotLatch();

	streamFactory.setWaiterLatch(waiterLatch);
	streamFactory.setBlockerLatch(blockerLatch);

	RunnableFuture<SnapshotResult<OperatorStateHandle>> runnableFuture =
			operatorStateBackend.snapshot(1, 1, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());

	ExecutorService executorService = Executors.newFixedThreadPool(1);

	executorService.submit(runnableFuture);

	// wait until the async checkpoint is in the stream's write code, then continue
	waiterLatch.await();

	// cancel the future, which should close the underlying stream
	runnableFuture.cancel(true);

	for (BlockingCheckpointOutputStream stream : streamFactory.getAllCreatedStreams()) {
		Assert.assertTrue(stream.isClosed());
	}

	// we allow the stream under test to proceed
	blockerLatch.trigger();

	try {
		runnableFuture.get(60, TimeUnit.SECONDS);
		Assert.fail();
	} catch (CancellationException ignore) {
	}
}
 

/**
 * The purpose of this test is to check that parallel snapshots are possible, and work even if a previous snapshot
 * is still running and blocking.
 */
@Test
public void testParallelAsyncSnapshots() throws Exception {
	OneShotLatch blocker = new OneShotLatch();
	OneShotLatch waiter = new OneShotLatch();
	BlockerCheckpointStreamFactory streamFactory = new BlockerCheckpointStreamFactory(1024 * 1024);
	streamFactory.setWaiterLatch(waiter);
	streamFactory.setBlockerLatch(blocker);
	streamFactory.setAfterNumberInvocations(10);

	final AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);

	try {

		if (!backend.supportsAsynchronousSnapshots()) {
			return;
		}

		// insert some data to the backend.
		InternalValueState<Integer, VoidNamespace, Integer> valueState = backend.createInternalState(
			VoidNamespaceSerializer.INSTANCE,
			new ValueStateDescriptor<>("test", IntSerializer.INSTANCE));

		valueState.setCurrentNamespace(VoidNamespace.INSTANCE);

		for (int i = 0; i < 10; ++i) {
			backend.setCurrentKey(i);
			valueState.update(i);
		}

		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot1 =
			backend.snapshot(0L, 0L, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());

		Thread runner1 = new Thread(snapshot1, "snapshot-1-runner");
		runner1.start();
		// after this call returns, we have a running snapshot-1 that is blocked in IO.
		waiter.await();

		// do some updates in between the snapshots.
		for (int i = 5; i < 15; ++i) {
			backend.setCurrentKey(i);
			valueState.update(i + 1);
		}

		// we don't want to block the second snapshot.
		streamFactory.setWaiterLatch(null);
		streamFactory.setBlockerLatch(null);

		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot2 =
			backend.snapshot(1L, 1L, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());

		Thread runner2 = new Thread(snapshot2,"snapshot-2-runner");
		runner2.start();
		// snapshot-2 should run and succeed, while snapshot-1 is still running and blocked in IO.
		snapshot2.get();

		// we release the blocking IO so that snapshot-1 can also finish and succeed.
		blocker.trigger();
		snapshot1.get();

	} finally {
		backend.dispose();
	}
}
 

@Test
public void testSharedIncrementalStateDeRegistration() throws Exception {
	if (enableIncrementalCheckpointing) {
		AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
		try {
			ValueStateDescriptor<String> kvId =
				new ValueStateDescriptor<>("id", String.class, null);

			kvId.initializeSerializerUnlessSet(new ExecutionConfig());

			ValueState<String> state =
				backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);

			Queue<IncrementalRemoteKeyedStateHandle> previousStateHandles = new LinkedList<>();
			SharedStateRegistry sharedStateRegistry = spy(new SharedStateRegistry());
			for (int checkpointId = 0; checkpointId < 3; ++checkpointId) {

				reset(sharedStateRegistry);

				backend.setCurrentKey(checkpointId);
				state.update("Hello-" + checkpointId);

				RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot = backend.snapshot(
					checkpointId,
					checkpointId,
					createStreamFactory(),
					CheckpointOptions.forCheckpointWithDefaultLocation());

				snapshot.run();

				SnapshotResult<KeyedStateHandle> snapshotResult = snapshot.get();

				IncrementalRemoteKeyedStateHandle stateHandle =
					(IncrementalRemoteKeyedStateHandle) snapshotResult.getJobManagerOwnedSnapshot();

				Map<StateHandleID, StreamStateHandle> sharedState =
					new HashMap<>(stateHandle.getSharedState());

				stateHandle.registerSharedStates(sharedStateRegistry);

				for (Map.Entry<StateHandleID, StreamStateHandle> e : sharedState.entrySet()) {
					verify(sharedStateRegistry).registerReference(
						stateHandle.createSharedStateRegistryKeyFromFileName(e.getKey()),
						e.getValue());
				}

				previousStateHandles.add(stateHandle);
				backend.notifyCheckpointComplete(checkpointId);

				//-----------------------------------------------------------------

				if (previousStateHandles.size() > 1) {
					checkRemove(previousStateHandles.remove(), sharedStateRegistry);
				}
			}

			while (!previousStateHandles.isEmpty()) {

				reset(sharedStateRegistry);

				checkRemove(previousStateHandles.remove(), sharedStateRegistry);
			}
		} finally {
			IOUtils.closeQuietly(backend);
			backend.dispose();
		}
	}
}
 

@Test
public void testSharedIncrementalStateDeRegistration() throws Exception {
	if (enableIncrementalCheckpointing) {
		AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
		try {
			ValueStateDescriptor<String> kvId =
				new ValueStateDescriptor<>("id", String.class, null);

			kvId.initializeSerializerUnlessSet(new ExecutionConfig());

			ValueState<String> state =
				backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);

			Queue<IncrementalRemoteKeyedStateHandle> previousStateHandles = new LinkedList<>();
			SharedStateRegistry sharedStateRegistry = spy(new SharedStateRegistry());
			for (int checkpointId = 0; checkpointId < 3; ++checkpointId) {

				reset(sharedStateRegistry);

				backend.setCurrentKey(checkpointId);
				state.update("Hello-" + checkpointId);

				RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot = backend.snapshot(
					checkpointId,
					checkpointId,
					createStreamFactory(),
					CheckpointOptions.forCheckpointWithDefaultLocation());

				snapshot.run();

				SnapshotResult<KeyedStateHandle> snapshotResult = snapshot.get();

				IncrementalRemoteKeyedStateHandle stateHandle =
					(IncrementalRemoteKeyedStateHandle) snapshotResult.getJobManagerOwnedSnapshot();

				Map<StateHandleID, StreamStateHandle> sharedState =
					new HashMap<>(stateHandle.getSharedState());

				stateHandle.registerSharedStates(sharedStateRegistry);

				for (Map.Entry<StateHandleID, StreamStateHandle> e : sharedState.entrySet()) {
					verify(sharedStateRegistry).registerReference(
						stateHandle.createSharedStateRegistryKeyFromFileName(e.getKey()),
						e.getValue());
				}

				previousStateHandles.add(stateHandle);
				backend.notifyCheckpointComplete(checkpointId);

				//-----------------------------------------------------------------

				if (previousStateHandles.size() > 1) {
					checkRemove(previousStateHandles.remove(), sharedStateRegistry);
				}
			}

			while (!previousStateHandles.isEmpty()) {

				reset(sharedStateRegistry);

				checkRemove(previousStateHandles.remove(), sharedStateRegistry);
			}
		} finally {
			IOUtils.closeQuietly(backend);
			backend.dispose();
		}
	}
}
 

@Test
public void testSnapshotAsyncClose() throws Exception {
	DefaultOperatorStateBackend operatorStateBackend =
		new DefaultOperatorStateBackendBuilder(
			OperatorStateBackendTest.class.getClassLoader(),
			new ExecutionConfig(),
			true,
			emptyStateHandles,
			new CloseableRegistry()).build();

	ListStateDescriptor<MutableType> stateDescriptor1 =
			new ListStateDescriptor<>("test1", new JavaSerializer<MutableType>());

	ListState<MutableType> listState1 = operatorStateBackend.getListState(stateDescriptor1);

	listState1.add(MutableType.of(42));
	listState1.add(MutableType.of(4711));

	MapStateDescriptor<MutableType, MutableType> broadcastStateDescriptor1 =
			new MapStateDescriptor<>("test4", new JavaSerializer<MutableType>(), new JavaSerializer<MutableType>());

	BroadcastState<MutableType, MutableType> broadcastState1 = operatorStateBackend.getBroadcastState(broadcastStateDescriptor1);
	broadcastState1.put(MutableType.of(1), MutableType.of(2));
	broadcastState1.put(MutableType.of(2), MutableType.of(5));

	BlockerCheckpointStreamFactory streamFactory = new BlockerCheckpointStreamFactory(1024 * 1024);

	OneShotLatch waiterLatch = new OneShotLatch();
	OneShotLatch blockerLatch = new OneShotLatch();

	streamFactory.setWaiterLatch(waiterLatch);
	streamFactory.setBlockerLatch(blockerLatch);

	RunnableFuture<SnapshotResult<OperatorStateHandle>> runnableFuture =
			operatorStateBackend.snapshot(1, 1, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());

	ExecutorService executorService = Executors.newFixedThreadPool(1);

	executorService.submit(runnableFuture);

	// wait until the async checkpoint is in the write code, then continue
	waiterLatch.await();

	operatorStateBackend.close();

	blockerLatch.trigger();

	try {
		runnableFuture.get(60, TimeUnit.SECONDS);
		Assert.fail();
	} catch (CancellationException expected) {
	}
}
 

@Test
public void testSnapshotAsyncCancel() throws Exception {
	DefaultOperatorStateBackend operatorStateBackend =
		new DefaultOperatorStateBackendBuilder(
			OperatorStateBackendTest.class.getClassLoader(),
			new ExecutionConfig(),
			true,
			emptyStateHandles,
			new CloseableRegistry()).build();

	ListStateDescriptor<MutableType> stateDescriptor1 =
			new ListStateDescriptor<>("test1", new JavaSerializer<MutableType>());

	ListState<MutableType> listState1 = operatorStateBackend.getListState(stateDescriptor1);

	listState1.add(MutableType.of(42));
	listState1.add(MutableType.of(4711));

	BlockerCheckpointStreamFactory streamFactory = new BlockerCheckpointStreamFactory(1024 * 1024);

	OneShotLatch waiterLatch = new OneShotLatch();
	OneShotLatch blockerLatch = new OneShotLatch();

	streamFactory.setWaiterLatch(waiterLatch);
	streamFactory.setBlockerLatch(blockerLatch);

	RunnableFuture<SnapshotResult<OperatorStateHandle>> runnableFuture =
			operatorStateBackend.snapshot(1, 1, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());

	ExecutorService executorService = Executors.newFixedThreadPool(1);

	executorService.submit(runnableFuture);

	// wait until the async checkpoint is in the stream's write code, then continue
	waiterLatch.await();

	// cancel the future, which should close the underlying stream
	runnableFuture.cancel(true);

	for (BlockingCheckpointOutputStream stream : streamFactory.getAllCreatedStreams()) {
		Assert.assertTrue(stream.isClosed());
	}

	// we allow the stream under test to proceed
	blockerLatch.trigger();

	try {
		runnableFuture.get(60, TimeUnit.SECONDS);
		Assert.fail();
	} catch (CancellationException ignore) {
	}
}
 

@Test
public void testSharedIncrementalStateDeRegistration() throws Exception {
	if (enableIncrementalCheckpointing) {
		AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
		try {
			ValueStateDescriptor<String> kvId =
				new ValueStateDescriptor<>("id", String.class, null);

			kvId.initializeSerializerUnlessSet(new ExecutionConfig());

			ValueState<String> state =
				backend.getPartitionedState(VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, kvId);

			Queue<IncrementalRemoteKeyedStateHandle> previousStateHandles = new LinkedList<>();
			SharedStateRegistry sharedStateRegistry = spy(new SharedStateRegistry());
			for (int checkpointId = 0; checkpointId < 3; ++checkpointId) {

				reset(sharedStateRegistry);

				backend.setCurrentKey(checkpointId);
				state.update("Hello-" + checkpointId);

				RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot = backend.snapshot(
					checkpointId,
					checkpointId,
					createStreamFactory(),
					CheckpointOptions.forCheckpointWithDefaultLocation());

				snapshot.run();

				SnapshotResult<KeyedStateHandle> snapshotResult = snapshot.get();

				IncrementalRemoteKeyedStateHandle stateHandle =
					(IncrementalRemoteKeyedStateHandle) snapshotResult.getJobManagerOwnedSnapshot();

				Map<StateHandleID, StreamStateHandle> sharedState =
					new HashMap<>(stateHandle.getSharedState());

				stateHandle.registerSharedStates(sharedStateRegistry);

				for (Map.Entry<StateHandleID, StreamStateHandle> e : sharedState.entrySet()) {
					verify(sharedStateRegistry).registerReference(
						stateHandle.createSharedStateRegistryKeyFromFileName(e.getKey()),
						e.getValue());
				}

				previousStateHandles.add(stateHandle);
				backend.notifyCheckpointComplete(checkpointId);

				//-----------------------------------------------------------------

				if (previousStateHandles.size() > 1) {
					checkRemove(previousStateHandles.remove(), sharedStateRegistry);
				}
			}

			while (!previousStateHandles.isEmpty()) {

				reset(sharedStateRegistry);

				checkRemove(previousStateHandles.remove(), sharedStateRegistry);
			}
		} finally {
			IOUtils.closeQuietly(backend);
			backend.dispose();
		}
	}
}
 

@Test
public void testAsyncSnapshotCancellation() throws Exception {
	OneShotLatch blocker = new OneShotLatch();
	OneShotLatch waiter = new OneShotLatch();
	BlockerCheckpointStreamFactory streamFactory = new BlockerCheckpointStreamFactory(1024 * 1024);
	streamFactory.setWaiterLatch(waiter);
	streamFactory.setBlockerLatch(blocker);
	streamFactory.setAfterNumberInvocations(10);

	final AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);

	try {

		if (!backend.supportsAsynchronousSnapshots()) {
			return;
		}

		InternalValueState<Integer, VoidNamespace, Integer> valueState = backend.createInternalState(
				VoidNamespaceSerializer.INSTANCE,
				new ValueStateDescriptor<>("test", IntSerializer.INSTANCE));

		valueState.setCurrentNamespace(VoidNamespace.INSTANCE);

		for (int i = 0; i < 10; ++i) {
			backend.setCurrentKey(i);
			valueState.update(i);
		}

		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot =
				backend.snapshot(0L, 0L, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());

		Thread runner = new Thread(snapshot);
		runner.start();

		// wait until the code reached some stream read
		waiter.await();

		// close the backend to see if the close is propagated to the stream
		IOUtils.closeQuietly(backend);

		//unblock the stream so that it can run into the IOException
		blocker.trigger();

		runner.join();

		try {
			snapshot.get();
			fail("Close was not propagated.");
		} catch (CancellationException ex) {
			//ignore
		}

	} finally {
		backend.dispose();
	}
}
 

@Test
public void testSnapshotAsyncClose() throws Exception {
	DefaultOperatorStateBackend operatorStateBackend =
		new DefaultOperatorStateBackendBuilder(
			OperatorStateBackendTest.class.getClassLoader(),
			new ExecutionConfig(),
			true,
			emptyStateHandles,
			new CloseableRegistry()).build();

	ListStateDescriptor<MutableType> stateDescriptor1 =
			new ListStateDescriptor<>("test1", new JavaSerializer<MutableType>());

	ListState<MutableType> listState1 = operatorStateBackend.getOperatorState(stateDescriptor1);

	listState1.add(MutableType.of(42));
	listState1.add(MutableType.of(4711));

	MapStateDescriptor<MutableType, MutableType> broadcastStateDescriptor1 =
			new MapStateDescriptor<>("test4", new JavaSerializer<MutableType>(), new JavaSerializer<MutableType>());

	BroadcastState<MutableType, MutableType> broadcastState1 = operatorStateBackend.getBroadcastState(broadcastStateDescriptor1);
	broadcastState1.put(MutableType.of(1), MutableType.of(2));
	broadcastState1.put(MutableType.of(2), MutableType.of(5));

	BlockerCheckpointStreamFactory streamFactory = new BlockerCheckpointStreamFactory(1024 * 1024);

	OneShotLatch waiterLatch = new OneShotLatch();
	OneShotLatch blockerLatch = new OneShotLatch();

	streamFactory.setWaiterLatch(waiterLatch);
	streamFactory.setBlockerLatch(blockerLatch);

	RunnableFuture<SnapshotResult<OperatorStateHandle>> runnableFuture =
			operatorStateBackend.snapshot(1, 1, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());

	ExecutorService executorService = Executors.newFixedThreadPool(1);

	executorService.submit(runnableFuture);

	// wait until the async checkpoint is in the write code, then continue
	waiterLatch.await();

	operatorStateBackend.close();

	blockerLatch.trigger();

	try {
		runnableFuture.get(60, TimeUnit.SECONDS);
		Assert.fail();
	} catch (CancellationException expected) {
	}
}
 

@Test
public void testSnapshotAsyncCancel() throws Exception {
	DefaultOperatorStateBackend operatorStateBackend =
		new DefaultOperatorStateBackendBuilder(
			OperatorStateBackendTest.class.getClassLoader(),
			new ExecutionConfig(),
			true,
			emptyStateHandles,
			new CloseableRegistry()).build();

	ListStateDescriptor<MutableType> stateDescriptor1 =
			new ListStateDescriptor<>("test1", new JavaSerializer<MutableType>());

	ListState<MutableType> listState1 = operatorStateBackend.getOperatorState(stateDescriptor1);

	listState1.add(MutableType.of(42));
	listState1.add(MutableType.of(4711));

	BlockerCheckpointStreamFactory streamFactory = new BlockerCheckpointStreamFactory(1024 * 1024);

	OneShotLatch waiterLatch = new OneShotLatch();
	OneShotLatch blockerLatch = new OneShotLatch();

	streamFactory.setWaiterLatch(waiterLatch);
	streamFactory.setBlockerLatch(blockerLatch);

	RunnableFuture<SnapshotResult<OperatorStateHandle>> runnableFuture =
			operatorStateBackend.snapshot(1, 1, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());

	ExecutorService executorService = Executors.newFixedThreadPool(1);

	executorService.submit(runnableFuture);

	// wait until the async checkpoint is in the stream's write code, then continue
	waiterLatch.await();

	// cancel the future, which should close the underlying stream
	runnableFuture.cancel(true);

	for (BlockingCheckpointOutputStream stream : streamFactory.getAllCreatedStreams()) {
		Assert.assertTrue(stream.isClosed());
	}

	// we allow the stream under test to proceed
	blockerLatch.trigger();

	try {
		runnableFuture.get(60, TimeUnit.SECONDS);
		Assert.fail();
	} catch (CancellationException ignore) {
	}
}