Add RocksIterator::Refresh (#6573)

Summary:
This PR exposes the `Iterator::Refresh` method to the Java API by adding it on the `RocksIteratorInterface` interface. There are three concrete implementations: `RocksIterator`, `SstFileReaderIterator`, and `WBWIRocksIterator`. For the first two cases, the JNI side simply delegates to the underlying `Iterator::Refresh` method; in the last case, as it doesn't share an ancestor, and per the discussion in https://github.com/facebook/rocksdb/issues/3465, a `Status::NotSupported` exception is thrown.

As the last PR had no activity in a while, I'm opening a new one - I'm completely fine with merging the previous PR if it gets completed before this is reviewed.

Let me know if there's anything missing or anything else I can do 👍
Pull Request resolved: https://github.com/facebook/rocksdb/pull/6573

Reviewed By: cheng-chang

Differential Revision: D20604666

Pulled By: pdillinger

fbshipit-source-id: 4de17df1180c3b87b76cfdd77b674b81fc0563f7

java/rocksjni/iterator.cc

@@ -81,6 +81,23 @@ void Java_org_rocksdb_RocksIterator_prev0(JNIEnv* /*env*/, jobject /*jobj*/,
   reinterpret_cast<ROCKSDB_NAMESPACE::Iterator*>(handle)->Prev();
 }
 
+/*
+ * Class:     org_rocksdb_RocksIterator
+ * Method:    refresh0
+ * Signature: (J)V
+ */
+void Java_org_rocksdb_RocksIterator_refresh0(JNIEnv* env, jobject /*jobj*/,
+                                            jlong handle) {
+  auto* it = reinterpret_cast<ROCKSDB_NAMESPACE::Iterator*>(handle);
+  ROCKSDB_NAMESPACE::Status s = it->Refresh();
+
+  if (s.ok()) {
+    return;
+  }
+
+  ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(env, s);
+}
+
 /*
  * Class:     org_rocksdb_RocksIterator
  * Method:    seek0

java/rocksjni/sst_file_reader_iterator.cc

@@ -251,3 +251,20 @@ void Java_org_rocksdb_SstFileReaderIterator_seekForPrevDirect0(
   ROCKSDB_NAMESPACE::JniUtil::k_op_direct(seekPrev, env, jtarget, jtarget_off,
                                           jtarget_len);
 }
+
+/*
+ * Class:     org_rocksdb_SstFileReaderIterator
+ * Method:    refresh0
+ * Signature: (J)V
+ */
+void Java_org_rocksdb_SstFileReaderIterator_refresh0(JNIEnv* env, jobject /*jobj*/,
+                                            jlong handle) {
+  auto* it = reinterpret_cast<ROCKSDB_NAMESPACE::Iterator*>(handle);
+  ROCKSDB_NAMESPACE::Status s = it->Refresh();
+
+  if (s.ok()) {
+    return;
+  }
+
+  ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(env, s);
+}

java/rocksjni/write_batch_with_index.cc

@@ -860,3 +860,13 @@ jlongArray Java_org_rocksdb_WBWIRocksIterator_entry1(JNIEnv* env,
 
   return jresults;
 }
+
+/*
+ * Class:     org_rocksdb_WBWIRocksIterator
+ * Method:    refresh0
+ * Signature: (J)V
+ */
+void Java_org_rocksdb_WBWIRocksIterator_refresh0(JNIEnv* env) {
+  ROCKSDB_NAMESPACE::Status s = ROCKSDB_NAMESPACE::Status::NotSupported("Refresh() is not supported");
+  ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(env, s);
+}

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

@@ -92,6 +92,12 @@ public abstract class AbstractRocksIterator<P extends RocksObject>
     prev0(nativeHandle_);
   }
 
+  @Override
+  public void refresh() throws RocksDBException {
+    assert (isOwningHandle());
+    refresh0(nativeHandle_);
+  }
+
   @Override
   public void status() throws RocksDBException {
     assert (isOwningHandle());
@@ -118,6 +124,7 @@ public abstract class AbstractRocksIterator<P extends RocksObject>
   abstract void seekToLast0(long handle);
   abstract void next0(long handle);
   abstract void prev0(long handle);
+  abstract void refresh0(long handle) throws RocksDBException;
   abstract void seek0(long handle, byte[] target, int targetLen);
   abstract void seekForPrev0(long handle, byte[] target, int targetLen);
   abstract void seekDirect0(long handle, ByteBuffer target, int targetOffset, int targetLen);

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

@@ -102,6 +102,7 @@ public class RocksIterator extends AbstractRocksIterator<RocksDB> {
   @Override final native void seekToLast0(long handle);
   @Override final native void next0(long handle);
   @Override final native void prev0(long handle);
+  @Override final native void refresh0(long handle);
   @Override final native void seek0(long handle, byte[] target, int targetLen);
   @Override final native void seekForPrev0(long handle, byte[] target, int targetLen);
   @Override

java/src/main/java/org/rocksdb/RocksIteratorInterface.java

@@ -114,4 +114,14 @@ public interface RocksIteratorInterface {
    *                          native library.
    */
   void status() throws RocksDBException;
+
+  /**
+   * <p>If supported, renew the iterator to represent the latest state. The iterator will be
+   * invalidated after the call. Not supported if {@link ReadOptions#setSnapshot(Snapshot)} was
+   * specified when creating the iterator.</p>
+   *
+   * @throws RocksDBException thrown if the operation is not supported or an error happens in the
+   *     underlying native library
+   */
+  void refresh() throws RocksDBException;
 }

java/src/main/java/org/rocksdb/SstFileReaderIterator.java

@@ -102,6 +102,7 @@ public class SstFileReaderIterator extends AbstractRocksIterator<SstFileReader>
   @Override final native void seekToLast0(long handle);
   @Override final native void next0(long handle);
   @Override final native void prev0(long handle);
+  @Override final native void refresh0(long handle) throws RocksDBException;
   @Override final native void seek0(long handle, byte[] target, int targetLen);
   @Override final native void seekForPrev0(long handle, byte[] target, int targetLen);
   @Override final native void status0(long handle) throws RocksDBException;

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

@@ -46,6 +46,7 @@ public class WBWIRocksIterator
   @Override final native void seekToLast0(long handle);
   @Override final native void next0(long handle);
   @Override final native void prev0(long handle);
+  @Override final native void refresh0(final long handle) throws RocksDBException;
   @Override final native void seek0(long handle, byte[] target, int targetLen);
   @Override final native void seekForPrev0(long handle, byte[] target, int targetLen);
   @Override final native void status0(long handle) throws RocksDBException;

java/src/test/java/org/rocksdb/RocksIteratorTest.java

@@ -147,6 +147,27 @@ public class RocksIteratorTest {
         assertThat(iterator.isValid()).isTrue();
         assertThat(iterator.key()).isEqualTo("key2".getBytes());
       }
+
+      try (final RocksIterator iterator = db.newIterator()) {
+        iterator.seekToFirst();
+        assertThat(iterator.isValid()).isTrue();
+
+        byte[] lastKey;
+        do {
+          lastKey = iterator.key();
+          iterator.next();
+        } while (iterator.isValid());
+
+        db.put("key3".getBytes(), "value3".getBytes());
+        assertThat(iterator.isValid()).isFalse();
+        iterator.refresh();
+        iterator.seek(lastKey);
+        assertThat(iterator.isValid()).isTrue();
+
+        iterator.next();
+        assertThat(iterator.isValid()).isTrue();
+        assertThat(iterator.key()).isEqualTo("key3".getBytes());
+      }
     }
   }
 

java/src/test/java/org/rocksdb/util/BytewiseComparatorTest.java

@@ -263,7 +263,7 @@ public class BytewiseComparatorTest {
       for (int i = 0; i < num_iter_ops; i++) {
         // Random walk and make sure iter and result_iter returns the
         // same key and value
-        final int type = rnd.nextInt(7);
+        final int type = rnd.nextInt(8);
         iter.status();
         switch (type) {
           case 0:
@@ -310,8 +310,15 @@ public class BytewiseComparatorTest {
               continue;
             }
             break;
+          case 6:
+            // Refresh
+            iter.refresh();
+            result_iter.refresh();
+            iter.seekToFirst();
+            result_iter.seekToFirst();
+            break;
           default: {
-            assert (type == 6);
+            assert (type == 7);
             final int key_idx = rnd.nextInt(source_strings.size());
             final String key = source_strings.get(key_idx);
             final byte[] result = db.get(readOptions, bytes(key));
@@ -472,6 +479,11 @@ public class BytewiseComparatorTest {
       }
     }
 
+    @Override
+    public void refresh() throws RocksDBException {
+      offset = -1;
+    }
+
     @Override
     public void status() throws RocksDBException {
       if(offset < 0 || offset >= entries.size()) {