ceph/src/rocksdb/java/src/test/java/org/rocksdb/util/BytewiseComparatorTest.java

   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
  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
  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
 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
 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
 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
 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
 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(8);
 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           case 6:
 314             // Refresh
 315             iter.refresh();
 316             result_iter.refresh();
 317             iter.seekToFirst();
 318             result_iter.seekToFirst();
 319             break;
 320           default: {
 321             assert (type == 7);
 322             final int key_idx = rnd.nextInt(source_strings.size());
 323             final String key = source_strings.get(key_idx);
 324             final byte[] result = db.get(readOptions, bytes(key));
 325             if (!map.containsKey(key)) {
 326               assertNull(result);
 327             } else {
 328               assertArrayEquals(bytes(map.get(key)), result);
 329             }
 330             break;
 331           }
 332         }
 333
 334         assertEquals(result_iter.isValid(), iter.isValid());
 335
 336         is_valid = iter.isValid();
 337
 338         if (is_valid) {
 339           assertArrayEquals(bytes(result_iter.key()), iter.key());
 340
 341           //note that calling value on a non-valid iterator from the Java API
 342           //results in a SIGSEGV
 343           assertArrayEquals(bytes(result_iter.value()), iter.value());
 344         }
 345       }
 346     }
 347   }
 348
 349   /**
 350    * Open the database using a C++ Comparator
 351    */
 352   private RocksDB openDatabase(
 353       final Path dbDir, final BuiltinComparator cppComparator)
 354       throws IOException, RocksDBException {
 355     final Options options = new Options()
 356         .setCreateIfMissing(true)
 357         .setComparator(cppComparator);
 358     return RocksDB.open(options, dbDir.toAbsolutePath().toString());
 359   }
 360
 361   /**
 362    * Open the database using a Java Comparator
 363    */
 364   private RocksDB openDatabase(
 365       final Path dbDir,
 366       final AbstractComparator javaComparator)
 367       throws IOException, RocksDBException {
 368     final Options options = new Options()
 369         .setCreateIfMissing(true)
 370         .setComparator(javaComparator);
 371     return RocksDB.open(options, dbDir.toAbsolutePath().toString());
 372   }
 373
 374   private java.util.Comparator<String> toJavaComparator(
 375       final AbstractComparator rocksComparator) {
 376     return new java.util.Comparator<String>() {
 377       @Override
 378       public int compare(final String s1, final String s2) {
 379         final ByteBuffer bufS1;
 380         final ByteBuffer bufS2;
 381         if (rocksComparator.usingDirectBuffers()) {
 382           bufS1 = ByteBuffer.allocateDirect(s1.length());
 383           bufS2 = ByteBuffer.allocateDirect(s2.length());
 384         } else {
 385           bufS1 = ByteBuffer.allocate(s1.length());
 386           bufS2 = ByteBuffer.allocate(s2.length());
 387         }
 388         bufS1.put(bytes(s1));
 389         bufS1.flip();
 390         bufS2.put(bytes(s2));
 391         bufS2.flip();
 392         return rocksComparator.compare(bufS1, bufS2);
 393       }
 394     };
 395   }
 396
 397   private static class KVIter<K, V> implements RocksIteratorInterface {
 398
 399     private final List<Map.Entry<K, V>> entries;
 400     private final java.util.Comparator<? super K> comparator;
 401     private int offset = -1;
 402
 403     private int lastPrefixMatchIdx = -1;
 404     private int lastPrefixMatch = 0;
 405
 406     public KVIter(final TreeMap<K, V> map) {
 407       this.entries = new ArrayList<>();
 408       entries.addAll(map.entrySet());
 409       this.comparator = map.comparator();
 410     }
 411
 412
 413     @Override
 414     public boolean isValid() {
 415       return offset > -1 && offset < entries.size();
 416     }
 417
 418     @Override
 419     public void seekToFirst() {
 420       offset = 0;
 421     }
 422
 423     @Override
 424     public void seekToLast() {
 425       offset = entries.size() - 1;
 426     }
 427
 428     @SuppressWarnings("unchecked")
 429     @Override
 430     public void seek(final byte[] target) {
 431       for(offset = 0; offset < entries.size(); offset++) {
 432         if(comparator.compare(entries.get(offset).getKey(),
 433             (K)new String(target, UTF_8)) >= 0) {
 434           return;
 435         }
 436       }
 437     }
 438
 439     @SuppressWarnings("unchecked")
 440     @Override
 441     public void seekForPrev(final byte[] target) {
 442       for(offset = entries.size()-1; offset >= 0; offset--) {
 443         if(comparator.compare(entries.get(offset).getKey(),
 444             (K)new String(target, UTF_8)) <= 0) {
 445           return;
 446         }
 447       }
 448     }
 449
 450     /**
 451      * Is `a` a prefix of `b`
 452      *
 453      * @return The length of the matching prefix, or 0 if it is not a prefix
 454      */
 455     private int isPrefix(final byte[] a, final byte[] b) {
 456       if(b.length >= a.length) {
 457         for(int i = 0; i < a.length; i++) {
 458           if(a[i] != b[i]) {
 459             return i;
 460           }
 461         }
 462         return a.length;
 463       } else {
 464         return 0;
 465       }
 466     }
 467
 468     @Override
 469     public void next() {
 470       if(offset < entries.size()) {
 471         offset++;
 472       }
 473     }
 474
 475     @Override
 476     public void prev() {
 477       if(offset >= 0) {
 478         offset--;
 479       }
 480     }
 481
 482     @Override
 483     public void refresh() throws RocksDBException {
 484       offset = -1;
 485     }
 486
 487     @Override
 488     public void status() throws RocksDBException {
 489       if(offset < 0 || offset >= entries.size()) {
 490         throw new RocksDBException("Index out of bounds. Size is: " +
 491             entries.size() + ", offset is: " + offset);
 492       }
 493     }
 494
 495     @SuppressWarnings("unchecked")
 496     public K key() {
 497       if(!isValid()) {
 498         if(entries.isEmpty()) {
 499           return (K)"";
 500         } else if(offset == -1){
 501           return entries.get(0).getKey();
 502         } else if(offset == entries.size()) {
 503           return entries.get(offset - 1).getKey();
 504         } else {
 505           return (K)"";
 506         }
 507       } else {
 508         return entries.get(offset).getKey();
 509       }
 510     }
 511
 512     @SuppressWarnings("unchecked")
 513     public V value() {
 514       if(!isValid()) {
 515         return (V)"";
 516       } else {
 517         return entries.get(offset).getValue();
 518       }
 519     }
 520
 521     @Override
 522     public void seek(ByteBuffer target) {
 523       throw new IllegalAccessError("Not implemented");
 524     }
 525
 526     @Override
 527     public void seekForPrev(ByteBuffer target) {
 528       throw new IllegalAccessError("Not implemented");
 529     }
 530   }
 531 }