From c9048021ad6fac842ab4d4eb66a9522a9ed1a543 Mon Sep 17 00:00:00 2001 From: Ben Clay Date: Mon, 8 Oct 2018 11:03:34 -0700 Subject: [PATCH] RocksJava: memory_util support (#4446) Summary: JNI passthrough for utilities/memory/memory_util.cc sagar0 adamretter Pull Request resolved: https://github.com/facebook/rocksdb/pull/4446 Differential Revision: D10174578 Pulled By: sagar0 fbshipit-source-id: d1d196d771dff22afb7ef7500f308233675696f8 --- java/CMakeLists.txt | 4 + java/Makefile | 3 + java/rocksjni/memory_util.cc | 100 ++++++++++++ java/rocksjni/portal.h | 78 +++++++++- .../java/org/rocksdb/MemoryUsageType.java | 72 +++++++++ .../src/main/java/org/rocksdb/MemoryUtil.java | 60 ++++++++ .../test/java/org/rocksdb/MemoryUtilTest.java | 143 ++++++++++++++++++ src.mk | 1 + 8 files changed, 460 insertions(+), 1 deletion(-) create mode 100644 java/rocksjni/memory_util.cc create mode 100644 java/src/main/java/org/rocksdb/MemoryUsageType.java create mode 100644 java/src/main/java/org/rocksdb/MemoryUtil.java create mode 100644 java/src/test/java/org/rocksdb/MemoryUtilTest.java diff --git a/java/CMakeLists.txt b/java/CMakeLists.txt index 96c08b231..986f22017 100644 --- a/java/CMakeLists.txt +++ b/java/CMakeLists.txt @@ -25,6 +25,7 @@ set(JNI_NATIVE_SOURCES rocksjni/jnicallback.cc rocksjni/loggerjnicallback.cc rocksjni/lru_cache.cc + rocksjni/memory_util.cc rocksjni/memtablejni.cc rocksjni/merge_operator.cc rocksjni/native_comparator_wrapper_test.cc @@ -96,6 +97,7 @@ set(NATIVE_JAVA_CLASSES org.rocksdb.IngestExternalFileOptions org.rocksdb.Logger org.rocksdb.LRUCache + org.rocksdb.MemoryUtil org.rocksdb.MemTableConfig org.rocksdb.NativeComparatorWrapper org.rocksdb.NativeLibraryLoader @@ -222,6 +224,8 @@ add_jar( src/main/java/org/rocksdb/IngestExternalFileOptions.java src/main/java/org/rocksdb/Logger.java src/main/java/org/rocksdb/LRUCache.java + src/main/java/org/rocksdb/MemoryUsageType.java + src/main/java/org/rocksdb/MemoryUtil.java src/main/java/org/rocksdb/MemTableConfig.java src/main/java/org/rocksdb/MergeOperator.java src/main/java/org/rocksdb/MutableColumnFamilyOptionsInterface.java diff --git a/java/Makefile b/java/Makefile index f58fff06e..112d47eb4 100644 --- a/java/Makefile +++ b/java/Makefile @@ -30,6 +30,8 @@ NATIVE_JAVA_CLASSES = org.rocksdb.AbstractCompactionFilter\ org.rocksdb.HashSkipListMemTableConfig\ org.rocksdb.Logger\ org.rocksdb.LRUCache\ + org.rocksdb.MemoryUsageType\ + org.rocksdb.MemoryUtil\ org.rocksdb.MergeOperator\ org.rocksdb.NativeComparatorWrapper\ org.rocksdb.OptimisticTransactionDB\ @@ -111,6 +113,7 @@ JAVA_TESTS = org.rocksdb.BackupableDBOptionsTest\ org.rocksdb.KeyMayExistTest\ org.rocksdb.LoggerTest\ org.rocksdb.LRUCacheTest\ + org.rocksdb.MemoryUtilTest\ org.rocksdb.MemTableTest\ org.rocksdb.MergeTest\ org.rocksdb.MixedOptionsTest\ diff --git a/java/rocksjni/memory_util.cc b/java/rocksjni/memory_util.cc new file mode 100644 index 000000000..9c2bfd04e --- /dev/null +++ b/java/rocksjni/memory_util.cc @@ -0,0 +1,100 @@ +// 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). + +#include +#include +#include +#include +#include + +#include "include/org_rocksdb_MemoryUtil.h" + +#include "rocksjni/portal.h" + +#include "rocksdb/utilities/memory_util.h" + + +/* + * Class: org_rocksdb_MemoryUtil + * Method: getApproximateMemoryUsageByType + * Signature: ([J[J)Ljava/util/Map; + */ +jobject Java_org_rocksdb_MemoryUtil_getApproximateMemoryUsageByType( + JNIEnv *env, jclass /*jclazz*/, jlongArray jdb_handles, jlongArray jcache_handles) { + + std::vector dbs; + jsize db_handle_count = env->GetArrayLength(jdb_handles); + if(db_handle_count > 0) { + jlong *ptr_jdb_handles = env->GetLongArrayElements(jdb_handles, nullptr); + if (ptr_jdb_handles == nullptr) { + // exception thrown: OutOfMemoryError + return nullptr; + } + for (jsize i = 0; i < db_handle_count; i++) { + dbs.push_back(reinterpret_cast(ptr_jdb_handles[i])); + } + env->ReleaseLongArrayElements(jdb_handles, ptr_jdb_handles, JNI_ABORT); + } + + std::unordered_set cache_set; + jsize cache_handle_count = env->GetArrayLength(jcache_handles); + if(cache_handle_count > 0) { + jlong *ptr_jcache_handles = env->GetLongArrayElements(jcache_handles, nullptr); + if (ptr_jcache_handles == nullptr) { + // exception thrown: OutOfMemoryError + return nullptr; + } + for (jsize i = 0; i < cache_handle_count; i++) { + auto *cache_ptr = + reinterpret_cast *>(ptr_jcache_handles[i]); + cache_set.insert(cache_ptr->get()); + } + env->ReleaseLongArrayElements(jcache_handles, ptr_jcache_handles, JNI_ABORT); + } + + std::map usage_by_type; + if(rocksdb::MemoryUtil::GetApproximateMemoryUsageByType(dbs, cache_set, &usage_by_type) != rocksdb::Status::OK()) { + // Non-OK status + return nullptr; + } + + jobject jusage_by_type = rocksdb::HashMapJni::construct( + env, static_cast(usage_by_type.size())); + if (jusage_by_type == nullptr) { + // exception occurred + return nullptr; + } + const rocksdb::HashMapJni::FnMapKV + fn_map_kv = + [env](const std::pair& pair) { + // Construct key + const jobject jusage_type = + rocksdb::ByteJni::valueOf(env, rocksdb::MemoryUsageTypeJni::toJavaMemoryUsageType(pair.first)); + if (jusage_type == nullptr) { + // an error occurred + return std::unique_ptr>(nullptr); + } + // Construct value + const jobject jusage_value = + rocksdb::LongJni::valueOf(env, pair.second); + if (jusage_value == nullptr) { + // an error occurred + return std::unique_ptr>(nullptr); + } + // Construct and return pointer to pair of jobjects + return std::unique_ptr>( + new std::pair(jusage_type, + jusage_value)); + }; + + if (!rocksdb::HashMapJni::putAll(env, jusage_by_type, usage_by_type.begin(), + usage_by_type.end(), fn_map_kv)) { + // exception occcurred + jusage_by_type = nullptr; + } + + return jusage_by_type; + +} diff --git a/java/rocksjni/portal.h b/java/rocksjni/portal.h index a0d1846a6..7a8d1f157 100644 --- a/java/rocksjni/portal.h +++ b/java/rocksjni/portal.h @@ -26,6 +26,7 @@ #include "rocksdb/rate_limiter.h" #include "rocksdb/status.h" #include "rocksdb/utilities/backupable_db.h" +#include "rocksdb/utilities/memory_util.h" #include "rocksdb/utilities/transaction_db.h" #include "rocksdb/utilities/write_batch_with_index.h" #include "rocksjni/compaction_filter_factory_jnicallback.h" @@ -2251,7 +2252,7 @@ class ByteJni : public JavaClass { * @param env A pointer to the Java environment * * @return The Java Method ID or nullptr if the class or method id could not - * be retieved + * be retrieved */ static jmethodID getByteValueMethod(JNIEnv* env) { jclass clazz = getJClass(env); @@ -2264,6 +2265,39 @@ class ByteJni : public JavaClass { assert(mid != nullptr); return mid; } + + /** + * Calls the Java Method: Byte#valueOf, returning a constructed Byte jobject + * + * @param env A pointer to the Java environment + * + * @return A constructing Byte object or nullptr if the class or method id could not + * be retrieved, or an exception occurred + */ + static jobject valueOf(JNIEnv* env, jbyte jprimitive_byte) { + jclass clazz = getJClass(env); + if (clazz == nullptr) { + // exception occurred accessing class + return nullptr; + } + + static jmethodID mid = + env->GetStaticMethodID(clazz, "valueOf", "(B)Ljava/lang/Byte;"); + if (mid == nullptr) { + // exception thrown: NoSuchMethodException or OutOfMemoryError + return nullptr; + } + + const jobject jbyte_obj = + env->CallStaticObjectMethod(clazz, mid, jprimitive_byte); + if (env->ExceptionCheck()) { + // exception occurred + return nullptr; + } + + return jbyte_obj; + } + }; // The portal class for java.lang.StringBuilder @@ -3795,6 +3829,48 @@ class RateLimiterModeJni { } }; +// The portal class for org.rocksdb.MemoryUsageType +class MemoryUsageTypeJni { +public: + // Returns the equivalent org.rocksdb.MemoryUsageType for the provided + // C++ rocksdb::MemoryUtil::UsageType enum + static jbyte toJavaMemoryUsageType( + const rocksdb::MemoryUtil::UsageType& usage_type) { + switch(usage_type) { + case rocksdb::MemoryUtil::UsageType::kMemTableTotal: + return 0x0; + case rocksdb::MemoryUtil::UsageType::kMemTableUnFlushed: + return 0x1; + case rocksdb::MemoryUtil::UsageType::kTableReadersTotal: + return 0x2; + case rocksdb::MemoryUtil::UsageType::kCacheTotal: + return 0x3; + default: + // undefined: use kNumUsageTypes + return 0x4; + } + } + + // Returns the equivalent C++ rocksdb::MemoryUtil::UsageType enum for the + // provided Java org.rocksdb.MemoryUsageType + static rocksdb::MemoryUtil::UsageType toCppMemoryUsageType( + jbyte usage_type) { + switch(usage_type) { + case 0x0: + return rocksdb::MemoryUtil::UsageType::kMemTableTotal; + case 0x1: + return rocksdb::MemoryUtil::UsageType::kMemTableUnFlushed; + case 0x2: + return rocksdb::MemoryUtil::UsageType::kTableReadersTotal; + case 0x3: + return rocksdb::MemoryUtil::UsageType::kCacheTotal; + default: + // undefined/default: use kNumUsageTypes + return rocksdb::MemoryUtil::UsageType::kNumUsageTypes; + } + } +}; + // The portal class for org.rocksdb.Transaction class TransactionJni : public JavaClass { public: diff --git a/java/src/main/java/org/rocksdb/MemoryUsageType.java b/java/src/main/java/org/rocksdb/MemoryUsageType.java new file mode 100644 index 000000000..3523cd0ee --- /dev/null +++ b/java/src/main/java/org/rocksdb/MemoryUsageType.java @@ -0,0 +1,72 @@ +// 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; + +/** + * MemoryUsageType + * + *

The value will be used as a key to indicate the type of memory usage + * described

+ */ +public enum MemoryUsageType { + /** + * Memory usage of all the mem-tables. + */ + kMemTableTotal((byte) 0), + /** + * Memory usage of those un-flushed mem-tables. + */ + kMemTableUnFlushed((byte) 1), + /** + * Memory usage of all the table readers. + */ + kTableReadersTotal((byte) 2), + /** + * Memory usage by Cache. + */ + kCacheTotal((byte) 3), + /** + * Max usage types - copied to keep 1:1 with native. + */ + kNumUsageTypes((byte) 4); + + /** + * Returns the byte value of the enumerations value + * + * @return byte representation + */ + public byte getValue() { + return value_; + } + + /** + *

Get the MemoryUsageType enumeration value by + * passing the byte identifier to this method.

+ * + * @param byteIdentifier of MemoryUsageType. + * + * @return MemoryUsageType instance. + * + * @throws IllegalArgumentException if the usage type for the byteIdentifier + * cannot be found + */ + public static MemoryUsageType getMemoryUsageType(final byte byteIdentifier) { + for (final MemoryUsageType MemoryUsageType : MemoryUsageType.values()) { + if (MemoryUsageType.getValue() == byteIdentifier) { + return MemoryUsageType; + } + } + + throw new IllegalArgumentException( + "Illegal value provided for MemoryUsageType."); + } + + private MemoryUsageType(byte value) { + value_ = value; + } + + private final byte value_; +} diff --git a/java/src/main/java/org/rocksdb/MemoryUtil.java b/java/src/main/java/org/rocksdb/MemoryUtil.java new file mode 100644 index 000000000..52b2175e6 --- /dev/null +++ b/java/src/main/java/org/rocksdb/MemoryUtil.java @@ -0,0 +1,60 @@ +// 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 java.util.*; + +/** + * JNI passthrough for MemoryUtil. + */ +public class MemoryUtil { + + /** + *

Returns the approximate memory usage of different types in the input + * list of DBs and Cache set. For instance, in the output map the key + * kMemTableTotal will be associated with the memory + * usage of all the mem-tables from all the input rocksdb instances.

+ * + *

Note that for memory usage inside Cache class, we will + * only report the usage of the input "cache_set" without + * including those Cache usage inside the input list "dbs" + * of DBs.

+ * + * @param dbs List of dbs to collect memory usage for. + * @param caches Set of caches to collect memory usage for. + * @return Map from {@link MemoryUsageType} to memory usage as a {@link Long}. + */ + public static Map getApproximateMemoryUsageByType(final List dbs, final Set caches) { + int dbCount = (dbs == null) ? 0 : dbs.size(); + int cacheCount = (caches == null) ? 0 : caches.size(); + long[] dbHandles = new long[dbCount]; + long[] cacheHandles = new long[cacheCount]; + if (dbCount > 0) { + ListIterator dbIter = dbs.listIterator(); + while (dbIter.hasNext()) { + dbHandles[dbIter.nextIndex()] = dbIter.next().nativeHandle_; + } + } + if (cacheCount > 0) { + // NOTE: This index handling is super ugly but I couldn't get a clean way to track both the + // index and the iterator simultaneously within a Set. + int i = 0; + for (Cache cache : caches) { + cacheHandles[i] = cache.nativeHandle_; + i++; + } + } + Map byteOutput = getApproximateMemoryUsageByType(dbHandles, cacheHandles); + Map output = new HashMap<>(); + for(Map.Entry longEntry : byteOutput.entrySet()) { + output.put(MemoryUsageType.getMemoryUsageType(longEntry.getKey()), longEntry.getValue()); + } + return output; + } + + private native static Map getApproximateMemoryUsageByType(final long[] dbHandles, + final long[] cacheHandles); +} diff --git a/java/src/test/java/org/rocksdb/MemoryUtilTest.java b/java/src/test/java/org/rocksdb/MemoryUtilTest.java new file mode 100644 index 000000000..73fcc87c3 --- /dev/null +++ b/java/src/test/java/org/rocksdb/MemoryUtilTest.java @@ -0,0 +1,143 @@ +// 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 org.junit.ClassRule; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; + +import java.nio.charset.StandardCharsets; +import java.util.*; + +import static org.assertj.core.api.Assertions.assertThat; + +public class MemoryUtilTest { + + private static final String MEMTABLE_SIZE = "rocksdb.size-all-mem-tables"; + private static final String UNFLUSHED_MEMTABLE_SIZE = "rocksdb.cur-size-all-mem-tables"; + private static final String TABLE_READERS = "rocksdb.estimate-table-readers-mem"; + + private final byte[] key = "some-key".getBytes(StandardCharsets.UTF_8); + private final byte[] value = "some-value".getBytes(StandardCharsets.UTF_8); + + @ClassRule + public static final RocksMemoryResource rocksMemoryResource = + new RocksMemoryResource(); + + @Rule public TemporaryFolder dbFolder1 = new TemporaryFolder(); + @Rule public TemporaryFolder dbFolder2 = new TemporaryFolder(); + + /** + * Test MemoryUtil.getApproximateMemoryUsageByType before and after a put + get + */ + @Test + public void getApproximateMemoryUsageByType() throws RocksDBException { + try (final Cache cache = new LRUCache(8 * 1024 * 1024); + final Options options = + new Options() + .setCreateIfMissing(true) + .setTableFormatConfig(new BlockBasedTableConfig().setBlockCache(cache)); + final FlushOptions flushOptions = + new FlushOptions().setWaitForFlush(true); + final RocksDB db = + RocksDB.open(options, dbFolder1.getRoot().getAbsolutePath())) { + + List dbs = new ArrayList<>(1); + dbs.add(db); + Set caches = new HashSet<>(1); + caches.add(cache); + Map usage = MemoryUtil.getApproximateMemoryUsageByType(dbs, caches); + + assertThat(usage.get(MemoryUsageType.kMemTableTotal)).isEqualTo( + db.getAggregatedLongProperty(MEMTABLE_SIZE)); + assertThat(usage.get(MemoryUsageType.kMemTableUnFlushed)).isEqualTo( + db.getAggregatedLongProperty(UNFLUSHED_MEMTABLE_SIZE)); + assertThat(usage.get(MemoryUsageType.kTableReadersTotal)).isEqualTo( + db.getAggregatedLongProperty(TABLE_READERS)); + assertThat(usage.get(MemoryUsageType.kCacheTotal)).isEqualTo(0); + + db.put(key, value); + db.flush(flushOptions); + db.get(key); + + usage = MemoryUtil.getApproximateMemoryUsageByType(dbs, caches); + assertThat(usage.get(MemoryUsageType.kMemTableTotal)).isGreaterThan(0); + assertThat(usage.get(MemoryUsageType.kMemTableTotal)).isEqualTo( + db.getAggregatedLongProperty(MEMTABLE_SIZE)); + assertThat(usage.get(MemoryUsageType.kMemTableUnFlushed)).isGreaterThan(0); + assertThat(usage.get(MemoryUsageType.kMemTableUnFlushed)).isEqualTo( + db.getAggregatedLongProperty(UNFLUSHED_MEMTABLE_SIZE)); + assertThat(usage.get(MemoryUsageType.kTableReadersTotal)).isGreaterThan(0); + assertThat(usage.get(MemoryUsageType.kTableReadersTotal)).isEqualTo( + db.getAggregatedLongProperty(TABLE_READERS)); + assertThat(usage.get(MemoryUsageType.kCacheTotal)).isGreaterThan(0); + + } + } + + /** + * Test MemoryUtil.getApproximateMemoryUsageByType with null inputs + */ + @Test + public void getApproximateMemoryUsageByTypeNulls() throws RocksDBException { + Map usage = MemoryUtil.getApproximateMemoryUsageByType(null, null); + + assertThat(usage.get(MemoryUsageType.kMemTableTotal)).isEqualTo(null); + assertThat(usage.get(MemoryUsageType.kMemTableUnFlushed)).isEqualTo(null); + assertThat(usage.get(MemoryUsageType.kTableReadersTotal)).isEqualTo(null); + assertThat(usage.get(MemoryUsageType.kCacheTotal)).isEqualTo(null); + } + + /** + * Test MemoryUtil.getApproximateMemoryUsageByType with two DBs and two caches + */ + @Test + public void getApproximateMemoryUsageByTypeMultiple() throws RocksDBException { + try (final Cache cache1 = new LRUCache(1 * 1024 * 1024); + final Options options1 = + new Options() + .setCreateIfMissing(true) + .setTableFormatConfig(new BlockBasedTableConfig().setBlockCache(cache1)); + final RocksDB db1 = + RocksDB.open(options1, dbFolder1.getRoot().getAbsolutePath()); + final Cache cache2 = new LRUCache(1 * 1024 * 1024); + final Options options2 = + new Options() + .setCreateIfMissing(true) + .setTableFormatConfig(new BlockBasedTableConfig().setBlockCache(cache2)); + final RocksDB db2 = + RocksDB.open(options2, dbFolder2.getRoot().getAbsolutePath()); + final FlushOptions flushOptions = + new FlushOptions().setWaitForFlush(true); + + ) { + List dbs = new ArrayList<>(1); + dbs.add(db1); + dbs.add(db2); + Set caches = new HashSet<>(1); + caches.add(cache1); + caches.add(cache2); + + for (RocksDB db: dbs) { + db.put(key, value); + db.flush(flushOptions); + db.get(key); + } + + Map usage = MemoryUtil.getApproximateMemoryUsageByType(dbs, caches); + assertThat(usage.get(MemoryUsageType.kMemTableTotal)).isEqualTo( + db1.getAggregatedLongProperty(MEMTABLE_SIZE) + db2.getAggregatedLongProperty(MEMTABLE_SIZE)); + assertThat(usage.get(MemoryUsageType.kMemTableUnFlushed)).isEqualTo( + db1.getAggregatedLongProperty(UNFLUSHED_MEMTABLE_SIZE) + db2.getAggregatedLongProperty(UNFLUSHED_MEMTABLE_SIZE)); + assertThat(usage.get(MemoryUsageType.kTableReadersTotal)).isEqualTo( + db1.getAggregatedLongProperty(TABLE_READERS) + db2.getAggregatedLongProperty(TABLE_READERS)); + assertThat(usage.get(MemoryUsageType.kCacheTotal)).isGreaterThan(0); + + } + } + +} diff --git a/src.mk b/src.mk index 9a0ce92ba..54ee1051e 100644 --- a/src.mk +++ b/src.mk @@ -437,6 +437,7 @@ JNI_NATIVE_SOURCES = \ java/rocksjni/loggerjnicallback.cc \ java/rocksjni/lru_cache.cc \ java/rocksjni/memtablejni.cc \ + java/rocksjni/memory_util.cc \ java/rocksjni/merge_operator.cc \ java/rocksjni/native_comparator_wrapper_test.cc \ java/rocksjni/optimistic_transaction_db.cc \