Java API Auto close iterators

.github/workflows/pr-jobs.yml

@@ -613,13 +613,15 @@ jobs:
     - name: PMD RocksDBJava
       run: make V=1 J=8 -j8 jpmd
     - uses: actions/[email protected]
+      if: failure()
       with:
         name: pmd-report
-        path: "${{ github.workspace }}/java/target/pmd.xml"
+        path: java/target/pmd.xml
     - uses: actions/[email protected]
+      if: failure()
       with:
         name: maven-site
-        path: "${{ github.workspace }}/java/target/site"
+        path: java/target/site
   build-linux-arm:
     if: ${{ github.repository_owner == 'facebook' }}
     runs-on:

java/CMakeLists.txt

@@ -348,6 +348,7 @@ set(JAVA_TEST_CLASSES
   src/test/java/org/rocksdb/ImportColumnFamilyTest.java
   src/test/java/org/rocksdb/InfoLogLevelTest.java
   src/test/java/org/rocksdb/IngestExternalFileOptionsTest.java
+  src/test/java/org/rocksdb/IteratorClosedDBTest.java
   src/test/java/org/rocksdb/KeyExistsTest.java
   src/test/java/org/rocksdb/KeyMayExistTest.java
   src/test/java/org/rocksdb/LRUCacheTest.java
@@ -466,6 +467,7 @@ set(JAVA_TEST_RUNNING_CLASSES
   org.rocksdb.ImportColumnFamilyTest
   org.rocksdb.InfoLogLevelTest
   org.rocksdb.IngestExternalFileOptionsTest
+  org.rocksdb.IteratorClosedDBTest
   org.rocksdb.KeyExistsTest
   org.rocksdb.KeyMayExistTest
   org.rocksdb.LRUCacheTest

java/Makefile

@@ -138,6 +138,7 @@ JAVA_TESTS = \
 	org.rocksdb.EnvOptionsTest\
 	org.rocksdb.EventListenerTest\
 	org.rocksdb.IngestExternalFileOptionsTest\
+	org.rocksdb.IteratorClosedDBTest\
 	org.rocksdb.util.IntComparatorTest\
 	org.rocksdb.util.JNIComparatorTest\
 	org.rocksdb.FilterTest\

java/src/main/java/org/rocksdb/AbstractRocksIterator.java

@@ -6,6 +6,8 @@
 package org.rocksdb;
 
 import java.nio.ByteBuffer;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
 
 /**
  * Base class implementation for Rocks Iterators
@@ -19,21 +21,28 @@
  * @param <P> The type of the Parent Object from which the Rocks Iterator was
  *          created. This is used by disposeInternal to avoid double-free
  *          issues with the underlying C++ object.
- * @see org.rocksdb.RocksObject
+ * @see RocksObject
  */
 public abstract class AbstractRocksIterator<P extends RocksObject>
     extends RocksObject implements RocksIteratorInterface {
   final P parent_;
+  final AtomicReference<Function<AbstractRocksIterator<P>, Boolean>> removeOnClosure_ =
+      new AtomicReference<>();
 
-  protected AbstractRocksIterator(final P parent,
-      final long nativeHandle) {
+  protected AbstractRocksIterator(final P parent, final long nativeHandle,
+      final Function<AbstractRocksIterator<P>, Boolean> removeOnClosure) {
     super(nativeHandle);
     // parent must point to a valid RocksDB instance.
     assert (parent != null);
     // RocksIterator must hold a reference to the related parent instance
     // to guarantee that while a GC cycle starts RocksIterator instances
     // are freed prior to parent instances.
     parent_ = parent;
+    removeOnClosure_.set(removeOnClosure);
+  }
+
+  protected AbstractRocksIterator(final P parent, final long nativeHandle) {
+    this(parent, nativeHandle, null);
   }
 
   @Override
@@ -120,6 +129,15 @@ public void status() throws RocksDBException {
     status0(nativeHandle_);
   }
 
+  @Override
+  public void close() {
+    super.close();
+    Function<AbstractRocksIterator<P>, Boolean> removeOnClosure = removeOnClosure_.getAndSet(null);
+    if (removeOnClosure != null) {
+      removeOnClosure.apply(this);
+    }
+  }
+
   /**
    * <p>Deletes underlying C++ iterator pointer.</p>
    *

java/src/main/java/org/rocksdb/RocksDB.java

@@ -43,6 +43,8 @@ private enum LibraryState {
 
   final List<ColumnFamilyHandle> ownedColumnFamilyHandles = new ArrayList<>();
 
+  final List<RocksIterator> ownedIterators = new ArrayList<>();
+
   /**
    * Loads the necessary library files.
    * Calling this method twice will have no effect.
@@ -671,9 +673,15 @@ public void closeE() throws RocksDBException {
    * <p>
    * See also {@link #close()}.
    */
-  @SuppressWarnings("PMD.EmptyCatchBlock")
+  @SuppressWarnings({"PMD.EmptyCatchBlock", "PMD.CloseResource"})
   @Override
   public void close() {
+    final List<RocksIterator> removeIterators = new ArrayList<>(ownedIterators);
+    for (final RocksIterator rocksIterator : removeIterators) {
+      // The lambda supplied to rocksIterator on construction will remove it from ownedIterators
+      rocksIterator.close();
+    }
+
     for (final ColumnFamilyHandle columnFamilyHandle : // NOPMD - CloseResource
         ownedColumnFamilyHandles) {
       columnFamilyHandle.close();
@@ -3262,9 +3270,8 @@ public KeyMayExist keyMayExist(final ColumnFamilyHandle columnFamilyHandle,
    * @return instance of iterator object.
    */
   public RocksIterator newIterator() {
-    return new RocksIterator(this,
-        iterator(nativeHandle_, defaultColumnFamilyHandle_.nativeHandle_,
-            defaultReadOptions_.nativeHandle_));
+    return makeIterator(iterator(nativeHandle_, defaultColumnFamilyHandle_.nativeHandle_,
+        defaultReadOptions_.nativeHandle_));
   }
 
   /**
@@ -3281,9 +3288,8 @@ public RocksIterator newIterator() {
    * @return instance of iterator object.
    */
   public RocksIterator newIterator(final ReadOptions readOptions) {
-    return new RocksIterator(this,
-        iterator(
-            nativeHandle_, defaultColumnFamilyHandle_.nativeHandle_, readOptions.nativeHandle_));
+    return makeIterator(iterator(
+        nativeHandle_, defaultColumnFamilyHandle_.nativeHandle_, readOptions.nativeHandle_));
   }
 
   /**
@@ -3302,9 +3308,8 @@ public RocksIterator newIterator(final ReadOptions readOptions) {
    */
   public RocksIterator newIterator(
       final ColumnFamilyHandle columnFamilyHandle) {
-    return new RocksIterator(this,
-        iterator(
-            nativeHandle_, columnFamilyHandle.nativeHandle_, defaultReadOptions_.nativeHandle_));
+    return makeIterator(iterator(
+        nativeHandle_, columnFamilyHandle.nativeHandle_, defaultReadOptions_.nativeHandle_));
   }
 
   /**
@@ -3324,8 +3329,8 @@ public RocksIterator newIterator(
    */
   public RocksIterator newIterator(final ColumnFamilyHandle columnFamilyHandle,
       final ReadOptions readOptions) {
-    return new RocksIterator(
-        this, iterator(nativeHandle_, columnFamilyHandle.nativeHandle_, readOptions.nativeHandle_));
+    return makeIterator(
+        iterator(nativeHandle_, columnFamilyHandle.nativeHandle_, readOptions.nativeHandle_));
   }
 
   /**
@@ -3376,7 +3381,7 @@ public List<RocksIterator> newIterators(
     final List<RocksIterator> iterators = new ArrayList<>(
         columnFamilyHandleList.size());
     for (int i=0; i<columnFamilyHandleList.size(); i++){
-      iterators.add(new RocksIterator(this, iteratorRefs[i]));
+      iterators.add(makeIterator(iteratorRefs[i]));
     }
     return iterators;
   }
@@ -4746,6 +4751,21 @@ public void deleteFilesInRanges(final ColumnFamilyHandle columnFamily, final Lis
         rangesArray, includeEnd);
   }
 
+  /**
+   * Wrap a (newly created) native iterator in a RocksIterator,
+   * and record it in our list of owned iterators for consistent removal prior to db close.
+   *
+   * @param nativeIterator the native iterator to wrap
+   * @return the wrapped iterator
+   */
+  private RocksIterator makeIterator(final long nativeIterator) {
+    final RocksIterator rocksIterator =
+        new RocksIterator(this, nativeIterator, ownedIterators::remove);
+    ownedIterators.add(rocksIterator);
+
+    return rocksIterator;
+  }
+
   /**
    * Static method to destroy the contents of the specified database.
    * Be very careful using this method.

java/src/main/java/org/rocksdb/RocksIterator.java

@@ -8,6 +8,7 @@
 import static org.rocksdb.util.BufferUtil.CheckBounds;
 
 import java.nio.ByteBuffer;
+import java.util.function.Function;
 
 /**
  * <p>An iterator that yields a sequence of key/value pairs from a source.
@@ -27,6 +28,11 @@ protected RocksIterator(final RocksDB rocksDB, final long nativeHandle) {
     super(rocksDB, nativeHandle);
   }
 
+  protected RocksIterator(final RocksDB rocksDB, final long nativeHandle,
+      final Function<AbstractRocksIterator<RocksDB>, Boolean> removeOnClosure) {
+    super(rocksDB, nativeHandle, removeOnClosure);
+  }
+
   /**
    * <p>Return the key for the current entry.  The underlying storage for
    * the returned slice is valid only until the next modification of

java/src/test/java/org/rocksdb/IteratorClosedDBTest.java

@@ -0,0 +1,76 @@
+package org.rocksdb;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import org.junit.ClassRule;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+
+public class IteratorClosedDBTest {
+  @ClassRule
+  public static final RocksNativeLibraryResource ROCKS_NATIVE_LIBRARY_RESOURCE =
+      new RocksNativeLibraryResource();
+
+  @Rule public TemporaryFolder dbFolder = new TemporaryFolder();
+
+  @Test
+  public void ownedIterators() throws RocksDBException {
+    try (Options options = new Options().setCreateIfMissing(true);
+         RocksDB db = RocksDB.open(options, dbFolder.getRoot().getAbsolutePath())) {
+      byte[] key = {0x1};
+      byte[] value = {0x2};
+      db.put(key, value);
+
+      RocksIterator it = db.newIterator();
+      it.seekToFirst();
+      assertThat(it.key()).isEqualTo(key);
+      assertThat(it.value()).isEqualTo(value);
+
+      it.next();
+      assertThat(it.isValid()).isFalse();
+    } // if iterator were still open when we close the DB, we would see a C++ assertion in
+      // DEBUG_LEVEL=1
+  }
+
+  @Test
+  public void shouldNotCrashJavaRocks() throws RocksDBException {
+    try (Options options = new Options().setCreateIfMissing(true)) {
+      RocksDB db = RocksDB.open(options, dbFolder.getRoot().getAbsolutePath());
+      byte[] key = {0x1};
+      byte[] value = {0x2};
+      db.put(key, value);
+
+      RocksIterator it = db.newIterator();
+      assertThat(it.isValid()).isFalse();
+      it.seekToFirst();
+      assertThat(it.isValid()).isTrue();
+
+      byte[] valueOK = it.value();
+      assertThat(valueOK).isEqualTo(value);
+
+      // Close should work because iterator references are now cleaned up
+      // Previously would have thrown an exception here (assertion failure in C++) -
+      // when built with DEBUG_LEVEL=1
+      // Because the outstanding iterator has a reference to the column family which is being closed
+      db.close();
+
+      try {
+        byte[] valueShouldAssert = it.value();
+        throw new RuntimeException("it.value() should cause an assertion");
+      } catch (AssertionError ignored) {
+      }
+
+      // should assert
+      try {
+        boolean isValidShouldAssert = it.isValid();
+        throw new RuntimeException("it.isValid() should cause an assertion");
+      } catch (AssertionError ignored) {
+      }
+
+      // Multiple close() should be fine/no-op
+      it.close();
+      it.close();
+    }
+  }
+}