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()); + } + } +}