org.apache.flink.contrib.streaming.state.RocksDBKeyedStateBackend Java Examples

@Before
public void init() throws IOException {
	rocksDir = tmp.newFolder();

	// ensure the RocksDB library is loaded to a distinct location each retry
	NativeLibraryLoader.getInstance().loadLibrary(rocksDir.getAbsolutePath());

	options = new Options()
			.setCompactionStyle(CompactionStyle.LEVEL)
			.setLevelCompactionDynamicLevelBytes(true)
			.setIncreaseParallelism(4)
			.setUseFsync(false)
			.setMaxOpenFiles(-1)
			.setCreateIfMissing(true)
			.setMergeOperatorName(RocksDBKeyedStateBackend.MERGE_OPERATOR_NAME);

	writeOptions = new WriteOptions()
			.setSync(false)
			.setDisableWAL(true);
}
 

/**
 * Tests map serialization and deserialization match.
 *
 * @see KvStateRequestSerializerTest#testMapSerialization()
 * KvStateRequestSerializerTest#testMapSerialization() using the heap state back-end
 * test
 */
@Test
public void testMapSerialization() throws Exception {
	final long key = 0L;

	// objects for RocksDB state list serialisation
	final RocksDBKeyedStateBackend<Long> longHeapKeyedStateBackend = RocksDBTestUtils
		.builderForTestDefaults(temporaryFolder.getRoot(), LongSerializer.INSTANCE)
		.build();

	longHeapKeyedStateBackend.setCurrentKey(key);

	final InternalMapState<Long, VoidNamespace, Long, String> mapState =
			(InternalMapState<Long, VoidNamespace, Long, String>)
					longHeapKeyedStateBackend.getPartitionedState(
							VoidNamespace.INSTANCE,
							VoidNamespaceSerializer.INSTANCE,
							new MapStateDescriptor<>("test", LongSerializer.INSTANCE, StringSerializer.INSTANCE));

	KvStateRequestSerializerTest.testMapSerialization(key, mapState);
	longHeapKeyedStateBackend.dispose();
}
 

/**
 * Tests list serialization and deserialization match.
 *
 * @see KvStateRequestSerializerTest#testListSerialization()
 * KvStateRequestSerializerTest#testListSerialization() using the heap state back-end
 * test
 */
@Test
public void testListSerialization() throws Exception {
	final long key = 0L;

	final RocksDBKeyedStateBackend<Long> longHeapKeyedStateBackend = RocksDBTestUtils
		.builderForTestDefaults(temporaryFolder.getRoot(), LongSerializer.INSTANCE)
		.build();

	longHeapKeyedStateBackend.setCurrentKey(key);

	final InternalListState<Long, VoidNamespace, Long> listState = longHeapKeyedStateBackend.createInternalState(VoidNamespaceSerializer.INSTANCE,
			new ListStateDescriptor<>("test", LongSerializer.INSTANCE));

	KvStateRequestSerializerTest.testListSerialization(key, listState);
	longHeapKeyedStateBackend.dispose();
}
 

@Before
public void init() throws IOException {
	rocksDir = tmp.newFolder();

	// ensure the RocksDB library is loaded to a distinct location each retry
	NativeLibraryLoader.getInstance().loadLibrary(rocksDir.getAbsolutePath());

	options = new Options()
			.setCompactionStyle(CompactionStyle.LEVEL)
			.setLevelCompactionDynamicLevelBytes(true)
			.setIncreaseParallelism(4)
			.setUseFsync(false)
			.setMaxOpenFiles(-1)
			.setCreateIfMissing(true)
			.setMergeOperatorName(RocksDBKeyedStateBackend.MERGE_OPERATOR_NAME);

	writeOptions = new WriteOptions()
			.setSync(false)
			.setDisableWAL(true);
}
 

@TearDown(Level.Iteration)
public void tearDownPerIteration() throws Exception {
    applyToAllKeys(
            keyedStateBackend,
            STATE_DESC,
            (k, state) -> {
                keyedStateBackend.setCurrentKey(k);
                state.clear();
            });
    // make the clearance effective, trigger compaction for RocksDB, and GC for heap.
    if (keyedStateBackend instanceof RocksDBKeyedStateBackend) {
        RocksDBKeyedStateBackend<Long> rocksDBKeyedStateBackend = (RocksDBKeyedStateBackend<Long>) keyedStateBackend;
        compactState(rocksDBKeyedStateBackend, STATE_DESC);
    } else {
        System.gc();
    }
    // wait a while for the clearance to take effect.
    Thread.sleep(1000);
}
 

@Before
public void init() throws IOException {
	rocksDir = tmp.newFolder();

	// ensure the RocksDB library is loaded to a distinct location each retry
	NativeLibraryLoader.getInstance().loadLibrary(rocksDir.getAbsolutePath());

	options = new Options()
			.setCompactionStyle(CompactionStyle.LEVEL)
			.setLevelCompactionDynamicLevelBytes(true)
			.setIncreaseParallelism(4)
			.setUseFsync(false)
			.setMaxOpenFiles(-1)
			.setCreateIfMissing(true)
			.setMergeOperatorName(RocksDBKeyedStateBackend.MERGE_OPERATOR_NAME);

	writeOptions = new WriteOptions()
			.setSync(false)
			.setDisableWAL(true);
}
 

/**
 * Tests list serialization and deserialization match.
 *
 * @see KvStateRequestSerializerTest#testListSerialization()
 * KvStateRequestSerializerTest#testListSerialization() using the heap state back-end
 * test
 */
@Test
public void testListSerialization() throws Exception {
	final long key = 0L;

	// objects for RocksDB state list serialisation
	DBOptions dbOptions = PredefinedOptions.DEFAULT.createDBOptions();
	dbOptions.setCreateIfMissing(true);
	ExecutionConfig executionConfig = new ExecutionConfig();
	final RocksDBKeyedStateBackend<Long> longHeapKeyedStateBackend =
		new RocksDBKeyedStateBackendBuilder<>(
			"no-op",
			ClassLoader.getSystemClassLoader(),
			temporaryFolder.getRoot(),
			dbOptions,
			stateName -> PredefinedOptions.DEFAULT.createColumnOptions(),
			mock(TaskKvStateRegistry.class),
			LongSerializer.INSTANCE,
			1,
			new KeyGroupRange(0, 0),
			executionConfig,
			TestLocalRecoveryConfig.disabled(),
			RocksDBStateBackend.PriorityQueueStateType.HEAP,
			TtlTimeProvider.DEFAULT,
			new UnregisteredMetricsGroup(),
			Collections.emptyList(),
			AbstractStateBackend.getCompressionDecorator(executionConfig),
			new CloseableRegistry()
		).build();
	longHeapKeyedStateBackend.setCurrentKey(key);

	final InternalListState<Long, VoidNamespace, Long> listState = longHeapKeyedStateBackend.createInternalState(VoidNamespaceSerializer.INSTANCE,
			new ListStateDescriptor<>("test", LongSerializer.INSTANCE));

	KvStateRequestSerializerTest.testListSerialization(key, listState);
	longHeapKeyedStateBackend.dispose();
}
 

@Setup(Level.Iteration)
public void setUpPerIteration() throws Exception {
    for (int i = 0; i < setupKeyCount; ++i) {
        keyedStateBackend.setCurrentKey((long) i);
        listState.add(random.nextLong());
    }
    // make sure only one sst file left, so all get invocation will access this single file,
    // to prevent the spike caused by different key distribution in multiple sst files,
    // the more access to the older sst file, the lower throughput will be.
    if (keyedStateBackend instanceof RocksDBKeyedStateBackend) {
        RocksDBKeyedStateBackend<Long> rocksDBKeyedStateBackend = (RocksDBKeyedStateBackend<Long>) keyedStateBackend;
        compactState(rocksDBKeyedStateBackend, STATE_DESC);
    }
}
 

private static RocksDBKeyedStateBackend<Long> createRocksDBKeyedStateBackend(File rootDir) throws IOException {
	File recoveryBaseDir = prepareDirectory(recoveryDirName, rootDir);
	File dbPathFile = prepareDirectory(dbDirName, rootDir);
	ExecutionConfig executionConfig = new ExecutionConfig();
	RocksDBResourceContainer resourceContainer = new RocksDBResourceContainer();
	RocksDBKeyedStateBackendBuilder<Long> builder = new RocksDBKeyedStateBackendBuilder<>(
		"Test",
		Thread.currentThread().getContextClassLoader(),
		dbPathFile,
		resourceContainer,
		stateName -> resourceContainer.getColumnOptions(),
		null,
		LongSerializer.INSTANCE,
		2,
		new KeyGroupRange(0, 1),
		executionConfig,
		new LocalRecoveryConfig(false, new LocalRecoveryDirectoryProviderImpl(recoveryBaseDir, new JobID(), new JobVertexID(), 0)),
		RocksDBStateBackend.PriorityQueueStateType.ROCKSDB,
		TtlTimeProvider.DEFAULT,
		new UnregisteredMetricsGroup(),
		Collections.emptyList(),
		AbstractStateBackend.getCompressionDecorator(executionConfig),
		new CloseableRegistry());
	try {
		return builder.build();
	} catch (Exception e) {
		IOUtils.closeQuietly(resourceContainer);
		throw e;
	}
}
 

@Test
public void testCompactState() throws Exception {
	KeyedStateBackend<Long> backend = createKeyedStateBackend(backendType);
	ListState<Long> listState = getListState(backend, listStateDescriptor);
	for (long i = 0; i < 10; i++) {
		backend.setCurrentKey(i);
		listState.add(i);
	}
	if (backend instanceof RocksDBKeyedStateBackend) {
		RocksDBKeyedStateBackend<Long> rocksDBKeyedStateBackend = (RocksDBKeyedStateBackend<Long>) backend;
		compactState(rocksDBKeyedStateBackend, listStateDescriptor);
	}
	cleanUp(backend);
}
 

/**
 * Tests list serialization and deserialization match.
 *
 * @see KvStateRequestSerializerTest#testListSerialization()
 * KvStateRequestSerializerTest#testListSerialization() using the heap state back-end
 * test
 */
@Test
public void testListSerialization() throws Exception {
	final long key = 0L;

	// objects for RocksDB state list serialisation
	DBOptions dbOptions = PredefinedOptions.DEFAULT.createDBOptions();
	dbOptions.setCreateIfMissing(true);
	ExecutionConfig executionConfig = new ExecutionConfig();
	final RocksDBKeyedStateBackend<Long> longHeapKeyedStateBackend =
		new RocksDBKeyedStateBackendBuilder<>(
			"no-op",
			ClassLoader.getSystemClassLoader(),
			temporaryFolder.getRoot(),
			dbOptions,
			stateName -> PredefinedOptions.DEFAULT.createColumnOptions(),
			mock(TaskKvStateRegistry.class),
			LongSerializer.INSTANCE,
			1,
			new KeyGroupRange(0, 0),
			executionConfig,
			TestLocalRecoveryConfig.disabled(),
			RocksDBStateBackend.PriorityQueueStateType.HEAP,
			TtlTimeProvider.DEFAULT,
			new UnregisteredMetricsGroup(),
			Collections.emptyList(),
			AbstractStateBackend.getCompressionDecorator(executionConfig),
			new CloseableRegistry()
		).build();
	longHeapKeyedStateBackend.setCurrentKey(key);

	final InternalListState<Long, VoidNamespace, Long> listState = longHeapKeyedStateBackend.createInternalState(VoidNamespaceSerializer.INSTANCE,
			new ListStateDescriptor<>("test", LongSerializer.INSTANCE));

	KvStateRequestSerializerTest.testListSerialization(key, listState);
	longHeapKeyedStateBackend.dispose();
}
 

@Test(timeout = 2000)
@RetryOnFailure(times = 3)
public void testRocksDbListStateAPIs() throws Exception {
	final File rocksDir = tmp.newFolder();

	// ensure the RocksDB library is loaded to a distinct location each retry
	NativeLibraryLoader.getInstance().loadLibrary(rocksDir.getAbsolutePath());

	final String key1 = "key1";
	final String key2 = "key2";
	final String value = "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ7890654321";

	final byte[] keyBytes1 = key1.getBytes(StandardCharsets.UTF_8);
	final byte[] keyBytes2 = key2.getBytes(StandardCharsets.UTF_8);
	final byte[] valueBytes = value.getBytes(StandardCharsets.UTF_8);

	// The number of values added to ListState. Can be changed for benchmarking
	final int num = 10;

	try (
		final Options options = new Options()
				.setCompactionStyle(CompactionStyle.LEVEL)
				.setLevelCompactionDynamicLevelBytes(true)
				.setIncreaseParallelism(4)
				.setUseFsync(false)
				.setMaxOpenFiles(-1)
				.setCreateIfMissing(true)
				.setMergeOperatorName(RocksDBKeyedStateBackend.MERGE_OPERATOR_NAME);

		final WriteOptions writeOptions = new WriteOptions()
				.setSync(false)
				.setDisableWAL(true);

		final RocksDB rocksDB = RocksDB.open(options, rocksDir.getAbsolutePath())) {

		// ----- add() API -----
		log.info("begin add");

		final long beginInsert1 = System.nanoTime();
		for (int i = 0; i < num; i++) {
			rocksDB.merge(writeOptions, keyBytes1, valueBytes);
		}
		final long endInsert1 = System.nanoTime();

		log.info("end add - duration: {} ns", (endInsert1 - beginInsert1));

		// ----- update() API -----

		List<byte[]> list = new ArrayList<>(num);
		for (int i = 0; i < num; i++) {
			list.add(valueBytes);
		}
		byte[] premerged = merge(list);

		log.info("begin update");

		final long beginInsert2 = System.nanoTime();
		rocksDB.merge(writeOptions, keyBytes2, premerged);
		final long endInsert2 = System.nanoTime();

		log.info("end update - duration: {} ns", (endInsert2 - beginInsert2));
	}
}
 

private void setTimeAndCompact(StateDescriptor<?, ?> stateDesc, long ts) throws RocksDBException {
	@SuppressWarnings("resource")
	RocksDBKeyedStateBackend<String> keyedBackend = sbetc.getKeyedStateBackend();
	timeProvider.time = ts;
	keyedBackend.compactState(stateDesc);
}
 

/**
 * Tests map serialization and deserialization match.
 *
 * @see KvStateRequestSerializerTest#testMapSerialization()
 * KvStateRequestSerializerTest#testMapSerialization() using the heap state back-end
 * test
 */
@Test
public void testMapSerialization() throws Exception {
	final long key = 0L;

	// objects for RocksDB state list serialisation
	DBOptions dbOptions = PredefinedOptions.DEFAULT.createDBOptions();
	dbOptions.setCreateIfMissing(true);
	ExecutionConfig executionConfig = new ExecutionConfig();
	final RocksDBKeyedStateBackend<Long> longHeapKeyedStateBackend =
		new RocksDBKeyedStateBackendBuilder<>(
			"no-op",
			ClassLoader.getSystemClassLoader(),
			temporaryFolder.getRoot(),
			dbOptions,
			stateName -> PredefinedOptions.DEFAULT.createColumnOptions(),
			mock(TaskKvStateRegistry.class),
			LongSerializer.INSTANCE,
			1,
			new KeyGroupRange(0, 0),
			executionConfig,
			TestLocalRecoveryConfig.disabled(),
			RocksDBStateBackend.PriorityQueueStateType.HEAP,
			TtlTimeProvider.DEFAULT,
			new UnregisteredMetricsGroup(),
			Collections.emptyList(),
			AbstractStateBackend.getCompressionDecorator(executionConfig),
			new CloseableRegistry()
		).build();
	longHeapKeyedStateBackend.setCurrentKey(key);

	final InternalMapState<Long, VoidNamespace, Long, String> mapState =
			(InternalMapState<Long, VoidNamespace, Long, String>)
					longHeapKeyedStateBackend.getPartitionedState(
							VoidNamespace.INSTANCE,
							VoidNamespaceSerializer.INSTANCE,
							new MapStateDescriptor<>("test", LongSerializer.INSTANCE, StringSerializer.INSTANCE));

	KvStateRequestSerializerTest.testMapSerialization(key, mapState);
	longHeapKeyedStateBackend.dispose();
}
 

/**
 * Tests map serialization and deserialization match.
 *
 * @see KvStateRequestSerializerTest#testMapSerialization()
 * KvStateRequestSerializerTest#testMapSerialization() using the heap state back-end
 * test
 */
@Test
public void testMapSerialization() throws Exception {
	final long key = 0L;

	// objects for RocksDB state list serialisation
	DBOptions dbOptions = PredefinedOptions.DEFAULT.createDBOptions();
	dbOptions.setCreateIfMissing(true);
	ExecutionConfig executionConfig = new ExecutionConfig();
	final RocksDBKeyedStateBackend<Long> longHeapKeyedStateBackend =
		new RocksDBKeyedStateBackendBuilder<>(
			"no-op",
			ClassLoader.getSystemClassLoader(),
			temporaryFolder.getRoot(),
			dbOptions,
			stateName -> PredefinedOptions.DEFAULT.createColumnOptions(),
			mock(TaskKvStateRegistry.class),
			LongSerializer.INSTANCE,
			1,
			new KeyGroupRange(0, 0),
			executionConfig,
			TestLocalRecoveryConfig.disabled(),
			RocksDBStateBackend.PriorityQueueStateType.HEAP,
			TtlTimeProvider.DEFAULT,
			new UnregisteredMetricsGroup(),
			Collections.emptyList(),
			AbstractStateBackend.getCompressionDecorator(executionConfig),
			new CloseableRegistry()
		).build();
	longHeapKeyedStateBackend.setCurrentKey(key);

	final InternalMapState<Long, VoidNamespace, Long, String> mapState =
			(InternalMapState<Long, VoidNamespace, Long, String>)
					longHeapKeyedStateBackend.getPartitionedState(
							VoidNamespace.INSTANCE,
							VoidNamespaceSerializer.INSTANCE,
							new MapStateDescriptor<>("test", LongSerializer.INSTANCE, StringSerializer.INSTANCE));

	KvStateRequestSerializerTest.testMapSerialization(key, mapState);
	longHeapKeyedStateBackend.dispose();
}
 

@Test(timeout = 2000)
@RetryOnFailure(times = 3)
public void testRocksDbListStateAPIs() throws Exception {
	final File rocksDir = tmp.newFolder();

	// ensure the RocksDB library is loaded to a distinct location each retry
	NativeLibraryLoader.getInstance().loadLibrary(rocksDir.getAbsolutePath());

	final String key1 = "key1";
	final String key2 = "key2";
	final String value = "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ7890654321";

	final byte[] keyBytes1 = key1.getBytes(StandardCharsets.UTF_8);
	final byte[] keyBytes2 = key2.getBytes(StandardCharsets.UTF_8);
	final byte[] valueBytes = value.getBytes(StandardCharsets.UTF_8);

	// The number of values added to ListState. Can be changed for benchmarking
	final int num = 10;

	try (
		final Options options = new Options()
				.setCompactionStyle(CompactionStyle.LEVEL)
				.setLevelCompactionDynamicLevelBytes(true)
				.setIncreaseParallelism(4)
				.setUseFsync(false)
				.setMaxOpenFiles(-1)
				.setCreateIfMissing(true)
				.setMergeOperatorName(RocksDBKeyedStateBackend.MERGE_OPERATOR_NAME);

		final WriteOptions writeOptions = new WriteOptions()
				.setSync(false)
				.setDisableWAL(true);

		final RocksDB rocksDB = RocksDB.open(options, rocksDir.getAbsolutePath())) {

		// ----- add() API -----
		log.info("begin add");

		final long beginInsert1 = System.nanoTime();
		for (int i = 0; i < num; i++) {
			rocksDB.merge(writeOptions, keyBytes1, valueBytes);
		}
		final long endInsert1 = System.nanoTime();

		log.info("end add - duration: {} ns", (endInsert1 - beginInsert1));

		// ----- update() API -----

		List<byte[]> list = new ArrayList<>(num);
		for (int i = 0; i < num; i++) {
			list.add(valueBytes);
		}
		byte[] premerged = merge(list);

		log.info("begin update");

		final long beginInsert2 = System.nanoTime();
		rocksDB.merge(writeOptions, keyBytes2, premerged);
		final long endInsert2 = System.nanoTime();

		log.info("end update - duration: {} ns", (endInsert2 - beginInsert2));
	}
}
 

public static <K, S extends State, T> void compactState(
	RocksDBKeyedStateBackend<K> backend,
	StateDescriptor<S, T> stateDescriptor) throws RocksDBException {
	backend.compactState(stateDescriptor);
}
 

private void setTimeAndCompact(StateDescriptor<?, ?> stateDesc, long ts) throws RocksDBException {
	@SuppressWarnings("resource")
	RocksDBKeyedStateBackend<String> keyedBackend = sbetc.getKeyedStateBackend();
	timeProvider.time = ts;
	keyedBackend.compactState(stateDesc);
}
 

private void setTimeAndCompact(StateDescriptor<?, ?> stateDesc, long ts) throws RocksDBException {
	@SuppressWarnings("resource")
	RocksDBKeyedStateBackend<String> keyedBackend = sbetc.getKeyedStateBackend();
	timeProvider.time = ts;
	keyedBackend.compactState(stateDesc);
}
 

@Test(timeout = 2000)
@RetryOnFailure(times = 3)
public void testRocksDbListStateAPIs() throws Exception {
	final File rocksDir = tmp.newFolder();

	// ensure the RocksDB library is loaded to a distinct location each retry
	NativeLibraryLoader.getInstance().loadLibrary(rocksDir.getAbsolutePath());

	final String key1 = "key1";
	final String key2 = "key2";
	final String value = "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ7890654321";

	final byte[] keyBytes1 = key1.getBytes(StandardCharsets.UTF_8);
	final byte[] keyBytes2 = key2.getBytes(StandardCharsets.UTF_8);
	final byte[] valueBytes = value.getBytes(StandardCharsets.UTF_8);

	// The number of values added to ListState. Can be changed for benchmarking
	final int num = 10;

	try (
		final Options options = new Options()
				.setCompactionStyle(CompactionStyle.LEVEL)
				.setLevelCompactionDynamicLevelBytes(true)
				.setIncreaseParallelism(4)
				.setUseFsync(false)
				.setMaxOpenFiles(-1)
				.setCreateIfMissing(true)
				.setMergeOperatorName(RocksDBKeyedStateBackend.MERGE_OPERATOR_NAME);

		final WriteOptions writeOptions = new WriteOptions()
				.setSync(false)
				.setDisableWAL(true);

		final RocksDB rocksDB = RocksDB.open(options, rocksDir.getAbsolutePath())) {

		// ----- add() API -----
		log.info("begin add");

		final long beginInsert1 = System.nanoTime();
		for (int i = 0; i < num; i++) {
			rocksDB.merge(writeOptions, keyBytes1, valueBytes);
		}
		final long endInsert1 = System.nanoTime();

		log.info("end add - duration: {} ns", (endInsert1 - beginInsert1));

		// ----- update() API -----

		List<byte[]> list = new ArrayList<>(num);
		for (int i = 0; i < num; i++) {
			list.add(valueBytes);
		}
		byte[] premerged = merge(list);

		log.info("begin update");

		final long beginInsert2 = System.nanoTime();
		rocksDB.merge(writeOptions, keyBytes2, premerged);
		final long endInsert2 = System.nanoTime();

		log.info("end update - duration: {} ns", (endInsert2 - beginInsert2));
	}
}