Mooncake/mooncake-store/tests/health_check_test.cpp

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

#include <chrono>
#include <memory>
#include <string>
#include <thread>

#include <csignal>
#include <ylt/coro_http/coro_http_client.hpp>

#include "client_service.h"
#include "real_client.h"
#include "test_server_helpers.h"
#include "default_config.h"

DEFINE_string(protocol, "tcp", "Transfer protocol: rdma|tcp");
DEFINE_string(device_name, "", "Device name to use, valid if protocol=rdma");

DECLARE_bool(enable_http_server);
DECLARE_int32(http_port);

namespace mooncake {
namespace testing {

class HealthCheckTest : public ::testing::Test {
   protected:
    static void SetUpTestSuite() {
        google::InitGoogleLogging("HealthCheckTest");
        FLAGS_logtostderr = 1;
        FLAGS_enable_http_server = true;
        if (getenv("PROTOCOL")) FLAGS_protocol = getenv("PROTOCOL");
        if (getenv("DEVICE_NAME")) FLAGS_device_name = getenv("DEVICE_NAME");
    }

    static void TearDownTestSuite() { google::ShutdownGoogleLogging(); }

    std::shared_ptr<RealClient> py_client_;
    InProcMaster master_;
    std::string master_address_;

    struct HealthResponse {
        int http_status;
        std::string body;
    };

    HealthResponse fetch_health(int port) {
        coro_http::coro_http_client client;
        std::string url =
            "http://127.0.0.1:" + std::to_string(port) + "/health";
        auto result = client.get(url);
        return {result.status, std::string(result.resp_body)};
    }

    struct HttpResponse {
        int http_status;
        std::string body;
    };

    HttpResponse fetch_url(int port, const std::string &path) {
        coro_http::coro_http_client client;
        std::string url = "http://127.0.0.1:" + std::to_string(port) + path;
        auto result = client.get(url);
        return {result.status, std::string(result.resp_body)};
    }

    void SetUp() override { py_client_ = RealClient::create(); }

    // Start master and set up the client on the given port.
    // Returns 0 on success.
    int StartMasterAndSetupClient(int port) {
        if (!master_.Start(InProcMasterConfigBuilder().build())) return -1;
        master_address_ = master_.master_address();

        const std::string rdma_devices =
            (FLAGS_protocol == "rdma") ? FLAGS_device_name : "";
        return py_client_->setup_real("localhost:" + std::to_string(port),
                                      "P2PHANDSHAKE", 16 * 1024 * 1024,
                                      16 * 1024 * 1024, FLAGS_protocol,
                                      rdma_devices, master_address_);
    }

    bool WaitForHealthCode(
        int expected_code,
        std::chrono::milliseconds timeout = std::chrono::seconds(10),
        std::chrono::milliseconds interval = std::chrono::milliseconds(100)) {
        const auto deadline = std::chrono::steady_clock::now() + timeout;
        while (std::chrono::steady_clock::now() < deadline) {
            if (py_client_->health_check() == expected_code) {
                return true;
            }
            std::this_thread::sleep_for(interval);
        }
        return py_client_->health_check() == expected_code;
    }
};

// Test 1: health_check returns HC_NOT_INITIALIZED before setup
TEST_F(HealthCheckTest, ReturnsOneBeforeSetup) {
    EXPECT_EQ(py_client_->health_check(), HC_NOT_INITIALIZED);
}

// Test 2: health_check returns HC_HEALTHY after successful setup with master
TEST_F(HealthCheckTest, ReturnsZeroWhenHealthy) {
    ASSERT_EQ(StartMasterAndSetupClient(18900), 0) << "setup_real failed";

    // Wait for ping thread to complete at least one cycle
    std::this_thread::sleep_for(std::chrono::seconds(2));

    EXPECT_EQ(py_client_->health_check(), HC_HEALTHY);

    py_client_->tearDownAll();
    master_.Stop();
}

// Test 3: health_check returns HC_NOT_INITIALIZED after teardown
TEST_F(HealthCheckTest, ReturnsOneAfterTeardown) {
    ASSERT_EQ(StartMasterAndSetupClient(18901), 0);

    py_client_->tearDownAll();
    EXPECT_EQ(py_client_->health_check(), HC_NOT_INITIALIZED);

    master_.Stop();
}

// Test 4: health_check returns HC_MASTER_UNREACHABLE after master stops
TEST_F(HealthCheckTest, ReturnsTwoWhenMasterDown) {
    ASSERT_EQ(StartMasterAndSetupClient(18902), 0);

    // Wait for ping to succeed
    std::this_thread::sleep_for(std::chrono::seconds(2));
    EXPECT_EQ(py_client_->health_check(), HC_HEALTHY);

    // Stop master, wait for ping to fail
    master_.Stop();
    EXPECT_TRUE(WaitForHealthCode(HC_MASTER_UNREACHABLE))
        << "Timed out waiting for HC_MASTER_UNREACHABLE";

    py_client_->tearDownAll();
}

// Test 5: HTTP /health returns 200 when healthy
TEST_F(HealthCheckTest, HttpReturns200WhenHealthy) {
    int http_port = getFreeTcpPort();
    FLAGS_http_port = http_port;
    ASSERT_EQ(StartMasterAndSetupClient(18910), 0) << "setup_real failed";

    std::this_thread::sleep_for(std::chrono::seconds(2));

    auto resp = fetch_health(http_port);
    EXPECT_EQ(resp.http_status, 200);
    EXPECT_NE(resp.body.find("\"healthy\""), std::string::npos);
    EXPECT_NE(resp.body.find("\"code\":0"), std::string::npos);

    py_client_->tearDownAll();
    master_.Stop();
}

// Test 6: HTTP /health returns 503 when master unreachable
TEST_F(HealthCheckTest, HttpReturns503WhenMasterDown) {
    int http_port = getFreeTcpPort();
    FLAGS_http_port = http_port;
    ASSERT_EQ(StartMasterAndSetupClient(18911), 0);

    std::this_thread::sleep_for(std::chrono::seconds(2));
    EXPECT_EQ(py_client_->health_check(), HC_HEALTHY);

    master_.Stop();
    ASSERT_TRUE(WaitForHealthCode(HC_MASTER_UNREACHABLE))
        << "Timed out waiting for HC_MASTER_UNREACHABLE";

    auto resp = fetch_health(http_port);
    EXPECT_EQ(resp.http_status, 503);
    EXPECT_NE(resp.body.find("\"master_unreachable\""), std::string::npos);
    EXPECT_NE(resp.body.find("\"code\":2"), std::string::npos);

    py_client_->tearDownAll();
}

// Test 7: HTTP /metrics and /metrics/summary return 200 and reflect
// transfer stats after put/get operations
TEST_F(HealthCheckTest, MetricsEndpointsReturnCorrectData) {
    int http_port = getFreeTcpPort();
    FLAGS_http_port = http_port;
    ASSERT_EQ(StartMasterAndSetupClient(18920), 0) << "setup_real failed";

    std::this_thread::sleep_for(std::chrono::seconds(2));

    // Phase 1: Verify endpoints return 200 before any transfers
    auto metrics_resp = fetch_url(http_port, "/metrics");
    EXPECT_EQ(metrics_resp.http_status, 200);
    EXPECT_EQ(metrics_resp.body.find("metrics not available"),
              std::string::npos);

    auto summary_resp = fetch_url(http_port, "/metrics/summary");
    EXPECT_EQ(summary_resp.http_status, 200);
    EXPECT_NE(summary_resp.body.find("Transfer Metrics Summary"),
              std::string::npos);

    // Phase 2: Perform put/get and verify metrics reflect the operations
    const std::string test_data(1024, 'A');  // 1KB of data
    const std::string key = "metrics_test_key";
    std::span<const char> data_span(test_data.data(), test_data.size());
    ReplicateConfig config;
    config.replica_num = 1;
    ASSERT_EQ(py_client_->put(key, data_span, config), 0) << "put failed";

    auto buffer_handle = py_client_->get_buffer(key);
    ASSERT_NE(buffer_handle, nullptr) << "get_buffer failed";
    EXPECT_EQ(buffer_handle->size(), test_data.size());

    // Verify /metrics now reports transfer metrics
    auto resp = fetch_url(http_port, "/metrics");
    LOG(INFO) << "=== /metrics output ===\n" << resp.body;
    LOG(INFO) << "=== end /metrics ===";
    EXPECT_EQ(resp.http_status, 200);
    EXPECT_NE(resp.body.find("mooncake_transfer_write_bytes"),
              std::string::npos)
        << "write_bytes metric missing";
    EXPECT_NE(resp.body.find("mooncake_transfer_read_bytes"), std::string::npos)
        << "read_bytes metric missing";
    EXPECT_NE(resp.body.find("mooncake_transfer_put_latency_count"),
              std::string::npos)
        << "put_latency histogram missing";
    EXPECT_NE(resp.body.find("mooncake_transfer_get_latency_count"),
              std::string::npos)
        << "get_latency histogram missing";

    // Verify /metrics/summary shows Put/Get data
    summary_resp = fetch_url(http_port, "/metrics/summary");
    LOG(INFO) << "=== /metrics/summary output ===\n" << summary_resp.body;
    LOG(INFO) << "=== end /metrics/summary ===";
    EXPECT_EQ(summary_resp.http_status, 200);
    EXPECT_NE(summary_resp.body.find("Put:"), std::string::npos)
        << "Put summary missing";
    EXPECT_NE(summary_resp.body.find("Get:"), std::string::npos)
        << "Get summary missing";

    py_client_->tearDownAll();
    master_.Stop();
}

}  // namespace testing
}  // namespace mooncake