package util.db; import static util.ByteUtils.*; import java.io.File; import java.io.IOException; import java.io.UnsupportedEncodingException; import java.time.LocalDateTime; import java.util.HashMap; import java.util.Map; import java.util.concurrent.atomic.AtomicLong; import org.apache.commons.io.FileUtils; import org.rocksdb.*; import util.TimeWatch; public class RDB { private RocksDB db; private String path; private static final String UTF_8 = "UTF-8"; public RDB() { // different dbs i ncase of concurrent calculations this.path = System.getProperty("java.io.tmpdir") .concat(File.separator) .concat(String.format("corpusAnalyzer_db%d", LocalDateTime.now().toString().hashCode())); this.db = createDB(); } private RocksDB createDB() { RocksDB.loadLibrary(); // the Options class contains a set of configurable DB options // that determines the behaviour of the database. try (final Options options = new Options()) { options.setCreateIfMissing(true); // a factory method that returns a RocksDB instance try (final RocksDB rdb = RocksDB.open(options, path)) { if (db != null) { return rdb; } else { this.db = rdb; } } } catch (RocksDBException e) { // do some error handling } return null; } public void writeBatch(Map results) throws UnsupportedEncodingException { RocksDB.loadLibrary(); // a factory method that returns a RocksDB instance try (final RocksDB rdb = RocksDB.open(new Options(), path)) { final WriteBatch wb = new WriteBatch(); for (Map.Entry entry : results.entrySet()) { byte[] key = entry.getKey().getBytes(UTF_8); long resultValue = entry.getValue().longValue(); try { final byte[] dbValue = rdb.get(key); if (dbValue != null) { // value == null if key does not exist in db. wb.put(key, longToBytes(bytesToLong(dbValue) + resultValue)); } else { wb.put(key, longToBytes(entry.getValue().longValue())); } } catch (RocksDBException e) { // TODO: error handling } } TimeWatch watch = TimeWatch.start(); rdb.write(new WriteOptions(), wb); System.out.println(String.format("Writing %d entries took: %s", wb.count(), watch.toFullTime())); } catch (RocksDBException e) { // do some error handling } } public Map getDump() throws UnsupportedEncodingException { Map dump = new HashMap<>(); RocksDB.loadLibrary(); // the Options class contains a set of configurable DB options // that determines the behaviour of the database. // a factory method that returns a RocksDB instance try (final RocksDB rdb = RocksDB.open(new Options(), path)) { try (RocksIterator it = rdb.newIterator()) { it.seekToFirst(); // it.next(); while (it.isValid()) { byte[] key = it.key(); byte[] value = it.value(); dump.put(new String(key, UTF_8), new AtomicLong(bytesToLong(value))); it.next(); } } } catch (RocksDBException e) { e.printStackTrace(); } return dump; } public void delete() { try { FileUtils.deleteDirectory(new File(path)); } catch (IOException e) { e.printStackTrace(); } } }