andreacavalli/rocksdb
0
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

java / jni io_uring support (#9224)

Browse Source 
Summary:
Existing multiGet() in java calls multi_get_helper() which then calls DB::std::vector MultiGet(). This doesn't take advantage of io_uring.

This change adds another JNI level method that runs a parallel code path using the DB::void MultiGet(), using ByteBuffers at the JNI level. We call it multiGetDirect(). In addition to using the io_uring path, this code internally returns pinned slices which we can copy out of into our direct byte buffers; this should reduce the overall number of copies in the code path to/from Java. Some jmh benchmark runs (100k keys, 1000 key multiGet) suggest that for value sizes > 1k, we see about a 20% performance improvement, although performance is slightly reduced for small value sizes, there's a little bit more overhead in the JNI methods.

Closes https://github.com/facebook/rocksdb/issues/8407

Pull Request resolved: https://github.com/facebook/rocksdb/pull/9224

Reviewed By: mrambacher

Differential Revision: D32951754

Pulled By: jay-zhuang

fbshipit-source-id: 1f70df7334be2b6c42a9c8f92725f67c71631690
This commit is contained in:
Alan Paxton 2021-12-15 18:07:37 -08:00 committed by Facebook GitHub Bot
parent 7ac3a5d406
commit c1ec0b28eb
10 changed files with 1086 additions and 73 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
java/CMakeLists.txt
View File

@ -112,6 +112,7 @@ set(JAVA_MAIN_CLASSES
src/main/java/org/rocksdb/BlockBasedTableConfig.java src/main/java/org/rocksdb/BlockBasedTableConfig.java
src/main/java/org/rocksdb/BloomFilter.java src/main/java/org/rocksdb/BloomFilter.java
src/main/java/org/rocksdb/BuiltinComparator.java src/main/java/org/rocksdb/BuiltinComparator.java
src/main/java/org/rocksdb/ByteBufferGetStatus.java
src/main/java/org/rocksdb/Cache.java src/main/java/org/rocksdb/Cache.java
src/main/java/org/rocksdb/CassandraCompactionFilter.java src/main/java/org/rocksdb/CassandraCompactionFilter.java
src/main/java/org/rocksdb/CassandraValueMergeOperator.java src/main/java/org/rocksdb/CassandraValueMergeOperator.java

6
java/jmh/README.md
View File

@ -6,6 +6,12 @@ These are micro-benchmarks for RocksJava functionality, using [JMH (Java Microbe
**Note**: This uses a specific build of RocksDB that is set in the `<version>` element of the `dependencies` section of the `pom.xml` file. If you are testing local changes you should build and install a SNAPSHOT version of rocksdbjni, and update the `pom.xml` of rocksdbjni-jmh file to test with this. **Note**: This uses a specific build of RocksDB that is set in the `<version>` element of the `dependencies` section of the `pom.xml` file. If you are testing local changes you should build and install a SNAPSHOT version of rocksdbjni, and update the `pom.xml` of rocksdbjni-jmh file to test with this.
For instance, this is how to install the OSX jar you just built for 6.26.0
```bash
$ mvn install:install-file -Dfile=./java/target/rocksdbjni-6.26.0-SNAPSHOT-osx.jar -DgroupId=org.rocksdb -DartifactId=rocksdbjni -Dversion=6.26.0-SNAPSHOT -Dpackaging=jar
```
```bash ```bash
$ mvn package $ mvn package
``` ```

2
java/jmh/pom.xml
View File

@ -50,7 +50,7 @@
<dependency> <dependency>
<groupId>org.rocksdb</groupId> <groupId>org.rocksdb</groupId>
<artifactId>rocksdbjni</artifactId> <artifactId>rocksdbjni</artifactId>
<version>6.6.0-SNAPSHOT</version> <version>6.27.0-SNAPSHOT</version>
</dependency> </dependency>
<dependency> <dependency>

110
java/jmh/src/main/java/org/rocksdb/jmh/MultiGetBenchmarks.java
View File

@ -6,23 +6,26 @@
*/ */
package org.rocksdb.jmh; package org.rocksdb.jmh;
import org.openjdk.jmh.annotations.*;
import org.rocksdb.*;
import org.rocksdb.util.FileUtils;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;
import static org.rocksdb.util.KVUtils.ba; import static org.rocksdb.util.KVUtils.ba;
import static org.rocksdb.util.KVUtils.keys; import static org.rocksdb.util.KVUtils.keys;
@State(Scope.Benchmark) import java.io.IOException;
public class MultiGetBenchmarks { import java.nio.ByteBuffer;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;
import org.openjdk.jmh.annotations.*;
import org.openjdk.jmh.runner.Runner;
import org.openjdk.jmh.runner.RunnerException;
import org.openjdk.jmh.runner.options.OptionsBuilder;
import org.rocksdb.*;
import org.rocksdb.util.FileUtils;
@State(Scope.Thread)
public class MultiGetBenchmarks {
@Param({ @Param({
"no_column_family", "no_column_family",
"1_column_family", "1_column_family",
@ -31,8 +34,7 @@ public class MultiGetBenchmarks {
}) })
String columnFamilyTestType; String columnFamilyTestType;
@Param("100000") @Param({"10000", "25000", "100000"}) int keyCount;
int keyCount;
@Param({ @Param({
"10", "10",
@ -42,6 +44,9 @@ public class MultiGetBenchmarks {
}) })
int multiGetSize; int multiGetSize;
@Param({"16", "64", "250", "1000", "4000", "16000"}) int valueSize;
@Param({"16"}) int keySize; // big enough
Path dbDir; Path dbDir;
DBOptions options; DBOptions options;
int cfs = 0; // number of column families int cfs = 0; // number of column families
@ -85,7 +90,8 @@ public class MultiGetBenchmarks {
// store initial data for retrieving via get // store initial data for retrieving via get
for (int i = 0; i < cfs; i++) { for (int i = 0; i < cfs; i++) {
for (int j = 0; j < keyCount; j++) { for (int j = 0; j < keyCount; j++) {
db.put(cfHandles[i], ba("key" + j), ba("value" + j)); final byte[] paddedValue = Arrays.copyOf(ba("value" + j), valueSize);
db.put(cfHandles[i], ba("key" + j), paddedValue);
} }
} }
@ -149,10 +155,78 @@ public class MultiGetBenchmarks {
} }
} }
ByteBuffer keysBuffer;
ByteBuffer valuesBuffer;
List<ByteBuffer> valueBuffersList;
List<ByteBuffer> keyBuffersList;
@Setup
public void allocateSliceBuffers() {
keysBuffer = ByteBuffer.allocateDirect(keyCount * valueSize);
valuesBuffer = ByteBuffer.allocateDirect(keyCount * valueSize);
valueBuffersList = new ArrayList<>();
keyBuffersList = new ArrayList<>();
for (int i = 0; i < keyCount; i++) {
valueBuffersList.add(valuesBuffer.slice());
valuesBuffer.position(i * valueSize);
keyBuffersList.add(keysBuffer.slice());
keysBuffer.position(i * keySize);
}
}
@TearDown
public void freeSliceBuffers() {
valueBuffersList.clear();
}
@Benchmark @Benchmark
public List<byte[]> multiGet10() throws RocksDBException { public List<byte[]> multiGet10() throws RocksDBException {
final int fromKeyIdx = next(multiGetSize, keyCount); final int fromKeyIdx = next(multiGetSize, keyCount);
final List<byte[]> keys = keys(fromKeyIdx, fromKeyIdx + multiGetSize); if (fromKeyIdx >= 0) {
return db.multiGetAsList(keys); final List<byte[]> keys = keys(fromKeyIdx, fromKeyIdx + multiGetSize);
final List<byte[]> valueResults = db.multiGetAsList(keys);
for (final byte[] result : valueResults) {
if (result.length != valueSize)
throw new RuntimeException("Test valueSize assumption wrong");
}
}
return new ArrayList<>();
}
@Benchmark
public List<RocksDB.MultiGetInstance> multiGetDirect10() throws RocksDBException {
final int fromKeyIdx = next(multiGetSize, keyCount);
if (fromKeyIdx >= 0) {
final List<ByteBuffer> keys = keys(keyBuffersList, fromKeyIdx, fromKeyIdx + multiGetSize);
final List<RocksDB.MultiGetInstance> results = db.multiGetByteBuffers(
keys, valueBuffersList.subList(fromKeyIdx, fromKeyIdx + multiGetSize));
for (final RocksDB.MultiGetInstance result : results) {
if (result.status.getCode() != Status.Code.Ok)
throw new RuntimeException("Test status assumption wrong");
if (result.valueSize != valueSize)
throw new RuntimeException("Test valueSize assumption wrong");
}
return results;
}
return new ArrayList<>();
}
public static void main(final String[] args) throws RunnerException {
final org.openjdk.jmh.runner.options.Options opt =
new OptionsBuilder()
.include(MultiGetBenchmarks.class.getSimpleName())
.forks(1)
.jvmArgs("-ea")
.warmupIterations(1)
.measurementIterations(2)
.forks(2)
.param("columnFamilyTestType=", "1_column_family")
.param("multiGetSize=", "10", "1000")
.param("keyCount=", "1000")
.output("jmh_output")
.build();
new Runner(opt).run();
} }
} }

18
java/jmh/src/main/java/org/rocksdb/util/KVUtils.java
View File

@ -6,11 +6,12 @@
*/ */
package org.rocksdb.util; package org.rocksdb.util;
import static java.nio.charset.StandardCharsets.UTF_8;
import java.nio.ByteBuffer;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.List; import java.util.List;
import static java.nio.charset.StandardCharsets.UTF_8;
public final class KVUtils { public final class KVUtils {
/** /**
@ -55,4 +56,17 @@ public final class KVUtils {
} }
return keys; return keys;
} }
public static List<ByteBuffer> keys(
final List<ByteBuffer> keyBuffers, final int from, final int to) {
final List<ByteBuffer> keys = new ArrayList<>(to - from);
for (int i = from; i < to; i++) {
final ByteBuffer key = keyBuffers.get(i);
key.clear();
key.put(ba("key" + i));
key.flip();
keys.add(key);
}
return keys;
}
} }

302
java/rocksjni/rocksjni.cc
View File

@ -1685,16 +1685,18 @@ inline void multi_get_helper_release_keys(std::vector<jbyte*>& keys_to_free) {
} }
/** /**
* cf multi get * @brief fill a native array of cf handles from java handles
* *
* @return byte[][] of values or nullptr if an exception occurs * @param env
* @param cf_handles to fill from the java variants
* @param jcolumn_family_handles
* @return true if the copy succeeds
* @return false if a JNI exception is generated
*/ */
jobjectArray multi_get_helper(JNIEnv* env, jobject, ROCKSDB_NAMESPACE::DB* db, inline bool cf_handles_from_jcf_handles(
const ROCKSDB_NAMESPACE::ReadOptions& rOpt, JNIEnv* env,
jobjectArray jkeys, jintArray jkey_offs, std::vector<ROCKSDB_NAMESPACE::ColumnFamilyHandle*>& cf_handles,
jintArray jkey_lens, jlongArray jcolumn_family_handles) {
jlongArray jcolumn_family_handles) {
std::vector<ROCKSDB_NAMESPACE::ColumnFamilyHandle*> cf_handles;
if (jcolumn_family_handles != nullptr) { if (jcolumn_family_handles != nullptr) {
const jsize len_cols = env->GetArrayLength(jcolumn_family_handles); const jsize len_cols = env->GetArrayLength(jcolumn_family_handles);
@ -1704,7 +1706,7 @@ jobjectArray multi_get_helper(JNIEnv* env, jobject, ROCKSDB_NAMESPACE::DB* db,
jclass exception_cls = (env)->FindClass("java/lang/OutOfMemoryError"); jclass exception_cls = (env)->FindClass("java/lang/OutOfMemoryError");
(env)->ThrowNew(exception_cls, (env)->ThrowNew(exception_cls,
"Insufficient Memory for CF handle array."); "Insufficient Memory for CF handle array.");
return nullptr; return false;
} }
for (jsize i = 0; i < len_cols; i++) { for (jsize i = 0; i < len_cols; i++) {
@ -1714,13 +1716,30 @@ jobjectArray multi_get_helper(JNIEnv* env, jobject, ROCKSDB_NAMESPACE::DB* db,
} }
env->ReleaseLongArrayElements(jcolumn_family_handles, jcfh, JNI_ABORT); env->ReleaseLongArrayElements(jcolumn_family_handles, jcfh, JNI_ABORT);
} }
return true;
}
/**
* @brief copy keys from JNI into vector of slices for Rocks API
*
* @param keys to instantiate
* @param jkeys
* @param jkey_offs
* @param jkey_lens
* @return true if the copy succeeds
* @return false if a JNI exception is raised
*/
inline bool keys_from_jkeys(JNIEnv* env,
std::vector<ROCKSDB_NAMESPACE::Slice>& keys,
std::vector<jbyte*>& keys_to_free,
jobjectArray jkeys, jintArray jkey_offs,
jintArray jkey_lens) {
jint* jkey_off = env->GetIntArrayElements(jkey_offs, nullptr); jint* jkey_off = env->GetIntArrayElements(jkey_offs, nullptr);
if (jkey_off == nullptr) { if (jkey_off == nullptr) {
// exception thrown: OutOfMemoryError // exception thrown: OutOfMemoryError
jclass exception_cls = (env)->FindClass("java/lang/OutOfMemoryError"); jclass exception_cls = (env)->FindClass("java/lang/OutOfMemoryError");
(env)->ThrowNew(exception_cls, "Insufficient Memory for key offset array."); (env)->ThrowNew(exception_cls, "Insufficient Memory for key offset array.");
return nullptr; return false;
} }
jint* jkey_len = env->GetIntArrayElements(jkey_lens, nullptr); jint* jkey_len = env->GetIntArrayElements(jkey_lens, nullptr);
@ -1729,12 +1748,10 @@ jobjectArray multi_get_helper(JNIEnv* env, jobject, ROCKSDB_NAMESPACE::DB* db,
env->ReleaseIntArrayElements(jkey_offs, jkey_off, JNI_ABORT); env->ReleaseIntArrayElements(jkey_offs, jkey_off, JNI_ABORT);
jclass exception_cls = (env)->FindClass("java/lang/OutOfMemoryError"); jclass exception_cls = (env)->FindClass("java/lang/OutOfMemoryError");
(env)->ThrowNew(exception_cls, "Insufficient Memory for key length array."); (env)->ThrowNew(exception_cls, "Insufficient Memory for key length array.");
return nullptr; return false;
} }
const jsize len_keys = env->GetArrayLength(jkeys); const jsize len_keys = env->GetArrayLength(jkeys);
std::vector<ROCKSDB_NAMESPACE::Slice> keys;
std::vector<jbyte*> keys_to_free;
for (jsize i = 0; i < len_keys; i++) { for (jsize i = 0; i < len_keys; i++) {
jobject jkey = env->GetObjectArrayElement(jkeys, i); jobject jkey = env->GetObjectArrayElement(jkeys, i);
if (env->ExceptionCheck()) { if (env->ExceptionCheck()) {
@ -1745,7 +1762,7 @@ jobjectArray multi_get_helper(JNIEnv* env, jobject, ROCKSDB_NAMESPACE::DB* db,
jclass exception_cls = (env)->FindClass("java/lang/OutOfMemoryError"); jclass exception_cls = (env)->FindClass("java/lang/OutOfMemoryError");
(env)->ThrowNew(exception_cls, (env)->ThrowNew(exception_cls,
"Insufficient Memory for key object array."); "Insufficient Memory for key object array.");
return nullptr; return false;
} }
jbyteArray jkey_ba = reinterpret_cast<jbyteArray>(jkey); jbyteArray jkey_ba = reinterpret_cast<jbyteArray>(jkey);
@ -1763,7 +1780,7 @@ jobjectArray multi_get_helper(JNIEnv* env, jobject, ROCKSDB_NAMESPACE::DB* db,
jclass exception_cls = jclass exception_cls =
(env)->FindClass("java/lang/ArrayIndexOutOfBoundsException"); (env)->FindClass("java/lang/ArrayIndexOutOfBoundsException");
(env)->ThrowNew(exception_cls, "Invalid byte array region index."); (env)->ThrowNew(exception_cls, "Invalid byte array region index.");
return nullptr; return false;
} }
ROCKSDB_NAMESPACE::Slice key_slice(reinterpret_cast<char*>(key), len_key); ROCKSDB_NAMESPACE::Slice key_slice(reinterpret_cast<char*>(key), len_key);
@ -1777,6 +1794,68 @@ jobjectArray multi_get_helper(JNIEnv* env, jobject, ROCKSDB_NAMESPACE::DB* db,
env->ReleaseIntArrayElements(jkey_lens, jkey_len, JNI_ABORT); env->ReleaseIntArrayElements(jkey_lens, jkey_len, JNI_ABORT);
env->ReleaseIntArrayElements(jkey_offs, jkey_off, JNI_ABORT); env->ReleaseIntArrayElements(jkey_offs, jkey_off, JNI_ABORT);
return true;
}
inline bool keys_from_bytebuffers(JNIEnv* env,
std::vector<ROCKSDB_NAMESPACE::Slice>& keys,
jobjectArray jkeys, jintArray jkey_offs,
jintArray jkey_lens) {
jint* jkey_off = env->GetIntArrayElements(jkey_offs, nullptr);
if (jkey_off == nullptr) {
// exception thrown: OutOfMemoryError
jclass exception_cls = (env)->FindClass("java/lang/OutOfMemoryError");
(env)->ThrowNew(exception_cls, "Insufficient Memory for key offset array.");
return false;
}
jint* jkey_len = env->GetIntArrayElements(jkey_lens, nullptr);
if (jkey_len == nullptr) {
// exception thrown: OutOfMemoryError
env->ReleaseIntArrayElements(jkey_offs, jkey_off, JNI_ABORT);
jclass exception_cls = (env)->FindClass("java/lang/OutOfMemoryError");
(env)->ThrowNew(exception_cls, "Insufficient Memory for key length array.");
return false;
}
const jsize len_keys = env->GetArrayLength(jkeys);
for (jsize i = 0; i < len_keys; i++) {
jobject jkey = env->GetObjectArrayElement(jkeys, i);
if (env->ExceptionCheck()) {
// exception thrown: ArrayIndexOutOfBoundsException
return false;
}
char* key = reinterpret_cast<char*>(env->GetDirectBufferAddress(jkey));
ROCKSDB_NAMESPACE::Slice key_slice(key + jkey_off[i], jkey_len[i]);
keys.push_back(key_slice);
env->DeleteLocalRef(jkey);
}
return true;
}
/**
* cf multi get
*
* @return byte[][] of values or nullptr if an
* exception occurs
*/
jobjectArray multi_get_helper(JNIEnv* env, jobject, ROCKSDB_NAMESPACE::DB* db,
const ROCKSDB_NAMESPACE::ReadOptions& rOpt,
jobjectArray jkeys, jintArray jkey_offs,
jintArray jkey_lens,
jlongArray jcolumn_family_handles) {
std::vector<ROCKSDB_NAMESPACE::ColumnFamilyHandle*> cf_handles;
if (!cf_handles_from_jcf_handles(env, cf_handles, jcolumn_family_handles)) {
return nullptr;
}
std::vector<ROCKSDB_NAMESPACE::Slice> keys;
std::vector<jbyte*> keys_to_free;
if (!keys_from_jkeys(env, keys, keys_to_free, jkeys, jkey_offs, jkey_lens)) {
return nullptr;
}
std::vector<std::string> values; std::vector<std::string> values;
std::vector<ROCKSDB_NAMESPACE::Status> s; std::vector<ROCKSDB_NAMESPACE::Status> s;
if (cf_handles.size() == 0) { if (cf_handles.size() == 0) {
@ -1814,14 +1893,16 @@ jobjectArray multi_get_helper(JNIEnv* env, jobject, ROCKSDB_NAMESPACE::DB* db,
jentry_value, 0, static_cast<jsize>(jvalue_len), jentry_value, 0, static_cast<jsize>(jvalue_len),
const_cast<jbyte*>(reinterpret_cast<const jbyte*>(value->c_str()))); const_cast<jbyte*>(reinterpret_cast<const jbyte*>(value->c_str())));
if (env->ExceptionCheck()) { if (env->ExceptionCheck()) {
// exception thrown: ArrayIndexOutOfBoundsException // exception thrown:
// ArrayIndexOutOfBoundsException
env->DeleteLocalRef(jentry_value); env->DeleteLocalRef(jentry_value);
return nullptr; return nullptr;
} }
env->SetObjectArrayElement(jresults, static_cast<jsize>(i), jentry_value); env->SetObjectArrayElement(jresults, static_cast<jsize>(i), jentry_value);
if (env->ExceptionCheck()) { if (env->ExceptionCheck()) {
// exception thrown: ArrayIndexOutOfBoundsException // exception thrown:
// ArrayIndexOutOfBoundsException
env->DeleteLocalRef(jentry_value); env->DeleteLocalRef(jentry_value);
return nullptr; return nullptr;
} }
@ -1833,14 +1914,129 @@ jobjectArray multi_get_helper(JNIEnv* env, jobject, ROCKSDB_NAMESPACE::DB* db,
return jresults; return jresults;
} }
/**
* cf multi get
*
* fill supplied native buffers, or raise JNI
* exception on a problem
*/
/**
* @brief multi_get_helper_direct for fast-path multiget (io_uring) on Linux
*
* @param env
* @param db
* @param rOpt read options
* @param jcolumn_family_handles 0, 1, or n column family handles
* @param jkeys
* @param jkey_offsets
* @param jkey_lengths
* @param jvalues byte buffers to receive values
* @param jvalue_sizes returned actual sizes of data values for keys
* @param jstatuses returned java RocksDB status values for per key
*/
void multi_get_helper_direct(JNIEnv* env, jobject, ROCKSDB_NAMESPACE::DB* db,
const ROCKSDB_NAMESPACE::ReadOptions& rOpt,
jlongArray jcolumn_family_handles,
jobjectArray jkeys, jintArray jkey_offsets,
jintArray jkey_lengths, jobjectArray jvalues,
jintArray jvalue_sizes, jobjectArray jstatuses) {
const jsize num_keys = env->GetArrayLength(jkeys);
std::vector<ROCKSDB_NAMESPACE::Slice> keys;
if (!keys_from_bytebuffers(env, keys, jkeys, jkey_offsets, jkey_lengths)) {
return;
}
std::vector<ROCKSDB_NAMESPACE::PinnableSlice> values(num_keys);
std::vector<ROCKSDB_NAMESPACE::ColumnFamilyHandle*> cf_handles;
if (!cf_handles_from_jcf_handles(env, cf_handles, jcolumn_family_handles)) {
return;
}
std::vector<ROCKSDB_NAMESPACE::Status> s(num_keys);
if (cf_handles.size() == 0) {
// we can use the more efficient call here
auto cf_handle = db->DefaultColumnFamily();
db->MultiGet(rOpt, cf_handle, num_keys, keys.data(), values.data(),
s.data());
} else if (cf_handles.size() == 1) {
// we can use the more efficient call here
auto cf_handle = cf_handles[0];
db->MultiGet(rOpt, cf_handle, num_keys, keys.data(), values.data(),
s.data());
} else {
// multiple CFs version
db->MultiGet(rOpt, num_keys, cf_handles.data(), keys.data(), values.data(),
s.data());
}
// prepare the results
jobjectArray jresults = ROCKSDB_NAMESPACE::ByteJni::new2dByteArray(
env, static_cast<jsize>(s.size()));
if (jresults == nullptr) {
// exception occurred
jclass exception_cls = (env)->FindClass("java/lang/OutOfMemoryError");
(env)->ThrowNew(exception_cls, "Insufficient Memory for results.");
return;
}
std::vector<jint> value_size;
for (int i = 0; i < num_keys; i++) {
auto jstatus = ROCKSDB_NAMESPACE::StatusJni::construct(env, s[i]);
if (jstatus == nullptr) {
// exception in context
return;
}
env->SetObjectArrayElement(jstatuses, i, jstatus);
if (s[i].ok()) {
jobject jvalue_bytebuf = env->GetObjectArrayElement(jvalues, i);
if (env->ExceptionCheck()) {
// ArrayIndexOutOfBoundsException is thrown
return;
}
jlong jvalue_capacity = env->GetDirectBufferCapacity(jvalue_bytebuf);
if (jvalue_capacity == -1) {
ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(
env,
"Invalid value(s) argument (argument is not a valid direct "
"ByteBuffer)");
return;
}
void* jvalue_address = env->GetDirectBufferAddress(jvalue_bytebuf);
if (jvalue_address == nullptr) {
ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(
env,
"Invalid value(s) argument (argument is not a valid direct "
"ByteBuffer)");
return;
}
// record num returned, push back that number, which may be bigger then
// the ByteBuffer supplied. then copy as much as fits in the ByteBuffer.
value_size.push_back(static_cast<jint>(values[i].size()));
auto copy_bytes =
std::min(static_cast<jlong>(values[i].size()), jvalue_capacity);
memcpy(jvalue_address, values[i].data(), copy_bytes);
} else {
// bad status for this
value_size.push_back(0);
}
}
env->SetIntArrayRegion(jvalue_sizes, 0, num_keys, value_size.data());
}
/* /*
* Class: org_rocksdb_RocksDB * Class: org_rocksdb_RocksDB
* Method: multiGet * Method: multiGet
* Signature: (J[[B[I[I)[[B * Signature: (J[[B[I[I)[[B
*/ */
jobjectArray Java_org_rocksdb_RocksDB_multiGet__J_3_3B_3I_3I( jobjectArray Java_org_rocksdb_RocksDB_multiGet__J_3_3B_3I_3I(
JNIEnv* env, jobject jdb, jlong jdb_handle, JNIEnv* env, jobject jdb, jlong jdb_handle, jobjectArray jkeys,
jobjectArray jkeys, jintArray jkey_offs, jintArray jkey_lens) { jintArray jkey_offs, jintArray jkey_lens) {
return multi_get_helper( return multi_get_helper(
env, jdb, reinterpret_cast<ROCKSDB_NAMESPACE::DB*>(jdb_handle), env, jdb, reinterpret_cast<ROCKSDB_NAMESPACE::DB*>(jdb_handle),
ROCKSDB_NAMESPACE::ReadOptions(), jkeys, jkey_offs, jkey_lens, nullptr); ROCKSDB_NAMESPACE::ReadOptions(), jkeys, jkey_offs, jkey_lens, nullptr);
@ -1852,8 +2048,8 @@ jobjectArray Java_org_rocksdb_RocksDB_multiGet__J_3_3B_3I_3I(
* Signature: (J[[B[I[I[J)[[B * Signature: (J[[B[I[I[J)[[B
*/ */
jobjectArray Java_org_rocksdb_RocksDB_multiGet__J_3_3B_3I_3I_3J( jobjectArray Java_org_rocksdb_RocksDB_multiGet__J_3_3B_3I_3I_3J(
JNIEnv* env, jobject jdb, jlong jdb_handle, JNIEnv* env, jobject jdb, jlong jdb_handle, jobjectArray jkeys,
jobjectArray jkeys, jintArray jkey_offs, jintArray jkey_lens, jintArray jkey_offs, jintArray jkey_lens,
jlongArray jcolumn_family_handles) { jlongArray jcolumn_family_handles) {
return multi_get_helper(env, jdb, return multi_get_helper(env, jdb,
reinterpret_cast<ROCKSDB_NAMESPACE::DB*>(jdb_handle), reinterpret_cast<ROCKSDB_NAMESPACE::DB*>(jdb_handle),
@ -1890,38 +2086,60 @@ jobjectArray Java_org_rocksdb_RocksDB_multiGet__JJ_3_3B_3I_3I_3J(
jkey_offs, jkey_lens, jcolumn_family_handles); jkey_offs, jkey_lens, jcolumn_family_handles);
} }
/*
* Class: org_rocksdb_RocksDB
* Method: multiGet
* Signature:
* (JJ[J[Ljava/nio/ByteBuffer;[I[I[Ljava/nio/ByteBuffer;[I[Lorg/rocksdb/Status;)V
*/
void Java_org_rocksdb_RocksDB_multiGet__JJ_3J_3Ljava_nio_ByteBuffer_2_3I_3I_3Ljava_nio_ByteBuffer_2_3I_3Lorg_rocksdb_Status_2(
JNIEnv* env, jobject jdb, jlong jdb_handle, jlong jropt_handle,
jlongArray jcolumn_family_handles, jobjectArray jkeys,
jintArray jkey_offsets, jintArray jkey_lengths, jobjectArray jvalues,
jintArray jvalues_sizes, jobjectArray jstatus_objects) {
return multi_get_helper_direct(
env, jdb, reinterpret_cast<ROCKSDB_NAMESPACE::DB*>(jdb_handle),
*reinterpret_cast<ROCKSDB_NAMESPACE::ReadOptions*>(jropt_handle),
jcolumn_family_handles, jkeys, jkey_offsets, jkey_lengths, jvalues,
jvalues_sizes, jstatus_objects);
}
// private native void
// multiGet(final long dbHandle, final long rOptHandle,
// final long[] columnFamilyHandles, final ByteBuffer[] keysArray,
// final ByteBuffer[] valuesArray);
////////////////////////////////////////////////////////////////////////////// //////////////////////////////////////////////////////////////////////////////
// ROCKSDB_NAMESPACE::DB::KeyMayExist // ROCKSDB_NAMESPACE::DB::KeyMayExist
bool key_may_exist_helper(JNIEnv* env, jlong jdb_handle, jlong jcf_handle, bool key_may_exist_helper(JNIEnv* env, jlong jdb_handle, jlong jcf_handle,
jlong jread_opts_handle, jlong jread_opts_handle, jbyteArray jkey,
jbyteArray jkey, jint jkey_offset, jint jkey_len, jint jkey_offset, jint jkey_len, bool* has_exception,
bool* has_exception, std::string* value, bool* value_found) { std::string* value, bool* value_found) {
auto* db = reinterpret_cast<ROCKSDB_NAMESPACE::DB*>(jdb_handle); auto* db = reinterpret_cast<ROCKSDB_NAMESPACE::DB*>(jdb_handle);
ROCKSDB_NAMESPACE::ColumnFamilyHandle* cf_handle; ROCKSDB_NAMESPACE::ColumnFamilyHandle* cf_handle;
if (jcf_handle == 0) { if (jcf_handle == 0) {
cf_handle = db->DefaultColumnFamily(); cf_handle = db->DefaultColumnFamily();
} else { } else {
cf_handle = cf_handle =
reinterpret_cast<ROCKSDB_NAMESPACE::ColumnFamilyHandle*>(jcf_handle); reinterpret_cast<ROCKSDB_NAMESPACE::ColumnFamilyHandle*>(jcf_handle);
} }
ROCKSDB_NAMESPACE::ReadOptions read_opts = ROCKSDB_NAMESPACE::ReadOptions read_opts =
jread_opts_handle == 0 jread_opts_handle == 0
? ROCKSDB_NAMESPACE::ReadOptions() ? ROCKSDB_NAMESPACE::ReadOptions()
: *(reinterpret_cast<ROCKSDB_NAMESPACE::ReadOptions*>( : *(reinterpret_cast<ROCKSDB_NAMESPACE::ReadOptions*>(
jread_opts_handle)); jread_opts_handle));
jbyte* key = new jbyte[jkey_len]; jbyte* key = new jbyte[jkey_len];
env->GetByteArrayRegion(jkey, jkey_offset, jkey_len, key); env->GetByteArrayRegion(jkey, jkey_offset, jkey_len, key);
if (env->ExceptionCheck()) { if (env->ExceptionCheck()) {
// exception thrown: ArrayIndexOutOfBoundsException // exception thrown: ArrayIndexOutOfBoundsException
delete[] key; delete[] key;
*has_exception = true; *has_exception = true;
return false; return false;
} }
ROCKSDB_NAMESPACE::Slice key_slice(reinterpret_cast<char*>(key), jkey_len); ROCKSDB_NAMESPACE::Slice key_slice(reinterpret_cast<char*>(key), jkey_len);
const bool exists = db->KeyMayExist( const bool exists =
read_opts, cf_handle, key_slice, value, value_found); db->KeyMayExist(read_opts, cf_handle, key_slice, value, value_found);
// cleanup // cleanup
delete[] key; delete[] key;

44
java/src/main/java/org/rocksdb/ByteBufferGetStatus.java Normal file
View File

@ -0,0 +1,44 @@
package org.rocksdb;
import java.nio.ByteBuffer;
import java.util.List;
/**
* A ByteBuffer containing fetched data, together with a result for the fetch
* and the total size of the object fetched.
*
* Used for the individual results of
* {@link RocksDB#multiGetByteBuffers(List, List)}
* {@link RocksDB#multiGetByteBuffers(List, List, List)}
* {@link RocksDB#multiGetByteBuffers(ReadOptions, List, List)}
* {@link RocksDB#multiGetByteBuffers(ReadOptions, List, List, List)}
*/
public class ByteBufferGetStatus {
public final Status status;
public final int requiredSize;
public final ByteBuffer value;
/**
* Constructor used for success status, when the value is contained in the buffer
*
* @param status the status of the request to fetch into the buffer
* @param requiredSize the size of the data, which may be bigger than the buffer
* @param value the buffer containing as much of the value as fits
*/
ByteBufferGetStatus(final Status status, final int requiredSize, final ByteBuffer value) {
this.status = status;
this.requiredSize = requiredSize;
this.value = value;
}
/**
* Constructor used for a failure status, when no value is filled in
*
* @param status the status of the request to fetch into the buffer
*/
ByteBufferGetStatus(final Status status) {
this.status = status;
this.requiredSize = 0;
this.value = null;
}
}

166
java/src/main/java/org/rocksdb/RocksDB.java
View File

@ -31,14 +31,14 @@ public class RocksDB extends RocksObject {
LOADED LOADED
} }
private static AtomicReference<LibraryState> libraryLoaded private static final AtomicReference<LibraryState> libraryLoaded =
= new AtomicReference<>(LibraryState.NOT_LOADED); new AtomicReference<>(LibraryState.NOT_LOADED);
static { static {
RocksDB.loadLibrary(); RocksDB.loadLibrary();
} }
private List<ColumnFamilyHandle> ownedColumnFamilyHandles = new ArrayList<>(); private final List<ColumnFamilyHandle> ownedColumnFamilyHandles = new ArrayList<>();
/** /**
* Loads the necessary library files. * Loads the necessary library files.
@ -2544,6 +2544,154 @@ public class RocksDB extends RocksObject {
keysArray, keyOffsets, keyLengths, cfHandles)); keysArray, keyOffsets, keyLengths, cfHandles));
} }
/**
* Fetches a list of values for the given list of keys, all from the default column family.
*
* @param keys list of keys for which values need to be retrieved.
* @param values list of buffers to return retrieved values in
* @return list of number of bytes in DB for each requested key
* this can be more than the size of the corresponding buffer; then the buffer will be filled
* with the appropriate truncation of the database value.
* @throws RocksDBException if error happens in underlying native library.
* @throws IllegalArgumentException thrown if the number of passed keys and passed values
* do not match.
*/
public List<ByteBufferGetStatus> multiGetByteBuffers(
final List<ByteBuffer> keys, final List<ByteBuffer> values) throws RocksDBException {
final ReadOptions readOptions = new ReadOptions();
final List<ColumnFamilyHandle> columnFamilyHandleList = new ArrayList<>(1);
columnFamilyHandleList.add(getDefaultColumnFamily());
return multiGetByteBuffers(readOptions, columnFamilyHandleList, keys, values);
}
/**
* Fetches a list of values for the given list of keys, all from the default column family.
*
* @param readOptions Read options
* @param keys list of keys for which values need to be retrieved.
* @param values list of buffers to return retrieved values in
* @throws RocksDBException if error happens in underlying native library.
* @throws IllegalArgumentException thrown if the number of passed keys and passed values
* do not match.
*/
public List<ByteBufferGetStatus> multiGetByteBuffers(final ReadOptions readOptions,
final List<ByteBuffer> keys, final List<ByteBuffer> values) throws RocksDBException {
final List<ColumnFamilyHandle> columnFamilyHandleList = new ArrayList<>(1);
columnFamilyHandleList.add(getDefaultColumnFamily());
return multiGetByteBuffers(readOptions, columnFamilyHandleList, keys, values);
}
/**
* Fetches a list of values for the given list of keys.
* <p>
* Note: Every key needs to have a related column family name in
* {@code columnFamilyHandleList}.
* </p>
*
* @param columnFamilyHandleList {@link java.util.List} containing
* {@link org.rocksdb.ColumnFamilyHandle} instances.
* @param keys list of keys for which values need to be retrieved.
* @param values list of buffers to return retrieved values in
* @throws RocksDBException if error happens in underlying native library.
* @throws IllegalArgumentException thrown if the number of passed keys, passed values and
* passed column family handles do not match.
*/
public List<ByteBufferGetStatus> multiGetByteBuffers(
final List<ColumnFamilyHandle> columnFamilyHandleList, final List<ByteBuffer> keys,
final List<ByteBuffer> values) throws RocksDBException {
final ReadOptions readOptions = new ReadOptions();
return multiGetByteBuffers(readOptions, columnFamilyHandleList, keys, values);
}
/**
* Fetches a list of values for the given list of keys.
* <p>
* Note: Every key needs to have a related column family name in
* {@code columnFamilyHandleList}.
* </p>
*
* @param readOptions Read options
* @param columnFamilyHandleList {@link java.util.List} containing
* {@link org.rocksdb.ColumnFamilyHandle} instances.
* @param keys list of keys for which values need to be retrieved.
* @param values list of buffers to return retrieved values in
* @throws RocksDBException if error happens in underlying native library.
* @throws IllegalArgumentException thrown if the number of passed keys, passed values and
* passed column family handles do not match.
*/
public List<ByteBufferGetStatus> multiGetByteBuffers(final ReadOptions readOptions,
final List<ColumnFamilyHandle> columnFamilyHandleList, final List<ByteBuffer> keys,
final List<ByteBuffer> values) throws RocksDBException {
assert (keys.size() != 0);
// Check if key size equals cfList size. If not a exception must be
// thrown. If not a Segmentation fault happens.
if (keys.size() != columnFamilyHandleList.size() && columnFamilyHandleList.size() > 1) {
throw new IllegalArgumentException(
"Wrong number of ColumnFamilyHandle(s) supplied. Provide 0, 1, or as many as there are key/value(s)");
}
// Check if key size equals cfList size. If not a exception must be
// thrown. If not a Segmentation fault happens.
if (values.size() != keys.size()) {
throw new IllegalArgumentException("For each key there must be a corresponding value.");
}
// TODO (AP) support indirect buffers
for (final ByteBuffer key : keys) {
if (!key.isDirect()) {
throw new IllegalArgumentException("All key buffers must be direct byte buffers");
}
}
// TODO (AP) support indirect buffers, though probably via a less efficient code path
for (final ByteBuffer value : values) {
if (!value.isDirect()) {
throw new IllegalArgumentException("All value buffers must be direct byte buffers");
}
}
final int numCFHandles = columnFamilyHandleList.size();
final long[] cfHandles = new long[numCFHandles];
for (int i = 0; i < numCFHandles; i++) {
cfHandles[i] = columnFamilyHandleList.get(i).nativeHandle_;
}
final int numValues = keys.size();
final ByteBuffer[] keysArray = keys.toArray(new ByteBuffer[0]);
final int[] keyOffsets = new int[numValues];
final int[] keyLengths = new int[numValues];
for (int i = 0; i < numValues; i++) {
// TODO (AP) add keysArray[i].arrayOffset() if the buffer is indirect
// TODO (AP) because in that case we have to pass the array directly,
// so that the JNI C++ code will not know to compensate for the array offset
keyOffsets[i] = keysArray[i].position();
keyLengths[i] = keysArray[i].limit();
}
final ByteBuffer[] valuesArray = values.toArray(new ByteBuffer[0]);
final int[] valuesSizeArray = new int[numValues];
final Status[] statusArray = new Status[numValues];
multiGet(nativeHandle_, readOptions.nativeHandle_, cfHandles, keysArray, keyOffsets, keyLengths,
valuesArray, valuesSizeArray, statusArray);
final List<ByteBufferGetStatus> results = new ArrayList<>();
for (int i = 0; i < numValues; i++) {
final Status status = statusArray[i];
if (status.getCode() == Status.Code.Ok) {
final ByteBuffer value = valuesArray[i];
value.position(Math.min(valuesSizeArray[i], value.capacity()));
value.flip(); // prepare for read out
results.add(new ByteBufferGetStatus(status, valuesSizeArray[i], value));
} else {
results.add(new ByteBufferGetStatus(status));
}
}
return results;
}
/** /**
* If the key definitely does not exist in the database, then this method * If the key definitely does not exist in the database, then this method
* returns false, otherwise it returns true if the key might exist. * returns false, otherwise it returns true if the key might exist.
@ -4841,6 +4989,12 @@ public class RocksDB extends RocksObject {
private native byte[][] multiGet(final long dbHandle, final long rOptHandle, private native byte[][] multiGet(final long dbHandle, final long rOptHandle,
final byte[][] keys, final int[] keyOffsets, final int[] keyLengths, final byte[][] keys, final int[] keyOffsets, final int[] keyLengths,
final long[] columnFamilyHandles); final long[] columnFamilyHandles);
private native void multiGet(final long dbHandle, final long rOptHandle,
final long[] columnFamilyHandles, final ByteBuffer[] keysArray, final int[] keyOffsets,
final int[] keyLengths, final ByteBuffer[] valuesArray, final int[] valuesSizeArray,
final Status[] statusArray);
private native boolean keyMayExist( private native boolean keyMayExist(
final long handle, final long cfHandle, final long readOptHandle, final long handle, final long cfHandle, final long readOptHandle,
final byte[] key, final int keyOffset, final int keyLength); final byte[] key, final int keyOffset, final int keyLength);
@ -4887,9 +5041,9 @@ public class RocksDB extends RocksObject {
private native long[] getApproximateSizes(final long nativeHandle, private native long[] getApproximateSizes(final long nativeHandle,
final long columnFamilyHandle, final long[] rangeSliceHandles, final long columnFamilyHandle, final long[] rangeSliceHandles,
final byte includeFlags); final byte includeFlags);
private final native long[] getApproximateMemTableStats( private native long[] getApproximateMemTableStats(final long nativeHandle,
final long nativeHandle, final long columnFamilyHandle, final long columnFamilyHandle, final long rangeStartSliceHandle,
final long rangeStartSliceHandle, final long rangeLimitSliceHandle); final long rangeLimitSliceHandle);
private native void compactRange(final long handle, private native void compactRange(final long handle,
/* @Nullable */ final byte[] begin, final int beginLen, /* @Nullable */ final byte[] begin, final int beginLen,
/* @Nullable */ final byte[] end, final int endLen, /* @Nullable */ final byte[] end, final int endLen,

491
java/src/test/java/org/rocksdb/MultiGetTest.java
View File

@ -5,12 +5,16 @@
package org.rocksdb; package org.rocksdb;
import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.fail;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays; import java.util.Arrays;
import java.util.List; import java.util.List;
import org.junit.Rule; import org.junit.Rule;
import org.junit.Test; import org.junit.Test;
import org.junit.rules.TemporaryFolder; import org.junit.rules.TemporaryFolder;
import org.rocksdb.util.TestUtil;
public class MultiGetTest { public class MultiGetTest {
@Rule public TemporaryFolder dbFolder = new TemporaryFolder(); @Rule public TemporaryFolder dbFolder = new TemporaryFolder();
@ -31,4 +35,491 @@ public class MultiGetTest {
assertThat(values.get(2)).isEqualTo("value3ForKey3".getBytes()); assertThat(values.get(2)).isEqualTo("value3ForKey3".getBytes());
} }
} }
@Test
public void putNThenMultiGetDirect() throws RocksDBException {
try (final Options opt = new Options().setCreateIfMissing(true);
final RocksDB db = RocksDB.open(opt, dbFolder.getRoot().getAbsolutePath())) {
db.put("key1".getBytes(), "value1ForKey1".getBytes());
db.put("key2".getBytes(), "value2ForKey2".getBytes());
db.put("key3".getBytes(), "value3ForKey3".getBytes());
final List<ByteBuffer> keys = new ArrayList<>();
keys.add(ByteBuffer.allocateDirect(12).put("key1".getBytes()));
keys.add(ByteBuffer.allocateDirect(12).put("key2".getBytes()));
keys.add(ByteBuffer.allocateDirect(12).put("key3".getBytes()));
// Java8 and lower flip() returns Buffer not ByteBuffer, so can't chain above /\/\
for (final ByteBuffer key : keys) {
key.flip();
}
final List<ByteBuffer> values = new ArrayList<>();
for (int i = 0; i < keys.size(); i++) {
values.add(ByteBuffer.allocateDirect(24));
}
{
final List<ByteBufferGetStatus> results = db.multiGetByteBuffers(keys, values);
assertThat(results.get(0).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(1).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(2).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(0).requiredSize).isEqualTo("value1ForKey1".getBytes().length);
assertThat(results.get(1).requiredSize).isEqualTo("value2ForKey2".getBytes().length);
assertThat(results.get(2).requiredSize).isEqualTo("value3ForKey3".getBytes().length);
assertThat(TestUtil.bufferBytes(results.get(0).value))
.isEqualTo("value1ForKey1".getBytes());
assertThat(TestUtil.bufferBytes(results.get(1).value))
.isEqualTo("value2ForKey2".getBytes());
assertThat(TestUtil.bufferBytes(results.get(2).value))
.isEqualTo("value3ForKey3".getBytes());
}
{
final List<ByteBufferGetStatus> results =
db.multiGetByteBuffers(new ReadOptions(), keys, values);
assertThat(results.get(0).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(1).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(2).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(0).requiredSize).isEqualTo("value1ForKey1".getBytes().length);
assertThat(results.get(1).requiredSize).isEqualTo("value2ForKey2".getBytes().length);
assertThat(results.get(2).requiredSize).isEqualTo("value3ForKey3".getBytes().length);
assertThat(TestUtil.bufferBytes(results.get(0).value))
.isEqualTo("value1ForKey1".getBytes());
assertThat(TestUtil.bufferBytes(results.get(1).value))
.isEqualTo("value2ForKey2".getBytes());
assertThat(TestUtil.bufferBytes(results.get(2).value))
.isEqualTo("value3ForKey3".getBytes());
}
}
}
@Test
public void putNThenMultiGetDirectSliced() throws RocksDBException {
try (final Options opt = new Options().setCreateIfMissing(true);
final RocksDB db = RocksDB.open(opt, dbFolder.getRoot().getAbsolutePath())) {
db.put("key1".getBytes(), "value1ForKey1".getBytes());
db.put("key2".getBytes(), "value2ForKey2".getBytes());
db.put("key3".getBytes(), "value3ForKey3".getBytes());
final List<ByteBuffer> keys = new ArrayList<>();
keys.add(ByteBuffer.allocateDirect(12).put("key2".getBytes()));
keys.add(ByteBuffer.allocateDirect(12).put("key3".getBytes()));
keys.add(
ByteBuffer.allocateDirect(12).put("prefix1".getBytes()).slice().put("key1".getBytes()));
// Java8 and lower flip() returns Buffer not ByteBuffer, so can't chain above /\/\
for (final ByteBuffer key : keys) {
key.flip();
}
final List<ByteBuffer> values = new ArrayList<>();
for (int i = 0; i < keys.size(); i++) {
values.add(ByteBuffer.allocateDirect(24));
}
{
final List<ByteBufferGetStatus> results = db.multiGetByteBuffers(keys, values);
assertThat(results.get(0).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(1).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(2).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(1).requiredSize).isEqualTo("value2ForKey2".getBytes().length);
assertThat(results.get(2).requiredSize).isEqualTo("value3ForKey3".getBytes().length);
assertThat(results.get(0).requiredSize).isEqualTo("value1ForKey1".getBytes().length);
assertThat(TestUtil.bufferBytes(results.get(0).value))
.isEqualTo("value2ForKey2".getBytes());
assertThat(TestUtil.bufferBytes(results.get(1).value))
.isEqualTo("value3ForKey3".getBytes());
assertThat(TestUtil.bufferBytes(results.get(2).value))
.isEqualTo("value1ForKey1".getBytes());
}
}
}
@Test
public void putNThenMultiGetDirectBadValuesArray() throws RocksDBException {
try (final Options opt = new Options().setCreateIfMissing(true);
final RocksDB db = RocksDB.open(opt, dbFolder.getRoot().getAbsolutePath())) {
db.put("key1".getBytes(), "value1ForKey1".getBytes());
db.put("key2".getBytes(), "value2ForKey2".getBytes());
db.put("key3".getBytes(), "value3ForKey3".getBytes());
final List<ByteBuffer> keys = new ArrayList<>();
keys.add(ByteBuffer.allocateDirect(12).put("key1".getBytes()));
keys.add(ByteBuffer.allocateDirect(12).put("key2".getBytes()));
keys.add(ByteBuffer.allocateDirect(12).put("key3".getBytes()));
// Java8 and lower flip() returns Buffer not ByteBuffer, so can't chain above /\/\
for (final ByteBuffer key : keys) {
key.flip();
}
{
final List<ByteBuffer> values = new ArrayList<>();
for (int i = 0; i < keys.size(); i++) {
values.add(ByteBuffer.allocateDirect(24));
}
values.remove(0);
try {
db.multiGetByteBuffers(keys, values);
fail("Expected exception when not enough value ByteBuffers supplied");
} catch (final IllegalArgumentException e) {
assertThat(e.getMessage()).contains("For each key there must be a corresponding value");
}
}
{
final List<ByteBuffer> values = new ArrayList<>();
for (int i = 0; i < keys.size(); i++) {
values.add(ByteBuffer.allocateDirect(24));
}
values.add(ByteBuffer.allocateDirect(24));
try {
db.multiGetByteBuffers(keys, values);
fail("Expected exception when too many value ByteBuffers supplied");
} catch (final IllegalArgumentException e) {
assertThat(e.getMessage()).contains("For each key there must be a corresponding value");
}
}
}
}
@Test
public void putNThenMultiGetDirectShortValueBuffers() throws RocksDBException {
try (final Options opt = new Options().setCreateIfMissing(true);
final RocksDB db = RocksDB.open(opt, dbFolder.getRoot().getAbsolutePath())) {
db.put("key1".getBytes(), "value1ForKey1".getBytes());
db.put("key2".getBytes(), "value2ForKey2".getBytes());
db.put("key3".getBytes(), "value3ForKey3".getBytes());
final List<ByteBuffer> keys = new ArrayList<>();
keys.add(ByteBuffer.allocateDirect(12).put("key1".getBytes()));
keys.add(ByteBuffer.allocateDirect(12).put("key2".getBytes()));
keys.add(ByteBuffer.allocateDirect(12).put("key3".getBytes()));
// Java8 and lower flip() returns Buffer not ByteBuffer, so can't chain above /\/\
for (final ByteBuffer key : keys) {
key.flip();
}
{
final List<ByteBuffer> values = new ArrayList<>();
for (int i = 0; i < keys.size(); i++) {
values.add(ByteBuffer.allocateDirect(4));
}
final List<ByteBufferGetStatus> statii = db.multiGetByteBuffers(keys, values);
assertThat(statii.size()).isEqualTo(values.size());
for (final ByteBufferGetStatus status : statii) {
assertThat(status.status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(status.requiredSize).isEqualTo("value3ForKey3".getBytes().length);
final ByteBuffer expected =
ByteBuffer.allocateDirect(24).put(Arrays.copyOf("valueX".getBytes(), 4));
expected.flip();
assertThat(status.value).isEqualTo(expected);
}
}
}
}
@Test
public void putNThenMultiGetDirectNondefaultCF() throws RocksDBException {
try (final Options opt = new Options().setCreateIfMissing(true);
final RocksDB db = RocksDB.open(opt, dbFolder.getRoot().getAbsolutePath())) {
final List<ColumnFamilyDescriptor> cfDescriptors = new ArrayList<>(0);
cfDescriptors.add(new ColumnFamilyDescriptor("cf0".getBytes()));
cfDescriptors.add(new ColumnFamilyDescriptor("cf1".getBytes()));
cfDescriptors.add(new ColumnFamilyDescriptor("cf2".getBytes()));
final List<ColumnFamilyHandle> cf = db.createColumnFamilies(cfDescriptors);
db.put(cf.get(0), "key1".getBytes(), "value1ForKey1".getBytes());
db.put(cf.get(0), "key2".getBytes(), "value2ForKey2".getBytes());
db.put(cf.get(0), "key3".getBytes(), "value3ForKey3".getBytes());
final List<ByteBuffer> keys = new ArrayList<>();
keys.add(ByteBuffer.allocateDirect(12).put("key1".getBytes()));
keys.add(ByteBuffer.allocateDirect(12).put("key2".getBytes()));
keys.add(ByteBuffer.allocateDirect(12).put("key3".getBytes()));
// Java8 and lower flip() returns Buffer not ByteBuffer, so can't chain above /\/\
for (final ByteBuffer key : keys) {
key.flip();
}
final List<ByteBuffer> values = new ArrayList<>();
for (int i = 0; i < keys.size(); i++) {
values.add(ByteBuffer.allocateDirect(24));
}
{
final List<ByteBufferGetStatus> results = db.multiGetByteBuffers(keys, values);
assertThat(results.get(0).status.getCode()).isEqualTo(Status.Code.NotFound);
assertThat(results.get(1).status.getCode()).isEqualTo(Status.Code.NotFound);
assertThat(results.get(2).status.getCode()).isEqualTo(Status.Code.NotFound);
}
{
final List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>();
columnFamilyHandles.add(cf.get(0));
final List<ByteBufferGetStatus> results =
db.multiGetByteBuffers(columnFamilyHandles, keys, values);
assertThat(results.get(0).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(1).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(2).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(0).requiredSize).isEqualTo("value1ForKey1".getBytes().length);
assertThat(results.get(1).requiredSize).isEqualTo("value2ForKey2".getBytes().length);
assertThat(results.get(2).requiredSize).isEqualTo("value3ForKey3".getBytes().length);
assertThat(TestUtil.bufferBytes(results.get(0).value))
.isEqualTo("value1ForKey1".getBytes());
assertThat(TestUtil.bufferBytes(results.get(1).value))
.isEqualTo("value2ForKey2".getBytes());
assertThat(TestUtil.bufferBytes(results.get(2).value))
.isEqualTo("value3ForKey3".getBytes());
}
{
final List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>();
columnFamilyHandles.add(cf.get(0));
columnFamilyHandles.add(cf.get(0));
columnFamilyHandles.add(cf.get(0));
final List<ByteBufferGetStatus> results =
db.multiGetByteBuffers(columnFamilyHandles, keys, values);
assertThat(results.get(0).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(1).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(2).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(0).requiredSize).isEqualTo("value1ForKey1".getBytes().length);
assertThat(results.get(1).requiredSize).isEqualTo("value2ForKey2".getBytes().length);
assertThat(results.get(2).requiredSize).isEqualTo("value3ForKey3".getBytes().length);
assertThat(TestUtil.bufferBytes(results.get(0).value))
.isEqualTo("value1ForKey1".getBytes());
assertThat(TestUtil.bufferBytes(results.get(1).value))
.isEqualTo("value2ForKey2".getBytes());
assertThat(TestUtil.bufferBytes(results.get(2).value))
.isEqualTo("value3ForKey3".getBytes());
}
}
}
@Test
public void putNThenMultiGetDirectCFParams() throws RocksDBException {
try (final Options opt = new Options().setCreateIfMissing(true);
final RocksDB db = RocksDB.open(opt, dbFolder.getRoot().getAbsolutePath())) {
db.put("key1".getBytes(), "value1ForKey1".getBytes());
db.put("key2".getBytes(), "value2ForKey2".getBytes());
db.put("key3".getBytes(), "value3ForKey3".getBytes());
final List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>();
columnFamilyHandles.add(db.getDefaultColumnFamily());
columnFamilyHandles.add(db.getDefaultColumnFamily());
final List<ByteBuffer> keys = new ArrayList<>();
keys.add(ByteBuffer.allocateDirect(12).put("key1".getBytes()));
keys.add(ByteBuffer.allocateDirect(12).put("key2".getBytes()));
keys.add(ByteBuffer.allocateDirect(12).put("key3".getBytes()));
// Java8 and lower flip() returns Buffer not ByteBuffer, so can't chain above /\/\
for (final ByteBuffer key : keys) {
key.flip();
}
final List<ByteBuffer> values = new ArrayList<>();
for (int i = 0; i < keys.size(); i++) {
values.add(ByteBuffer.allocateDirect(24));
}
try {
db.multiGetByteBuffers(columnFamilyHandles, keys, values);
fail("Expected exception when 2 column families supplied");
} catch (final IllegalArgumentException e) {
assertThat(e.getMessage()).contains("Wrong number of ColumnFamilyHandle(s) supplied");
}
columnFamilyHandles.clear();
columnFamilyHandles.add(db.getDefaultColumnFamily());
final List<ByteBufferGetStatus> results =
db.multiGetByteBuffers(columnFamilyHandles, keys, values);
assertThat(results.get(0).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(1).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(2).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(0).requiredSize).isEqualTo("value1ForKey1".getBytes().length);
assertThat(results.get(1).requiredSize).isEqualTo("value2ForKey2".getBytes().length);
assertThat(results.get(2).requiredSize).isEqualTo("value3ForKey3".getBytes().length);
assertThat(TestUtil.bufferBytes(results.get(0).value)).isEqualTo("value1ForKey1".getBytes());
assertThat(TestUtil.bufferBytes(results.get(1).value)).isEqualTo("value2ForKey2".getBytes());
assertThat(TestUtil.bufferBytes(results.get(2).value)).isEqualTo("value3ForKey3".getBytes());
}
}
@Test
public void putNThenMultiGetDirectMixedCF() throws RocksDBException {
try (final Options opt = new Options().setCreateIfMissing(true);
final RocksDB db = RocksDB.open(opt, dbFolder.getRoot().getAbsolutePath())) {
final List<ColumnFamilyDescriptor> cfDescriptors = new ArrayList<>();
cfDescriptors.add(new ColumnFamilyDescriptor("cf0".getBytes()));
cfDescriptors.add(new ColumnFamilyDescriptor("cf1".getBytes()));
cfDescriptors.add(new ColumnFamilyDescriptor("cf2".getBytes()));
cfDescriptors.add(new ColumnFamilyDescriptor("cf3".getBytes()));
final List<ColumnFamilyHandle> cf = db.createColumnFamilies(cfDescriptors);
db.put(cf.get(1), "key1".getBytes(), "value1ForKey1".getBytes());
db.put("key2".getBytes(), "value2ForKey2".getBytes());
db.put(cf.get(3), "key3".getBytes(), "value3ForKey3".getBytes());
final List<ByteBuffer> keys = new ArrayList<>();
keys.add(ByteBuffer.allocateDirect(12).put("key1".getBytes()));
keys.add(ByteBuffer.allocateDirect(12).put("key2".getBytes()));
keys.add(ByteBuffer.allocateDirect(12).put("key3".getBytes()));
// Java8 and lower flip() returns Buffer not ByteBuffer, so can't chain above /\/\
for (final ByteBuffer key : keys) {
key.flip();
}
final List<ByteBuffer> values = new ArrayList<>();
for (int i = 0; i < keys.size(); i++) {
values.add(ByteBuffer.allocateDirect(24));
}
{
final List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>();
columnFamilyHandles.add(db.getDefaultColumnFamily());
final List<ByteBufferGetStatus> results =
db.multiGetByteBuffers(columnFamilyHandles, keys, values);
assertThat(results.get(0).status.getCode()).isEqualTo(Status.Code.NotFound);
assertThat(results.get(1).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(2).status.getCode()).isEqualTo(Status.Code.NotFound);
assertThat(results.get(1).requiredSize).isEqualTo("value2ForKey2".getBytes().length);
assertThat(TestUtil.bufferBytes(results.get(1).value))
.isEqualTo("value2ForKey2".getBytes());
}
{
final List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>();
columnFamilyHandles.add(cf.get(1));
final List<ByteBufferGetStatus> results =
db.multiGetByteBuffers(columnFamilyHandles, keys, values);
assertThat(results.get(0).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(1).status.getCode()).isEqualTo(Status.Code.NotFound);
assertThat(results.get(2).status.getCode()).isEqualTo(Status.Code.NotFound);
assertThat(results.get(0).requiredSize).isEqualTo("value2ForKey2".getBytes().length);
assertThat(TestUtil.bufferBytes(results.get(0).value))
.isEqualTo("value1ForKey1".getBytes());
}
{
final List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>();
columnFamilyHandles.add(cf.get(1));
columnFamilyHandles.add(db.getDefaultColumnFamily());
columnFamilyHandles.add(cf.get(3));
final List<ByteBufferGetStatus> results =
db.multiGetByteBuffers(columnFamilyHandles, keys, values);
assertThat(results.get(0).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(1).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(2).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(0).requiredSize).isEqualTo("value1ForKey1".getBytes().length);
assertThat(results.get(1).requiredSize).isEqualTo("value2ForKey2".getBytes().length);
assertThat(results.get(2).requiredSize).isEqualTo("value3ForKey3".getBytes().length);
assertThat(TestUtil.bufferBytes(results.get(0).value))
.isEqualTo("value1ForKey1".getBytes());
assertThat(TestUtil.bufferBytes(results.get(1).value))
.isEqualTo("value2ForKey2".getBytes());
assertThat(TestUtil.bufferBytes(results.get(2).value))
.isEqualTo("value3ForKey3".getBytes());
}
{
final List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>();
columnFamilyHandles.add(db.getDefaultColumnFamily());
columnFamilyHandles.add(cf.get(1));
columnFamilyHandles.add(cf.get(3));
final List<ByteBufferGetStatus> results =
db.multiGetByteBuffers(columnFamilyHandles, keys, values);
assertThat(results.get(0).status.getCode()).isEqualTo(Status.Code.NotFound);
assertThat(results.get(1).status.getCode()).isEqualTo(Status.Code.NotFound);
assertThat(results.get(2).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(2).requiredSize).isEqualTo("value3ForKey3".getBytes().length);
assertThat(TestUtil.bufferBytes(results.get(2).value))
.isEqualTo("value3ForKey3".getBytes());
}
}
}
@Test
public void putNThenMultiGetDirectTruncateCF() throws RocksDBException {
try (final Options opt = new Options().setCreateIfMissing(true);
final RocksDB db = RocksDB.open(opt, dbFolder.getRoot().getAbsolutePath())) {
final List<ColumnFamilyDescriptor> cfDescriptors = new ArrayList<>();
cfDescriptors.add(new ColumnFamilyDescriptor("cf0".getBytes()));
final List<ColumnFamilyHandle> cf = db.createColumnFamilies(cfDescriptors);
db.put(cf.get(0), "key1".getBytes(), "value1ForKey1".getBytes());
db.put(cf.get(0), "key2".getBytes(), "value2ForKey2WithLotsOfTrailingGarbage".getBytes());
db.put(cf.get(0), "key3".getBytes(), "value3ForKey3".getBytes());
final List<ByteBuffer> keys = new ArrayList<>();
keys.add(ByteBuffer.allocateDirect(12).put("key1".getBytes()));
keys.add(ByteBuffer.allocateDirect(12).put("key2".getBytes()));
keys.add(ByteBuffer.allocateDirect(12).put("key3".getBytes()));
// Java8 and lower flip() returns Buffer not ByteBuffer, so can't chain above /\/\
for (final ByteBuffer key : keys) {
key.flip();
}
final List<ByteBuffer> values = new ArrayList<>();
for (int i = 0; i < keys.size(); i++) {
values.add(ByteBuffer.allocateDirect(24));
}
{
final List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>();
columnFamilyHandles.add(cf.get(0));
final List<ByteBufferGetStatus> results =
db.multiGetByteBuffers(columnFamilyHandles, keys, values);
assertThat(results.get(0).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(1).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(2).status.getCode()).isEqualTo(Status.Code.Ok);
assertThat(results.get(0).requiredSize).isEqualTo("value1ForKey1".getBytes().length);
assertThat(results.get(1).requiredSize)
.isEqualTo("value2ForKey2WithLotsOfTrailingGarbage".getBytes().length);
assertThat(results.get(2).requiredSize).isEqualTo("value3ForKey3".getBytes().length);
assertThat(TestUtil.bufferBytes(results.get(0).value))
.isEqualTo("value1ForKey1".getBytes());
assertThat(TestUtil.bufferBytes(results.get(1).value))
.isEqualTo("valu e2Fo rKey 2Wit hLot sOfT".replace(" ", "").getBytes());
assertThat(TestUtil.bufferBytes(results.get(2).value))
.isEqualTo("value3ForKey3".getBytes());
}
}
}
} }

19
java/src/test/java/org/rocksdb/util/TestUtil.java
View File

@ -5,14 +5,14 @@
package org.rocksdb.util; package org.rocksdb.util;
import static java.nio.charset.StandardCharsets.UTF_8;
import java.nio.ByteBuffer;
import java.util.Random;
import org.rocksdb.CompactionPriority; import org.rocksdb.CompactionPriority;
import org.rocksdb.Options; import org.rocksdb.Options;
import org.rocksdb.WALRecoveryMode; import org.rocksdb.WALRecoveryMode;
import java.util.Random;
import static java.nio.charset.StandardCharsets.UTF_8;
/** /**
* General test utilities. * General test utilities.
*/ */
@ -58,4 +58,15 @@ public class TestUtil {
random.nextBytes(str); random.nextBytes(str);
return str; return str;
} }
/**
* Copy a {@link ByteBuffer} into an array for shorthand ease of test coding
* @param byteBuffer the buffer to copy
* @return a {@link byte[]} containing the same bytes as the input
*/
public static byte[] bufferBytes(final ByteBuffer byteBuffer) {
final byte[] result = new byte[byteBuffer.limit()];
byteBuffer.get(result);
return result;
}
} }