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