]>
Commit | Line | Data |
---|---|---|
7c673cae | 1 | // Copyright (c) 2011-present, Facebook, Inc. All rights reserved. |
11fdf7f2 TL |
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). | |
7c673cae FG |
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 < 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 | } |