forked from mooncake-track/Mooncake
112 lines
3.4 KiB
C++
112 lines
3.4 KiB
C++
// Copyright 2024 KVCache.AI
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
#include "tier/memory_tier.h"
|
|
|
|
#include <gtest/gtest.h>
|
|
|
|
#include <cstdint>
|
|
#include <cstring>
|
|
#include <vector>
|
|
|
|
using namespace mooncake::store;
|
|
|
|
namespace {
|
|
|
|
constexpr uint64_t BS = 4096;
|
|
|
|
void fill(std::vector<uint8_t> &b, uint64_t id) {
|
|
for (size_t i = 0; i < b.size(); ++i) b[i] = (uint8_t)(id * 7 + i);
|
|
std::memcpy(b.data(), &id, sizeof(id));
|
|
}
|
|
bool verify(const std::vector<uint8_t> &b, uint64_t id) {
|
|
std::vector<uint8_t> ref(b.size());
|
|
fill(ref, id);
|
|
return std::memcmp(b.data(), ref.data(), b.size()) == 0;
|
|
}
|
|
|
|
std::string tmpPath(const char *name) {
|
|
const char *dir = std::getenv("TMPDIR");
|
|
return std::string(dir ? dir : "/tmp") + "/mooncake_" + name;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
TEST(MemoryTier, DramRoundtripAndCapacity) {
|
|
auto dram = makeDramTier(BS, 4);
|
|
ASSERT_TRUE(dram->ready());
|
|
std::vector<uint8_t> b(BS), o(BS);
|
|
for (uint64_t i = 0; i < 4; ++i) {
|
|
fill(b, i);
|
|
EXPECT_TRUE(dram->put(i, b.data()));
|
|
}
|
|
fill(b, 99);
|
|
EXPECT_FALSE(dram->put(99, b.data())) << "put beyond capacity must fail";
|
|
for (uint64_t i = 0; i < 4; ++i) {
|
|
EXPECT_TRUE(dram->get(i, o.data()));
|
|
EXPECT_TRUE(verify(o, i));
|
|
}
|
|
EXPECT_FALSE(dram->get(99, o.data()));
|
|
}
|
|
|
|
TEST(MemoryTier, CxlEmulatedAndNvmeBackends) {
|
|
auto cxl = makeCxlTier(BS, 4);
|
|
ASSERT_TRUE(cxl->ready());
|
|
EXPECT_EQ(cxl->kind(), TierKind::CXL);
|
|
EXPECT_TRUE(cxl->emulated()); // no /dev/dax on most hosts
|
|
|
|
auto nvme = makeNvmeTier(BS, 4, tmpPath("tier_nvme.bin"));
|
|
ASSERT_TRUE(nvme->ready());
|
|
std::vector<uint8_t> b(BS), o(BS);
|
|
fill(b, 7);
|
|
EXPECT_TRUE(nvme->put(7, b.data()));
|
|
EXPECT_TRUE(nvme->get(7, o.data()));
|
|
EXPECT_TRUE(verify(o, 7));
|
|
}
|
|
|
|
TEST(TieredStore, IntegrityPromotionSingleCopy) {
|
|
const uint64_t W = 64;
|
|
std::vector<std::shared_ptr<MemoryTier>> tiers = {
|
|
makeDramTier(BS, 8), makeCxlTier(BS, 16),
|
|
makeNvmeTier(BS, W, tmpPath("tier_store.bin"))};
|
|
TieredStore store(std::move(tiers), /*promote_on_hit=*/true);
|
|
|
|
std::vector<uint8_t> b(BS), o(BS);
|
|
for (uint64_t i = 0; i < W; ++i) {
|
|
fill(b, i);
|
|
store.put(i, b.data());
|
|
}
|
|
|
|
bool all_ok = true;
|
|
for (uint64_t i = 0; i < W; ++i)
|
|
if (!store.get(i, o.data()) || !verify(o, i)) all_ok = false;
|
|
EXPECT_TRUE(all_ok) << "all blocks must read back byte-exact across tiers";
|
|
|
|
// A freshly read cold block should be promoted into DRAM (tier 0).
|
|
uint64_t cold = 3;
|
|
store.get(cold, o.data());
|
|
EXPECT_TRUE(store.tiers()[0]->contains(cold));
|
|
|
|
// Single-copy invariant: each live id lives in exactly one tier.
|
|
bool single = true;
|
|
for (uint64_t i = 0; i < W; ++i) {
|
|
int count = 0;
|
|
for (auto &t : store.tiers())
|
|
if (t->contains(i)) ++count;
|
|
if (count != 1) single = false;
|
|
}
|
|
EXPECT_TRUE(single);
|
|
EXPECT_GT(store.stats().promotions, 0u);
|
|
}
|