dolysis
/
rocksdb
mirror of https://github.com/facebook/rocksdb.git
3
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

RocksJava API - fix Transaction.multiGet() size limit, remove bogus EnsureLocalCapacity() calls (#10674) 

Summary:
Resolves see https://github.com/facebook/rocksdb/issues/9006

Fixes 2 related issues with JNI local references in the RocksJava API.

1. Some instances of RocksJava API JNI code appear to have misunderstood the reason for `JNIEnv->EnsureLocalCapacity()` and are carrying out bogus checks which happen to fail with some larger parameter values (many column families in a single call, very long key names or values). Remove these checks and add some regression tests for the previous failures.

2. The helper for Transaction multiGet operations (`multiGet()`, `multiGetForUpdate()`,...) is limited in the number of keys it can `get()` for because it requires a corresponding number of live local references. Refactor the helper slightly, copying out the key contents within a loop so that the references don't have to exist at the same time.

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

Reviewed By: ajkr

Differential Revision: D40515361

Pulled By: jay-zhuang

fbshipit-source-id: f1be0126181a698b3ad27c0945a39c54d950aa25
This commit is contained in:
Alan Paxton 2022-10-26 17:25:33 -07:00 committed by Facebook GitHub Bot
parent bf78380851
commit 17553bdd5e
7 changed files with 526 additions and 97 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
java/Makefile
View File

@ -146,6 +146,7 @@ JAVA_TESTS = \
org.rocksdb.MemoryUtilTest\
org.rocksdb.MemTableTest\
org.rocksdb.MergeTest\
org.rocksdb.MultiColumnRegressionTest \
org.rocksdb.MultiGetManyKeysTest\
org.rocksdb.MultiGetTest\
org.rocksdb.MixedOptionsTest\

14
java/rocksjni/optimistic_transaction_db.cc
View File

@ -63,12 +63,6 @@ Java_org_rocksdb_OptimisticTransactionDB_open__JLjava_lang_String_2_3_3B_3J(
std::vector<ROCKSDB_NAMESPACE::ColumnFamilyDescriptor> column_families;
const jsize len_cols = env->GetArrayLength(jcolumn_names);
if (len_cols > 0) {
if (env->EnsureLocalCapacity(len_cols) != 0) {
// out of memory
env->ReleaseStringUTFChars(jdb_path, db_path);
return nullptr;
}
jlong* jco = env->GetLongArrayElements(jcolumn_options_handles, nullptr);
if (jco == nullptr) {
// exception thrown: OutOfMemoryError
@ -87,14 +81,6 @@ Java_org_rocksdb_OptimisticTransactionDB_open__JLjava_lang_String_2_3_3B_3J(
const jbyteArray jcn_ba = reinterpret_cast<jbyteArray>(jcn);
const jsize jcf_name_len = env->GetArrayLength(jcn_ba);
if (env->EnsureLocalCapacity(jcf_name_len) != 0) {
// out of memory
env->DeleteLocalRef(jcn);
env->ReleaseLongArrayElements(jcolumn_options_handles, jco, JNI_ABORT);
env->ReleaseStringUTFChars(jdb_path, db_path);
return nullptr;
}
jbyte* jcf_name = env->GetByteArrayElements(jcn_ba, nullptr);
if (jcf_name == nullptr) {
// exception thrown: OutOfMemoryError

77
java/rocksjni/transaction.cc
View File

@ -252,12 +252,6 @@ std::vector<ROCKSDB_NAMESPACE::ColumnFamilyHandle*> txn_column_families_helper(
if (jcolumn_family_handles != nullptr) {
const jsize len_cols = env->GetArrayLength(jcolumn_family_handles);
if (len_cols > 0) {
if (env->EnsureLocalCapacity(len_cols) != 0) {
// out of memory
*has_exception = JNI_TRUE;
return std::vector<ROCKSDB_NAMESPACE::ColumnFamilyHandle*>();
}
jlong* jcfh = env->GetLongArrayElements(jcolumn_family_handles, nullptr);
if (jcfh == nullptr) {
// exception thrown: OutOfMemoryError
@ -293,47 +287,48 @@ void free_parts(
}
}
void free_key_values(std::vector<jbyte*>& keys_to_free) {
for (auto& key : keys_to_free) {
delete[] key;
}
}
// TODO(AR) consider refactoring to share this between here and rocksjni.cc
// cf multi get
jobjectArray txn_multi_get_helper(JNIEnv* env, const FnMultiGet& fn_multi_get,
const jlong& jread_options_handle,
const jobjectArray& jkey_parts) {
const jsize len_key_parts = env->GetArrayLength(jkey_parts);
if (env->EnsureLocalCapacity(len_key_parts) != 0) {
// out of memory
return nullptr;
}
std::vector<ROCKSDB_NAMESPACE::Slice> key_parts;
std::vector<std::tuple<jbyteArray, jbyte*, jobject>> key_parts_to_free;
std::vector<jbyte*> keys_to_free;
for (int i = 0; i < len_key_parts; i++) {
const jobject jk = env->GetObjectArrayElement(jkey_parts, i);
if (env->ExceptionCheck()) {
// exception thrown: ArrayIndexOutOfBoundsException
free_parts(env, key_parts_to_free);
free_key_values(keys_to_free);
return nullptr;
}
jbyteArray jk_ba = reinterpret_cast<jbyteArray>(jk);
const jsize len_key = env->GetArrayLength(jk_ba);
if (env->EnsureLocalCapacity(len_key) != 0) {
// out of memory
env->DeleteLocalRef(jk);
free_parts(env, key_parts_to_free);
return nullptr;
}
jbyte* jk_val = env->GetByteArrayElements(jk_ba, nullptr);
jbyte* jk_val = new jbyte[len_key];
if (jk_val == nullptr) {
// exception thrown: OutOfMemoryError
env->DeleteLocalRef(jk);
free_parts(env, key_parts_to_free);
free_key_values(keys_to_free);
jclass exception_cls = (env)->FindClass("java/lang/OutOfMemoryError");
(env)->ThrowNew(exception_cls,
"Insufficient Memory for CF handle array.");
return nullptr;
}
env->GetByteArrayRegion(jk_ba, 0, len_key, jk_val);
ROCKSDB_NAMESPACE::Slice key_slice(reinterpret_cast<char*>(jk_val),
len_key);
key_parts.push_back(key_slice);
key_parts_to_free.push_back(std::make_tuple(jk_ba, jk_val, jk));
keys_to_free.push_back(jk_val);
env->DeleteLocalRef(jk);
}
auto* read_options =
@ -343,7 +338,7 @@ jobjectArray txn_multi_get_helper(JNIEnv* env, const FnMultiGet& fn_multi_get,
fn_multi_get(*read_options, key_parts, &value_parts);
// free up allocated byte arrays
free_parts(env, key_parts_to_free);
free_key_values(keys_to_free);
// prepare the results
const jclass jcls_ba = env->FindClass("[B");
@ -658,6 +653,20 @@ void txn_write_kv_parts_helper(JNIEnv* env,
auto value_parts = std::vector<ROCKSDB_NAMESPACE::Slice>();
auto jparts_to_free = std::vector<std::tuple<jbyteArray, jbyte*, jobject>>();
// Since this is fundamentally a gather write at the RocksDB level,
// it seems wrong to refactor it by copying (gathering) keys and data here,
// in order to avoid the local reference limit.
// The user needs to be a aware that there is a limit to the number of parts
// which can be gathered.
if (env->EnsureLocalCapacity(jkey_parts_len + jvalue_parts_len) != 0) {
// no space for all the jobjects we store up
env->ExceptionClear();
ROCKSDB_NAMESPACE::RocksDBExceptionJni::ThrowNew(
env, "Insufficient JNI local references for " +
std::to_string(jkey_parts_len) + " key/value parts");
return;
}
// convert java key_parts/value_parts byte[][] to Slice(s)
for (jsize i = 0; i < jkey_parts_len; ++i) {
const jobject jobj_key_part = env->GetObjectArrayElement(jkey_parts, i);
@ -676,13 +685,6 @@ void txn_write_kv_parts_helper(JNIEnv* env,
const jbyteArray jba_key_part = reinterpret_cast<jbyteArray>(jobj_key_part);
const jsize jkey_part_len = env->GetArrayLength(jba_key_part);
if (env->EnsureLocalCapacity(jkey_part_len) != 0) {
// out of memory
env->DeleteLocalRef(jobj_value_part);
env->DeleteLocalRef(jobj_key_part);
free_parts(env, jparts_to_free);
return;
}
jbyte* jkey_part = env->GetByteArrayElements(jba_key_part, nullptr);
if (jkey_part == nullptr) {
// exception thrown: OutOfMemoryError
@ -695,18 +697,9 @@ void txn_write_kv_parts_helper(JNIEnv* env,
const jbyteArray jba_value_part =
reinterpret_cast<jbyteArray>(jobj_value_part);
const jsize jvalue_part_len = env->GetArrayLength(jba_value_part);
if (env->EnsureLocalCapacity(jvalue_part_len) != 0) {
// out of memory
env->DeleteLocalRef(jobj_value_part);
env->DeleteLocalRef(jobj_key_part);
env->ReleaseByteArrayElements(jba_key_part, jkey_part, JNI_ABORT);
free_parts(env, jparts_to_free);
return;
}
jbyte* jvalue_part = env->GetByteArrayElements(jba_value_part, nullptr);
if (jvalue_part == nullptr) {
// exception thrown: OutOfMemoryError
env->ReleaseByteArrayElements(jba_value_part, jvalue_part, JNI_ABORT);
env->DeleteLocalRef(jobj_value_part);
env->DeleteLocalRef(jobj_key_part);
env->ReleaseByteArrayElements(jba_key_part, jkey_part, JNI_ABORT);
@ -911,12 +904,6 @@ void txn_write_k_parts_helper(JNIEnv* env,
const jbyteArray jba_key_part = reinterpret_cast<jbyteArray>(jobj_key_part);
const jsize jkey_part_len = env->GetArrayLength(jba_key_part);
if (env->EnsureLocalCapacity(jkey_part_len) != 0) {
// out of memory
env->DeleteLocalRef(jobj_key_part);
free_parts(env, jkey_parts_to_free);
return;
}
jbyte* jkey_part = env->GetByteArrayElements(jba_key_part, nullptr);
if (jkey_part == nullptr) {
// exception thrown: OutOfMemoryError

14
java/rocksjni/transaction_db.cc
View File

@ -67,12 +67,6 @@ jlongArray Java_org_rocksdb_TransactionDB_open__JJLjava_lang_String_2_3_3B_3J(
}
const jsize len_cols = env->GetArrayLength(jcolumn_names);
if (env->EnsureLocalCapacity(len_cols) != 0) {
// out of memory
env->ReleaseStringUTFChars(jdb_path, db_path);
return nullptr;
}
jlong* jco = env->GetLongArrayElements(jcolumn_options_handles, nullptr);
if (jco == nullptr) {
// exception thrown: OutOfMemoryError
@ -99,14 +93,6 @@ jlongArray Java_org_rocksdb_TransactionDB_open__JJLjava_lang_String_2_3_3B_3J(
}
const int jcf_name_len = env->GetArrayLength(jcn_ba);
if (env->EnsureLocalCapacity(jcf_name_len) != 0) {
// out of memory
env->ReleaseByteArrayElements(jcn_ba, jcf_name, JNI_ABORT);
env->DeleteLocalRef(jcn);
env->ReleaseLongArrayElements(jcolumn_options_handles, jco, JNI_ABORT);
env->ReleaseStringUTFChars(jdb_path, db_path);
return nullptr;
}
const std::string cf_name(reinterpret_cast<char*>(jcf_name), jcf_name_len);
const ROCKSDB_NAMESPACE::ColumnFamilyOptions* cf_options =
reinterpret_cast<ROCKSDB_NAMESPACE::ColumnFamilyOptions*>(jco[i]);

140
java/src/test/java/org/rocksdb/MultiColumnRegressionTest.java Normal file
View File

@ -0,0 +1,140 @@
package org.rocksdb;
import static org.assertj.core.api.Assertions.assertThat;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;
/**
* Test for changes made by
* <a link="https://github.com/facebook/rocksdb/issues/9006">transactional multiGet problem</a>
* the tests here were previously broken by the nonsense removed by that change.
*/
@RunWith(Parameterized.class)
public class MultiColumnRegressionTest {
@Parameterized.Parameters
public static List<Params> data() {
return Arrays.asList(new Params(3, 100), new Params(3, 1000000));
}
public static class Params {
final int numColumns;
final int keySize;
public Params(final int numColumns, final int keySize) {
this.numColumns = numColumns;
this.keySize = keySize;
}
}
@Rule public TemporaryFolder dbFolder = new TemporaryFolder();
private final Params params;
public MultiColumnRegressionTest(final Params params) {
this.params = params;
}
@Test
public void transactionDB() throws RocksDBException {
final List<ColumnFamilyDescriptor> columnFamilyDescriptors = new ArrayList<>();
for (int i = 0; i < params.numColumns; i++) {
StringBuilder sb = new StringBuilder();
sb.append("cf" + i);
for (int j = 0; j < params.keySize; j++) sb.append("_cf");
columnFamilyDescriptors.add(new ColumnFamilyDescriptor(sb.toString().getBytes()));
}
try (final Options opt = new Options().setCreateIfMissing(true);
final RocksDB db = RocksDB.open(opt, dbFolder.getRoot().getAbsolutePath())) {
final List<ColumnFamilyHandle> columnFamilyHandles =
db.createColumnFamilies(columnFamilyDescriptors);
}
columnFamilyDescriptors.add(new ColumnFamilyDescriptor("default".getBytes()));
final List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>();
try (final TransactionDB tdb = TransactionDB.open(new DBOptions().setCreateIfMissing(true),
new TransactionDBOptions(), dbFolder.getRoot().getAbsolutePath(),
columnFamilyDescriptors, columnFamilyHandles)) {
final WriteOptions writeOptions = new WriteOptions();
try (Transaction transaction = tdb.beginTransaction(writeOptions)) {
for (int i = 0; i < params.numColumns; i++) {
transaction.put(
columnFamilyHandles.get(i), ("key" + i).getBytes(), ("value" + (i - 7)).getBytes());
}
transaction.put("key".getBytes(), "value".getBytes());
transaction.commit();
}
for (ColumnFamilyHandle columnFamilyHandle : columnFamilyHandles) {
columnFamilyHandle.close();
}
}
final List<ColumnFamilyHandle> columnFamilyHandles2 = new ArrayList<>();
try (final TransactionDB tdb = TransactionDB.open(new DBOptions().setCreateIfMissing(true),
new TransactionDBOptions(), dbFolder.getRoot().getAbsolutePath(),
columnFamilyDescriptors, columnFamilyHandles2)) {
try (Transaction transaction = tdb.beginTransaction(new WriteOptions())) {
final ReadOptions readOptions = new ReadOptions();
for (int i = 0; i < params.numColumns; i++) {
final byte[] value =
transaction.get(columnFamilyHandles2.get(i), readOptions, ("key" + i).getBytes());
assertThat(value).isEqualTo(("value" + (i - 7)).getBytes());
}
transaction.commit();
}
for (ColumnFamilyHandle columnFamilyHandle : columnFamilyHandles2) {
columnFamilyHandle.close();
}
}
}
@Test
public void optimisticDB() throws RocksDBException {
final List<ColumnFamilyDescriptor> columnFamilyDescriptors = new ArrayList<>();
for (int i = 0; i < params.numColumns; i++) {
columnFamilyDescriptors.add(new ColumnFamilyDescriptor("default".getBytes()));
}
columnFamilyDescriptors.add(new ColumnFamilyDescriptor("default".getBytes()));
final List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>();
try (final OptimisticTransactionDB otdb = OptimisticTransactionDB.open(
new DBOptions().setCreateIfMissing(true), dbFolder.getRoot().getAbsolutePath(),
columnFamilyDescriptors, columnFamilyHandles)) {
try (Transaction transaction = otdb.beginTransaction(new WriteOptions())) {
for (int i = 0; i < params.numColumns; i++) {
transaction.put(
columnFamilyHandles.get(i), ("key" + i).getBytes(), ("value" + (i - 7)).getBytes());
}
transaction.put("key".getBytes(), "value".getBytes());
transaction.commit();
}
for (ColumnFamilyHandle columnFamilyHandle : columnFamilyHandles) {
columnFamilyHandle.close();
}
}
final List<ColumnFamilyHandle> columnFamilyHandles2 = new ArrayList<>();
try (final OptimisticTransactionDB otdb = OptimisticTransactionDB.open(
new DBOptions().setCreateIfMissing(true), dbFolder.getRoot().getAbsolutePath(),
columnFamilyDescriptors, columnFamilyHandles2)) {
try (Transaction transaction = otdb.beginTransaction(new WriteOptions())) {
final ReadOptions readOptions = new ReadOptions();
for (int i = 0; i < params.numColumns; i++) {
final byte[] value =
transaction.get(columnFamilyHandles2.get(i), readOptions, ("key" + i).getBytes());
assertThat(value).isEqualTo(("value" + (i - 7)).getBytes());
}
transaction.commit();
}
for (ColumnFamilyHandle columnFamilyHandle : columnFamilyHandles2) {
columnFamilyHandle.close();
}
}
}
}

219
java/src/test/java/org/rocksdb/MultiGetManyKeysTest.java
View File

@ -6,7 +6,6 @@ package org.rocksdb;
import static org.assertj.core.api.Assertions.assertThat;
import java.nio.charset.StandardCharsets;
import java.util.*;
import org.junit.Rule;
import org.junit.Test;
@ -18,53 +17,225 @@ import org.junit.runners.Parameterized;
public class MultiGetManyKeysTest {
@Parameterized.Parameters
public static List<Integer> data() {
return Arrays.asList(3, 250, 60000, 70000, 150000, 750000);
return Arrays.asList(2, 3, 250, 60000, 70000, 150000, 750000);
}
@Rule public TemporaryFolder dbFolder = new TemporaryFolder();
private final int keySize;
private final int numKeys;
public MultiGetManyKeysTest(final Integer keySize) {
this.keySize = keySize;
public MultiGetManyKeysTest(final Integer numKeys) {
this.numKeys = numKeys;
}
/**
* Test for https://github.com/facebook/rocksdb/issues/8039
* Test for <a link="https://github.com/facebook/rocksdb/issues/8039">multiGet problem</a>
*/
@Test
public void multiGetAsListLarge() throws RocksDBException {
final List<byte[]> keys = generateRandomKeys(numKeys);
final Map<Key, byte[]> keyValues = generateRandomKeyValues(keys, 10);
putKeysAndValues(keyValues);
try (final Options opt = new Options().setCreateIfMissing(true);
final RocksDB db = RocksDB.open(opt, dbFolder.getRoot().getAbsolutePath())) {
final List<byte[]> values = db.multiGetAsList(keys);
assertKeysAndValues(keys, keyValues, values);
}
}
/**
* Test for <a link="https://github.com/facebook/rocksdb/issues/9006">transactional multiGet
* problem</a>
*/
@Test
public void multiGetAsListLargeTransactional() throws RocksDBException {
final List<byte[]> keys = generateRandomKeys(numKeys);
final Map<Key, byte[]> keyValues = generateRandomKeyValues(keys, 10);
putKeysAndValues(keyValues);
try (final Options options = new Options().setCreateIfMissing(true);
final TransactionDBOptions txnDbOptions = new TransactionDBOptions();
final TransactionDB txnDB =
TransactionDB.open(options, txnDbOptions, dbFolder.getRoot().getAbsolutePath())) {
try (final Transaction transaction = txnDB.beginTransaction(new WriteOptions())) {
final List<byte[]> values = transaction.multiGetAsList(new ReadOptions(), keys);
assertKeysAndValues(keys, keyValues, values);
}
}
}
/**
* Test for <a link="https://github.com/facebook/rocksdb/issues/9006">transactional multiGet
* problem</a>
*/
@Test
public void multiGetForUpdateAsListLargeTransactional() throws RocksDBException {
final List<byte[]> keys = generateRandomKeys(numKeys);
final Map<Key, byte[]> keyValues = generateRandomKeyValues(keys, 10);
putKeysAndValues(keyValues);
try (final Options options = new Options().setCreateIfMissing(true);
final TransactionDBOptions txnDbOptions = new TransactionDBOptions();
final TransactionDB txnDB =
TransactionDB.open(options, txnDbOptions, dbFolder.getRoot().getAbsolutePath())) {
try (final Transaction transaction = txnDB.beginTransaction(new WriteOptions())) {
final List<byte[]> values = transaction.multiGetForUpdateAsList(new ReadOptions(), keys);
assertKeysAndValues(keys, keyValues, values);
}
}
}
/**
* Test for <a link="https://github.com/facebook/rocksdb/issues/9006">transactional multiGet
* problem</a>
*/
@Test
public void multiGetAsListLargeTransactionalCF() throws RocksDBException {
final List<byte[]> keys = generateRandomKeys(numKeys);
final Map<Key, byte[]> keyValues = generateRandomKeyValues(keys, 10);
final ColumnFamilyDescriptor columnFamilyDescriptor =
new ColumnFamilyDescriptor("cfTest".getBytes());
putKeysAndValues(columnFamilyDescriptor, keyValues);
final List<ColumnFamilyDescriptor> columnFamilyDescriptors = new ArrayList<>();
columnFamilyDescriptors.add(columnFamilyDescriptor);
columnFamilyDescriptors.add(new ColumnFamilyDescriptor("default".getBytes()));
final List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>();
try (final Options options = new Options().setCreateIfMissing(true);
final TransactionDBOptions txnDbOptions = new TransactionDBOptions();
final TransactionDB txnDB = TransactionDB.open(new DBOptions(options), txnDbOptions,
dbFolder.getRoot().getAbsolutePath(), columnFamilyDescriptors, columnFamilyHandles)) {
final List<ColumnFamilyHandle> columnFamilyHandlesForMultiGet = new ArrayList<>(numKeys);
for (int i = 0; i < numKeys; i++)
columnFamilyHandlesForMultiGet.add(columnFamilyHandles.get(0));
try (final Transaction transaction = txnDB.beginTransaction(new WriteOptions())) {
final List<byte[]> values =
transaction.multiGetAsList(new ReadOptions(), columnFamilyHandlesForMultiGet, keys);
assertKeysAndValues(keys, keyValues, values);
}
for (ColumnFamilyHandle columnFamilyHandle : columnFamilyHandles) {
columnFamilyHandle.close();
}
}
}
/**
* Test for <a link="https://github.com/facebook/rocksdb/issues/9006">transactional multiGet
* problem</a>
*/
@Test
public void multiGetForUpdateAsListLargeTransactionalCF() throws RocksDBException {
final List<byte[]> keys = generateRandomKeys(numKeys);
final Map<Key, byte[]> keyValues = generateRandomKeyValues(keys, 10);
final ColumnFamilyDescriptor columnFamilyDescriptor =
new ColumnFamilyDescriptor("cfTest".getBytes());
putKeysAndValues(columnFamilyDescriptor, keyValues);
final List<ColumnFamilyDescriptor> columnFamilyDescriptors = new ArrayList<>();
columnFamilyDescriptors.add(columnFamilyDescriptor);
columnFamilyDescriptors.add(new ColumnFamilyDescriptor("default".getBytes()));
final List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>();
try (final Options options = new Options().setCreateIfMissing(true);
final TransactionDBOptions txnDbOptions = new TransactionDBOptions();
final TransactionDB txnDB = TransactionDB.open(new DBOptions(options), txnDbOptions,
dbFolder.getRoot().getAbsolutePath(), columnFamilyDescriptors, columnFamilyHandles)) {
final List<ColumnFamilyHandle> columnFamilyHandlesForMultiGet = new ArrayList<>(numKeys);
for (int i = 0; i < numKeys; i++)
columnFamilyHandlesForMultiGet.add(columnFamilyHandles.get(0));
try (final Transaction transaction = txnDB.beginTransaction(new WriteOptions())) {
final List<byte[]> values = transaction.multiGetForUpdateAsList(
new ReadOptions(), columnFamilyHandlesForMultiGet, keys);
assertKeysAndValues(keys, keyValues, values);
}
for (ColumnFamilyHandle columnFamilyHandle : columnFamilyHandles) {
columnFamilyHandle.close();
}
}
}
private List<byte[]> generateRandomKeys(final int numKeys) {
final Random rand = new Random();
final List<byte[]> keys = new ArrayList<>();
for (int i = 0; i < keySize; i++) {
for (int i = 0; i < numKeys; i++) {
final byte[] key = new byte[4];
rand.nextBytes(key);
keys.add(key);
}
return keys;
}
try (final Options opt = new Options().setCreateIfMissing(true);
final RocksDB db = RocksDB.open(opt, dbFolder.getRoot().getAbsolutePath())) {
final List<byte[]> values = db.multiGetAsList(keys);
assertThat(values.size()).isEqualTo(keys.size());
private Map<Key, byte[]> generateRandomKeyValues(final List<byte[]> keys, final int percent) {
final Random rand = new Random();
final Map<Key, byte[]> keyValues = new HashMap<>();
for (int i = 0; i < numKeys; i++) {
if (rand.nextInt(100) < percent) {
final byte[] value = new byte[1024];
rand.nextBytes(value);
keyValues.put(new Key(keys.get(i)), value);
}
}
return keyValues;
}
private void putKeysAndValues(Map<Key, byte[]> keyValues) throws RocksDBException {
try (final Options options = new Options().setCreateIfMissing(true);
final RocksDB db = RocksDB.open(options, dbFolder.getRoot().getAbsolutePath())) {
for (Map.Entry<Key, byte[]> keyValue : keyValues.entrySet()) {
db.put(keyValue.getKey().get(), keyValue.getValue());
}
}
}
@Test
public void multiGetAsListCheckResults() throws RocksDBException {
try (final Options opt = new Options().setCreateIfMissing(true);
final RocksDB db = RocksDB.open(opt, dbFolder.getRoot().getAbsolutePath())) {
final List<byte[]> keys = new ArrayList<>();
for (int i = 0; i < keySize; i++) {
byte[] key = ("key" + i + ":").getBytes();
keys.add(key);
db.put(key, ("value" + i + ":").getBytes());
private void putKeysAndValues(ColumnFamilyDescriptor columnFamilyDescriptor,
Map<Key, byte[]> keyValues) throws RocksDBException {
try (final Options options = new Options().setCreateIfMissing(true);
final RocksDB db = RocksDB.open(options, dbFolder.getRoot().getAbsolutePath());
final ColumnFamilyHandle columnFamilyHandle =
db.createColumnFamily(columnFamilyDescriptor)) {
for (Map.Entry<Key, byte[]> keyValue : keyValues.entrySet()) {
db.put(columnFamilyHandle, keyValue.getKey().get(), keyValue.getValue());
}
}
}
final List<byte[]> values = db.multiGetAsList(keys);
assertThat(values.size()).isEqualTo(keys.size());
for (int i = 0; i < keySize; i++) {
assertThat(values.get(i)).isEqualTo(("value" + i + ":").getBytes());
private void assertKeysAndValues(
final List<byte[]> keys, final Map<Key, byte[]> keyValues, final List<byte[]> values) {
assertThat(values.size()).isEqualTo(keys.size());
for (int i = 0; i < numKeys; i++) {
final Key key = new Key(keys.get(i));
final byte[] value = values.get(i);
if (keyValues.containsKey(key)) {
assertThat(value).isEqualTo(keyValues.get(key));
} else {
assertThat(value).isNull();
}
}
}
static private class Key {
private final byte[] bytes;
public Key(byte[] bytes) {
this.bytes = bytes;
}
public byte[] get() {
return this.bytes;
}
@Override
public boolean equals(Object o) {
if (this == o)
return true;
if (o == null || getClass() != o.getClass())
return false;
Key key = (Key) o;
return Arrays.equals(bytes, key.bytes);
}
@Override
public int hashCode() {
return Arrays.hashCode(bytes);
}
}
}

158
java/src/test/java/org/rocksdb/PutMultiplePartsTest.java Normal file
View File

@ -0,0 +1,158 @@
package org.rocksdb;
import static org.assertj.core.api.Assertions.assertThat;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;
@RunWith(Parameterized.class)
public class PutMultiplePartsTest {
@Parameterized.Parameters
public static List<Integer> data() {
return Arrays.asList(2, 3, 250, 20000);
}
@Rule public TemporaryFolder dbFolder = new TemporaryFolder();
private final int numParts;
public PutMultiplePartsTest(final Integer numParts) {
this.numParts = numParts;
}
@Test
public void putUntracked() throws RocksDBException {
try (final Options options = new Options().setCreateIfMissing(true);
final TransactionDBOptions txnDbOptions = new TransactionDBOptions();
final TransactionDB txnDB =
TransactionDB.open(options, txnDbOptions, dbFolder.getRoot().getAbsolutePath())) {
try (final Transaction transaction = txnDB.beginTransaction(new WriteOptions())) {
final byte[][] keys = generateItems("key", ":", numParts);
final byte[][] values = generateItems("value", "", numParts);
transaction.putUntracked(keys, values);
transaction.commit();
}
txnDB.syncWal();
}
validateResults();
}
@Test
public void put() throws RocksDBException {
try (final Options options = new Options().setCreateIfMissing(true);
final TransactionDBOptions txnDbOptions = new TransactionDBOptions();
final TransactionDB txnDB =
TransactionDB.open(options, txnDbOptions, dbFolder.getRoot().getAbsolutePath())) {
try (final Transaction transaction = txnDB.beginTransaction(new WriteOptions())) {
final byte[][] keys = generateItems("key", ":", numParts);
final byte[][] values = generateItems("value", "", numParts);
transaction.put(keys, values);
transaction.commit();
}
txnDB.syncWal();
}
validateResults();
}
@Test
public void putUntrackedCF() throws RocksDBException {
try (final Options options = new Options().setCreateIfMissing(true);
final TransactionDBOptions txnDbOptions = new TransactionDBOptions();
final TransactionDB txnDB =
TransactionDB.open(options, txnDbOptions, dbFolder.getRoot().getAbsolutePath());
final ColumnFamilyHandle columnFamilyHandle =
txnDB.createColumnFamily(new ColumnFamilyDescriptor("cfTest".getBytes()))) {
try (final Transaction transaction = txnDB.beginTransaction(new WriteOptions())) {
final byte[][] keys = generateItems("key", ":", numParts);
final byte[][] values = generateItems("value", "", numParts);
transaction.putUntracked(columnFamilyHandle, keys, values);
transaction.commit();
}
txnDB.syncWal();
}
validateResultsCF();
}
@Test
public void putCF() throws RocksDBException {
try (final Options options = new Options().setCreateIfMissing(true);
final TransactionDBOptions txnDbOptions = new TransactionDBOptions();
final TransactionDB txnDB =
TransactionDB.open(options, txnDbOptions, dbFolder.getRoot().getAbsolutePath());
final ColumnFamilyHandle columnFamilyHandle =
txnDB.createColumnFamily(new ColumnFamilyDescriptor("cfTest".getBytes()))) {
try (final Transaction transaction = txnDB.beginTransaction(new WriteOptions())) {
final byte[][] keys = generateItems("key", ":", numParts);
final byte[][] values = generateItems("value", "", numParts);
transaction.put(columnFamilyHandle, keys, values);
transaction.commit();
}
txnDB.syncWal();
}
validateResultsCF();
}
private void validateResults() throws RocksDBException {
try (final RocksDB db = RocksDB.open(new Options(), dbFolder.getRoot().getAbsolutePath())) {
final List<byte[]> keys = generateItemsAsList("key", ":", numParts);
final byte[][] values = generateItems("value", "", numParts);
StringBuilder singleKey = new StringBuilder();
for (int i = 0; i < numParts; i++) {
singleKey.append(new String(keys.get(i), StandardCharsets.UTF_8));
}
final byte[] result = db.get(singleKey.toString().getBytes());
StringBuilder singleValue = new StringBuilder();
for (int i = 0; i < numParts; i++) {
singleValue.append(new String(values[i], StandardCharsets.UTF_8));
}
assertThat(result).isEqualTo(singleValue.toString().getBytes());
}
}
private void validateResultsCF() throws RocksDBException {
final List<ColumnFamilyDescriptor> columnFamilyDescriptors = new ArrayList<>();
columnFamilyDescriptors.add(new ColumnFamilyDescriptor("cfTest".getBytes()));
columnFamilyDescriptors.add(new ColumnFamilyDescriptor("default".getBytes()));
final List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>();
try (final RocksDB db = RocksDB.open(new DBOptions(), dbFolder.getRoot().getAbsolutePath(),
columnFamilyDescriptors, columnFamilyHandles)) {
final List<byte[]> keys = generateItemsAsList("key", ":", numParts);
final byte[][] values = generateItems("value", "", numParts);
StringBuilder singleKey = new StringBuilder();
for (int i = 0; i < numParts; i++) {
singleKey.append(new String(keys.get(i), StandardCharsets.UTF_8));
}
final byte[] result = db.get(columnFamilyHandles.get(0), singleKey.toString().getBytes());
StringBuilder singleValue = new StringBuilder();
for (int i = 0; i < numParts; i++) {
singleValue.append(new String(values[i], StandardCharsets.UTF_8));
}
assertThat(result).isEqualTo(singleValue.toString().getBytes());
}
}
private byte[][] generateItems(final String prefix, final String suffix, final int numItems) {
return generateItemsAsList(prefix, suffix, numItems).toArray(new byte[0][0]);
}
private List<byte[]> generateItemsAsList(
final String prefix, final String suffix, final int numItems) {
final List<byte[]> items = new ArrayList<>();
for (int i = 0; i < numItems; i++) {
items.add((prefix + i + suffix).getBytes());
}
return items;
}
}