RocksDBStateBackend stores large state in Flink.

Features and Use Cases

As discussed in the previous article, large state makes FsStateBackend vulnerable to GC pauses. RocksDBStateBackend is a more sensible choice in this scenario, although its drawbacks are also clear:

• More overhead than FsStateBackend
  • Serialization and deserialization of Java objects
  • Disk I/O
  • RocksDB compaction
• Poor support for Read-Modify-Write workloads (discussed below)

Its advantage is the ability to store much larger state. A common example is matching clicks and impressions in advertising: storing individual records requires several orders of magnitude more state than storing aggregates. RocksDBStateBackend generally demands fast disks. Without SSDs, I would advise against it. (Performance with HDDs can be hundreds of times worse.)

State Storage Layout

RocksDBStateBackend uses RocksDB underneath. During initialization, each Task creates a RocksDB instance on local disk. Its keys combine the record key and namespace, while each column family identifies a state. For example, tracking clicks and impressions for each user produces the following layout. (Namespaces are omitted here.)

State operations can be understood as operations on this RocksDB instance. The Read-Modify-Write limitation mentioned above comes from RocksDB’s append-only design. As illustrated below, data is read, updated, written into the write buffer, and flushed to disk. The next read loads it from disk again, continually churning the RocksDB cache.

Optimizations

Given these performance concerns, RocksDBStateBackend has received various optimizations since its release. Here are several commonly used ones.

Using RocksDB Merge for ListState

RocksDB introduced Merge Operation to avoid Read-Modify-Write overhead. A callback describes the modification, and the operation is written directly to RocksDB. At read time, RocksDB combines the original value with all subsequent merge operations and returns the result. For example, an unoptimized List add operation looks like this:

1
public void add(T element) {
2
  List<T> list = read from rocksdb and deserialize
3
  list.add(element)
4
  serialize list and write into rocksdb
5
}

This is a typical Read-Modify-Write operation. The optimized implementation only needs to call merge. (Because the merge operation does not expose a Java API, Flink implements this through JNI.)

Cleaning Up Expired Data with RocksDB Compaction

Flink provides three strategies for removing expired data:

• FULL_STATE_SCAN_SNAPSHOT: A full snapshot already traverses all state, so expired entries can be removed during that scan. RocksDBStateBackend normally uses incremental snapshots, making this strategy inapplicable here.
• INCREMENTAL_CLEANUP: Whenever state is touched, read a few entries incrementally and remove expired ones.
• ROCKSDB_COMPACTION_FILTER: This takes advantage of RocksDB’s compaction filter. RocksDB already scans data during compaction, and Flink implements this feature through JNI.

Other Optimizations

Other optimizations include incremental snapshots and MapState caching; I will not cover each of them here.

Snapshots and Recovery

Incremental Snapshots

Because RocksDB is append-only, operations are persisted into a sequence of SST files. At each snapshot, we can flush the write buffer to disk and record the latest SST file identifier. The next snapshot can then identify newly added files. This raises a question: how should historical files be removed after compaction? See incremental checkpointing for details.

Flink uses reference counting to remove historical files. The diagram above illustrates retained-checkpoints=2:

• After CP1 completes, sst-1 and sst-2 each have a reference count of 1.
• After CP2 completes, sst-1 and sst-2 each have a count of 2. Newly added sst-3 and sst-4 each have a count of 1.
• After CP3 completes, sst-1 through sst-3 have been compacted into new files. Since only two checkpoints are retained, CP1 is discarded, reducing the reference counts of sst-1 and sst-2 by 1. The count for sst-3 is unchanged because it is absent from this checkpoint.
• After CP4 completes, CP2 is discarded, reducing the counts of sst-1 and sst-2 by 1 again. No checkpoint depends on these files anymore, so they can be removed.

Recovery

FsStateBackend eagerly loads all state into memory during recovery. RocksDBStateBackend instead uses lazy recovery: it first copies SST files back from HDFS, then loads state as the task runs. Rescaling requires redistributing SST files, and there is no simple solution. If a task’s KeyGroupRange changes from (4-6) to (3-7), it needs its previous SST files plus RocksDB data from two other tasks to restore the state.

Flink handles this as follows:

1. Create a temporary directory, tmp_dir, and a new RocksDB instance, newDB.
2. Create three RocksDB instances to read the files from the three tasks. Use RocksDB range operations to read data with KeyGroupRange in (3-7), and write it through newDB into the temporary directory.
3. Move the temporary data into the final directory to serve as runtime state.

As this shows, recovering RocksDB state after rescaling is also expensive.

Other Work

The community has many further improvements planned for RocksDBStateBackend, including:

1. Merging small SST files before uploading them to HDFS.
2. Adding a cache above RocksDBStateBackend to address Read-Modify-Write workloads.

References

• https://flink.apache.org/features/2018/01/30/incremental-checkpointing.html
• https://github.com/facebook/rocksdb