Add new API CacheReservationManager::GetTotalMemoryUsage() (#9071)

Summary:
Note: This PR is the 2nd PR of a bigger PR stack (https://github.com/facebook/rocksdb/pull/9073).

Context:
`CacheReservationManager::UpdateCacheReservation(std::size_t new_memory_used)` accepts an accumulated total memory used (e.g, used 10MB so far) instead of usage change (e.g, increase by 5 MB, decrease by 5 MB). It has benefits including consolidating API for increase and decrease as described in https://github.com/facebook/rocksdb/pull/8506.

However, not every `CacheReservationManager` user keeps track of this accumulated total memory usage. For example, Bloom/Ribbon Filter construction (e.g, [here](822d729fcd/table/block_based/filter_policy.cc (L587)) in https://github.com/facebook/rocksdb/pull/9073) does not  while WriteBufferManager and compression dictionary buffering do.

Considering future users might or might not keep track of this counter and implementing this counter within `CacheReservationManager` is easy due to the passed-in `std::size_t new_memory_used` in calling `CacheReservationManager::UpdateCacheReservation(std::size_t new_memory_used)`, it is proposed to add a new API `CacheReservationManager::GetTotalMemoryUsage()`.

As noted in the API comments,   since `CacheReservationManager` is NOT thread-safe, external synchronization is
 needed in calling `UpdateCacheReservation()` if you want `GetTotalMemoryUsed()` returns the indeed latest memory used.
- Added and updated private counter `memory_used_` every time `CacheReservationManager::UpdateCacheReservation(std::size_t new_memory_used)` is called regardless if the call returns non-okay status
- Added `CacheReservationManager::GetTotalMemoryUsage()` to return `memory_used_`

Pull Request resolved: https://github.com/facebook/rocksdb/pull/9071

Test Plan:
- Passing new tests
- Passing existing tests

Reviewed By: ajkr

Differential Revision: D31887813

Pulled By: hx235

fbshipit-source-id: 9a09f0c8683822673260362894c878b61ee60ceb
main
Hui Xiao 3 years ago committed by Facebook GitHub Bot
parent efaef9b40a
commit ffd6085e1f
  1. 9
      cache/cache_reservation_manager.cc
  2. 36
      cache/cache_reservation_manager.h
  3. 29
      cache/cache_reservation_manager_test.cc

@ -23,7 +23,9 @@
namespace ROCKSDB_NAMESPACE { namespace ROCKSDB_NAMESPACE {
CacheReservationManager::CacheReservationManager(std::shared_ptr<Cache> cache, CacheReservationManager::CacheReservationManager(std::shared_ptr<Cache> cache,
bool delayed_decrease) bool delayed_decrease)
: delayed_decrease_(delayed_decrease), cache_allocated_size_(0) { : delayed_decrease_(delayed_decrease),
cache_allocated_size_(0),
memory_used_(0) {
assert(cache != nullptr); assert(cache != nullptr);
cache_ = cache; cache_ = cache;
std::memset(cache_key_, 0, kCacheKeyPrefixSize + kMaxVarint64Length); std::memset(cache_key_, 0, kCacheKeyPrefixSize + kMaxVarint64Length);
@ -39,6 +41,7 @@ CacheReservationManager::~CacheReservationManager() {
template <CacheEntryRole R> template <CacheEntryRole R>
Status CacheReservationManager::UpdateCacheReservation( Status CacheReservationManager::UpdateCacheReservation(
std::size_t new_mem_used) { std::size_t new_mem_used) {
memory_used_ = new_mem_used;
std::size_t cur_cache_allocated_size = std::size_t cur_cache_allocated_size =
cache_allocated_size_.load(std::memory_order_relaxed); cache_allocated_size_.load(std::memory_order_relaxed);
if (new_mem_used == cur_cache_allocated_size) { if (new_mem_used == cur_cache_allocated_size) {
@ -118,6 +121,10 @@ std::size_t CacheReservationManager::GetTotalReservedCacheSize() {
return cache_allocated_size_.load(std::memory_order_relaxed); return cache_allocated_size_.load(std::memory_order_relaxed);
} }
std::size_t CacheReservationManager::GetTotalMemoryUsed() {
return memory_used_;
}
Slice CacheReservationManager::GetNextCacheKey() { Slice CacheReservationManager::GetNextCacheKey() {
// Calling this function will have the side-effect of changing the // Calling this function will have the side-effect of changing the
// underlying cache_key_ that is shared among other keys generated from this // underlying cache_key_ that is shared among other keys generated from this

@ -26,7 +26,9 @@ namespace ROCKSDB_NAMESPACE {
// CacheReservationManager is for reserving cache space for the memory used // CacheReservationManager is for reserving cache space for the memory used
// through inserting/releasing dummy entries in the cache. // through inserting/releasing dummy entries in the cache.
// This class is not thread-safe. //
// This class is NOT thread-safe, except that GetTotalReservedCacheSize()
// can be called without external synchronization.
class CacheReservationManager { class CacheReservationManager {
public: public:
// Construct a CacheReservationManager // Construct a CacheReservationManager
@ -53,8 +55,8 @@ class CacheReservationManager {
template <CacheEntryRole R> template <CacheEntryRole R>
// Insert and release dummy entries in the cache to // Insert and release dummy entries in the cache to
// match the size of total dummy entries with the smallest multiple of // match the size of total dummy entries with the least multiple of
// kSizeDummyEntry that is greater than or equal to new_mem_used // kSizeDummyEntry greater than or equal to new_mem_used
// //
// Insert dummy entries if new_memory_used > cache_allocated_size_; // Insert dummy entries if new_memory_used > cache_allocated_size_;
// //
@ -68,13 +70,34 @@ class CacheReservationManager {
// is set true. // is set true.
// //
// @param new_memory_used The number of bytes used by new memory // @param new_memory_used The number of bytes used by new memory
// The most recent new_memoy_used passed in will be returned
// in GetTotalMemoryUsed() even when the call return non-ok status.
//
// Since the class is NOT thread-safe, external synchronization on the
// order of calling UpdateCacheReservation() is needed if you want
// GetTotalMemoryUsed() indeed returns the latest memory used.
//
// @return On inserting dummy entries, it returns Status::OK() if all dummy // @return On inserting dummy entries, it returns Status::OK() if all dummy
// entry insertions succeed. Otherwise, it returns the first non-ok status; // entry insertions succeed.
// On releasing dummy entries, it always returns Status::OK(). // Otherwise, it returns the first non-ok status;
// On keeping dummy entries the same, it always returns Status::OK(). // On releasing dummy entries, it always returns Status::OK().
// On keeping dummy entries the same, it always returns Status::OK().
Status UpdateCacheReservation(std::size_t new_memory_used); Status UpdateCacheReservation(std::size_t new_memory_used);
// Return the size of the cache (which is a multiple of kSizeDummyEntry)
// successfully reserved by calling UpdateCacheReservation().
//
// When UpdateCacheReservation() returns non-ok status,
// calling GetTotalReservedCacheSize() after that might return a slightly
// smaller number than the actual reserved cache size due to
// the returned number will always be a multiple of kSizeDummyEntry
// and cache full might happen in the middle of inserting a dummy entry.
std::size_t GetTotalReservedCacheSize(); std::size_t GetTotalReservedCacheSize();
// Return the latest total memory used indicated by the most recent call of
// UpdateCacheReservation(std::size_t new_memory_used);
std::size_t GetTotalMemoryUsed();
static constexpr std::size_t GetDummyEntrySize() { return kSizeDummyEntry; } static constexpr std::size_t GetDummyEntrySize() { return kSizeDummyEntry; }
private: private:
@ -92,6 +115,7 @@ class CacheReservationManager {
std::shared_ptr<Cache> cache_; std::shared_ptr<Cache> cache_;
bool delayed_decrease_; bool delayed_decrease_;
std::atomic<std::size_t> cache_allocated_size_; std::atomic<std::size_t> cache_allocated_size_;
std::size_t memory_used_;
std::vector<Cache::Handle *> dummy_handles_; std::vector<Cache::Handle *> dummy_handles_;
std::uint64_t next_cache_key_id_ = 0; std::uint64_t next_cache_key_id_ = 0;
// The non-prefix part will be updated according to the ID to use. // The non-prefix part will be updated according to the ID to use.

@ -81,6 +81,7 @@ TEST_F(CacheReservationManagerTest, KeepCacheReservationTheSame) {
ASSERT_EQ(s, Status::OK()); ASSERT_EQ(s, Status::OK());
ASSERT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), ASSERT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(),
1 * kSizeDummyEntry); 1 * kSizeDummyEntry);
ASSERT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used);
std::size_t initial_pinned_usage = cache->GetPinnedUsage(); std::size_t initial_pinned_usage = cache->GetPinnedUsage();
ASSERT_GE(initial_pinned_usage, 1 * kSizeDummyEntry); ASSERT_GE(initial_pinned_usage, 1 * kSizeDummyEntry);
ASSERT_LT(initial_pinned_usage, ASSERT_LT(initial_pinned_usage,
@ -96,6 +97,9 @@ TEST_F(CacheReservationManagerTest, KeepCacheReservationTheSame) {
1 * kSizeDummyEntry) 1 * kSizeDummyEntry)
<< "Failed to bookkeep correctly when new_mem_used equals to current " << "Failed to bookkeep correctly when new_mem_used equals to current "
"cache reservation"; "cache reservation";
EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
<< "Failed to bookkeep the used memory correctly when new_mem_used "
"equals to current cache reservation";
EXPECT_EQ(cache->GetPinnedUsage(), initial_pinned_usage) EXPECT_EQ(cache->GetPinnedUsage(), initial_pinned_usage)
<< "Failed to keep underlying dummy entries the same when new_mem_used " << "Failed to keep underlying dummy entries the same when new_mem_used "
"equals to current cache reservation"; "equals to current cache reservation";
@ -113,6 +117,8 @@ TEST_F(CacheReservationManagerTest,
EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(),
2 * kSizeDummyEntry) 2 * kSizeDummyEntry)
<< "Failed to bookkeep cache reservation increase correctly"; << "Failed to bookkeep cache reservation increase correctly";
EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
<< "Failed to bookkeep the used memory correctly";
EXPECT_GE(cache->GetPinnedUsage(), 2 * kSizeDummyEntry) EXPECT_GE(cache->GetPinnedUsage(), 2 * kSizeDummyEntry)
<< "Failed to increase underlying dummy entries in cache correctly"; << "Failed to increase underlying dummy entries in cache correctly";
EXPECT_LT(cache->GetPinnedUsage(), EXPECT_LT(cache->GetPinnedUsage(),
@ -132,6 +138,8 @@ TEST_F(CacheReservationManagerTest,
EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(),
3 * kSizeDummyEntry) 3 * kSizeDummyEntry)
<< "Failed to bookkeep cache reservation increase correctly"; << "Failed to bookkeep cache reservation increase correctly";
EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
<< "Failed to bookkeep the used memory correctly";
EXPECT_GE(cache->GetPinnedUsage(), 3 * kSizeDummyEntry) EXPECT_GE(cache->GetPinnedUsage(), 3 * kSizeDummyEntry)
<< "Failed to increase underlying dummy entries in cache correctly"; << "Failed to increase underlying dummy entries in cache correctly";
EXPECT_LT(cache->GetPinnedUsage(), EXPECT_LT(cache->GetPinnedUsage(),
@ -173,6 +181,8 @@ TEST(CacheReservationManagerIncreaseReservcationOnFullCacheTest,
<< "Failed to bookkeep correctly (i.e, bookkeep only successful dummy " << "Failed to bookkeep correctly (i.e, bookkeep only successful dummy "
"entry insertions) when encountering cache resevation failure due to " "entry insertions) when encountering cache resevation failure due to "
"full cache"; "full cache";
EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
<< "Failed to bookkeep the used memory correctly";
EXPECT_GE(cache->GetPinnedUsage(), 1 * kSizeDummyEntry) EXPECT_GE(cache->GetPinnedUsage(), 1 * kSizeDummyEntry)
<< "Failed to insert underlying dummy entries correctly when " << "Failed to insert underlying dummy entries correctly when "
"encountering cache resevation failure due to full cache"; "encountering cache resevation failure due to full cache";
@ -191,6 +201,8 @@ TEST(CacheReservationManagerIncreaseReservcationOnFullCacheTest,
2 * kSizeDummyEntry) 2 * kSizeDummyEntry)
<< "Failed to bookkeep cache reservation decrease correctly after " << "Failed to bookkeep cache reservation decrease correctly after "
"encountering cache reservation due to full cache"; "encountering cache reservation due to full cache";
EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
<< "Failed to bookkeep the used memory correctly";
EXPECT_GE(cache->GetPinnedUsage(), 2 * kSizeDummyEntry) EXPECT_GE(cache->GetPinnedUsage(), 2 * kSizeDummyEntry)
<< "Failed to release underlying dummy entries correctly on cache " << "Failed to release underlying dummy entries correctly on cache "
"reservation decrease after encountering cache resevation failure due " "reservation decrease after encountering cache resevation failure due "
@ -218,6 +230,8 @@ TEST(CacheReservationManagerIncreaseReservcationOnFullCacheTest,
<< "Failed to bookkeep correctly (i.e, bookkeep only successful dummy " << "Failed to bookkeep correctly (i.e, bookkeep only successful dummy "
"entry insertions) when encountering cache resevation failure due to " "entry insertions) when encountering cache resevation failure due to "
"full cache"; "full cache";
EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
<< "Failed to bookkeep the used memory correctly";
EXPECT_GE(cache->GetPinnedUsage(), 1 * kSizeDummyEntry) EXPECT_GE(cache->GetPinnedUsage(), 1 * kSizeDummyEntry)
<< "Failed to insert underlying dummy entries correctly when " << "Failed to insert underlying dummy entries correctly when "
"encountering cache resevation failure due to full cache"; "encountering cache resevation failure due to full cache";
@ -239,6 +253,8 @@ TEST(CacheReservationManagerIncreaseReservcationOnFullCacheTest,
5 * kSizeDummyEntry) 5 * kSizeDummyEntry)
<< "Failed to bookkeep cache reservation increase correctly after " << "Failed to bookkeep cache reservation increase correctly after "
"increasing cache capacity and mitigating cache full error"; "increasing cache capacity and mitigating cache full error";
EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
<< "Failed to bookkeep the used memory correctly";
EXPECT_GE(cache->GetPinnedUsage(), 5 * kSizeDummyEntry) EXPECT_GE(cache->GetPinnedUsage(), 5 * kSizeDummyEntry)
<< "Failed to insert underlying dummy entries correctly after increasing " << "Failed to insert underlying dummy entries correctly after increasing "
"cache capacity and mitigating cache full error"; "cache capacity and mitigating cache full error";
@ -258,6 +274,7 @@ TEST_F(CacheReservationManagerTest,
ASSERT_EQ(s, Status::OK()); ASSERT_EQ(s, Status::OK());
ASSERT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), ASSERT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(),
2 * kSizeDummyEntry); 2 * kSizeDummyEntry);
ASSERT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used);
ASSERT_GE(cache->GetPinnedUsage(), 2 * kSizeDummyEntry); ASSERT_GE(cache->GetPinnedUsage(), 2 * kSizeDummyEntry);
ASSERT_LT(cache->GetPinnedUsage(), ASSERT_LT(cache->GetPinnedUsage(),
2 * kSizeDummyEntry + kMetaDataChargeOverhead); 2 * kSizeDummyEntry + kMetaDataChargeOverhead);
@ -271,6 +288,8 @@ TEST_F(CacheReservationManagerTest,
EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(),
1 * kSizeDummyEntry) 1 * kSizeDummyEntry)
<< "Failed to bookkeep cache reservation decrease correctly"; << "Failed to bookkeep cache reservation decrease correctly";
EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
<< "Failed to bookkeep the used memory correctly";
EXPECT_GE(cache->GetPinnedUsage(), 1 * kSizeDummyEntry) EXPECT_GE(cache->GetPinnedUsage(), 1 * kSizeDummyEntry)
<< "Failed to decrease underlying dummy entries in cache correctly"; << "Failed to decrease underlying dummy entries in cache correctly";
EXPECT_LT(cache->GetPinnedUsage(), EXPECT_LT(cache->GetPinnedUsage(),
@ -288,6 +307,7 @@ TEST_F(CacheReservationManagerTest,
ASSERT_EQ(s, Status::OK()); ASSERT_EQ(s, Status::OK());
ASSERT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), ASSERT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(),
2 * kSizeDummyEntry); 2 * kSizeDummyEntry);
ASSERT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used);
ASSERT_GE(cache->GetPinnedUsage(), 2 * kSizeDummyEntry); ASSERT_GE(cache->GetPinnedUsage(), 2 * kSizeDummyEntry);
ASSERT_LT(cache->GetPinnedUsage(), ASSERT_LT(cache->GetPinnedUsage(),
2 * kSizeDummyEntry + kMetaDataChargeOverhead); 2 * kSizeDummyEntry + kMetaDataChargeOverhead);
@ -301,6 +321,8 @@ TEST_F(CacheReservationManagerTest,
EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(),
1 * kSizeDummyEntry) 1 * kSizeDummyEntry)
<< "Failed to bookkeep cache reservation decrease correctly"; << "Failed to bookkeep cache reservation decrease correctly";
EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
<< "Failed to bookkeep the used memory correctly";
EXPECT_GE(cache->GetPinnedUsage(), 1 * kSizeDummyEntry) EXPECT_GE(cache->GetPinnedUsage(), 1 * kSizeDummyEntry)
<< "Failed to decrease underlying dummy entries in cache correctly"; << "Failed to decrease underlying dummy entries in cache correctly";
EXPECT_LT(cache->GetPinnedUsage(), EXPECT_LT(cache->GetPinnedUsage(),
@ -330,6 +352,7 @@ TEST(CacheReservationManagerWithDelayedDecreaseTest,
ASSERT_EQ(s, Status::OK()); ASSERT_EQ(s, Status::OK());
ASSERT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), ASSERT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(),
8 * kSizeDummyEntry); 8 * kSizeDummyEntry);
ASSERT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used);
std::size_t initial_pinned_usage = cache->GetPinnedUsage(); std::size_t initial_pinned_usage = cache->GetPinnedUsage();
ASSERT_GE(initial_pinned_usage, 8 * kSizeDummyEntry); ASSERT_GE(initial_pinned_usage, 8 * kSizeDummyEntry);
ASSERT_LT(initial_pinned_usage, ASSERT_LT(initial_pinned_usage,
@ -344,6 +367,8 @@ TEST(CacheReservationManagerWithDelayedDecreaseTest,
8 * kSizeDummyEntry) 8 * kSizeDummyEntry)
<< "Failed to bookkeep correctly when delaying cache reservation " << "Failed to bookkeep correctly when delaying cache reservation "
"decrease"; "decrease";
EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
<< "Failed to bookkeep the used memory correctly";
EXPECT_EQ(cache->GetPinnedUsage(), initial_pinned_usage) EXPECT_EQ(cache->GetPinnedUsage(), initial_pinned_usage)
<< "Failed to delay decreasing underlying dummy entries in cache"; << "Failed to delay decreasing underlying dummy entries in cache";
@ -356,6 +381,8 @@ TEST(CacheReservationManagerWithDelayedDecreaseTest,
8 * kSizeDummyEntry) 8 * kSizeDummyEntry)
<< "Failed to bookkeep correctly when delaying cache reservation " << "Failed to bookkeep correctly when delaying cache reservation "
"decrease"; "decrease";
EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
<< "Failed to bookkeep the used memory correctly";
EXPECT_EQ(cache->GetPinnedUsage(), initial_pinned_usage) EXPECT_EQ(cache->GetPinnedUsage(), initial_pinned_usage)
<< "Failed to delay decreasing underlying dummy entries in cache"; << "Failed to delay decreasing underlying dummy entries in cache";
@ -370,6 +397,8 @@ TEST(CacheReservationManagerWithDelayedDecreaseTest,
6 * kSizeDummyEntry) 6 * kSizeDummyEntry)
<< "Failed to bookkeep correctly when new_mem_used < " << "Failed to bookkeep correctly when new_mem_used < "
"GetTotalReservedCacheSize() * 3 / 4 on delayed decrease mode"; "GetTotalReservedCacheSize() * 3 / 4 on delayed decrease mode";
EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used)
<< "Failed to bookkeep the used memory correctly";
EXPECT_GE(cache->GetPinnedUsage(), 6 * kSizeDummyEntry) EXPECT_GE(cache->GetPinnedUsage(), 6 * kSizeDummyEntry)
<< "Failed to decrease underlying dummy entries in cache when " << "Failed to decrease underlying dummy entries in cache when "
"new_mem_used < GetTotalReservedCacheSize() * 3 / 4 on delayed " "new_mem_used < GetTotalReservedCacheSize() * 3 / 4 on delayed "

Loading…
Cancel
Save