org.apache.flink.runtime.state.memory.MemCheckpointStreamFactory Java Examples

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

/**
 * Checks that the state snapshot context is closed after a successful snapshot operation.
 */
@Test
public void testSnapshotMethod() throws Exception {
	final long checkpointId = 42L;
	final long timestamp = 1L;

	final CloseableRegistry closeableRegistry = new CloseableRegistry();

	StateSnapshotContextSynchronousImpl context = spy(new StateSnapshotContextSynchronousImpl(0L, 0L));

	whenNew(StateSnapshotContextSynchronousImpl.class).withAnyArguments().thenReturn(context);

	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();
	doReturn(containingTask).when(operator).getContainingTask();

	operator.snapshotState(
			checkpointId,
			timestamp,
			CheckpointOptions.forCheckpointWithDefaultLocation(),
			new MemCheckpointStreamFactory(Integer.MAX_VALUE));

}
 

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

/**
 * Checks that the state snapshot context is closed after a successful snapshot operation.
 */
@Test
public void testSnapshotMethod() throws Exception {
	final long checkpointId = 42L;
	final long timestamp = 1L;

	final CloseableRegistry closeableRegistry = new CloseableRegistry();

	StateSnapshotContextSynchronousImpl context = spy(new StateSnapshotContextSynchronousImpl(0L, 0L));

	whenNew(StateSnapshotContextSynchronousImpl.class).withAnyArguments().thenReturn(context);

	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();
	doReturn(containingTask).when(operator).getContainingTask();

	operator.snapshotState(
			checkpointId,
			timestamp,
			CheckpointOptions.forCheckpointWithDefaultLocation(),
			new MemCheckpointStreamFactory(Integer.MAX_VALUE));

	verify(context).close();
}
 

@Before
public void setUp() throws Exception {
	CloseableRegistry closableRegistry = new CloseableRegistry();
	CheckpointStreamFactory streamFactory = new MemCheckpointStreamFactory(1024);
	KeyGroupRange keyGroupRange = new KeyGroupRange(0, 2);
	this.snapshotContext = new StateSnapshotContextSynchronousImpl(42, 4711, streamFactory, keyGroupRange, closableRegistry);
}
 

public TestStateSnapshotContextSynchronousImpl(
		long checkpointId,
		long timestamp,
		CloseableRegistry closeableRegistry) {
	super(checkpointId, timestamp, new MemCheckpointStreamFactory(1024), new KeyGroupRange(0, 2), closeableRegistry);
	this.keyedStateCheckpointClosingFuture = new CancelableFuture<>();
	this.operatorStateCheckpointClosingFuture = new CancelableFuture<>();
}
 

@Override
public CheckpointStateOutputStream createCheckpointStateOutputStream(
	CheckpointedStateScope scope) throws IOException {

	BlockingCheckpointOutputStream blockingStream = new BlockingCheckpointOutputStream(
		new MemCheckpointStreamFactory.MemoryCheckpointOutputStream(maxSize),
		waiter,
		blocker,
		afterNumberInvocations);

	allCreatedStreams.add(blockingStream);

	return blockingStream;
}
 

@Before
public void setUp() throws Exception {
	CloseableRegistry closableRegistry = new CloseableRegistry();
	CheckpointStreamFactory streamFactory = new MemCheckpointStreamFactory(1024);
	KeyGroupRange keyGroupRange = new KeyGroupRange(0, 2);
	this.snapshotContext = new StateSnapshotContextSynchronousImpl(42, 4711, streamFactory, keyGroupRange, closableRegistry);
}
 

/**
 * Tests that the created StateSnapshotContextSynchronousImpl is closed in case of a failing
 * Operator#snapshotState(StateSnapshotContextSynchronousImpl) call.
 */
@Test
public void testFailingSnapshotMethod() throws Exception {
	final long checkpointId = 42L;
	final long timestamp = 1L;

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

	final CloseableRegistry closeableRegistry = new CloseableRegistry();

	StateSnapshotContextSynchronousImpl context = mock(StateSnapshotContextSynchronousImpl.class);

	whenNew(StateSnapshotContextSynchronousImpl.class).withAnyArguments().thenReturn(context);

	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();
	doReturn(containingTask).when(operator).getContainingTask();

	// lets fail when calling the actual snapshotState method
	doThrow(failingException).when(operator).snapshotState(eq(context));

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

@Override
public CheckpointStateOutputStream createCheckpointStateOutputStream(
	CheckpointedStateScope scope) throws IOException {

	BlockingCheckpointOutputStream blockingStream = new BlockingCheckpointOutputStream(
		new MemCheckpointStreamFactory.MemoryCheckpointOutputStream(maxSize),
		waiter,
		blocker,
		afterNumberInvocations);

	allCreatedStreams.add(blockingStream);

	return blockingStream;
}
 

@Override
public CheckpointStateOutputStream createCheckpointStateOutputStream(
	CheckpointedStateScope scope) throws IOException {

	BlockingCheckpointOutputStream blockingStream = new BlockingCheckpointOutputStream(
		new MemCheckpointStreamFactory.MemoryCheckpointOutputStream(maxSize),
		waiter,
		blocker,
		afterNumberInvocations);

	allCreatedStreams.add(blockingStream);

	return blockingStream;
}
 

/**
 * Tests that the created StateSnapshotContextSynchronousImpl is closed in case of a failing
 * Operator#snapshotState(StateSnapshotContextSynchronousImpl) call.
 */
@Test
public void testFailingSnapshotMethod() throws Exception {
	final long checkpointId = 42L;
	final long timestamp = 1L;

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

	final CloseableRegistry closeableRegistry = new CloseableRegistry();

	StateSnapshotContextSynchronousImpl context = mock(StateSnapshotContextSynchronousImpl.class);

	whenNew(StateSnapshotContextSynchronousImpl.class).withAnyArguments().thenReturn(context);

	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();
	doReturn(containingTask).when(operator).getContainingTask();

	// lets fail when calling the actual snapshotState method
	doThrow(failingException).when(operator).snapshotState(eq(context));

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

	verify(context).close();
}
 

@Before
public void setUp() throws Exception {
	CloseableRegistry closableRegistry = new CloseableRegistry();
	CheckpointStreamFactory streamFactory = new MemCheckpointStreamFactory(1024);
	KeyGroupRange keyGroupRange = new KeyGroupRange(0, 2);
	this.snapshotContext = new StateSnapshotContextSynchronousImpl(42, 4711, streamFactory, keyGroupRange, closableRegistry);
}
 

@Test
public void testMapStateMigrationAfterHashMapSerRemoval() throws Exception {
	ClassLoader cl = getClass().getClassLoader();
	URL resource = cl.getResource("heap_keyed_statebackend_1_5_map.snapshot");

	Preconditions.checkNotNull(resource, "Binary snapshot resource not found!");

	final SnapshotResult<KeyedStateHandle> stateHandles;
	try (BufferedInputStream bis = new BufferedInputStream((new FileInputStream(resource.getFile())))) {
		stateHandles = InstantiationUtil.deserializeObject(bis, Thread.currentThread().getContextClassLoader());
	}
	final KeyedStateHandle stateHandle = stateHandles.getJobManagerOwnedSnapshot();
	try (final HeapKeyedStateBackend<String> keyedBackend = createKeyedBackend(StateObjectCollection.singleton(stateHandle))) {
		final Integer namespace1 = 1;
		final Integer namespace2 = 2;
		final Integer namespace3 = 3;

		final MapStateDescriptor<Long, Long> stateDescr = new MapStateDescriptor<>("my-map-state", Long.class, Long.class);
		stateDescr.initializeSerializerUnlessSet(new ExecutionConfig());

		InternalMapState<String, Integer, Long, Long> state = keyedBackend.createInternalState(IntSerializer.INSTANCE, stateDescr);

		keyedBackend.setCurrentKey("abc");
		state.setCurrentNamespace(namespace1);
		assertEquals(33L, (long) state.get(33L));
		assertEquals(55L, (long) state.get(55L));
		assertEquals(2, getStateSize(state));

		state.setCurrentNamespace(namespace2);
		assertEquals(22L, (long) state.get(22L));
		assertEquals(11L, (long) state.get(11L));
		assertEquals(2, getStateSize(state));

		state.setCurrentNamespace(namespace3);
		assertEquals(44L, (long) state.get(44L));
		assertEquals(1, getStateSize(state));

		keyedBackend.setCurrentKey("def");
		state.setCurrentNamespace(namespace1);
		assertEquals(11L, (long) state.get(11L));
		assertEquals(44L, (long) state.get(44L));
		assertEquals(2, getStateSize(state));

		state.setCurrentNamespace(namespace3);
		assertEquals(22L, (long) state.get(22L));
		assertEquals(55L, (long) state.get(55L));
		assertEquals(33L, (long) state.get(33L));
		assertEquals(3, getStateSize(state));

		keyedBackend.setCurrentKey("jkl");
		state.setCurrentNamespace(namespace1);
		assertEquals(11L, (long) state.get(11L));
		assertEquals(22L, (long) state.get(22L));
		assertEquals(33L, (long) state.get(33L));
		assertEquals(44L, (long) state.get(44L));
		assertEquals(55L, (long) state.get(55L));
		assertEquals(5, getStateSize(state));

		keyedBackend.setCurrentKey("mno");
		state.setCurrentNamespace(namespace3);
		assertEquals(11L, (long) state.get(11L));
		assertEquals(22L, (long) state.get(22L));
		assertEquals(33L, (long) state.get(33L));
		assertEquals(44L, (long) state.get(44L));
		assertEquals(55L, (long) state.get(55L));
		assertEquals(5, getStateSize(state));

		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot = keyedBackend.snapshot(
				1L,
				1L,
				new MemCheckpointStreamFactory(4 * 1024 * 1024),
				CheckpointOptions.forCheckpointWithDefaultLocation());

		snapshot.run();
	}
}
 

/**
 * 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();
}
 

private CheckpointStreamFactory createCheckpointStreamFactory() {
	return new MemCheckpointStreamFactory(16 * 1024);
}
 

@Override
public CheckpointStreamFactory resolveCheckpointStorageLocation(long checkpointId, CheckpointStorageLocationReference reference) {
	return new MemCheckpointStreamFactory(maxStateSize);
}
 

private TestCheckpointStorageWorkerView(int maxStateSize, CheckpointedStateScope taskOwnedStateScope) {
	this.maxStateSize = maxStateSize;
	this.taskOwnedCheckpointStreamFactory = new MemCheckpointStreamFactory(maxStateSize);
	this.taskOwnedStateScope = taskOwnedStateScope;
}
 

@Override
public CheckpointStreamFactory.CheckpointStateOutputStream createTaskOwnedStateStream() {
	return new MemCheckpointStreamFactory.MemoryCheckpointOutputStream(Integer.MAX_VALUE);
}
 

@Override
public CheckpointStreamFactory resolveCheckpointStorageLocation(
		long checkpointId,
		CheckpointStorageLocationReference reference) {
	return new MemCheckpointStreamFactory(Integer.MAX_VALUE);
}
 

@Test
public void testMapStateMigrationAfterHashMapSerRemoval() throws Exception {
	ClassLoader cl = getClass().getClassLoader();
	URL resource = cl.getResource("heap_keyed_statebackend_1_5_map.snapshot");

	Preconditions.checkNotNull(resource, "Binary snapshot resource not found!");

	final SnapshotResult<KeyedStateHandle> stateHandles;
	try (BufferedInputStream bis = new BufferedInputStream((new FileInputStream(resource.getFile())))) {
		stateHandles = InstantiationUtil.deserializeObject(bis, Thread.currentThread().getContextClassLoader());
	}
	final KeyedStateHandle stateHandle = stateHandles.getJobManagerOwnedSnapshot();
	try (final HeapKeyedStateBackend<String> keyedBackend = createKeyedBackend(StateObjectCollection.singleton(stateHandle))) {
		final Integer namespace1 = 1;
		final Integer namespace2 = 2;
		final Integer namespace3 = 3;

		final MapStateDescriptor<Long, Long> stateDescr = new MapStateDescriptor<>("my-map-state", Long.class, Long.class);
		stateDescr.initializeSerializerUnlessSet(new ExecutionConfig());

		InternalMapState<String, Integer, Long, Long> state = keyedBackend.createInternalState(IntSerializer.INSTANCE, stateDescr);

		keyedBackend.setCurrentKey("abc");
		state.setCurrentNamespace(namespace1);
		assertEquals(33L, (long) state.get(33L));
		assertEquals(55L, (long) state.get(55L));
		assertEquals(2, getStateSize(state));

		state.setCurrentNamespace(namespace2);
		assertEquals(22L, (long) state.get(22L));
		assertEquals(11L, (long) state.get(11L));
		assertEquals(2, getStateSize(state));

		state.setCurrentNamespace(namespace3);
		assertEquals(44L, (long) state.get(44L));
		assertEquals(1, getStateSize(state));

		keyedBackend.setCurrentKey("def");
		state.setCurrentNamespace(namespace1);
		assertEquals(11L, (long) state.get(11L));
		assertEquals(44L, (long) state.get(44L));
		assertEquals(2, getStateSize(state));

		state.setCurrentNamespace(namespace3);
		assertEquals(22L, (long) state.get(22L));
		assertEquals(55L, (long) state.get(55L));
		assertEquals(33L, (long) state.get(33L));
		assertEquals(3, getStateSize(state));

		keyedBackend.setCurrentKey("jkl");
		state.setCurrentNamespace(namespace1);
		assertEquals(11L, (long) state.get(11L));
		assertEquals(22L, (long) state.get(22L));
		assertEquals(33L, (long) state.get(33L));
		assertEquals(44L, (long) state.get(44L));
		assertEquals(55L, (long) state.get(55L));
		assertEquals(5, getStateSize(state));

		keyedBackend.setCurrentKey("mno");
		state.setCurrentNamespace(namespace3);
		assertEquals(11L, (long) state.get(11L));
		assertEquals(22L, (long) state.get(22L));
		assertEquals(33L, (long) state.get(33L));
		assertEquals(44L, (long) state.get(44L));
		assertEquals(55L, (long) state.get(55L));
		assertEquals(5, getStateSize(state));

		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot = keyedBackend.snapshot(
				1L,
				1L,
				new MemCheckpointStreamFactory(4 * 1024 * 1024),
				CheckpointOptions.forCheckpointWithDefaultLocation());

		snapshot.run();
	}
}
 

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

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

/**
 * 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();
}
 

@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
public void testMapStateMigrationAfterHashMapSerRemoval() throws Exception {
	ClassLoader cl = getClass().getClassLoader();
	URL resource = cl.getResource("heap_keyed_statebackend_1_5_map.snapshot");

	Preconditions.checkNotNull(resource, "Binary snapshot resource not found!");

	final SnapshotResult<KeyedStateHandle> stateHandles;
	try (BufferedInputStream bis = new BufferedInputStream((new FileInputStream(resource.getFile())))) {
		stateHandles = InstantiationUtil.deserializeObject(bis, Thread.currentThread().getContextClassLoader());
	}
	final KeyedStateHandle stateHandle = stateHandles.getJobManagerOwnedSnapshot();
	try (final HeapKeyedStateBackend<String> keyedBackend = createKeyedBackend(StateObjectCollection.singleton(stateHandle))) {
		final Integer namespace1 = 1;
		final Integer namespace2 = 2;
		final Integer namespace3 = 3;

		final MapStateDescriptor<Long, Long> stateDescr = new MapStateDescriptor<>("my-map-state", Long.class, Long.class);
		stateDescr.initializeSerializerUnlessSet(new ExecutionConfig());

		InternalMapState<String, Integer, Long, Long> state = keyedBackend.createInternalState(IntSerializer.INSTANCE, stateDescr);

		keyedBackend.setCurrentKey("abc");
		state.setCurrentNamespace(namespace1);
		assertEquals(33L, (long) state.get(33L));
		assertEquals(55L, (long) state.get(55L));
		assertEquals(2, getStateSize(state));

		state.setCurrentNamespace(namespace2);
		assertEquals(22L, (long) state.get(22L));
		assertEquals(11L, (long) state.get(11L));
		assertEquals(2, getStateSize(state));

		state.setCurrentNamespace(namespace3);
		assertEquals(44L, (long) state.get(44L));
		assertEquals(1, getStateSize(state));

		keyedBackend.setCurrentKey("def");
		state.setCurrentNamespace(namespace1);
		assertEquals(11L, (long) state.get(11L));
		assertEquals(44L, (long) state.get(44L));
		assertEquals(2, getStateSize(state));

		state.setCurrentNamespace(namespace3);
		assertEquals(22L, (long) state.get(22L));
		assertEquals(55L, (long) state.get(55L));
		assertEquals(33L, (long) state.get(33L));
		assertEquals(3, getStateSize(state));

		keyedBackend.setCurrentKey("jkl");
		state.setCurrentNamespace(namespace1);
		assertEquals(11L, (long) state.get(11L));
		assertEquals(22L, (long) state.get(22L));
		assertEquals(33L, (long) state.get(33L));
		assertEquals(44L, (long) state.get(44L));
		assertEquals(55L, (long) state.get(55L));
		assertEquals(5, getStateSize(state));

		keyedBackend.setCurrentKey("mno");
		state.setCurrentNamespace(namespace3);
		assertEquals(11L, (long) state.get(11L));
		assertEquals(22L, (long) state.get(22L));
		assertEquals(33L, (long) state.get(33L));
		assertEquals(44L, (long) state.get(44L));
		assertEquals(55L, (long) state.get(55L));
		assertEquals(5, getStateSize(state));

		RunnableFuture<SnapshotResult<KeyedStateHandle>> snapshot = keyedBackend.snapshot(
				1L,
				1L,
				new MemCheckpointStreamFactory(4 * 1024 * 1024),
				CheckpointOptions.forCheckpointWithDefaultLocation());

		snapshot.run();
	}
}
 

private CheckpointStreamFactory createCheckpointStreamFactory() {
	return new MemCheckpointStreamFactory(16 * 1024);
}
 

private CheckpointStreamFactory createCheckpointStreamFactory() {
	return new MemCheckpointStreamFactory(16 * 1024);
}