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