| // Copyright (c) 2011-present, Facebook, Inc. All rights reserved. |
| // This source code is licensed under both the GPLv2 (found in the |
| // COPYING file in the root directory) and Apache 2.0 License |
| // (found in the LICENSE.Apache file in the root directory). |
| // |
| // Copyright (c) 2011 The LevelDB Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. See the AUTHORS file for names of contributors. |
| |
| package org.rocksdb; |
| |
| import org.junit.ClassRule; |
| import org.junit.Rule; |
| import org.junit.Test; |
| import org.junit.rules.TemporaryFolder; |
| |
| import java.io.UnsupportedEncodingException; |
| import java.nio.ByteBuffer; |
| import java.util.Arrays; |
| |
| import static org.assertj.core.api.Assertions.assertThat; |
| |
| |
| public class WriteBatchWithIndexTest { |
| |
| @ClassRule |
| public static final RocksMemoryResource rocksMemoryResource = |
| new RocksMemoryResource(); |
| |
| @Rule |
| public TemporaryFolder dbFolder = new TemporaryFolder(); |
| |
| @Test |
| public void readYourOwnWrites() throws RocksDBException { |
| try (final Options options = new Options().setCreateIfMissing(true); |
| final RocksDB db = RocksDB.open(options, |
| dbFolder.getRoot().getAbsolutePath())) { |
| |
| final byte[] k1 = "key1".getBytes(); |
| final byte[] v1 = "value1".getBytes(); |
| final byte[] k2 = "key2".getBytes(); |
| final byte[] v2 = "value2".getBytes(); |
| |
| db.put(k1, v1); |
| db.put(k2, v2); |
| |
| try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex(true); |
| final RocksIterator base = db.newIterator(); |
| final RocksIterator it = wbwi.newIteratorWithBase(base)) { |
| |
| it.seek(k1); |
| assertThat(it.isValid()).isTrue(); |
| assertThat(it.key()).isEqualTo(k1); |
| assertThat(it.value()).isEqualTo(v1); |
| |
| it.seek(k2); |
| assertThat(it.isValid()).isTrue(); |
| assertThat(it.key()).isEqualTo(k2); |
| assertThat(it.value()).isEqualTo(v2); |
| |
| //put data to the write batch and make sure we can read it. |
| final byte[] k3 = "key3".getBytes(); |
| final byte[] v3 = "value3".getBytes(); |
| wbwi.put(k3, v3); |
| it.seek(k3); |
| assertThat(it.isValid()).isTrue(); |
| assertThat(it.key()).isEqualTo(k3); |
| assertThat(it.value()).isEqualTo(v3); |
| |
| //update k2 in the write batch and check the value |
| final byte[] v2Other = "otherValue2".getBytes(); |
| wbwi.put(k2, v2Other); |
| it.seek(k2); |
| assertThat(it.isValid()).isTrue(); |
| assertThat(it.key()).isEqualTo(k2); |
| assertThat(it.value()).isEqualTo(v2Other); |
| |
| //remove k1 and make sure we can read back the write |
| wbwi.remove(k1); |
| it.seek(k1); |
| assertThat(it.key()).isNotEqualTo(k1); |
| |
| //reinsert k1 and make sure we see the new value |
| final byte[] v1Other = "otherValue1".getBytes(); |
| wbwi.put(k1, v1Other); |
| it.seek(k1); |
| assertThat(it.isValid()).isTrue(); |
| assertThat(it.key()).isEqualTo(k1); |
| assertThat(it.value()).isEqualTo(v1Other); |
| } |
| } |
| } |
| |
| @Test |
| public void write_writeBatchWithIndex() throws RocksDBException { |
| try (final Options options = new Options().setCreateIfMissing(true); |
| final RocksDB db = RocksDB.open(options, |
| dbFolder.getRoot().getAbsolutePath())) { |
| |
| final byte[] k1 = "key1".getBytes(); |
| final byte[] v1 = "value1".getBytes(); |
| final byte[] k2 = "key2".getBytes(); |
| final byte[] v2 = "value2".getBytes(); |
| |
| try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex()) { |
| wbwi.put(k1, v1); |
| wbwi.put(k2, v2); |
| |
| db.write(new WriteOptions(), wbwi); |
| } |
| |
| assertThat(db.get(k1)).isEqualTo(v1); |
| assertThat(db.get(k2)).isEqualTo(v2); |
| } |
| } |
| |
| @Test |
| public void iterator() throws RocksDBException { |
| try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex(true)) { |
| |
| final String k1 = "key1"; |
| final String v1 = "value1"; |
| final String k2 = "key2"; |
| final String v2 = "value2"; |
| final String k3 = "key3"; |
| final String v3 = "value3"; |
| final byte[] k1b = k1.getBytes(); |
| final byte[] v1b = v1.getBytes(); |
| final byte[] k2b = k2.getBytes(); |
| final byte[] v2b = v2.getBytes(); |
| final byte[] k3b = k3.getBytes(); |
| final byte[] v3b = v3.getBytes(); |
| |
| //add put records |
| wbwi.put(k1b, v1b); |
| wbwi.put(k2b, v2b); |
| wbwi.put(k3b, v3b); |
| |
| //add a deletion record |
| final String k4 = "key4"; |
| final byte[] k4b = k4.getBytes(); |
| wbwi.remove(k4b); |
| |
| final WBWIRocksIterator.WriteEntry[] expected = { |
| new WBWIRocksIterator.WriteEntry(WBWIRocksIterator.WriteType.PUT, |
| new DirectSlice(k1), new DirectSlice(v1)), |
| new WBWIRocksIterator.WriteEntry(WBWIRocksIterator.WriteType.PUT, |
| new DirectSlice(k2), new DirectSlice(v2)), |
| new WBWIRocksIterator.WriteEntry(WBWIRocksIterator.WriteType.PUT, |
| new DirectSlice(k3), new DirectSlice(v3)), |
| new WBWIRocksIterator.WriteEntry(WBWIRocksIterator.WriteType.DELETE, |
| new DirectSlice(k4), DirectSlice.NONE) |
| }; |
| |
| try (final WBWIRocksIterator it = wbwi.newIterator()) { |
| //direct access - seek to key offsets |
| final int[] testOffsets = {2, 0, 1, 3}; |
| |
| for (int i = 0; i < testOffsets.length; i++) { |
| final int testOffset = testOffsets[i]; |
| final byte[] key = toArray(expected[testOffset].getKey().data()); |
| |
| it.seek(key); |
| assertThat(it.isValid()).isTrue(); |
| |
| final WBWIRocksIterator.WriteEntry entry = it.entry(); |
| assertThat(entry.equals(expected[testOffset])).isTrue(); |
| } |
| |
| //forward iterative access |
| int i = 0; |
| for (it.seekToFirst(); it.isValid(); it.next()) { |
| assertThat(it.entry().equals(expected[i++])).isTrue(); |
| } |
| |
| //reverse iterative access |
| i = expected.length - 1; |
| for (it.seekToLast(); it.isValid(); it.prev()) { |
| assertThat(it.entry().equals(expected[i--])).isTrue(); |
| } |
| } |
| } |
| } |
| |
| @Test |
| public void zeroByteTests() { |
| try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex(true)) { |
| final byte[] zeroByteValue = new byte[]{0, 0}; |
| //add zero byte value |
| wbwi.put(zeroByteValue, zeroByteValue); |
| |
| final ByteBuffer buffer = ByteBuffer.allocateDirect(zeroByteValue.length); |
| buffer.put(zeroByteValue); |
| |
| final WBWIRocksIterator.WriteEntry expected = |
| new WBWIRocksIterator.WriteEntry(WBWIRocksIterator.WriteType.PUT, |
| new DirectSlice(buffer, zeroByteValue.length), |
| new DirectSlice(buffer, zeroByteValue.length)); |
| |
| try (final WBWIRocksIterator it = wbwi.newIterator()) { |
| it.seekToFirst(); |
| final WBWIRocksIterator.WriteEntry actual = it.entry(); |
| assertThat(actual.equals(expected)).isTrue(); |
| assertThat(it.entry().hashCode() == expected.hashCode()).isTrue(); |
| } |
| } |
| } |
| |
| @Test |
| public void savePoints() |
| throws UnsupportedEncodingException, RocksDBException { |
| try (final Options options = new Options().setCreateIfMissing(true); |
| final RocksDB db = RocksDB.open(options, |
| dbFolder.getRoot().getAbsolutePath())) { |
| try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex(true); |
| final ReadOptions readOptions = new ReadOptions()) { |
| wbwi.put("k1".getBytes(), "v1".getBytes()); |
| wbwi.put("k2".getBytes(), "v2".getBytes()); |
| wbwi.put("k3".getBytes(), "v3".getBytes()); |
| |
| assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k1")) |
| .isEqualTo("v1"); |
| assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k2")) |
| .isEqualTo("v2"); |
| assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k3")) |
| .isEqualTo("v3"); |
| |
| |
| wbwi.setSavePoint(); |
| |
| wbwi.remove("k2".getBytes()); |
| wbwi.put("k3".getBytes(), "v3-2".getBytes()); |
| |
| assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k2")) |
| .isNull(); |
| assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k3")) |
| .isEqualTo("v3-2"); |
| |
| |
| wbwi.setSavePoint(); |
| |
| wbwi.put("k3".getBytes(), "v3-3".getBytes()); |
| wbwi.put("k4".getBytes(), "v4".getBytes()); |
| |
| assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k3")) |
| .isEqualTo("v3-3"); |
| assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k4")) |
| .isEqualTo("v4"); |
| |
| |
| wbwi.rollbackToSavePoint(); |
| |
| assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k2")) |
| .isNull(); |
| assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k3")) |
| .isEqualTo("v3-2"); |
| assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k4")) |
| .isNull(); |
| |
| |
| wbwi.rollbackToSavePoint(); |
| |
| assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k1")) |
| .isEqualTo("v1"); |
| assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k2")) |
| .isEqualTo("v2"); |
| assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k3")) |
| .isEqualTo("v3"); |
| assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k4")) |
| .isNull(); |
| } |
| } |
| } |
| |
| @Test(expected = RocksDBException.class) |
| public void restorePoints_withoutSavePoints() throws RocksDBException { |
| try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex()) { |
| wbwi.rollbackToSavePoint(); |
| } |
| } |
| |
| @Test(expected = RocksDBException.class) |
| public void restorePoints_withoutSavePoints_nested() throws RocksDBException { |
| try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex()) { |
| |
| wbwi.setSavePoint(); |
| wbwi.rollbackToSavePoint(); |
| |
| // without previous corresponding setSavePoint |
| wbwi.rollbackToSavePoint(); |
| } |
| } |
| |
| private static String getFromWriteBatchWithIndex(final RocksDB db, |
| final ReadOptions readOptions, final WriteBatchWithIndex wbwi, |
| final String skey) { |
| final byte[] key = skey.getBytes(); |
| try(final RocksIterator baseIterator = db.newIterator(readOptions); |
| final RocksIterator iterator = wbwi.newIteratorWithBase(baseIterator)) { |
| iterator.seek(key); |
| |
| // Arrays.equals(key, iterator.key()) ensures an exact match in Rocks, |
| // instead of a nearest match |
| return iterator.isValid() && |
| Arrays.equals(key, iterator.key()) ? |
| new String(iterator.value()) : null; |
| } |
| } |
| |
| @Test |
| public void getFromBatch() throws RocksDBException { |
| final byte[] k1 = "k1".getBytes(); |
| final byte[] k2 = "k2".getBytes(); |
| final byte[] k3 = "k3".getBytes(); |
| final byte[] k4 = "k4".getBytes(); |
| |
| final byte[] v1 = "v1".getBytes(); |
| final byte[] v2 = "v2".getBytes(); |
| final byte[] v3 = "v3".getBytes(); |
| |
| try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex(true); |
| final DBOptions dbOptions = new DBOptions()) { |
| wbwi.put(k1, v1); |
| wbwi.put(k2, v2); |
| wbwi.put(k3, v3); |
| |
| assertThat(wbwi.getFromBatch(dbOptions, k1)).isEqualTo(v1); |
| assertThat(wbwi.getFromBatch(dbOptions, k2)).isEqualTo(v2); |
| assertThat(wbwi.getFromBatch(dbOptions, k3)).isEqualTo(v3); |
| assertThat(wbwi.getFromBatch(dbOptions, k4)).isNull(); |
| |
| wbwi.remove(k2); |
| |
| assertThat(wbwi.getFromBatch(dbOptions, k2)).isNull(); |
| } |
| } |
| |
| @Test |
| public void getFromBatchAndDB() throws RocksDBException { |
| final byte[] k1 = "k1".getBytes(); |
| final byte[] k2 = "k2".getBytes(); |
| final byte[] k3 = "k3".getBytes(); |
| final byte[] k4 = "k4".getBytes(); |
| |
| final byte[] v1 = "v1".getBytes(); |
| final byte[] v2 = "v2".getBytes(); |
| final byte[] v3 = "v3".getBytes(); |
| final byte[] v4 = "v4".getBytes(); |
| |
| try (final Options options = new Options().setCreateIfMissing(true); |
| final RocksDB db = RocksDB.open(options, |
| dbFolder.getRoot().getAbsolutePath())) { |
| |
| db.put(k1, v1); |
| db.put(k2, v2); |
| db.put(k4, v4); |
| |
| try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex(true); |
| final DBOptions dbOptions = new DBOptions(); |
| final ReadOptions readOptions = new ReadOptions()) { |
| |
| assertThat(wbwi.getFromBatch(dbOptions, k1)).isNull(); |
| assertThat(wbwi.getFromBatch(dbOptions, k2)).isNull(); |
| assertThat(wbwi.getFromBatch(dbOptions, k4)).isNull(); |
| |
| wbwi.put(k3, v3); |
| |
| assertThat(wbwi.getFromBatch(dbOptions, k3)).isEqualTo(v3); |
| |
| assertThat(wbwi.getFromBatchAndDB(db, readOptions, k1)).isEqualTo(v1); |
| assertThat(wbwi.getFromBatchAndDB(db, readOptions, k2)).isEqualTo(v2); |
| assertThat(wbwi.getFromBatchAndDB(db, readOptions, k3)).isEqualTo(v3); |
| assertThat(wbwi.getFromBatchAndDB(db, readOptions, k4)).isEqualTo(v4); |
| |
| wbwi.remove(k4); |
| |
| assertThat(wbwi.getFromBatchAndDB(db, readOptions, k4)).isNull(); |
| } |
| } |
| } |
| private byte[] toArray(final ByteBuffer buf) { |
| final byte[] ary = new byte[buf.remaining()]; |
| buf.get(ary); |
| return ary; |
| } |
| |
| @Test |
| public void deleteRange() throws RocksDBException { |
| try (final RocksDB db = RocksDB.open(dbFolder.getRoot().getAbsolutePath()); |
| final WriteOptions wOpt = new WriteOptions()) { |
| db.put("key1".getBytes(), "value".getBytes()); |
| db.put("key2".getBytes(), "12345678".getBytes()); |
| db.put("key3".getBytes(), "abcdefg".getBytes()); |
| db.put("key4".getBytes(), "xyz".getBytes()); |
| assertThat(db.get("key1".getBytes())).isEqualTo("value".getBytes()); |
| assertThat(db.get("key2".getBytes())).isEqualTo("12345678".getBytes()); |
| assertThat(db.get("key3".getBytes())).isEqualTo("abcdefg".getBytes()); |
| assertThat(db.get("key4".getBytes())).isEqualTo("xyz".getBytes()); |
| |
| WriteBatch batch = new WriteBatch(); |
| batch.deleteRange("key2".getBytes(), "key4".getBytes()); |
| db.write(new WriteOptions(), batch); |
| |
| assertThat(db.get("key1".getBytes())).isEqualTo("value".getBytes()); |
| assertThat(db.get("key2".getBytes())).isNull(); |
| assertThat(db.get("key3".getBytes())).isNull(); |
| assertThat(db.get("key4".getBytes())).isEqualTo("xyz".getBytes()); |
| } |
| } |
| } |