From 23f1100d773330a70415ea9ce192c70e2b89b2b3 Mon Sep 17 00:00:00 2001 From: Daniel Byrne Date: Wed, 22 Nov 2023 12:13:48 -0800 Subject: [PATCH] serialize multi-tier allocator with unit test --- cachelib/allocator/CacheAllocator.h | 77 ++++++---- cachelib/allocator/serialize/objects.thrift | 17 ++- .../allocator/tests/AllocatorTypeTest.cpp | 4 + cachelib/allocator/tests/BaseAllocatorTest.h | 135 ++++++++++++++++++ 4 files changed, 202 insertions(+), 31 deletions(-) diff --git a/cachelib/allocator/CacheAllocator.h b/cachelib/allocator/CacheAllocator.h index 7422c1c61c..096d8dbdd2 100644 --- a/cachelib/allocator/CacheAllocator.h +++ b/cachelib/allocator/CacheAllocator.h @@ -395,6 +395,7 @@ class CacheAllocator : public CacheBase { using MMSerializationTypeContainer = typename MMType::SerializationTypeContainer; using AccessSerializationType = typename AccessType::SerializationType; + using AllocatorsSerializationType = serialization::MemoryAllocatorCollection; using ShmManager = facebook::cachelib::ShmManager; @@ -2153,7 +2154,7 @@ class CacheAllocator : public CacheBase { PrivateSegmentOpts createPrivateSegmentOpts(TierId tid); std::unique_ptr createPrivateAllocator(TierId tid); std::unique_ptr createNewMemoryAllocator(TierId tid); - std::unique_ptr restoreMemoryAllocator(TierId tid); + std::unique_ptr restoreMemoryAllocator(TierId tid, const serialization::MemoryAllocatorObject& sAllocator); std::unique_ptr restoreCCacheManager(TierId tid); PoolIds filterCompactCachePools(const PoolIds& poolIds) const; @@ -2698,9 +2699,10 @@ CacheAllocator::createNewMemoryAllocator(TierId tid) { template std::unique_ptr -CacheAllocator::restoreMemoryAllocator(TierId tid) { +CacheAllocator::restoreMemoryAllocator(TierId tid, + const serialization::MemoryAllocatorObject& sAllocator) { return std::make_unique( - deserializer_->deserialize(), + sAllocator, shmManager_ ->attachShm(detail::kShmCacheName + std::to_string(tid), config_.slabMemoryBaseAddr, createShmCacheOpts(tid)).addr, @@ -2732,8 +2734,11 @@ template std::vector> CacheAllocator::restoreAllocators() { std::vector> allocators; + const auto allocatorCollection = + deserializer_->deserialize(); + auto allocMap = *allocatorCollection.allocators(); for (int tid = 0; tid < getNumTiers(); tid++) { - allocators.emplace_back(restoreMemoryAllocator(tid)); + allocators.emplace_back(restoreMemoryAllocator(tid,allocMap[tid])); } return allocators; } @@ -6407,26 +6412,43 @@ folly::IOBufQueue CacheAllocator::saveStateToIOBuf() { *metadata_.numChainedChildItems() = stats_.numChainedChildItems.get(); *metadata_.numAbortedSlabReleases() = stats_.numAbortedSlabReleases.get(); + const auto numTiers = getNumTiers(); // TODO: implement serialization for multiple tiers - auto serializeMMContainers = [](MMContainers& mmContainers) { - MMSerializationTypeContainer state; - for (unsigned int i = 0; i < 1 /* TODO: */ ; ++i) { + auto serializeMMContainers = [numTiers](MMContainers& mmContainers) { + std::map containers; + for (unsigned int i = 0; i < numTiers; ++i) { for (unsigned int j = 0; j < mmContainers[i].size(); ++j) { for (unsigned int k = 0; k < mmContainers[i][j].size(); ++k) { if (mmContainers[i][j][k]) { - state.pools_ref()[j][k] = mmContainers[i][j][k]->saveState(); + serialization::MemoryDescriptorObject md; + md.tid_ref() = i; + md.pid_ref() = j; + md.cid_ref() = k; + containers[md] = mmContainers[i][j][k]->saveState(); } } } } + MMSerializationTypeContainer state; + state.containers_ref() = containers; return state; }; MMSerializationTypeContainer mmContainersState = serializeMMContainers(mmContainers_); AccessSerializationType accessContainerState = accessContainer_->saveState(); - // TODO: foreach allocator - MemoryAllocator::SerializationType allocatorState = allocator_[0]->saveState(); + + auto serializeAllocators = [numTiers,this]() { + AllocatorsSerializationType state; + std::map allocators; + for (int i = 0; i < numTiers; ++i) { + allocators[i] = allocator_[i]->saveState(); + } + state.allocators_ref() = allocators; + return state; + }; + AllocatorsSerializationType allocatorsState = serializeAllocators(); + CCacheManager::SerializationType ccState = compactCacheManager_->saveState(); AccessSerializationType chainedItemAccessContainerState = @@ -6436,7 +6458,7 @@ folly::IOBufQueue CacheAllocator::saveStateToIOBuf() { // results into a single buffer. folly::IOBufQueue queue; Serializer::serializeToIOBufQueue(queue, metadata_); - Serializer::serializeToIOBufQueue(queue, allocatorState); + Serializer::serializeToIOBufQueue(queue, allocatorsState); Serializer::serializeToIOBufQueue(queue, ccState); Serializer::serializeToIOBufQueue(queue, mmContainersState); Serializer::serializeToIOBufQueue(queue, accessContainerState); @@ -6559,23 +6581,22 @@ CacheAllocator::deserializeMMContainers( * only works for a single (topmost) tier. */ MMContainers mmContainers{getNumTiers()}; - for (auto& kvPool : *container.pools_ref()) { - auto i = static_cast(kvPool.first); - auto& pool = getPool(i); - for (auto& kv : kvPool.second) { - auto j = static_cast(kv.first); - for (TierId tid = 0; tid < getNumTiers(); tid++) { - MMContainerPtr ptr = - std::make_unique(kv.second, - compressor); - auto config = ptr->getConfig(); - config.addExtraConfig(config_.trackTailHits - ? pool.getAllocationClass(j).getAllocsPerSlab() - : 0); - ptr->setConfig(config); - mmContainers[tid][i][j] = std::move(ptr); - } - } + std::map containerMap = + *container.containers(); + for (auto md : containerMap) { + uint32_t tid = *md.first.tid(); + uint32_t pid = *md.first.pid(); + uint32_t cid = *md.first.cid(); + auto& pool = getPoolByTid(pid,tid); + MMContainerPtr ptr = + std::make_unique(md.second, + compressor); + auto config = ptr->getConfig(); + config.addExtraConfig(config_.trackTailHits + ? pool.getAllocationClass(cid).getAllocsPerSlab() + : 0); + ptr->setConfig(config); + mmContainers[tid][pid][cid] = std::move(ptr); } // We need to drop the unevictableMMContainer in the desierializer. // TODO: remove this at version 17. diff --git a/cachelib/allocator/serialize/objects.thrift b/cachelib/allocator/serialize/objects.thrift index 61297cdf1d..c702f3e335 100644 --- a/cachelib/allocator/serialize/objects.thrift +++ b/cachelib/allocator/serialize/objects.thrift @@ -17,11 +17,22 @@ namespace cpp2 facebook.cachelib.serialization include "cachelib/allocator/datastruct/serialize/objects.thrift" +include "cachelib/allocator/memory/serialize/objects.thrift" // Adding a new "required" field will cause the cache to be dropped // in the next release for our users. If the field needs to be required, // make sure to communicate that with our users. +struct MemoryAllocatorCollection { + 1: required map allocators; +} + +struct MemoryDescriptorObject { + 1: required i32 tid; + 2: required i32 pid; + 3: required i32 cid; +} + struct CacheAllocatorMetadata { 1: required i64 allocatorVersion; // version of cache alloctor 2: i64 cacheCreationTime = 0; // time when the cache was created. @@ -80,7 +91,7 @@ struct MMLruObject { } struct MMLruCollection { - 1: required map> pools; + 1: required map containers; } struct MM2QConfig { @@ -106,7 +117,7 @@ struct MM2QObject { } struct MM2QCollection { - 1: required map> pools; + 1: required map containers; } struct MMTinyLFUConfig { @@ -130,7 +141,7 @@ struct MMTinyLFUObject { } struct MMTinyLFUCollection { - 1: required map> pools; + 1: required map containers; } struct ChainedHashTableObject { diff --git a/cachelib/allocator/tests/AllocatorTypeTest.cpp b/cachelib/allocator/tests/AllocatorTypeTest.cpp index 28c145b39d..69076eb6ee 100644 --- a/cachelib/allocator/tests/AllocatorTypeTest.cpp +++ b/cachelib/allocator/tests/AllocatorTypeTest.cpp @@ -117,10 +117,14 @@ TYPED_TEST(BaseAllocatorTest, DropFile) { this->testDropFile(); } TYPED_TEST(BaseAllocatorTest, ShmTemporary) { this->testShmTemporary(); } TYPED_TEST(BaseAllocatorTest, Serialization) { this->testSerialization(); } +TYPED_TEST(BaseAllocatorTest, MultiTierSerialization) { this->testMultiTierSerialization(); } TYPED_TEST(BaseAllocatorTest, SerializationMMConfig) { this->testSerializationMMConfig(); } +TYPED_TEST(BaseAllocatorTest, MultiTierSerializationMMConfig) { + this->testMultiTierSerializationMMConfig(); +} TYPED_TEST(BaseAllocatorTest, testSerializationWithFragmentation) { this->testSerializationWithFragmentation(); diff --git a/cachelib/allocator/tests/BaseAllocatorTest.h b/cachelib/allocator/tests/BaseAllocatorTest.h index 16d3c03ccd..2986987784 100644 --- a/cachelib/allocator/tests/BaseAllocatorTest.h +++ b/cachelib/allocator/tests/BaseAllocatorTest.h @@ -1710,6 +1710,141 @@ class BaseAllocatorTest : public AllocatorTest { testShmIsRemoved(config); } + void testMultiTierSerialization() { + std::set evictedKeys; + auto removeCb = + [&evictedKeys](const typename AllocatorT::RemoveCbData& data) { + if (data.context == RemoveContext::kEviction) { + const auto key = data.item.getKey(); + evictedKeys.insert({key.data(), key.size()}); + } + }; + + const size_t nSlabs = 40; + const size_t size = nSlabs * Slab::kSize; + const unsigned int nSizes = 1; + const unsigned int keyLen = 100; + + std::vector sizes; + uint8_t poolId; + + // Test allocations. These allocations should remain after save/restore. + // Original lru allocator - with two tiers + typename AllocatorT::Config config; + config.setCacheSize(size); + config.enableCachePersistence(this->cacheDir_); + config.enablePoolRebalancing(nullptr, std::chrono::seconds{0}); + config.configureMemoryTiers({ + MemoryTierCacheConfig::fromShm() + .setRatio(1).setMemBind(std::string("0")), + MemoryTierCacheConfig::fromShm() + .setRatio(1).setMemBind(std::string("0"))}); + std::vector keys; + { + AllocatorT alloc(AllocatorT::SharedMemNew, config); + const size_t numBytes = alloc.getCacheMemoryStats().ramCacheSize; + poolId = alloc.addPool("foobar", numBytes); + sizes = this->getValidAllocSizes(alloc, poolId, nSlabs, keyLen); + this->fillUpPoolUntilEvictions(alloc, 0, poolId, sizes, keyLen); + this->fillUpPoolUntilEvictions(alloc, 1, poolId, sizes, keyLen); + for (const auto& item : alloc) { + auto key = item.getKey(); + keys.push_back(key.str()); + } + + // save + alloc.shutDown(); + } + + testShmIsNotRemoved(config); + // Restored lru allocator + { + AllocatorT alloc(AllocatorT::SharedMemAttach, config); + for (auto& key : keys) { + auto handle = alloc.find(typename AllocatorT::Key{key}); + ASSERT_NE(nullptr, handle.get()); + } + } + + testShmIsRemoved(config); + // Test LRU eviction and length before and after save/restore + // Original lru allocator + typename AllocatorT::Config config2; + config2.setCacheSize(size); + config2.setRemoveCallback(removeCb); + config2.enableCachePersistence(this->cacheDir_); + config2.configureMemoryTiers({ + MemoryTierCacheConfig::fromShm() + .setRatio(1).setMemBind(std::string("0")), + MemoryTierCacheConfig::fromShm() + .setRatio(1).setMemBind(std::string("0"))}); + { + AllocatorT alloc(AllocatorT::SharedMemNew, config2); + const size_t numBytes = alloc.getCacheMemoryStats().ramCacheSize; + poolId = alloc.addPool("foobar", numBytes); + + sizes = this->getValidAllocSizes(alloc, poolId, nSizes, keyLen); + + this->testLruLength(alloc, poolId, sizes, keyLen, evictedKeys); + + // save + alloc.shutDown(); + } + evictedKeys.clear(); + + testShmIsNotRemoved(config2); + // Restored lru allocator + { + AllocatorT alloc(AllocatorT::SharedMemAttach, config2); + this->testLruLength(alloc, poolId, sizes, keyLen, evictedKeys); + } + + testShmIsRemoved(config2); + } + + void testMultiTierSerializationMMConfig() { + typename AllocatorT::Config config; + config.setCacheSize(20 * Slab::kSize); + config.enableCachePersistence(this->cacheDir_); + config.enablePoolRebalancing(nullptr, std::chrono::seconds{0}); + config.configureMemoryTiers({ + MemoryTierCacheConfig::fromShm() + .setRatio(1).setMemBind(std::string("0")), + MemoryTierCacheConfig::fromShm() + .setRatio(1).setMemBind(std::string("0"))}); + double ratio = 0.2; + + // start allocator + { + AllocatorT alloc(AllocatorT::SharedMemNew, config); + const size_t numBytes = alloc.getCacheMemoryStats().ramCacheSize; + { + typename AllocatorT::MMConfig mmConfig; + mmConfig.lruRefreshRatio = ratio; + auto pid = + alloc.addPool("foobar", numBytes, /* allocSizes = */ {}, mmConfig); + auto handle = util::allocateAccessible(alloc, pid, "key", 10); + ASSERT_NE(nullptr, handle); + auto& container = alloc.getMMContainer(*handle); + EXPECT_DOUBLE_EQ(ratio, container.getConfig().lruRefreshRatio); + } + + // save + alloc.shutDown(); + } + testShmIsNotRemoved(config); + + // restore allocator and check lruRefreshRatio + { + AllocatorT alloc(AllocatorT::SharedMemAttach, config); + auto handle = alloc.find("key"); + ASSERT_NE(nullptr, handle); + auto& container = alloc.getMMContainer(*handle); + EXPECT_DOUBLE_EQ(ratio, container.getConfig().lruRefreshRatio); + } + testShmIsRemoved(config); + } + // Test temporary shared memory mode which is enabled when memory // monitoring is enabled. void testShmTemporary() {