Zero copy cursor put, #51

deephacks · Dec 16, 2015 · 1b3d929 · 1b3d929
1 parent 0b628e3
commit 1b3d929
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 3 deletions.
diff --git a/lmdbjni/src/main/java/org/fusesource/lmdbjni/Cursor.java b/lmdbjni/src/main/java/org/fusesource/lmdbjni/Cursor.java
@@ -267,6 +267,31 @@ private byte[] put(Value keySlice, Value valueSlice, int flags) {
     return valueSlice.toByteArray();
   }
 
+
+  /**
+   *
+   * @param key
+   * @param size
+   * @return
+   */
+  public DirectBuffer reserve(DirectBuffer key, int size) {
+    checkArgNotNull(key, "key");
+    if (buffer == null) {
+      buffer = new DirectBuffer(ByteBuffer.allocateDirect(Unsafe.ADDRESS_SIZE * 4));
+      bufferAddress = buffer.addressOffset();
+    }
+    Unsafe.putLong(bufferAddress, 0, key.capacity());
+    Unsafe.putLong(bufferAddress, 1, key.addressOffset());
+    Unsafe.putLong(bufferAddress, 2, size);
+    int rc = mdb_cursor_put_address(pointer(), bufferAddress, bufferAddress + 2 * Unsafe.ADDRESS_SIZE, Constants.RESERVE);
+    checkErrorCode(rc);
+    int valSize = (int) Unsafe.getLong(bufferAddress, 2);
+    long valAddress = Unsafe.getAddress(bufferAddress, 3);
+    DirectBuffer empty = new DirectBuffer(0, 0);
+    empty.wrap(valAddress, valSize);
+    return empty;
+  }
+
   /**
    * <p>
    *   Delete current key/data pair.

diff --git a/lmdbjni/src/test/java/org/fusesource/lmdbjni/ZeroCopyTest.java b/lmdbjni/src/test/java/org/fusesource/lmdbjni/ZeroCopyTest.java
@@ -8,14 +8,14 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.List;
 
 import static org.fusesource.lmdbjni.Constants.FIRST;
 import static org.fusesource.lmdbjni.Constants.NEXT;
 import static org.fusesource.lmdbjni.LMDBException.NOTFOUND;
 import static org.hamcrest.core.Is.is;
 import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThat;
 
 public class ZeroCopyTest {
@@ -165,8 +165,8 @@ public void testIteration() {
 
   @Test
   public void testReserve() {
-    byte[] key = new byte[] {1,2,3};
-    byte[] val = new byte[] {3,2,1};
+    byte[] key = new byte[]{1, 2, 3};
+    byte[] val = new byte[]{3, 2, 1};
 
     try (Transaction tx = env.createWriteTransaction()) {
       DirectBuffer keyBuf = new DirectBuffer(ByteBuffer.allocateDirect(key.length));
@@ -181,4 +181,42 @@ public void testReserve() {
       assertArrayEquals(result, val);
     }
   }
+
+  @Test
+  public void testReserveCursor() {
+    byte[] key = new byte[]{1, 1, 1};
+    byte[] val = new byte[]{3, 3, 3};
+
+    try (Transaction tx = env.createWriteTransaction()) {
+      try (Cursor cursor = db.openCursor(tx)) {
+        DirectBuffer keyBuf = new DirectBuffer(ByteBuffer.allocateDirect(key.length));
+        keyBuf.putBytes(0, key);
+        DirectBuffer valBuf = cursor.reserve(keyBuf, val.length);
+        valBuf.putBytes(0, val);
+      }
+      tx.commit();
+    }
+
+    try (Transaction tx = env.createReadTransaction()) {
+      byte[] result = db.get(tx, key);
+      assertArrayEquals(result, val);
+    }
+  }
+
+  @Test
+  public void testReserveCursorRollback() {
+    byte[] key = new byte[]{1, 1, 1};
+    byte[] val = new byte[]{3, 3, 3};
+
+    Transaction tx = env.createWriteTransaction();
+    Cursor cursor = db.openCursor(tx);
+    DirectBuffer keyBuf = new DirectBuffer(ByteBuffer.allocateDirect(key.length));
+    keyBuf.putBytes(0, key);
+    DirectBuffer valBuf = cursor.reserve(keyBuf, val.length);
+    valBuf.putBytes(0, val);
+    tx.abort();
+
+    byte[] result = db.get(key);
+    assertNull(result);
+  }
 }