You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by se...@apache.org on 2019/12/06 17:00:58 UTC
[flink] 02/03: [FLINK-14484][state-backend] Control memory usage of
RocksDB via Cache and WriteBufferManager
This is an automated email from the ASF dual-hosted git repository.
sewen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
commit 112e86558df345c533322021823fb3d01f225854
Author: Yun Tang <my...@live.com>
AuthorDate: Wed Dec 4 15:41:18 2019 +0800
[FLINK-14484][state-backend] Control memory usage of RocksDB via Cache and WriteBufferManager
This closes #10416
---
.../streaming/state/RocksDBKeyedStateBackend.java | 5 +
.../contrib/streaming/state/RocksDBOptions.java | 20 ++
.../streaming/state/RocksDBSharedObjects.java | 50 +++++
.../streaming/state/RocksDBStateBackend.java | 142 +++++++++++++-
.../RocksDBStateBackendBoundedMemoryTest.java | 204 +++++++++++++++++++++
.../state/RocksDBStateBackendConfigTest.java | 50 +++++
6 files changed, 469 insertions(+), 2 deletions(-)
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
index 16c8757..07b1fd8 100644
--- a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
+++ b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBKeyedStateBackend.java
@@ -403,6 +403,11 @@ public class RocksDBKeyedStateBackend<K> extends AbstractKeyedStateBackend<K> {
}
@VisibleForTesting
+ DBOptions getDbOptions() {
+ return dbOptions;
+ }
+
+ @VisibleForTesting
boolean isDisposed() {
return this.disposed;
}
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBOptions.java b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBOptions.java
index 8ffd142..928882f 100644
--- a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBOptions.java
+++ b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBOptions.java
@@ -86,4 +86,24 @@ public class RocksDBOptions {
"The default options factory is %s, and it would read the configured options which provided in 'RocksDBConfigurableOptions'.",
DefaultConfigurableOptionsFactory.class.getName()));
+ public static final ConfigOption<String> BOUNDED_MEMORY_SIZE = ConfigOptions
+ .key("state.backend.rocksdb.per-slot.total.memory")
+ .stringType()
+ .noDefaultValue()
+ .withDescription("The total memory size shared among all RocksDB instances per slot. " +
+ "This option has no default value which means no bounded memory limit.");
+
+ public static final ConfigOption<Double> WRITE_BUFFER_RATIO = ConfigOptions
+ .key("state.backend.rocksdb.write-buffer.ratio")
+ .doubleType()
+ .defaultValue(0.5)
+ .withDescription(String.format("This option would only take effect when %s is configured, " +
+ "all RocksDB instances would share the same write buffer manager with the ratio of a LRUCache.", BOUNDED_MEMORY_SIZE));
+
+ public static final ConfigOption<Double> HIGH_PRI_POOL_RATIO = ConfigOptions
+ .key("state.backend.rocksdb.high-pri-pool.ratio")
+ .doubleType()
+ .defaultValue(0.1)
+ .withDescription(String.format("This option would only take effect when %s is configured, " +
+ "all RocksDB instances would share the high priority ratio for index, filter, and compression dictionary blocks in a LRUCache.", BOUNDED_MEMORY_SIZE));
}
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBSharedObjects.java b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBSharedObjects.java
new file mode 100644
index 0000000..188c506
--- /dev/null
+++ b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBSharedObjects.java
@@ -0,0 +1,50 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.contrib.streaming.state;
+
+import org.rocksdb.Cache;
+import org.rocksdb.WriteBufferManager;
+
+/**
+ * Shared objects among RocksDB instances per slot.
+ */
+public class RocksDBSharedObjects implements AutoCloseable {
+
+ private final Cache cache;
+ private final WriteBufferManager writeBufferManager;
+
+ RocksDBSharedObjects(Cache cache, WriteBufferManager writeBufferManager) {
+ this.cache = cache;
+ this.writeBufferManager = writeBufferManager;
+ }
+
+ @Override
+ public void close() {
+ writeBufferManager.close();
+ cache.close();
+ }
+
+ public Cache getCache() {
+ return cache;
+ }
+
+ public WriteBufferManager getWriteBufferManager() {
+ return writeBufferManager;
+ }
+}
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
index 8563895..90e8583 100644
--- a/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
+++ b/flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java
@@ -24,10 +24,12 @@ import org.apache.flink.api.common.typeutils.TypeSerializer;
import org.apache.flink.configuration.CheckpointingOptions;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.IllegalConfigurationException;
+import org.apache.flink.configuration.MemorySize;
import org.apache.flink.core.fs.CloseableRegistry;
import org.apache.flink.core.fs.Path;
import org.apache.flink.metrics.MetricGroup;
import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.memory.MemoryManager;
import org.apache.flink.runtime.query.TaskKvStateRegistry;
import org.apache.flink.runtime.state.AbstractKeyedStateBackend;
import org.apache.flink.runtime.state.AbstractStateBackend;
@@ -50,10 +52,15 @@ import org.apache.flink.util.FlinkRuntimeException;
import org.apache.flink.util.Preconditions;
import org.apache.flink.util.TernaryBoolean;
+import org.rocksdb.BlockBasedTableConfig;
+import org.rocksdb.Cache;
import org.rocksdb.ColumnFamilyOptions;
import org.rocksdb.DBOptions;
+import org.rocksdb.LRUCache;
import org.rocksdb.NativeLibraryLoader;
import org.rocksdb.RocksDB;
+import org.rocksdb.TableFormatConfig;
+import org.rocksdb.WriteBufferManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -70,10 +77,15 @@ import java.util.Collection;
import java.util.List;
import java.util.Random;
import java.util.UUID;
+import java.util.function.Function;
+import java.util.function.Supplier;
+import static org.apache.flink.contrib.streaming.state.RocksDBOptions.BOUNDED_MEMORY_SIZE;
import static org.apache.flink.contrib.streaming.state.RocksDBOptions.CHECKPOINT_TRANSFER_THREAD_NUM;
+import static org.apache.flink.contrib.streaming.state.RocksDBOptions.HIGH_PRI_POOL_RATIO;
import static org.apache.flink.contrib.streaming.state.RocksDBOptions.TIMER_SERVICE_FACTORY;
import static org.apache.flink.contrib.streaming.state.RocksDBOptions.TTL_COMPACT_FILTER_ENABLED;
+import static org.apache.flink.contrib.streaming.state.RocksDBOptions.WRITE_BUFFER_RATIO;
import static org.apache.flink.util.Preconditions.checkNotNull;
/**
@@ -111,6 +123,8 @@ public class RocksDBStateBackend extends AbstractStateBackend implements Configu
private static final int UNDEFINED_NUMBER_OF_TRANSFER_THREADS = -1;
+ static final long UNDEFINED_VALUE = -1;
+
// ------------------------------------------------------------------------
// -- configuration values, set in the application / configuration
@@ -146,6 +160,18 @@ public class RocksDBStateBackend extends AbstractStateBackend implements Configu
*/
private TernaryBoolean enableTtlCompactionFilter;
+ /**
+ * Total memory for all rocksDB instances at this slot.
+ * Currently, we would use a LRU cache to serve as the total memory usage.
+ */
+ private long totalMemoryPerSlot;
+
+ /** The write buffer ratio which consumed by write-buffer manager from the shared cache. */
+ private double writeBufferRatio;
+
+ /** The high priority pool ratio which consumed by index&filter from the shared cache. */
+ private double highPriPoolRatio;
+
/** This determines the type of priority queue state. */
private final PriorityQueueStateType priorityQueueStateType;
@@ -267,6 +293,9 @@ public class RocksDBStateBackend extends AbstractStateBackend implements Configu
this.priorityQueueStateType = PriorityQueueStateType.HEAP;
this.defaultMetricOptions = new RocksDBNativeMetricOptions();
this.enableTtlCompactionFilter = TernaryBoolean.UNDEFINED;
+ this.totalMemoryPerSlot = UNDEFINED_VALUE;
+ this.writeBufferRatio = UNDEFINED_VALUE;
+ this.highPriPoolRatio = UNDEFINED_VALUE;
}
/**
@@ -312,6 +341,33 @@ public class RocksDBStateBackend extends AbstractStateBackend implements Configu
this.enableTtlCompactionFilter = original.enableTtlCompactionFilter
.resolveUndefined(config.getBoolean(TTL_COMPACT_FILTER_ENABLED));
+ configureBoundedMemory(
+ () -> (double) original.getTotalMemoryPerSlot(),
+ () -> {
+ if (config.getString(BOUNDED_MEMORY_SIZE) != null) {
+ setTotalMemoryPerSlot(config.getString(BOUNDED_MEMORY_SIZE));
+ } else {
+ // we still left total memory size per slot as undefined if no actual settings.
+ this.totalMemoryPerSlot = UNDEFINED_VALUE;
+ }
+ },
+ () -> setTotalMemoryPerSlot(String.valueOf(original.totalMemoryPerSlot)));
+
+ configureBoundedMemory(
+ original::getWriteBufferRatio,
+ () -> setWriteBufferRatio(config.getDouble(WRITE_BUFFER_RATIO)),
+ () -> setWriteBufferRatio(original.writeBufferRatio));
+
+ configureBoundedMemory(
+ original::getHighPriPoolRatio,
+ () -> setHighPriPoolRatio(config.getDouble(HIGH_PRI_POOL_RATIO)),
+ () -> setHighPriPoolRatio(original.highPriPoolRatio));
+
+ if (isBoundedMemoryEnabled()) {
+ Preconditions.checkArgument(this.writeBufferRatio + this.highPriPoolRatio < 1.0,
+ String.format("Illegal sum of writeBufferRatio %s and highPriPoolRatio %s, should be less than 1.0", this.writeBufferRatio, this.highPriPoolRatio));
+ }
+
final String priorityQueueTypeString = config.getString(TIMER_SERVICE_FACTORY);
this.priorityQueueStateType = priorityQueueTypeString.length() > 0 ?
@@ -356,6 +412,14 @@ public class RocksDBStateBackend extends AbstractStateBackend implements Configu
}
}
+ private void configureBoundedMemory(Supplier<Double> getOriginalValue, Runnable parseFromConfig, Runnable setFromOriginal) {
+ if (getOriginalValue.get() == UNDEFINED_VALUE) {
+ parseFromConfig.run();
+ } else {
+ setFromOriginal.run();
+ }
+ }
+
// ------------------------------------------------------------------------
// Reconfiguration
// ------------------------------------------------------------------------
@@ -491,14 +555,53 @@ public class RocksDBStateBackend extends AbstractStateBackend implements Configu
LocalRecoveryConfig localRecoveryConfig =
env.getTaskStateManager().createLocalRecoveryConfig();
+ DBOptions dbOptions = getDbOptions();
+ Function<String, ColumnFamilyOptions> createColumnOptions;
+
+ if (isBoundedMemoryEnabled()) {
+ RocksDBSharedObjects rocksDBSharedObjects;
+ MemoryManager memoryManager = env.getMemoryManager();
+ // only initialized LRUCache and write buffer manager once.
+ // we would not dispose the cache and write buffer manager during disposing state backend
+ // as the same objects are shared by every RocksDB instance per slot.
+ while ((rocksDBSharedObjects = (RocksDBSharedObjects) memoryManager.getStateBackendSharedObject()) == null) {
+ if (memoryManager.getLazyInitializeSharedObjectHelper().compareAndSet(false, true)) {
+ Cache lruCache = new LRUCache(getTotalMemoryPerSlot(), -1, false, getHighPriPoolRatio());
+ WriteBufferManager writeBufferManager = new WriteBufferManager((long) (getTotalMemoryPerSlot() * getWriteBufferRatio()), lruCache);
+
+ rocksDBSharedObjects = new RocksDBSharedObjects(lruCache, writeBufferManager);
+ memoryManager.setStateBackendSharedObject(rocksDBSharedObjects);
+ }
+ }
+
+ Cache blockCache = checkNotNull(rocksDBSharedObjects.getCache());
+ dbOptions.setWriteBufferManager(checkNotNull(rocksDBSharedObjects.getWriteBufferManager()));
+
+ createColumnOptions = stateName -> {
+ ColumnFamilyOptions columnOptions = getColumnOptions();
+ TableFormatConfig tableFormatConfig = columnOptions.tableFormatConfig();
+ Preconditions.checkArgument(tableFormatConfig instanceof BlockBasedTableConfig,
+ "We currently only support BlockBasedTableConfig When bounding total memory.");
+ BlockBasedTableConfig blockBasedTableConfig = (BlockBasedTableConfig) tableFormatConfig;
+ blockBasedTableConfig.setBlockCache(blockCache);
+ blockBasedTableConfig.setCacheIndexAndFilterBlocks(true);
+ blockBasedTableConfig.setCacheIndexAndFilterBlocksWithHighPriority(true);
+ blockBasedTableConfig.setPinL0FilterAndIndexBlocksInCache(true);
+ columnOptions.setTableFormatConfig(blockBasedTableConfig);
+ return columnOptions;
+ };
+ } else {
+ createColumnOptions = stateName -> getColumnOptions();
+ }
+
ExecutionConfig executionConfig = env.getExecutionConfig();
StreamCompressionDecorator keyGroupCompressionDecorator = getCompressionDecorator(executionConfig);
RocksDBKeyedStateBackendBuilder<K> builder = new RocksDBKeyedStateBackendBuilder<>(
operatorIdentifier,
env.getUserClassLoader(),
instanceBasePath,
- getDbOptions(),
- stateName -> getColumnOptions(),
+ dbOptions,
+ createColumnOptions,
kvStateRegistry,
keySerializer,
numberOfKeyGroups,
@@ -705,6 +808,13 @@ public class RocksDBStateBackend extends AbstractStateBackend implements Configu
}
/**
+ * Gets whether bounding memory is enabled for this state backend.
+ */
+ public boolean isBoundedMemoryEnabled() {
+ return totalMemoryPerSlot > 0;
+ }
+
+ /**
* Enable compaction filter to cleanup state with TTL is enabled.
*
* <p>Note: User can still decide in state TTL configuration in state descriptor
@@ -853,6 +963,34 @@ public class RocksDBStateBackend extends AbstractStateBackend implements Configu
setNumberOfTransferThreads(numberOfTransferingThreads);
}
+ public void setTotalMemoryPerSlot(String totalMemoryPerSlotStr) {
+ long totalMemoryPerSlot = MemorySize.parseBytes(totalMemoryPerSlotStr);
+ Preconditions.checkArgument(totalMemoryPerSlot > 0, String.format("Illegal total memory per slot %s for RocksDBs.", totalMemoryPerSlot));
+ this.totalMemoryPerSlot = totalMemoryPerSlot;
+ }
+
+ public long getTotalMemoryPerSlot() {
+ return totalMemoryPerSlot;
+ }
+
+ public void setWriteBufferRatio(double writeBufferRatio) {
+ Preconditions.checkArgument(writeBufferRatio > 0 && writeBufferRatio < 1.0, String.format("Illegal write buffer ratio %s for RocksDBs.", writeBufferRatio));
+ this.writeBufferRatio = writeBufferRatio;
+ }
+
+ public double getWriteBufferRatio() {
+ return writeBufferRatio;
+ }
+
+ public void setHighPriPoolRatio(double highPriPoolRatio) {
+ Preconditions.checkArgument(highPriPoolRatio > 0 && highPriPoolRatio < 1.0, String.format("Illegal high priority pool ratio %s for RocksDBs.", highPriPoolRatio));
+ this.highPriPoolRatio = highPriPoolRatio;
+ }
+
+ public double getHighPriPoolRatio() {
+ return highPriPoolRatio;
+ }
+
// ------------------------------------------------------------------------
// utilities
// ------------------------------------------------------------------------
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendBoundedMemoryTest.java b/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendBoundedMemoryTest.java
new file mode 100644
index 0000000..c564eb4
--- /dev/null
+++ b/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendBoundedMemoryTest.java
@@ -0,0 +1,204 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.contrib.streaming.state;
+
+import org.apache.flink.api.common.typeutils.base.IntSerializer;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.core.fs.CloseableRegistry;
+import org.apache.flink.core.memory.MemoryType;
+import org.apache.flink.metrics.groups.UnregisteredMetricsGroup;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.memory.MemoryManager;
+import org.apache.flink.runtime.memory.MemoryManagerBuilder;
+import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
+import org.apache.flink.runtime.state.AbstractKeyedStateBackend;
+import org.apache.flink.runtime.state.KeyGroupRange;
+import org.apache.flink.runtime.state.ttl.TtlTimeProvider;
+
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.rocksdb.Cache;
+import org.rocksdb.WriteBufferManager;
+
+import java.io.IOException;
+import java.util.Collections;
+
+import static org.junit.Assert.fail;
+
+/**
+ * Tests to verify memory bounded for rocksDB state backend.
+ */
+public class RocksDBStateBackendBoundedMemoryTest {
+
+ private static final int MEMORY_SIZE = 64 * 1024 * 1024; // 64 MiBytes
+
+ private MemoryManager memoryManager;
+
+ @Rule
+ public final TemporaryFolder tempFolder = new TemporaryFolder();
+
+ @Before
+ public void setUp() {
+ this.memoryManager = MemoryManagerBuilder
+ .newBuilder()
+ .setMemorySize(MemoryType.OFF_HEAP, MEMORY_SIZE)
+ .build();
+ }
+
+ @After
+ public void shutDown() {
+ this.memoryManager.shutdown();
+ }
+
+ @Test
+ public void testSharedObjectsInitializeOnlyOnce() throws IOException {
+ DummyEnvironment env = new DummyEnvironment();
+ env.setMemoryManager(memoryManager);
+
+ Assert.assertNull(memoryManager.getStateBackendSharedObject());
+
+ Configuration configuration = new Configuration();
+ configuration.setString(RocksDBOptions.BOUNDED_MEMORY_SIZE, "128MB");
+ RocksDBStateBackend originalStateBackend = new RocksDBStateBackend(tempFolder.newFolder().toURI());
+ originalStateBackend.setDbStoragePath(tempFolder.newFolder().getAbsolutePath());
+
+ RocksDBStateBackend rocksDBStateBackend1 = originalStateBackend.configure(configuration, Thread.currentThread().getContextClassLoader());
+ AbstractKeyedStateBackend keyedStateBackend1 = createKeyedStateBackend(rocksDBStateBackend1, env);
+
+ Assert.assertTrue("The shared object should be a RocksDBSharedObject instance but actually not",
+ memoryManager.getStateBackendSharedObject() instanceof RocksDBSharedObjects);
+ RocksDBSharedObjects sharedObjects = (RocksDBSharedObjects) memoryManager.getStateBackendSharedObject();
+
+ Cache lruCache = sharedObjects.getCache();
+ WriteBufferManager writeBufferManager = sharedObjects.getWriteBufferManager();
+
+ RocksDBStateBackend rocksDBStateBackend2 = originalStateBackend.configure(configuration, Thread.currentThread().getContextClassLoader());
+ AbstractKeyedStateBackend keyedStateBackend2 = createKeyedStateBackend(rocksDBStateBackend2, env);
+
+ // Another keyed state backend is created but only initialized once for cache and write buffer manager.
+ sharedObjects = (RocksDBSharedObjects) memoryManager.getStateBackendSharedObject();
+
+ try {
+ Assert.assertEquals(lruCache, sharedObjects.getCache());
+ Assert.assertEquals(writeBufferManager, sharedObjects.getWriteBufferManager());
+ } finally {
+ keyedStateBackend1.close();
+ keyedStateBackend2.close();
+ }
+ }
+
+ @Test
+ public void testSharedObjectsNotClosedAfterKeyedStateBackendClosed() throws IOException {
+ DummyEnvironment env = new DummyEnvironment();
+ env.setMemoryManager(memoryManager);
+ Assert.assertNull(memoryManager.getStateBackendSharedObject());
+
+ Configuration configuration = new Configuration();
+ configuration.setString(RocksDBOptions.BOUNDED_MEMORY_SIZE, "128MB");
+ RocksDBStateBackend originalStateBackend = new RocksDBStateBackend(tempFolder.newFolder().toURI());
+ originalStateBackend.setDbStoragePath(tempFolder.newFolder().getAbsolutePath());
+ RocksDBStateBackend rocksDBStateBackend = originalStateBackend.configure(configuration, Thread.currentThread().getContextClassLoader());
+ AbstractKeyedStateBackend keyedStateBackend = createKeyedStateBackend(rocksDBStateBackend, env);
+
+ Assert.assertTrue("The shared object should be a RocksDBSharedObject instance but actually not",
+ memoryManager.getStateBackendSharedObject() instanceof RocksDBSharedObjects);
+
+ RocksDBSharedObjects stateBackendSharedObject = (RocksDBSharedObjects) memoryManager.getStateBackendSharedObject();
+ Assert.assertTrue(stateBackendSharedObject.getCache().isOwningHandle());
+ Assert.assertTrue(stateBackendSharedObject.getWriteBufferManager().isOwningHandle());
+
+ keyedStateBackend.close();
+ // even keyed state backend closed, cache and write buffer manager would not be disposed.
+ Assert.assertTrue(stateBackendSharedObject.getCache().isOwningHandle());
+ Assert.assertTrue(stateBackendSharedObject.getWriteBufferManager().isOwningHandle());
+ }
+
+ @Test
+ public void testSharedObjectsInitializedConcurrently() throws IOException {
+ DummyEnvironment env = new DummyEnvironment();
+ env.setMemoryManager(memoryManager);
+
+ Assert.assertNull(memoryManager.getStateBackendSharedObject());
+
+ Configuration configuration = new Configuration();
+ configuration.setString(RocksDBOptions.BOUNDED_MEMORY_SIZE, "128MB");
+ RocksDBStateBackend originalStateBackend = new RocksDBStateBackend(tempFolder.newFolder().toURI());
+ originalStateBackend.setDbStoragePath(tempFolder.newFolder().getAbsolutePath());
+
+ int numThreads = 4;
+ Thread[] threads = new Thread[numThreads];
+ RocksDBKeyedStateBackend[] keyedStateBackends = new RocksDBKeyedStateBackend[numThreads];
+ for (int i = 0; i < numThreads; i++) {
+ int index = i;
+ threads[index] = new Thread(() -> {
+ RocksDBStateBackend rocksDBStateBackend = originalStateBackend.configure(configuration, Thread.currentThread().getContextClassLoader());
+ try {
+ keyedStateBackends[index] = (RocksDBKeyedStateBackend) createKeyedStateBackend(rocksDBStateBackend, env);
+ } catch (IOException ignored) {
+ }
+ });
+ threads[i].setDaemon(true);
+ }
+
+ // launch concurrently
+ for (int i = 0; i < numThreads; i++) {
+ threads[i].start();
+ }
+
+ try {
+ // sync
+ for (int i = 0; i < numThreads; i++) {
+ threads[i].join(5000);
+ }
+ } catch (InterruptedException e) {
+ fail(e.getMessage());
+ }
+
+ WriteBufferManager writeBufferManager = null;
+ for (RocksDBKeyedStateBackend keyedStateBackend : keyedStateBackends) {
+ if (writeBufferManager == null) {
+ writeBufferManager = keyedStateBackend.getDbOptions().writeBufferManager();
+ } else {
+ // different keyed state backends but share the same write buffer manager
+ Assert.assertEquals(writeBufferManager, keyedStateBackend.getDbOptions().writeBufferManager());
+ }
+ keyedStateBackend.close();
+ }
+ }
+
+ private AbstractKeyedStateBackend createKeyedStateBackend(RocksDBStateBackend rocksDBStateBackend, Environment env) throws IOException {
+ return rocksDBStateBackend.createKeyedStateBackend(
+ env,
+ env.getJobID(),
+ "test_op",
+ IntSerializer.INSTANCE,
+ 1,
+ new KeyGroupRange(0, 0),
+ env.getTaskKvStateRegistry(),
+ TtlTimeProvider.DEFAULT,
+ new UnregisteredMetricsGroup(),
+ Collections.emptyList(),
+ new CloseableRegistry());
+
+ }
+}
diff --git a/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java b/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java
index c00ffdb..9cd338f 100644
--- a/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java
+++ b/flink-state-backends/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackendConfigTest.java
@@ -632,6 +632,56 @@ public class RocksDBStateBackendConfigTest {
}
// ------------------------------------------------------------------------
+ // RocksDB Memory Control
+ // ------------------------------------------------------------------------
+
+ @Test
+ public void testDefaultMemoryControlParameters() {
+ StateBackend storageBackend = new MemoryStateBackend();
+ RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(storageBackend);
+ assertEquals(RocksDBStateBackend.UNDEFINED_VALUE, rocksDbBackend.getTotalMemoryPerSlot());
+ assertEquals(RocksDBStateBackend.UNDEFINED_VALUE, rocksDbBackend.getHighPriPoolRatio(), 0.0);
+ assertEquals(RocksDBStateBackend.UNDEFINED_VALUE, rocksDbBackend.getWriteBufferRatio(), 0.0);
+
+ RocksDBStateBackend configure = rocksDbBackend.configure(new Configuration(), Thread.currentThread().getContextClassLoader());
+ assertEquals(RocksDBStateBackend.UNDEFINED_VALUE, configure.getTotalMemoryPerSlot());
+ assertEquals(RocksDBOptions.HIGH_PRI_POOL_RATIO.defaultValue(), configure.getHighPriPoolRatio(), 0.0);
+ assertEquals(RocksDBOptions.WRITE_BUFFER_RATIO.defaultValue(), configure.getWriteBufferRatio(), 0.0);
+ }
+
+ @Test
+ public void testConfigureIllegalMemoryControlParameters() {
+ StateBackend storageBackend = new MemoryStateBackend();
+ RocksDBStateBackend rocksDbBackend = new RocksDBStateBackend(storageBackend);
+
+ verifySetParameter(() -> rocksDbBackend.setTotalMemoryPerSlot("-1B"));
+ verifySetParameter(() -> rocksDbBackend.setHighPriPoolRatio(-0.1));
+ verifySetParameter(() -> rocksDbBackend.setHighPriPoolRatio(1.1));
+ verifySetParameter(() -> rocksDbBackend.setWriteBufferRatio(-0.1));
+ verifySetParameter(() -> rocksDbBackend.setWriteBufferRatio(1.1));
+
+ rocksDbBackend.setTotalMemoryPerSlot("128MB");
+ rocksDbBackend.setWriteBufferRatio(0.6);
+ rocksDbBackend.setHighPriPoolRatio(0.6);
+ try {
+ // sum of writeBufferRatio and highPriPoolRatio larger than 1.0
+ rocksDbBackend.configure(new Configuration(), Thread.currentThread().getContextClassLoader());
+ fail("No expected IllegalArgumentException.");
+ } catch (IllegalArgumentException expected) {
+ // expected exception
+ }
+ }
+
+ private void verifySetParameter(Runnable setter) {
+ try {
+ setter.run();
+ fail("No expected IllegalArgumentException.");
+ } catch (IllegalArgumentException expected) {
+ // expected exception
+ }
+ }
+
+ // ------------------------------------------------------------------------
// Contained Non-partitioned State Backend
// ------------------------------------------------------------------------