Java Code Examples for org.apache.flink.runtime.state.AbstractKeyedStateBackend#getPartitionedState()

The following examples show how to use org.apache.flink.runtime.state.AbstractKeyedStateBackend#getPartitionedState() . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.

Example 1

Source File: KeyedStateInputFormat.java From flink with Apache License 2.0

6 votes

private Context(AbstractKeyedStateBackend<K> keyedStateBackend, InternalTimerService<VoidNamespace> timerService) throws Exception {
	eventTimers = keyedStateBackend.getPartitionedState(
		USER_TIMERS_NAME,
		StringSerializer.INSTANCE,
		new ListStateDescriptor<>(EVENT_TIMER_STATE, Types.LONG)
	);

	timerService.forEachEventTimeTimer((namespace, timer) -> {
		if (namespace.equals(VoidNamespace.INSTANCE)) {
			eventTimers.add(timer);
		}
	});

	procTimers = keyedStateBackend.getPartitionedState(
		USER_TIMERS_NAME,
		StringSerializer.INSTANCE,
		new ListStateDescriptor<>(PROC_TIMER_STATE, Types.LONG)
	);

	timerService.forEachProcessingTimeTimer((namespace, timer) -> {
		if (namespace.equals(VoidNamespace.INSTANCE)) {
			procTimers.add(timer);
		}
	});
}

Example 2

Source File: RocksDBStateBackendTest.java From Flink-CEPplus with Apache License 2.0

5 votes

@Test
public void testDisposeDeletesAllDirectories() throws Exception {
	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
	Collection<File> allFilesInDbDir =
		FileUtils.listFilesAndDirs(new File(dbPath), new AcceptAllFilter(), new AcceptAllFilter());
	try {
		ValueStateDescriptor<String> kvId =
			new ValueStateDescriptor<>("id", String.class, null);

		kvId.initializeSerializerUnlessSet(new ExecutionConfig());

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

		backend.setCurrentKey(1);
		state.update("Hello");

		// more than just the root directory
		assertTrue(allFilesInDbDir.size() > 1);
	} finally {
		IOUtils.closeQuietly(backend);
		backend.dispose();
	}
	allFilesInDbDir =
		FileUtils.listFilesAndDirs(new File(dbPath), new AcceptAllFilter(), new AcceptAllFilter());

	// just the root directory left
	assertEquals(1, allFilesInDbDir.size());
}

Example 3

Source File: RocksDBStateBackendTest.java From flink with Apache License 2.0

5 votes

@Test
public void testDisposeDeletesAllDirectories() throws Exception {
	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
	Collection<File> allFilesInDbDir =
		FileUtils.listFilesAndDirs(new File(dbPath), new AcceptAllFilter(), new AcceptAllFilter());
	try {
		ValueStateDescriptor<String> kvId =
			new ValueStateDescriptor<>("id", String.class, null);

		kvId.initializeSerializerUnlessSet(new ExecutionConfig());

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

		backend.setCurrentKey(1);
		state.update("Hello");

		// more than just the root directory
		assertTrue(allFilesInDbDir.size() > 1);
	} finally {
		IOUtils.closeQuietly(backend);
		backend.dispose();
	}
	allFilesInDbDir =
		FileUtils.listFilesAndDirs(new File(dbPath), new AcceptAllFilter(), new AcceptAllFilter());

	// just the root directory left
	assertEquals(1, allFilesInDbDir.size());
}

Example 4

Source File: RocksDBStateBackendTest.java From flink with Apache License 2.0

5 votes

@Test
public void testDisposeDeletesAllDirectories() throws Exception {
	AbstractKeyedStateBackend<Integer> backend = createKeyedBackend(IntSerializer.INSTANCE);
	Collection<File> allFilesInDbDir =
		FileUtils.listFilesAndDirs(new File(dbPath), new AcceptAllFilter(), new AcceptAllFilter());
	try {
		ValueStateDescriptor<String> kvId =
			new ValueStateDescriptor<>("id", String.class, null);

		kvId.initializeSerializerUnlessSet(new ExecutionConfig());

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

		backend.setCurrentKey(1);
		state.update("Hello");

		// more than just the root directory
		assertTrue(allFilesInDbDir.size() > 1);
	} finally {
		IOUtils.closeQuietly(backend);
		backend.dispose();
	}
	allFilesInDbDir =
		FileUtils.listFilesAndDirs(new File(dbPath), new AcceptAllFilter(), new AcceptAllFilter());

	// just the root directory left
	assertEquals(1, allFilesInDbDir.size());
}

Example 5

Source File: RocksDBStateBackendTest.java From Flink-CEPplus with Apache License 2.0

4 votes

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

Example 6

Source File: RocksDBAsyncSnapshotTest.java From Flink-CEPplus with Apache License 2.0

4 votes

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

Example 7

Source File: KvStateServerHandlerTest.java From Flink-CEPplus with Apache License 2.0

4 votes

/**
 * Tests a simple successful query via an EmbeddedChannel.
 */
@Test
public void testSimpleQuery() throws Exception {
	KvStateRegistry registry = new KvStateRegistry();
	AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();

	MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

	KvStateServerHandler handler = new KvStateServerHandler(testServer, registry, serializer, stats);
	EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);

	// Register state
	ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>("any", IntSerializer.INSTANCE);
	desc.setQueryable("vanilla");

	int numKeyGroups = 1;
	AbstractStateBackend abstractBackend = new MemoryStateBackend();
	DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
	dummyEnv.setKvStateRegistry(registry);
	AbstractKeyedStateBackend<Integer> backend = createKeyedStateBackend(registry, numKeyGroups, abstractBackend, dummyEnv);

	final TestRegistryListener registryListener = new TestRegistryListener();
	registry.registerListener(dummyEnv.getJobID(), registryListener);

	// Update the KvState and request it
	int expectedValue = 712828289;

	int key = 99812822;
	backend.setCurrentKey(key);
	ValueState<Integer> state = backend.getPartitionedState(
			VoidNamespace.INSTANCE,
			VoidNamespaceSerializer.INSTANCE,
			desc);

	state.update(expectedValue);

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key,
			IntSerializer.INSTANCE,
			VoidNamespace.INSTANCE,
			VoidNamespaceSerializer.INSTANCE);

	long requestId = Integer.MAX_VALUE + 182828L;

	assertTrue(registryListener.registrationName.equals("vanilla"));

	KvStateInternalRequest request = new KvStateInternalRequest(
			registryListener.kvStateId, serializedKeyAndNamespace);

	ByteBuf serRequest = MessageSerializer.serializeRequest(channel.alloc(), requestId, request);

	// Write the request and wait for the response
	channel.writeInbound(serRequest);

	ByteBuf buf = (ByteBuf) readInboundBlocking(channel);
	buf.skipBytes(4); // skip frame length

	// Verify the response
	assertEquals(MessageType.REQUEST_RESULT, MessageSerializer.deserializeHeader(buf));
	long deserRequestId = MessageSerializer.getRequestId(buf);
	KvStateResponse response = serializer.deserializeResponse(buf);

	assertEquals(requestId, deserRequestId);

	int actualValue = KvStateSerializer.deserializeValue(response.getContent(), IntSerializer.INSTANCE);
	assertEquals(expectedValue, actualValue);

	assertEquals(stats.toString(), 1, stats.getNumRequests());

	// Wait for async successful request report
	long deadline = System.nanoTime() + TimeUnit.NANOSECONDS.convert(30, TimeUnit.SECONDS);
	while (stats.getNumSuccessful() != 1L && System.nanoTime() <= deadline) {
		Thread.sleep(10L);
	}

	assertEquals(stats.toString(), 1L, stats.getNumSuccessful());
}

Example 8

Source File: KvStateServerHandlerTest.java From Flink-CEPplus with Apache License 2.0

4 votes

/**
 * Tests that large responses are chunked.
 */
@Test
public void testChunkedResponse() throws Exception {
	KvStateRegistry registry = new KvStateRegistry();
	KvStateRequestStats stats = new AtomicKvStateRequestStats();

	MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

	KvStateServerHandler handler = new KvStateServerHandler(testServer, registry, serializer, stats);
	EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);

	int numKeyGroups = 1;
	AbstractStateBackend abstractBackend = new MemoryStateBackend();
	DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
	dummyEnv.setKvStateRegistry(registry);
	AbstractKeyedStateBackend<Integer> backend = createKeyedStateBackend(registry, numKeyGroups, abstractBackend, dummyEnv);

	final TestRegistryListener registryListener = new TestRegistryListener();
	registry.registerListener(dummyEnv.getJobID(), registryListener);

	// Register state
	ValueStateDescriptor<byte[]> desc = new ValueStateDescriptor<>("any", BytePrimitiveArraySerializer.INSTANCE);
	desc.setQueryable("vanilla");

	ValueState<byte[]> state = backend.getPartitionedState(
			VoidNamespace.INSTANCE,
			VoidNamespaceSerializer.INSTANCE,
			desc);

	// Update KvState
	byte[] bytes = new byte[2 * channel.config().getWriteBufferHighWaterMark()];

	byte current = 0;
	for (int i = 0; i < bytes.length; i++) {
		bytes[i] = current++;
	}

	int key = 99812822;
	backend.setCurrentKey(key);
	state.update(bytes);

	// Request
	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key,
			IntSerializer.INSTANCE,
			VoidNamespace.INSTANCE,
			VoidNamespaceSerializer.INSTANCE);

	long requestId = Integer.MAX_VALUE + 182828L;

	assertTrue(registryListener.registrationName.equals("vanilla"));

	KvStateInternalRequest request = new KvStateInternalRequest(registryListener.kvStateId, serializedKeyAndNamespace);
	ByteBuf serRequest = MessageSerializer.serializeRequest(channel.alloc(), requestId, request);

	// Write the request and wait for the response
	channel.writeInbound(serRequest);

	Object msg = readInboundBlocking(channel);
	assertTrue("Not ChunkedByteBuf", msg instanceof ChunkedByteBuf);
}

Example 9

Source File: RocksDBStateBackendTest.java From flink with Apache License 2.0

4 votes

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

Example 10

Source File: RocksDBAsyncSnapshotTest.java From flink with Apache License 2.0

4 votes

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

Example 11

Source File: KvStateServerHandlerTest.java From flink with Apache License 2.0

4 votes

/**
 * Tests a simple successful query via an EmbeddedChannel.
 */
@Test
public void testSimpleQuery() throws Exception {
	KvStateRegistry registry = new KvStateRegistry();
	AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();

	MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

	KvStateServerHandler handler = new KvStateServerHandler(testServer, registry, serializer, stats);
	EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);

	// Register state
	ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>("any", IntSerializer.INSTANCE);
	desc.setQueryable("vanilla");

	int numKeyGroups = 1;
	AbstractStateBackend abstractBackend = new MemoryStateBackend();
	DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
	dummyEnv.setKvStateRegistry(registry);
	AbstractKeyedStateBackend<Integer> backend = createKeyedStateBackend(registry, numKeyGroups, abstractBackend, dummyEnv);

	final TestRegistryListener registryListener = new TestRegistryListener();
	registry.registerListener(dummyEnv.getJobID(), registryListener);

	// Update the KvState and request it
	int expectedValue = 712828289;

	int key = 99812822;
	backend.setCurrentKey(key);
	ValueState<Integer> state = backend.getPartitionedState(
			VoidNamespace.INSTANCE,
			VoidNamespaceSerializer.INSTANCE,
			desc);

	state.update(expectedValue);

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key,
			IntSerializer.INSTANCE,
			VoidNamespace.INSTANCE,
			VoidNamespaceSerializer.INSTANCE);

	long requestId = Integer.MAX_VALUE + 182828L;

	assertTrue(registryListener.registrationName.equals("vanilla"));

	KvStateInternalRequest request = new KvStateInternalRequest(
			registryListener.kvStateId, serializedKeyAndNamespace);

	ByteBuf serRequest = MessageSerializer.serializeRequest(channel.alloc(), requestId, request);

	// Write the request and wait for the response
	channel.writeInbound(serRequest);

	ByteBuf buf = (ByteBuf) readInboundBlocking(channel);
	buf.skipBytes(4); // skip frame length

	// Verify the response
	assertEquals(MessageType.REQUEST_RESULT, MessageSerializer.deserializeHeader(buf));
	long deserRequestId = MessageSerializer.getRequestId(buf);
	KvStateResponse response = serializer.deserializeResponse(buf);

	assertEquals(requestId, deserRequestId);

	int actualValue = KvStateSerializer.deserializeValue(response.getContent(), IntSerializer.INSTANCE);
	assertEquals(expectedValue, actualValue);

	assertEquals(stats.toString(), 1, stats.getNumRequests());

	// Wait for async successful request report
	long deadline = System.nanoTime() + TimeUnit.NANOSECONDS.convert(30, TimeUnit.SECONDS);
	while (stats.getNumSuccessful() != 1L && System.nanoTime() <= deadline) {
		Thread.sleep(10L);
	}

	assertEquals(stats.toString(), 1L, stats.getNumSuccessful());
}

Example 12

Source File: KvStateServerHandlerTest.java From flink with Apache License 2.0

4 votes

/**
 * Tests that large responses are chunked.
 */
@Test
public void testChunkedResponse() throws Exception {
	KvStateRegistry registry = new KvStateRegistry();
	KvStateRequestStats stats = new AtomicKvStateRequestStats();

	MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

	KvStateServerHandler handler = new KvStateServerHandler(testServer, registry, serializer, stats);
	EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);

	int numKeyGroups = 1;
	AbstractStateBackend abstractBackend = new MemoryStateBackend();
	DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
	dummyEnv.setKvStateRegistry(registry);
	AbstractKeyedStateBackend<Integer> backend = createKeyedStateBackend(registry, numKeyGroups, abstractBackend, dummyEnv);

	final TestRegistryListener registryListener = new TestRegistryListener();
	registry.registerListener(dummyEnv.getJobID(), registryListener);

	// Register state
	ValueStateDescriptor<byte[]> desc = new ValueStateDescriptor<>("any", BytePrimitiveArraySerializer.INSTANCE);
	desc.setQueryable("vanilla");

	ValueState<byte[]> state = backend.getPartitionedState(
			VoidNamespace.INSTANCE,
			VoidNamespaceSerializer.INSTANCE,
			desc);

	// Update KvState
	byte[] bytes = new byte[2 * channel.config().getWriteBufferHighWaterMark()];

	byte current = 0;
	for (int i = 0; i < bytes.length; i++) {
		bytes[i] = current++;
	}

	int key = 99812822;
	backend.setCurrentKey(key);
	state.update(bytes);

	// Request
	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key,
			IntSerializer.INSTANCE,
			VoidNamespace.INSTANCE,
			VoidNamespaceSerializer.INSTANCE);

	long requestId = Integer.MAX_VALUE + 182828L;

	assertTrue(registryListener.registrationName.equals("vanilla"));

	KvStateInternalRequest request = new KvStateInternalRequest(registryListener.kvStateId, serializedKeyAndNamespace);
	ByteBuf serRequest = MessageSerializer.serializeRequest(channel.alloc(), requestId, request);

	// Write the request and wait for the response
	channel.writeInbound(serRequest);

	Object msg = readInboundBlocking(channel);
	assertTrue("Not ChunkedByteBuf", msg instanceof ChunkedByteBuf);
}

Example 13

Source File: RocksDBStateBackendTest.java From flink with Apache License 2.0

4 votes

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

Example 14

Source File: RocksDBAsyncSnapshotTest.java From flink with Apache License 2.0

4 votes

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

Example 15

Source File: KvStateServerHandlerTest.java From flink with Apache License 2.0

4 votes

/**
 * Tests a simple successful query via an EmbeddedChannel.
 */
@Test
public void testSimpleQuery() throws Exception {
	KvStateRegistry registry = new KvStateRegistry();
	AtomicKvStateRequestStats stats = new AtomicKvStateRequestStats();

	MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

	KvStateServerHandler handler = new KvStateServerHandler(testServer, registry, serializer, stats);
	EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);

	// Register state
	ValueStateDescriptor<Integer> desc = new ValueStateDescriptor<>("any", IntSerializer.INSTANCE);
	desc.setQueryable("vanilla");

	int numKeyGroups = 1;
	AbstractStateBackend abstractBackend = new MemoryStateBackend();
	DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
	dummyEnv.setKvStateRegistry(registry);
	AbstractKeyedStateBackend<Integer> backend = createKeyedStateBackend(registry, numKeyGroups, abstractBackend, dummyEnv);

	final TestRegistryListener registryListener = new TestRegistryListener();
	registry.registerListener(dummyEnv.getJobID(), registryListener);

	// Update the KvState and request it
	int expectedValue = 712828289;

	int key = 99812822;
	backend.setCurrentKey(key);
	ValueState<Integer> state = backend.getPartitionedState(
			VoidNamespace.INSTANCE,
			VoidNamespaceSerializer.INSTANCE,
			desc);

	state.update(expectedValue);

	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key,
			IntSerializer.INSTANCE,
			VoidNamespace.INSTANCE,
			VoidNamespaceSerializer.INSTANCE);

	long requestId = Integer.MAX_VALUE + 182828L;

	assertTrue(registryListener.registrationName.equals("vanilla"));

	KvStateInternalRequest request = new KvStateInternalRequest(
			registryListener.kvStateId, serializedKeyAndNamespace);

	ByteBuf serRequest = MessageSerializer.serializeRequest(channel.alloc(), requestId, request);

	// Write the request and wait for the response
	channel.writeInbound(serRequest);

	ByteBuf buf = (ByteBuf) readInboundBlocking(channel);
	buf.skipBytes(4); // skip frame length

	// Verify the response
	assertEquals(MessageType.REQUEST_RESULT, MessageSerializer.deserializeHeader(buf));
	long deserRequestId = MessageSerializer.getRequestId(buf);
	KvStateResponse response = serializer.deserializeResponse(buf);
	buf.release();

	assertEquals(requestId, deserRequestId);

	int actualValue = KvStateSerializer.deserializeValue(response.getContent(), IntSerializer.INSTANCE);
	assertEquals(expectedValue, actualValue);

	assertEquals(stats.toString(), 1, stats.getNumRequests());

	// Wait for async successful request report
	long deadline = System.nanoTime() + TimeUnit.NANOSECONDS.convert(30, TimeUnit.SECONDS);
	while (stats.getNumSuccessful() != 1L && System.nanoTime() <= deadline) {
		Thread.sleep(10L);
	}

	assertEquals(stats.toString(), 1L, stats.getNumSuccessful());
}

Example 16

Source File: KvStateServerHandlerTest.java From flink with Apache License 2.0

4 votes

/**
 * Tests that large responses are chunked.
 */
@Test
public void testChunkedResponse() throws Exception {
	KvStateRegistry registry = new KvStateRegistry();
	KvStateRequestStats stats = new AtomicKvStateRequestStats();

	MessageSerializer<KvStateInternalRequest, KvStateResponse> serializer =
			new MessageSerializer<>(new KvStateInternalRequest.KvStateInternalRequestDeserializer(), new KvStateResponse.KvStateResponseDeserializer());

	KvStateServerHandler handler = new KvStateServerHandler(testServer, registry, serializer, stats);
	EmbeddedChannel channel = new EmbeddedChannel(getFrameDecoder(), handler);

	int numKeyGroups = 1;
	AbstractStateBackend abstractBackend = new MemoryStateBackend();
	DummyEnvironment dummyEnv = new DummyEnvironment("test", 1, 0);
	dummyEnv.setKvStateRegistry(registry);
	AbstractKeyedStateBackend<Integer> backend = createKeyedStateBackend(registry, numKeyGroups, abstractBackend, dummyEnv);

	final TestRegistryListener registryListener = new TestRegistryListener();
	registry.registerListener(dummyEnv.getJobID(), registryListener);

	// Register state
	ValueStateDescriptor<byte[]> desc = new ValueStateDescriptor<>("any", BytePrimitiveArraySerializer.INSTANCE);
	desc.setQueryable("vanilla");

	ValueState<byte[]> state = backend.getPartitionedState(
			VoidNamespace.INSTANCE,
			VoidNamespaceSerializer.INSTANCE,
			desc);

	// Update KvState
	byte[] bytes = new byte[2 * channel.config().getWriteBufferHighWaterMark()];

	byte current = 0;
	for (int i = 0; i < bytes.length; i++) {
		bytes[i] = current++;
	}

	int key = 99812822;
	backend.setCurrentKey(key);
	state.update(bytes);

	// Request
	byte[] serializedKeyAndNamespace = KvStateSerializer.serializeKeyAndNamespace(
			key,
			IntSerializer.INSTANCE,
			VoidNamespace.INSTANCE,
			VoidNamespaceSerializer.INSTANCE);

	long requestId = Integer.MAX_VALUE + 182828L;

	assertTrue(registryListener.registrationName.equals("vanilla"));

	KvStateInternalRequest request = new KvStateInternalRequest(registryListener.kvStateId, serializedKeyAndNamespace);
	ByteBuf serRequest = MessageSerializer.serializeRequest(channel.alloc(), requestId, request);

	// Write the request and wait for the response
	channel.writeInbound(serRequest);

	Object msg = readInboundBlocking(channel);
	assertTrue("Not ChunkedByteBuf", msg instanceof ChunkedByteBuf);
	((ChunkedByteBuf) msg).close();
}