Flink中KeyedStateStore实现--怎么做到一个Key对应一个State

背景

在Flink中有两种基本的状态：Keyed State和Operator State，Operator State很好理解，一个特定的Operator算子共享同一个state，这是实现层面很好做到的。
但是 Keyed State 是怎么实现的？一般来说，正常的人第一眼就会想到:一个task绑定一个Keyd State,从网上随便查找资料就能发现正确的答案是:对于每一个Key会绑定一个State，但是这在Flink中是怎么实现的呢？
注意:这里我们只讲Flink中是怎么实现一个Key对应一个State的，其他细节并不细说,且state的backend为RocksDB

闲说杂谈

我们以ValueState类型的Keyed State举例：

ValueStateDescriptor indexStateDesc =
        new ValueStateDescriptor<>(
            "indexState",
            TypeInformation.of(HoodieRecordGlobalLocation.class));
ValueState indexState = context.getKeyedStateStore().getState(indexStateDesc)
....
indexState.update((HoodieRecordGlobalLocation) indexRecord.getCurrentLocation())

• context.getKeyedStateStore().getState是获取对应key的State，最终的调用链如下：

   DefaultKeyedStateStore.getState -> getPartitionedState
                                                ||
                                                \/
   RocksDBKeyedStateBackend.getPartitionedState -> getOrCreateKeyedState -> createInternalState -> tryRegisterKvStateInformation
                                                                                                               ||
                                                                                                               \/
  
   RocksDBValueState.create(创建RocksDBValueState)                                                                             
                                        
  

  这里的 tryRegisterKvStateInformation会涉及到RocksDB ColumnFamily的创建：

  RocksDBOperationUtils.createStateInfo -> createColumnFamilyDescriptor 
  // createColumnFamilyDescriptor的部分代码：
  ColumnFamilyOptions options =
               createColumnFamilyOptions(columnFamilyOptionsFactory, metaInfoBase.getName());
  if (ttlCompactFiltersManager != null) {
      ttlCompactFiltersManager.setAndRegisterCompactFilterIfStateTtl(metaInfoBase, options);
  }
  byte[] nameBytes = metaInfoBase.getName().getBytes(ConfigConstants.DEFAULT_CHARSET);
  ...
  return new ColumnFamilyDescriptor(nameBytes, options);
  
  

  其实最终会发现RocksDB的ColumnFamily是跟ValueStateDescriptor也就是描述符的名字有关的，这就是为什么描述符必须是唯一的,关于RocksDB的ColumnFamily，可以参考RocksDB 简介
  注意此时返回是key对应的一个State的ColumnFamily，该Family包括该task所有的key的value值

• indexState.update 这里是更新indexState得值
  因为上一步得到只是该Task所对应的ColumanFamily所对应的所有的values，也就是* Flink中的Key-Groups*,(关于Key-Groups可以参考Apache-Flink深度解析-State)

    public void update(V value) {
       if (value == null) {
           clear();
           return;
       }
  
       try {
           backend.db.put(
                   columnFamily,
                   writeOptions,
                   serializeCurrentKeyWithGroupAndNamespace(),
                   serializeValue(value));
       } catch (Exception e) {
           throw new FlinkRuntimeException("Error while adding data to RocksDB", e);
       }
   }
  

  最终的调用链如下:

  RocksDBValueState.update -> serializeCurrentKeyWithGroupAndNamespace
          ||
          \/
  SerializedCompositeKeyBuilder.buildCompositeKeyNamespace
          ||
          \/
  serializeNamespace(namespace, namespaceSerializer) -> keyOutView.getCopyOfBuffer()   
  
  

  这里的keyOutView.getCopyOfBuffer是会获得的record的key，所以在backend.db.put方法中才会更新对应的Key值。
  但是什么时候Record的key信息会被写入到keyOutView中去呢？

• Record的key何时被写到keyOutView中

  AbstractStreamTaskNetworkInput.emitNext -> processElement
         ||
         \/
  OneInputStreamTask.emitRecord
         ||
         \/
  OneInputStreamOperator.setKeyContextElement -> setKeyContextElement1 -> setKeyContextElement
  
         ||
         \/
  AbstractStreamOperator.setCurrentKey
         ||
         \/
  StreamOperatorStateHandler.setCurrentKey
         ||
         \/
  RocksDBKeyedStateBackend.setCurrentKey
         ||
         \/
  SerializedCompositeKeyBuilder.setCurrentKey -> serializeKeyGroupAndKey
         ||
         \/
  keySerializer.serialize(key, keyOutView);    
  
  

  最后一步keySerializer.serialize(key, keyOutView)一个Record的key就被写到keyOutView中,也就是说对应的key是从每个record中获取的，所以在backend.db.put方法中就能获取到对应的Key

其他

对于keyedStateStore是在哪里初始化的，可以看AbstractStreamOperator中initializeState方法:

final StreamOperatorStateContext context =
             streamTaskStateManager.streamOperatorStateContext(
                     getOperatorID(),
                     getClass().getSimpleName(),
                     getProcessingTimeService(),
                     this,
                     keySerializer,
                     streamTaskCloseableRegistry,
                     metrics,
                     config.getManagedMemoryFractionOperatorUseCaseOfSlot(
                             ManagedMemoryUseCase.STATE_BACKEND,
                             runtimeContext.getTaskManagerRuntimeInfo().getConfiguration(),
                             runtimeContext.getUserCodeClassLoader()),
                     isUsingCustomRawKeyedState());

stateHandler =
        new StreamOperatorStateHandler(
                context, getExecutionConfig(), streamTaskCloseableRegistry);

这个方法里也包括了keyedStatedBackend和operatorStateBackend等初始化， 具体的细节后续再解析。