Flink StateBackend (2) - DefaultOperatorStateBackend

Liao Jiayi Liao Jiayi #Flink#StateBackend#Apache Flink StateBackend

DefaultOperatorStateBackend is a basic and relatively simple component of Flink's state backend.

Translated from Chinese with AI · Read the original

DefaultOperatorStateBackend is a basic and relatively simple component of Flink’s state backend.

Introduction

DefaultOperatorStateBackend manages non-keyed operator state. Common forms are ListState, UnionState, and BroadcastState, each with different uses.

State

ListState

ListState is the most widely used. Kafka offsets, for example, store each TopicPartition’s offset this way. Elements are redistributed during rescaling:

(1) Parallelism = 2
task-1 : List(TopicPartition(partition=1, offset=100), TopicPartition(partition=2, offset=100))
task-2 : List(TopicPartition(partition=3, offset=100), TopicPartition(partition=4, offset=100))
(2) Parallelism = 2 -> 3
task-1 : List(TopicPartition(partition=1, offset=100), TopicPartition(partition=4, offset=100))
task-2 : List(TopicPartition(partition=2, offset=100))
task-3 : List(TopicPartition(partition=3, offset=100))

After scaling out, list elements are evenly redistributed in round-robin order.

UnionState

UnionState is less common. We used it to solve a production watermark problem; see FLINK-5601. Using the Kafka-offset example, checkpoint recovery produces this distribution:

(1) Parallelism = 2
task-1 : List(TopicPartition(partition=1, offset=100), TopicPartition(partition=2, offset=100))
task-2 : List(TopicPartition(partition=3, offset=100), TopicPartition(partition=4, offset=100))
(2) Parallelism = 2 -> 3
task-1 : List(TopicPartition(partition=1, offset=100), TopicPartition(partition=2, offset=100), TopicPartition(partition=3, offset=100), TopicPartition(partition=4, offset=100))
task-2 : List(TopicPartition(partition=1, offset=100), TopicPartition(partition=2, offset=100), TopicPartition(partition=3, offset=100), TopicPartition(partition=4, offset=100))
task-3 : List(TopicPartition(partition=1, offset=100), TopicPartition(partition=2, offset=100), TopicPartition(partition=3, offset=100), TopicPartition(partition=4, offset=100))

Rescaling does not change UnionState’s distribution rule: each restored task receives the union of the lists previously held by every task.

BroadcastState

BroadcastState is usually accessed through BroadcastProcessFunction or KeyedBroadcastProcessFunction rather than directly. For example:

val descriptor = new MapStateDescriptor[Long, String](
      "broadcast-state",
      BasicTypeInfo.LONG_TYPE_INFO.asInstanceOf[TypeInformation[Long]],
      BasicTypeInfo.STRING_TYPE_INFO)


val srcOne = env.addSource(..)
val srcTwo = env.addSource(..)
val broadcast = srcTwo.broadcast(descriptor)
srcOne.connect(broadcast).process(new TestBroadcastProcessFunction)


class TestBroadcastProcessFunction extends KeyedBroadcastProcessFunction[Long, Long, String, String] {


  lazy val localDescriptor = new MapStateDescriptor[Long, String](
    "broadcast-state",
    BasicTypeInfo.LONG_TYPE_INFO.asInstanceOf[TypeInformation[Long]],
    BasicTypeInfo.STRING_TYPE_INFO)


  override def processBroadcastElement(
      value: String,
      ctx: KeyedBroadcastProcessFunction[Long, Long, String, String]#Context,
      out: Collector[String]): Unit = {


    val key = value.split(":")(1).toLong
    ctx.getBroadcastState(localDescriptor).put(key, value)
  }
}

Calling broadcast on srcTwo creates a BroadcastStream, which connects with srcOne. Every srcTwo record then reaches all connected downstream tasks. A KeyedBroadcastProcessFunction can receive the data and store it as BroadcastState, as shown above.

Operations

Modify

UnionState and ListState both present a list interface and use PartitionableListState internally. Its implementation wraps a java.util.ArrayList and supplies methods such as get(), add(String), and update(List).

BroadcastState similarly wraps a java.util.Map and exposes map operations.

Snapshot

Snapshot implementations implement SnapshotStrategy.snapshot, so we can go directly to DefaultOperatorStateBackendSnapshotStrategy.

if (!registeredOperatorStates.isEmpty()) {
  for (Map.Entry<String, PartitionableListState<?>> entry : registeredOperatorStates.entrySet()) {
    PartitionableListState<?> listState = entry.getValue();
    if (null != listState) {
      listState = listState.deepCopy();
    }
    registeredOperatorStatesDeepCopies.put(entry.getKey(), listState);
  }
}


if (!registeredBroadcastStates.isEmpty()) {
  for (Map.Entry<String, BackendWritableBroadcastState<?, ?>> entry : registeredBroadcastStates.entrySet()) {
    BackendWritableBroadcastState<?, ?> broadcastState = entry.getValue();
    if (null != broadcastState) {
      broadcastState = broadcastState.deepCopy();
    }
    registeredBroadcastStatesDeepCopies.put(entry.getKey(), broadcastState);
  }
}

First, deep-copy every PartitionableListState and BackendWritableBroadcastState. This may seem inefficient, and it is: Flink assumes operator state is lightweight. More efficient approaches will be discussed with HeapKeyedStateBackend and RocksDBKeyedStateBackend.

The copy enables subsequent asynchronous work.

AsyncSnapshotCallable<SnapshotResult<OperatorStateHandle>> snapshotCallable =
new AsyncSnapshotCallable<SnapshotResult<OperatorStateHandle>>() {


  @Override
  protected SnapshotResult<OperatorStateHandle> callInternal() throws Exception {


    for (Map.Entry<String, PartitionableListState<?>> entry :
      registeredOperatorStatesDeepCopies.entrySet()) {


      PartitionableListState<?> value = entry.getValue();
      long[] partitionOffsets = value.write(localOut);
      OperatorStateHandle.Mode mode = value.getStateMetaInfo().getAssignmentMode();
      writtenStatesMetaData.put(
        entry.getKey(),
        new OperatorStateHandle.StateMetaInfo(partitionOffsets, mode));
    }


    for (Map.Entry<String, BackendWritableBroadcastState<?, ?>> entry :
      registeredBroadcastStatesDeepCopies.entrySet()) {


      BackendWritableBroadcastState<?, ?> value = entry.getValue();
      long[] partitionOffsets = {value.write(localOut)};
      OperatorStateHandle.Mode mode = value.getStateMetaInfo().getAssignmentMode();
      writtenStatesMetaData.put(
        entry.getKey(),
        new OperatorStateHandle.StateMetaInfo(partitionOffsets, mode));
    }
  }
};

Skipping state metadata serialization, PartitionableListState and BackendWritableBroadcastState follow nearly identical flows. For ListState:

  • Write list values and return each element’s offset.
  • Add stateName and partitionOffsets to the metadata collection.

Offsets enable direct access during recovery. After writing state, write the collected metadata. This completes the task’s DefaultOperatorStateBackend snapshot and creates its own HDFS file.

Recovery

Recovery mainly assigns state handles through JobMaster’s CheckpointCoordinator. Tasks restore from these handles straightforwardly; see OperatorStateRestoreOperation.restore for details.

A completed checkpoint produces x HDFS files. When parallelism changes from p1 to p2, how are those files reassigned? The key code is reDistributePartitionableStates in StateAssignmentOperator:

OperatorStateRepartitioner opStateRepartitioner = RoundRobinOperatorStateRepartitioner.INSTANCE;


for (int operatorIndex = 0; operatorIndex < newOperatorIDs.size(); operatorIndex++) {
  OperatorState operatorState = oldOperatorStates.get(operatorIndex);
  int oldParallelism = operatorState.getParallelism();


  OperatorID operatorID = newOperatorIDs.get(operatorIndex);


  newManagedOperatorStates.putAll(applyRepartitioner(
    operatorID,
    opStateRepartitioner,
    oldManagedOperatorStates.get(operatorIndex),
    oldParallelism,
    newParallelism));
}

Every operator uses RoundRobinOperatorStateRepartitioner.INSTANCE. There are two cases.

Unchanged Parallelism

ListState and BroadcastState can restore directly from their files. UnionState must still combine lists from every task, so only it needs redistribution. First collect all UnionState metadata:

/**
 * Collect union states from given parallelSubtaskStates.
 */
private Map<String, List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>>> collectUnionStates(
  List<List<OperatorStateHandle>> parallelSubtaskStates) {


   // stateName -> List(每个并行度subtask的所有 union state)
  Map<String, List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>>> unionStates =
    new HashMap<>(parallelSubtaskStates.size());


    // 遍历所有的并行度 subtask state
  for (List<OperatorStateHandle> subTaskState : parallelSubtaskStates) {
    for (OperatorStateHandle operatorStateHandle : subTaskState) {
      if (operatorStateHandle == null) {
        continue;
      }


          // 获取 stateName -> StateMetaInfo 的映射集合
      final Set<Map.Entry<String, OperatorStateHandle.StateMetaInfo>> partitionOffsetEntries =
        operatorStateHandle.getStateNameToPartitionOffsets().entrySet();


          // 过滤 UNION 类型的 State
      partitionOffsetEntries.stream()
        .filter(entry -> entry.getValue().getDistributionMode().equals(OperatorStateHandle.Mode.UNION))
        .forEach(entry -> {
           // 把 UNION State 相关信息放入 unionStates 中，每个 stateName 对应的 List 大小为 并行度 * union state数
          List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>> stateLocations =
            unionStates.computeIfAbsent(entry.getKey(), k -> new ArrayList<>(parallelSubtaskStates.size() * partitionOffsetEntries.size()));


          stateLocations.add(Tuple2.of(operatorStateHandle.getDelegateStateHandle(), entry.getValue()));
        });
    }
  }


  return unionStates;
}

The List&lt;List&lt;OperatorStateHandle&gt;&gt; parameter parallelSubtaskStates contains every subtask’s state. The innermost List&lt;OperatorStateHandle&gt; can be treated as an OperatorStateHandle because its size is always one. Union states are then repartitioned:

/**
 * Repartition UNION state.
 */
private void repartitionUnionState(
    Map<String, List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>>> unionState,
    List<Map<StreamStateHandle, OperatorStateHandle>> mergeMapList) {


  for (Map<StreamStateHandle, OperatorStateHandle> mergeMap : mergeMapList) {
    for (Map.Entry<String, List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>>> e :
        unionState.entrySet()) {


      for (Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo> handleWithMetaInfo : e.getValue()) {
        OperatorStateHandle operatorStateHandle = mergeMap.get(handleWithMetaInfo.f0);
        if (operatorStateHandle == null) {
          operatorStateHandle = new OperatorStreamStateHandle(
            new HashMap<>(unionState.size()),
            handleWithMetaInfo.f0);
          mergeMap.put(handleWithMetaInfo.f0, operatorStateHandle);
        }
        operatorStateHandle.getStateNameToPartitionOffsets().put(e.getKey(), handleWithMetaInfo.f1);
      }
    }
  }
}

Changed Parallelism

ListState also needs redistribution when parallelism changes. Since snapshots record every list element’s offset, the repartitioner can determine the total element count and assign them under the new parallelism.

// 通过统计 snapshot 中的元信息算出总共需要分配的元素
int totalPartitions = 0;
for (Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo> offsets : current) {
  totalPartitions += offsets.f1.getOffsets().length;
}


// 每个 Task 需要分配的元素
int baseFraction = totalPartitions / newParallelism;
// 无法整除剩余的元素，采用先到先得的原则分配
int remainder = totalPartitions % newParallelism;