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: d1d196d771dff22afb7ef7500f308233675696f8main
Ben Clay
6 years ago
committed by
Facebook Github Bot
parent
21b51dfec4
commit
c9048021ad
@ -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; |
||||
|
||||
} |
||||
@ -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); |
||||
|
||||
} |
||||
} |
||||
|
||||
} |
||||
Loading…
Reference in new issue