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 package org.rocksdb.util;
7 
8 import org.junit.ClassRule;
9 import org.junit.Rule;
10 import org.junit.Test;
11 import org.junit.rules.TemporaryFolder;
12 import org.rocksdb.*;
13 
14 import java.io.IOException;
15 import java.nio.ByteBuffer;
16 import java.nio.file.*;
17 import java.util.*;
18 
19 import static java.nio.charset.StandardCharsets.UTF_8;
20 import static org.junit.Assert.*;
21 import static org.rocksdb.util.ByteUtil.bytes;
22 
23 /**
24  * This is a direct port of various C++
25  * tests from db/comparator_db_test.cc
26  * and some code to adapt it to RocksJava
27  */
28 public class BytewiseComparatorTest {
29 
30   @ClassRule
31   public static final RocksNativeLibraryResource ROCKS_NATIVE_LIBRARY_RESOURCE =
32       new RocksNativeLibraryResource();
33 
34   @Rule
35   public TemporaryFolder dbFolder = new TemporaryFolder();
36 
37   private List<String> source_strings = Arrays.asList("b", "d", "f", "h", "j", "l");
38   private List<String> interleaving_strings = Arrays.asList("a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m");
39 
40   /**
41    * Open the database using the C++ BytewiseComparatorImpl
42    * and test the results against our Java BytewiseComparator
43    */
44   @Test
java_vs_cpp_bytewiseComparator()45   public void java_vs_cpp_bytewiseComparator()
46       throws IOException, RocksDBException {
47     for(int rand_seed = 301; rand_seed < 306; rand_seed++) {
48       final Path dbDir =
49           FileSystems.getDefault().getPath(dbFolder.newFolder().getAbsolutePath());
50       try(final RocksDB db = openDatabase(dbDir,
51           BuiltinComparator.BYTEWISE_COMPARATOR)) {
52 
53         final Random rnd = new Random(rand_seed);
54         try(final ComparatorOptions copt2 = new ComparatorOptions()
55             .setUseDirectBuffer(false);
56             final AbstractComparator comparator2 = new BytewiseComparator(copt2)) {
57           final java.util.Comparator<String> jComparator = toJavaComparator(comparator2);
58           doRandomIterationTest(
59               db,
60               jComparator,
61               rnd,
62               8, 100, 3
63           );
64         }
65       }
66     }
67   }
68 
69   /**
70    * Open the database using the Java BytewiseComparator
71    * and test the results against another Java BytewiseComparator
72    */
73   @Test
java_vs_java_bytewiseComparator()74   public void java_vs_java_bytewiseComparator()
75       throws IOException, RocksDBException {
76     for(int rand_seed = 301; rand_seed < 306; rand_seed++) {
77       final Path dbDir =
78           FileSystems.getDefault().getPath(dbFolder.newFolder().getAbsolutePath());
79       try(final ComparatorOptions copt = new ComparatorOptions()
80           .setUseDirectBuffer(false);
81           final AbstractComparator comparator = new BytewiseComparator(copt);
82           final RocksDB db = openDatabase(dbDir, comparator)) {
83 
84         final Random rnd = new Random(rand_seed);
85         try(final ComparatorOptions copt2 = new ComparatorOptions()
86             .setUseDirectBuffer(false);
87             final AbstractComparator comparator2 = new BytewiseComparator(copt2)) {
88           final java.util.Comparator<String> jComparator = toJavaComparator(comparator2);
89           doRandomIterationTest(
90               db,
91               jComparator,
92               rnd,
93               8, 100, 3
94           );
95         }
96       }
97     }
98   }
99 
100   /**
101    * Open the database using the C++ BytewiseComparatorImpl
102    * and test the results against our Java DirectBytewiseComparator
103    */
104   @Test
java_vs_cpp_directBytewiseComparator()105   public void java_vs_cpp_directBytewiseComparator()
106       throws IOException, RocksDBException {
107     for(int rand_seed = 301; rand_seed < 306; rand_seed++) {
108       final Path dbDir =
109           FileSystems.getDefault().getPath(dbFolder.newFolder().getAbsolutePath());
110       try(final RocksDB db = openDatabase(dbDir,
111           BuiltinComparator.BYTEWISE_COMPARATOR)) {
112 
113         final Random rnd = new Random(rand_seed);
114         try(final ComparatorOptions copt2 = new ComparatorOptions()
115               .setUseDirectBuffer(true);
116             final AbstractComparator comparator2 = new BytewiseComparator(copt2)) {
117           final java.util.Comparator<String> jComparator = toJavaComparator(comparator2);
118           doRandomIterationTest(
119               db,
120               jComparator,
121               rnd,
122               8, 100, 3
123           );
124         }
125       }
126     }
127   }
128 
129   /**
130    * Open the database using the Java DirectBytewiseComparator
131    * and test the results against another Java DirectBytewiseComparator
132    */
133   @Test
java_vs_java_directBytewiseComparator()134   public void java_vs_java_directBytewiseComparator()
135       throws IOException, RocksDBException {
136     for(int rand_seed = 301; rand_seed < 306; rand_seed++) {
137       final Path dbDir =
138           FileSystems.getDefault().getPath(dbFolder.newFolder().getAbsolutePath());
139       try (final ComparatorOptions copt = new ComparatorOptions()
140            .setUseDirectBuffer(true);
141           final AbstractComparator comparator = new BytewiseComparator(copt);
142           final RocksDB db = openDatabase(dbDir, comparator)) {
143 
144         final Random rnd = new Random(rand_seed);
145         try(final ComparatorOptions copt2 = new ComparatorOptions()
146               .setUseDirectBuffer(true);
147             final AbstractComparator comparator2 = new BytewiseComparator(copt2)) {
148           final java.util.Comparator<String> jComparator = toJavaComparator(comparator2);
149           doRandomIterationTest(
150               db,
151               jComparator,
152               rnd,
153               8, 100, 3
154           );
155         }
156       }
157     }
158   }
159 
160   /**
161    * Open the database using the C++ ReverseBytewiseComparatorImpl
162    * and test the results against our Java ReverseBytewiseComparator
163    */
164   @Test
java_vs_cpp_reverseBytewiseComparator()165   public void java_vs_cpp_reverseBytewiseComparator()
166       throws IOException, RocksDBException {
167     for(int rand_seed = 301; rand_seed < 306; rand_seed++) {
168       final Path dbDir =
169           FileSystems.getDefault().getPath(dbFolder.newFolder().getAbsolutePath());
170       try(final RocksDB db = openDatabase(dbDir,
171           BuiltinComparator.REVERSE_BYTEWISE_COMPARATOR)) {
172 
173         final Random rnd = new Random(rand_seed);
174         try(final ComparatorOptions copt2 = new ComparatorOptions()
175             .setUseDirectBuffer(false);
176             final AbstractComparator comparator2 = new ReverseBytewiseComparator(copt2)) {
177           final java.util.Comparator<String> jComparator = toJavaComparator(comparator2);
178           doRandomIterationTest(
179               db,
180               jComparator,
181               rnd,
182               8, 100, 3
183           );
184         }
185       }
186     }
187   }
188 
189   /**
190    * Open the database using the Java ReverseBytewiseComparator
191    * and test the results against another Java ReverseBytewiseComparator
192    */
193   @Test
java_vs_java_reverseBytewiseComparator()194   public void java_vs_java_reverseBytewiseComparator()
195       throws IOException, RocksDBException {
196     for(int rand_seed = 301; rand_seed < 306; rand_seed++) {
197       final Path dbDir =
198           FileSystems.getDefault().getPath(dbFolder.newFolder().getAbsolutePath());
199       try (final ComparatorOptions copt = new ComparatorOptions()
200            .setUseDirectBuffer(false);
201            final AbstractComparator comparator = new ReverseBytewiseComparator(copt);
202            final RocksDB db = openDatabase(dbDir, comparator)) {
203 
204         final Random rnd = new Random(rand_seed);
205         try(final ComparatorOptions copt2 = new ComparatorOptions()
206             .setUseDirectBuffer(false);
207             final AbstractComparator comparator2 = new ReverseBytewiseComparator(copt2)) {
208           final java.util.Comparator<String> jComparator = toJavaComparator(comparator2);
209           doRandomIterationTest(
210               db,
211               jComparator,
212               rnd,
213               8, 100, 3
214           );
215         }
216       }
217     }
218   }
219 
doRandomIterationTest( final RocksDB db, final java.util.Comparator<String> javaComparator, final Random rnd, final int num_writes, final int num_iter_ops, final int num_trigger_flush)220   private void doRandomIterationTest(
221       final RocksDB db, final java.util.Comparator<String> javaComparator,
222       final Random rnd,
223       final int num_writes, final int num_iter_ops,
224       final int num_trigger_flush) throws RocksDBException {
225 
226     final TreeMap<String, String> map = new TreeMap<>(javaComparator);
227 
228     try (final FlushOptions flushOptions = new FlushOptions();
229          final WriteOptions writeOptions = new WriteOptions()) {
230       for (int i = 0; i < num_writes; i++) {
231         if (num_trigger_flush > 0 && i != 0 && i % num_trigger_flush == 0) {
232           db.flush(flushOptions);
233         }
234 
235         final int type = rnd.nextInt(2);
236         final int index = rnd.nextInt(source_strings.size());
237         final String key = source_strings.get(index);
238         switch (type) {
239           case 0:
240             // put
241             map.put(key, key);
242             db.put(writeOptions, bytes(key), bytes(key));
243             break;
244           case 1:
245             // delete
246             if (map.containsKey(key)) {
247               map.remove(key);
248             }
249             db.delete(writeOptions, bytes(key));
250             break;
251 
252           default:
253             fail("Should not be able to generate random outside range 1..2");
254         }
255       }
256     }
257 
258     try (final ReadOptions readOptions = new ReadOptions();
259          final RocksIterator iter = db.newIterator(readOptions)) {
260       final KVIter<String, String> result_iter = new KVIter<>(map);
261 
262       boolean is_valid = false;
263       for (int i = 0; i < num_iter_ops; i++) {
264         // Random walk and make sure iter and result_iter returns the
265         // same key and value
266         final int type = rnd.nextInt(7);
267         iter.status();
268         switch (type) {
269           case 0:
270             // Seek to First
271             iter.seekToFirst();
272             result_iter.seekToFirst();
273             break;
274           case 1:
275             // Seek to last
276             iter.seekToLast();
277             result_iter.seekToLast();
278             break;
279           case 2: {
280             // Seek to random (existing or non-existing) key
281             final int key_idx = rnd.nextInt(interleaving_strings.size());
282             final String key = interleaving_strings.get(key_idx);
283             iter.seek(bytes(key));
284             result_iter.seek(bytes(key));
285             break;
286           }
287           case 3: {
288             // SeekForPrev to random (existing or non-existing) key
289             final int key_idx = rnd.nextInt(interleaving_strings.size());
290             final String key = interleaving_strings.get(key_idx);
291             iter.seekForPrev(bytes(key));
292             result_iter.seekForPrev(bytes(key));
293             break;
294           }
295           case 4:
296             // Next
297             if (is_valid) {
298               iter.next();
299               result_iter.next();
300             } else {
301               continue;
302             }
303             break;
304           case 5:
305             // Prev
306             if (is_valid) {
307               iter.prev();
308               result_iter.prev();
309             } else {
310               continue;
311             }
312             break;
313           default: {
314             assert (type == 6);
315             final int key_idx = rnd.nextInt(source_strings.size());
316             final String key = source_strings.get(key_idx);
317             final byte[] result = db.get(readOptions, bytes(key));
318             if (!map.containsKey(key)) {
319               assertNull(result);
320             } else {
321               assertArrayEquals(bytes(map.get(key)), result);
322             }
323             break;
324           }
325         }
326 
327         assertEquals(result_iter.isValid(), iter.isValid());
328 
329         is_valid = iter.isValid();
330 
331         if (is_valid) {
332           assertArrayEquals(bytes(result_iter.key()), iter.key());
333 
334           //note that calling value on a non-valid iterator from the Java API
335           //results in a SIGSEGV
336           assertArrayEquals(bytes(result_iter.value()), iter.value());
337         }
338       }
339     }
340   }
341 
342   /**
343    * Open the database using a C++ Comparator
344    */
openDatabase( final Path dbDir, final BuiltinComparator cppComparator)345   private RocksDB openDatabase(
346       final Path dbDir, final BuiltinComparator cppComparator)
347       throws IOException, RocksDBException {
348     final Options options = new Options()
349         .setCreateIfMissing(true)
350         .setComparator(cppComparator);
351     return RocksDB.open(options, dbDir.toAbsolutePath().toString());
352   }
353 
354   /**
355    * Open the database using a Java Comparator
356    */
openDatabase( final Path dbDir, final AbstractComparator javaComparator)357   private RocksDB openDatabase(
358       final Path dbDir,
359       final AbstractComparator javaComparator)
360       throws IOException, RocksDBException {
361     final Options options = new Options()
362         .setCreateIfMissing(true)
363         .setComparator(javaComparator);
364     return RocksDB.open(options, dbDir.toAbsolutePath().toString());
365   }
366 
toJavaComparator( final AbstractComparator rocksComparator)367   private java.util.Comparator<String> toJavaComparator(
368       final AbstractComparator rocksComparator) {
369     return new java.util.Comparator<String>() {
370       @Override
371       public int compare(final String s1, final String s2) {
372         final ByteBuffer bufS1;
373         final ByteBuffer bufS2;
374         if (rocksComparator.usingDirectBuffers()) {
375           bufS1 = ByteBuffer.allocateDirect(s1.length());
376           bufS2 = ByteBuffer.allocateDirect(s2.length());
377         } else {
378           bufS1 = ByteBuffer.allocate(s1.length());
379           bufS2 = ByteBuffer.allocate(s2.length());
380         }
381         bufS1.put(bytes(s1));
382         bufS1.flip();
383         bufS2.put(bytes(s2));
384         bufS2.flip();
385         return rocksComparator.compare(bufS1, bufS2);
386       }
387     };
388   }
389 
390   private static class KVIter<K, V> implements RocksIteratorInterface {
391 
392     private final List<Map.Entry<K, V>> entries;
393     private final java.util.Comparator<? super K> comparator;
394     private int offset = -1;
395 
396     private int lastPrefixMatchIdx = -1;
397     private int lastPrefixMatch = 0;
398 
399     public KVIter(final TreeMap<K, V> map) {
400       this.entries = new ArrayList<>();
401       entries.addAll(map.entrySet());
402       this.comparator = map.comparator();
403     }
404 
405 
406     @Override
407     public boolean isValid() {
408       return offset > -1 && offset < entries.size();
409     }
410 
411     @Override
412     public void seekToFirst() {
413       offset = 0;
414     }
415 
416     @Override
417     public void seekToLast() {
418       offset = entries.size() - 1;
419     }
420 
421     @SuppressWarnings("unchecked")
422     @Override
423     public void seek(final byte[] target) {
424       for(offset = 0; offset < entries.size(); offset++) {
425         if(comparator.compare(entries.get(offset).getKey(),
426             (K)new String(target, UTF_8)) >= 0) {
427           return;
428         }
429       }
430     }
431 
432     @SuppressWarnings("unchecked")
433     @Override
434     public void seekForPrev(final byte[] target) {
435       for(offset = entries.size()-1; offset >= 0; offset--) {
436         if(comparator.compare(entries.get(offset).getKey(),
437             (K)new String(target, UTF_8)) <= 0) {
438           return;
439         }
440       }
441     }
442 
443     /**
444      * Is `a` a prefix of `b`
445      *
446      * @return The length of the matching prefix, or 0 if it is not a prefix
447      */
448     private int isPrefix(final byte[] a, final byte[] b) {
449       if(b.length >= a.length) {
450         for(int i = 0; i < a.length; i++) {
451           if(a[i] != b[i]) {
452             return i;
453           }
454         }
455         return a.length;
456       } else {
457         return 0;
458       }
459     }
460 
461     @Override
462     public void next() {
463       if(offset < entries.size()) {
464         offset++;
465       }
466     }
467 
468     @Override
469     public void prev() {
470       if(offset >= 0) {
471         offset--;
472       }
473     }
474 
475     @Override
476     public void status() throws RocksDBException {
477       if(offset < 0 || offset >= entries.size()) {
478         throw new RocksDBException("Index out of bounds. Size is: " +
479             entries.size() + ", offset is: " + offset);
480       }
481     }
482 
483     @SuppressWarnings("unchecked")
484     public K key() {
485       if(!isValid()) {
486         if(entries.isEmpty()) {
487           return (K)"";
488         } else if(offset == -1){
489           return entries.get(0).getKey();
490         } else if(offset == entries.size()) {
491           return entries.get(offset - 1).getKey();
492         } else {
493           return (K)"";
494         }
495       } else {
496         return entries.get(offset).getKey();
497       }
498     }
499 
500     @SuppressWarnings("unchecked")
501     public V value() {
502       if(!isValid()) {
503         return (V)"";
504       } else {
505         return entries.get(offset).getValue();
506       }
507     }
508 
509     @Override
510     public void seek(ByteBuffer target) {
511       throw new IllegalAccessError("Not implemented");
512     }
513 
514     @Override
515     public void seekForPrev(ByteBuffer target) {
516       throw new IllegalAccessError("Not implemented");
517     }
518   }
519 }
520