ceph/src/rocksdb/java/src/test/java/org/rocksdb/AbstractComparatorTest.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;
   7
   8 import java.io.IOException;
   9 import java.nio.file.*;
  10 import java.util.ArrayList;
  11 import java.util.Arrays;
  12 import java.util.List;
  13 import java.util.Random;
  14
  15 import static org.assertj.core.api.Assertions.assertThat;
  16 import static org.rocksdb.Types.byteToInt;
  17 import static org.rocksdb.Types.intToByte;
  18
  19 /**
  20  * Abstract tests for both Comparator and DirectComparator
  21  */
  22 public abstract class AbstractComparatorTest {
  23
  24   /**
  25    * Get a comparator which will expect Integer keys
  26    * and determine an ascending order
  27    *
  28    * @return An integer ascending order key comparator
  29    */
  30   public abstract AbstractComparator getAscendingIntKeyComparator();
  31
  32   /**
  33    * Test which stores random keys into the database
  34    * using an @see getAscendingIntKeyComparator
  35    * it then checks that these keys are read back in
  36    * ascending order
  37    *
  38    * @param db_path A path where we can store database
  39    *                files temporarily
  40    *
  41    * @throws java.io.IOException if IO error happens.
  42    */
  43   public void testRoundtrip(final Path db_path) throws IOException,
  44       RocksDBException {
  45     try (final AbstractComparator comparator = getAscendingIntKeyComparator();
  46          final Options opt = new Options()
  47              .setCreateIfMissing(true)
  48              .setComparator(comparator)) {
  49
  50       // store 10,000 random integer keys
  51       final int ITERATIONS = 10000;
  52       try (final RocksDB db = RocksDB.open(opt, db_path.toString())) {
  53         final Random random = new Random();
  54         for (int i = 0; i < ITERATIONS; i++) {
  55           final byte key[] = intToByte(random.nextInt());
  56           // does key already exist (avoid duplicates)
  57           if (i > 0 && db.get(key) != null) {
  58             i--; // generate a different key
  59           } else {
  60             db.put(key, "value".getBytes());
  61           }
  62         }
  63       }
  64
  65       // re-open db and read from start to end
  66       // integer keys should be in ascending
  67       // order as defined by SimpleIntComparator
  68       try (final RocksDB db = RocksDB.open(opt, db_path.toString());
  69            final RocksIterator it = db.newIterator()) {
  70         it.seekToFirst();
  71         int lastKey = Integer.MIN_VALUE;
  72         int count = 0;
  73         for (it.seekToFirst(); it.isValid(); it.next()) {
  74           final int thisKey = byteToInt(it.key());
  75           assertThat(thisKey).isGreaterThan(lastKey);
  76           lastKey = thisKey;
  77           count++;
  78         }
  79         assertThat(count).isEqualTo(ITERATIONS);
  80       }
  81     }
  82   }
  83
  84   /**
  85    * Test which stores random keys into a column family
  86    * in the database
  87    * using an @see getAscendingIntKeyComparator
  88    * it then checks that these keys are read back in
  89    * ascending order
  90    *
  91    * @param db_path A path where we can store database
  92    *                files temporarily
  93    *
  94    * @throws java.io.IOException if IO error happens.
  95    */
  96   public void testRoundtripCf(final Path db_path) throws IOException,
  97       RocksDBException {
  98
  99     try(final AbstractComparator comparator = getAscendingIntKeyComparator()) {
 100       final List<ColumnFamilyDescriptor> cfDescriptors = Arrays.asList(
 101           new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY),
 102           new ColumnFamilyDescriptor("new_cf".getBytes(),
 103               new ColumnFamilyOptions().setComparator(comparator))
 104       );
 105
 106       final List<ColumnFamilyHandle> cfHandles = new ArrayList<>();
 107
 108       try (final DBOptions opt = new DBOptions().
 109           setCreateIfMissing(true).
 110           setCreateMissingColumnFamilies(true)) {
 111
 112         // store 10,000 random integer keys
 113         final int ITERATIONS = 10000;
 114
 115         try (final RocksDB db = RocksDB.open(opt, db_path.toString(),
 116             cfDescriptors, cfHandles)) {
 117           try {
 118             assertThat(cfDescriptors.size()).isEqualTo(2);
 119             assertThat(cfHandles.size()).isEqualTo(2);
 120
 121             final Random random = new Random();
 122             for (int i = 0; i < ITERATIONS; i++) {
 123               final byte key[] = intToByte(random.nextInt());
 124               if (i > 0 && db.get(cfHandles.get(1), key) != null) {
 125                 // does key already exist (avoid duplicates)
 126                 i--; // generate a different key
 127               } else {
 128                 db.put(cfHandles.get(1), key, "value".getBytes());
 129               }
 130             }
 131           } finally {
 132             for (final ColumnFamilyHandle handle : cfHandles) {
 133               handle.close();
 134             }
 135           }
 136           cfHandles.clear();
 137         }
 138
 139         // re-open db and read from start to end
 140         // integer keys should be in ascending
 141         // order as defined by SimpleIntComparator
 142         try (final RocksDB db = RocksDB.open(opt, db_path.toString(),
 143             cfDescriptors, cfHandles);
 144              final RocksIterator it = db.newIterator(cfHandles.get(1))) {
 145           try {
 146             assertThat(cfDescriptors.size()).isEqualTo(2);
 147             assertThat(cfHandles.size()).isEqualTo(2);
 148
 149             it.seekToFirst();
 150             int lastKey = Integer.MIN_VALUE;
 151             int count = 0;
 152             for (it.seekToFirst(); it.isValid(); it.next()) {
 153               final int thisKey = byteToInt(it.key());
 154               assertThat(thisKey).isGreaterThan(lastKey);
 155               lastKey = thisKey;
 156               count++;
 157             }
 158
 159             assertThat(count).isEqualTo(ITERATIONS);
 160
 161           } finally {
 162             for (final ColumnFamilyHandle handle : cfHandles) {
 163               handle.close();
 164             }
 165           }
 166           cfHandles.clear();
 167         }
 168       }
 169     }
 170   }
 171
 172   /**
 173    * Compares integer keys
 174    * so that they are in ascending order
 175    *
 176    * @param a 4-bytes representing an integer key
 177    * @param b 4-bytes representing an integer key
 178    *
 179    * @return negative if a &lt; b, 0 if a == b, positive otherwise
 180    */
 181   protected final int compareIntKeys(final byte[] a, final byte[] b) {
 182
 183     final int iA = byteToInt(a);
 184     final int iB = byteToInt(b);
 185
 186     // protect against int key calculation overflow
 187     final double diff = (double)iA - iB;
 188     final int result;
 189     if (diff < Integer.MIN_VALUE) {
 190       result = Integer.MIN_VALUE;
 191     } else if(diff > Integer.MAX_VALUE) {
 192       result = Integer.MAX_VALUE;
 193     } else {
 194       result = (int)diff;
 195     }
 196
 197     return result;
 198   }
 199 }