From 85d4362f438617e8ae46eb27e6dcb142cc766f85 Mon Sep 17 00:00:00 2001 From: Jigar Bhati Date: Wed, 17 Oct 2018 11:47:51 -0700 Subject: [PATCH 1/3] Plumb WriteBufferManager through JNI (#4492) Summary: Allow rocks java to explicitly create WriteBufferManager by plumbing it to the native code through JNI. Pull Request resolved: https://github.com/facebook/rocksdb/pull/4492 Differential Revision: D10428506 Pulled By: sagar0 fbshipit-source-id: cd9dd8c2ef745a0303416b44e2080547bdcca1fd --- java/CMakeLists.txt | 3 ++ java/Makefile | 1 + java/rocksjni/options.cc | 14 +++++++ java/rocksjni/write_buffer_manager.cc | 38 +++++++++++++++++++ java/src/main/java/org/rocksdb/DBOptions.java | 9 +++++ .../java/org/rocksdb/DBOptionsInterface.java | 13 +++++++ java/src/main/java/org/rocksdb/Options.java | 9 +++++ .../java/org/rocksdb/WriteBufferManager.java | 30 +++++++++++++++ .../test/java/org/rocksdb/DBOptionsTest.java | 18 +++++++++ .../test/java/org/rocksdb/OptionsTest.java | 18 +++++++++ src.mk | 3 +- 11 files changed, 155 insertions(+), 1 deletion(-) create mode 100644 java/rocksjni/write_buffer_manager.cc create mode 100644 java/src/main/java/org/rocksdb/WriteBufferManager.java diff --git a/java/CMakeLists.txt b/java/CMakeLists.txt index b79eae3b7..9df9ea8d8 100644 --- a/java/CMakeLists.txt +++ b/java/CMakeLists.txt @@ -58,6 +58,7 @@ set(JNI_NATIVE_SOURCES rocksjni/writebatchhandlerjnicallback.cc rocksjni/write_batch_test.cc rocksjni/write_batch_with_index.cc + rocksjni/write_buffer_manager.cc ) set(NATIVE_JAVA_CLASSES @@ -145,6 +146,7 @@ set(NATIVE_JAVA_CLASSES org.rocksdb.SnapshotTest org.rocksdb.WriteBatchTest org.rocksdb.WriteBatchTestInternalHelper + org.rocksdb.WriteBufferManager ) include(FindJava) @@ -283,6 +285,7 @@ add_jar( src/main/java/org/rocksdb/WriteBatch.java src/main/java/org/rocksdb/WriteBatchWithIndex.java src/main/java/org/rocksdb/WriteOptions.java + src/main/java/org/rocksdb/WriteBufferManager.java src/main/java/org/rocksdb/util/BytewiseComparator.java src/main/java/org/rocksdb/util/DirectBytewiseComparator.java src/main/java/org/rocksdb/util/Environment.java diff --git a/java/Makefile b/java/Makefile index 5f78094ca..b8eb39b53 100644 --- a/java/Makefile +++ b/java/Makefile @@ -66,6 +66,7 @@ NATIVE_JAVA_CLASSES = org.rocksdb.AbstractCompactionFilter\ org.rocksdb.WriteBatch.Handler\ org.rocksdb.WriteOptions\ org.rocksdb.WriteBatchWithIndex\ + org.rocksdb.WriteBufferManager\ org.rocksdb.WBWIRocksIterator NATIVE_JAVA_TEST_CLASSES = org.rocksdb.RocksDBExceptionTest\ diff --git a/java/rocksjni/options.cc b/java/rocksjni/options.cc index 9aed80e1e..ab213862b 100644 --- a/java/rocksjni/options.cc +++ b/java/rocksjni/options.cc @@ -250,6 +250,20 @@ void Java_org_rocksdb_Options_setWriteBufferSize(JNIEnv* env, jobject /*jobj*/, } } +/* + * Class: org_rocksdb_Options + * Method: setWriteBufferManager + * Signature: (JJ)V + */ +void Java_org_rocksdb_Options_setWriteBufferManager(JNIEnv* /*env*/, jobject /*jobj*/, + jlong joptions_handle, + jlong jwrite_buffer_manager_handle) { + auto* write_buffer_manager = + reinterpret_cast *>(jwrite_buffer_manager_handle); + reinterpret_cast(joptions_handle)->write_buffer_manager = + *write_buffer_manager; +} + /* * Class: org_rocksdb_Options * Method: writeBufferSize diff --git a/java/rocksjni/write_buffer_manager.cc b/java/rocksjni/write_buffer_manager.cc new file mode 100644 index 000000000..043f69031 --- /dev/null +++ b/java/rocksjni/write_buffer_manager.cc @@ -0,0 +1,38 @@ +// 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/org_rocksdb_WriteBufferManager.h" + +#include "rocksdb/cache.h" +#include "rocksdb/write_buffer_manager.h" + +/* + * Class: org_rocksdb_WriteBufferManager + * Method: newWriteBufferManager + * Signature: (JJ)J + */ +jlong Java_org_rocksdb_WriteBufferManager_newWriteBufferManager( + JNIEnv* /*env*/, jclass /*jclazz*/, jlong jbuffer_size, jlong jcache_handle) { + auto* cache_ptr = + reinterpret_cast *>(jcache_handle); + auto* write_buffer_manager = new std::shared_ptr( + std::make_shared(jbuffer_size, *cache_ptr)); + return reinterpret_cast(write_buffer_manager); +} + +/* + * Class: org_rocksdb_WriteBufferManager + * Method: disposeInternal + * Signature: (J)V + */ +void Java_org_rocksdb_WriteBufferManager_disposeInternal( + JNIEnv* /*env*/, jobject /*jobj*/, jlong jhandle) { + auto* write_buffer_manager = + reinterpret_cast *>(jhandle); + assert(write_buffer_manager != nullptr); + delete write_buffer_manager; +} diff --git a/java/src/main/java/org/rocksdb/DBOptions.java b/java/src/main/java/org/rocksdb/DBOptions.java index c32329388..9623a1042 100644 --- a/java/src/main/java/org/rocksdb/DBOptions.java +++ b/java/src/main/java/org/rocksdb/DBOptions.java @@ -667,6 +667,13 @@ public DBOptions setDbWriteBufferSize(final long dbWriteBufferSize) { return this; } + @Override + public DBOptions setWriteBufferManager(final WriteBufferManager writeBufferManager) { + assert(isOwningHandle()); + setWriteBufferManager(nativeHandle_, writeBufferManager.nativeHandle_); + return this; + } + @Override public long dbWriteBufferSize() { assert(isOwningHandle()); @@ -1087,6 +1094,8 @@ private native void setAdviseRandomOnOpen( private native boolean adviseRandomOnOpen(long handle); private native void setDbWriteBufferSize(final long handle, final long dbWriteBufferSize); + private native void setWriteBufferManager(final long dbOptionsHandle, + final long writeBufferManagerHandle); private native long dbWriteBufferSize(final long handle); private native void setAccessHintOnCompactionStart(final long handle, final byte accessHintOnCompactionStart); diff --git a/java/src/main/java/org/rocksdb/DBOptionsInterface.java b/java/src/main/java/org/rocksdb/DBOptionsInterface.java index 7c406eaf8..d3a33b757 100644 --- a/java/src/main/java/org/rocksdb/DBOptionsInterface.java +++ b/java/src/main/java/org/rocksdb/DBOptionsInterface.java @@ -991,6 +991,19 @@ public interface DBOptionsInterface { */ T setDbWriteBufferSize(long dbWriteBufferSize); + /** + * Use passed {@link WriteBufferManager} to control memory usage across + * multiple column families and/or DB instances. + * + * Check + * https://github.com/facebook/rocksdb/wiki/Write-Buffer-Manager + * for more details on when to use it + * + * @param writeBufferManager The WriteBufferManager to use + * @return the reference of the current options. + */ + T setWriteBufferManager(final WriteBufferManager writeBufferManager); + /** * Amount of data to build up in memtables across all column * families before writing to disk. diff --git a/java/src/main/java/org/rocksdb/Options.java b/java/src/main/java/org/rocksdb/Options.java index cac4fc5a3..e409d3466 100644 --- a/java/src/main/java/org/rocksdb/Options.java +++ b/java/src/main/java/org/rocksdb/Options.java @@ -723,6 +723,13 @@ public Options setDbWriteBufferSize(final long dbWriteBufferSize) { return this; } + @Override + public Options setWriteBufferManager(final WriteBufferManager writeBufferManager) { + assert(isOwningHandle()); + setWriteBufferManager(nativeHandle_, writeBufferManager.nativeHandle_); + return this; + } + @Override public long dbWriteBufferSize() { assert(isOwningHandle()); @@ -1690,6 +1697,8 @@ private native void setAdviseRandomOnOpen( private native boolean adviseRandomOnOpen(long handle); private native void setDbWriteBufferSize(final long handle, final long dbWriteBufferSize); + private native void setWriteBufferManager(final long handle, + final long writeBufferManagerHandle); private native long dbWriteBufferSize(final long handle); private native void setAccessHintOnCompactionStart(final long handle, final byte accessHintOnCompactionStart); diff --git a/java/src/main/java/org/rocksdb/WriteBufferManager.java b/java/src/main/java/org/rocksdb/WriteBufferManager.java new file mode 100644 index 000000000..a5f80644f --- /dev/null +++ b/java/src/main/java/org/rocksdb/WriteBufferManager.java @@ -0,0 +1,30 @@ +package org.rocksdb; + +import org.rocksdb.Cache; + +/** + * Java wrapper over native write_buffer_manager class + */ +public class WriteBufferManager extends RocksObject { + static { + RocksDB.loadLibrary(); + } + + /** + * Construct a new instance of WriteBufferManager. + * + * Check + * https://github.com/facebook/rocksdb/wiki/Write-Buffer-Manager + * for more details on when to use it + * + * @param bufferSizeBytes buffer size(in bytes) to use for native write_buffer_manager + * @param cache cache whose memory should be bounded by this write buffer manager + */ + public WriteBufferManager(final long bufferSizeBytes, final Cache cache){ + super(newWriteBufferManager(bufferSizeBytes, cache.nativeHandle_)); + } + + private native static long newWriteBufferManager(final long bufferSizeBytes, final long cacheHandle); + @Override + protected native void disposeInternal(final long handle); +} diff --git a/java/src/test/java/org/rocksdb/DBOptionsTest.java b/java/src/test/java/org/rocksdb/DBOptionsTest.java index 453639d57..79cee3950 100644 --- a/java/src/test/java/org/rocksdb/DBOptionsTest.java +++ b/java/src/test/java/org/rocksdb/DBOptionsTest.java @@ -424,6 +424,24 @@ public void dbWriteBufferSize() { } } + @Test + public void setWriteBufferManager() throws RocksDBException { + try (final Options opt = new Options(); + final Cache cache = new LRUCache(1 * 1024 * 1024); + final WriteBufferManager writeBufferManager = new WriteBufferManager(2000l, cache)) { + opt.setWriteBufferManager(writeBufferManager); + } + } + + @Test + public void setWriteBufferManagerWithZeroBufferSize() throws RocksDBException { + try (final Options opt = new Options(); + final Cache cache = new LRUCache(1 * 1024 * 1024); + final WriteBufferManager writeBufferManager = new WriteBufferManager(0l, cache)) { + opt.setWriteBufferManager(writeBufferManager); + } + } + @Test public void accessHintOnCompactionStart() { try(final DBOptions opt = new DBOptions()) { diff --git a/java/src/test/java/org/rocksdb/OptionsTest.java b/java/src/test/java/org/rocksdb/OptionsTest.java index 7f7679d73..44d70c16c 100644 --- a/java/src/test/java/org/rocksdb/OptionsTest.java +++ b/java/src/test/java/org/rocksdb/OptionsTest.java @@ -645,6 +645,24 @@ public void dbWriteBufferSize() { } } + @Test + public void setWriteBufferManager() throws RocksDBException { + try (final Options opt = new Options(); + final Cache cache = new LRUCache(1 * 1024 * 1024); + final WriteBufferManager writeBufferManager = new WriteBufferManager(2000l, cache)) { + opt.setWriteBufferManager(writeBufferManager); + } + } + + @Test + public void setWriteBufferManagerWithZeroBufferSize() throws RocksDBException { + try (final Options opt = new Options(); + final Cache cache = new LRUCache(1 * 1024 * 1024); + final WriteBufferManager writeBufferManager = new WriteBufferManager(0l, cache)) { + opt.setWriteBufferManager(writeBufferManager); + } + } + @Test public void accessHintOnCompactionStart() { try (final Options opt = new Options()) { diff --git a/src.mk b/src.mk index 0f4284aa4..cf0df2ac0 100644 --- a/src.mk +++ b/src.mk @@ -472,4 +472,5 @@ JNI_NATIVE_SOURCES = \ java/rocksjni/write_batch.cc \ java/rocksjni/writebatchhandlerjnicallback.cc \ java/rocksjni/write_batch_test.cc \ - java/rocksjni/write_batch_with_index.cc + java/rocksjni/write_batch_with_index.cc \ + java/rocksjni/write_buffer_manager.cc From ebab391d9f32cbffd03c1a2f7d627d135edae173 Mon Sep 17 00:00:00 2001 From: Jigar Bhati Date: Wed, 24 Oct 2018 12:39:07 -0700 Subject: [PATCH 2/3] WriteBufferManager JNI fixes (#4579) Summary: 1. `WriteBufferManager` should have a reference alive in Java side through `Options`/`DBOptions` otherwise, if it's GC'ed at java side, native side can seg fault. 2. native method `setWriteBufferManager()` in `DBOptions.java` doesn't have it's jni method invocation in rocksdbjni which is added in this PR 3. `DBOptionsTest.java` is referencing object of `Options`. Instead it should be testing against `DBOptions`. Seems like a copy paste error. 4. Add a getter for WriteBufferManager. Pull Request resolved: https://github.com/facebook/rocksdb/pull/4579 Differential Revision: D10561150 Pulled By: sagar0 fbshipit-source-id: 139a15c7f051a9f77b4200215b88267b48fbc487 --- java/rocksjni/options.cc | 14 ++++++++++++++ java/src/main/java/org/rocksdb/DBOptions.java | 9 +++++++++ .../main/java/org/rocksdb/DBOptionsInterface.java | 9 +++++++++ java/src/main/java/org/rocksdb/Options.java | 9 +++++++++ java/src/test/java/org/rocksdb/DBOptionsTest.java | 6 ++++-- java/src/test/java/org/rocksdb/OptionsTest.java | 2 ++ 6 files changed, 47 insertions(+), 2 deletions(-) diff --git a/java/rocksjni/options.cc b/java/rocksjni/options.cc index ab213862b..b4c18285e 100644 --- a/java/rocksjni/options.cc +++ b/java/rocksjni/options.cc @@ -5532,6 +5532,20 @@ void Java_org_rocksdb_DBOptions_setDbWriteBufferSize( opt->db_write_buffer_size = static_cast(jdb_write_buffer_size); } +/* + * Class: org_rocksdb_DBOptions + * Method: setWriteBufferManager + * Signature: (JJ)V + */ +void Java_org_rocksdb_DBOptions_setWriteBufferManager(JNIEnv* /*env*/, jobject /*jobj*/, + jlong jdb_options_handle, + jlong jwrite_buffer_manager_handle) { + auto* write_buffer_manager = + reinterpret_cast *>(jwrite_buffer_manager_handle); + reinterpret_cast(jdb_options_handle)->write_buffer_manager = + *write_buffer_manager; +} + /* * Class: org_rocksdb_DBOptions * Method: dbWriteBufferSize diff --git a/java/src/main/java/org/rocksdb/DBOptions.java b/java/src/main/java/org/rocksdb/DBOptions.java index 9623a1042..280623a20 100644 --- a/java/src/main/java/org/rocksdb/DBOptions.java +++ b/java/src/main/java/org/rocksdb/DBOptions.java @@ -46,6 +46,7 @@ public DBOptions(DBOptions other) { this.numShardBits_ = other.numShardBits_; this.rateLimiter_ = other.rateLimiter_; this.rowCache_ = other.rowCache_; + this.writeBufferManager_ = other.writeBufferManager_; } /** @@ -671,10 +672,17 @@ public DBOptions setDbWriteBufferSize(final long dbWriteBufferSize) { public DBOptions setWriteBufferManager(final WriteBufferManager writeBufferManager) { assert(isOwningHandle()); setWriteBufferManager(nativeHandle_, writeBufferManager.nativeHandle_); + this.writeBufferManager_ = writeBufferManager; return this; } @Override + public WriteBufferManager writeBufferManager() { + assert(isOwningHandle()); + return this.writeBufferManager_; + } + + @Override public long dbWriteBufferSize() { assert(isOwningHandle()); return dbWriteBufferSize(nativeHandle_); @@ -1167,4 +1175,5 @@ private native void setAvoidFlushDuringShutdown(final long handle, private int numShardBits_; private RateLimiter rateLimiter_; private Cache rowCache_; + private WriteBufferManager writeBufferManager_; } diff --git a/java/src/main/java/org/rocksdb/DBOptionsInterface.java b/java/src/main/java/org/rocksdb/DBOptionsInterface.java index d3a33b757..accfb4c59 100644 --- a/java/src/main/java/org/rocksdb/DBOptionsInterface.java +++ b/java/src/main/java/org/rocksdb/DBOptionsInterface.java @@ -1004,6 +1004,15 @@ public interface DBOptionsInterface { */ T setWriteBufferManager(final WriteBufferManager writeBufferManager); + /** + * Reference to {@link WriteBufferManager} used by it.
+ * + * Default: null (Disabled) + * + * @return a reference to WriteBufferManager + */ + WriteBufferManager writeBufferManager(); + /** * Amount of data to build up in memtables across all column * families before writing to disk. diff --git a/java/src/main/java/org/rocksdb/Options.java b/java/src/main/java/org/rocksdb/Options.java index e409d3466..2ff4ec120 100644 --- a/java/src/main/java/org/rocksdb/Options.java +++ b/java/src/main/java/org/rocksdb/Options.java @@ -70,6 +70,7 @@ public Options(Options other) { this.compactionOptionsFIFO_ = other.compactionOptionsFIFO_; this.compressionOptions_ = other.compressionOptions_; this.rowCache_ = other.rowCache_; + this.writeBufferManager_ = other.writeBufferManager_; } @Override @@ -727,10 +728,17 @@ public Options setDbWriteBufferSize(final long dbWriteBufferSize) { public Options setWriteBufferManager(final WriteBufferManager writeBufferManager) { assert(isOwningHandle()); setWriteBufferManager(nativeHandle_, writeBufferManager.nativeHandle_); + this.writeBufferManager_ = writeBufferManager; return this; } @Override + public WriteBufferManager writeBufferManager() { + assert(isOwningHandle()); + return this.writeBufferManager_; + } + + @Override public long dbWriteBufferSize() { assert(isOwningHandle()); return dbWriteBufferSize(nativeHandle_); @@ -1918,4 +1926,5 @@ private native void setForceConsistencyChecks(final long handle, private CompactionOptionsFIFO compactionOptionsFIFO_; private CompressionOptions compressionOptions_; private Cache rowCache_; + private WriteBufferManager writeBufferManager_; } diff --git a/java/src/test/java/org/rocksdb/DBOptionsTest.java b/java/src/test/java/org/rocksdb/DBOptionsTest.java index 79cee3950..bad01c435 100644 --- a/java/src/test/java/org/rocksdb/DBOptionsTest.java +++ b/java/src/test/java/org/rocksdb/DBOptionsTest.java @@ -426,19 +426,21 @@ public void dbWriteBufferSize() { @Test public void setWriteBufferManager() throws RocksDBException { - try (final Options opt = new Options(); + try (final DBOptions opt = new DBOptions(); final Cache cache = new LRUCache(1 * 1024 * 1024); final WriteBufferManager writeBufferManager = new WriteBufferManager(2000l, cache)) { opt.setWriteBufferManager(writeBufferManager); + assertThat(opt.writeBufferManager()).isEqualTo(writeBufferManager); } } @Test public void setWriteBufferManagerWithZeroBufferSize() throws RocksDBException { - try (final Options opt = new Options(); + try (final DBOptions opt = new DBOptions(); final Cache cache = new LRUCache(1 * 1024 * 1024); final WriteBufferManager writeBufferManager = new WriteBufferManager(0l, cache)) { opt.setWriteBufferManager(writeBufferManager); + assertThat(opt.writeBufferManager()).isEqualTo(writeBufferManager); } } diff --git a/java/src/test/java/org/rocksdb/OptionsTest.java b/java/src/test/java/org/rocksdb/OptionsTest.java index 44d70c16c..2571c3e26 100644 --- a/java/src/test/java/org/rocksdb/OptionsTest.java +++ b/java/src/test/java/org/rocksdb/OptionsTest.java @@ -651,6 +651,7 @@ public void setWriteBufferManager() throws RocksDBException { final Cache cache = new LRUCache(1 * 1024 * 1024); final WriteBufferManager writeBufferManager = new WriteBufferManager(2000l, cache)) { opt.setWriteBufferManager(writeBufferManager); + assertThat(opt.writeBufferManager()).isEqualTo(writeBufferManager); } } @@ -660,6 +661,7 @@ public void setWriteBufferManagerWithZeroBufferSize() throws RocksDBException { final Cache cache = new LRUCache(1 * 1024 * 1024); final WriteBufferManager writeBufferManager = new WriteBufferManager(0l, cache)) { opt.setWriteBufferManager(writeBufferManager); + assertThat(opt.writeBufferManager()).isEqualTo(writeBufferManager); } } From b5d0382429be612c8c60868b1cf81f66678d92e3 Mon Sep 17 00:00:00 2001 From: Siying Dong Date: Sun, 18 Nov 2018 16:51:15 -0800 Subject: [PATCH 3/3] WriteBufferManger doens't cost to cache if no limit is set (#4695) Summary: WriteBufferManger is not invoked when allocating memory for memtable if the limit is not set even if a cache is passed. It is inconsistent from the comment syas. Fix it. Pull Request resolved: https://github.com/facebook/rocksdb/pull/4695 Differential Revision: D13112722 Pulled By: siying fbshipit-source-id: 0b27eef63867f679cd06033ea56907c0569597f4 --- db/db_test2.cc | 16 ++++++++++++++++ db/memtable.cc | 13 +++++++------ include/rocksdb/write_buffer_manager.h | 2 ++ memtable/alloc_tracker.cc | 9 ++++++--- 4 files changed, 31 insertions(+), 9 deletions(-) diff --git a/db/db_test2.cc b/db/db_test2.cc index cc30ef85b..287654742 100644 --- a/db/db_test2.cc +++ b/db/db_test2.cc @@ -454,6 +454,22 @@ TEST_F(DBTest2, SharedWriteBufferLimitAcrossDB) { rocksdb::SyncPoint::GetInstance()->DisableProcessing(); } +TEST_F(DBTest2, TestWriteBufferNoLimitWithCache) { + Options options = CurrentOptions(); + options.arena_block_size = 4096; + std::shared_ptr cache = + NewLRUCache(LRUCacheOptions(10000000, 1, false, 0.0)); + options.write_buffer_size = 50000; // this is never hit + // Use a write buffer total size so that the soft limit is about + // 105000. + options.write_buffer_manager.reset(new WriteBufferManager(0, cache)); + Reopen(options); + + ASSERT_OK(Put("foo", "bar")); + // One dummy entry is 1MB. + ASSERT_GT(cache->GetUsage(), 500000); +} + namespace { void ValidateKeyExistence(DB* db, const std::vector& keys_must_exist, const std::vector& keys_must_not_exist) { diff --git a/db/memtable.cc b/db/memtable.cc index 70e6d9da5..0b30e33c2 100644 --- a/db/memtable.cc +++ b/db/memtable.cc @@ -68,12 +68,13 @@ MemTable::MemTable(const InternalKeyComparator& cmp, refs_(0), kArenaBlockSize(OptimizeBlockSize(moptions_.arena_block_size)), mem_tracker_(write_buffer_manager), - arena_( - moptions_.arena_block_size, - (write_buffer_manager != nullptr && write_buffer_manager->enabled()) - ? &mem_tracker_ - : nullptr, - mutable_cf_options.memtable_huge_page_size), + arena_(moptions_.arena_block_size, + (write_buffer_manager != nullptr && + (write_buffer_manager->enabled() || + write_buffer_manager->cost_to_cache())) + ? &mem_tracker_ + : nullptr, + mutable_cf_options.memtable_huge_page_size), table_(ioptions.memtable_factory->CreateMemTableRep( comparator_, &arena_, mutable_cf_options.prefix_extractor.get(), ioptions.info_log, column_family_id)), diff --git a/include/rocksdb/write_buffer_manager.h b/include/rocksdb/write_buffer_manager.h index 856cf4b24..dea904c18 100644 --- a/include/rocksdb/write_buffer_manager.h +++ b/include/rocksdb/write_buffer_manager.h @@ -30,6 +30,8 @@ class WriteBufferManager { bool enabled() const { return buffer_size_ != 0; } + bool cost_to_cache() const { return cache_rep_ != nullptr; } + // Only valid if enabled() size_t memory_usage() const { return memory_used_.load(std::memory_order_relaxed); diff --git a/memtable/alloc_tracker.cc b/memtable/alloc_tracker.cc index 9889cc423..a1fa4938c 100644 --- a/memtable/alloc_tracker.cc +++ b/memtable/alloc_tracker.cc @@ -24,7 +24,8 @@ AllocTracker::~AllocTracker() { FreeMem(); } void AllocTracker::Allocate(size_t bytes) { assert(write_buffer_manager_ != nullptr); - if (write_buffer_manager_->enabled()) { + if (write_buffer_manager_->enabled() || + write_buffer_manager_->cost_to_cache()) { bytes_allocated_.fetch_add(bytes, std::memory_order_relaxed); write_buffer_manager_->ReserveMem(bytes); } @@ -32,7 +33,8 @@ void AllocTracker::Allocate(size_t bytes) { void AllocTracker::DoneAllocating() { if (write_buffer_manager_ != nullptr && !done_allocating_) { - if (write_buffer_manager_->enabled()) { + if (write_buffer_manager_->enabled() || + write_buffer_manager_->cost_to_cache()) { write_buffer_manager_->ScheduleFreeMem( bytes_allocated_.load(std::memory_order_relaxed)); } else { @@ -47,7 +49,8 @@ void AllocTracker::FreeMem() { DoneAllocating(); } if (write_buffer_manager_ != nullptr && !freed_) { - if (write_buffer_manager_->enabled()) { + if (write_buffer_manager_->enabled() || + write_buffer_manager_->cost_to_cache()) { write_buffer_manager_->FreeMem( bytes_allocated_.load(std::memory_order_relaxed)); } else {