Skip to content

Multi tier shm attach #99

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jun 13, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
77 changes: 49 additions & 28 deletions cachelib/allocator/CacheAllocator.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -2153,7 +2154,7 @@ class CacheAllocator : public CacheBase {
PrivateSegmentOpts createPrivateSegmentOpts(TierId tid);
std::unique_ptr<MemoryAllocator> createPrivateAllocator(TierId tid);
std::unique_ptr<MemoryAllocator> createNewMemoryAllocator(TierId tid);
std::unique_ptr<MemoryAllocator> restoreMemoryAllocator(TierId tid);
std::unique_ptr<MemoryAllocator> restoreMemoryAllocator(TierId tid, const serialization::MemoryAllocatorObject& sAllocator);
std::unique_ptr<CCacheManager> restoreCCacheManager(TierId tid);

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

template <typename CacheTrait>
std::unique_ptr<MemoryAllocator>
CacheAllocator<CacheTrait>::restoreMemoryAllocator(TierId tid) {
CacheAllocator<CacheTrait>::restoreMemoryAllocator(TierId tid,
const serialization::MemoryAllocatorObject& sAllocator) {
return std::make_unique<MemoryAllocator>(
deserializer_->deserialize<MemoryAllocator::SerializationType>(),
sAllocator,
shmManager_
->attachShm(detail::kShmCacheName + std::to_string(tid),
config_.slabMemoryBaseAddr, createShmCacheOpts(tid)).addr,
Expand Down Expand Up @@ -2732,8 +2734,11 @@ template <typename CacheTrait>
std::vector<std::unique_ptr<MemoryAllocator>>
CacheAllocator<CacheTrait>::restoreAllocators() {
std::vector<std::unique_ptr<MemoryAllocator>> allocators;
const auto allocatorCollection =
deserializer_->deserialize<AllocatorsSerializationType>();
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;
}
Expand Down Expand Up @@ -6407,26 +6412,43 @@ folly::IOBufQueue CacheAllocator<CacheTrait>::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<serialization::MemoryDescriptorObject,MMSerializationType> 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<int,MemoryAllocator::SerializationType> 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 =
Expand All @@ -6436,7 +6458,7 @@ folly::IOBufQueue CacheAllocator<CacheTrait>::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);
Expand Down Expand Up @@ -6559,23 +6581,22 @@ CacheAllocator<CacheTrait>::deserializeMMContainers(
* only works for a single (topmost) tier. */
MMContainers mmContainers{getNumTiers()};

for (auto& kvPool : *container.pools_ref()) {
auto i = static_cast<PoolId>(kvPool.first);
auto& pool = getPool(i);
for (auto& kv : kvPool.second) {
auto j = static_cast<ClassId>(kv.first);
for (TierId tid = 0; tid < getNumTiers(); tid++) {
MMContainerPtr ptr =
std::make_unique<typename MMContainerPtr::element_type>(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<serialization::MemoryDescriptorObject,MMSerializationType> 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<typename MMContainerPtr::element_type>(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.
Expand Down
17 changes: 14 additions & 3 deletions cachelib/allocator/serialize/objects.thrift
Original file line number Diff line number Diff line change
Expand Up @@ -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<i32, MemoryAllocatorObject> 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.
Expand Down Expand Up @@ -80,7 +91,7 @@ struct MMLruObject {
}

struct MMLruCollection {
1: required map<i32, map<i32, MMLruObject>> pools;
1: required map<MemoryDescriptorObject, MMLruObject> containers;
}

struct MM2QConfig {
Expand All @@ -106,7 +117,7 @@ struct MM2QObject {
}

struct MM2QCollection {
1: required map<i32, map<i32, MM2QObject>> pools;
1: required map<MemoryDescriptorObject, MM2QObject> containers;
}

struct MMTinyLFUConfig {
Expand All @@ -130,7 +141,7 @@ struct MMTinyLFUObject {
}

struct MMTinyLFUCollection {
1: required map<i32, map<i32, MMTinyLFUObject>> pools;
1: required map<MemoryDescriptorObject, MMTinyLFUObject> containers;
}

struct ChainedHashTableObject {
Expand Down
4 changes: 4 additions & 0 deletions cachelib/allocator/tests/AllocatorTypeTest.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
135 changes: 135 additions & 0 deletions cachelib/allocator/tests/BaseAllocatorTest.h
Original file line number Diff line number Diff line change
Expand Up @@ -1710,6 +1710,141 @@ class BaseAllocatorTest : public AllocatorTest<AllocatorT> {
testShmIsRemoved(config);
}

void testMultiTierSerialization() {
std::set<std::string> 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<uint32_t> 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<std::string> 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() {
Expand Down