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 package org.rocksdb;
10 
11 import static java.nio.charset.StandardCharsets.UTF_8;
12 import static org.assertj.core.api.Assertions.assertThat;
13 import static org.rocksdb.util.CapturingWriteBatchHandler.Action.DELETE;
14 import static org.rocksdb.util.CapturingWriteBatchHandler.Action.DELETE_RANGE;
15 import static org.rocksdb.util.CapturingWriteBatchHandler.Action.LOG;
16 import static org.rocksdb.util.CapturingWriteBatchHandler.Action.MERGE;
17 import static org.rocksdb.util.CapturingWriteBatchHandler.Action.PUT;
18 import static org.rocksdb.util.CapturingWriteBatchHandler.Action.SINGLE_DELETE;
19 
20 import java.io.UnsupportedEncodingException;
21 import java.nio.ByteBuffer;
22 import org.junit.ClassRule;
23 import org.junit.Rule;
24 import org.junit.Test;
25 import org.junit.rules.TemporaryFolder;
26 import org.rocksdb.util.CapturingWriteBatchHandler;
27 import org.rocksdb.util.CapturingWriteBatchHandler.Event;
28 import org.rocksdb.util.WriteBatchGetter;
29 
30 /**
31  * This class mimics the db/write_batch_test.cc
32  * in the c++ rocksdb library.
33  */
34 public class WriteBatchTest {
35   @ClassRule
36   public static final RocksNativeLibraryResource ROCKS_NATIVE_LIBRARY_RESOURCE =
37       new RocksNativeLibraryResource();
38 
39   @Rule
40   public TemporaryFolder dbFolder = new TemporaryFolder();
41 
42   @Test
emptyWriteBatch()43   public void emptyWriteBatch() {
44     try (final WriteBatch batch = new WriteBatch()) {
45       assertThat(batch.count()).isEqualTo(0);
46     }
47   }
48 
49   @Test
multipleBatchOperations()50   public void multipleBatchOperations()
51       throws RocksDBException {
52 
53     final byte[] foo = "foo".getBytes(UTF_8);
54     final byte[] bar = "bar".getBytes(UTF_8);
55     final byte[] box = "box".getBytes(UTF_8);
56     final byte[] baz = "baz".getBytes(UTF_8);
57     final byte[] boo = "boo".getBytes(UTF_8);
58     final byte[] hoo = "hoo".getBytes(UTF_8);
59     final byte[] hello = "hello".getBytes(UTF_8);
60 
61     try (final WriteBatch batch = new WriteBatch()) {
62       batch.put(foo, bar);
63       batch.delete(box);
64       batch.put(baz, boo);
65       batch.merge(baz, hoo);
66       batch.singleDelete(foo);
67       batch.deleteRange(baz, foo);
68       batch.putLogData(hello);
69 
70       try(final CapturingWriteBatchHandler handler =
71               new CapturingWriteBatchHandler()) {
72         batch.iterate(handler);
73 
74         assertThat(handler.getEvents().size()).isEqualTo(7);
75 
76         assertThat(handler.getEvents().get(0)).isEqualTo(new Event(PUT, foo, bar));
77         assertThat(handler.getEvents().get(1)).isEqualTo(new Event(DELETE, box, null));
78         assertThat(handler.getEvents().get(2)).isEqualTo(new Event(PUT, baz, boo));
79         assertThat(handler.getEvents().get(3)).isEqualTo(new Event(MERGE, baz, hoo));
80         assertThat(handler.getEvents().get(4)).isEqualTo(new Event(SINGLE_DELETE, foo, null));
81         assertThat(handler.getEvents().get(5)).isEqualTo(new Event(DELETE_RANGE, baz, foo));
82         assertThat(handler.getEvents().get(6)).isEqualTo(new Event(LOG, null, hello));
83       }
84     }
85   }
86 
87   @Test
multipleBatchOperationsDirect()88   public void multipleBatchOperationsDirect()
89       throws UnsupportedEncodingException, RocksDBException {
90     try (WriteBatch batch = new WriteBatch()) {
91       ByteBuffer key = ByteBuffer.allocateDirect(16);
92       ByteBuffer value = ByteBuffer.allocateDirect(16);
93       key.put("foo".getBytes("US-ASCII")).flip();
94       value.put("bar".getBytes("US-ASCII")).flip();
95       batch.put(key, value);
96       assertThat(key.position()).isEqualTo(3);
97       assertThat(key.limit()).isEqualTo(3);
98       assertThat(value.position()).isEqualTo(3);
99       assertThat(value.limit()).isEqualTo(3);
100 
101       key.clear();
102       key.put("box".getBytes("US-ASCII")).flip();
103       batch.remove(key);
104       assertThat(key.position()).isEqualTo(3);
105       assertThat(key.limit()).isEqualTo(3);
106 
107       batch.put("baz".getBytes("US-ASCII"), "boo".getBytes("US-ASCII"));
108 
109       WriteBatchTestInternalHelper.setSequence(batch, 100);
110       assertThat(WriteBatchTestInternalHelper.sequence(batch)).isNotNull().isEqualTo(100);
111       assertThat(batch.count()).isEqualTo(3);
112       assertThat(new String(getContents(batch), "US-ASCII"))
113           .isEqualTo("Put(baz, boo)@102"
114               + "Delete(box)@101"
115               + "Put(foo, bar)@100");
116     }
117   }
118 
119   @Test
testAppendOperation()120   public void testAppendOperation()
121       throws RocksDBException {
122     try (final WriteBatch b1 = new WriteBatch();
123          final WriteBatch b2 = new WriteBatch()) {
124       WriteBatchTestInternalHelper.setSequence(b1, 200);
125       WriteBatchTestInternalHelper.setSequence(b2, 300);
126       WriteBatchTestInternalHelper.append(b1, b2);
127       assertThat(getContents(b1).length).isEqualTo(0);
128       assertThat(b1.count()).isEqualTo(0);
129       b2.put("a".getBytes(UTF_8), "va".getBytes(UTF_8));
130       WriteBatchTestInternalHelper.append(b1, b2);
131       assertThat("Put(a, va)@200".equals(new String(getContents(b1),
132           UTF_8)));
133       assertThat(b1.count()).isEqualTo(1);
134       b2.clear();
135       b2.put("b".getBytes(UTF_8), "vb".getBytes(UTF_8));
136       WriteBatchTestInternalHelper.append(b1, b2);
137       assertThat(("Put(a, va)@200" +
138           "Put(b, vb)@201")
139           .equals(new String(getContents(b1), UTF_8)));
140       assertThat(b1.count()).isEqualTo(2);
141       b2.delete("foo".getBytes(UTF_8));
142       WriteBatchTestInternalHelper.append(b1, b2);
143       assertThat(("Put(a, va)@200" +
144           "Put(b, vb)@202" +
145           "Put(b, vb)@201" +
146           "Delete(foo)@203")
147           .equals(new String(getContents(b1), UTF_8)));
148       assertThat(b1.count()).isEqualTo(4);
149     }
150   }
151 
152   @Test
blobOperation()153   public void blobOperation()
154       throws RocksDBException {
155     try (final WriteBatch batch = new WriteBatch()) {
156       batch.put("k1".getBytes(UTF_8), "v1".getBytes(UTF_8));
157       batch.put("k2".getBytes(UTF_8), "v2".getBytes(UTF_8));
158       batch.put("k3".getBytes(UTF_8), "v3".getBytes(UTF_8));
159       batch.putLogData("blob1".getBytes(UTF_8));
160       batch.delete("k2".getBytes(UTF_8));
161       batch.putLogData("blob2".getBytes(UTF_8));
162       batch.merge("foo".getBytes(UTF_8), "bar".getBytes(UTF_8));
163       assertThat(batch.count()).isEqualTo(5);
164       assertThat(("Merge(foo, bar)@4" +
165           "Put(k1, v1)@0" +
166           "Delete(k2)@3" +
167           "Put(k2, v2)@1" +
168           "Put(k3, v3)@2")
169           .equals(new String(getContents(batch), UTF_8)));
170     }
171   }
172 
173   @Test
savePoints()174   public void savePoints()
175       throws RocksDBException {
176     try (final WriteBatch batch = new WriteBatch()) {
177       batch.put("k1".getBytes(UTF_8), "v1".getBytes(UTF_8));
178       batch.put("k2".getBytes(UTF_8), "v2".getBytes(UTF_8));
179       batch.put("k3".getBytes(UTF_8), "v3".getBytes(UTF_8));
180 
181       assertThat(getFromWriteBatch(batch, "k1")).isEqualTo("v1");
182       assertThat(getFromWriteBatch(batch, "k2")).isEqualTo("v2");
183       assertThat(getFromWriteBatch(batch, "k3")).isEqualTo("v3");
184 
185       batch.setSavePoint();
186 
187       batch.delete("k2".getBytes(UTF_8));
188       batch.put("k3".getBytes(UTF_8), "v3-2".getBytes(UTF_8));
189 
190       assertThat(getFromWriteBatch(batch, "k2")).isNull();
191       assertThat(getFromWriteBatch(batch, "k3")).isEqualTo("v3-2");
192 
193 
194       batch.setSavePoint();
195 
196       batch.put("k3".getBytes(UTF_8), "v3-3".getBytes(UTF_8));
197       batch.put("k4".getBytes(UTF_8), "v4".getBytes(UTF_8));
198 
199       assertThat(getFromWriteBatch(batch, "k3")).isEqualTo("v3-3");
200       assertThat(getFromWriteBatch(batch, "k4")).isEqualTo("v4");
201 
202 
203       batch.rollbackToSavePoint();
204 
205       assertThat(getFromWriteBatch(batch, "k2")).isNull();
206       assertThat(getFromWriteBatch(batch, "k3")).isEqualTo("v3-2");
207       assertThat(getFromWriteBatch(batch, "k4")).isNull();
208 
209 
210       batch.rollbackToSavePoint();
211 
212       assertThat(getFromWriteBatch(batch, "k1")).isEqualTo("v1");
213       assertThat(getFromWriteBatch(batch, "k2")).isEqualTo("v2");
214       assertThat(getFromWriteBatch(batch, "k3")).isEqualTo("v3");
215       assertThat(getFromWriteBatch(batch, "k4")).isNull();
216     }
217   }
218 
219   @Test
deleteRange()220   public void deleteRange() throws RocksDBException {
221     try (final RocksDB db = RocksDB.open(dbFolder.getRoot().getAbsolutePath());
222          final WriteBatch batch = new WriteBatch();
223          final WriteOptions wOpt = new WriteOptions()) {
224       db.put("key1".getBytes(), "value".getBytes());
225       db.put("key2".getBytes(), "12345678".getBytes());
226       db.put("key3".getBytes(), "abcdefg".getBytes());
227       db.put("key4".getBytes(), "xyz".getBytes());
228       assertThat(db.get("key1".getBytes())).isEqualTo("value".getBytes());
229       assertThat(db.get("key2".getBytes())).isEqualTo("12345678".getBytes());
230       assertThat(db.get("key3".getBytes())).isEqualTo("abcdefg".getBytes());
231       assertThat(db.get("key4".getBytes())).isEqualTo("xyz".getBytes());
232 
233       batch.deleteRange("key2".getBytes(), "key4".getBytes());
234       db.write(wOpt, batch);
235 
236       assertThat(db.get("key1".getBytes())).isEqualTo("value".getBytes());
237       assertThat(db.get("key2".getBytes())).isNull();
238       assertThat(db.get("key3".getBytes())).isNull();
239       assertThat(db.get("key4".getBytes())).isEqualTo("xyz".getBytes());
240     }
241   }
242 
243   @Test
restorePoints()244   public void restorePoints() throws RocksDBException {
245     try (final WriteBatch batch = new WriteBatch()) {
246 
247       batch.put("k1".getBytes(), "v1".getBytes());
248       batch.put("k2".getBytes(), "v2".getBytes());
249 
250       batch.setSavePoint();
251 
252       batch.put("k1".getBytes(), "123456789".getBytes());
253       batch.delete("k2".getBytes());
254 
255       batch.rollbackToSavePoint();
256 
257       try(final CapturingWriteBatchHandler handler = new CapturingWriteBatchHandler()) {
258         batch.iterate(handler);
259 
260         assertThat(handler.getEvents().size()).isEqualTo(2);
261         assertThat(handler.getEvents().get(0)).isEqualTo(new Event(PUT, "k1".getBytes(), "v1".getBytes()));
262         assertThat(handler.getEvents().get(1)).isEqualTo(new Event(PUT, "k2".getBytes(), "v2".getBytes()));
263       }
264     }
265   }
266 
267   @Test(expected = RocksDBException.class)
restorePoints_withoutSavePoints()268   public void restorePoints_withoutSavePoints() throws RocksDBException {
269     try (final WriteBatch batch = new WriteBatch()) {
270       batch.rollbackToSavePoint();
271     }
272   }
273 
274   @Test(expected = RocksDBException.class)
restorePoints_withoutSavePoints_nested()275   public void restorePoints_withoutSavePoints_nested() throws RocksDBException {
276     try (final WriteBatch batch = new WriteBatch()) {
277 
278       batch.setSavePoint();
279       batch.rollbackToSavePoint();
280 
281       // without previous corresponding setSavePoint
282       batch.rollbackToSavePoint();
283     }
284   }
285 
286   @Test
popSavePoint()287   public void popSavePoint() throws RocksDBException {
288     try (final WriteBatch batch = new WriteBatch()) {
289 
290       batch.put("k1".getBytes(), "v1".getBytes());
291       batch.put("k2".getBytes(), "v2".getBytes());
292 
293       batch.setSavePoint();
294 
295       batch.put("k1".getBytes(), "123456789".getBytes());
296       batch.delete("k2".getBytes());
297 
298       batch.setSavePoint();
299 
300       batch.popSavePoint();
301 
302       batch.rollbackToSavePoint();
303 
304       try(final CapturingWriteBatchHandler handler = new CapturingWriteBatchHandler()) {
305         batch.iterate(handler);
306 
307         assertThat(handler.getEvents().size()).isEqualTo(2);
308         assertThat(handler.getEvents().get(0)).isEqualTo(new Event(PUT, "k1".getBytes(), "v1".getBytes()));
309         assertThat(handler.getEvents().get(1)).isEqualTo(new Event(PUT, "k2".getBytes(), "v2".getBytes()));
310       }
311     }
312   }
313 
314   @Test(expected = RocksDBException.class)
popSavePoint_withoutSavePoints()315   public void popSavePoint_withoutSavePoints() throws RocksDBException {
316     try (final WriteBatch batch = new WriteBatch()) {
317       batch.popSavePoint();
318     }
319   }
320 
321   @Test(expected = RocksDBException.class)
popSavePoint_withoutSavePoints_nested()322   public void popSavePoint_withoutSavePoints_nested() throws RocksDBException {
323     try (final WriteBatch batch = new WriteBatch()) {
324 
325       batch.setSavePoint();
326       batch.popSavePoint();
327 
328       // without previous corresponding setSavePoint
329       batch.popSavePoint();
330     }
331   }
332 
333   @Test
maxBytes()334   public void maxBytes() throws RocksDBException {
335     try (final WriteBatch batch = new WriteBatch()) {
336       batch.setMaxBytes(19);
337 
338       batch.put("k1".getBytes(), "v1".getBytes());
339     }
340   }
341 
342   @Test(expected = RocksDBException.class)
maxBytes_over()343   public void maxBytes_over() throws RocksDBException {
344     try (final WriteBatch batch = new WriteBatch()) {
345       batch.setMaxBytes(1);
346 
347       batch.put("k1".getBytes(), "v1".getBytes());
348     }
349   }
350 
351   @Test
data()352   public void data() throws RocksDBException {
353     try (final WriteBatch batch1 = new WriteBatch()) {
354       batch1.delete("k0".getBytes());
355       batch1.put("k1".getBytes(), "v1".getBytes());
356       batch1.put("k2".getBytes(), "v2".getBytes());
357       batch1.put("k3".getBytes(), "v3".getBytes());
358       batch1.putLogData("log1".getBytes());
359       batch1.merge("k2".getBytes(), "v22".getBytes());
360       batch1.delete("k3".getBytes());
361 
362       final byte[] serialized = batch1.data();
363 
364       try(final WriteBatch batch2 = new WriteBatch(serialized)) {
365         assertThat(batch2.count()).isEqualTo(batch1.count());
366 
367         try(final CapturingWriteBatchHandler handler1 = new CapturingWriteBatchHandler()) {
368           batch1.iterate(handler1);
369 
370           try (final CapturingWriteBatchHandler handler2 = new CapturingWriteBatchHandler()) {
371             batch2.iterate(handler2);
372 
373             assertThat(handler1.getEvents().equals(handler2.getEvents())).isTrue();
374           }
375         }
376       }
377     }
378   }
379 
380   @Test
dataSize()381   public void dataSize() throws RocksDBException {
382     try (final WriteBatch batch = new WriteBatch()) {
383       batch.put("k1".getBytes(), "v1".getBytes());
384 
385       assertThat(batch.getDataSize()).isEqualTo(19);
386     }
387   }
388 
389   @Test
hasPut()390   public void hasPut() throws RocksDBException {
391     try (final WriteBatch batch = new WriteBatch()) {
392       assertThat(batch.hasPut()).isFalse();
393 
394       batch.put("k1".getBytes(), "v1".getBytes());
395 
396       assertThat(batch.hasPut()).isTrue();
397     }
398   }
399 
400   @Test
hasDelete()401   public void hasDelete() throws RocksDBException {
402     try (final WriteBatch batch = new WriteBatch()) {
403       assertThat(batch.hasDelete()).isFalse();
404 
405       batch.delete("k1".getBytes());
406 
407       assertThat(batch.hasDelete()).isTrue();
408     }
409   }
410 
411   @Test
hasSingleDelete()412   public void hasSingleDelete() throws RocksDBException {
413     try (final WriteBatch batch = new WriteBatch()) {
414       assertThat(batch.hasSingleDelete()).isFalse();
415 
416       batch.singleDelete("k1".getBytes());
417 
418       assertThat(batch.hasSingleDelete()).isTrue();
419     }
420   }
421 
422   @Test
hasDeleteRange()423   public void hasDeleteRange() throws RocksDBException {
424     try (final WriteBatch batch = new WriteBatch()) {
425       assertThat(batch.hasDeleteRange()).isFalse();
426 
427       batch.deleteRange("k1".getBytes(), "k2".getBytes());
428 
429       assertThat(batch.hasDeleteRange()).isTrue();
430     }
431   }
432 
433   @Test
hasBeginPrepareRange()434   public void hasBeginPrepareRange() throws RocksDBException {
435     try (final WriteBatch batch = new WriteBatch()) {
436       assertThat(batch.hasBeginPrepare()).isFalse();
437     }
438   }
439 
440   @Test
hasEndPrepareRange()441   public void hasEndPrepareRange() throws RocksDBException {
442     try (final WriteBatch batch = new WriteBatch()) {
443       assertThat(batch.hasEndPrepare()).isFalse();
444     }
445   }
446 
447   @Test
hasCommit()448   public void hasCommit() throws RocksDBException {
449     try (final WriteBatch batch = new WriteBatch()) {
450       assertThat(batch.hasCommit()).isFalse();
451     }
452   }
453 
454   @Test
hasRollback()455   public void hasRollback() throws RocksDBException {
456     try (final WriteBatch batch = new WriteBatch()) {
457       assertThat(batch.hasRollback()).isFalse();
458     }
459   }
460 
461   @Test
walTerminationPoint()462   public void walTerminationPoint() throws RocksDBException {
463     try (final WriteBatch batch = new WriteBatch()) {
464       WriteBatch.SavePoint walTerminationPoint = batch.getWalTerminationPoint();
465       assertThat(walTerminationPoint.isCleared()).isTrue();
466 
467       batch.put("k1".getBytes(UTF_8), "v1".getBytes(UTF_8));
468 
469       batch.markWalTerminationPoint();
470 
471       walTerminationPoint = batch.getWalTerminationPoint();
472       assertThat(walTerminationPoint.getSize()).isEqualTo(19);
473       assertThat(walTerminationPoint.getCount()).isEqualTo(1);
474       assertThat(walTerminationPoint.getContentFlags()).isEqualTo(2);
475     }
476   }
477 
478   @Test
getWriteBatch()479   public void getWriteBatch() {
480     try (final WriteBatch batch = new WriteBatch()) {
481       assertThat(batch.getWriteBatch()).isEqualTo(batch);
482     }
483   }
484 
getContents(final WriteBatch wb)485   static byte[] getContents(final WriteBatch wb) {
486     return getContents(wb.nativeHandle_);
487   }
488 
getFromWriteBatch(final WriteBatch wb, final String key)489   static String getFromWriteBatch(final WriteBatch wb, final String key)
490       throws RocksDBException {
491     final WriteBatchGetter getter =
492         new WriteBatchGetter(key.getBytes(UTF_8));
493     wb.iterate(getter);
494     if(getter.getValue() != null) {
495       return new String(getter.getValue(), UTF_8);
496     } else {
497       return null;
498     }
499   }
500 
getContents(final long writeBatchHandle)501   private static native byte[] getContents(final long writeBatchHandle);
502 }
503 
504 /**
505  * Package-private class which provides java api to access
506  * c++ WriteBatchInternal.
507  */
508 class WriteBatchTestInternalHelper {
setSequence(final WriteBatch wb, final long sn)509   static void setSequence(final WriteBatch wb, final long sn) {
510     setSequence(wb.nativeHandle_, sn);
511   }
512 
sequence(final WriteBatch wb)513   static long sequence(final WriteBatch wb) {
514     return sequence(wb.nativeHandle_);
515   }
516 
append(final WriteBatch wb1, final WriteBatch wb2)517   static void append(final WriteBatch wb1, final WriteBatch wb2) {
518     append(wb1.nativeHandle_, wb2.nativeHandle_);
519   }
520 
setSequence(final long writeBatchHandle, final long sn)521   private static native void setSequence(final long writeBatchHandle,
522       final long sn);
523 
sequence(final long writeBatchHandle)524   private static native long sequence(final long writeBatchHandle);
525 
append(final long writeBatchHandle1, final long writeBatchHandle2)526   private static native void append(final long writeBatchHandle1,
527       final long writeBatchHandle2);
528 }
529