Fix a memory leak of Slice objects from org.rocksdb.WBWIRocksIterator#entry1

java/rocksjni/comparatorjnicallback.cc

@@ -60,8 +60,8 @@ int BaseComparatorJniCallback::Compare(const Slice& a, const Slice& b) const {
   // performance.
   mtx_compare->Lock();
 
-  AbstractSliceJni::setHandle(m_env, m_jSliceA, &a);
-  AbstractSliceJni::setHandle(m_env, m_jSliceB, &b);
+  AbstractSliceJni::setHandle(m_env, m_jSliceA, &a, JNI_FALSE);
+  AbstractSliceJni::setHandle(m_env, m_jSliceB, &b, JNI_FALSE);
   jint result =
     m_env->CallIntMethod(m_jComparator, m_jCompareMethodId, m_jSliceA,
       m_jSliceB);
@@ -89,7 +89,7 @@ void BaseComparatorJniCallback::FindShortestSeparator(
   // performance.
   mtx_findShortestSeparator->Lock();
 
-  AbstractSliceJni::setHandle(m_env, m_jSliceLimit, &limit);
+  AbstractSliceJni::setHandle(m_env, m_jSliceLimit, &limit, JNI_FALSE);
   jstring jsResultStart =
     (jstring)m_env->CallObjectMethod(m_jComparator,
       m_jFindShortestSeparatorMethodId, jsStart, m_jSliceLimit);

java/rocksjni/portal.h

@@ -49,25 +49,22 @@ template<class PTR, class DERIVED> class RocksDBNativeClass {
     assert(jclazz != nullptr);
     return jclazz;
   }
-
-  // Get the field id of the member variable to store
-  // the ptr
-  static jfieldID getHandleFieldID(JNIEnv* env) {
-    static jfieldID fid = env->GetFieldID(
-        DERIVED::getJClass(env), "nativeHandle_", "J");
-    assert(fid != nullptr);
-    return fid;
-  }
 };
 
 // Native class template for sub-classes of RocksMutableObject
 template<class PTR, class DERIVED> class NativeRocksMutableObject : public RocksDBNativeClass<PTR, DERIVED> {
  public:
+
+  static jmethodID getSetNativeHandleMethod(JNIEnv* env) {
+    static jmethodID mid = env->GetMethodID(
+        DERIVED::getJClass(env), "setNativeHandle", "(JZ)V");
+    assert(mid != nullptr);
+    return mid;
+  }
+
   // Pass the pointer to the java side.
-  static void setHandle(JNIEnv* env, jobject jdb, PTR ptr) {
-    env->SetLongField(
-        jdb, RocksDBNativeClass<PTR, DERIVED>::getHandleFieldID(env),
-        reinterpret_cast<jlong>(ptr));
+  static void setHandle(JNIEnv* env, jobject jobj, PTR ptr, jboolean java_owns_handle) {
+    env->CallVoidMethod(jobj, getSetNativeHandleMethod(env), reinterpret_cast<jlong>(ptr), java_owns_handle);
   }
 };
 
@@ -647,67 +644,6 @@ class WriteEntryJni {
       assert(jclazz != nullptr);
       return jclazz;
     }
-
-    static void setWriteType(JNIEnv* env, jobject jwrite_entry,
-        WriteType write_type) {
-      jobject jwrite_type;
-      switch (write_type) {
-        case kPutRecord:
-          jwrite_type = WriteTypeJni::PUT(env);
-          break;
-
-        case kMergeRecord:
-          jwrite_type = WriteTypeJni::MERGE(env);
-          break;
-
-        case kDeleteRecord:
-          jwrite_type = WriteTypeJni::DELETE(env);
-          break;
-
-        case kLogDataRecord:
-          jwrite_type = WriteTypeJni::LOG(env);
-          break;
-
-        default:
-          jwrite_type = nullptr;
-      }
-      assert(jwrite_type != nullptr);
-      env->SetObjectField(jwrite_entry, getWriteTypeField(env), jwrite_type);
-    }
-
-    static void setKey(JNIEnv* env, jobject jwrite_entry,
-        const rocksdb::Slice* slice) {
-      jobject jkey = env->GetObjectField(jwrite_entry, getKeyField(env));
-      AbstractSliceJni::setHandle(env, jkey, slice);
-    }
-
-    static void setValue(JNIEnv* env, jobject jwrite_entry,
-        const rocksdb::Slice* slice) {
-      jobject jvalue = env->GetObjectField(jwrite_entry, getValueField(env));
-      AbstractSliceJni::setHandle(env, jvalue, slice);
-    }
-
- private:
-    static jfieldID getWriteTypeField(JNIEnv* env) {
-      static jfieldID fid = env->GetFieldID(
-          getJClass(env), "type", "Lorg/rocksdb/WBWIRocksIterator$WriteType;");
-        assert(fid != nullptr);
-        return fid;
-    }
-
-    static jfieldID getKeyField(JNIEnv* env) {
-      static jfieldID fid = env->GetFieldID(
-          getJClass(env), "key", "Lorg/rocksdb/DirectSlice;");
-      assert(fid != nullptr);
-      return fid;
-    }
-
-    static jfieldID getValueField(JNIEnv* env) {
-      static jfieldID fid = env->GetFieldID(
-          getJClass(env), "value", "Lorg/rocksdb/DirectSlice;");
-      assert(fid != nullptr);
-      return fid;
-    }
 };
 
 class InfoLogLevelJni {

java/rocksjni/write_batch_with_index.cc

@@ -353,27 +353,57 @@ void Java_org_rocksdb_WBWIRocksIterator_status0(
 /*
  * Class:     org_rocksdb_WBWIRocksIterator
  * Method:    entry1
- * Signature: (JLorg/rocksdb/WBWIRocksIterator/WriteEntry;)V
+ * Signature: (J)[J
  */
-void Java_org_rocksdb_WBWIRocksIterator_entry1(
-    JNIEnv* env, jobject jobj, jlong handle, jobject jwrite_entry) {
+jlongArray Java_org_rocksdb_WBWIRocksIterator_entry1(
+    JNIEnv* env, jobject jobj, jlong handle) {
   auto* it = reinterpret_cast<rocksdb::WBWIIterator*>(handle);
   const rocksdb::WriteEntry& we = it->Entry();
-  jobject jwe = rocksdb::WBWIRocksIteratorJni::getWriteEntry(env, jobj);
-  rocksdb::WriteEntryJni::setWriteType(env, jwe, we.type);
 
+  jlong results[3];
+
+  //set the type of the write entry
+  switch (we.type) {
+    case rocksdb::kPutRecord:
+      results[0] = 0x1;
+      break;
+
+    case rocksdb::kMergeRecord:
+      results[0] = 0x2;
+      break;
+
+    case rocksdb::kDeleteRecord:
+      results[0] = 0x4;
+      break;
+
+    case rocksdb::kLogDataRecord:
+      results[0] = 0x8;
+      break;
+
+    default:
+      results[0] = 0x0;
+  }
+
+  //TODO(AR) do we leak buf and value_buf?
+
+  //set the pointer to the key slice
   char* buf = new char[we.key.size()];
   memcpy(buf, we.key.data(), we.key.size());
   auto* key_slice = new rocksdb::Slice(buf, we.key.size());
-  rocksdb::WriteEntryJni::setKey(env, jwe, key_slice);
+  results[1] = reinterpret_cast<jlong>(key_slice);
 
+  //set the pointer to the value slice
   if (we.type == rocksdb::kDeleteRecord || we.type == rocksdb::kLogDataRecord) {
     // set native handle of value slice to null if no value available
-    rocksdb::WriteEntryJni::setValue(env, jwe, nullptr);
+    results[2] = 0;
   } else {
     char* value_buf = new char[we.value.size()];
     memcpy(value_buf, we.value.data(), we.value.size());
     auto* value_slice = new rocksdb::Slice(value_buf, we.value.size());
-    rocksdb::WriteEntryJni::setValue(env, jwe, value_slice);
+    results[2] = reinterpret_cast<jlong>(value_slice);
   }
+
+  jlongArray jresults = env->NewLongArray(3);
+  env->SetLongArrayRegion(jresults, 0, 3, results);
+  return jresults;
 }

java/src/main/java/org/rocksdb/AbstractSlice.java

@@ -42,8 +42,7 @@ abstract class AbstractSlice<T> extends RocksMutableObject {
    *   @see org.rocksdb.AbstractSlice#data0(long)
    */
   public T data() {
-    assert (isOwningHandle());
-    return data0(nativeHandle_);
+    return data0(getNativeHandle());
   }
 
   /**
@@ -64,8 +63,7 @@ abstract class AbstractSlice<T> extends RocksMutableObject {
    * @return The length in bytes.
    */
   public int size() {
-    assert (isOwningHandle());
-    return size0(nativeHandle_);
+    return size0(getNativeHandle());
   }
 
   /**
@@ -75,8 +73,7 @@ abstract class AbstractSlice<T> extends RocksMutableObject {
    * @return true if there is no data, false otherwise.
    */
   public boolean empty() {
-    assert (isOwningHandle());
-    return empty0(nativeHandle_);
+    return empty0(getNativeHandle());
   }
 
   /**
@@ -88,8 +85,7 @@ abstract class AbstractSlice<T> extends RocksMutableObject {
    * @return The string representation of the data.
    */
   public String toString(final boolean hex) {
-    assert (isOwningHandle());
-    return toString0(nativeHandle_, hex);
+    return toString0(getNativeHandle(), hex);
   }
 
   @Override
@@ -109,8 +105,15 @@ abstract class AbstractSlice<T> extends RocksMutableObject {
    */
   public int compare(final AbstractSlice<?> other) {
     assert (other != null);
-    assert (isOwningHandle());
-    return compare0(nativeHandle_, other.nativeHandle_);
+    if(!isOwningHandle()) {
+      return other.isOwningHandle() ? -1 : 0;
+    } else {
+      if(!other.isOwningHandle()) {
+        return 1;
+      } else {
+        return compare0(getNativeHandle(), other.getNativeHandle());
+      }
+    }
   }
 
   @Override
@@ -149,8 +152,7 @@ abstract class AbstractSlice<T> extends RocksMutableObject {
    */
   public boolean startsWith(final AbstractSlice<?> prefix) {
     if (prefix != null) {
-      assert (isOwningHandle());
-      return startsWith0(nativeHandle_, prefix.nativeHandle_);
+      return startsWith0(getNativeHandle(), prefix.getNativeHandle());
     } else {
       return false;
     }

java/src/main/java/org/rocksdb/DirectSlice.java

@@ -24,10 +24,11 @@ public class DirectSlice extends AbstractSlice<ByteBuffer> {
    * at creation time.
    *
    * Note: You should be aware that it is intentionally marked as
-   * package-private. This is so that developers cannot construct their own default
-   * DirectSlice objects (at present). As developers cannot construct their own
-   * DirectSlice objects through this, they are not creating underlying C++
-   * DirectSlice objects, and so there is nothing to free (dispose) from Java.
+   * package-private. This is so that developers cannot construct their own
+   * default DirectSlice objects (at present). As developers cannot construct
+   * their own DirectSlice objects through this, they are not creating
+   * underlying C++ DirectSlice objects, and so there is nothing to free
+   * (dispose) from Java.
    */
   DirectSlice() {
     super();
@@ -68,7 +69,8 @@ public class DirectSlice extends AbstractSlice<ByteBuffer> {
   }
 
   private static ByteBuffer ensureDirect(final ByteBuffer data) {
-    //TODO(AR) consider throwing a checked exception, as if it's not direct this can SIGSEGV
+    // TODO(AR) consider throwing a checked exception, as if it's not direct
+    // this can SIGSEGV
     assert(data.isDirect());
     return data;
   }
@@ -82,16 +84,14 @@ public class DirectSlice extends AbstractSlice<ByteBuffer> {
    * @return the requested byte
    */
   public byte get(int offset) {
-    assert (isOwningHandle());
-    return get0(nativeHandle_, offset);
+    return get0(getNativeHandle(), offset);
   }
 
   /**
    * Clears the backing slice
    */
   public void clear() {
-    assert (isOwningHandle());
-    clear0(nativeHandle_);
+    clear0(getNativeHandle());
   }
 
   /**
@@ -102,8 +102,7 @@ public class DirectSlice extends AbstractSlice<ByteBuffer> {
    * @param n The number of bytes to drop
    */
   public void removePrefix(final int n) {
-    assert (isOwningHandle());
-    removePrefix0(nativeHandle_, n);
+    removePrefix0(getNativeHandle(), n);
   }
 
   private native static long createNewDirectSlice0(final ByteBuffer data,

java/src/main/java/org/rocksdb/RocksMutableObject.java

@@ -1,30 +1,47 @@
 package org.rocksdb;
 
-public abstract class RocksMutableObject extends NativeReference {
+public abstract class RocksMutableObject /*extends NativeReference*/ {
 
-    private final boolean shouldOwnHandle;
-    protected volatile long nativeHandle_;
+    private long nativeHandle_;
+    private boolean owningHandle_;
 
     protected RocksMutableObject() {
-        super(false);
-        this.shouldOwnHandle = false;
     }
 
     protected RocksMutableObject(final long nativeHandle) {
-        super(true);
-        this.shouldOwnHandle = true;
         this.nativeHandle_ = nativeHandle;
+        this.owningHandle_ = true;
     }
 
-    @Override
-    public boolean isOwningHandle() {
-        return ((!shouldOwnHandle) || super.isOwningHandle()) && nativeHandle_ != 0;
+    public synchronized void setNativeHandle(final long nativeHandle, final boolean owningNativeHandle) {
+        this.nativeHandle_ = nativeHandle;
+        this.owningHandle_ = owningNativeHandle;
+    }
+
+    //@Override
+    protected synchronized boolean isOwningHandle() {
+        return this.owningHandle_;
+    }
+
+    protected synchronized long getNativeHandle() {
+        assert(this.nativeHandle_ != 0);
+        return this.nativeHandle_;
+    }
+
+    public synchronized final void dispose() {
+        if(isOwningHandle()) {
+            disposeInternal();
+            this.owningHandle_ = false;
+            this.nativeHandle_ = 0;
+        }
     }
 
-    /**
-     * Deletes underlying C++ object pointer.
-     */
     @Override
+    protected void finalize() throws Throwable {
+        dispose();
+        super.finalize();
+    }
+
     protected void disposeInternal() {
         disposeInternal(nativeHandle_);
     }

java/src/main/java/org/rocksdb/Slice.java

@@ -73,8 +73,9 @@ public class Slice extends AbstractSlice<byte[]> {
    */
   @Override
   protected void disposeInternal() {
-    disposeInternalBuf(nativeHandle_);
-    super.disposeInternal();
+    final long nativeHandle = getNativeHandle();
+    disposeInternalBuf(nativeHandle);
+    super.disposeInternal(nativeHandle);
   }
 
   @Override protected final native byte[] data0(long handle);

java/src/main/java/org/rocksdb/WBWIRocksIterator.java

@@ -5,10 +5,12 @@
 
 package org.rocksdb;
 
-public class WBWIRocksIterator extends AbstractRocksIterator<WriteBatchWithIndex> {
+public class WBWIRocksIterator
+    extends AbstractRocksIterator<WriteBatchWithIndex> {
   private final WriteEntry entry = new WriteEntry();
 
-  protected WBWIRocksIterator(final WriteBatchWithIndex wbwi, final long nativeHandle) {
+  protected WBWIRocksIterator(final WriteBatchWithIndex wbwi,
+      final long nativeHandle) {
     super(wbwi, nativeHandle);
   }
 
@@ -20,12 +22,20 @@ public class WBWIRocksIterator extends AbstractRocksIterator<WriteBatchWithIndex
    * If you want to keep the WriteEntry across iterator
    * movements, you must make a copy of its data!
    *
+   * Note - This method is not thread-safe with respect to the WriteEntry
+   * as it performs a non-atomic update across the fields of the WriteEntry
+   *
    * @return The WriteEntry of the current entry
    */
   public WriteEntry entry() {
     assert(isOwningHandle());
     assert(entry != null);
-    entry1(nativeHandle_, entry);
+    final long ptrs[] = entry1(nativeHandle_);
+
+    entry.type = WriteType.fromId((byte)ptrs[0]);
+    entry.key.setNativeHandle(ptrs[1], true);
+    entry.value.setNativeHandle(ptrs[2], ptrs[2] != 0);
+
     return entry;
   }
 
@@ -38,17 +48,31 @@ public class WBWIRocksIterator extends AbstractRocksIterator<WriteBatchWithIndex
   @Override final native void seek0(long handle, byte[] target, int targetLen);
   @Override final native void status0(long handle) throws RocksDBException;
 
-  private native void entry1(long handle, WriteEntry entry);
+  private native long[] entry1(final long handle);
 
   /**
    * Enumeration of the Write operation
    * that created the record in the Write Batch
    */
   public enum WriteType {
-    PUT,
-    MERGE,
-    DELETE,
-    LOG
+    PUT((byte)0x1),
+    MERGE((byte)0x2),
+    DELETE((byte)0x4),
+    LOG((byte)0x8);
+
+    final byte id;
+    WriteType(final byte id) {
+      this.id = id;
+    }
+
+    public static WriteType fromId(final byte id) {
+      for(final WriteType wt : WriteType.values()) {
+        if(id == wt.id) {
+          return wt;
+        }
+      }
+      throw new IllegalArgumentException("No WriteType with id=" + id);
+    }
   }
 
   /**
@@ -139,8 +163,7 @@ public class WBWIRocksIterator extends AbstractRocksIterator<WriteBatchWithIndex
         final WriteEntry otherWriteEntry = (WriteEntry)other;
         return type.equals(otherWriteEntry.type)
             && key.equals(otherWriteEntry.key)
-            && (value.isOwningHandle() ? value.equals(otherWriteEntry.value)
-                : !otherWriteEntry.value.isOwningHandle());
+            && value.equals(otherWriteEntry.value);
       } else {
         return false;
       }