1 // Copyright (c) 2011-present, Facebook, Inc. All rights reserved. 2 // This source code is licensed under both the GPLv2 (found in the 3 // COPYING file in the root directory) and Apache 2.0 License 4 // (found in the LICENSE.Apache file in the root directory). 5 // 6 // Copyright (c) 2011 The LevelDB Authors. All rights reserved. 7 // Use of this source code is governed by a BSD-style license that can be 8 // found in the LICENSE file. See the AUTHORS file for names of contributors. 9 10 package org.rocksdb; 11 12 import static java.nio.charset.StandardCharsets.UTF_8; 13 import static org.assertj.core.api.Assertions.assertThat; 14 15 import java.nio.ByteBuffer; 16 import java.util.Arrays; 17 import org.junit.ClassRule; 18 import org.junit.Rule; 19 import org.junit.Test; 20 import org.junit.rules.TemporaryFolder; 21 22 public class WriteBatchWithIndexTest { 23 24 @ClassRule 25 public static final RocksNativeLibraryResource ROCKS_NATIVE_LIBRARY_RESOURCE = 26 new RocksNativeLibraryResource(); 27 28 @Rule 29 public TemporaryFolder dbFolder = new TemporaryFolder(); 30 31 @Test readYourOwnWrites()32 public void readYourOwnWrites() throws RocksDBException { 33 try (final Options options = new Options().setCreateIfMissing(true); 34 final RocksDB db = RocksDB.open(options, 35 dbFolder.getRoot().getAbsolutePath())) { 36 37 final byte[] k1 = "key1".getBytes(); 38 final byte[] v1 = "value1".getBytes(); 39 final byte[] k2 = "key2".getBytes(); 40 final byte[] v2 = "value2".getBytes(); 41 42 db.put(k1, v1); 43 db.put(k2, v2); 44 45 try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex(true); 46 final RocksIterator base = db.newIterator(); 47 final RocksIterator it = wbwi.newIteratorWithBase(base)) { 48 it.seek(k1); 49 assertThat(it.isValid()).isTrue(); 50 assertThat(it.key()).isEqualTo(k1); 51 assertThat(it.value()).isEqualTo(v1); 52 53 it.seek(k2); 54 assertThat(it.isValid()).isTrue(); 55 assertThat(it.key()).isEqualTo(k2); 56 assertThat(it.value()).isEqualTo(v2); 57 58 //put data to the write batch and make sure we can read it. 59 final byte[] k3 = "key3".getBytes(); 60 final byte[] v3 = "value3".getBytes(); 61 wbwi.put(k3, v3); 62 it.seek(k3); 63 assertThat(it.isValid()).isTrue(); 64 assertThat(it.key()).isEqualTo(k3); 65 assertThat(it.value()).isEqualTo(v3); 66 67 //update k2 in the write batch and check the value 68 final byte[] v2Other = "otherValue2".getBytes(); 69 wbwi.put(k2, v2Other); 70 it.seek(k2); 71 assertThat(it.isValid()).isTrue(); 72 assertThat(it.key()).isEqualTo(k2); 73 assertThat(it.value()).isEqualTo(v2Other); 74 75 //delete k1 and make sure we can read back the write 76 wbwi.delete(k1); 77 it.seek(k1); 78 assertThat(it.key()).isNotEqualTo(k1); 79 80 //reinsert k1 and make sure we see the new value 81 final byte[] v1Other = "otherValue1".getBytes(); 82 wbwi.put(k1, v1Other); 83 it.seek(k1); 84 assertThat(it.isValid()).isTrue(); 85 assertThat(it.key()).isEqualTo(k1); 86 assertThat(it.value()).isEqualTo(v1Other); 87 88 //single remove k3 and make sure we can read back the write 89 wbwi.singleDelete(k3); 90 it.seek(k3); 91 assertThat(it.isValid()).isEqualTo(false); 92 93 //reinsert k3 and make sure we see the new value 94 final byte[] v3Other = "otherValue3".getBytes(); 95 wbwi.put(k3, v3Other); 96 it.seek(k3); 97 assertThat(it.isValid()).isTrue(); 98 assertThat(it.key()).isEqualTo(k3); 99 assertThat(it.value()).isEqualTo(v3Other); 100 } 101 } 102 } 103 104 @Test writeBatchWithIndex()105 public void writeBatchWithIndex() throws RocksDBException { 106 try (final Options options = new Options().setCreateIfMissing(true); 107 final RocksDB db = RocksDB.open(options, 108 dbFolder.getRoot().getAbsolutePath())) { 109 110 final byte[] k1 = "key1".getBytes(); 111 final byte[] v1 = "value1".getBytes(); 112 final byte[] k2 = "key2".getBytes(); 113 final byte[] v2 = "value2".getBytes(); 114 115 try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex(); 116 final WriteOptions wOpt = new WriteOptions()) { 117 wbwi.put(k1, v1); 118 wbwi.put(k2, v2); 119 120 db.write(wOpt, wbwi); 121 } 122 123 assertThat(db.get(k1)).isEqualTo(v1); 124 assertThat(db.get(k2)).isEqualTo(v2); 125 } 126 } 127 128 @Test write_writeBatchWithIndexDirect()129 public void write_writeBatchWithIndexDirect() throws RocksDBException { 130 try (final Options options = new Options().setCreateIfMissing(true); 131 final RocksDB db = RocksDB.open(options, dbFolder.getRoot().getAbsolutePath())) { 132 ByteBuffer k1 = ByteBuffer.allocateDirect(16); 133 ByteBuffer v1 = ByteBuffer.allocateDirect(16); 134 ByteBuffer k2 = ByteBuffer.allocateDirect(16); 135 ByteBuffer v2 = ByteBuffer.allocateDirect(16); 136 k1.put("key1".getBytes()).flip(); 137 v1.put("value1".getBytes()).flip(); 138 k2.put("key2".getBytes()).flip(); 139 v2.put("value2".getBytes()).flip(); 140 141 try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex()) { 142 wbwi.put(k1, v1); 143 assertThat(k1.position()).isEqualTo(4); 144 assertThat(k1.limit()).isEqualTo(4); 145 assertThat(v1.position()).isEqualTo(6); 146 assertThat(v1.limit()).isEqualTo(6); 147 148 wbwi.put(k2, v2); 149 150 db.write(new WriteOptions(), wbwi); 151 } 152 153 assertThat(db.get("key1".getBytes())).isEqualTo("value1".getBytes()); 154 assertThat(db.get("key2".getBytes())).isEqualTo("value2".getBytes()); 155 } 156 } 157 158 @Test iterator()159 public void iterator() throws RocksDBException { 160 try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex(true)) { 161 162 final String k1 = "key1"; 163 final String v1 = "value1"; 164 final String k2 = "key2"; 165 final String v2 = "value2"; 166 final String k3 = "key3"; 167 final String v3 = "value3"; 168 final String k4 = "key4"; 169 final String k5 = "key5"; 170 final String k6 = "key6"; 171 final String k7 = "key7"; 172 final String v8 = "value8"; 173 final byte[] k1b = k1.getBytes(UTF_8); 174 final byte[] v1b = v1.getBytes(UTF_8); 175 final byte[] k2b = k2.getBytes(UTF_8); 176 final byte[] v2b = v2.getBytes(UTF_8); 177 final byte[] k3b = k3.getBytes(UTF_8); 178 final byte[] v3b = v3.getBytes(UTF_8); 179 final byte[] k4b = k4.getBytes(UTF_8); 180 final byte[] k5b = k5.getBytes(UTF_8); 181 final byte[] k6b = k6.getBytes(UTF_8); 182 final byte[] k7b = k7.getBytes(UTF_8); 183 final byte[] v8b = v8.getBytes(UTF_8); 184 185 // add put records 186 wbwi.put(k1b, v1b); 187 wbwi.put(k2b, v2b); 188 wbwi.put(k3b, v3b); 189 190 // add a deletion record 191 wbwi.delete(k4b); 192 193 // add a single deletion record 194 wbwi.singleDelete(k5b); 195 196 // add a log record 197 wbwi.putLogData(v8b); 198 199 final WBWIRocksIterator.WriteEntry[] expected = { 200 new WBWIRocksIterator.WriteEntry(WBWIRocksIterator.WriteType.PUT, 201 new DirectSlice(k1), new DirectSlice(v1)), 202 new WBWIRocksIterator.WriteEntry(WBWIRocksIterator.WriteType.PUT, 203 new DirectSlice(k2), new DirectSlice(v2)), 204 new WBWIRocksIterator.WriteEntry(WBWIRocksIterator.WriteType.PUT, 205 new DirectSlice(k3), new DirectSlice(v3)), 206 new WBWIRocksIterator.WriteEntry(WBWIRocksIterator.WriteType.DELETE, 207 new DirectSlice(k4), DirectSlice.NONE), 208 new WBWIRocksIterator.WriteEntry(WBWIRocksIterator.WriteType.SINGLE_DELETE, 209 new DirectSlice(k5), DirectSlice.NONE), 210 }; 211 212 try (final WBWIRocksIterator it = wbwi.newIterator()) { 213 //direct access - seek to key offsets 214 final int[] testOffsets = {2, 0, 3, 4, 1}; 215 216 for (int i = 0; i < testOffsets.length; i++) { 217 final int testOffset = testOffsets[i]; 218 final byte[] key = toArray(expected[testOffset].getKey().data()); 219 220 it.seek(key); 221 assertThat(it.isValid()).isTrue(); 222 223 final WBWIRocksIterator.WriteEntry entry = it.entry(); 224 assertThat(entry).isEqualTo(expected[testOffset]); 225 226 // Direct buffer seek 227 expected[testOffset].getKey().data().mark(); 228 ByteBuffer db = expected[testOffset].getKey().data(); 229 it.seek(db); 230 assertThat(db.position()).isEqualTo(key.length); 231 assertThat(it.isValid()).isTrue(); 232 } 233 234 //forward iterative access 235 int i = 0; 236 for (it.seekToFirst(); it.isValid(); it.next()) { 237 assertThat(it.entry()).isEqualTo(expected[i++]); 238 } 239 240 //reverse iterative access 241 i = expected.length - 1; 242 for (it.seekToLast(); it.isValid(); it.prev()) { 243 assertThat(it.entry()).isEqualTo(expected[i--]); 244 } 245 } 246 } 247 } 248 249 @Test zeroByteTests()250 public void zeroByteTests() throws RocksDBException { 251 try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex(true)) { 252 final byte[] zeroByteValue = new byte[]{0, 0}; 253 //add zero byte value 254 wbwi.put(zeroByteValue, zeroByteValue); 255 256 final ByteBuffer buffer = ByteBuffer.allocateDirect(zeroByteValue.length); 257 buffer.put(zeroByteValue); 258 259 final WBWIRocksIterator.WriteEntry expected = 260 new WBWIRocksIterator.WriteEntry(WBWIRocksIterator.WriteType.PUT, 261 new DirectSlice(buffer, zeroByteValue.length), 262 new DirectSlice(buffer, zeroByteValue.length)); 263 264 try (final WBWIRocksIterator it = wbwi.newIterator()) { 265 it.seekToFirst(); 266 final WBWIRocksIterator.WriteEntry actual = it.entry(); 267 assertThat(actual.equals(expected)).isTrue(); 268 assertThat(it.entry().hashCode() == expected.hashCode()).isTrue(); 269 } 270 } 271 } 272 273 @Test savePoints()274 public void savePoints() throws RocksDBException { 275 try (final Options options = new Options().setCreateIfMissing(true); 276 final RocksDB db = RocksDB.open(options, 277 dbFolder.getRoot().getAbsolutePath())) { 278 try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex(true); 279 final ReadOptions readOptions = new ReadOptions()) { 280 wbwi.put("k1".getBytes(), "v1".getBytes()); 281 wbwi.put("k2".getBytes(), "v2".getBytes()); 282 wbwi.put("k3".getBytes(), "v3".getBytes()); 283 284 assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k1")) 285 .isEqualTo("v1"); 286 assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k2")) 287 .isEqualTo("v2"); 288 assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k3")) 289 .isEqualTo("v3"); 290 291 292 wbwi.setSavePoint(); 293 294 wbwi.delete("k2".getBytes()); 295 wbwi.put("k3".getBytes(), "v3-2".getBytes()); 296 297 assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k2")) 298 .isNull(); 299 assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k3")) 300 .isEqualTo("v3-2"); 301 302 303 wbwi.setSavePoint(); 304 305 wbwi.put("k3".getBytes(), "v3-3".getBytes()); 306 wbwi.put("k4".getBytes(), "v4".getBytes()); 307 308 assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k3")) 309 .isEqualTo("v3-3"); 310 assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k4")) 311 .isEqualTo("v4"); 312 313 314 wbwi.rollbackToSavePoint(); 315 316 assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k2")) 317 .isNull(); 318 assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k3")) 319 .isEqualTo("v3-2"); 320 assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k4")) 321 .isNull(); 322 323 324 wbwi.rollbackToSavePoint(); 325 326 assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k1")) 327 .isEqualTo("v1"); 328 assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k2")) 329 .isEqualTo("v2"); 330 assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k3")) 331 .isEqualTo("v3"); 332 assertThat(getFromWriteBatchWithIndex(db, readOptions, wbwi, "k4")) 333 .isNull(); 334 } 335 } 336 } 337 338 @Test restorePoints()339 public void restorePoints() throws RocksDBException { 340 try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex()) { 341 342 wbwi.put("k1".getBytes(UTF_8), "v1".getBytes(UTF_8)); 343 wbwi.put("k2".getBytes(UTF_8), "v2".getBytes(UTF_8)); 344 345 wbwi.setSavePoint(); 346 347 wbwi.put("k1".getBytes(UTF_8), "123456789".getBytes(UTF_8)); 348 wbwi.delete("k2".getBytes(UTF_8)); 349 350 wbwi.rollbackToSavePoint(); 351 352 try(final DBOptions options = new DBOptions()) { 353 assertThat(wbwi.getFromBatch(options,"k1".getBytes(UTF_8))).isEqualTo("v1".getBytes()); 354 assertThat(wbwi.getFromBatch(options,"k2".getBytes(UTF_8))).isEqualTo("v2".getBytes()); 355 } 356 } 357 } 358 359 @Test(expected = RocksDBException.class) restorePoints_withoutSavePoints()360 public void restorePoints_withoutSavePoints() throws RocksDBException { 361 try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex()) { 362 wbwi.rollbackToSavePoint(); 363 } 364 } 365 366 @Test(expected = RocksDBException.class) restorePoints_withoutSavePoints_nested()367 public void restorePoints_withoutSavePoints_nested() throws RocksDBException { 368 try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex()) { 369 370 wbwi.setSavePoint(); 371 wbwi.rollbackToSavePoint(); 372 373 // without previous corresponding setSavePoint 374 wbwi.rollbackToSavePoint(); 375 } 376 } 377 378 @Test popSavePoint()379 public void popSavePoint() throws RocksDBException { 380 try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex()) { 381 382 wbwi.put("k1".getBytes(), "v1".getBytes()); 383 wbwi.put("k2".getBytes(), "v2".getBytes()); 384 385 wbwi.setSavePoint(); 386 387 wbwi.put("k1".getBytes(), "123456789".getBytes()); 388 wbwi.delete("k2".getBytes()); 389 390 wbwi.setSavePoint(); 391 392 wbwi.popSavePoint(); 393 394 wbwi.rollbackToSavePoint(); 395 396 try(final DBOptions options = new DBOptions()) { 397 assertThat(wbwi.getFromBatch(options,"k1".getBytes(UTF_8))).isEqualTo("v1".getBytes()); 398 assertThat(wbwi.getFromBatch(options,"k2".getBytes(UTF_8))).isEqualTo("v2".getBytes()); 399 } 400 } 401 } 402 403 @Test(expected = RocksDBException.class) popSavePoint_withoutSavePoints()404 public void popSavePoint_withoutSavePoints() throws RocksDBException { 405 try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex()) { 406 wbwi.popSavePoint(); 407 } 408 } 409 410 @Test(expected = RocksDBException.class) popSavePoint_withoutSavePoints_nested()411 public void popSavePoint_withoutSavePoints_nested() throws RocksDBException { 412 try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex()) { 413 414 wbwi.setSavePoint(); 415 wbwi.popSavePoint(); 416 417 // without previous corresponding setSavePoint 418 wbwi.popSavePoint(); 419 } 420 } 421 422 @Test maxBytes()423 public void maxBytes() throws RocksDBException { 424 try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex()) { 425 wbwi.setMaxBytes(19); 426 427 wbwi.put("k1".getBytes(), "v1".getBytes()); 428 } 429 } 430 431 @Test(expected = RocksDBException.class) maxBytes_over()432 public void maxBytes_over() throws RocksDBException { 433 try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex()) { 434 wbwi.setMaxBytes(1); 435 436 wbwi.put("k1".getBytes(), "v1".getBytes()); 437 } 438 } 439 440 @Test getWriteBatch()441 public void getWriteBatch() { 442 try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex()) { 443 444 final WriteBatch wb = wbwi.getWriteBatch(); 445 assertThat(wb).isNotNull(); 446 assertThat(wb.isOwningHandle()).isFalse(); 447 } 448 } 449 getFromWriteBatchWithIndex(final RocksDB db, final ReadOptions readOptions, final WriteBatchWithIndex wbwi, final String skey)450 private static String getFromWriteBatchWithIndex(final RocksDB db, 451 final ReadOptions readOptions, final WriteBatchWithIndex wbwi, 452 final String skey) { 453 final byte[] key = skey.getBytes(); 454 try (final RocksIterator baseIterator = db.newIterator(readOptions); 455 final RocksIterator iterator = wbwi.newIteratorWithBase(baseIterator)) { 456 iterator.seek(key); 457 458 // Arrays.equals(key, iterator.key()) ensures an exact match in Rocks, 459 // instead of a nearest match 460 return iterator.isValid() && 461 Arrays.equals(key, iterator.key()) ? 462 new String(iterator.value()) : null; 463 } 464 } 465 466 @Test getFromBatch()467 public void getFromBatch() throws RocksDBException { 468 final byte[] k1 = "k1".getBytes(); 469 final byte[] k2 = "k2".getBytes(); 470 final byte[] k3 = "k3".getBytes(); 471 final byte[] k4 = "k4".getBytes(); 472 473 final byte[] v1 = "v1".getBytes(); 474 final byte[] v2 = "v2".getBytes(); 475 final byte[] v3 = "v3".getBytes(); 476 477 try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex(true); 478 final DBOptions dbOptions = new DBOptions()) { 479 wbwi.put(k1, v1); 480 wbwi.put(k2, v2); 481 wbwi.put(k3, v3); 482 483 assertThat(wbwi.getFromBatch(dbOptions, k1)).isEqualTo(v1); 484 assertThat(wbwi.getFromBatch(dbOptions, k2)).isEqualTo(v2); 485 assertThat(wbwi.getFromBatch(dbOptions, k3)).isEqualTo(v3); 486 assertThat(wbwi.getFromBatch(dbOptions, k4)).isNull(); 487 488 wbwi.delete(k2); 489 490 assertThat(wbwi.getFromBatch(dbOptions, k2)).isNull(); 491 } 492 } 493 494 @Test getFromBatchAndDB()495 public void getFromBatchAndDB() throws RocksDBException { 496 final byte[] k1 = "k1".getBytes(); 497 final byte[] k2 = "k2".getBytes(); 498 final byte[] k3 = "k3".getBytes(); 499 final byte[] k4 = "k4".getBytes(); 500 501 final byte[] v1 = "v1".getBytes(); 502 final byte[] v2 = "v2".getBytes(); 503 final byte[] v3 = "v3".getBytes(); 504 final byte[] v4 = "v4".getBytes(); 505 506 try (final Options options = new Options().setCreateIfMissing(true); 507 final RocksDB db = RocksDB.open(options, 508 dbFolder.getRoot().getAbsolutePath())) { 509 510 db.put(k1, v1); 511 db.put(k2, v2); 512 db.put(k4, v4); 513 514 try (final WriteBatchWithIndex wbwi = new WriteBatchWithIndex(true); 515 final DBOptions dbOptions = new DBOptions(); 516 final ReadOptions readOptions = new ReadOptions()) { 517 518 assertThat(wbwi.getFromBatch(dbOptions, k1)).isNull(); 519 assertThat(wbwi.getFromBatch(dbOptions, k2)).isNull(); 520 assertThat(wbwi.getFromBatch(dbOptions, k4)).isNull(); 521 522 wbwi.put(k3, v3); 523 524 assertThat(wbwi.getFromBatch(dbOptions, k3)).isEqualTo(v3); 525 526 assertThat(wbwi.getFromBatchAndDB(db, readOptions, k1)).isEqualTo(v1); 527 assertThat(wbwi.getFromBatchAndDB(db, readOptions, k2)).isEqualTo(v2); 528 assertThat(wbwi.getFromBatchAndDB(db, readOptions, k3)).isEqualTo(v3); 529 assertThat(wbwi.getFromBatchAndDB(db, readOptions, k4)).isEqualTo(v4); 530 531 wbwi.delete(k4); 532 533 assertThat(wbwi.getFromBatchAndDB(db, readOptions, k4)).isNull(); 534 } 535 } 536 } toArray(final ByteBuffer buf)537 private byte[] toArray(final ByteBuffer buf) { 538 final byte[] ary = new byte[buf.remaining()]; 539 buf.get(ary); 540 return ary; 541 } 542 543 @Test deleteRange()544 public void deleteRange() throws RocksDBException { 545 try (final RocksDB db = RocksDB.open(dbFolder.getRoot().getAbsolutePath()); 546 final WriteBatch batch = new WriteBatch(); 547 final WriteOptions wOpt = new WriteOptions()) { 548 db.put("key1".getBytes(), "value".getBytes()); 549 db.put("key2".getBytes(), "12345678".getBytes()); 550 db.put("key3".getBytes(), "abcdefg".getBytes()); 551 db.put("key4".getBytes(), "xyz".getBytes()); 552 assertThat(db.get("key1".getBytes())).isEqualTo("value".getBytes()); 553 assertThat(db.get("key2".getBytes())).isEqualTo("12345678".getBytes()); 554 assertThat(db.get("key3".getBytes())).isEqualTo("abcdefg".getBytes()); 555 assertThat(db.get("key4".getBytes())).isEqualTo("xyz".getBytes()); 556 557 batch.deleteRange("key2".getBytes(), "key4".getBytes()); 558 db.write(wOpt, batch); 559 560 assertThat(db.get("key1".getBytes())).isEqualTo("value".getBytes()); 561 assertThat(db.get("key2".getBytes())).isNull(); 562 assertThat(db.get("key3".getBytes())).isNull(); 563 assertThat(db.get("key4".getBytes())).isEqualTo("xyz".getBytes()); 564 } 565 } 566 } 567