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
main
Ben Clay 6 years ago committed by Facebook Github Bot
parent 21b51dfec4
commit c9048021ad
  1. 4
      java/CMakeLists.txt
  2. 3
      java/Makefile
  3. 100
      java/rocksjni/memory_util.cc
  4. 78
      java/rocksjni/portal.h
  5. 72
      java/src/main/java/org/rocksdb/MemoryUsageType.java
  6. 60
      java/src/main/java/org/rocksdb/MemoryUtil.java
  7. 143
      java/src/test/java/org/rocksdb/MemoryUtilTest.java
  8. 1
      src.mk

@ -25,6 +25,7 @@ set(JNI_NATIVE_SOURCES
rocksjni/jnicallback.cc rocksjni/jnicallback.cc
rocksjni/loggerjnicallback.cc rocksjni/loggerjnicallback.cc
rocksjni/lru_cache.cc rocksjni/lru_cache.cc
rocksjni/memory_util.cc
rocksjni/memtablejni.cc rocksjni/memtablejni.cc
rocksjni/merge_operator.cc rocksjni/merge_operator.cc
rocksjni/native_comparator_wrapper_test.cc rocksjni/native_comparator_wrapper_test.cc
@ -96,6 +97,7 @@ set(NATIVE_JAVA_CLASSES
org.rocksdb.IngestExternalFileOptions org.rocksdb.IngestExternalFileOptions
org.rocksdb.Logger org.rocksdb.Logger
org.rocksdb.LRUCache org.rocksdb.LRUCache
org.rocksdb.MemoryUtil
org.rocksdb.MemTableConfig org.rocksdb.MemTableConfig
org.rocksdb.NativeComparatorWrapper org.rocksdb.NativeComparatorWrapper
org.rocksdb.NativeLibraryLoader org.rocksdb.NativeLibraryLoader
@ -222,6 +224,8 @@ add_jar(
src/main/java/org/rocksdb/IngestExternalFileOptions.java src/main/java/org/rocksdb/IngestExternalFileOptions.java
src/main/java/org/rocksdb/Logger.java src/main/java/org/rocksdb/Logger.java
src/main/java/org/rocksdb/LRUCache.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/MemTableConfig.java
src/main/java/org/rocksdb/MergeOperator.java src/main/java/org/rocksdb/MergeOperator.java
src/main/java/org/rocksdb/MutableColumnFamilyOptionsInterface.java src/main/java/org/rocksdb/MutableColumnFamilyOptionsInterface.java

@ -30,6 +30,8 @@ NATIVE_JAVA_CLASSES = org.rocksdb.AbstractCompactionFilter\
org.rocksdb.HashSkipListMemTableConfig\ org.rocksdb.HashSkipListMemTableConfig\
org.rocksdb.Logger\ org.rocksdb.Logger\
org.rocksdb.LRUCache\ org.rocksdb.LRUCache\
org.rocksdb.MemoryUsageType\
org.rocksdb.MemoryUtil\
org.rocksdb.MergeOperator\ org.rocksdb.MergeOperator\
org.rocksdb.NativeComparatorWrapper\ org.rocksdb.NativeComparatorWrapper\
org.rocksdb.OptimisticTransactionDB\ org.rocksdb.OptimisticTransactionDB\
@ -111,6 +113,7 @@ JAVA_TESTS = org.rocksdb.BackupableDBOptionsTest\
org.rocksdb.KeyMayExistTest\ org.rocksdb.KeyMayExistTest\
org.rocksdb.LoggerTest\ org.rocksdb.LoggerTest\
org.rocksdb.LRUCacheTest\ org.rocksdb.LRUCacheTest\
org.rocksdb.MemoryUtilTest\
org.rocksdb.MemTableTest\ org.rocksdb.MemTableTest\
org.rocksdb.MergeTest\ org.rocksdb.MergeTest\
org.rocksdb.MixedOptionsTest\ org.rocksdb.MixedOptionsTest\

@ -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 <jni.h>
#include <map>
#include <string>
#include <unordered_set>
#include <vector>
#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<rocksdb::DB*> 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<rocksdb::DB *>(ptr_jdb_handles[i]));
}
env->ReleaseLongArrayElements(jdb_handles, ptr_jdb_handles, JNI_ABORT);
}
std::unordered_set<const rocksdb::Cache*> 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<std::shared_ptr<rocksdb::Cache> *>(ptr_jcache_handles[i]);
cache_set.insert(cache_ptr->get());
}
env->ReleaseLongArrayElements(jcache_handles, ptr_jcache_handles, JNI_ABORT);
}
std::map<rocksdb::MemoryUtil::UsageType, uint64_t> 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<uint32_t>(usage_by_type.size()));
if (jusage_by_type == nullptr) {
// exception occurred
return nullptr;
}
const rocksdb::HashMapJni::FnMapKV<const rocksdb::MemoryUtil::UsageType, const uint64_t>
fn_map_kv =
[env](const std::pair<rocksdb::MemoryUtil::UsageType, uint64_t>& 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<std::pair<jobject, jobject>>(nullptr);
}
// Construct value
const jobject jusage_value =
rocksdb::LongJni::valueOf(env, pair.second);
if (jusage_value == nullptr) {
// an error occurred
return std::unique_ptr<std::pair<jobject, jobject>>(nullptr);
}
// Construct and return pointer to pair of jobjects
return std::unique_ptr<std::pair<jobject, jobject>>(
new std::pair<jobject, jobject>(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;
}

@ -26,6 +26,7 @@
#include "rocksdb/rate_limiter.h" #include "rocksdb/rate_limiter.h"
#include "rocksdb/status.h" #include "rocksdb/status.h"
#include "rocksdb/utilities/backupable_db.h" #include "rocksdb/utilities/backupable_db.h"
#include "rocksdb/utilities/memory_util.h"
#include "rocksdb/utilities/transaction_db.h" #include "rocksdb/utilities/transaction_db.h"
#include "rocksdb/utilities/write_batch_with_index.h" #include "rocksdb/utilities/write_batch_with_index.h"
#include "rocksjni/compaction_filter_factory_jnicallback.h" #include "rocksjni/compaction_filter_factory_jnicallback.h"
@ -2251,7 +2252,7 @@ class ByteJni : public JavaClass {
* @param env A pointer to the Java environment * @param env A pointer to the Java environment
* *
* @return The Java Method ID or nullptr if the class or method id could not * @return The Java Method ID or nullptr if the class or method id could not
* be retieved * be retrieved
*/ */
static jmethodID getByteValueMethod(JNIEnv* env) { static jmethodID getByteValueMethod(JNIEnv* env) {
jclass clazz = getJClass(env); jclass clazz = getJClass(env);
@ -2264,6 +2265,39 @@ class ByteJni : public JavaClass {
assert(mid != nullptr); assert(mid != nullptr);
return mid; 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 // 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 // The portal class for org.rocksdb.Transaction
class TransactionJni : public JavaClass { class TransactionJni : public JavaClass {
public: public:

@ -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
*
* <p>The value will be used as a key to indicate the type of memory usage
* described</p>
*/
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_;
}
/**
* <p>Get the MemoryUsageType enumeration value by
* passing the byte identifier to this method.</p>
*
* @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_;
}

@ -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 {
/**
* <p>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.</p>
*
* <p>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.</p>
*
* @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<MemoryUsageType, Long> getApproximateMemoryUsageByType(final List<RocksDB> dbs, final Set<Cache> 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<RocksDB> 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<Byte, Long> byteOutput = getApproximateMemoryUsageByType(dbHandles, cacheHandles);
Map<MemoryUsageType, Long> output = new HashMap<>();
for(Map.Entry<Byte, Long> longEntry : byteOutput.entrySet()) {
output.put(MemoryUsageType.getMemoryUsageType(longEntry.getKey()), longEntry.getValue());
}
return output;
}
private native static Map<Byte, Long> getApproximateMemoryUsageByType(final long[] dbHandles,
final long[] cacheHandles);
}

@ -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<RocksDB> dbs = new ArrayList<>(1);
dbs.add(db);
Set<Cache> caches = new HashSet<>(1);
caches.add(cache);
Map<MemoryUsageType, Long> 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<MemoryUsageType, Long> 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<RocksDB> dbs = new ArrayList<>(1);
dbs.add(db1);
dbs.add(db2);
Set<Cache> 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<MemoryUsageType, Long> 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);
}
}
}

@ -437,6 +437,7 @@ JNI_NATIVE_SOURCES = \
java/rocksjni/loggerjnicallback.cc \ java/rocksjni/loggerjnicallback.cc \
java/rocksjni/lru_cache.cc \ java/rocksjni/lru_cache.cc \
java/rocksjni/memtablejni.cc \ java/rocksjni/memtablejni.cc \
java/rocksjni/memory_util.cc \
java/rocksjni/merge_operator.cc \ java/rocksjni/merge_operator.cc \
java/rocksjni/native_comparator_wrapper_test.cc \ java/rocksjni/native_comparator_wrapper_test.cc \
java/rocksjni/optimistic_transaction_db.cc \ java/rocksjni/optimistic_transaction_db.cc \

Loading…
Cancel
Save