[FLINK-19180][state backends] Make rocksdb state backend respect mana…

…ged memory fraction.

This closes apache#13500.

flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/BootstrapTransformation.java

@@ -28,6 +28,7 @@
 import org.apache.flink.api.java.operators.MapPartitionOperator;
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.core.fs.Path;
+import org.apache.flink.core.memory.ManagedMemoryUseCase;
 import org.apache.flink.runtime.checkpoint.OperatorState;
 import org.apache.flink.runtime.jobgraph.OperatorID;
 import org.apache.flink.runtime.state.StateBackend;
@@ -194,6 +195,7 @@ StreamConfig getConfig(OperatorID operatorID, StateBackend stateBackend, StreamO
  config.setOperatorName(operatorID.toHexString());
  config.setOperatorID(operatorID);
  config.setStateBackend(stateBackend);
+ config.setManagedMemoryFractionOperatorOfUseCase(ManagedMemoryUseCase.STATE_BACKEND, 1.0);
  return config;
  }
 

flink-libraries/flink-state-processing-api/src/main/java/org/apache/flink/state/api/input/KeyedStateInputFormat.java

@@ -183,7 +183,8 @@ private StreamOperatorStateContext getStreamOperatorStateContext(Environment env
  operator,
  operator.getKeyType().createSerializer(environment.getExecutionConfig()),
  registry,
- getRuntimeContext().getMetricGroup());
+ getRuntimeContext().getMetricGroup(),
+ 1.0);
  } catch (Exception e) {
  throw new IOException("Failed to restore state backend", e);
  }

flink-runtime/src/main/java/org/apache/flink/runtime/memory/MemoryManager.java

@@ -470,17 +470,6 @@ public void releaseAllMemory(Object owner) {
  // Shared opaque memory resources
  // ------------------------------------------------------------------------
 
- /**
- * Acquires a shared memory resource, that uses all the memory of this memory manager.
- * This method behaves otherwise exactly as {@link #getSharedMemoryResourceForManagedMemory(String, LongFunctionWithException, double)}.
- */
- public <T extends AutoCloseable> OpaqueMemoryResource<T> getSharedMemoryResourceForManagedMemory(
- String type,
- LongFunctionWithException<T, Exception> initializer) throws Exception {
-
- return getSharedMemoryResourceForManagedMemory(type, initializer, 1.0);
- }
-
  /**
  * Acquires a shared memory resource, identified by a type string. If the resource already exists, this
  * returns a descriptor to the resource. If the resource does not yet exist, the given memory fraction

flink-runtime/src/main/java/org/apache/flink/runtime/state/AbstractManagedMemoryStateBackend.java

@@ -0,0 +1,58 @@
+/*
+ * 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
+ *
+ * https://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.runtime.state;
+
+import org.apache.flink.api.common.JobID;
+import org.apache.flink.api.common.typeutils.TypeSerializer;
+import org.apache.flink.core.fs.CloseableRegistry;
+import org.apache.flink.metrics.MetricGroup;
+import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.query.TaskKvStateRegistry;
+import org.apache.flink.runtime.state.ttl.TtlTimeProvider;
+
+import javax.annotation.Nonnull;
+import java.util.Collection;
+
+/**
+ * Abstract base class for state backends that use managed memory.
+ */
+public abstract class AbstractManagedMemoryStateBackend extends AbstractStateBackend {
+
+ private static final long serialVersionUID = 1L;
+
+ @Override
+ public abstract <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
+ Environment env,
+ JobID jobID,
+ String operatorIdentifier,
+ TypeSerializer<K> keySerializer,
+ int numberOfKeyGroups,
+ KeyGroupRange keyGroupRange,
+ TaskKvStateRegistry kvStateRegistry,
+ TtlTimeProvider ttlTimeProvider,
+ MetricGroup metricGroup,
+ @Nonnull Collection<KeyedStateHandle> stateHandles,
+ CloseableRegistry cancelStreamRegistry,
+ double managedMemoryFraction) throws Exception;
+
+ @Override
+ public boolean useManagedMemory() {
+ return true;
+ }
+}

flink-runtime/src/main/java/org/apache/flink/runtime/state/StateBackend.java

@@ -155,6 +155,39 @@ <K> CheckpointableKeyedStateBackend<K> createKeyedStateBackend(
  MetricGroup metricGroup,
  @Nonnull Collection<KeyedStateHandle> stateHandles,
  CloseableRegistry cancelStreamRegistry) throws Exception;
+
+ /**
+ * Creates a new {@link CheckpointableKeyedStateBackend} with the given managed memory fraction.
+ * Backends that use managed memory are required to implement this interface.
+ */
+ default <K> CheckpointableKeyedStateBackend<K> createKeyedStateBackend(
+ Environment env,
+ JobID jobID,
+ String operatorIdentifier,
+ TypeSerializer<K> keySerializer,
+ int numberOfKeyGroups,
+ KeyGroupRange keyGroupRange,
+ TaskKvStateRegistry kvStateRegistry,
+ TtlTimeProvider ttlTimeProvider,
+ MetricGroup metricGroup,
+ @Nonnull Collection<KeyedStateHandle> stateHandles,
+ CloseableRegistry cancelStreamRegistry,
+ double managedMemoryFraction) throws Exception {
+
+ // ignore managed memory fraction by default
+ return createKeyedStateBackend(
+ env,
+ jobID,
+ operatorIdentifier,
+ keySerializer,
+ numberOfKeyGroups,
+ keyGroupRange,
+ kvStateRegistry,
+ ttlTimeProvider,
+ metricGroup,
+ stateHandles,
+ cancelStreamRegistry);
+ }
 
  /**
  * Creates a new {@link OperatorStateBackend} that can be used for storing operator state.

flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBOperationUtils.java

@@ -190,6 +190,7 @@ public static void addColumnFamilyOptionsToCloseLater(
  public static OpaqueMemoryResource<RocksDBSharedResources> allocateSharedCachesIfConfigured(
  RocksDBMemoryConfiguration memoryConfig,
  MemoryManager memoryManager,
+ double memoryFraction,
  Logger logger) throws IOException {
 
  if (!memoryConfig.isUsingFixedMemoryPerSlot() && !memoryConfig.isUsingManagedMemory()) {
@@ -212,7 +213,7 @@ public static OpaqueMemoryResource<RocksDBSharedResources> allocateSharedCachesI
  }
  else {
  logger.info("Getting managed memory shared cache for RocksDB.");
- return memoryManager.getSharedMemoryResourceForManagedMemory(MANAGED_MEMORY_RESOURCE_ID, allocator);
+ return memoryManager.getSharedMemoryResourceForManagedMemory(MANAGED_MEMORY_RESOURCE_ID, allocator, memoryFraction);
  }
  }
  catch (Exception e) {

flink-state-backends/flink-statebackend-rocksdb/src/main/java/org/apache/flink/contrib/streaming/state/RocksDBStateBackend.java

@@ -31,6 +31,7 @@
 import org.apache.flink.runtime.memory.OpaqueMemoryResource;
 import org.apache.flink.runtime.query.TaskKvStateRegistry;
 import org.apache.flink.runtime.state.AbstractKeyedStateBackend;
+import org.apache.flink.runtime.state.AbstractManagedMemoryStateBackend;
 import org.apache.flink.runtime.state.AbstractStateBackend;
 import org.apache.flink.runtime.state.CheckpointStorage;
 import org.apache.flink.runtime.state.CompletedCheckpointStorageLocation;
@@ -90,7 +91,7 @@
  * using the methods {@link #setPredefinedOptions(PredefinedOptions)} and
  * {@link #setRocksDBOptions(RocksDBOptionsFactory)}.
  */
-public class RocksDBStateBackend extends AbstractStateBackend implements ConfigurableStateBackend {
+public class RocksDBStateBackend extends AbstractManagedMemoryStateBackend implements ConfigurableStateBackend {
 
  /**
  * The options to chose for the type of priority queue state.
@@ -451,11 +452,6 @@ private File getNextStoragePath() {
  return initializedDbBasePaths[ni];
  }
 
- @Override
- public boolean useManagedMemory() {
- return true;
- }
-
  // ------------------------------------------------------------------------
  // Checkpoint initialization and persistent storage
  // ------------------------------------------------------------------------
@@ -474,6 +470,35 @@ public CheckpointStorage createCheckpointStorage(JobID jobId) throws IOException
  // State holding data structures
  // ------------------------------------------------------------------------
 
+ @Override
+ public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
+ Environment env,
+ JobID jobID,
+ String operatorIdentifier,
+ TypeSerializer<K> keySerializer,
+ int numberOfKeyGroups,
+ KeyGroupRange keyGroupRange,
+ TaskKvStateRegistry kvStateRegistry,
+ TtlTimeProvider ttlTimeProvider,
+ MetricGroup metricGroup,
+ @Nonnull Collection<KeyedStateHandle> stateHandles,
+ CloseableRegistry cancelStreamRegistry) throws IOException {
+ return createKeyedStateBackend(
+ env,
+ jobID,
+ operatorIdentifier,
+ keySerializer,
+ numberOfKeyGroups,
+ keyGroupRange,
+ kvStateRegistry,
+ ttlTimeProvider,
+ metricGroup,
+ stateHandles,
+ cancelStreamRegistry,
+ 1.0
+ );
+ }
+
  @Override
  public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
  Environment env,
@@ -486,7 +511,8 @@ public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
  TtlTimeProvider ttlTimeProvider,
  MetricGroup metricGroup,
  @Nonnull Collection<KeyedStateHandle> stateHandles,
- CloseableRegistry cancelStreamRegistry) throws IOException {
+ CloseableRegistry cancelStreamRegistry,
+ double managedMemoryFraction) throws IOException {
 
  // first, make sure that the RocksDB JNI library is loaded
  // we do this explicitly here to have better error handling
@@ -506,7 +532,7 @@ public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
  env.getTaskStateManager().createLocalRecoveryConfig();
 
  final OpaqueMemoryResource<RocksDBSharedResources> sharedResources = RocksDBOperationUtils
- .allocateSharedCachesIfConfigured(memoryConfiguration, env.getMemoryManager(), LOG);
+ .allocateSharedCachesIfConfigured(memoryConfiguration, env.getMemoryManager(), managedMemoryFraction, LOG);
  if (sharedResources != null) {
  LOG.info("Obtained shared RocksDB cache of size {} bytes", sharedResources.getSize());
  }

flink-streaming-java/src/main/java/org/apache/flink/streaming/api/graph/StreamConfig.java

@@ -523,6 +523,11 @@ public void setStateBackend(StateBackend backend) {
  }
  }
 
+ @VisibleForTesting
+ public void setStateBackendUsesManagedMemory(boolean usesManagedMemory) {
+ this.config.setBoolean(STATE_BACKEND_USE_MANAGED_MEMORY, usesManagedMemory);
+ }
+
  public StateBackend getStateBackend(ClassLoader cl) {
  try {
  return InstantiationUtil.readObjectFromConfig(this.config, STATE_BACKEND, cl);

flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperator.java

@@ -29,6 +29,7 @@
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.configuration.MetricOptions;
 import org.apache.flink.core.fs.CloseableRegistry;
+import org.apache.flink.core.memory.ManagedMemoryUseCase;
 import org.apache.flink.metrics.MetricGroup;
 import org.apache.flink.runtime.checkpoint.CheckpointOptions;
 import org.apache.flink.runtime.execution.Environment;
@@ -250,7 +251,11 @@ public final void initializeState(StreamTaskStateInitializer streamTaskStateMana
  this,
  keySerializer,
  streamTaskCloseableRegistry,
- metrics);
+ metrics,
+ config.getManagedMemoryFractionOperatorUseCaseOfSlot(
+ ManagedMemoryUseCase.STATE_BACKEND,
+ runtimeContext.getTaskManagerRuntimeInfo().getConfiguration(),
+ runtimeContext.getUserCodeClassLoader()));
 
  stateHandler = new StreamOperatorStateHandler(context, getExecutionConfig(), streamTaskCloseableRegistry);
  timeServiceManager = context.internalTimerServiceManager();

flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/AbstractStreamOperatorV2.java

@@ -29,6 +29,7 @@
 import org.apache.flink.configuration.Configuration;
 import org.apache.flink.configuration.MetricOptions;
 import org.apache.flink.core.fs.CloseableRegistry;
+import org.apache.flink.core.memory.ManagedMemoryUseCase;
 import org.apache.flink.metrics.MetricGroup;
 import org.apache.flink.runtime.checkpoint.CheckpointOptions;
 import org.apache.flink.runtime.execution.Environment;
@@ -200,7 +201,11 @@ public final void initializeState(StreamTaskStateInitializer streamTaskStateMana
  this,
  keySerializer,
  cancelables,
- metrics);
+ metrics,
+ config.getManagedMemoryFractionOperatorUseCaseOfSlot(
+ ManagedMemoryUseCase.STATE_BACKEND,
+ runtimeContext.getTaskManagerRuntimeInfo().getConfiguration(),
+ runtimeContext.getUserCodeClassLoader()));
 
  stateHandler = new StreamOperatorStateHandler(context, getExecutionConfig(), cancelables);
  timeServiceManager = context.internalTimerServiceManager();

flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamTaskStateInitializer.java

@@ -45,6 +45,7 @@ public interface StreamTaskStateInitializer {
  * @param keySerializer the key-serializer for the operator. Can be null.
  * @param streamTaskCloseableRegistry the closeable registry to which created closeable objects will be registered.
  * @param metricGroup the parent metric group for all statebackend metrics
+ * @param managedMemoryFraction the managed memory fraction of the operator for state backend
  * @return a context from which the given operator can initialize everything related to state.
  * @throws Exception when something went wrong while creating the context.
  */
@@ -55,5 +56,6 @@ StreamOperatorStateContext streamOperatorStateContext(
  @Nonnull KeyContext keyContext,
  @Nullable TypeSerializer<?> keySerializer,
  @Nonnull CloseableRegistry streamTaskCloseableRegistry,
- @Nonnull MetricGroup metricGroup) throws Exception;
+ @Nonnull MetricGroup metricGroup,
+ double managedMemoryFraction) throws Exception;
 }

flink-streaming-java/src/main/java/org/apache/flink/streaming/api/operators/StreamTaskStateInitializerImpl.java

@@ -123,7 +123,8 @@ public StreamOperatorStateContext streamOperatorStateContext(
  @Nonnull KeyContext keyContext,
  @Nullable TypeSerializer<?> keySerializer,
  @Nonnull CloseableRegistry streamTaskCloseableRegistry,
- @Nonnull MetricGroup metricGroup) throws Exception {
+ @Nonnull MetricGroup metricGroup,
+ double managedMemoryFraction) throws Exception {
 
  TaskInfo taskInfo = environment.getTaskInfo();
  OperatorSubtaskDescriptionText operatorSubtaskDescription =
@@ -152,7 +153,8 @@ public StreamOperatorStateContext streamOperatorStateContext(
  operatorIdentifierText,
  prioritizedOperatorSubtaskStates,
  streamTaskCloseableRegistry,
- metricGroup);
+ metricGroup,
+ managedMemoryFraction);
 
  // -------------- Operator State Backend --------------
  operatorStateBackend = operatorStateBackend(
@@ -253,7 +255,8 @@ protected <K> CheckpointableKeyedStateBackend<K> keyedStatedBackend(
  String operatorIdentifierText,
  PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates,
  CloseableRegistry backendCloseableRegistry,
- MetricGroup metricGroup) throws Exception {
+ MetricGroup metricGroup,
+ double managedMemoryFraction) throws Exception {
 
  if (keySerializer == null) {
  return null;
@@ -286,7 +289,8 @@ protected <K> CheckpointableKeyedStateBackend<K> keyedStatedBackend(
  ttlTimeProvider,
  metricGroup,
  stateHandles,
- cancelStreamRegistryForRestore),
+ cancelStreamRegistryForRestore,
+ managedMemoryFraction),
  backendCloseableRegistry,
  logDescription);
 

flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StateInitializationContextImplTest.java

@@ -201,7 +201,8 @@ public <K> InternalTimeServiceManager<K> create(
  // consumed by the timer service.
  IntSerializer.INSTANCE,
  closableRegistry,
- new UnregisteredMetricsGroup());
+ new UnregisteredMetricsGroup(),
+ 1.0);
 
  this.initializationContext =
  new StateInitializationContextImpl(

flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamOperatorStateHandlerTest.java

@@ -104,7 +104,8 @@ public void testFailingBackendSnapshotMethod() throws Exception {
  new UnUsedKeyContext(),
  IntSerializer.INSTANCE,
  closeableRegistry,
- new InterceptingOperatorMetricGroup());
+ new InterceptingOperatorMetricGroup(),
+ 1.0);
  StreamOperatorStateHandler stateHandler = new StreamOperatorStateHandler(stateContext, new ExecutionConfig(), closeableRegistry);
 
  final String keyedStateField = "keyedStateField";

flink-streaming-java/src/test/java/org/apache/flink/streaming/api/operators/StreamTaskStateInitializerImplTest.java

@@ -101,7 +101,8 @@ public void testNoRestore() throws Exception {
  streamOperator,
  typeSerializer,
  closeableRegistry,
- new UnregisteredMetricsGroup());
+ new UnregisteredMetricsGroup(),
+ 1.0);
 
  OperatorStateBackend operatorStateBackend = stateContext.operatorStateBackend();
  CheckpointableKeyedStateBackend<?> keyedStateBackend = stateContext.keyedStateBackend();
@@ -212,7 +213,8 @@ public OperatorStateBackend createOperatorStateBackend(
  streamOperator,
  typeSerializer,
  closeableRegistry,
- new UnregisteredMetricsGroup());
+ new UnregisteredMetricsGroup(),
+ 1.0);
 
  OperatorStateBackend operatorStateBackend = stateContext.operatorStateBackend();
  CheckpointableKeyedStateBackend<?> keyedStateBackend = stateContext.keyedStateBackend();