如何在Flink流任务中正确访问RocksDB？

Hey folks, after digging into how to use RocksDB for caching within a Flink ProcessFunction, I've found only two reliable approaches that actually work in practice:

Approach 1: Load Data to RocksDB in open() and Close the Handle Afterward #

This approach uses Flink's lifecycle methods to pre-load all required cache data upfront. You'll connect to your data source (like MySQL) in the open() method, bulk-load the data into RocksDB, and then close the RocksDB handle right after the load completes.

It’s ideal for scenarios where your cache data doesn’t update frequently—you get fast local access without the overhead of opening/closing the RocksDB handle every time you process an element.

Example Code Snippet #

import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.ProcessFunction;
import org.apache.flink.util.Collector;
import org.rocksdb.Options;
import org.rocksdb.RocksDB;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.nio.charset.StandardCharsets;

public static class MatchFunction extends ProcessFunction<TaxiRide, TaxiRide> {
    private transient RocksDB rocksDB;
    private Options options;

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        // Initialize RocksDB configuration
        options = new Options().setCreateIfMissing(true);
        rocksDB = RocksDB.open(options, "/path/to/rocksdb/storage");

        // Load cache data from MySQL (pseudo-code for DB connection)
        try (Connection conn = DriverManager.getConnection("jdbc:mysql://host:port/db", "user", "pass")) {
            String query = "SELECT ride_id, ride_details FROM cached_ride_data";
            try (PreparedStatement stmt = conn.prepareStatement(query);
                 ResultSet rs = stmt.executeQuery()) {
                while (rs.next()) {
                    byte[] key = rs.getString("ride_id").getBytes(StandardCharsets.UTF_8);
                    byte[] value = rs.getString("ride_details").getBytes(StandardCharsets.UTF_8);
                    rocksDB.put(key, value);
                }
            }
        } finally {
            // Close RocksDB handle after pre-loading is done
            if (rocksDB != null) {
                rocksDB.close();
            }
        }
    }

    @Override
    public void processElement(TaxiRide ride, Context ctx, Collector<TaxiRide> out) throws Exception {
        // Re-open RocksDB if we need to access cached data during element processing
        try (RocksDB db = RocksDB.open(options, "/path/to/rocksdb/storage")) {
            byte[] cachedDetails = db.get(ride.getRideId().getBytes(StandardCharsets.UTF_8));
            if (cachedDetails != null) {
                // Use cached data for business logic
                System.out.println("Cached details for ride " + ride.getRideId() + ": " + new String(cachedDetails));
            }
            out.collect(ride);
        }
    }

    @Override
    public void close() throws Exception {
        super.close();
        if (options != null) {
            options.close();
        }
    }
}

Approach 2: Open and Close RocksDB Handle in Each processElement() Call #

Here, you open the RocksDB handle at the start of every processElement() call, perform your cache read/write operations, and close the handle immediately afterward (using try-with-resources to ensure proper cleanup).

This is a good fit if your cache data changes frequently and you need to access the latest version every time you process an element. Just keep in mind: opening/closing the handle for every element can add performance overhead, so use this only when necessary.

Example Code Snippet #

import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.ProcessFunction;
import org.apache.flink.util.Collector;
import org.rocksdb.Options;
import org.rocksdb.RocksDB;
import org.apache.commons.lang3.SerializationUtils;
import java.nio.charset.StandardCharsets;

public static class MatchFunction extends ProcessFunction<TaxiRide, TaxiRide> {
    private transient Options options;

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        // Initialize RocksDB options once to avoid re-creating them per element
        options = new Options().setCreateIfMissing(true);
    }

    @Override
    public void processElement(TaxiRide ride, Context ctx, Collector<TaxiRide> out) throws Exception {
        // Open RocksDB handle for this element's processing
        try (RocksDB rocksDB = RocksDB.open(options, "/path/to/rocksdb/storage")) {
            String cacheKey = "END_EVENT_" + ride.getRideId();
            byte[] cachedEndEvent = rocksDB.get(cacheKey.getBytes(StandardCharsets.UTF_8));
            
            if (ride.isEndEvent()) {
                // Save end event to RocksDB cache
                rocksDB.put(cacheKey.getBytes(StandardCharsets.UTF_8), SerializationUtils.serialize(ride));
            } else {
                // Check if we have a matching end event cached
                if (cachedEndEvent != null) {
                    TaxiRide matchedEndRide = (TaxiRide) SerializationUtils.deserialize(cachedEndEvent);
                    // Execute matching logic
                    System.out.println("Matched start/end rides for ID: " + ride.getRideId());
                }
            }
            
            out.collect(ride);
        } // RocksDB handle auto-closes here via try-with-resources
    }

    @Override
    public void close() throws Exception {
        super.close();
        if (options != null) {
            options.close();
        }
    }
}

内容的提问来源于stack exchange，提问作者James Yu