Mooncake/mooncake-store/tests/master_service_ssd_test.cpp

#include "master_service.h"

#include <glog/logging.h>
#include <gtest/gtest.h>

#include <atomic>
#include <memory>
#include <random>
#include <thread>
#include <vector>

#include "types.h"

namespace mooncake::test {

std::unique_ptr<MasterService> CreateMasterServiceWithSSDFeat(
    const std::string& root_fs_dir) {
    return std::make_unique<MasterService>(
        false,                         // enable_gc
        DEFAULT_DEFAULT_KV_LEASE_TTL,  // default_kv_lease_ttl
        DEFAULT_KV_SOFT_PIN_TTL_MS,    // default_kv_soft_pin_ttl
        DEFAULT_ALLOW_EVICT_SOFT_PINNED_OBJECTS, DEFAULT_EVICTION_RATIO,
        DEFAULT_EVICTION_HIGH_WATERMARK_RATIO,
        0,  // view_version
        DEFAULT_CLIENT_LIVE_TTL_SEC,
        false,  // enable_ha
        DEFAULT_CLUSTER_ID, root_fs_dir, BufferAllocatorType::CACHELIB);
}

class MasterServiceSSDTest : public ::testing::Test {
   protected:
    void SetUp() override {
        google::InitGoogleLogging("MasterServiceTest");
        FLAGS_logtostderr = true;
    }

    void TearDown() override { google::ShutdownGoogleLogging(); }
};

TEST_F(MasterServiceSSDTest, PutEndBothReplica) {
    auto service_ = CreateMasterServiceWithSSDFeat("/mnt/ssd");

    constexpr size_t buffer = 0x300000000;
    constexpr size_t size = 1024 * 1024 * 64;
    std::string segment_name = "test_segment";
    Segment segment(generate_uuid(), segment_name, buffer, size);
    UUID client_id = generate_uuid();

    auto mount_result = service_->MountSegment(segment, client_id);
    ASSERT_TRUE(mount_result.has_value());

    std::string key = "disk_key";
    std::vector<uint64_t> slice_lengths = {1024};
    ReplicateConfig config;
    config.replica_num = 1;

    auto put_start_result = service_->PutStart(key, slice_lengths, config);
    ASSERT_TRUE(put_start_result.has_value());
    auto replicas = put_start_result.value();
    ASSERT_EQ(2, replicas.size());

    bool has_mem = false, has_disk = false;
    for (const auto& r : replicas) {
        if (r.is_memory_replica()) has_mem = true;
        if (r.is_disk_replica()) has_disk = true;
    }
    EXPECT_TRUE(has_mem);
    EXPECT_TRUE(has_disk);

    auto get_result = service_->GetReplicaList(key);
    ASSERT_FALSE(get_result.has_value());
    EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error());

    // PutEnd for both memory and disk
    EXPECT_TRUE(service_->PutEnd(key, ReplicaType::MEMORY).has_value());
    EXPECT_TRUE(service_->PutEnd(key, ReplicaType::DISK).has_value());

    get_result = service_->GetReplicaList(key);
    ASSERT_TRUE(get_result.has_value());
    EXPECT_EQ(2, get_result.value().size());

    for (const auto& r : get_result.value()) {
        EXPECT_EQ(ReplicaStatus::COMPLETE, r.status);
    }
}

TEST_F(MasterServiceSSDTest, PutRevokeDiskReplica) {
    auto service_ = CreateMasterServiceWithSSDFeat("/mnt/ssd");

    constexpr size_t buffer = 0x300000000;
    constexpr size_t size = 1024 * 1024 * 64;
    std::string segment_name = "test_segment";
    Segment segment(generate_uuid(), segment_name, buffer, size);
    UUID client_id = generate_uuid();

    auto mount_result = service_->MountSegment(segment, client_id);
    ASSERT_TRUE(mount_result.has_value());

    std::string key = "revoke_key";
    std::vector<uint64_t> slice_lengths = {1024};
    ReplicateConfig config;
    config.replica_num = 1;

    ASSERT_TRUE(service_->PutStart(key, slice_lengths, config).has_value());
    EXPECT_TRUE(service_->PutEnd(key, ReplicaType::MEMORY).has_value());

    auto get_result = service_->GetReplicaList(key);
    ASSERT_TRUE(get_result.has_value());
    EXPECT_EQ(1, get_result.value().size());
    ASSERT_TRUE(get_result.value()[0].is_memory_replica());

    EXPECT_TRUE(service_->PutRevoke(key, ReplicaType::DISK).has_value());

    get_result = service_->GetReplicaList(key);
    ASSERT_TRUE(get_result.has_value());
    EXPECT_EQ(1, get_result.value().size());
    ASSERT_TRUE(get_result.value()[0].is_memory_replica());
}

TEST_F(MasterServiceSSDTest, PutRevokeMemoryReplica) {
    auto service_ = CreateMasterServiceWithSSDFeat("/mnt/ssd");

    constexpr size_t buffer = 0x300000000;
    constexpr size_t size = 1024 * 1024 * 64;
    std::string segment_name = "test_segment";
    Segment segment(generate_uuid(), segment_name, buffer, size);
    UUID client_id = generate_uuid();

    auto mount_result = service_->MountSegment(segment, client_id);
    ASSERT_TRUE(mount_result.has_value());

    std::string key = "revoke_key";
    std::vector<uint64_t> slice_lengths = {1024};
    ReplicateConfig config;
    config.replica_num = 1;

    ASSERT_TRUE(service_->PutStart(key, slice_lengths, config).has_value());
    EXPECT_TRUE(service_->PutRevoke(key, ReplicaType::MEMORY).has_value());

    auto get_result = service_->GetReplicaList(key);
    ASSERT_FALSE(get_result.has_value());
    EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error());

    EXPECT_TRUE(service_->PutEnd(key, ReplicaType::DISK).has_value());
    get_result = service_->GetReplicaList(key);
    ASSERT_TRUE(get_result.has_value());
    EXPECT_EQ(1, get_result.value().size());
    ASSERT_TRUE(get_result.value()[0].is_disk_replica());
}

TEST_F(MasterServiceSSDTest, PutRevokeBothReplica) {
    auto service_ = CreateMasterServiceWithSSDFeat("/mnt/ssd");

    constexpr size_t buffer = 0x300000000;
    constexpr size_t size = 1024 * 1024 * 64;
    std::string segment_name = "test_segment";
    Segment segment(generate_uuid(), segment_name, buffer, size);
    UUID client_id = generate_uuid();

    auto mount_result = service_->MountSegment(segment, client_id);
    ASSERT_TRUE(mount_result.has_value());

    std::string key = "revoke_key";
    std::vector<uint64_t> slice_lengths = {1024};
    ReplicateConfig config;
    config.replica_num = 1;

    ASSERT_TRUE(service_->PutStart(key, slice_lengths, config).has_value());
    EXPECT_TRUE(service_->PutRevoke(key, ReplicaType::DISK).has_value());

    auto get_result = service_->GetReplicaList(key);
    ASSERT_FALSE(get_result.has_value());
    EXPECT_EQ(ErrorCode::REPLICA_IS_NOT_READY, get_result.error());

    EXPECT_TRUE(service_->PutRevoke(key, ReplicaType::MEMORY).has_value());
    get_result = service_->GetReplicaList(key);
    ASSERT_FALSE(get_result.has_value());
    EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
}

TEST_F(MasterServiceSSDTest, RemoveKey) {
    auto service_ = CreateMasterServiceWithSSDFeat("/mnt/ssd");

    constexpr size_t buffer = 0x300000000;
    constexpr size_t size = 1024 * 1024 * 64;
    std::string segment_name = "test_segment";
    Segment segment(generate_uuid(), segment_name, buffer, size);
    UUID client_id = generate_uuid();

    auto mount_result = service_->MountSegment(segment, client_id);
    ASSERT_TRUE(mount_result.has_value());

    std::string key = "remove_key";
    std::vector<uint64_t> slice_lengths = {1024};
    ReplicateConfig config;
    config.replica_num = 1;

    ASSERT_TRUE(service_->PutStart(key, slice_lengths, config).has_value());
    EXPECT_TRUE(service_->PutEnd(key, ReplicaType::MEMORY).has_value());
    EXPECT_TRUE(service_->PutEnd(key, ReplicaType::DISK).has_value());

    EXPECT_TRUE(service_->Remove(key).has_value());

    auto get_result = service_->GetReplicaList(key);
    EXPECT_FALSE(get_result.has_value());
    EXPECT_EQ(ErrorCode::OBJECT_NOT_FOUND, get_result.error());
}

TEST_F(MasterServiceSSDTest, EvictObject) {
    auto service_ = CreateMasterServiceWithSSDFeat("/mnt/ssd");
    // Mount a segment that can hold about 1024 * 16 objects.
    // As the eviction is processed separately for each shard,
    // we need to fill each shard with enough objects to thoroughly
    // test the eviction process.
    constexpr size_t buffer = 0x300000000;
    constexpr size_t size = 1024 * 1024 * 16 * 15;
    constexpr size_t object_size = 1024 * 15;
    std::string segment_name = "test_segment";
    Segment segment(generate_uuid(), segment_name, buffer, size);
    UUID client_id = generate_uuid();
    auto mount_result = service_->MountSegment(segment, client_id);
    ASSERT_TRUE(mount_result.has_value());

    // Verify if we can put objects more than the segment can hold
    int success_puts = 0;
    for (int i = 0; i < 1024 * 16 + 50; ++i) {
        std::string key = "test_key" + std::to_string(i);
        std::vector<uint64_t> slice_lengths = {object_size};
        ReplicateConfig config;
        config.replica_num = 1;
        auto put_start_result = service_->PutStart(key, slice_lengths, config);
        if (put_start_result.has_value()) {
            auto put_end_mem_result =
                service_->PutEnd(key, ReplicaType::MEMORY);
            auto put_end_disk_result = service_->PutEnd(key, ReplicaType::DISK);
            ASSERT_TRUE(put_end_mem_result.has_value());
            ASSERT_TRUE(put_end_disk_result.has_value());
            success_puts++;
        } else {
            // wait for gc thread to work
            std::this_thread::sleep_for(std::chrono::milliseconds(50));
        }
    }
    ASSERT_GT(success_puts, 1024 * 16);

    // Verify if we can get objects more than the segment can hold
    int success_gets = 0;
    for (int i = 0; i < 1024 * 16 + 50; ++i) {
        std::string key = "test_key" + std::to_string(i);
        auto get_result = service_->GetReplicaList(key);
        if (get_result.has_value()) {
            success_gets++;
        }
    }
    ASSERT_GT(success_gets, 1024 * 16);

    std::this_thread::sleep_for(
        std::chrono::milliseconds(DEFAULT_DEFAULT_KV_LEASE_TTL));
    service_->RemoveAll();
}

}  // namespace mooncake::test

int main(int argc, char** argv) {
    ::testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}