RAII support for per cache reservation through handle (#9130)
Summary:
Note: This PR is the 3rd PR of a bigger PR stack (https://github.com/facebook/rocksdb/issues/9073) and depends on the second PR (https://github.com/facebook/rocksdb/pull/9071). **See changes from this PR only 00447324d0
**
Context:
pdillinger brought up a good [point](https://github.com/facebook/rocksdb/pull/9073#discussion_r741478309) about lacking RAII support for per cache reservation in `CacheReservationManager` when reviewing https://github.com/facebook/rocksdb/pull/9073.
To summarize the discussion, the current API `CacheReservationManager::UpdateCacheReservation()` requires callers to explicitly calculate and pass in a correct`new_mem_used` to release a cache reservation (if they don't want to rely on the clean-up during `CacheReservationManager`'s destruction - such as they want to release it earlier).
While this implementation has convenience in some use-case such as `WriteBufferManager`, where [reservation](https://github.com/facebook/rocksdb/blob/main/memtable/write_buffer_manager.cc#L69-L91) and [release](https://github.com/facebook/rocksdb/blob/main/memtable/write_buffer_manager.cc#L109-L129) amounts do not necessarily correspond symmetrically and thus a flexible `new_mem_used` inputing is needed, it can be prone to caller's calculation error as well as cause a mass of codes in releasing cache in other use-case such as filter construction, where reservation and release amounts do correspond symmetrically and many code paths requiring a cache release, as [pointed](https://github.com/facebook/rocksdb/pull/9073#discussion_r741478309) out by pdillinger.
Therefore we decided to provide a new API in `CacheReservationManager` to update reservation with better RAII support for per cache reservation, using a handle to manage the life time of that particular cache reservation.
- Added a new class `CacheReservationHandle`
- Added a new API `CacheReservationManager::MakeCacheReservation()` that outputs a `CacheReservationHandle` for managing the reservation
- Updated class comments to clarify two different cache reservation methods
Tests:
- Passing new tests
Pull Request resolved: https://github.com/facebook/rocksdb/pull/9130
Reviewed By: pdillinger
Differential Revision: D32199446
Pulled By: hx235
fbshipit-source-id: 1cba7c636e5ecfb55b0c1e0c2d218cc9b5b30b4e
This commit is contained in:
parent
ffd6085e1f
commit
2fbe32b0c1
40
cache/cache_reservation_manager.cc
vendored
40
cache/cache_reservation_manager.cc
vendored
@ -79,6 +79,23 @@ template Status CacheReservationManager::UpdateCacheReservation<
|
|||||||
template Status CacheReservationManager::UpdateCacheReservation<
|
template Status CacheReservationManager::UpdateCacheReservation<
|
||||||
CacheEntryRole::kMisc>(std::size_t new_mem_used);
|
CacheEntryRole::kMisc>(std::size_t new_mem_used);
|
||||||
|
|
||||||
|
template <CacheEntryRole R>
|
||||||
|
Status CacheReservationManager::MakeCacheReservation(
|
||||||
|
std::size_t incremental_memory_used,
|
||||||
|
std::unique_ptr<CacheReservationHandle<R>>* handle) {
|
||||||
|
assert(handle != nullptr);
|
||||||
|
Status s =
|
||||||
|
UpdateCacheReservation<R>(GetTotalMemoryUsed() + incremental_memory_used);
|
||||||
|
(*handle).reset(new CacheReservationHandle<R>(incremental_memory_used,
|
||||||
|
shared_from_this()));
|
||||||
|
return s;
|
||||||
|
}
|
||||||
|
|
||||||
|
template Status
|
||||||
|
CacheReservationManager::MakeCacheReservation<CacheEntryRole::kMisc>(
|
||||||
|
std::size_t incremental_memory_used,
|
||||||
|
std::unique_ptr<CacheReservationHandle<CacheEntryRole::kMisc>>* handle);
|
||||||
|
|
||||||
template <CacheEntryRole R>
|
template <CacheEntryRole R>
|
||||||
Status CacheReservationManager::IncreaseCacheReservation(
|
Status CacheReservationManager::IncreaseCacheReservation(
|
||||||
std::size_t new_mem_used) {
|
std::size_t new_mem_used) {
|
||||||
@ -135,4 +152,27 @@ Slice CacheReservationManager::GetNextCacheKey() {
|
|||||||
EncodeVarint64(cache_key_ + kCacheKeyPrefixSize, next_cache_key_id_++);
|
EncodeVarint64(cache_key_ + kCacheKeyPrefixSize, next_cache_key_id_++);
|
||||||
return Slice(cache_key_, static_cast<std::size_t>(end - cache_key_));
|
return Slice(cache_key_, static_cast<std::size_t>(end - cache_key_));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
template <CacheEntryRole R>
|
||||||
|
CacheReservationHandle<R>::CacheReservationHandle(
|
||||||
|
std::size_t incremental_memory_used,
|
||||||
|
std::shared_ptr<CacheReservationManager> cache_res_mgr)
|
||||||
|
: incremental_memory_used_(incremental_memory_used) {
|
||||||
|
assert(cache_res_mgr != nullptr);
|
||||||
|
cache_res_mgr_ = cache_res_mgr;
|
||||||
|
}
|
||||||
|
|
||||||
|
template <CacheEntryRole R>
|
||||||
|
CacheReservationHandle<R>::~CacheReservationHandle() {
|
||||||
|
assert(cache_res_mgr_ != nullptr);
|
||||||
|
assert(cache_res_mgr_->GetTotalMemoryUsed() >= incremental_memory_used_);
|
||||||
|
|
||||||
|
Status s = cache_res_mgr_->UpdateCacheReservation<R>(
|
||||||
|
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<CacheEntryRole::kMisc>;
|
||||||
} // namespace ROCKSDB_NAMESPACE
|
} // namespace ROCKSDB_NAMESPACE
|
||||||
|
69
cache/cache_reservation_manager.h
vendored
69
cache/cache_reservation_manager.h
vendored
@ -24,12 +24,16 @@
|
|||||||
|
|
||||||
namespace ROCKSDB_NAMESPACE {
|
namespace ROCKSDB_NAMESPACE {
|
||||||
|
|
||||||
|
template <CacheEntryRole R>
|
||||||
|
class CacheReservationHandle;
|
||||||
|
|
||||||
// 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, except that GetTotalReservedCacheSize()
|
// This class is NOT thread-safe, except that GetTotalReservedCacheSize()
|
||||||
// can be called without external synchronization.
|
// can be called without external synchronization.
|
||||||
class CacheReservationManager {
|
class CacheReservationManager
|
||||||
|
: public std::enable_shared_from_this<CacheReservationManager> {
|
||||||
public:
|
public:
|
||||||
// Construct a CacheReservationManager
|
// Construct a CacheReservationManager
|
||||||
// @param cache The cache where dummy entries are inserted and released for
|
// @param cache The cache where dummy entries are inserted and released for
|
||||||
@ -54,6 +58,10 @@ class CacheReservationManager {
|
|||||||
|
|
||||||
template <CacheEntryRole R>
|
template <CacheEntryRole R>
|
||||||
|
|
||||||
|
// 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
|
// Insert and release dummy entries in the cache to
|
||||||
// match the size of total dummy entries with the least multiple of
|
// match the size of total dummy entries with the least multiple of
|
||||||
// kSizeDummyEntry greater than or equal to new_mem_used
|
// 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().
|
// 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);
|
||||||
|
|
||||||
|
// 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<CacheReservationHandle<R>> 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 <CacheEntryRole R>
|
||||||
|
Status MakeCacheReservation(
|
||||||
|
std::size_t incremental_memory_used,
|
||||||
|
std::unique_ptr<CacheReservationHandle<R>> *handle);
|
||||||
|
|
||||||
// Return the size of the cache (which is a multiple of kSizeDummyEntry)
|
// Return the size of the cache (which is a multiple of kSizeDummyEntry)
|
||||||
// successfully reserved by calling UpdateCacheReservation().
|
// successfully reserved by calling UpdateCacheReservation().
|
||||||
//
|
//
|
||||||
@ -121,4 +171,21 @@ class CacheReservationManager {
|
|||||||
// 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.
|
||||||
char cache_key_[kCacheKeyPrefixSize + kMaxVarint64Length];
|
char cache_key_[kCacheKeyPrefixSize + kMaxVarint64Length];
|
||||||
};
|
};
|
||||||
|
|
||||||
|
// CacheReservationHandle is for managing the lifetime of a cache reservation
|
||||||
|
// This class is NOT thread-safe
|
||||||
|
template <CacheEntryRole R>
|
||||||
|
class CacheReservationHandle {
|
||||||
|
public:
|
||||||
|
// REQUIRES: cache_res_mgr != nullptr
|
||||||
|
explicit CacheReservationHandle(
|
||||||
|
std::size_t incremental_memory_used,
|
||||||
|
std::shared_ptr<CacheReservationManager> cache_res_mgr);
|
||||||
|
|
||||||
|
~CacheReservationHandle();
|
||||||
|
|
||||||
|
private:
|
||||||
|
std::size_t incremental_memory_used_;
|
||||||
|
std::shared_ptr<CacheReservationManager> cache_res_mgr_;
|
||||||
|
};
|
||||||
} // namespace ROCKSDB_NAMESPACE
|
} // namespace ROCKSDB_NAMESPACE
|
||||||
|
68
cache/cache_reservation_manager_test.cc
vendored
68
cache/cache_reservation_manager_test.cc
vendored
@ -438,6 +438,74 @@ TEST(CacheReservationManagerDestructorTest,
|
|||||||
<< "Failed to release remaining underlying dummy entries in cache in "
|
<< "Failed to release remaining underlying dummy entries in cache in "
|
||||||
"CacheReservationManager's destructor";
|
"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> cache = NewLRUCache(lo);
|
||||||
|
|
||||||
|
std::shared_ptr<CacheReservationManager> test_cache_rev_mng(
|
||||||
|
std::make_shared<CacheReservationManager>(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<CacheReservationHandle<CacheEntryRole::kMisc>> 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<CacheEntryRole::kMisc>(
|
||||||
|
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<CacheEntryRole::kMisc>(
|
||||||
|
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
|
} // namespace ROCKSDB_NAMESPACE
|
||||||
|
|
||||||
int main(int argc, char** argv) {
|
int main(int argc, char** argv) {
|
||||||
|
Loading…
Reference in New Issue
Block a user