[FLINK-34558][state] Support tracking state size

Author: hejufang

flink-annotations/src/main/java/org/apache/flink/annotation/docs/Documentation.java

@@ -79,6 +79,7 @@ public static final class Sections {
         public static final String STATE_BACKEND_ROCKSDB = "state_backend_rocksdb";
 
         public static final String STATE_LATENCY_TRACKING = "state_latency_tracking";
+        public static final String STATE_SIZE_TRACKING = "state_size_tracking";
 
         public static final String STATE_CHANGELOG = "state_changelog";
 

flink-core/src/main/java/org/apache/flink/configuration/StateSizeTrackOptions.java

@@ -0,0 +1,62 @@
+/*
+ * 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.configuration;
+
+import org.apache.flink.annotation.PublicEvolving;
+import org.apache.flink.annotation.docs.Documentation;
+
+/** A collection of all configuration options that relate to the size tracking for state access. */
+@PublicEvolving
+public class StateSizeTrackOptions {
+
+    @Documentation.Section(Documentation.Sections.STATE_SIZE_TRACKING)
+    public static final ConfigOption<Boolean> SIZE_TRACK_ENABLED =
+            ConfigOptions.key("state.backend.size-track.keyed-state-enabled")
+                    .booleanType()
+                    .defaultValue(false)
+                    .withDescription(
+                            "Whether to track size of keyed state operations, e.g value state put/get/clear.");
+
+    @Documentation.Section(Documentation.Sections.STATE_SIZE_TRACKING)
+    public static final ConfigOption<Integer> SIZE_TRACK_SAMPLE_INTERVAL =
+            ConfigOptions.key("state.backend.size-track.sample-interval")
+                    .intType()
+                    .defaultValue(100)
+                    .withDescription(
+                            String.format(
+                                    "The sample interval of size track once '%s' is enabled. "
+                                            + "The default value is 100, which means we would track the size every 100 access requests.",
+                                    SIZE_TRACK_ENABLED.key()));
+
+    @Documentation.Section(Documentation.Sections.STATE_SIZE_TRACKING)
+    public static final ConfigOption<Integer> SIZE_TRACK_HISTORY_SIZE =
+            ConfigOptions.key("state.backend.size-track.history-size")
+                    .intType()
+                    .defaultValue(128)
+                    .withDescription(
+                            "Defines the number of measured size to maintain at each state access operation.");
+
+    @Documentation.Section(Documentation.Sections.STATE_SIZE_TRACKING)
+    public static final ConfigOption<Boolean> SIZE_TRACK_STATE_NAME_AS_VARIABLE =
+            ConfigOptions.key("state.backend.size-track.state-name-as-variable")
+                    .booleanType()
+                    .defaultValue(true)
+                    .withDescription(
+                            "Whether to expose state name as a variable if tracking size.");
+}

flink-libraries/flink-state-processing-api/src/test/java/org/apache/flink/state/api/input/MultiStateKeyIteratorTest.java

@@ -52,6 +52,7 @@
 import org.apache.flink.runtime.state.VoidNamespaceSerializer;
 import org.apache.flink.runtime.state.heap.HeapPriorityQueueElement;
 import org.apache.flink.runtime.state.metrics.LatencyTrackingStateConfig;
+import org.apache.flink.runtime.state.metrics.SizeTrackingStateConfig;
 import org.apache.flink.runtime.state.ttl.TtlTimeProvider;
 import org.apache.flink.runtime.state.ttl.mock.MockRestoreOperation;
 import org.apache.flink.runtime.state.ttl.mock.MockStateBackend;
@@ -132,6 +133,7 @@ private static CountingKeysKeyedStateBackend createCountingKeysKeyedStateBackend
                 env.getExecutionConfig(),
                 ttlTimeProvider,
                 LatencyTrackingStateConfig.disabled(),
+                SizeTrackingStateConfig.disabled(),
                 cancelStreamRegistry,
                 new InternalKeyContextImpl<>(keyGroupRange, numberOfKeyGroups));
     }
@@ -269,6 +271,7 @@ public CountingKeysKeyedStateBackend(
                 ExecutionConfig executionConfig,
                 TtlTimeProvider ttlTimeProvider,
                 LatencyTrackingStateConfig latencyTrackingStateConfig,
+                SizeTrackingStateConfig sizeTrackingStateConfig,
                 CloseableRegistry cancelStreamRegistry,
                 InternalKeyContext<Integer> keyContext) {
             super(
@@ -278,6 +281,7 @@ public CountingKeysKeyedStateBackend(
                     executionConfig,
                     ttlTimeProvider,
                     latencyTrackingStateConfig,
+                    sizeTrackingStateConfig,
                     cancelStreamRegistry,
                     keyContext);
             this.numberOfKeysGenerated = numberOfKeysGenerated;

flink-queryable-state/flink-queryable-state-runtime/src/test/java/org/apache/flink/queryablestate/network/KvStateRequestSerializerTest.java

@@ -44,6 +44,7 @@
 import org.apache.flink.runtime.state.internal.InternalListState;
 import org.apache.flink.runtime.state.internal.InternalMapState;
 import org.apache.flink.runtime.state.metrics.LatencyTrackingStateConfig;
+import org.apache.flink.runtime.state.metrics.SizeTrackingStateConfig;
 import org.apache.flink.runtime.state.ttl.TtlTimeProvider;
 
 import org.junit.jupiter.api.Test;
@@ -336,6 +337,7 @@ private HeapKeyedStateBackend<Long> getLongHeapKeyedStateBackend(final long key,
                                 executionConfig,
                                 TtlTimeProvider.DEFAULT,
                                 LatencyTrackingStateConfig.disabled(),
+                                SizeTrackingStateConfig.disabled(),
                                 Collections.emptyList(),
                                 AbstractStateBackend.getCompressionDecorator(executionConfig),
                                 TestLocalRecoveryConfig.disabled(),

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

@@ -31,6 +31,7 @@
 import org.apache.flink.runtime.state.internal.InternalKvState;
 import org.apache.flink.runtime.state.metrics.LatencyTrackingStateConfig;
 import org.apache.flink.runtime.state.metrics.LatencyTrackingStateFactory;
+import org.apache.flink.runtime.state.metrics.SizeTrackingStateConfig;
 import org.apache.flink.runtime.state.ttl.TtlStateFactory;
 import org.apache.flink.runtime.state.ttl.TtlTimeProvider;
 import org.apache.flink.util.IOUtils;
@@ -93,6 +94,8 @@ public abstract class AbstractKeyedStateBackend<K>
 
     protected final LatencyTrackingStateConfig latencyTrackingStateConfig;
 
+    protected final SizeTrackingStateConfig sizeTrackingStateConfig;
+
     /** Decorates the input and output streams to write key-groups compressed. */
     protected final StreamCompressionDecorator keyGroupCompressionDecorator;
 
@@ -106,6 +109,7 @@ public AbstractKeyedStateBackend(
             ExecutionConfig executionConfig,
             TtlTimeProvider ttlTimeProvider,
             LatencyTrackingStateConfig latencyTrackingStateConfig,
+            SizeTrackingStateConfig sizeTrackingStateConfig,
             CloseableRegistry cancelStreamRegistry,
             InternalKeyContext<K> keyContext) {
         this(
@@ -115,6 +119,7 @@ public AbstractKeyedStateBackend(
                 executionConfig,
                 ttlTimeProvider,
                 latencyTrackingStateConfig,
+                sizeTrackingStateConfig,
                 cancelStreamRegistry,
                 determineStreamCompression(executionConfig),
                 keyContext);
@@ -127,6 +132,7 @@ public AbstractKeyedStateBackend(
             ExecutionConfig executionConfig,
             TtlTimeProvider ttlTimeProvider,
             LatencyTrackingStateConfig latencyTrackingStateConfig,
+            SizeTrackingStateConfig sizeTrackingStateConfig,
             CloseableRegistry cancelStreamRegistry,
             StreamCompressionDecorator keyGroupCompressionDecorator,
             InternalKeyContext<K> keyContext) {
@@ -137,6 +143,7 @@ public AbstractKeyedStateBackend(
                 executionConfig,
                 ttlTimeProvider,
                 latencyTrackingStateConfig,
+                sizeTrackingStateConfig,
                 cancelStreamRegistry,
                 keyGroupCompressionDecorator,
                 Preconditions.checkNotNull(keyContext),
@@ -157,6 +164,7 @@ protected AbstractKeyedStateBackend(AbstractKeyedStateBackend<K> abstractKeyedSt
                 abstractKeyedStateBackend.executionConfig,
                 abstractKeyedStateBackend.ttlTimeProvider,
                 abstractKeyedStateBackend.latencyTrackingStateConfig,
+                abstractKeyedStateBackend.sizeTrackingStateConfig,
                 abstractKeyedStateBackend.cancelStreamRegistry,
                 abstractKeyedStateBackend.keyGroupCompressionDecorator,
                 abstractKeyedStateBackend.keyContext,
@@ -176,6 +184,7 @@ private AbstractKeyedStateBackend(
             ExecutionConfig executionConfig,
             TtlTimeProvider ttlTimeProvider,
             LatencyTrackingStateConfig latencyTrackingStateConfig,
+            SizeTrackingStateConfig sizeTrackingStateConfig,
             CloseableRegistry cancelStreamRegistry,
             StreamCompressionDecorator keyGroupCompressionDecorator,
             InternalKeyContext<K> keyContext,
@@ -206,6 +215,7 @@ private AbstractKeyedStateBackend(
         this.keyGroupCompressionDecorator = keyGroupCompressionDecorator;
         this.ttlTimeProvider = Preconditions.checkNotNull(ttlTimeProvider);
         this.latencyTrackingStateConfig = Preconditions.checkNotNull(latencyTrackingStateConfig);
+        this.sizeTrackingStateConfig = Preconditions.checkNotNull(sizeTrackingStateConfig);
         this.keySelectionListeners = keySelectionListeners;
         this.lastState = lastState;
         this.lastName = lastName;
@@ -383,8 +393,10 @@ public <N, S extends State, V> S getOrCreateKeyedState(
                     LatencyTrackingStateFactory.createStateAndWrapWithLatencyTrackingIfEnabled(
                             TtlStateFactory.createStateAndWrapWithTtlIfEnabled(
                                     namespaceSerializer, stateDescriptor, this, ttlTimeProvider),
+                            this,
                             stateDescriptor,
-                            latencyTrackingStateConfig);
+                            latencyTrackingStateConfig,
+                            sizeTrackingStateConfig);
             keyValueStatesByName.put(stateDescriptor.getName(), kvState);
             publishQueryableStateIfEnabled(stateDescriptor, kvState);
         }
@@ -451,6 +463,10 @@ public LatencyTrackingStateConfig getLatencyTrackingStateConfig() {
         return latencyTrackingStateConfig;
     }
 
+    public SizeTrackingStateConfig getSizeTrackingStateConfig() {
+        return sizeTrackingStateConfig;
+    }
+
     @VisibleForTesting
     public StreamCompressionDecorator getKeyGroupCompressionDecorator() {
         return keyGroupCompressionDecorator;

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

@@ -23,6 +23,7 @@
 import org.apache.flink.core.fs.CloseableRegistry;
 import org.apache.flink.runtime.query.TaskKvStateRegistry;
 import org.apache.flink.runtime.state.metrics.LatencyTrackingStateConfig;
+import org.apache.flink.runtime.state.metrics.SizeTrackingStateConfig;
 import org.apache.flink.runtime.state.ttl.TtlTimeProvider;
 
 import org.slf4j.Logger;
@@ -45,6 +46,7 @@ public abstract class AbstractKeyedStateBackendBuilder<K>
     protected final ExecutionConfig executionConfig;
     protected final TtlTimeProvider ttlTimeProvider;
     protected final LatencyTrackingStateConfig latencyTrackingStateConfig;
+    protected final SizeTrackingStateConfig sizeTrackingStateConfig;
     protected final StreamCompressionDecorator keyGroupCompressionDecorator;
     protected final Collection<KeyedStateHandle> restoreStateHandles;
     protected final CloseableRegistry cancelStreamRegistry;
@@ -58,6 +60,7 @@ public AbstractKeyedStateBackendBuilder(
             ExecutionConfig executionConfig,
             TtlTimeProvider ttlTimeProvider,
             LatencyTrackingStateConfig latencyTrackingStateConfig,
+            SizeTrackingStateConfig sizeTrackingStateConfig,
             @Nonnull Collection<KeyedStateHandle> stateHandles,
             StreamCompressionDecorator keyGroupCompressionDecorator,
             CloseableRegistry cancelStreamRegistry) {
@@ -70,6 +73,7 @@ public AbstractKeyedStateBackendBuilder(
         this.executionConfig = executionConfig;
         this.ttlTimeProvider = ttlTimeProvider;
         this.latencyTrackingStateConfig = latencyTrackingStateConfig;
+        this.sizeTrackingStateConfig = sizeTrackingStateConfig;
         this.keyGroupCompressionDecorator = keyGroupCompressionDecorator;
         this.restoreStateHandles = stateHandles;
         this.cancelStreamRegistry = cancelStreamRegistry;

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

@@ -21,6 +21,7 @@
 import org.apache.flink.annotation.PublicEvolving;
 import org.apache.flink.api.common.ExecutionConfig;
 import org.apache.flink.runtime.state.metrics.LatencyTrackingStateConfig;
+import org.apache.flink.runtime.state.metrics.SizeTrackingStateConfig;
 
 import java.io.IOException;
 
@@ -47,6 +48,9 @@ public static StreamCompressionDecorator getCompressionDecorator(
     protected LatencyTrackingStateConfig.Builder latencyTrackingConfigBuilder =
             LatencyTrackingStateConfig.newBuilder();
 
+    protected SizeTrackingStateConfig.Builder sizeTrackingConfigBuilder =
+            SizeTrackingStateConfig.newBuilder();
+
     // ------------------------------------------------------------------------
     //  State Backend - State-Holding Backends
     // ------------------------------------------------------------------------

flink-runtime/src/main/java/org/apache/flink/runtime/state/hashmap/HashMapStateBackend.java

@@ -33,6 +33,7 @@
 import org.apache.flink.runtime.state.heap.HeapKeyedStateBackendBuilder;
 import org.apache.flink.runtime.state.heap.HeapPriorityQueueSetFactory;
 import org.apache.flink.runtime.state.metrics.LatencyTrackingStateConfig;
+import org.apache.flink.runtime.state.metrics.SizeTrackingStateConfig;
 
 import java.io.IOException;
 
@@ -71,6 +72,7 @@ public HashMapStateBackend() {}
     private HashMapStateBackend(HashMapStateBackend original, ReadableConfig config) {
         // configure latency tracking
         latencyTrackingConfigBuilder = original.latencyTrackingConfigBuilder.configure(config);
+        sizeTrackingConfigBuilder = original.sizeTrackingConfigBuilder.configure(config);
     }
 
     @Override
@@ -101,6 +103,8 @@ public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
 
         LatencyTrackingStateConfig latencyTrackingStateConfig =
                 latencyTrackingConfigBuilder.setMetricGroup(parameters.getMetricGroup()).build();
+        SizeTrackingStateConfig sizeTrackingStateConfig =
+                sizeTrackingConfigBuilder.setMetricGroup(parameters.getMetricGroup()).build();
         return new HeapKeyedStateBackendBuilder<>(
                         parameters.getKvStateRegistry(),
                         parameters.getKeySerializer(),
@@ -110,6 +114,7 @@ public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(
                         parameters.getEnv().getExecutionConfig(),
                         parameters.getTtlTimeProvider(),
                         latencyTrackingStateConfig,
+                        sizeTrackingStateConfig,
                         parameters.getStateHandles(),
                         getCompressionDecorator(parameters.getEnv().getExecutionConfig()),
                         localRecoveryConfig,

flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackend.java

@@ -51,6 +51,7 @@
 import org.apache.flink.runtime.state.StateSnapshotTransformers;
 import org.apache.flink.runtime.state.StreamCompressionDecorator;
 import org.apache.flink.runtime.state.metrics.LatencyTrackingStateConfig;
+import org.apache.flink.runtime.state.metrics.SizeTrackingStateConfig;
 import org.apache.flink.runtime.state.ttl.TtlTimeProvider;
 import org.apache.flink.util.FlinkRuntimeException;
 import org.apache.flink.util.StateMigrationException;
@@ -141,6 +142,7 @@ public HeapKeyedStateBackend(
             ExecutionConfig executionConfig,
             TtlTimeProvider ttlTimeProvider,
             LatencyTrackingStateConfig latencyTrackingStateConfig,
+            SizeTrackingStateConfig sizeTrackingStateConfig,
             CloseableRegistry cancelStreamRegistry,
             StreamCompressionDecorator keyGroupCompressionDecorator,
             Map<String, StateTable<K, ?, ?>> registeredKVStates,
@@ -158,6 +160,7 @@ public HeapKeyedStateBackend(
                 executionConfig,
                 ttlTimeProvider,
                 latencyTrackingStateConfig,
+                sizeTrackingStateConfig,
                 cancelStreamRegistry,
                 keyGroupCompressionDecorator,
                 keyContext);

flink-runtime/src/main/java/org/apache/flink/runtime/state/heap/HeapKeyedStateBackendBuilder.java

@@ -33,6 +33,7 @@
 import org.apache.flink.runtime.state.SavepointKeyedStateHandle;
 import org.apache.flink.runtime.state.StreamCompressionDecorator;
 import org.apache.flink.runtime.state.metrics.LatencyTrackingStateConfig;
+import org.apache.flink.runtime.state.metrics.SizeTrackingStateConfig;
 import org.apache.flink.runtime.state.ttl.TtlTimeProvider;
 
 import javax.annotation.Nonnull;
@@ -69,6 +70,7 @@ public HeapKeyedStateBackendBuilder(
             ExecutionConfig executionConfig,
             TtlTimeProvider ttlTimeProvider,
             LatencyTrackingStateConfig latencyTrackingStateConfig,
+            SizeTrackingStateConfig sizeTrackingStateConfig,
             @Nonnull Collection<KeyedStateHandle> stateHandles,
             StreamCompressionDecorator keyGroupCompressionDecorator,
             LocalRecoveryConfig localRecoveryConfig,
@@ -84,6 +86,7 @@ public HeapKeyedStateBackendBuilder(
                 executionConfig,
                 ttlTimeProvider,
                 latencyTrackingStateConfig,
+                sizeTrackingStateConfig,
                 stateHandles,
                 keyGroupCompressionDecorator,
                 cancelStreamRegistry);
@@ -115,6 +118,7 @@ public HeapKeyedStateBackend<K> build() throws BackendBuildingException {
                 executionConfig,
                 ttlTimeProvider,
                 latencyTrackingStateConfig,
+                sizeTrackingStateConfig,
                 cancelStreamRegistryForBackend,
                 keyGroupCompressionDecorator,
                 registeredKVStates,