Project copied

src/main/java/util/db/RDB.java

@@ -0,0 +1,132 @@
+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<String, AtomicLong> 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<String, AtomicLong> 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 byte[] atomicIntToByteArray(final AtomicLong i) {
+	// 	BigInteger bigInt = BigInteger.valueOf(i.intValue());
+	//
+	// 	return bigInt.toByteArray();
+	// }
+
+	public RocksDB getDb() {
+		return db;
+	}
+
+	public Map<String, AtomicLong> getDump() throws UnsupportedEncodingException {
+		Map<String, AtomicLong> 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();
+		}
+	}
+}