Java Code Examples for org.apache.flink.runtime.concurrent.FutureUtils#runIfNotDoneAndGet()

/**
 * Discards the given state future by first trying to cancel it. If this is not possible, then
 * the state object contained in the future is calculated and afterwards discarded.
 *
 * @param stateFuture to be discarded
 * @throws Exception if the discard operation failed
 */
public static void discardStateFuture(RunnableFuture<? extends StateObject> stateFuture) throws Exception {
	if (null != stateFuture) {
		if (!stateFuture.cancel(true)) {

			try {
				// We attempt to get a result, in case the future completed before cancellation.
				StateObject stateObject = FutureUtils.runIfNotDoneAndGet(stateFuture);

				if (null != stateObject) {
					stateObject.discardState();
				}
			} catch (CancellationException | ExecutionException ex) {
				LOG.debug("Cancelled execution of snapshot future runnable. Cancellation produced the following " +
					"exception, which is expected an can be ignored.", ex);
			}
		}
	}
}
 

@Test
public void testSnapshotEmpty() throws Exception {
	final AbstractStateBackend abstractStateBackend = new MemoryStateBackend(4096);
	CloseableRegistry cancelStreamRegistry = new CloseableRegistry();

	final OperatorStateBackend operatorStateBackend =
			abstractStateBackend.createOperatorStateBackend(createMockEnvironment(), "testOperator", emptyStateHandles, cancelStreamRegistry);

	CheckpointStreamFactory streamFactory = new MemCheckpointStreamFactory(4096);

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

	SnapshotResult<OperatorStateHandle> snapshotResult = FutureUtils.runIfNotDoneAndGet(snapshot);
	OperatorStateHandle stateHandle = snapshotResult.getJobManagerOwnedSnapshot();
	assertNull(stateHandle);
}
 

/**
 * Discards the given state future by first trying to cancel it. If this is not possible, then
 * the state object contained in the future is calculated and afterwards discarded.
 *
 * @param stateFuture to be discarded
 * @throws Exception if the discard operation failed
 */
public static void discardStateFuture(RunnableFuture<? extends StateObject> stateFuture) throws Exception {
	if (null != stateFuture) {
		if (!stateFuture.cancel(true)) {

			try {
				// We attempt to get a result, in case the future completed before cancellation.
				StateObject stateObject = FutureUtils.runIfNotDoneAndGet(stateFuture);

				if (null != stateObject) {
					stateObject.discardState();
				}
			} catch (CancellationException | ExecutionException ex) {
				LOG.debug("Cancelled execution of snapshot future runnable. Cancellation produced the following " +
					"exception, which is expected an can be ignored.", ex);
			}
		}
	}
}
 

@Test
public void testSnapshotEmpty() throws Exception {
	final AbstractStateBackend abstractStateBackend = new MemoryStateBackend(4096);
	CloseableRegistry cancelStreamRegistry = new CloseableRegistry();

	final OperatorStateBackend operatorStateBackend =
			abstractStateBackend.createOperatorStateBackend(createMockEnvironment(), "testOperator", emptyStateHandles, cancelStreamRegistry);

	CheckpointStreamFactory streamFactory = new MemCheckpointStreamFactory(4096);

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

	SnapshotResult<OperatorStateHandle> snapshotResult = FutureUtils.runIfNotDoneAndGet(snapshot);
	OperatorStateHandle stateHandle = snapshotResult.getJobManagerOwnedSnapshot();
	assertNull(stateHandle);
}
 

@Test
public void testSnapshotEmpty() throws Exception {
	final AbstractStateBackend abstractStateBackend = new MemoryStateBackend(4096);
	CloseableRegistry cancelStreamRegistry = new CloseableRegistry();

	final OperatorStateBackend operatorStateBackend =
			abstractStateBackend.createOperatorStateBackend(createMockEnvironment(), "testOperator", emptyStateHandles, cancelStreamRegistry);

	CheckpointStreamFactory streamFactory = new MemCheckpointStreamFactory(4096);

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

	SnapshotResult<OperatorStateHandle> snapshotResult = FutureUtils.runIfNotDoneAndGet(snapshot);
	OperatorStateHandle stateHandle = snapshotResult.getJobManagerOwnedSnapshot();
	assertNull(stateHandle);
}
 

public OperatorSnapshotFinalizer(
	@Nonnull OperatorSnapshotFutures snapshotFutures) throws ExecutionException, InterruptedException {

	SnapshotResult<KeyedStateHandle> keyedManaged =
		FutureUtils.runIfNotDoneAndGet(snapshotFutures.getKeyedStateManagedFuture());

	SnapshotResult<KeyedStateHandle> keyedRaw =
		FutureUtils.runIfNotDoneAndGet(snapshotFutures.getKeyedStateRawFuture());

	SnapshotResult<OperatorStateHandle> operatorManaged =
		FutureUtils.runIfNotDoneAndGet(snapshotFutures.getOperatorStateManagedFuture());

	SnapshotResult<OperatorStateHandle> operatorRaw =
		FutureUtils.runIfNotDoneAndGet(snapshotFutures.getOperatorStateRawFuture());

	jobManagerOwnedState = new OperatorSubtaskState(
		operatorManaged.getJobManagerOwnedSnapshot(),
		operatorRaw.getJobManagerOwnedSnapshot(),
		keyedManaged.getJobManagerOwnedSnapshot(),
		keyedRaw.getJobManagerOwnedSnapshot()
	);

	taskLocalState = new OperatorSubtaskState(
		operatorManaged.getTaskLocalSnapshot(),
		operatorRaw.getTaskLocalSnapshot(),
		keyedManaged.getTaskLocalSnapshot(),
		keyedRaw.getTaskLocalSnapshot()
	);
}
 

public OperatorSnapshotFinalizer(
	@Nonnull OperatorSnapshotFutures snapshotFutures) throws ExecutionException, InterruptedException {

	SnapshotResult<KeyedStateHandle> keyedManaged =
		FutureUtils.runIfNotDoneAndGet(snapshotFutures.getKeyedStateManagedFuture());

	SnapshotResult<KeyedStateHandle> keyedRaw =
		FutureUtils.runIfNotDoneAndGet(snapshotFutures.getKeyedStateRawFuture());

	SnapshotResult<OperatorStateHandle> operatorManaged =
		FutureUtils.runIfNotDoneAndGet(snapshotFutures.getOperatorStateManagedFuture());

	SnapshotResult<OperatorStateHandle> operatorRaw =
		FutureUtils.runIfNotDoneAndGet(snapshotFutures.getOperatorStateRawFuture());

	jobManagerOwnedState = new OperatorSubtaskState(
		operatorManaged.getJobManagerOwnedSnapshot(),
		operatorRaw.getJobManagerOwnedSnapshot(),
		keyedManaged.getJobManagerOwnedSnapshot(),
		keyedRaw.getJobManagerOwnedSnapshot()
	);

	taskLocalState = new OperatorSubtaskState(
		operatorManaged.getTaskLocalSnapshot(),
		operatorRaw.getTaskLocalSnapshot(),
		keyedManaged.getTaskLocalSnapshot(),
		keyedRaw.getTaskLocalSnapshot()
	);
}
 

public OperatorSnapshotFinalizer(
	@Nonnull OperatorSnapshotFutures snapshotFutures) throws ExecutionException, InterruptedException {

	SnapshotResult<KeyedStateHandle> keyedManaged =
		FutureUtils.runIfNotDoneAndGet(snapshotFutures.getKeyedStateManagedFuture());

	SnapshotResult<KeyedStateHandle> keyedRaw =
		FutureUtils.runIfNotDoneAndGet(snapshotFutures.getKeyedStateRawFuture());

	SnapshotResult<OperatorStateHandle> operatorManaged =
		FutureUtils.runIfNotDoneAndGet(snapshotFutures.getOperatorStateManagedFuture());

	SnapshotResult<OperatorStateHandle> operatorRaw =
		FutureUtils.runIfNotDoneAndGet(snapshotFutures.getOperatorStateRawFuture());

	SnapshotResult<StateObjectCollection<InputChannelStateHandle>> inputChannel = snapshotFutures.getInputChannelStateFuture().get();

	SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>> resultSubpartition = snapshotFutures.getResultSubpartitionStateFuture().get();

	jobManagerOwnedState = new OperatorSubtaskState(
		operatorManaged.getJobManagerOwnedSnapshot(),
		operatorRaw.getJobManagerOwnedSnapshot(),
		keyedManaged.getJobManagerOwnedSnapshot(),
		keyedRaw.getJobManagerOwnedSnapshot(),
		inputChannel.getJobManagerOwnedSnapshot(),
		resultSubpartition.getJobManagerOwnedSnapshot()
	);

	taskLocalState = new OperatorSubtaskState(
		operatorManaged.getTaskLocalSnapshot(),
		operatorRaw.getTaskLocalSnapshot(),
		keyedManaged.getTaskLocalSnapshot(),
		keyedRaw.getTaskLocalSnapshot(),
		inputChannel.getTaskLocalSnapshot(),
		resultSubpartition.getTaskLocalSnapshot()
	);
}
 

/**
 * Test that the snapshot files are cleaned up in case of a failure during the snapshot
 * procedure.
 */
@Test
public void testCleanupOfSnapshotsInFailureCase() throws Exception {
	long checkpointId = 1L;
	long timestamp = 42L;

	Environment env = new DummyEnvironment("test task", 1, 0);

	final IOException testException = new IOException("Test exception");
	CheckpointStateOutputStream outputStream = spy(new FailingStream(testException));

	RocksDBStateBackend backend = new RocksDBStateBackend((StateBackend) new MemoryStateBackend());

	backend.setDbStoragePath(temporaryFolder.newFolder().toURI().toString());

	AbstractKeyedStateBackend<Void> keyedStateBackend = backend.createKeyedStateBackend(
		env,
		new JobID(),
		"test operator",
		VoidSerializer.INSTANCE,
		1,
		new KeyGroupRange(0, 0),
		null,
		TtlTimeProvider.DEFAULT,
		new UnregisteredMetricsGroup(),
		Collections.emptyList(),
		new CloseableRegistry());

	try {
		// register a state so that the state backend has to checkpoint something
		keyedStateBackend.getPartitionedState(
			"namespace",
			StringSerializer.INSTANCE,
			new ValueStateDescriptor<>("foobar", String.class));

		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotFuture = keyedStateBackend.snapshot(
			checkpointId, timestamp,
			new TestCheckpointStreamFactory(() -> outputStream),
			CheckpointOptions.forCheckpointWithDefaultLocation());

		try {
			FutureUtils.runIfNotDoneAndGet(snapshotFuture);
			fail("Expected an exception to be thrown here.");
		} catch (ExecutionException e) {
			Assert.assertEquals(testException, e.getCause());
		}

		verify(outputStream).close();
	} finally {
		IOUtils.closeQuietly(keyedStateBackend);
		keyedStateBackend.dispose();
	}
}
 

@SuppressWarnings("unchecked")
@Test
public void testCorrectClassLoaderUsedOnSnapshot() throws Exception {

	AbstractStateBackend abstractStateBackend = new MemoryStateBackend(4096);

	final Environment env = createMockEnvironment();
	CloseableRegistry cancelStreamRegistry = new CloseableRegistry();
	OperatorStateBackend operatorStateBackend = abstractStateBackend.createOperatorStateBackend(env, "test-op-name", emptyStateHandles, cancelStreamRegistry);

	AtomicInteger copyCounter = new AtomicInteger(0);
	TypeSerializer<Integer> serializer = new VerifyingIntSerializer(env.getUserClassLoader(), copyCounter);

	// write some state
	ListStateDescriptor<Integer> stateDescriptor = new ListStateDescriptor<>("test", serializer);
	ListState<Integer> listState = operatorStateBackend.getListState(stateDescriptor);

	listState.add(42);

	AtomicInteger keyCopyCounter = new AtomicInteger(0);
	AtomicInteger valueCopyCounter = new AtomicInteger(0);

	TypeSerializer<Integer> keySerializer = new VerifyingIntSerializer(env.getUserClassLoader(), keyCopyCounter);
	TypeSerializer<Integer> valueSerializer = new VerifyingIntSerializer(env.getUserClassLoader(), valueCopyCounter);

	MapStateDescriptor<Integer, Integer> broadcastStateDesc = new MapStateDescriptor<>(
			"test-broadcast", keySerializer, valueSerializer);

	BroadcastState<Integer, Integer> broadcastState = operatorStateBackend.getBroadcastState(broadcastStateDesc);
	broadcastState.put(1, 2);
	broadcastState.put(3, 4);
	broadcastState.put(5, 6);

	CheckpointStreamFactory streamFactory = new MemCheckpointStreamFactory(4096);
	RunnableFuture<SnapshotResult<OperatorStateHandle>> runnableFuture =
		operatorStateBackend.snapshot(1, 1, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());
	FutureUtils.runIfNotDoneAndGet(runnableFuture);

	// make sure that the copy method has been called
	assertTrue(copyCounter.get() > 0);
	assertTrue(keyCopyCounter.get() > 0);
	assertTrue(valueCopyCounter.get() > 0);
}
 

/**
 * Test that the snapshot files are cleaned up in case of a failure during the snapshot
 * procedure.
 */
@Test
public void testCleanupOfSnapshotsInFailureCase() throws Exception {
	long checkpointId = 1L;
	long timestamp = 42L;

	Environment env = new DummyEnvironment("test task", 1, 0);

	final IOException testException = new IOException("Test exception");
	CheckpointStateOutputStream outputStream = spy(new FailingStream(testException));

	RocksDBStateBackend backend = new RocksDBStateBackend((StateBackend) new MemoryStateBackend());

	backend.setDbStoragePath(temporaryFolder.newFolder().toURI().toString());

	AbstractKeyedStateBackend<Void> keyedStateBackend = backend.createKeyedStateBackend(
		env,
		new JobID(),
		"test operator",
		VoidSerializer.INSTANCE,
		1,
		new KeyGroupRange(0, 0),
		null,
		TtlTimeProvider.DEFAULT,
		new UnregisteredMetricsGroup(),
		Collections.emptyList(),
		new CloseableRegistry());

	try {
		// register a state so that the state backend has to checkpoint something
		keyedStateBackend.getPartitionedState(
			"namespace",
			StringSerializer.INSTANCE,
			new ValueStateDescriptor<>("foobar", String.class));

		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotFuture = keyedStateBackend.snapshot(
			checkpointId, timestamp,
			new TestCheckpointStreamFactory(() -> outputStream),
			CheckpointOptions.forCheckpointWithDefaultLocation());

		try {
			FutureUtils.runIfNotDoneAndGet(snapshotFuture);
			fail("Expected an exception to be thrown here.");
		} catch (ExecutionException e) {
			Assert.assertEquals(testException, e.getCause());
		}

		verify(outputStream).close();
	} finally {
		IOUtils.closeQuietly(keyedStateBackend);
		keyedStateBackend.dispose();
	}
}
 

@SuppressWarnings("unchecked")
@Test
public void testCorrectClassLoaderUsedOnSnapshot() throws Exception {

	AbstractStateBackend abstractStateBackend = new MemoryStateBackend(4096);

	final Environment env = createMockEnvironment();
	CloseableRegistry cancelStreamRegistry = new CloseableRegistry();
	OperatorStateBackend operatorStateBackend = abstractStateBackend.createOperatorStateBackend(env, "test-op-name", emptyStateHandles, cancelStreamRegistry);

	AtomicInteger copyCounter = new AtomicInteger(0);
	TypeSerializer<Integer> serializer = new VerifyingIntSerializer(env.getUserClassLoader(), copyCounter);

	// write some state
	ListStateDescriptor<Integer> stateDescriptor = new ListStateDescriptor<>("test", serializer);
	ListState<Integer> listState = operatorStateBackend.getListState(stateDescriptor);

	listState.add(42);

	AtomicInteger keyCopyCounter = new AtomicInteger(0);
	AtomicInteger valueCopyCounter = new AtomicInteger(0);

	TypeSerializer<Integer> keySerializer = new VerifyingIntSerializer(env.getUserClassLoader(), keyCopyCounter);
	TypeSerializer<Integer> valueSerializer = new VerifyingIntSerializer(env.getUserClassLoader(), valueCopyCounter);

	MapStateDescriptor<Integer, Integer> broadcastStateDesc = new MapStateDescriptor<>(
			"test-broadcast", keySerializer, valueSerializer);

	BroadcastState<Integer, Integer> broadcastState = operatorStateBackend.getBroadcastState(broadcastStateDesc);
	broadcastState.put(1, 2);
	broadcastState.put(3, 4);
	broadcastState.put(5, 6);

	CheckpointStreamFactory streamFactory = new MemCheckpointStreamFactory(4096);
	RunnableFuture<SnapshotResult<OperatorStateHandle>> runnableFuture =
		operatorStateBackend.snapshot(1, 1, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());
	FutureUtils.runIfNotDoneAndGet(runnableFuture);

	// make sure that the copy method has been called
	assertTrue(copyCounter.get() > 0);
	assertTrue(keyCopyCounter.get() > 0);
	assertTrue(valueCopyCounter.get() > 0);
}
 

/**
 * Test that the snapshot files are cleaned up in case of a failure during the snapshot
 * procedure.
 */
@Test
public void testCleanupOfSnapshotsInFailureCase() throws Exception {
	long checkpointId = 1L;
	long timestamp = 42L;

	MockEnvironment env = MockEnvironment.builder().build();

	final IOException testException = new IOException("Test exception");
	CheckpointStateOutputStream outputStream = spy(new FailingStream(testException));

	RocksDBStateBackend backend = new RocksDBStateBackend((StateBackend) new MemoryStateBackend());

	backend.setDbStoragePath(temporaryFolder.newFolder().toURI().toString());

	AbstractKeyedStateBackend<Void> keyedStateBackend = backend.createKeyedStateBackend(
		env,
		new JobID(),
		"test operator",
		VoidSerializer.INSTANCE,
		1,
		new KeyGroupRange(0, 0),
		null,
		TtlTimeProvider.DEFAULT,
		new UnregisteredMetricsGroup(),
		Collections.emptyList(),
		new CloseableRegistry());

	try {
		// register a state so that the state backend has to checkpoint something
		keyedStateBackend.getPartitionedState(
			"namespace",
			StringSerializer.INSTANCE,
			new ValueStateDescriptor<>("foobar", String.class));

		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshotFuture = keyedStateBackend.snapshot(
			checkpointId, timestamp,
			new TestCheckpointStreamFactory(() -> outputStream),
			CheckpointOptions.forCheckpointWithDefaultLocation());

		try {
			FutureUtils.runIfNotDoneAndGet(snapshotFuture);
			fail("Expected an exception to be thrown here.");
		} catch (ExecutionException e) {
			Assert.assertEquals(testException, e.getCause());
		}

		verify(outputStream).close();
	} finally {
		IOUtils.closeQuietly(keyedStateBackend);
		keyedStateBackend.dispose();
		IOUtils.closeQuietly(env);
	}
}
 

@SuppressWarnings("unchecked")
@Test
public void testCorrectClassLoaderUsedOnSnapshot() throws Exception {

	AbstractStateBackend abstractStateBackend = new MemoryStateBackend(4096);

	final Environment env = createMockEnvironment();
	CloseableRegistry cancelStreamRegistry = new CloseableRegistry();
	OperatorStateBackend operatorStateBackend = abstractStateBackend.createOperatorStateBackend(env, "test-op-name", emptyStateHandles, cancelStreamRegistry);

	AtomicInteger copyCounter = new AtomicInteger(0);
	TypeSerializer<Integer> serializer = new VerifyingIntSerializer(env.getUserClassLoader(), copyCounter);

	// write some state
	ListStateDescriptor<Integer> stateDescriptor = new ListStateDescriptor<>("test", serializer);
	ListState<Integer> listState = operatorStateBackend.getListState(stateDescriptor);

	listState.add(42);

	AtomicInteger keyCopyCounter = new AtomicInteger(0);
	AtomicInteger valueCopyCounter = new AtomicInteger(0);

	TypeSerializer<Integer> keySerializer = new VerifyingIntSerializer(env.getUserClassLoader(), keyCopyCounter);
	TypeSerializer<Integer> valueSerializer = new VerifyingIntSerializer(env.getUserClassLoader(), valueCopyCounter);

	MapStateDescriptor<Integer, Integer> broadcastStateDesc = new MapStateDescriptor<>(
			"test-broadcast", keySerializer, valueSerializer);

	BroadcastState<Integer, Integer> broadcastState = operatorStateBackend.getBroadcastState(broadcastStateDesc);
	broadcastState.put(1, 2);
	broadcastState.put(3, 4);
	broadcastState.put(5, 6);

	CheckpointStreamFactory streamFactory = new MemCheckpointStreamFactory(4096);
	RunnableFuture<SnapshotResult<OperatorStateHandle>> runnableFuture =
		operatorStateBackend.snapshot(1, 1, streamFactory, CheckpointOptions.forCheckpointWithDefaultLocation());
	FutureUtils.runIfNotDoneAndGet(runnableFuture);

	// make sure that the copy method has been called
	assertTrue(copyCounter.get() > 0);
	assertTrue(keyCopyCounter.get() > 0);
	assertTrue(valueCopyCounter.get() > 0);
}