RocksBdbGetSliceBench.java

package com.bic.jmh;

import java.io.*;
import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.function.Supplier;

import com.sleepycat.je.Cursor;
import com.sleepycat.je.Database;
import com.sleepycat.je.DatabaseConfig;
import com.sleepycat.je.DatabaseEntry;
import com.sleepycat.je.Environment;
import com.sleepycat.je.EnvironmentConfig;
import com.sleepycat.je.Get;
import com.sleepycat.je.LockMode;
import com.sleepycat.je.OperationResult;
import com.sleepycat.je.ReadOptions;
import com.sleepycat.je.Transaction;
import com.sleepycat.je.TransactionConfig;
import org.apache.commons.io.FileUtils;
import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.Fork;
import org.openjdk.jmh.annotations.OutputTimeUnit;
import org.openjdk.jmh.annotations.Param;
import org.openjdk.jmh.annotations.Scope;
import org.openjdk.jmh.annotations.Setup;
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.annotations.TearDown;
import org.openjdk.jmh.infra.Blackhole;
import org.openjdk.jmh.profile.StackProfiler;
import org.openjdk.jmh.runner.Runner;
import org.openjdk.jmh.runner.RunnerException;
import org.openjdk.jmh.runner.options.Options;
import org.openjdk.jmh.runner.options.OptionsBuilder;
import org.openjdk.jmh.runner.options.TimeValue;
import org.rocksdb.BlockBasedTableConfig;
import org.rocksdb.BloomFilter;
import org.rocksdb.ColumnFamilyDescriptor;
import org.rocksdb.ColumnFamilyHandle;
import org.rocksdb.ColumnFamilyOptions;
import org.rocksdb.CompressionType;
import org.rocksdb.LRUCache;
import org.rocksdb.RocksDBException;
import org.rocksdb.RocksIterator;
import org.rocksdb.TransactionDB;
import org.rocksdb.TransactionDBOptions;
import org.rocksdb.TransactionOptions;
import org.rocksdb.WriteOptions;

@State(Scope.Benchmark)
@OutputTimeUnit(TimeUnit.MILLISECONDS)
@Fork(jvmArgsAppend = "-Xmx1G")
public class RocksBdbGetSliceBench {
    public static final int NUM_KEYS = 1_000_000;
    public static final int CACHE_PERCENTAGE = 30;
    private final List<byte[]> keys = new ArrayList<>();
    private Random random;

    public Supplier<Object> startTx;
    public BiFunction<Object, byte[], List<byte[]>> getSlice;
    public Consumer<Object> rollbackTx;
    public Consumer<Object> onClose;

    @Param({ "rocksdb", "rocksdb-opt", "bdb" })
    String type;

    @Param({"7", "10", "13"})
    Integer seed = 7;

    // @Param({"100", "1000", "10000"})
    Integer numReads = 1000;

    // @Param({"10", "20"})
    Integer keyLength = 10;

    // @Param({"5", "100"})
    Integer valueLength = 5;

    @Setup
    public void setup() throws RocksDBException {
        random = new Random(seed);
        File directory = new File("/tmp/test" + type);
        try {
            if (directory.exists()) {
                FileUtils.cleanDirectory(directory);
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        if ("bdb".equals(type)) {
            EnvironmentConfig envConfig = new EnvironmentConfig();
            envConfig.setAllowCreate(true);
            envConfig.setTransactional(true);
            envConfig.setCachePercent(CACHE_PERCENTAGE);
            envConfig.setSharedCache(true);
            Environment environment = new Environment(directory, envConfig);

            DatabaseConfig dbConfig = new DatabaseConfig();
            dbConfig.setAllowCreate(true);
            dbConfig.setKeyPrefixing(true);
            dbConfig.setTransactional(true);

            Database db = environment.openDatabase(null, "db", dbConfig);

            startTx = () -> environment.beginTransaction(null, new TransactionConfig().setReadUncommitted(true));
            rollbackTx = (t) -> ((Transaction) t).abort();
            getSlice = (t, key) -> {
                DatabaseEntry foundKey = new DatabaseEntry(key);
                DatabaseEntry foundData = new DatabaseEntry();

                Cursor cursor = db.openCursor((Transaction) t, null);

                ReadOptions options = new ReadOptions().setLockMode(LockMode.READ_UNCOMMITTED);
                OperationResult status = cursor.get(foundKey, foundData, Get.SEARCH_GTE, options);

                List<byte[]> result = new ArrayList<>();
                while (status != null) {
                    result.add(foundData.getData());
                    status = cursor.get(foundKey, foundData, Get.NEXT, options);
                    if (Arrays.compare(foundKey.getData(), key) >= 0) {
                        break;
                    }
                }
                cursor.close();
                return result;
            };
            onClose = (ignored) -> {
                db.close();
                environment.close();
            };

            Transaction tx = environment.beginTransaction(null, new TransactionConfig());
            for (int i = 0; i < NUM_KEYS; i++) {
                byte[] key = getBytes(keyLength);
                db.put(tx, new DatabaseEntry(key), new DatabaseEntry(getBytes(valueLength)));
                if (i % 10_000 == 0) {
                    tx.commit();
                    tx = environment.beginTransaction(null, new TransactionConfig());
                }
            }
            tx.commitSync();
        } else if (type.startsWith("rocksdb")) {
            org.rocksdb.Options options = new org.rocksdb.Options().setCreateIfMissing(true);
            final BlockBasedTableConfig config = new BlockBasedTableConfig();
            if (type.contains("opt")) {
                config.setPinL0FilterAndIndexBlocksInCache(true)
                      .setCacheIndexAndFilterBlocks(true);
                BloomFilter filterPolicy = new BloomFilter(10, false);
                config.setFilterPolicy(filterPolicy)
                      .setOptimizeFiltersForMemory(true);
                config.setEnableIndexCompression(false);
            }
            long cacheSizeInBytes = Runtime.getRuntime().maxMemory() * CACHE_PERCENTAGE / 100;
            LRUCache cache = new LRUCache(cacheSizeInBytes);
            config.setBlockCache(cache);

            options.setTableFormatConfig(config);
            TransactionDBOptions txDbOptions = new TransactionDBOptions();
            TransactionDB db = TransactionDB.open(options, txDbOptions, directory.getAbsolutePath());

            final ColumnFamilyOptions columnFamilyOptions = new ColumnFamilyOptions();
            if (type.contains("opt")) {
                columnFamilyOptions.setTargetFileSizeBase(128 * 1024 * 1024);
                columnFamilyOptions.setTargetFileSizeMultiplier(10);
                columnFamilyOptions.setCompressionType(CompressionType.NO_COMPRESSION);
                // columnFamilyOptions.setCompactionStyle(CompactionStyle.UNIVERSAL);
                // columnFamilyOptions.setCompactionStyle(CompactionStyle.NONE);
                // columnFamilyOptions.setDisableAutoCompactions(true);
            }
            ColumnFamilyHandle cf = db.createColumnFamily(new ColumnFamilyDescriptor("db".getBytes(), columnFamilyOptions));

            startTx = () -> {
                TransactionOptions transactionOptions = new TransactionOptions();
                WriteOptions writeOptions = new WriteOptions();
                return db.beginTransaction(writeOptions, transactionOptions);
            };
            rollbackTx = (t) -> {
                try {
                    ((org.rocksdb.Transaction) t).rollback();
                    ((org.rocksdb.Transaction) t).close();
                } catch (RocksDBException e) {
                    throw new RuntimeException(e);
                }
            };
            getSlice = (t, key) -> {
                org.rocksdb.ReadOptions readOptions = new org.rocksdb.ReadOptions();
                RocksIterator iterator = ((org.rocksdb.Transaction) t).getIterator(readOptions, cf);
                iterator.seek(key);

                List<byte[]> result = new ArrayList<>();
                while (iterator.isValid()) {
                    result.add(iterator.value());
                    iterator.next();
                    if (Arrays.compare(iterator.key(), key) >= 0) {
                        break;
                    }
                }
                iterator.close();
                readOptions.close();
                return result;
            };
            onClose = (ignored) -> {
                cf.close();
                db.close();
                options.close();
                txDbOptions.close();
            };

            org.rocksdb.Transaction tx = db.beginTransaction(new WriteOptions(), new TransactionOptions());
            for (int i = 0; i < NUM_KEYS; i++) {
                byte[] key = getBytes(keyLength);
                tx.put(cf, key, getBytes(valueLength));
                if (i % 10_000 == 0) {
                    tx.commit();
                    tx.close();
                    tx = db.beginTransaction(new WriteOptions(), new TransactionOptions());
                }
            }
            tx.commit();
            tx.close();
        } else {
            throw new UnsupportedOperationException(type);
        }

        for (int i = 0; i < numReads; i++) {
            keys.add(getBytes(keyLength));
        }
    }

    private byte[] getBytes(final int count) {
        byte[] key = new byte[count];
        random.nextBytes(key);
        return key;
    }

    @Benchmark
    public void getSlice(Blackhole bh) {
        Object tx = startTx.get();
        for (int i = 0; i < numReads; i++) {
            List<byte[]> apply = getSlice.apply(tx, keys.get(i));
            bh.consume(apply);
        }
        rollbackTx.accept(tx);
    }

    @TearDown
    public void tearDown() {
        onClose.accept(1);
    }

    public static void main(String[] args) throws RunnerException, RocksDBException {
        Options opt = new OptionsBuilder()
            .include(RocksBdbGetSliceBench.class.getSimpleName())
            .addProfiler(StackProfiler.class,
                "lines=15;top=10;period=1;excludePackages=true;detailLine=true;"
                + "excludePackageNames=java.,sun.,com.sun.,jdk.internal.,org.openjdk.,com.intellij.,")
            .forks(1)
            .timeout(TimeValue.minutes(60))
            .measurementIterations(3)
            .measurementTime(TimeValue.seconds(5))
            .warmupIterations(0)
            .warmupTime(TimeValue.seconds(5))
            .build();
        new Runner(opt).run();
    }
}