Fix multiget throwing NPE for num of keys > 70k (#9012)

Summary:
closes https://github.com/facebook/rocksdb/issues/8039

Unnecessary use of multiple local JNI references at the same time, 1 per key, was limiting the size of the key array. The local references don't need to be held simultaneously, so if we rearrange the code we can make it work for bigger key arrays.

Incidentally, make errors throw helpful exception messages rather than returning a null pointer.

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

Reviewed By: mrambacher

Differential Revision: D31580862

Pulled By: jay-zhuang

fbshipit-source-id: ce05831d52ede332e1b20e74d2dc621d219b9616
main
Alan Paxton 3 years ago committed by Facebook GitHub Bot
parent ffc48b6cad
commit f5526af8ed
  1. 2
      java/Makefile
  2. 45
      java/rocksjni/rocksjni.cc
  3. 70
      java/src/test/java/org/rocksdb/MultiGetManyKeysTest.java
  4. 34
      java/src/test/java/org/rocksdb/MultiGetTest.java

@ -145,6 +145,8 @@ JAVA_TESTS = \
org.rocksdb.MemoryUtilTest\
org.rocksdb.MemTableTest\
org.rocksdb.MergeTest\
org.rocksdb.MultiGetManyKeysTest\
org.rocksdb.MultiGetTest\
org.rocksdb.MixedOptionsTest\
org.rocksdb.MutableColumnFamilyOptionsTest\
org.rocksdb.MutableDBOptionsTest\

@ -1676,12 +1676,10 @@ jint Java_org_rocksdb_RocksDB_get__JJ_3BII_3BIIJ(
}
}
inline void multi_get_helper_release_keys(
JNIEnv* env, std::vector<std::pair<jbyte*, jobject>>& keys_to_free) {
inline void multi_get_helper_release_keys(std::vector<jbyte*>& keys_to_free) {
auto end = keys_to_free.end();
for (auto it = keys_to_free.begin(); it != end; ++it) {
delete[] it->first;
env->DeleteLocalRef(it->second);
delete[] * it;
}
keys_to_free.clear();
}
@ -1703,6 +1701,9 @@ jobjectArray multi_get_helper(JNIEnv* env, jobject, ROCKSDB_NAMESPACE::DB* db,
jlong* jcfh = env->GetLongArrayElements(jcolumn_family_handles, nullptr);
if (jcfh == nullptr) {
// exception thrown: OutOfMemoryError
jclass exception_cls = (env)->FindClass("java/lang/OutOfMemoryError");
(env)->ThrowNew(exception_cls,
"Insufficient Memory for CF handle array.");
return nullptr;
}
@ -1714,15 +1715,11 @@ jobjectArray multi_get_helper(JNIEnv* env, jobject, ROCKSDB_NAMESPACE::DB* db,
env->ReleaseLongArrayElements(jcolumn_family_handles, jcfh, JNI_ABORT);
}
const jsize len_keys = env->GetArrayLength(jkeys);
if (env->EnsureLocalCapacity(len_keys) != 0) {
// exception thrown: OutOfMemoryError
return nullptr;
}
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 nullptr;
}
@ -1730,18 +1727,24 @@ jobjectArray multi_get_helper(JNIEnv* env, jobject, ROCKSDB_NAMESPACE::DB* db,
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 nullptr;
}
const jsize len_keys = env->GetArrayLength(jkeys);
std::vector<ROCKSDB_NAMESPACE::Slice> keys;
std::vector<std::pair<jbyte*, jobject>> keys_to_free;
std::vector<jbyte*> keys_to_free;
for (jsize i = 0; i < len_keys; i++) {
jobject jkey = env->GetObjectArrayElement(jkeys, i);
if (env->ExceptionCheck()) {
// exception thrown: ArrayIndexOutOfBoundsException
env->ReleaseIntArrayElements(jkey_lens, jkey_len, JNI_ABORT);
env->ReleaseIntArrayElements(jkey_offs, jkey_off, JNI_ABORT);
multi_get_helper_release_keys(env, keys_to_free);
multi_get_helper_release_keys(keys_to_free);
jclass exception_cls = (env)->FindClass("java/lang/OutOfMemoryError");
(env)->ThrowNew(exception_cls,
"Insufficient Memory for key object array.");
return nullptr;
}
@ -1756,14 +1759,18 @@ jobjectArray multi_get_helper(JNIEnv* env, jobject, ROCKSDB_NAMESPACE::DB* db,
env->DeleteLocalRef(jkey);
env->ReleaseIntArrayElements(jkey_lens, jkey_len, JNI_ABORT);
env->ReleaseIntArrayElements(jkey_offs, jkey_off, JNI_ABORT);
multi_get_helper_release_keys(env, keys_to_free);
multi_get_helper_release_keys(keys_to_free);
jclass exception_cls =
(env)->FindClass("java/lang/ArrayIndexOutOfBoundsException");
(env)->ThrowNew(exception_cls, "Invalid byte array region index.");
return nullptr;
}
ROCKSDB_NAMESPACE::Slice key_slice(reinterpret_cast<char*>(key), len_key);
keys.push_back(key_slice);
keys_to_free.push_back(std::pair<jbyte*, jobject>(key, jkey));
env->DeleteLocalRef(jkey);
keys_to_free.push_back(key);
}
// cleanup jkey_off and jken_len
@ -1779,22 +1786,18 @@ jobjectArray multi_get_helper(JNIEnv* env, jobject, ROCKSDB_NAMESPACE::DB* db,
}
// free up allocated byte arrays
multi_get_helper_release_keys(env, keys_to_free);
multi_get_helper_release_keys(keys_to_free);
// 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 nullptr;
}
// TODO(AR) it is not clear to me why EnsureLocalCapacity is needed for the
// loop as we cleanup references with env->DeleteLocalRef(jentry_value);
if (env->EnsureLocalCapacity(static_cast<jint>(s.size())) != 0) {
// exception thrown: OutOfMemoryError
return nullptr;
}
// add to the jresults
for (std::vector<ROCKSDB_NAMESPACE::Status>::size_type i = 0; i != s.size();
i++) {

@ -0,0 +1,70 @@
// Copyright (c) 2011-present, Facebook, Inc. All rights reserved.
// This source code is licensed under both the GPLv2 (found in the
// COPYING file in the root directory) and Apache 2.0 License
// (found in the LICENSE.Apache file in the root directory).
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;
import org.junit.rules.TemporaryFolder;
import org.junit.runner.RunWith;
import org.junit.runners.Parameterized;
@RunWith(Parameterized.class)
public class MultiGetManyKeysTest {
@Parameterized.Parameters
public static List<Integer> data() {
return Arrays.asList(3, 250, 60000, 70000, 150000, 750000);
}
@Rule public TemporaryFolder dbFolder = new TemporaryFolder();
private final int keySize;
public MultiGetManyKeysTest(final Integer keySize) {
this.keySize = keySize;
}
/**
* Test for https://github.com/facebook/rocksdb/issues/8039
*/
@Test
public void multiGetAsListLarge() throws RocksDBException {
final Random rand = new Random();
final List<byte[]> keys = new ArrayList<>();
for (int i = 0; i < keySize; i++) {
final byte[] key = new byte[4];
rand.nextBytes(key);
keys.add(key);
}
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());
}
}
@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());
}
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());
}
}
}
}

@ -0,0 +1,34 @@
// Copyright (c) 2011-present, Facebook, Inc. All rights reserved.
// This source code is licensed under both the GPLv2 (found in the
// COPYING file in the root directory) and Apache 2.0 License
// (found in the LICENSE.Apache file in the root directory).
package org.rocksdb;
import static org.assertj.core.api.Assertions.assertThat;
import java.util.Arrays;
import java.util.List;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
public class MultiGetTest {
@Rule public TemporaryFolder dbFolder = new TemporaryFolder();
@Test
public void putNThenMultiGet() 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<byte[]> keys =
Arrays.asList("key1".getBytes(), "key2".getBytes(), "key3".getBytes());
final List<byte[]> values = db.multiGetAsList(keys);
assertThat(values.size()).isEqualTo(keys.size());
assertThat(values.get(0)).isEqualTo("value1ForKey1".getBytes());
assertThat(values.get(1)).isEqualTo("value2ForKey2".getBytes());
assertThat(values.get(2)).isEqualTo("value3ForKey3".getBytes());
}
}
}
Loading…
Cancel
Save