1 package org.rocksdb.util;
2 
3 import org.junit.ClassRule;
4 import org.junit.Rule;
5 import org.junit.Test;
6 import org.junit.rules.TemporaryFolder;
7 import org.junit.runner.RunWith;
8 import org.junit.runners.Parameterized;
9 import org.junit.runners.Parameterized.Parameter;
10 import org.junit.runners.Parameterized.Parameters;
11 import org.rocksdb.*;
12 
13 import java.io.IOException;
14 import java.nio.ByteBuffer;
15 import java.nio.file.*;
16 import java.util.Arrays;
17 
18 import static org.assertj.core.api.Assertions.assertThat;
19 
20 @RunWith(Parameterized.class)
21 public class JNIComparatorTest {
22 
23   @Parameters(name = "{0}")
parameters()24   public static Iterable<Object[]> parameters() {
25     return Arrays.asList(new Object[][] {
26         { "bytewise_non-direct", BuiltinComparator.BYTEWISE_COMPARATOR, false },
27         { "bytewise_direct", BuiltinComparator.BYTEWISE_COMPARATOR, true },
28         { "reverse-bytewise_non-direct", BuiltinComparator.REVERSE_BYTEWISE_COMPARATOR, false },
29         { "reverse-bytewise_direct", BuiltinComparator.REVERSE_BYTEWISE_COMPARATOR, true },
30     });
31   }
32 
33   @Parameter(0)
34   public String name;
35 
36   @Parameter(1)
37   public BuiltinComparator builtinComparator;
38 
39   @Parameter(2)
40   public boolean useDirectBuffer;
41 
42   @ClassRule
43   public static final RocksNativeLibraryResource ROCKS_NATIVE_LIBRARY_RESOURCE =
44       new RocksNativeLibraryResource();
45 
46   @Rule
47   public TemporaryFolder dbFolder = new TemporaryFolder();
48 
49   private static final int MIN = Short.MIN_VALUE - 1;
50   private static final int MAX = Short.MAX_VALUE + 1;
51 
52   @Test
java_comparator_equals_cpp_comparator()53   public void java_comparator_equals_cpp_comparator() throws RocksDBException, IOException {
54     final int[] javaKeys;
55     try (final ComparatorOptions comparatorOptions = new ComparatorOptions();
56          final AbstractComparator comparator = builtinComparator == BuiltinComparator.BYTEWISE_COMPARATOR
57              ? new BytewiseComparator(comparatorOptions)
58              : new ReverseBytewiseComparator(comparatorOptions)) {
59       final Path javaDbDir =
60           FileSystems.getDefault().getPath(dbFolder.newFolder().getAbsolutePath());
61       storeWithJavaComparator(javaDbDir, comparator);
62       javaKeys = readAllWithJavaComparator(javaDbDir, comparator);
63     }
64 
65     final Path cppDbDir =
66         FileSystems.getDefault().getPath(dbFolder.newFolder().getAbsolutePath());
67     storeWithCppComparator(cppDbDir, builtinComparator);
68     final int[] cppKeys =
69         readAllWithCppComparator(cppDbDir, builtinComparator);
70 
71     assertThat(javaKeys).isEqualTo(cppKeys);
72   }
73 
storeWithJavaComparator(final Path dir, final AbstractComparator comparator)74   private void storeWithJavaComparator(final Path dir,
75       final AbstractComparator comparator) throws RocksDBException {
76     final ByteBuffer buf = ByteBuffer.allocate(4);
77     try (final Options options = new Options()
78              .setCreateIfMissing(true)
79              .setComparator(comparator);
80          final RocksDB db =
81              RocksDB.open(options, dir.toAbsolutePath().toString())) {
82       for (int i = MIN; i < MAX; i++) {
83         buf.putInt(i);
84         buf.flip();
85 
86         db.put(buf.array(), buf.array());
87 
88         buf.clear();
89       }
90     }
91   }
92 
storeWithCppComparator(final Path dir, final BuiltinComparator builtinComparator)93   private void storeWithCppComparator(final Path dir,
94       final BuiltinComparator builtinComparator) throws RocksDBException {
95     try (final Options options = new Options()
96              .setCreateIfMissing(true)
97              .setComparator(builtinComparator);
98          final RocksDB db =
99              RocksDB.open(options, dir.toAbsolutePath().toString())) {
100 
101       final ByteBuffer buf = ByteBuffer.allocate(4);
102       for (int i = MIN; i < MAX; i++) {
103         buf.putInt(i);
104         buf.flip();
105 
106         db.put(buf.array(), buf.array());
107 
108         buf.clear();
109       }
110     }
111   }
112 
readAllWithJavaComparator(final Path dir, final AbstractComparator comparator)113   private int[] readAllWithJavaComparator(final Path dir,
114       final AbstractComparator comparator) throws RocksDBException {
115     try (final Options options = new Options()
116         .setCreateIfMissing(true)
117         .setComparator(comparator);
118          final RocksDB db =
119              RocksDB.open(options, dir.toAbsolutePath().toString())) {
120 
121       try (final RocksIterator it = db.newIterator()) {
122         it.seekToFirst();
123 
124         final ByteBuffer buf = ByteBuffer.allocate(4);
125         final int[] keys = new int[MAX - MIN];
126         int idx = 0;
127         while (it.isValid()) {
128           buf.put(it.key());
129           buf.flip();
130 
131           final int thisKey = buf.getInt();
132           keys[idx++] = thisKey;
133 
134           buf.clear();
135 
136           it.next();
137         }
138 
139         return keys;
140       }
141     }
142   }
143 
readAllWithCppComparator(final Path dir, final BuiltinComparator comparator)144   private int[] readAllWithCppComparator(final Path dir,
145       final BuiltinComparator comparator) throws RocksDBException {
146     try (final Options options = new Options()
147         .setCreateIfMissing(true)
148         .setComparator(comparator);
149          final RocksDB db =
150              RocksDB.open(options, dir.toAbsolutePath().toString())) {
151 
152       try (final RocksIterator it = db.newIterator()) {
153         it.seekToFirst();
154 
155         final ByteBuffer buf = ByteBuffer.allocate(4);
156         final int[] keys = new int[MAX - MIN];
157         int idx = 0;
158         while (it.isValid()) {
159           buf.put(it.key());
160           buf.flip();
161 
162           final int thisKey = buf.getInt();
163           keys[idx++] = thisKey;
164 
165           buf.clear();
166 
167           it.next();
168         }
169 
170         return keys;
171       }
172     }
173   }
174 }
175