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

@@ -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

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\

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

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:

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
+ *
+ * <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_;
+}

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 {
+
+  /**
+   * <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);
+}

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

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                  \