diff --git a/cache/cache_reservation_manager.cc b/cache/cache_reservation_manager.cc index a4fd48fbc..dd69da785 100644 --- a/cache/cache_reservation_manager.cc +++ b/cache/cache_reservation_manager.cc @@ -79,6 +79,23 @@ template Status CacheReservationManager::UpdateCacheReservation< template Status CacheReservationManager::UpdateCacheReservation< CacheEntryRole::kMisc>(std::size_t new_mem_used); +template +Status CacheReservationManager::MakeCacheReservation( + std::size_t incremental_memory_used, + std::unique_ptr>* handle) { + assert(handle != nullptr); + Status s = + UpdateCacheReservation(GetTotalMemoryUsed() + incremental_memory_used); + (*handle).reset(new CacheReservationHandle(incremental_memory_used, + shared_from_this())); + return s; +} + +template Status +CacheReservationManager::MakeCacheReservation( + std::size_t incremental_memory_used, + std::unique_ptr>* handle); + template Status CacheReservationManager::IncreaseCacheReservation( std::size_t new_mem_used) { @@ -135,4 +152,27 @@ Slice CacheReservationManager::GetNextCacheKey() { EncodeVarint64(cache_key_ + kCacheKeyPrefixSize, next_cache_key_id_++); return Slice(cache_key_, static_cast(end - cache_key_)); } + +template +CacheReservationHandle::CacheReservationHandle( + std::size_t incremental_memory_used, + std::shared_ptr cache_res_mgr) + : incremental_memory_used_(incremental_memory_used) { + assert(cache_res_mgr != nullptr); + cache_res_mgr_ = cache_res_mgr; +} + +template +CacheReservationHandle::~CacheReservationHandle() { + assert(cache_res_mgr_ != nullptr); + assert(cache_res_mgr_->GetTotalMemoryUsed() >= incremental_memory_used_); + + Status s = cache_res_mgr_->UpdateCacheReservation( + cache_res_mgr_->GetTotalMemoryUsed() - incremental_memory_used_); + s.PermitUncheckedError(); +} + +// Explicitly instantiate templates for "CacheEntryRole" values we use. +// This makes it possible to keep the template definitions in the .cc file. +template class CacheReservationHandle; } // namespace ROCKSDB_NAMESPACE diff --git a/cache/cache_reservation_manager.h b/cache/cache_reservation_manager.h index f248b7bfd..faa516217 100644 --- a/cache/cache_reservation_manager.h +++ b/cache/cache_reservation_manager.h @@ -24,12 +24,16 @@ namespace ROCKSDB_NAMESPACE { +template +class CacheReservationHandle; + // CacheReservationManager is for reserving cache space for the memory used // through inserting/releasing dummy entries in the cache. // // This class is NOT thread-safe, except that GetTotalReservedCacheSize() // can be called without external synchronization. -class CacheReservationManager { +class CacheReservationManager + : public std::enable_shared_from_this { public: // Construct a CacheReservationManager // @param cache The cache where dummy entries are inserted and released for @@ -54,6 +58,10 @@ class CacheReservationManager { template + // One of the two ways of reserving/releasing cache, + // see CacheReservationManager::MakeCacheReservation() for the other. + // Use ONLY one of them to prevent unexpected behavior. + // // Insert and release dummy entries in the cache to // match the size of total dummy entries with the least multiple of // kSizeDummyEntry greater than or equal to new_mem_used @@ -84,6 +92,48 @@ class CacheReservationManager { // On keeping dummy entries the same, it always returns Status::OK(). Status UpdateCacheReservation(std::size_t new_memory_used); + // One of the two ways of reserving/releasing cache, + // see CacheReservationManager::UpdateCacheReservation() for the other. + // Use ONLY one of them to prevent unexpected behavior. + // + // Insert dummy entries in the cache for the incremental memory usage + // to match the size of total dummy entries with the least multiple of + // kSizeDummyEntry greater than or equal to the total memory used. + // + // A CacheReservationHandle is returned as an output parameter. + // The reserved dummy entries are automatically released on the destruction of + // this handle, which achieves better RAII per cache reservation. + // + // WARNING: Deallocate all the handles of the CacheReservationManager object + // before deallocating the object to prevent unexpected behavior. + // + // @param incremental_memory_used The number of bytes increased in memory + // usage. + // + // Calling GetTotalMemoryUsed() afterward will return the total memory + // increased by this number, even when calling MakeCacheReservation() + // returns non-ok status. + // + // Since the class is NOT thread-safe, external synchronization in + // calling MakeCacheReservation() is needed if you want + // GetTotalMemoryUsed() indeed returns the latest memory used. + // + // @param handle An pointer to std::unique_ptr> that + // manages the lifetime of the handle and its cache reservation. + // + // @return It returns Status::OK() if all dummy + // entry insertions succeed. + // Otherwise, it returns the first non-ok status; + // + // REQUIRES: handle != nullptr + // REQUIRES: The CacheReservationManager object is NOT managed by + // std::unique_ptr as CacheReservationHandle needs to + // shares ownership to the CacheReservationManager object. + template + Status MakeCacheReservation( + std::size_t incremental_memory_used, + std::unique_ptr> *handle); + // Return the size of the cache (which is a multiple of kSizeDummyEntry) // successfully reserved by calling UpdateCacheReservation(). // @@ -121,4 +171,21 @@ class CacheReservationManager { // The non-prefix part will be updated according to the ID to use. char cache_key_[kCacheKeyPrefixSize + kMaxVarint64Length]; }; + +// CacheReservationHandle is for managing the lifetime of a cache reservation +// This class is NOT thread-safe +template +class CacheReservationHandle { + public: + // REQUIRES: cache_res_mgr != nullptr + explicit CacheReservationHandle( + std::size_t incremental_memory_used, + std::shared_ptr cache_res_mgr); + + ~CacheReservationHandle(); + + private: + std::size_t incremental_memory_used_; + std::shared_ptr cache_res_mgr_; +}; } // namespace ROCKSDB_NAMESPACE diff --git a/cache/cache_reservation_manager_test.cc b/cache/cache_reservation_manager_test.cc index 5886dbd58..9aa367302 100644 --- a/cache/cache_reservation_manager_test.cc +++ b/cache/cache_reservation_manager_test.cc @@ -438,6 +438,74 @@ TEST(CacheReservationManagerDestructorTest, << "Failed to release remaining underlying dummy entries in cache in " "CacheReservationManager's destructor"; } + +TEST(CacheReservationHandleTest, HandleTest) { + constexpr std::size_t kOneGigabyte = 1024 * 1024 * 1024; + constexpr std::size_t kSizeDummyEntry = 256 * 1024; + constexpr std::size_t kMetaDataChargeOverhead = 10000; + + LRUCacheOptions lo; + lo.capacity = kOneGigabyte; + lo.num_shard_bits = 0; + std::shared_ptr cache = NewLRUCache(lo); + + std::shared_ptr test_cache_rev_mng( + std::make_shared(cache)); + + std::size_t mem_used = 0; + const std::size_t incremental_mem_used_handle_1 = 1 * kSizeDummyEntry; + const std::size_t incremental_mem_used_handle_2 = 2 * kSizeDummyEntry; + std::unique_ptr> handle_1, + handle_2; + + // To test consecutive CacheReservationManager::MakeCacheReservation works + // correctly in terms of returning the handle as well as updating cache + // reservation and the latest total memory used + Status s = test_cache_rev_mng->MakeCacheReservation( + incremental_mem_used_handle_1, &handle_1); + mem_used = mem_used + incremental_mem_used_handle_1; + ASSERT_EQ(s, Status::OK()); + EXPECT_TRUE(handle_1 != nullptr); + EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), mem_used); + EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), mem_used); + EXPECT_GE(cache->GetPinnedUsage(), mem_used); + EXPECT_LT(cache->GetPinnedUsage(), mem_used + kMetaDataChargeOverhead); + + s = test_cache_rev_mng->MakeCacheReservation( + incremental_mem_used_handle_2, &handle_2); + mem_used = mem_used + incremental_mem_used_handle_2; + ASSERT_EQ(s, Status::OK()); + EXPECT_TRUE(handle_2 != nullptr); + EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), mem_used); + EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), mem_used); + EXPECT_GE(cache->GetPinnedUsage(), mem_used); + EXPECT_LT(cache->GetPinnedUsage(), mem_used + kMetaDataChargeOverhead); + + // To test CacheReservationHandle::~CacheReservationHandle() works correctly + // in releasing the cache reserved for the handle + handle_1.reset(); + EXPECT_TRUE(handle_1 == nullptr); + mem_used = mem_used - incremental_mem_used_handle_1; + EXPECT_EQ(test_cache_rev_mng->GetTotalReservedCacheSize(), mem_used); + EXPECT_EQ(test_cache_rev_mng->GetTotalMemoryUsed(), mem_used); + EXPECT_GE(cache->GetPinnedUsage(), mem_used); + EXPECT_LT(cache->GetPinnedUsage(), mem_used + kMetaDataChargeOverhead); + + // To test the actual CacheReservationManager object won't be deallocated + // as long as there remain handles pointing to it. + // We strongly recommend deallocating CacheReservationManager object only + // after all its handles are deallocated to keep things easy to reasonate + test_cache_rev_mng.reset(); + EXPECT_GE(cache->GetPinnedUsage(), mem_used); + EXPECT_LT(cache->GetPinnedUsage(), mem_used + kMetaDataChargeOverhead); + + handle_2.reset(); + // The CacheReservationManager object is now deallocated since all the handles + // and its original pointer is gone + mem_used = mem_used - incremental_mem_used_handle_2; + EXPECT_EQ(mem_used, 0); + EXPECT_EQ(cache->GetPinnedUsage(), mem_used); +} } // namespace ROCKSDB_NAMESPACE int main(int argc, char** argv) {