[JNI] Add java api and java tests for WriteBatch and WriteOptions, add put() and remove() to RocksDB.

Summary:
* Add java api for rocksdb::WriteBatch and rocksdb::WriteOptions, which are necessary components
  for running benchmark.
* Add java test for org.rocksdb.WriteBatch and org.rocksdb.WriteOptions.
* Add remove() to org.rocksdb.RocksDB, and add put() and remove() to RocksDB which take
  org.rocksdb.WriteOptions.

Test Plan: make jtest

Reviewers: haobo, sdong, dhruba

Reviewed By: sdong

CC: leveldb

Differential Revision: https://reviews.facebook.net/D17373

java/rocksjni/options.cc

@@ -11,8 +11,10 @@
 #include <string>
 
 #include "include/org_rocksdb_Options.h"
+#include "include/org_rocksdb_WriteOptions.h"
 #include "rocksjni/portal.h"
 #include "rocksdb/db.h"
+#include "rocksdb/options.h"
 
 /*
  * Class:     org_rocksdb_Options
@@ -55,3 +57,72 @@ jboolean Java_org_rocksdb_Options_createIfMissing(
     JNIEnv* env, jobject jobj, jlong jhandle) {
   return reinterpret_cast<rocksdb::Options*>(jhandle)->create_if_missing;
 }
+
+//////////////////////////////////////////////////////////////////////////////
+// WriteOptions
+
+/*
+ * Class:     org_rocksdb_WriteOptions
+ * Method:    newWriteOptions
+ * Signature: ()V
+ */
+void Java_org_rocksdb_WriteOptions_newWriteOptions(
+    JNIEnv* env, jobject jwrite_options) {
+  rocksdb::WriteOptions* op = new rocksdb::WriteOptions();
+  rocksdb::WriteOptionsJni::setHandle(env, jwrite_options, op);
+}
+
+/*
+ * Class:     org_rocksdb_WriteOptions
+ * Method:    dispose0
+ * Signature: ()V
+ */
+void Java_org_rocksdb_WriteOptions_dispose0(
+    JNIEnv* env, jobject jwrite_options, jlong jhandle) {
+  auto write_options = reinterpret_cast<rocksdb::WriteOptions*>(jhandle);
+  delete write_options;
+
+  rocksdb::WriteOptionsJni::setHandle(env, jwrite_options, nullptr);
+}
+
+/*
+ * Class:     org_rocksdb_WriteOptions
+ * Method:    setSync
+ * Signature: (JZ)V
+ */
+void Java_org_rocksdb_WriteOptions_setSync(
+  JNIEnv* env, jobject jwrite_options, jlong jhandle, jboolean jflag) {
+  reinterpret_cast<rocksdb::WriteOptions*>(jhandle)->sync = jflag;
+}
+
+/*
+ * Class:     org_rocksdb_WriteOptions
+ * Method:    sync
+ * Signature: (J)Z
+ */
+jboolean Java_org_rocksdb_WriteOptions_sync(
+    JNIEnv* env, jobject jwrite_options, jlong jhandle) {
+  return reinterpret_cast<rocksdb::WriteOptions*>(jhandle)->sync;
+}
+
+/*
+ * Class:     org_rocksdb_WriteOptions
+ * Method:    setDisableWAL
+ * Signature: (JZ)V
+ */
+void Java_org_rocksdb_WriteOptions_setDisableWAL(
+    JNIEnv* env, jobject jwrite_options, jlong jhandle, jboolean jflag) {
+  reinterpret_cast<rocksdb::WriteOptions*>(jhandle)->disableWAL = jflag;
+}
+
+/*
+ * Class:     org_rocksdb_WriteOptions
+ * Method:    disableWAL
+ * Signature: (J)Z
+ */
+jboolean Java_org_rocksdb_WriteOptions_disableWAL(
+    JNIEnv* env, jobject jwrite_options, jlong jhandle) {
+  return reinterpret_cast<rocksdb::WriteOptions*>(jhandle)->disableWAL;
+}
+
+

java/rocksjni/portal.h

@@ -109,5 +109,66 @@ class OptionsJni {
   }
 };
 
+class WriteOptionsJni {
+ public:
+  // Get the java class id of org.rocksdb.WriteOptions.
+  static jclass getJClass(JNIEnv* env) {
+    static jclass jclazz = env->FindClass("org/rocksdb/WriteOptions");
+    assert(jclazz != nullptr);
+    return jclazz;
+  }
+
+  // Get the field id of the member variable of org.rocksdb.WriteOptions
+  // that stores the pointer to rocksdb::WriteOptions
+  static jfieldID getHandleFieldID(JNIEnv* env) {
+    static jfieldID fid = env->GetFieldID(
+        getJClass(env), "nativeHandle_", "J");
+    assert(fid != nullptr);
+    return fid;
+  }
+
+  // Get the pointer to rocksdb::WriteOptions
+  static rocksdb::WriteOptions* getHandle(JNIEnv* env, jobject jobj) {
+    return reinterpret_cast<rocksdb::WriteOptions*>(
+        env->GetLongField(jobj, getHandleFieldID(env)));
+  }
+
+  // Pass the rocksdb::WriteOptions pointer to the java side.
+  static void setHandle(JNIEnv* env, jobject jobj, rocksdb::WriteOptions* op) {
+    env->SetLongField(
+        jobj, getHandleFieldID(env),
+        reinterpret_cast<jlong>(op));
+  }
+};
+
+class WriteBatchJni {
+ public:
+  static jclass getJClass(JNIEnv* env) {
+    static jclass jclazz = env->FindClass("org/rocksdb/WriteBatch");
+    assert(jclazz != nullptr);
+    return jclazz;
+  }
+
+  static jfieldID getHandleFieldID(JNIEnv* env) {
+    static jfieldID fid = env->GetFieldID(
+        getJClass(env), "nativeHandle_", "J");
+    assert(fid != nullptr);
+    return fid;
+  }
+
+  // Get the pointer to rocksdb::WriteBatch of the specified
+  // org.rocksdb.WriteBatch.
+  static rocksdb::WriteBatch* getHandle(JNIEnv* env, jobject jwb) {
+    return reinterpret_cast<rocksdb::WriteBatch*>(
+        env->GetLongField(jwb, getHandleFieldID(env)));
+  }
+
+  // Pass the rocksdb::WriteBatch pointer to the java side.
+  static void setHandle(JNIEnv* env, jobject jwb, rocksdb::WriteBatch* wb) {
+    env->SetLongField(
+        jwb, getHandleFieldID(env),
+        reinterpret_cast<jlong>(wb));
+  }
+};
 }  // namespace rocksdb
 #endif  // JAVA_ROCKSJNI_PORTAL_H_

java/rocksjni/rocksjni.cc

@@ -15,6 +15,9 @@
 #include "rocksjni/portal.h"
 #include "rocksdb/db.h"
 
+//////////////////////////////////////////////////////////////////////////////
+// rocksdb::DB::Open
+
 void rocksdb_open_helper(
   JNIEnv* env, jobject java_db, jstring jdb_path, const rocksdb::Options& opt) {
   rocksdb::DB* db;
@@ -54,27 +57,20 @@ void Java_org_rocksdb_RocksDB_open(
   rocksdb_open_helper(env, jdb, jdb_path, *options);
 }
 
-/*
- * Class:     org_rocksdb_RocksDB
- * Method:    put
- * Signature: ([BI[BI)V
- */
-void Java_org_rocksdb_RocksDB_put(
-    JNIEnv* env, jobject jdb,
+//////////////////////////////////////////////////////////////////////////////
+// rocksdb::DB::Put
+
+void rocksdb_put_helper(
+    JNIEnv* env, rocksdb::DB* db, const rocksdb::WriteOptions& write_options,
     jbyteArray jkey, jint jkey_len,
     jbyteArray jvalue, jint jvalue_len) {
-  rocksdb::DB* db = rocksdb::RocksDBJni::getHandle(env, jdb);
 
-  jboolean isCopy;
-  jbyte* key = env->GetByteArrayElements(jkey, &isCopy);
-  jbyte* value = env->GetByteArrayElements(jvalue, &isCopy);
-  rocksdb::Slice key_slice(
-      reinterpret_cast<char*>(key), jkey_len);
-  rocksdb::Slice value_slice(
-      reinterpret_cast<char*>(value), jvalue_len);
+  jbyte* key = env->GetByteArrayElements(jkey, 0);
+  jbyte* value = env->GetByteArrayElements(jvalue, 0);
+  rocksdb::Slice key_slice(reinterpret_cast<char*>(key), jkey_len);
+  rocksdb::Slice value_slice(reinterpret_cast<char*>(value), jvalue_len);
 
-  rocksdb::Status s = db->Put(
-      rocksdb::WriteOptions(), key_slice, value_slice);
+  rocksdb::Status s = db->Put(write_options, key_slice, value_slice);
 
   // trigger java unref on key and value.
   // by passing JNI_ABORT, it will simply release the reference without
@@ -88,6 +84,45 @@ void Java_org_rocksdb_RocksDB_put(
   rocksdb::RocksDBExceptionJni::ThrowNew(env, s);
 }
 
+/*
+ * Class:     org_rocksdb_RocksDB
+ * Method:    put
+ * Signature: ([BI[BI)V
+ */
+void Java_org_rocksdb_RocksDB_put___3BI_3BI(
+    JNIEnv* env, jobject jdb,
+    jbyteArray jkey, jint jkey_len,
+    jbyteArray jvalue, jint jvalue_len) {
+  rocksdb::DB* db = rocksdb::RocksDBJni::getHandle(env, jdb);
+  static const rocksdb::WriteOptions default_write_options =
+      rocksdb::WriteOptions();
+
+  rocksdb_put_helper(env, db, default_write_options,
+                     jkey, jkey_len,
+                     jvalue, jvalue_len);
+}
+
+/*
+ * Class:     org_rocksdb_RocksDB
+ * Method:    put
+ * Signature: (J[BI[BI)V
+ */
+void Java_org_rocksdb_RocksDB_put__J_3BI_3BI(
+    JNIEnv* env, jobject jdb, jlong jwrite_options_handle,
+    jbyteArray jkey, jint jkey_len,
+    jbyteArray jvalue, jint jvalue_len) {
+  rocksdb::DB* db = rocksdb::RocksDBJni::getHandle(env, jdb);
+  auto write_options = reinterpret_cast<rocksdb::WriteOptions*>(
+      jwrite_options_handle);
+
+  rocksdb_put_helper(env, db, *write_options,
+                     jkey, jkey_len,
+                     jvalue, jvalue_len);
+}
+
+//////////////////////////////////////////////////////////////////////////////
+// rocksdb::DB::Get
+
 /*
  * Class:     org_rocksdb_RocksDB
  * Method:    get
@@ -184,6 +219,59 @@ jint Java_org_rocksdb_RocksDB_get___3BI_3BI(
   return static_cast<jint>(cvalue_len);
 }
 
+//////////////////////////////////////////////////////////////////////////////
+// rocksdb::DB::Delete()
+void rocksdb_remove_helper(
+    JNIEnv* env, rocksdb::DB* db, const rocksdb::WriteOptions& write_options,
+    jbyteArray jkey, jint jkey_len) {
+  jbyte* key = env->GetByteArrayElements(jkey, 0);
+  rocksdb::Slice key_slice(reinterpret_cast<char*>(key), jkey_len);
+
+  rocksdb::Status s = db->Delete(write_options, key_slice);
+
+  // trigger java unref on key and value.
+  // by passing JNI_ABORT, it will simply release the reference without
+  // copying the result back to the java byte array.
+  env->ReleaseByteArrayElements(jkey, key, JNI_ABORT);
+
+  if (!s.ok()) {
+    rocksdb::RocksDBExceptionJni::ThrowNew(env, s);
+  }
+  return;
+}
+
+
+/*
+ * Class:     org_rocksdb_RocksDB
+ * Method:    remove
+ * Signature: ([BI)V
+ */
+void Java_org_rocksdb_RocksDB_remove___3BI(
+    JNIEnv* env, jobject jdb, jbyteArray jkey, jint jkey_len) {
+  rocksdb::DB* db = rocksdb::RocksDBJni::getHandle(env, jdb);
+  static const rocksdb::WriteOptions default_write_options =
+      rocksdb::WriteOptions();
+
+  rocksdb_remove_helper(env, db, default_write_options, jkey, jkey_len);
+}
+
+/*
+ * Class:     org_rocksdb_RocksDB
+ * Method:    remove
+ * Signature: (J[BI)V
+ */
+void Java_org_rocksdb_RocksDB_remove__J_3BI(
+    JNIEnv* env, jobject jdb, jlong jwrite_options,
+    jbyteArray jkey, jint jkey_len) {
+  rocksdb::DB* db = rocksdb::RocksDBJni::getHandle(env, jdb);
+  auto write_options = reinterpret_cast<rocksdb::WriteOptions*>(jwrite_options);
+
+  rocksdb_remove_helper(env, db, *write_options, jkey, jkey_len);
+}
+
+//////////////////////////////////////////////////////////////////////////////
+// rocksdb::DB::~DB()
+
 /*
  * Class:     org_rocksdb_RocksDB
  * Method:    close0
@@ -192,8 +280,8 @@ jint Java_org_rocksdb_RocksDB_get___3BI_3BI(
 void Java_org_rocksdb_RocksDB_close0(
     JNIEnv* env, jobject java_db) {
   rocksdb::DB* db = rocksdb::RocksDBJni::getHandle(env, java_db);
+  assert(db != nullptr);
   delete db;
-  db = nullptr;
 
-  rocksdb::RocksDBJni::setHandle(env, java_db, db);
+  rocksdb::RocksDBJni::setHandle(env, java_db, nullptr);
 }

java/rocksjni/write_batch.cc

@@ -0,0 +1,263 @@
+// Copyright (c) 2014, Facebook, Inc.  All rights reserved.
+// This source code is licensed under the BSD-style license found in the
+// LICENSE file in the root directory of this source tree. An additional grant
+// of patent rights can be found in the PATENTS file in the same directory.
+//
+// This file implements the "bridge" between Java and C++ and enables
+// calling c++ rocksdb::WriteBatch methods from Java side.
+#include <memory>
+
+#include "include/org_rocksdb_WriteBatch.h"
+#include "include/org_rocksdb_WriteBatchInternal.h"
+#include "include/org_rocksdb_WriteBatchTest.h"
+#include "rocksjni/portal.h"
+#include "rocksdb/db.h"
+#include "db/memtable.h"
+#include "rocksdb/write_batch.h"
+#include "db/write_batch_internal.h"
+#include "rocksdb/env.h"
+#include "rocksdb/memtablerep.h"
+#include "util/logging.h"
+#include "util/testharness.h"
+
+/*
+ * Class:     org_rocksdb_WriteBatch
+ * Method:    newWriteBatch
+ * Signature: (I)V
+ */
+void Java_org_rocksdb_WriteBatch_newWriteBatch(
+    JNIEnv* env, jobject jobj, jint jreserved_bytes) {
+  rocksdb::WriteBatch* wb = new rocksdb::WriteBatch(
+      static_cast<size_t>(jreserved_bytes));
+
+  rocksdb::WriteBatchJni::setHandle(env, jobj, wb);
+}
+
+/*
+ * Class:     org_rocksdb_WriteBatch
+ * Method:    count
+ * Signature: ()I
+ */
+jint Java_org_rocksdb_WriteBatch_count(JNIEnv* env, jobject jobj) {
+  rocksdb::WriteBatch* wb = rocksdb::WriteBatchJni::getHandle(env, jobj);
+  assert(wb != nullptr);
+
+  return static_cast<jint>(wb->Count());
+}
+
+/*
+ * Class:     org_rocksdb_WriteBatch
+ * Method:    clear
+ * Signature: ()V
+ */
+void Java_org_rocksdb_WriteBatch_clear(JNIEnv* env, jobject jobj) {
+  rocksdb::WriteBatch* wb = rocksdb::WriteBatchJni::getHandle(env, jobj);
+  assert(wb != nullptr);
+
+  wb->Clear();
+}
+
+/*
+ * Class:     org_rocksdb_WriteBatch
+ * Method:    put
+ * Signature: ([BI[BI)V
+ */
+void Java_org_rocksdb_WriteBatch_put(
+    JNIEnv* env, jobject jobj,
+    jbyteArray jkey, jint jkey_len,
+    jbyteArray jvalue, jint jvalue_len) {
+  rocksdb::WriteBatch* wb = rocksdb::WriteBatchJni::getHandle(env, jobj);
+  assert(wb != nullptr);
+
+  jbyte* key = env->GetByteArrayElements(jkey, nullptr);
+  jbyte* value = env->GetByteArrayElements(jvalue, nullptr);
+  rocksdb::Slice key_slice(reinterpret_cast<char*>(key), jkey_len);
+  rocksdb::Slice value_slice(reinterpret_cast<char*>(value), jvalue_len);
+  wb->Put(key_slice, value_slice);
+  env->ReleaseByteArrayElements(jkey, key, JNI_ABORT);
+  env->ReleaseByteArrayElements(jvalue, value, JNI_ABORT);
+}
+
+/*
+ * Class:     org_rocksdb_WriteBatch
+ * Method:    merge
+ * Signature: ([BI[BI)V
+ */
+JNIEXPORT void JNICALL Java_org_rocksdb_WriteBatch_merge(
+    JNIEnv* env, jobject jobj,
+    jbyteArray jkey, jint jkey_len,
+    jbyteArray jvalue, jint jvalue_len) {
+  rocksdb::WriteBatch* wb = rocksdb::WriteBatchJni::getHandle(env, jobj);
+  assert(wb != nullptr);
+
+  jbyte* key = env->GetByteArrayElements(jkey, nullptr);
+  jbyte* value = env->GetByteArrayElements(jvalue, nullptr);
+  rocksdb::Slice key_slice(reinterpret_cast<char*>(key), jkey_len);
+  rocksdb::Slice value_slice(reinterpret_cast<char*>(value), jvalue_len);
+  wb->Merge(key_slice, value_slice);
+  env->ReleaseByteArrayElements(jkey, key, JNI_ABORT);
+  env->ReleaseByteArrayElements(jvalue, value, JNI_ABORT);
+}
+
+/*
+ * Class:     org_rocksdb_WriteBatch
+ * Method:    remove
+ * Signature: ([BI)V
+ */
+JNIEXPORT void JNICALL Java_org_rocksdb_WriteBatch_remove(
+    JNIEnv* env, jobject jobj,
+    jbyteArray jkey, jint jkey_len) {
+  rocksdb::WriteBatch* wb = rocksdb::WriteBatchJni::getHandle(env, jobj);
+  assert(wb != nullptr);
+
+  jbyte* key = env->GetByteArrayElements(jkey, nullptr);
+  rocksdb::Slice key_slice(reinterpret_cast<char*>(key), jkey_len);
+  wb->Delete(key_slice);
+  env->ReleaseByteArrayElements(jkey, key, JNI_ABORT);
+}
+
+/*
+ * Class:     org_rocksdb_WriteBatch
+ * Method:    putLogData
+ * Signature: ([BI)V
+ */
+void Java_org_rocksdb_WriteBatch_putLogData(
+    JNIEnv* env, jobject jobj, jbyteArray jblob, jint jblob_len) {
+  rocksdb::WriteBatch* wb = rocksdb::WriteBatchJni::getHandle(env, jobj);
+  assert(wb != nullptr);
+
+  jbyte* blob = env->GetByteArrayElements(jblob, nullptr);
+  rocksdb::Slice blob_slice(reinterpret_cast<char*>(blob), jblob_len);
+  wb->PutLogData(blob_slice);
+  env->ReleaseByteArrayElements(jblob, blob, JNI_ABORT);
+}
+
+/*
+ * Class:     org_rocksdb_WriteBatch
+ * Method:    dispose0
+ * Signature: ()V
+ */
+void Java_org_rocksdb_WriteBatch_dispose0(JNIEnv* env, jobject jobj) {
+  rocksdb::WriteBatch* wb = rocksdb::WriteBatchJni::getHandle(env, jobj);
+  assert(wb != nullptr);
+  delete wb;
+
+  rocksdb::WriteBatchJni::setHandle(env, jobj, nullptr);
+}
+
+/*
+ * Class:     org_rocksdb_WriteBatchInternal
+ * Method:    setSequence
+ * Signature: (Lorg/rocksdb/WriteBatch;J)V
+ */
+void Java_org_rocksdb_WriteBatchInternal_setSequence(
+    JNIEnv* env, jclass jclazz, jobject jobj, jlong jsn) {
+  rocksdb::WriteBatch* wb = rocksdb::WriteBatchJni::getHandle(env, jobj);
+  assert(wb != nullptr);
+
+  rocksdb::WriteBatchInternal::SetSequence(
+      wb, static_cast<rocksdb::SequenceNumber>(jsn));
+}
+
+/*
+ * Class:     org_rocksdb_WriteBatchInternal
+ * Method:    sequence
+ * Signature: (Lorg/rocksdb/WriteBatch;)J
+ */
+jlong Java_org_rocksdb_WriteBatchInternal_sequence(
+    JNIEnv* env, jclass jclazz, jobject jobj) {
+  rocksdb::WriteBatch* wb = rocksdb::WriteBatchJni::getHandle(env, jobj);
+  assert(wb != nullptr);
+
+  return static_cast<jlong>(rocksdb::WriteBatchInternal::Sequence(wb));
+}
+
+/*
+ * Class:     org_rocksdb_WriteBatchInternal
+ * Method:    append
+ * Signature: (Lorg/rocksdb/WriteBatch;Lorg/rocksdb/WriteBatch;)V
+ */
+void Java_org_rocksdb_WriteBatchInternal_append(
+    JNIEnv* env, jclass jclazz, jobject jwb1, jobject jwb2) {
+  rocksdb::WriteBatch* wb1 = rocksdb::WriteBatchJni::getHandle(env, jwb1);
+  assert(wb1 != nullptr);
+  rocksdb::WriteBatch* wb2 = rocksdb::WriteBatchJni::getHandle(env, jwb2);
+  assert(wb2 != nullptr);
+
+  rocksdb::WriteBatchInternal::Append(wb1, wb2);
+}
+
+/*
+ * Class:     org_rocksdb_WriteBatchTest
+ * Method:    getContents
+ * Signature: (Lorg/rocksdb/WriteBatch;)[B
+ */
+jbyteArray Java_org_rocksdb_WriteBatchTest_getContents(
+    JNIEnv* env, jclass jclazz, jobject jobj) {
+  rocksdb::WriteBatch* b = rocksdb::WriteBatchJni::getHandle(env, jobj);
+  assert(b != nullptr);
+
+  // todo: Currently the following code is directly copied from
+  // db/write_bench_test.cc.  It could be implemented in java once
+  // all the necessary components can be accessed via jni api.
+
+  rocksdb::InternalKeyComparator cmp(rocksdb::BytewiseComparator());
+  auto factory = std::make_shared<rocksdb::SkipListFactory>();
+  rocksdb::Options options;
+  options.memtable_factory = factory;
+  rocksdb::MemTable* mem = new rocksdb::MemTable(cmp, options);
+  mem->Ref();
+  std::string state;
+  rocksdb::Status s = rocksdb::WriteBatchInternal::InsertInto(b, mem, &options);
+  int count = 0;
+  rocksdb::Iterator* iter = mem->NewIterator();
+  for (iter->SeekToFirst(); iter->Valid(); iter->Next()) {
+    rocksdb::ParsedInternalKey ikey;
+    memset(reinterpret_cast<void*>(&ikey), 0, sizeof(ikey));
+    ASSERT_TRUE(rocksdb::ParseInternalKey(iter->key(), &ikey));
+    switch (ikey.type) {
+      case rocksdb::kTypeValue:
+        state.append("Put(");
+        state.append(ikey.user_key.ToString());
+        state.append(", ");
+        state.append(iter->value().ToString());
+        state.append(")");
+        count++;
+        break;
+      case rocksdb::kTypeMerge:
+        state.append("Merge(");
+        state.append(ikey.user_key.ToString());
+        state.append(", ");
+        state.append(iter->value().ToString());
+        state.append(")");
+        count++;
+        break;
+      case rocksdb::kTypeDeletion:
+        state.append("Delete(");
+        state.append(ikey.user_key.ToString());
+        state.append(")");
+        count++;
+        break;
+      default:
+        assert(false);
+        break;
+    }
+    state.append("@");
+    state.append(rocksdb::NumberToString(ikey.sequence));
+  }
+  delete iter;
+  if (!s.ok()) {
+    state.append(s.ToString());
+  } else if (count != rocksdb::WriteBatchInternal::Count(b)) {
+    state.append("CountMismatch()");
+  }
+  delete mem->Unref();
+
+  jbyteArray jstate = env->NewByteArray(state.size());
+  env->SetByteArrayRegion(
+      jstate, 0, state.size(),
+      reinterpret_cast<const jbyte*>(state.c_str()));
+
+  return jstate;
+}
+