Skip to content

Commit 36b5517

Browse files
byrnedjbyrnedj
committed
serialize multi-tier allocator with unit test (#99)
1 parent d86cc23 commit 36b5517

File tree

4 files changed

+202
-31
lines changed

4 files changed

+202
-31
lines changed

‎cachelib/allocator/CacheAllocator.h

+49-28
Original file line numberDiff line numberDiff line change
@@ -395,6 +395,7 @@ class CacheAllocator : public CacheBase {
395395
using MMSerializationTypeContainer =
396396
typename MMType::SerializationTypeContainer;
397397
using AccessSerializationType = typename AccessType::SerializationType;
398+
using AllocatorsSerializationType = serialization::MemoryAllocatorCollection;
398399

399400
using ShmManager = facebook::cachelib::ShmManager;
400401

@@ -2153,7 +2154,7 @@ class CacheAllocator : public CacheBase {
21532154
PrivateSegmentOpts createPrivateSegmentOpts(TierId tid);
21542155
std::unique_ptr<MemoryAllocator> createPrivateAllocator(TierId tid);
21552156
std::unique_ptr<MemoryAllocator> createNewMemoryAllocator(TierId tid);
2156-
std::unique_ptr<MemoryAllocator> restoreMemoryAllocator(TierId tid);
2157+
std::unique_ptr<MemoryAllocator> restoreMemoryAllocator(TierId tid, const serialization::MemoryAllocatorObject& sAllocator);
21572158
std::unique_ptr<CCacheManager> restoreCCacheManager(TierId tid);
21582159

21592160
PoolIds filterCompactCachePools(const PoolIds& poolIds) const;
@@ -2698,9 +2699,10 @@ CacheAllocator<CacheTrait>::createNewMemoryAllocator(TierId tid) {
26982699

26992700
template <typename CacheTrait>
27002701
std::unique_ptr<MemoryAllocator>
2701-
CacheAllocator<CacheTrait>::restoreMemoryAllocator(TierId tid) {
2702+
CacheAllocator<CacheTrait>::restoreMemoryAllocator(TierId tid,
2703+
const serialization::MemoryAllocatorObject& sAllocator) {
27022704
return std::make_unique<MemoryAllocator>(
2703-
deserializer_->deserialize<MemoryAllocator::SerializationType>(),
2705+
sAllocator,
27042706
shmManager_
27052707
->attachShm(detail::kShmCacheName + std::to_string(tid),
27062708
config_.slabMemoryBaseAddr, createShmCacheOpts(tid)).addr,
@@ -2732,8 +2734,11 @@ template <typename CacheTrait>
27322734
std::vector<std::unique_ptr<MemoryAllocator>>
27332735
CacheAllocator<CacheTrait>::restoreAllocators() {
27342736
std::vector<std::unique_ptr<MemoryAllocator>> allocators;
2737+
const auto allocatorCollection =
2738+
deserializer_->deserialize<AllocatorsSerializationType>();
2739+
auto allocMap = *allocatorCollection.allocators();
27352740
for (int tid = 0; tid < getNumTiers(); tid++) {
2736-
allocators.emplace_back(restoreMemoryAllocator(tid));
2741+
allocators.emplace_back(restoreMemoryAllocator(tid,allocMap[tid]));
27372742
}
27382743
return allocators;
27392744
}
@@ -6410,26 +6415,43 @@ folly::IOBufQueue CacheAllocator<CacheTrait>::saveStateToIOBuf() {
64106415
*metadata_.numChainedChildItems() = stats_.numChainedChildItems.get();
64116416
*metadata_.numAbortedSlabReleases() = stats_.numAbortedSlabReleases.get();
64126417

6418+
const auto numTiers = getNumTiers();
64136419
// TODO: implement serialization for multiple tiers
6414-
auto serializeMMContainers = [](MMContainers& mmContainers) {
6415-
MMSerializationTypeContainer state;
6416-
for (unsigned int i = 0; i < 1 /* TODO: */ ; ++i) {
6420+
auto serializeMMContainers = [numTiers](MMContainers& mmContainers) {
6421+
std::map<serialization::MemoryDescriptorObject,MMSerializationType> containers;
6422+
for (unsigned int i = 0; i < numTiers; ++i) {
64176423
for (unsigned int j = 0; j < mmContainers[i].size(); ++j) {
64186424
for (unsigned int k = 0; k < mmContainers[i][j].size(); ++k) {
64196425
if (mmContainers[i][j][k]) {
6420-
state.pools_ref()[j][k] = mmContainers[i][j][k]->saveState();
6426+
serialization::MemoryDescriptorObject md;
6427+
md.tid_ref() = i;
6428+
md.pid_ref() = j;
6429+
md.cid_ref() = k;
6430+
containers[md] = mmContainers[i][j][k]->saveState();
64216431
}
64226432
}
64236433
}
64246434
}
6435+
MMSerializationTypeContainer state;
6436+
state.containers_ref() = containers;
64256437
return state;
64266438
};
64276439
MMSerializationTypeContainer mmContainersState =
64286440
serializeMMContainers(mmContainers_);
64296441

64306442
AccessSerializationType accessContainerState = accessContainer_->saveState();
6431-
// TODO: foreach allocator
6432-
MemoryAllocator::SerializationType allocatorState = allocator_[0]->saveState();
6443+
6444+
auto serializeAllocators = [numTiers,this]() {
6445+
AllocatorsSerializationType state;
6446+
std::map<int,MemoryAllocator::SerializationType> allocators;
6447+
for (int i = 0; i < numTiers; ++i) {
6448+
allocators[i] = allocator_[i]->saveState();
6449+
}
6450+
state.allocators_ref() = allocators;
6451+
return state;
6452+
};
6453+
AllocatorsSerializationType allocatorsState = serializeAllocators();
6454+
64336455
CCacheManager::SerializationType ccState = compactCacheManager_->saveState();
64346456

64356457
AccessSerializationType chainedItemAccessContainerState =
@@ -6439,7 +6461,7 @@ folly::IOBufQueue CacheAllocator<CacheTrait>::saveStateToIOBuf() {
64396461
// results into a single buffer.
64406462
folly::IOBufQueue queue;
64416463
Serializer::serializeToIOBufQueue(queue, metadata_);
6442-
Serializer::serializeToIOBufQueue(queue, allocatorState);
6464+
Serializer::serializeToIOBufQueue(queue, allocatorsState);
64436465
Serializer::serializeToIOBufQueue(queue, ccState);
64446466
Serializer::serializeToIOBufQueue(queue, mmContainersState);
64456467
Serializer::serializeToIOBufQueue(queue, accessContainerState);
@@ -6562,23 +6584,22 @@ CacheAllocator<CacheTrait>::deserializeMMContainers(
65626584
* only works for a single (topmost) tier. */
65636585
MMContainers mmContainers{getNumTiers()};
65646586

6565-
for (auto& kvPool : *container.pools_ref()) {
6566-
auto i = static_cast<PoolId>(kvPool.first);
6567-
auto& pool = getPool(i);
6568-
for (auto& kv : kvPool.second) {
6569-
auto j = static_cast<ClassId>(kv.first);
6570-
for (TierId tid = 0; tid < getNumTiers(); tid++) {
6571-
MMContainerPtr ptr =
6572-
std::make_unique<typename MMContainerPtr::element_type>(kv.second,
6573-
compressor);
6574-
auto config = ptr->getConfig();
6575-
config.addExtraConfig(config_.trackTailHits
6576-
? pool.getAllocationClass(j).getAllocsPerSlab()
6577-
: 0);
6578-
ptr->setConfig(config);
6579-
mmContainers[tid][i][j] = std::move(ptr);
6580-
}
6581-
}
6587+
std::map<serialization::MemoryDescriptorObject,MMSerializationType> containerMap =
6588+
*container.containers();
6589+
for (auto md : containerMap) {
6590+
uint32_t tid = *md.first.tid();
6591+
uint32_t pid = *md.first.pid();
6592+
uint32_t cid = *md.first.cid();
6593+
auto& pool = getPoolByTid(pid,tid);
6594+
MMContainerPtr ptr =
6595+
std::make_unique<typename MMContainerPtr::element_type>(md.second,
6596+
compressor);
6597+
auto config = ptr->getConfig();
6598+
config.addExtraConfig(config_.trackTailHits
6599+
? pool.getAllocationClass(cid).getAllocsPerSlab()
6600+
: 0);
6601+
ptr->setConfig(config);
6602+
mmContainers[tid][pid][cid] = std::move(ptr);
65826603
}
65836604
// We need to drop the unevictableMMContainer in the desierializer.
65846605
// TODO: remove this at version 17.

‎cachelib/allocator/serialize/objects.thrift

+14-3
Original file line numberDiff line numberDiff line change
@@ -17,11 +17,22 @@
1717
namespace cpp2 facebook.cachelib.serialization
1818

1919
include "cachelib/allocator/datastruct/serialize/objects.thrift"
20+
include "cachelib/allocator/memory/serialize/objects.thrift"
2021

2122
// Adding a new "required" field will cause the cache to be dropped
2223
// in the next release for our users. If the field needs to be required,
2324
// make sure to communicate that with our users.
2425

26+
struct MemoryAllocatorCollection {
27+
1: required map<i32, MemoryAllocatorObject> allocators;
28+
}
29+
30+
struct MemoryDescriptorObject {
31+
1: required i32 tid;
32+
2: required i32 pid;
33+
3: required i32 cid;
34+
}
35+
2536
struct CacheAllocatorMetadata {
2637
1: required i64 allocatorVersion; // version of cache alloctor
2738
2: i64 cacheCreationTime = 0; // time when the cache was created.
@@ -80,7 +91,7 @@ struct MMLruObject {
8091
}
8192

8293
struct MMLruCollection {
83-
1: required map<i32, map<i32, MMLruObject>> pools;
94+
1: required map<MemoryDescriptorObject, MMLruObject> containers;
8495
}
8596

8697
struct MM2QConfig {
@@ -106,7 +117,7 @@ struct MM2QObject {
106117
}
107118

108119
struct MM2QCollection {
109-
1: required map<i32, map<i32, MM2QObject>> pools;
120+
1: required map<MemoryDescriptorObject, MM2QObject> containers;
110121
}
111122

112123
struct MMTinyLFUConfig {
@@ -132,7 +143,7 @@ struct MMTinyLFUObject {
132143
}
133144

134145
struct MMTinyLFUCollection {
135-
1: required map<i32, map<i32, MMTinyLFUObject>> pools;
146+
1: required map<MemoryDescriptorObject, MMTinyLFUObject> containers;
136147
}
137148

138149
struct ChainedHashTableObject {

‎cachelib/allocator/tests/AllocatorTypeTest.cpp

+4
Original file line numberDiff line numberDiff line change
@@ -117,10 +117,14 @@ TYPED_TEST(BaseAllocatorTest, DropFile) { this->testDropFile(); }
117117
TYPED_TEST(BaseAllocatorTest, ShmTemporary) { this->testShmTemporary(); }
118118

119119
TYPED_TEST(BaseAllocatorTest, Serialization) { this->testSerialization(); }
120+
TYPED_TEST(BaseAllocatorTest, MultiTierSerialization) { this->testMultiTierSerialization(); }
120121

121122
TYPED_TEST(BaseAllocatorTest, SerializationMMConfig) {
122123
this->testSerializationMMConfig();
123124
}
125+
TYPED_TEST(BaseAllocatorTest, MultiTierSerializationMMConfig) {
126+
this->testMultiTierSerializationMMConfig();
127+
}
124128

125129
TYPED_TEST(BaseAllocatorTest, testSerializationWithFragmentation) {
126130
this->testSerializationWithFragmentation();

‎cachelib/allocator/tests/BaseAllocatorTest.h

+135
Original file line numberDiff line numberDiff line change
@@ -1710,6 +1710,141 @@ class BaseAllocatorTest : public AllocatorTest<AllocatorT> {
17101710
testShmIsRemoved(config);
17111711
}
17121712

1713+
void testMultiTierSerialization() {
1714+
std::set<std::string> evictedKeys;
1715+
auto removeCb =
1716+
[&evictedKeys](const typename AllocatorT::RemoveCbData& data) {
1717+
if (data.context == RemoveContext::kEviction) {
1718+
const auto key = data.item.getKey();
1719+
evictedKeys.insert({key.data(), key.size()});
1720+
}
1721+
};
1722+
1723+
const size_t nSlabs = 40;
1724+
const size_t size = nSlabs * Slab::kSize;
1725+
const unsigned int nSizes = 1;
1726+
const unsigned int keyLen = 100;
1727+
1728+
std::vector<uint32_t> sizes;
1729+
uint8_t poolId;
1730+
1731+
// Test allocations. These allocations should remain after save/restore.
1732+
// Original lru allocator - with two tiers
1733+
typename AllocatorT::Config config;
1734+
config.setCacheSize(size);
1735+
config.enableCachePersistence(this->cacheDir_);
1736+
config.enablePoolRebalancing(nullptr, std::chrono::seconds{0});
1737+
config.configureMemoryTiers({
1738+
MemoryTierCacheConfig::fromShm()
1739+
.setRatio(1).setMemBind(std::string("0")),
1740+
MemoryTierCacheConfig::fromShm()
1741+
.setRatio(1).setMemBind(std::string("0"))});
1742+
std::vector<std::string> keys;
1743+
{
1744+
AllocatorT alloc(AllocatorT::SharedMemNew, config);
1745+
const size_t numBytes = alloc.getCacheMemoryStats().ramCacheSize;
1746+
poolId = alloc.addPool("foobar", numBytes);
1747+
sizes = this->getValidAllocSizes(alloc, poolId, nSlabs, keyLen);
1748+
this->fillUpPoolUntilEvictions(alloc, 0, poolId, sizes, keyLen);
1749+
this->fillUpPoolUntilEvictions(alloc, 1, poolId, sizes, keyLen);
1750+
for (const auto& item : alloc) {
1751+
auto key = item.getKey();
1752+
keys.push_back(key.str());
1753+
}
1754+
1755+
// save
1756+
alloc.shutDown();
1757+
}
1758+
1759+
testShmIsNotRemoved(config);
1760+
// Restored lru allocator
1761+
{
1762+
AllocatorT alloc(AllocatorT::SharedMemAttach, config);
1763+
for (auto& key : keys) {
1764+
auto handle = alloc.find(typename AllocatorT::Key{key});
1765+
ASSERT_NE(nullptr, handle.get());
1766+
}
1767+
}
1768+
1769+
testShmIsRemoved(config);
1770+
// Test LRU eviction and length before and after save/restore
1771+
// Original lru allocator
1772+
typename AllocatorT::Config config2;
1773+
config2.setCacheSize(size);
1774+
config2.setRemoveCallback(removeCb);
1775+
config2.enableCachePersistence(this->cacheDir_);
1776+
config2.configureMemoryTiers({
1777+
MemoryTierCacheConfig::fromShm()
1778+
.setRatio(1).setMemBind(std::string("0")),
1779+
MemoryTierCacheConfig::fromShm()
1780+
.setRatio(1).setMemBind(std::string("0"))});
1781+
{
1782+
AllocatorT alloc(AllocatorT::SharedMemNew, config2);
1783+
const size_t numBytes = alloc.getCacheMemoryStats().ramCacheSize;
1784+
poolId = alloc.addPool("foobar", numBytes);
1785+
1786+
sizes = this->getValidAllocSizes(alloc, poolId, nSizes, keyLen);
1787+
1788+
this->testLruLength(alloc, poolId, sizes, keyLen, evictedKeys);
1789+
1790+
// save
1791+
alloc.shutDown();
1792+
}
1793+
evictedKeys.clear();
1794+
1795+
testShmIsNotRemoved(config2);
1796+
// Restored lru allocator
1797+
{
1798+
AllocatorT alloc(AllocatorT::SharedMemAttach, config2);
1799+
this->testLruLength(alloc, poolId, sizes, keyLen, evictedKeys);
1800+
}
1801+
1802+
testShmIsRemoved(config2);
1803+
}
1804+
1805+
void testMultiTierSerializationMMConfig() {
1806+
typename AllocatorT::Config config;
1807+
config.setCacheSize(20 * Slab::kSize);
1808+
config.enableCachePersistence(this->cacheDir_);
1809+
config.enablePoolRebalancing(nullptr, std::chrono::seconds{0});
1810+
config.configureMemoryTiers({
1811+
MemoryTierCacheConfig::fromShm()
1812+
.setRatio(1).setMemBind(std::string("0")),
1813+
MemoryTierCacheConfig::fromShm()
1814+
.setRatio(1).setMemBind(std::string("0"))});
1815+
double ratio = 0.2;
1816+
1817+
// start allocator
1818+
{
1819+
AllocatorT alloc(AllocatorT::SharedMemNew, config);
1820+
const size_t numBytes = alloc.getCacheMemoryStats().ramCacheSize;
1821+
{
1822+
typename AllocatorT::MMConfig mmConfig;
1823+
mmConfig.lruRefreshRatio = ratio;
1824+
auto pid =
1825+
alloc.addPool("foobar", numBytes, /* allocSizes = */ {}, mmConfig);
1826+
auto handle = util::allocateAccessible(alloc, pid, "key", 10);
1827+
ASSERT_NE(nullptr, handle);
1828+
auto& container = alloc.getMMContainer(*handle);
1829+
EXPECT_DOUBLE_EQ(ratio, container.getConfig().lruRefreshRatio);
1830+
}
1831+
1832+
// save
1833+
alloc.shutDown();
1834+
}
1835+
testShmIsNotRemoved(config);
1836+
1837+
// restore allocator and check lruRefreshRatio
1838+
{
1839+
AllocatorT alloc(AllocatorT::SharedMemAttach, config);
1840+
auto handle = alloc.find("key");
1841+
ASSERT_NE(nullptr, handle);
1842+
auto& container = alloc.getMMContainer(*handle);
1843+
EXPECT_DOUBLE_EQ(ratio, container.getConfig().lruRefreshRatio);
1844+
}
1845+
testShmIsRemoved(config);
1846+
}
1847+
17131848
// Test temporary shared memory mode which is enabled when memory
17141849
// monitoring is enabled.
17151850
void testShmTemporary() {

0 commit comments

Comments
 (0)