Range search in LmdbJava

IntegerKeyTest2.java

import org.lmdbjava.*;

import java.io.File;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.file.Files;
import java.time.LocalDateTime;
import java.time.Month;
import java.time.ZoneOffset;
import java.util.Comparator;

import static java.nio.ByteBuffer.allocateDirect;
import static java.nio.ByteOrder.nativeOrder;
import static java.util.Comparator.comparing;
import static org.lmdbjava.DbiFlags.*;

public final class IntegerKeyTest2 {
    private static final Comparator<ByteBuffer> cmp = comparing(a -> a.order(nativeOrder()));
    private static final int keyCount = 10;

    private final ByteBuffer KEY = allocateDirect(8).order(nativeOrder());
    private final ByteBuffer VALUE = allocateDirect(4).order(nativeOrder());

    private final Env<ByteBuffer> env;
    private final Dbi<ByteBuffer> db;

    public IntegerKeyTest2() throws IOException {
        File dir = Files.createTempDirectory(IntegerKeyTest2.class.getSimpleName()).toFile();
        env = Env.create().open(dir);
        db = env.openDbi(IntegerKeyTest2.class.getSimpleName(),
                MDB_CREATE, MDB_INTEGERKEY, MDB_DUPSORT);
    }

    public void execute() {
        long[] keys = createSortedKeys();
        for (int i = 0; i < keys.length; i++) {
            add(keys[i], i);
        }

        rangeTest(KeyRange.all(), keys.length);

        long min = keys[0];
        long max = keys[keys.length - 1];

        KeyRange<ByteBuffer> range = KeyRange.closed(bb(min), bb(max * keyCount));
        rangeTest(range, keys.length);

        KeyRange<ByteBuffer> range2 = KeyRange.closed(bb(min), bb(max));
        rangeTest(range2, keys.length); // fails for this case
    }

    private void rangeTest(KeyRange<ByteBuffer> range, int expected) {
        try (Txn<ByteBuffer> txn = env.txnRead();
             CursorIterator<ByteBuffer> itr = db.iterate(txn, range, cmp)) {
            int count = 0;
            for (CursorIterator.KeyVal<ByteBuffer> kv : itr.iterable()) {
                count++;
            }

            assert count == expected : "Expected: " + expected + ", got: " + count;
        }
    }

    private static long[] createSortedKeys() {
        long[] xs = new long[keyCount];
        LocalDateTime start = LocalDateTime.of(2018, Month.JANUARY, 1, 0, 0);
        for (int i = 0; i < xs.length; i++) {
            xs[i] = toMillis(start.plusHours(i));
        }

        return xs;
    }

    private void add(long key, int value) {
        KEY.clear();
        KEY.putLong(key).flip();

        VALUE.clear();
        VALUE.putInt(value).flip();

        db.put(KEY, VALUE);
    }

    private static ByteBuffer bb(long a) {
        ByteBuffer buffer = allocateDirect(8).order(nativeOrder());
        buffer.putLong(a).flip();
        return buffer;
    }

    private static long toMillis(LocalDateTime time) {
        return time.toInstant(ZoneOffset.UTC).toEpochMilli();
    }

    public static void main(String[] args) throws IOException {
        IntegerKeyTest2 map = new IntegerKeyTest2();
        map.execute();
    }
}