ceph/src/rocksdb/java/src/test/java/org/rocksdb/util/IntComparatorTest.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.BeforeClass;
   9 import org.junit.ClassRule;
  10 import org.junit.Rule;
  11 import org.junit.Test;
  12 import org.junit.rules.TemporaryFolder;
  13 import org.junit.runner.RunWith;
  14 import org.junit.runners.Parameterized;
  15 import org.junit.runners.Parameterized.Parameter;
  16 import org.junit.runners.Parameterized.Parameters;
  17 import org.rocksdb.*;
  18
  19 import java.nio.ByteBuffer;
  20 import java.nio.file.*;
  21 import java.util.ArrayList;
  22 import java.util.Arrays;
  23 import java.util.List;
  24 import java.util.Random;
  25
  26 import static java.nio.charset.StandardCharsets.UTF_8;
  27 import static org.assertj.core.api.Assertions.assertThat;
  28
  29 /**
  30  * Tests for IntComparator, but more generally
  31  * also for rocksdb::ComparatorJniCallback implementation.
  32  */
  33 @RunWith(Parameterized.class)
  34 public class IntComparatorTest {
  35
  36   // test with 500 random integer keys
  37   private static final int TOTAL_KEYS = 500;
  38   private static final byte[][] keys = new byte[TOTAL_KEYS][4];
  39
  40   @BeforeClass
  41   public static void prepareKeys() {
  42     final ByteBuffer buf = ByteBuffer.allocate(4);
  43     final Random random = new Random();
  44     for (int i = 0; i < TOTAL_KEYS; i++) {
  45       final int ri = random.nextInt();
  46       buf.putInt(ri);
  47       buf.flip();
  48       final byte[] key = buf.array();
  49
  50       // does key already exist (avoid duplicates)
  51       if (keyExists(key, i)) {
  52         i--; // loop round and generate a different key
  53       } else {
  54         System.arraycopy(key, 0, keys[i], 0, 4);
  55       }
  56     }
  57   }
  58
  59   private static boolean keyExists(final byte[] key, final int limit) {
  60     for (int j = 0; j < limit; j++) {
  61       if (Arrays.equals(key, keys[j])) {
  62         return true;
  63       }
  64     }
  65     return false;
  66   }
  67
  68   @Parameters(name = "{0}")
  69   public static Iterable<Object[]> parameters() {
  70     return Arrays.asList(new Object[][] {
  71         { "non-direct_reused64_mutex", false, 64, ReusedSynchronisationType.MUTEX },
  72         { "direct_reused64_mutex", true, 64, ReusedSynchronisationType.MUTEX },
  73         { "non-direct_reused64_adaptive-mutex", false, 64, ReusedSynchronisationType.ADAPTIVE_MUTEX },
  74         { "direct_reused64_adaptive-mutex", true, 64, ReusedSynchronisationType.ADAPTIVE_MUTEX },
  75         { "non-direct_reused64_thread-local", false, 64, ReusedSynchronisationType.THREAD_LOCAL },
  76         { "direct_reused64_thread-local", true, 64, ReusedSynchronisationType.THREAD_LOCAL },
  77         { "non-direct_noreuse", false, -1, null },
  78         { "direct_noreuse", true, -1, null }
  79     });
  80   }
  81
  82   @Parameter(0)
  83   public String name;
  84
  85   @Parameter(1)
  86   public boolean useDirectBuffer;
  87
  88   @Parameter(2)
  89   public int maxReusedBufferSize;
  90
  91   @Parameter(3)
  92   public ReusedSynchronisationType reusedSynchronisationType;
  93
  94   @ClassRule
  95   public static final RocksNativeLibraryResource ROCKS_NATIVE_LIBRARY_RESOURCE =
  96       new RocksNativeLibraryResource();
  97
  98   @Rule
  99   public TemporaryFolder dbFolder = new TemporaryFolder();
 100
 101
 102   @Test
 103   public void javaComparatorDefaultCf() throws RocksDBException {
 104     try (final ComparatorOptions options = new ComparatorOptions()
 105         .setUseDirectBuffer(useDirectBuffer)
 106         .setMaxReusedBufferSize(maxReusedBufferSize)
 107         // if reusedSynchronisationType == null we assume that maxReusedBufferSize <= 0 and so we just set ADAPTIVE_MUTEX, even though it won't be used
 108         .setReusedSynchronisationType(reusedSynchronisationType == null ? ReusedSynchronisationType.ADAPTIVE_MUTEX : reusedSynchronisationType);
 109         final IntComparator comparator = new IntComparator(options)) {
 110
 111       // test the round-tripability of keys written and read with the Comparator
 112       testRoundtrip(FileSystems.getDefault().getPath(
 113           dbFolder.getRoot().getAbsolutePath()), comparator);
 114     }
 115   }
 116
 117   @Test
 118   public void javaComparatorNamedCf() throws RocksDBException {
 119     try (final ComparatorOptions options = new ComparatorOptions()
 120         .setUseDirectBuffer(useDirectBuffer)
 121         .setMaxReusedBufferSize(maxReusedBufferSize)
 122         // if reusedSynchronisationType == null we assume that maxReusedBufferSize <= 0 and so we just set ADAPTIVE_MUTEX, even though it won't be used
 123         .setReusedSynchronisationType(reusedSynchronisationType == null ? ReusedSynchronisationType.ADAPTIVE_MUTEX : reusedSynchronisationType);
 124       final IntComparator comparator = new IntComparator(options)) {
 125
 126       // test the round-tripability of keys written and read with the Comparator
 127       testRoundtripCf(FileSystems.getDefault().getPath(
 128           dbFolder.getRoot().getAbsolutePath()), comparator);
 129     }
 130   }
 131
 132   /**
 133    * Test which stores random keys into the database
 134    * using an {@link IntComparator}
 135    * it then checks that these keys are read back in
 136    * ascending order
 137    *
 138    * @param db_path A path where we can store database
 139    *                files temporarily
 140    *
 141    * @param comparator the comparator
 142    *
 143    * @throws RocksDBException if a database error happens.
 144    */
 145   private void testRoundtrip(final Path db_path,
 146       final AbstractComparator comparator) throws RocksDBException {
 147     try (final Options opt = new Options()
 148              .setCreateIfMissing(true)
 149              .setComparator(comparator)) {
 150
 151       // store TOTAL_KEYS into the db
 152       try (final RocksDB db = RocksDB.open(opt, db_path.toString())) {
 153         for (int i = 0; i < TOTAL_KEYS; i++) {
 154               db.put(keys[i], "value".getBytes(UTF_8));
 155         }
 156       }
 157
 158       // re-open db and read from start to end
 159       // integer keys should be in ascending
 160       // order as defined by IntComparator
 161       final ByteBuffer key = ByteBuffer.allocate(4);
 162       try (final RocksDB db = RocksDB.open(opt, db_path.toString());
 163            final RocksIterator it = db.newIterator()) {
 164         it.seekToFirst();
 165         int lastKey = Integer.MIN_VALUE;
 166         int count = 0;
 167         for (it.seekToFirst(); it.isValid(); it.next()) {
 168           key.put(it.key());
 169           key.flip();
 170           final int thisKey = key.getInt();
 171           key.clear();
 172           assertThat(thisKey).isGreaterThan(lastKey);
 173           lastKey = thisKey;
 174           count++;
 175         }
 176         assertThat(count).isEqualTo(TOTAL_KEYS);
 177       }
 178     }
 179   }
 180
 181   /**
 182    * Test which stores random keys into a column family
 183    * in the database
 184    * using an {@link IntComparator}
 185    * it then checks that these keys are read back in
 186    * ascending order
 187    *
 188    * @param db_path A path where we can store database
 189    *                files temporarily
 190    *
 191    * @param comparator the comparator
 192    *
 193    * @throws RocksDBException if a database error happens.
 194    */
 195   private void testRoundtripCf(final Path db_path,
 196       final AbstractComparator comparator) throws RocksDBException {
 197
 198     final List<ColumnFamilyDescriptor> cfDescriptors = Arrays.asList(
 199         new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY),
 200         new ColumnFamilyDescriptor("new_cf".getBytes(),
 201             new ColumnFamilyOptions()
 202                 .setComparator(comparator))
 203     );
 204
 205     final List<ColumnFamilyHandle> cfHandles = new ArrayList<>();
 206
 207     try (final DBOptions opt = new DBOptions()
 208         .setCreateIfMissing(true)
 209         .setCreateMissingColumnFamilies(true)) {
 210
 211       try (final RocksDB db = RocksDB.open(opt, db_path.toString(),
 212           cfDescriptors, cfHandles)) {
 213         try {
 214           assertThat(cfDescriptors.size()).isEqualTo(2);
 215           assertThat(cfHandles.size()).isEqualTo(2);
 216
 217           for (int i = 0; i < TOTAL_KEYS; i++) {
 218             db.put(cfHandles.get(1), keys[i], "value".getBytes(UTF_8));
 219           }
 220         } finally {
 221           for (final ColumnFamilyHandle cfHandle : cfHandles) {
 222             cfHandle.close();
 223           }
 224           cfHandles.clear();
 225         }
 226       }
 227
 228       // re-open db and read from start to end
 229       // integer keys should be in ascending
 230       // order as defined by SimpleIntComparator
 231       final ByteBuffer key = ByteBuffer.allocate(4);
 232       try (final RocksDB db = RocksDB.open(opt, db_path.toString(),
 233           cfDescriptors, cfHandles);
 234            final RocksIterator it = db.newIterator(cfHandles.get(1))) {
 235         try {
 236           assertThat(cfDescriptors.size()).isEqualTo(2);
 237           assertThat(cfHandles.size()).isEqualTo(2);
 238
 239           it.seekToFirst();
 240           int lastKey = Integer.MIN_VALUE;
 241           int count = 0;
 242           for (it.seekToFirst(); it.isValid(); it.next()) {
 243             key.put(it.key());
 244             key.flip();
 245             final int thisKey = key.getInt();
 246             key.clear();
 247             assertThat(thisKey).isGreaterThan(lastKey);
 248             lastKey = thisKey;
 249             count++;
 250           }
 251
 252           assertThat(count).isEqualTo(TOTAL_KEYS);
 253
 254         } finally {
 255           for (final ColumnFamilyHandle cfHandle : cfHandles) {
 256             cfHandle.close();
 257           }
 258           cfHandles.clear();
 259           for (final ColumnFamilyDescriptor cfDescriptor : cfDescriptors) {
 260             cfDescriptor.getOptions().close();
 261           }
 262         }
 263       }
 264     }
 265   }
 266 }