diff --git a/cache/cache_reservation_manager.cc b/cache/cache_reservation_manager.cc index 6a0074871..a4fd48fbc 100644 --- a/cache/cache_reservation_manager.cc +++ b/cache/cache_reservation_manager.cc @@ -23,7 +23,9 @@ namespace ROCKSDB_NAMESPACE { CacheReservationManager::CacheReservationManager(std::shared_ptr cache, 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); cache_ = cache; std::memset(cache_key_, 0, kCacheKeyPrefixSize + kMaxVarint64Length); @@ -39,6 +41,7 @@ CacheReservationManager::~CacheReservationManager() { template Status CacheReservationManager::UpdateCacheReservation( std::size_t new_mem_used) { + memory_used_ = new_mem_used; std::size_t cur_cache_allocated_size = cache_allocated_size_.load(std::memory_order_relaxed); 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); } +std::size_t CacheReservationManager::GetTotalMemoryUsed() { + return memory_used_; +} + Slice CacheReservationManager::GetNextCacheKey() { // Calling this function will have the side-effect of changing the // underlying cache_key_ that is shared among other keys generated from this @@ -128,4 +135,4 @@ Slice CacheReservationManager::GetNextCacheKey() { EncodeVarint64(cache_key_ + kCacheKeyPrefixSize, next_cache_key_id_++); return Slice(cache_key_, static_cast(end - cache_key_)); } -} // namespace ROCKSDB_NAMESPACE \ No newline at end of file +} // namespace ROCKSDB_NAMESPACE diff --git a/cache/cache_reservation_manager.h b/cache/cache_reservation_manager.h index cc69de8d8..f248b7bfd 100644 --- a/cache/cache_reservation_manager.h +++ b/cache/cache_reservation_manager.h @@ -26,7 +26,9 @@ namespace ROCKSDB_NAMESPACE { // CacheReservationManager is for reserving cache space for the memory used // 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 { public: // Construct a CacheReservationManager @@ -53,8 +55,8 @@ class CacheReservationManager { template // Insert and release dummy entries in the cache to - // match the size of total dummy entries with the smallest multiple of - // kSizeDummyEntry that is greater than or equal to new_mem_used + // match the size of total dummy entries with the least multiple of + // kSizeDummyEntry greater than or equal to new_mem_used // // Insert dummy entries if new_memory_used > cache_allocated_size_; // @@ -68,13 +70,34 @@ class CacheReservationManager { // is set true. // // @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 - // entry insertions succeed. Otherwise, it returns the first non-ok status; - // On releasing dummy entries, it always returns Status::OK(). - // On keeping dummy entries the same, it always returns Status::OK(). + // entry insertions succeed. + // Otherwise, it returns the first non-ok status; + // 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); + + // 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(); + // 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; } private: @@ -92,6 +115,7 @@ class CacheReservationManager { std::shared_ptr cache_; bool delayed_decrease_; std::atomic cache_allocated_size_; + std::size_t memory_used_; std::vector dummy_handles_; std::uint64_t next_cache_key_id_ = 0; // The non-prefix part will be updated according to the ID to use. diff --git a/cache/cache_reservation_manager_test.cc b/cache/cache_reservation_manager_test.cc index 03afdb689..5886dbd58 100644 --- a/cache/cache_reservation_manager_test.cc +++ b/cache/cache_reservation_manager_test.cc @@ -81,6 +81,7 @@ TEST_F(CacheReservationManagerTest, KeepCacheReservationTheSame) { ASSERT_EQ(s, Status::OK()); ASSERT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), 1 * kSizeDummyEntry); + ASSERT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used); std::size_t initial_pinned_usage = cache->GetPinnedUsage(); ASSERT_GE(initial_pinned_usage, 1 * kSizeDummyEntry); ASSERT_LT(initial_pinned_usage, @@ -96,6 +97,9 @@ TEST_F(CacheReservationManagerTest, KeepCacheReservationTheSame) { 1 * kSizeDummyEntry) << "Failed to bookkeep correctly when new_mem_used equals to current " "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) << "Failed to keep underlying dummy entries the same when new_mem_used " "equals to current cache reservation"; @@ -113,6 +117,8 @@ TEST_F(CacheReservationManagerTest, EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), 2 * kSizeDummyEntry) << "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) << "Failed to increase underlying dummy entries in cache correctly"; EXPECT_LT(cache->GetPinnedUsage(), @@ -132,6 +138,8 @@ TEST_F(CacheReservationManagerTest, EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), 3 * kSizeDummyEntry) << "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) << "Failed to increase underlying dummy entries in cache correctly"; EXPECT_LT(cache->GetPinnedUsage(), @@ -173,6 +181,8 @@ TEST(CacheReservationManagerIncreaseReservcationOnFullCacheTest, << "Failed to bookkeep correctly (i.e, bookkeep only successful dummy " "entry insertions) when encountering cache resevation failure 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(), 1 * kSizeDummyEntry) << "Failed to insert underlying dummy entries correctly when " "encountering cache resevation failure due to full cache"; @@ -191,6 +201,8 @@ TEST(CacheReservationManagerIncreaseReservcationOnFullCacheTest, 2 * kSizeDummyEntry) << "Failed to bookkeep cache reservation decrease correctly after " "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) << "Failed to release underlying dummy entries correctly on cache " "reservation decrease after encountering cache resevation failure due " @@ -218,6 +230,8 @@ TEST(CacheReservationManagerIncreaseReservcationOnFullCacheTest, << "Failed to bookkeep correctly (i.e, bookkeep only successful dummy " "entry insertions) when encountering cache resevation failure 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(), 1 * kSizeDummyEntry) << "Failed to insert underlying dummy entries correctly when " "encountering cache resevation failure due to full cache"; @@ -239,6 +253,8 @@ TEST(CacheReservationManagerIncreaseReservcationOnFullCacheTest, 5 * kSizeDummyEntry) << "Failed to bookkeep cache reservation increase correctly after " "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) << "Failed to insert underlying dummy entries correctly after increasing " "cache capacity and mitigating cache full error"; @@ -258,6 +274,7 @@ TEST_F(CacheReservationManagerTest, ASSERT_EQ(s, Status::OK()); ASSERT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), 2 * kSizeDummyEntry); + ASSERT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used); ASSERT_GE(cache->GetPinnedUsage(), 2 * kSizeDummyEntry); ASSERT_LT(cache->GetPinnedUsage(), 2 * kSizeDummyEntry + kMetaDataChargeOverhead); @@ -271,6 +288,8 @@ TEST_F(CacheReservationManagerTest, EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), 1 * kSizeDummyEntry) << "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) << "Failed to decrease underlying dummy entries in cache correctly"; EXPECT_LT(cache->GetPinnedUsage(), @@ -288,6 +307,7 @@ TEST_F(CacheReservationManagerTest, ASSERT_EQ(s, Status::OK()); ASSERT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), 2 * kSizeDummyEntry); + ASSERT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used); ASSERT_GE(cache->GetPinnedUsage(), 2 * kSizeDummyEntry); ASSERT_LT(cache->GetPinnedUsage(), 2 * kSizeDummyEntry + kMetaDataChargeOverhead); @@ -301,6 +321,8 @@ TEST_F(CacheReservationManagerTest, EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), 1 * kSizeDummyEntry) << "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) << "Failed to decrease underlying dummy entries in cache correctly"; EXPECT_LT(cache->GetPinnedUsage(), @@ -330,6 +352,7 @@ TEST(CacheReservationManagerWithDelayedDecreaseTest, ASSERT_EQ(s, Status::OK()); ASSERT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), 8 * kSizeDummyEntry); + ASSERT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), new_mem_used); std::size_t initial_pinned_usage = cache->GetPinnedUsage(); ASSERT_GE(initial_pinned_usage, 8 * kSizeDummyEntry); ASSERT_LT(initial_pinned_usage, @@ -344,6 +367,8 @@ TEST(CacheReservationManagerWithDelayedDecreaseTest, 8 * kSizeDummyEntry) << "Failed to bookkeep correctly when delaying cache reservation " "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) << "Failed to delay decreasing underlying dummy entries in cache"; @@ -356,6 +381,8 @@ TEST(CacheReservationManagerWithDelayedDecreaseTest, 8 * kSizeDummyEntry) << "Failed to bookkeep correctly when delaying cache reservation " "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) << "Failed to delay decreasing underlying dummy entries in cache"; @@ -370,6 +397,8 @@ TEST(CacheReservationManagerWithDelayedDecreaseTest, 6 * kSizeDummyEntry) << "Failed to bookkeep correctly when new_mem_used < " "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) << "Failed to decrease underlying dummy entries in cache when " "new_mem_used < GetTotalReservedCacheSize() * 3 / 4 on delayed "