Addresed comments from code review https://reviews.facebook.net/D27567

java/org/rocksdb/WriteBatch.java

@@ -105,8 +105,13 @@ public class WriteBatch extends RocksObject {
 
   /**
    * Support for iterating over the contents of a batch.
+   *
+   * @param handler A handler that is called back for each
+   *                update present in the batch
+   *
+   * @throws RocksDBException If we cannot iterate over the batch
    */
-  public void iterate(Handler handler) {
+  public void iterate(Handler handler) throws RocksDBException {
     iterate(handler.nativeHandle_);
   }
 
@@ -138,7 +143,7 @@ public class WriteBatch extends RocksObject {
   private native void remove(byte[] key, int keyLen,
                             long cfHandle);
   private native void putLogData(byte[] blob, int blobLen);
-  private native void iterate(long handlerHandle);
+  private native void iterate(long handlerHandle) throws RocksDBException;
   private native void disposeInternal(long handle);
 
   /**
@@ -157,7 +162,7 @@ public class WriteBatch extends RocksObject {
 
     /**
      * shouldContinue is called by the underlying iterator
-     * (WriteBatch::Iterate.If it returns false,
+     * WriteBatch::Iterate. If it returns false,
      * iteration is halted. Otherwise, it continues
      * iterating. The default implementation always
      * returns true.

java/rocksjni/write_batch.cc

@@ -317,7 +317,7 @@ void Java_org_rocksdb_WriteBatch_00024Handler_createNewHandler0(
  */
 void Java_org_rocksdb_WriteBatch_00024Handler_disposeInternal(
     JNIEnv* env, jobject jobj, jlong handle) {
-    delete reinterpret_cast<rocksdb::WriteBatchHandlerJniCallback*>(handle);
+  delete reinterpret_cast<rocksdb::WriteBatchHandlerJniCallback*>(handle);
 }
 
 /*

java/rocksjni/writebatchhandlerjnicallback.cc

@@ -11,12 +11,8 @@
 
 namespace rocksdb {
 WriteBatchHandlerJniCallback::WriteBatchHandlerJniCallback(
-    JNIEnv* env, jobject jWriteBatchHandler) {
+    JNIEnv* env, jobject jWriteBatchHandler)
-
+    : m_env(env) {
-  // Note: WriteBatchHandler methods may be accessed by multiple threads,
-  // so we ref the jvm not the env
-  const jint rs = env->GetJavaVM(&m_jvm);
-  assert(rs == JNI_OK);
 
   // Note: we want to access the Java WriteBatchHandler instance
   // across multiple method calls, so we create a global ref
@@ -29,70 +25,80 @@ WriteBatchHandlerJniCallback::WriteBatchHandlerJniCallback(
   m_jContinueMethodId = WriteBatchHandlerJni::getContinueMethodId(env);
 }
 
-/**
- * Attach/Get a JNIEnv for the current native thread
- */
-JNIEnv* WriteBatchHandlerJniCallback::getJniEnv() const {
-  JNIEnv *env;
-  jint rs = m_jvm->AttachCurrentThread(reinterpret_cast<void **>(&env), NULL);
-  assert(rs == JNI_OK);
-  return env;
-}
-
 void WriteBatchHandlerJniCallback::Put(const Slice& key, const Slice& value) {
-  getJniEnv()->CallVoidMethod(
+  const jbyteArray j_key = sliceToJArray(key);
+  const jbyteArray j_value = sliceToJArray(value);
+
+  m_env->CallVoidMethod(
       m_jWriteBatchHandler,
       m_jPutMethodId,
-      sliceToJArray(key),
+      j_key,
-      sliceToJArray(value));
+      j_value);
+
+  m_env->DeleteLocalRef(j_value);
+  m_env->DeleteLocalRef(j_key);
 }
 
 void WriteBatchHandlerJniCallback::Merge(const Slice& key, const Slice& value) {
-  getJniEnv()->CallVoidMethod(
+  const jbyteArray j_key = sliceToJArray(key);
+  const jbyteArray j_value = sliceToJArray(value);
+
+  m_env->CallVoidMethod(
       m_jWriteBatchHandler,
       m_jMergeMethodId,
-      sliceToJArray(key),
+      j_key,
-      sliceToJArray(value));
+      j_value);
+
+  m_env->DeleteLocalRef(j_value);
+  m_env->DeleteLocalRef(j_key);
 }
 
 void WriteBatchHandlerJniCallback::Delete(const Slice& key) {
-  getJniEnv()->CallVoidMethod(
+  const jbyteArray j_key = sliceToJArray(key);
+
+  m_env->CallVoidMethod(
       m_jWriteBatchHandler,
       m_jDeleteMethodId,
-      sliceToJArray(key));
+      j_key);
+
+  m_env->DeleteLocalRef(j_key);
 }
 
 void WriteBatchHandlerJniCallback::LogData(const Slice& blob) {
-  getJniEnv()->CallVoidMethod(
+  const jbyteArray j_blob = sliceToJArray(blob);
+
+  m_env->CallVoidMethod(
       m_jWriteBatchHandler,
       m_jLogDataMethodId,
-      sliceToJArray(blob));
+      j_blob);
+
+  m_env->DeleteLocalRef(j_blob);
 }
 
 bool WriteBatchHandlerJniCallback::Continue() {
-  jboolean jContinue = getJniEnv()->CallBooleanMethod(
+  jboolean jContinue = m_env->CallBooleanMethod(
       m_jWriteBatchHandler,
       m_jContinueMethodId);
 
   return static_cast<bool>(jContinue == JNI_TRUE);
 }
 
+/*
+ * Creates a Java Byte Array from the data in a Slice
+ *
+ * When calling this function
+ * you must remember to call env->DeleteLocalRef
+ * on the result after you have finished with it
+ */
 jbyteArray WriteBatchHandlerJniCallback::sliceToJArray(const Slice& s) {
-  jbyteArray ja = getJniEnv()->NewByteArray(s.size());
+  jbyteArray ja = m_env->NewByteArray(s.size());
-  getJniEnv()->SetByteArrayRegion(
+  m_env->SetByteArrayRegion(
       ja, 0, s.size(),
       reinterpret_cast<const jbyte*>(s.data()));
   return ja;
 }
 
 WriteBatchHandlerJniCallback::~WriteBatchHandlerJniCallback() {
-  JNIEnv* m_env = getJniEnv();
-
   m_env->DeleteGlobalRef(m_jWriteBatchHandler);
-
-  // Note: do not need to explicitly detach, as this function is effectively
-  // called from the Java class's disposeInternal method, and so already
-  // has an attached thread, getJniEnv above is just a no-op Attach to get
-  // the env jvm->DetachCurrentThread();
 }
 }  // namespace rocksdb

java/rocksjni/writebatchhandlerjnicallback.h

@@ -17,8 +17,8 @@ namespace rocksdb {
  * This class acts as a bridge between C++
  * and Java. The methods in this class will be
  * called back from the RocksDB storage engine (C++)
- * we then callback to the appropriate Java method
+ * which calls the appropriate Java method.
- * this enables Write Batch Handlers to be implemented in Java.
+ * This enables Write Batch Handlers to be implemented in Java.
  */
 class WriteBatchHandlerJniCallback : public WriteBatch::Handler {
  public:
@@ -32,9 +32,8 @@ class WriteBatchHandlerJniCallback : public WriteBatch::Handler {
     bool Continue();
 
  private:
-    JavaVM* m_jvm;
+    JNIEnv* m_env;
     jobject m_jWriteBatchHandler;
-    JNIEnv* getJniEnv() const;
     jbyteArray sliceToJArray(const Slice& s);
     jmethodID m_jPutMethodId;
     jmethodID m_jMergeMethodId;