Expose save points in Java WriteBatch and WBWI

include/rocksdb/write_batch.h

@@ -125,7 +125,7 @@ class WriteBatch : public WriteBatchBase {
   // most recent call to SetSavePoint() and removes the most recent save point.
   // If there is no previous call to SetSavePoint(), Status::NotFound()
   // will be returned.
-  // Oterwise returns Status::OK().
+  // Otherwise returns Status::OK().
   Status RollbackToSavePoint() override;
 
   // Support for iterating over the contents of a batch.

java/rocksjni/write_batch.cc

@@ -62,6 +62,37 @@ void Java_org_rocksdb_WriteBatch_clear0(JNIEnv* env, jobject jobj,
   wb->Clear();
 }
 
+/*
+ * Class:     org_rocksdb_WriteBatch
+ * Method:    setSavePoint0
+ * Signature: (J)V
+ */
+void Java_org_rocksdb_WriteBatch_setSavePoint0(
+    JNIEnv* env, jobject jobj, jlong jwb_handle) {
+  auto* wb = reinterpret_cast<rocksdb::WriteBatch*>(jwb_handle);
+  assert(wb != nullptr);
+
+  wb->SetSavePoint();
+}
+
+/*
+ * Class:     org_rocksdb_WriteBatch
+ * Method:    rollbackToSavePoint0
+ * Signature: (J)V
+ */
+void Java_org_rocksdb_WriteBatch_rollbackToSavePoint0(
+    JNIEnv* env, jobject jobj, jlong jwb_handle) {
+  auto* wb = reinterpret_cast<rocksdb::WriteBatch*>(jwb_handle);
+  assert(wb != nullptr);
+
+  auto s = wb->RollbackToSavePoint();
+
+  if (s.ok()) {
+    return;
+  }
+  rocksdb::RocksDBExceptionJni::ThrowNew(env, s);
+}
+
 /*
  * Class:     org_rocksdb_WriteBatch
  * Method:    put

java/rocksjni/write_batch_with_index.cc

@@ -198,7 +198,39 @@ void Java_org_rocksdb_WriteBatchWithIndex_clear0(
   auto* wbwi = reinterpret_cast<rocksdb::WriteBatchWithIndex*>(jwbwi_handle);
   assert(wbwi != nullptr);
 
-  wbwi->GetWriteBatch()->Clear();
+  wbwi->Clear();
+}
+
+/*
+ * Class:     org_rocksdb_WriteBatchWithIndex
+ * Method:    setSavePoint0
+ * Signature: (J)V
+ */
+void Java_org_rocksdb_WriteBatchWithIndex_setSavePoint0(
+    JNIEnv* env, jobject jobj, jlong jwbwi_handle) {
+  auto* wbwi = reinterpret_cast<rocksdb::WriteBatchWithIndex*>(jwbwi_handle);
+  assert(wbwi != nullptr);
+
+  wbwi->SetSavePoint();
+}
+
+/*
+ * Class:     org_rocksdb_WriteBatchWithIndex
+ * Method:    rollbackToSavePoint0
+ * Signature: (J)V
+ */
+void Java_org_rocksdb_WriteBatchWithIndex_rollbackToSavePoint0(
+    JNIEnv* env, jobject jobj, jlong jwbwi_handle) {
+  auto* wbwi = reinterpret_cast<rocksdb::WriteBatchWithIndex*>(jwbwi_handle);
+  assert(wbwi != nullptr);
+
+  auto s = wbwi->RollbackToSavePoint();
+
+  if (s.ok()) {
+    return;
+  }
+
+  rocksdb::RocksDBExceptionJni::ThrowNew(env, s);
 }
 
 /*

java/src/main/java/org/rocksdb/AbstractWriteBatch.java

@@ -70,6 +70,18 @@ public void clear() {
     clear0(nativeHandle_);
   }
 
+  @Override
+  public void setSavePoint() {
+    assert (isOwningHandle());
+    setSavePoint0(nativeHandle_);
+  }
+
+  @Override
+  public void rollbackToSavePoint() throws RocksDBException {
+    assert (isOwningHandle());
+    rollbackToSavePoint0(nativeHandle_);
+  }
+
   abstract int count0(final long handle);
 
   abstract void put(final long handle, final byte[] key, final int keyLen,
@@ -94,4 +106,8 @@ abstract void putLogData(final long handle, final byte[] blob,
       final int blobLen);
 
   abstract void clear0(final long handle);
+
+  abstract void setSavePoint0(final long handle);
+
+  abstract void rollbackToSavePoint0(final long handle);
 }

java/src/main/java/org/rocksdb/WriteBatch.java

@@ -82,6 +82,8 @@ public void iterate(final Handler handler) throws RocksDBException {
   @Override final native void putLogData(final long handle,
       final byte[] blob, final int blobLen);
   @Override final native void clear0(final long handle);
+  @Override final native void setSavePoint0(final long handle);
+  @Override final native void rollbackToSavePoint0(final long handle);
 
   private native static long newWriteBatch(final int reserved_bytes);
   private native void iterate(final long handle, final long handlerHandle)

java/src/main/java/org/rocksdb/WriteBatchInterface.java

@@ -95,4 +95,19 @@ void merge(ColumnFamilyHandle columnFamilyHandle,
      * Clear all updates buffered in this batch
      */
     void clear();
+
+    /**
+     * Records the state of the batch for future calls to RollbackToSavePoint().
+     * May be called multiple times to set multiple save points.
+     */
+    void setSavePoint();
+
+    /**
+     * Remove all entries in this batch (Put, Merge, Delete, PutLogData) since
+     * the most recent call to SetSavePoint() and removes the most recent save
+     * point.
+     *
+     * @throws RocksDBException if there is no previous call to SetSavePoint()
+     */
+    void rollbackToSavePoint() throws RocksDBException;
 }

java/src/main/java/org/rocksdb/WriteBatchWithIndex.java

@@ -155,6 +155,8 @@ public RocksIterator newIteratorWithBase(final RocksIterator baseIterator) {
   @Override final native void putLogData(final long handle, final byte[] blob,
       final int blobLen);
   @Override final native void clear0(final long handle);
+  @Override final native void setSavePoint0(final long handle);
+  @Override final native void rollbackToSavePoint0(final long handle);
 
   private native static long newWriteBatchWithIndex();
   private native static long newWriteBatchWithIndex(final boolean overwriteKey);

java/src/test/java/org/rocksdb/WriteBatchTest.java

@@ -14,6 +14,7 @@
 import org.junit.rules.TemporaryFolder;
 
 import java.io.UnsupportedEncodingException;
+import java.util.Arrays;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -115,11 +116,128 @@ public void blobOperation()
     }
   }
 
+  @Test
+  public void savePoints()
+      throws UnsupportedEncodingException, RocksDBException {
+    try (final WriteBatch batch = new WriteBatch()) {
+      batch.put("k1".getBytes("US-ASCII"), "v1".getBytes("US-ASCII"));
+      batch.put("k2".getBytes("US-ASCII"), "v2".getBytes("US-ASCII"));
+      batch.put("k3".getBytes("US-ASCII"), "v3".getBytes("US-ASCII"));
+
+      assertThat(getFromWriteBatch(batch, "k1")).isEqualTo("v1");
+      assertThat(getFromWriteBatch(batch, "k2")).isEqualTo("v2");
+      assertThat(getFromWriteBatch(batch, "k3")).isEqualTo("v3");
+
+
+      batch.setSavePoint();
+
+      batch.remove("k2".getBytes("US-ASCII"));
+      batch.put("k3".getBytes("US-ASCII"), "v3-2".getBytes("US-ASCII"));
+
+      assertThat(getFromWriteBatch(batch, "k2")).isNull();
+      assertThat(getFromWriteBatch(batch, "k3")).isEqualTo("v3-2");
+
+
+      batch.setSavePoint();
+
+      batch.put("k3".getBytes("US-ASCII"), "v3-3".getBytes("US-ASCII"));
+      batch.put("k4".getBytes("US-ASCII"), "v4".getBytes("US-ASCII"));
+
+      assertThat(getFromWriteBatch(batch, "k3")).isEqualTo("v3-3");
+      assertThat(getFromWriteBatch(batch, "k4")).isEqualTo("v4");
+
+
+      batch.rollbackToSavePoint();
+
+      assertThat(getFromWriteBatch(batch, "k2")).isNull();
+      assertThat(getFromWriteBatch(batch, "k3")).isEqualTo("v3-2");
+      assertThat(getFromWriteBatch(batch, "k4")).isNull();
+
+
+      batch.rollbackToSavePoint();
+
+      assertThat(getFromWriteBatch(batch, "k1")).isEqualTo("v1");
+      assertThat(getFromWriteBatch(batch, "k2")).isEqualTo("v2");
+      assertThat(getFromWriteBatch(batch, "k3")).isEqualTo("v3");
+      assertThat(getFromWriteBatch(batch, "k4")).isNull();
+    }
+  }
+
+  @Test(expected = RocksDBException.class)
+  public void restorePoints_withoutSavePoints() throws RocksDBException {
+    try (final WriteBatch batch = new WriteBatch()) {
+      batch.rollbackToSavePoint();
+    }
+  }
+
+  @Test(expected = RocksDBException.class)
+  public void restorePoints_withoutSavePoints_nested() throws RocksDBException {
+    try (final WriteBatch batch = new WriteBatch()) {
+
+      batch.setSavePoint();
+      batch.rollbackToSavePoint();
+
+      // without previous corresponding setSavePoint
+      batch.rollbackToSavePoint();
+    }
+  }
+
   static byte[] getContents(final WriteBatch wb) {
     return getContents(wb.nativeHandle_);
   }
 
+  static String getFromWriteBatch(final WriteBatch wb, final String key)
+      throws RocksDBException, UnsupportedEncodingException {
+    final WriteBatchGetter getter =
+        new WriteBatchGetter(key.getBytes("US-ASCII"));
+    wb.iterate(getter);
+    if(getter.getValue() != null) {
+      return new String(getter.getValue(), "US-ASCII");
+    } else {
+      return null;
+    }
+  }
+
   private static native byte[] getContents(final long writeBatchHandle);
+
+  private static class WriteBatchGetter extends WriteBatch.Handler {
+
+    private final byte[] key;
+    private byte[] value;
+
+    public WriteBatchGetter(final byte[] key) {
+      this.key = key;
+    }
+
+    public byte[] getValue() {
+      return value;
+    }
+
+    @Override
+    public void put(final byte[] key, final byte[] value) {
+      if(Arrays.equals(this.key, key)) {
+        this.value = value;
+      }
+    }
+
+    @Override
+    public void merge(final byte[] key, final byte[] value) {
+      if(Arrays.equals(this.key, key)) {
+        throw new UnsupportedOperationException();
+      }
+    }
+
+    @Override
+    public void delete(final byte[] key) {
+      if(Arrays.equals(this.key, key)) {
+        this.value = null;
+      }
+    }
+
+    @Override
+    public void logData(final byte[] blob) {
+    }
+  }
 }
 
 /**