/*************************************************************************
* *
* This file is part of the 20n/act project. *
* 20n/act enables DNA prediction for synthetic biology/bioengineering. *
* Copyright (C) 2017 20n Labs, Inc. *
* *
* Please direct all queries to act@20n.com. *
* *
* This program is free software: you can redistribute it and/or modify *
* it under the terms of the GNU General Public License as published by *
* the Free Software Foundation, either version 3 of the License, or *
* (at your option) any later version. *
* *
* This program is distributed in the hope that it will be useful, *
* but WITHOUT ANY WARRANTY; without even the implied warranty of *
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the *
* GNU General Public License for more details. *
* *
* You should have received a copy of the GNU General Public License *
* along with this program. If not, see <http://www.gnu.org/licenses/>. *
* *
*************************************************************************/
package com.act.utils.rocksdb;
import org.rocksdb.ColumnFamilyHandle;
import org.rocksdb.FlushOptions;
import org.rocksdb.RocksDB;
import org.rocksdb.RocksDBException;
import org.rocksdb.RocksIterator;
import org.rocksdb.WriteBatch;
import org.rocksdb.WriteOptions;
import java.util.Map;
public class RocksDBAndHandles<T extends ColumnFamilyEnumeration<T>> {
RocksDB db;
Map<T, ColumnFamilyHandle> columnFamilyHandleMap;
WriteOptions writeOptions = null;
// Here to hijack the DB interface for easy testing or alternate implementations.
protected RocksDBAndHandles() {
}
public RocksDBAndHandles(RocksDB db, Map<T, ColumnFamilyHandle> columnFamilyHandleMap) {
this.db = db;
this.columnFamilyHandleMap = columnFamilyHandleMap;
}
public void close() {
this.db.close();
}
public WriteOptions getWriteOptions() {
return writeOptions;
}
public void setWriteOptions(WriteOptions writeOptions) {
this.writeOptions = writeOptions;
}
public RocksDB getDb() {
return db;
}
protected Map<T, ColumnFamilyHandle> getColumnFamilyHandleMap() {
return columnFamilyHandleMap;
}
protected ColumnFamilyHandle getHandle(T columnFamilyLabel) {
return this.columnFamilyHandleMap.get(columnFamilyLabel);
}
public void put(T columnFamily, byte[] key, byte[] val) throws RocksDBException {
if (writeOptions != null) {
this.db.put(getHandle(columnFamily), writeOptions, key, val);
} else {
this.db.put(getHandle(columnFamily), key, val);
}
}
public boolean keyMayExist(T columnFamily, byte[] key) throws RocksDBException {
StringBuffer buffer = new StringBuffer();
return this.db.keyMayExist(getHandle(columnFamily), key, buffer);
}
public byte[] get(T columnFamily, byte[] key) throws RocksDBException {
return this.db.get(getHandle(columnFamily), key);
}
public void flush(boolean waitForFlush) throws RocksDBException {
FlushOptions options = new FlushOptions();
options.setWaitForFlush(waitForFlush);
db.flush(options);
}
/* Important: don't expose merge(), as it appears to be broken in RocksDB JNI.
*
* RocksDB supports "merge" functionality, where a new value can be merged at the DB level (as opposed to the client
* level) into an existing value for a given key using a predefined function.
* See https://github.com/facebook/rocksdb/wiki/Merge-Operator.
*
* This ought to be super fast, as it's being done in the same stroke as the DB lookup--the old and new values should
* be in the CPU cache (assuming they fit) and the merge is done in a layer of the library that should have some
* notion of what operations will do right by the storage system. So I tried it!
*
* Alas, calling merge() from Java causes the native layer to throw an exception that complains about a merge
* function not being defined. When you specify a built-in merge function in the DB constructor, the exception
* *still* gets thrown. Sad face!
*
* The workable alternative is to bubble the value all the way up to Java land, merge there, and then send the
* resulting value back to the DB. This tends to be incredibly slow, however, so just don't expose it at all.
*/
// Wrap cursors for easier testing.
public RocksDBIterator newIterator(T columnFamily) throws RocksDBException {
return new RocksDBIterator(this.db.newIterator(getHandle(columnFamily)));
}
// Wrap write batches for easier CF management and testing.
public RocksDBWriteBatch<T> makeWriteBatch() {
return new RocksDBWriteBatch<T>(this, RocksDBWriteBatch.RESERVED_BYTES);
}
public RocksDBWriteBatch<T> makeWriteBatch(int reservedBytes) {
return new RocksDBWriteBatch<T>(this, reservedBytes);
}
/* ----------------------------------------
* Proxy classes for write batches and cursors. These proxies allow us to condense the API to the parts we care about
* and create hooks we can override for testing without using an actual DB.
*/
public static class RocksDBWriteBatch<T extends ColumnFamilyEnumeration<T>> {
protected static final int RESERVED_BYTES = 1 << 18;
private static final WriteOptions DEFAULT_WRITE_OPTIONS;
static {
// Note: the WriteOptions constructor requires a native library call, so make sure RocksDB is loaded first.
RocksDB.loadLibrary();
DEFAULT_WRITE_OPTIONS = new WriteOptions();
}
WriteBatch batch;
RocksDBAndHandles<T> parent;
protected RocksDBWriteBatch() {
// Just for testing.
}
protected RocksDBWriteBatch(RocksDBAndHandles<T> parent, int reservedBytes) {
this.parent = parent;
this.batch = new WriteBatch(reservedBytes);
}
public void put(T columnFamily, byte[] key, byte[] val) throws RocksDBException {
batch.put(parent.getHandle(columnFamily), key, val);
}
public void write() throws RocksDBException {
WriteOptions writeOptions = parent.getWriteOptions() != null ? parent.getWriteOptions() : DEFAULT_WRITE_OPTIONS;
parent.getDb().write(writeOptions, batch);
}
}
/* RocksDB's iterators don't implement Iterable because they talk about byte arrays instead of objects (hence no
* Object subtype to bind to T). It's API is also a little backwards: rather than calling `hasNext` and then `next`,
* you call `next` and then `isValid` to check if you've gone off the end of the index. We don't do anything to
* simplify that right now, but we could in the future. */
public static class RocksDBIterator {
RocksIterator rocksIter;
protected RocksDBIterator() {
// Just for testing.
}
protected RocksDBIterator(RocksIterator rocksIter) {
this.rocksIter = rocksIter;
this.rocksIter.seekToFirst(); // How does this not throw an exception? Oh well.
}
public void reset() { // Easy synonym so users don't have to think about seeking.
seekToFirst();
}
public void seekToFirst() {
rocksIter.seekToFirst();
}
public void seekToLast() {
rocksIter.seekToLast();
}
// Don't expose seek-to-byte yet. Not sure what that byte offset means.
public void next() {
rocksIter.next();
}
public void prev() {
rocksIter.prev();
}
public boolean isValid() {
return rocksIter.isValid();
}
public byte[] value() {
return rocksIter.value();
}
public byte[] key() {
return rocksIter.key();
}
}
}