From f5526af8ed028e8a8c768ed5914b6fa92fd8fc57 Mon Sep 17 00:00:00 2001 From: Alan Paxton Date: Thu, 14 Oct 2021 11:46:59 -0700 Subject: [PATCH] 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 --- java/Makefile | 2 + java/rocksjni/rocksjni.cc | 45 ++++++------ .../org/rocksdb/MultiGetManyKeysTest.java | 70 +++++++++++++++++++ .../test/java/org/rocksdb/MultiGetTest.java | 34 +++++++++ 4 files changed, 130 insertions(+), 21 deletions(-) create mode 100644 java/src/test/java/org/rocksdb/MultiGetManyKeysTest.java create mode 100644 java/src/test/java/org/rocksdb/MultiGetTest.java diff --git a/java/Makefile b/java/Makefile index 1c7023712..03ff58651 100644 --- a/java/Makefile +++ b/java/Makefile @@ -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\ diff --git a/java/rocksjni/rocksjni.cc b/java/rocksjni/rocksjni.cc index d00b676da..96d4177fb 100644 --- a/java/rocksjni/rocksjni.cc +++ b/java/rocksjni/rocksjni.cc @@ -1676,12 +1676,10 @@ jint Java_org_rocksdb_RocksDB_get__JJ_3BII_3BIIJ( } } -inline void multi_get_helper_release_keys( - JNIEnv* env, std::vector>& keys_to_free) { +inline void multi_get_helper_release_keys(std::vector& 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 keys; - std::vector> keys_to_free; + std::vector 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(key), len_key); keys.push_back(key_slice); - keys_to_free.push_back(std::pair(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(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(s.size())) != 0) { - // exception thrown: OutOfMemoryError - return nullptr; - } // add to the jresults for (std::vector::size_type i = 0; i != s.size(); i++) { diff --git a/java/src/test/java/org/rocksdb/MultiGetManyKeysTest.java b/java/src/test/java/org/rocksdb/MultiGetManyKeysTest.java new file mode 100644 index 000000000..9a23be788 --- /dev/null +++ b/java/src/test/java/org/rocksdb/MultiGetManyKeysTest.java @@ -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 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 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 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 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 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()); + } + } + } +} diff --git a/java/src/test/java/org/rocksdb/MultiGetTest.java b/java/src/test/java/org/rocksdb/MultiGetTest.java new file mode 100644 index 000000000..ce8212409 --- /dev/null +++ b/java/src/test/java/org/rocksdb/MultiGetTest.java @@ -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 keys = + Arrays.asList("key1".getBytes(), "key2".getBytes(), "key3".getBytes()); + final List 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()); + } + } +}