diff --git a/.gitignore b/.gitignore index 6323e02..aaf426b 100644 --- a/.gitignore +++ b/.gitignore @@ -28,3 +28,4 @@ compile_commands.json # clang cache .cache/ +docs/design/ diff --git a/docs/admin_tool.md b/docs/admin_tool.md index 5fad967..e98a10a 100644 --- a/docs/admin_tool.md +++ b/docs/admin_tool.md @@ -1,79 +1,125 @@ # SiMM Admin Control Tool -Unified admin tool to operate maintanence works on SiMM Cluster, Nodes, Shards, Global Flags, KV entry. + +Unified admin tool to operate maintenance works on SiMM Cluster, Nodes, Shards, Global Flags, Tracing, and process status. ## Binary Output + ```bash # source code -tools/simm_ctl_admin.cc +tools/simm_ctl_admin.cc # binary output path after build build/${build_mode}/bin/tools/simmctl ``` +## Two Communication Modes + +simmctl supports two modes for communicating with CM/DS processes: + +| Mode | Options | Transport | Use Case | +|------|---------|-----------|----------| +| RPC mode | `--ip`, `--port` | SiCL RPC (RDMA) | Remote admin from any host with RDMA | +| UDS mode | `--pid` or `--proc` | Unix domain socket | Local admin on the same host, no RDMA needed | + +**UDS mode** connects to the AdminServer inside CM/DS via a local Unix domain socket (`/run/simm/admin_cm..sock` or `/run/simm/admin_ds..sock`). It must run on the same host as the target process. + +`--pid` and `--proc` are mutually exclusive. When `--proc` is used, simmctl resolves the PID automatically via `pgrep`. If multiple processes match, an error is returned. ## Help Output -```bash -./simmctl -h -Usage: ./simmctl [OPTIONS] SUBCOMMAND [ARGS] + +``` +Usage: simmctl [OPTIONS] SUBCOMMAND [ARGS] OPTIONS: -SiMMCtl - SiMM RPC Management Tool: -h [ --help ] Show help message - -i [ --ip ] arg Target Cluster Manager IP address - -p [ --port ] arg (=30002) Target Cluster Manager port + -i [ --ip ] arg Target IP address for RPC-based commands + -p [ --port ] arg (=30002) Target port for RPC-based commands + -P [ --pid ] arg (=-1) Target process PID for UDS-based commands + --proc arg Target process name for UDS-based commands + (cluster_manager or data_server) -v [ --verbose ] Enable verbose output SUBCOMMANDS: node list [OPTIONS] List all nodes + node summary [OPTIONS] Show cluster-wide node resource summary + node stat Show detailed resource stats for one node node set Set node status (0=DEAD, 1=RUNNING) + cm status Query CM internal status via UDS + ds status Query DS internal status via UDS shard list [OPTIONS] List all shards gflag list [OPTIONS] List all gflags gflag get [OPTIONS] Get a gflag value gflag set Set a gflag value + trace Set tracing status (0=OFF, 1=ON) + +UDS options (--pid or --proc, mutually exclusive): + --pid Target process by PID + --proc Target process by name (cluster_manager or data_server) ``` +## Subcommand Mode Support + +| Subcommand | RPC mode | UDS mode | +|------------|----------|----------| +| `node list` | yes | yes | +| `node summary` | yes | no | +| `node stat` | yes | no | +| `node set` | yes | no | +| `cm status` | no | yes | +| `ds status` | no | yes | +| `shard list` | yes | yes | +| `gflag list` | yes | yes | +| `gflag get` | yes | yes | +| `gflag set` | yes | yes | +| `trace` | yes | yes | + ## Cheatsheet -```xml -simm-ctl -i -p - ├── node - │ ├── list [-v] - │ └── set [-v] - | - ├── shard - │ └── list [-v] - | - |── kv - │ └── exist - │ └── get - │ └── del - | - └── gflag - ├── list [-v] - ├── get [-v] - └── set [-v] -``` -## Command Examples ``` -./simmctl [OPTIONS] TYPE SUBCOMMAND +simmctl + ├── [RPC] -i -p + │ ├── node + │ │ ├── list [-v] + │ │ ├── summary [-v] + │ │ ├── stat [-v] + │ │ └── set + │ ├── shard + │ │ └── list [-v] + │ ├── gflag + │ │ ├── list [-v] + │ │ ├── get [-v] + │ │ └── set + │ └── trace <0|1> + │ + └── [UDS] --pid | --proc + ├── node + │ └── list [-v] + ├── cm + │ └── status + ├── ds + │ └── status + ├── shard + │ └── list [-v] + ├── gflag + │ ├── list [-v] + │ ├── get [-v] + │ └── set + └── trace <0|1> ``` -### OPTIONS -- -i / --ip (Required): Specify ClusterManager IP -- -p / --port（Optional, default is 30002): Specify ClusterManager Port to receive RPC requests -- -v / --verbose （Optional): Enable verbose output +## Command Examples ### Node Subcommand -#### List DataServers list of cluster + +#### List all nodes (RPC mode) + ```bash -./simmctl -i [IP] node list [-v] +simmctl --ip=172.168.1.1 node list ``` -example -```bash -./simmctl --ip=172.168.1.1 node list + +output: + ``` -output -```bash +--------------------+------------+ | Node Address | Status | +--------------------+------------+ @@ -83,32 +129,89 @@ output +--------------------+------------+ ``` -#### Set status of single DataServer +#### List all nodes (UDS mode) + ```bash -./simmctl -i [IP] node set [IP:PORT] [STATUS] +# By PID +simmctl --pid 12345 node list + +# By process name +simmctl --proc cluster_manager node list ``` -**STATUS** types include ***RUNNING***, ***DEAD*** -⚠️ **CAUTION**: Set DataServer Status to ***DEAD*** may trigger background shards migration! +#### Show cluster resource summary (RPC mode) -example ```bash -# notify ClusterManager to add DataServer(172.18.11.43) into blacklist +simmctl --ip=172.168.1.1 node summary +``` + +#### Show single node resource stats (RPC mode) -./simmctl --ip=172.168.1.1 node set 172.18.11.43:40002 DEAD +```bash +simmctl --ip=172.168.1.1 node stat 172.18.77.65:40000 ``` -### Shard Subcommand -#### Query cluster shards distribution +#### Set status of single DataServer (RPC mode) + +```bash +simmctl --ip=172.168.1.1 node set 172.18.11.43:40000 DEAD +``` + +**STATUS** types: **RUNNING**, **DEAD** + +⚠️ **CAUTION**: Setting a DataServer status to **DEAD** may trigger background shard migration! + +### CM / DS Status Subcommand (UDS only) + +#### Query CM internal status + ```bash -./simmctl -i [IP] shard list [-v] +# By PID +simmctl --pid 12345 cm status + +# By process name +simmctl --proc cluster_manager cm status ``` -example + +#### Query DS internal status + ```bash -./simmctl --ip=172.168.1.1 shard list +# By PID +simmctl --pid 67890 ds status + +# By process name +simmctl --proc data_server ds status +``` + +output: + +``` ++----------------------------+-------+ +| Field | Value | ++----------------------------+-------+ +| is_registered | true | +| cm_ready | true | +| heartbeat_failure_count | 0 | ++----------------------------+-------+ ``` -output + +| Field | Meaning | +|-------|---------| +| `is_registered` | DS has completed handshake with CM | +| `cm_ready` | DS considers CM reachable (false when consecutive HB failures reach tolerance) | +| `heartbeat_failure_count` | Consecutive heartbeat failure counter (resets on success or re-registration) | + +### Shard Subcommand + +#### Query cluster shard distribution (RPC mode) + ```bash +simmctl --ip=172.168.1.1 shard list +``` + +output: + +``` +--------------------+---------------+ | Data Server | Shard Count | +--------------------+---------------+ @@ -117,84 +220,106 @@ output | 172.18.77.65:40000 | 8192 | +--------------------+---------------+ ``` -The command output means cluster has two data servers, and one half of shards loaded on 172.18.11.43 and another half of shards loaded on 172.18.77.65 -Use -v / --verbose to see complete shards distribution list (one shard per line) -### Global Flag Subcommand -#### Get complete global flags of single node +Use `-v` / `--verbose` to see complete shard distribution list (one shard per line). + +#### Query cluster shard distribution (UDS mode) + ```bash -./simmctl -i [IP] gflag list [-v] +# By PID +simmctl --pid 12345 shard list + +# By process name +simmctl --proc cluster_manager shard list ``` -example + +### Global Flag Subcommand + +#### List all gflags (RPC mode) + ```bash -./simmctl --ip=172.168.1.1 --port=30002 gflag list +simmctl --ip=172.168.1.1 --port=30002 gflag list +``` + +output: + ``` -output -```bash +------------------------------+------------------------------+ | Flag Name | VALUE | +------------------------------+------------------------------+ | alsologtoemail | | +------------------------------+------------------------------+ ... -+------------------------------+------------------------------+ -| folly_hazptr_use_executor | true | -+------------------------------+------------------------------+ ``` -Use -v/--verbose to display Key/Value/Default/Type/Description of global flags -#### Get single global flag of single node +Use `-v` / `--verbose` to display Key/Value/Default/Type/Description. + +#### List all gflags (UDS mode) + ```bash -/simmctl -i [IP] gflag get [FLAG] [-v] +# By PID — query CM's gflags +simmctl --pid 12345 gflag list + +# By process name — query DS's gflags +simmctl --proc data_server gflag list ``` -example + +#### Get a single gflag + ```bash -# get CM heartbeat timeout(with DataServer) flag value +# RPC mode +simmctl --ip=172.168.1.1 --port=30002 gflag get cm_heartbeat_timeout_inSecs + +# UDS mode — by PID +simmctl --pid 12345 gflag get cm_heartbeat_timeout_inSecs + +# UDS mode — by process name +simmctl --proc cluster_manager gflag get cm_heartbeat_timeout_inSecs +``` + +output: -./simmctl --ip=172.168.1.1 --port=30002 gflag get cm_heartbeat_timeout_inSecs ``` -output -```bash +--------------------+--------------------------------------------------+ | Flag Name | cm_heartbeat_timeout_inSecs | +--------------------+--------------------------------------------------+ -| VALUE | 5 | +| VALUE | 30 | +--------------------+--------------------------------------------------+ ``` -#### Set single global flag value +#### Set a single gflag + ```bash -./simmctl -i [IP] gflag set [FLAG] [VALUE] +# RPC mode +simmctl --ip=172.168.1.1 --port=30002 gflag set cm_heartbeat_timeout_inSecs 10 + +# UDS mode — by PID +simmctl --pid 12345 gflag set cm_heartbeat_timeout_inSecs 10 + +# UDS mode — by process name +simmctl --proc cluster_manager gflag set cm_heartbeat_timeout_inSecs 10 ``` -example + +### Trace Subcommand + ```bash -# origin value of flag ：30 -./simmctl -i 172.168.1.1 -p 30002 gflag get cm_heartbeat_timeout_inSecs -+--------------------+--------------------------------------------------+ -| Flag Name | cm_heartbeat_timeout_inSecs | -+--------------------+--------------------------------------------------+ -| VALUE | 30 | -+--------------------+--------------------------------------------------+ +# Enable tracing (RPC mode) +simmctl --ip=172.168.1.1 --port=30002 trace 1 -# Set flag value -./simmctl -i 172.168.1.1 -p 30002 gflag set cm_heartbeat_timeout_inSecs 10 -+--------------------+--------------------------------------------------+ -| Flag Name | cm_heartbeat_timeout_inSecs | -+--------------------+--------------------------------------------------+ -| VALUE | 10 | -+--------------------+--------------------------------------------------+ +# Disable tracing (UDS mode — by PID) +simmctl --pid 12345 trace 0 -# Get flag value for double check -./simmctl -i 172.168.1.1 -p 30002 gflag get cm_heartbeat_timeout_inSecs -+--------------------+--------------------------------------------------+ -| Flag Name | cm_heartbeat_timeout_inSecs | -+--------------------+--------------------------------------------------+ -| VALUE | 10 | -+--------------------+--------------------------------------------------+ +# Disable tracing (UDS mode — by process name) +simmctl --proc data_server trace 0 ``` -### KV Subcommand -🚧 +## Limitations + +- **UDS mode is local-only**: `simmctl --pid` / `--proc` connects via Unix domain socket on the same host. For remote hosts, SSH to the target host first, or use RPC mode. +- **`--proc` requires unique process**: fails if zero or multiple processes match. Use `--pid` for disambiguation. +- **Payload size**: admin requests with payload larger than 1 MB (server-side `--admin_max_payload_bytes` flag) are rejected. +- **RPC mode requires RDMA**: `--ip` / `--port` uses SiCL RPC, which requires RDMA hardware. ## Issues -If you find any issues about simmctl tool, feel free to create issue to report it, thx! + +If you find any issues about simmctl tool, feel free to create an issue to report it! diff --git a/src/cluster_manager/cm_main.cc b/src/cluster_manager/cm_main.cc index f9dde86..202fd6a 100644 --- a/src/cluster_manager/cm_main.cc +++ b/src/cluster_manager/cm_main.cc @@ -8,6 +8,7 @@ #include #include +#include "common/admin/admin_server.h" #include "common/base/common_types.h" #include "common/errcode/errcode_def.h" #include "common/logging/logging.h" @@ -80,6 +81,12 @@ int main(int argc, char *argv[]) { // TODO(ytji): load configuration from file, e.g. cm_conf.json + auto admin_server = std::make_unique(simm::common::kCmAdminUdsBasePath); + if (admin_server == nullptr || !admin_server->isRunning()) { + MLOG_CRITICAL("Failed to init AdminServer at {}", simm::common::kCmAdminUdsBasePath); + return CmErr::InitFailed; + } + // Register signal handlers MLOG_INFO("Register signal handers for SIGINT/SIGTERM/SIGSEGV"); std::signal(SIGINT, signalHandler); @@ -109,6 +116,12 @@ int main(int argc, char *argv[]) { goto exit; } + rc = cm_service_ptr->RegisterAdminHandlers(admin_server.get()); + if (rc != CommonErr::OK) { + MLOG_CRITICAL("Failed to register CM admin handlers, rc:{}", rc); + goto exit; + } + MLOG_INFO("ClusterManager main process starts successfully!"); // Simulate as a long-running process diff --git a/src/cluster_manager/cm_service.cc b/src/cluster_manager/cm_service.cc index 6ab0daf..2665aad 100644 --- a/src/cluster_manager/cm_service.cc +++ b/src/cluster_manager/cm_service.cc @@ -6,6 +6,7 @@ #include "cm_rpc_handler.h" #include "cm_service.h" +#include "common/base/common_types.h" #include "common/logging/logging.h" #include "common/rpc_handlers/common_rpc_handlers.h" #include "proto/cm_clnt_rpcs.pb.h" @@ -200,5 +201,129 @@ error_code_t ClusterManagerService::StopRPCServices() { return CommonErr::OK; } +error_code_t ClusterManagerService::RegisterAdminHandlers( + simm::common::AdminServer* admin_server) { + if (admin_server == nullptr || !admin_server->isRunning()) { + MLOG_ERROR("RegisterAdminHandlers: AdminServer is null or not running"); + return CommonErr::InvalidState; + } + + admin_server->registerHandler( + simm::common::AdminMsgType::CM_STATUS, + [this](const std::string& /* payload */) -> std::string { + proto::common::AdmCmStatusResponsePB resp; + resp.set_ret_code(CommonErr::OK); + resp.set_is_running(is_running_.load()); + resp.set_service_ready( + simm::common::ModuleServiceState::GetInstance().IsServiceReady()); + + auto allStatus = node_manager_->GetAllNodeStatus(); + uint32_t aliveCount = 0; + uint32_t deadCount = 0; + for (const auto& [addr, status] : allStatus) { + if (status == simm::common::RUNNING) { + ++aliveCount; + } else if (status == simm::common::DEAD) { + ++deadCount; + } + } + resp.set_alive_node_count(aliveCount); + resp.set_dead_node_count(deadCount); + resp.set_total_shard_count(shard_manager_->GetTotalShardNum()); + + std::string buf; + resp.SerializeToString(&buf); + return buf; + }); + + admin_server->registerHandler( + simm::common::AdminMsgType::NODE_LIST, + [this](const std::string& /* payload */) -> std::string { + ListNodesResponsePB resp; + + if (!simm::common::ModuleServiceState::GetInstance().IsServiceReady()) { + resp.set_ret_code(CmErr::InitInGracePeriod); + std::string buf; + resp.SerializeToString(&buf); + return buf; + } + + const auto nodeStatList = node_manager_->GetAllNodeStatus(); + const auto nodeResList = node_manager_->GetAllNodeResource(); + + std::unordered_map> resMap; + for (const auto& [addrStr, resource] : nodeResList) { + resMap[addrStr] = resource; + } + + for (const auto& [addrStr, status] : nodeStatList) { + auto nodeAddr = simm::common::NodeAddress::ParseFromString(addrStr); + if (!nodeAddr) { + continue; + } + auto* nodeInfo = resp.add_nodes(); + auto* addrPb = nodeInfo->mutable_node_address(); + addrPb->set_ip(nodeAddr->node_ip_); + addrPb->set_port(nodeAddr->node_port_); + nodeInfo->set_node_status(static_cast(status)); + + auto resIt = resMap.find(addrStr); + if (resIt != resMap.end() && resIt->second) { + auto* resPb = nodeInfo->mutable_resource(); + resPb->set_mem_free_bytes(resIt->second->mem_free_bytes_); + resPb->set_mem_total_bytes(resIt->second->mem_total_bytes_); + resPb->set_mem_used_bytes(resIt->second->mem_used_bytes_); + } + } + + resp.set_ret_code(CommonErr::OK); + std::string buf; + resp.SerializeToString(&buf); + return buf; + }); + + admin_server->registerHandler( + simm::common::AdminMsgType::SHARD_LIST, + [this](const std::string& /* payload */) -> std::string { + QueryShardRoutingTableAllResponsePB resp; + + if (!simm::common::ModuleServiceState::GetInstance().IsServiceReady()) { + resp.set_ret_code(CmErr::InitInGracePeriod); + std::string buf; + resp.SerializeToString(&buf); + return buf; + } + + auto queryRes = shard_manager_->QueryAllShardRoutingInfos(); + + // Group shards by data server address + using NodeAddrPtr = std::shared_ptr; + std::unordered_map> dsShards; + for (const auto& [shardId, nodeAddr] : queryRes) { + if (nodeAddr) { + dsShards[nodeAddr].push_back(shardId); + } + } + for (const auto& [nodeAddr, shardIds] : dsShards) { + auto* entry = resp.add_shard_info(); + auto* dsAddr = entry->mutable_data_server_address(); + dsAddr->set_ip(nodeAddr->node_ip_); + dsAddr->set_port(nodeAddr->node_port_); + for (auto sid : shardIds) { + entry->add_shard_ids(sid); + } + } + + resp.set_ret_code(CommonErr::OK); + std::string buf; + resp.SerializeToString(&buf); + return buf; + }); + + MLOG_INFO("CM admin handlers registered"); + return CommonErr::OK; +} + } // namespace cm } // namespace simm diff --git a/src/cluster_manager/cm_service.h b/src/cluster_manager/cm_service.h index e1286c7..35d3718 100644 --- a/src/cluster_manager/cm_service.h +++ b/src/cluster_manager/cm_service.h @@ -10,6 +10,7 @@ #include "cm_hb_monitor.h" #include "cm_node_manager.h" #include "cm_shard_manager.h" +#include "common/admin/admin_server.h" #include "common/errcode/errcode_def.h" namespace simm { @@ -47,6 +48,10 @@ class ClusterManagerService { error_code_t Start(); error_code_t Stop(); + // Register CM-specific admin handlers to the UDS AdminServer. + // Called from cm_main after service is initialized. + error_code_t RegisterAdminHandlers(simm::common::AdminServer* admin_server); + bool IsRunning() const { return is_running_.load(); } bool IsStopped() const { return !is_running_.load(); } diff --git a/src/common/CMakeLists.txt b/src/common/CMakeLists.txt index 2b05ae9..8f566ad 100644 --- a/src/common/CMakeLists.txt +++ b/src/common/CMakeLists.txt @@ -1,3 +1,4 @@ +add_subdirectory(admin) add_subdirectory(errcode) add_subdirectory(flags) add_subdirectory(logging) @@ -9,6 +10,7 @@ add_subdirectory(hashkit) add_subdirectory(metrics) set(SIMM_COMMON_SOURCES + ${COMMON_ADMIN_SRC} ${ERRCODE_SRC} ${COMMON_FLAG_SRC} ${COMMON_LOGGING_SRC} diff --git a/src/common/admin/CMakeLists.txt b/src/common/admin/CMakeLists.txt new file mode 100644 index 0000000..c7bfa22 --- /dev/null +++ b/src/common/admin/CMakeLists.txt @@ -0,0 +1,4 @@ +# Collect admin sources to be consumed by simm_common +set(COMMON_ADMIN_SRC + "${CMAKE_CURRENT_SOURCE_DIR}/admin_server.cc" + PARENT_SCOPE) diff --git a/src/common/admin/admin_msg_types.h b/src/common/admin/admin_msg_types.h new file mode 100644 index 0000000..7d96a4b --- /dev/null +++ b/src/common/admin/admin_msg_types.h @@ -0,0 +1,28 @@ +#pragma once + +#include + +namespace simm { +namespace common { + +// UDS socket base path constants. +// Socket path = ..sock +inline constexpr const char* kCmAdminUdsBasePath = "/run/simm/admin_cm"; +inline constexpr const char* kDsAdminUdsBasePath = "/run/simm/admin_ds"; + +// Shared message types for UDS admin protocol. +// Used by AdminServer (server side) and simm_ctl_admin (client side). +// Wire format: [uint32_t frame_len][uint16_t type][payload] +enum class AdminMsgType : uint16_t { + TRACE_TOGGLE = 1, + GFLAG_LIST = 2, + GFLAG_GET = 3, + GFLAG_SET = 4, + DS_STATUS = 5, + CM_STATUS = 6, + NODE_LIST = 7, + SHARD_LIST = 8, +}; + +} // namespace common +} // namespace simm diff --git a/src/common/admin/admin_server.cc b/src/common/admin/admin_server.cc new file mode 100644 index 0000000..d409312 --- /dev/null +++ b/src/common/admin/admin_server.cc @@ -0,0 +1,417 @@ +#include "common/admin/admin_server.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include +#include + +#include "common/errcode/errcode_def.h" +#include "common/logging/logging.h" +#include "proto/common.pb.h" + +#ifdef SIMM_ENABLE_TRACE +#include "common/trace/trace.h" +#endif + +DEFINE_uint32(admin_max_payload_bytes, 1 << 20, + "Max payload size in bytes for admin UDS requests (default 1MB)"); + +DECLARE_LOG_MODULE("admin_server"); + +namespace simm { +namespace common { + +// --------------------------------------------------------------------------- +// Construction: create socket, bind, listen, spawn serve thread +// --------------------------------------------------------------------------- + +AdminServer::AdminServer(std::string basePath) + : basePath_(std::move(basePath)), listenFd_(-1), running_(false) { + // Derive directory from basePath and ensure it exists + std::string dirStr; + auto pos = basePath_.find_last_of('/'); + if (pos == std::string::npos) { + dirStr = "."; + } else if (pos == 0) { + dirStr = "/"; + } else { + dirStr = basePath_.substr(0, pos); + } + struct stat st; + if (::stat(dirStr.c_str(), &st) != 0) { + if (::mkdir(dirStr.c_str(), 0777) != 0 && errno != EEXIST) { + MLOG_ERROR("mkdir({}) failed, errno={}", dirStr, errno); + return; + } + } else if (!S_ISDIR(st.st_mode)) { + MLOG_ERROR("{} exists but is not a directory", dirStr); + return; + } + + // Socket path: ..sock + socketPath_ = basePath_ + "." + std::to_string(::getpid()) + ".sock"; + ::unlink(socketPath_.c_str()); + + // Create self-pipe for clean shutdown + if (::pipe(shutdownPipe_) < 0) { + MLOG_ERROR("pipe() failed for shutdown self-pipe, errno={}", errno); + goto err_exit; + } + // Make read end non-blocking + ::fcntl(shutdownPipe_[0], F_SETFL, O_NONBLOCK); + + // Create listen socket + listenFd_ = ::socket(AF_UNIX, SOCK_STREAM, 0); + if (listenFd_ < 0) { + MLOG_ERROR("socket(AF_UNIX) failed, errno={}", errno); + goto err_exit; + } + + { + sockaddr_un addr; + std::memset(&addr, 0, sizeof(addr)); + addr.sun_family = AF_UNIX; + std::strncpy(addr.sun_path, socketPath_.c_str(), sizeof(addr.sun_path) - 1); + addr.sun_path[sizeof(addr.sun_path) - 1] = '\0'; + + socklen_t addrLen = static_cast( + offsetof(sockaddr_un, sun_path) + std::strlen(addr.sun_path)); + + if (::bind(listenFd_, reinterpret_cast(&addr), addrLen) < 0) { + MLOG_ERROR("bind({}) failed, errno={}", socketPath_, errno); + goto err_exit; + } + } + + if (::listen(listenFd_, 16) < 0) { + MLOG_ERROR("listen({}) failed, errno={}", socketPath_, errno); + goto err_exit; + } + + // Register built-in handlers (gflag + trace, available for all SiMM components) + handlers_[static_cast(AdminMsgType::GFLAG_LIST)] = + [this](const std::string& p) { return handleGFlagList(p); }; + handlers_[static_cast(AdminMsgType::GFLAG_GET)] = + [this](const std::string& p) { return handleGFlagGet(p); }; + handlers_[static_cast(AdminMsgType::GFLAG_SET)] = + [this](const std::string& p) { return handleGFlagSet(p); }; + handlers_[static_cast(AdminMsgType::TRACE_TOGGLE)] = + [this](const std::string& p) { return handleTraceToggle(p); }; + + // Spawn serve thread + running_.store(true); + worker_ = std::thread(&AdminServer::serveLoop, this); + + MLOG_INFO("AdminServer listening on {}", socketPath_); + return; + +err_exit: + shutdown(); +} + +// --------------------------------------------------------------------------- +// Destruction: signal thread to exit, join, close fds, unlink socket +// --------------------------------------------------------------------------- + +AdminServer::~AdminServer() { + shutdown(); +} + +void AdminServer::shutdown() { + if (!running_.exchange(false)) { + // Already shut down or never started — just clean up fds + if (listenFd_ >= 0) { ::close(listenFd_); listenFd_ = -1; } + if (shutdownPipe_[0] >= 0) { ::close(shutdownPipe_[0]); shutdownPipe_[0] = -1; } + if (shutdownPipe_[1] >= 0) { ::close(shutdownPipe_[1]); shutdownPipe_[1] = -1; } + if (!socketPath_.empty()) { ::unlink(socketPath_.c_str()); } + return; + } + + // Wake the serve loop via self-pipe + if (shutdownPipe_[1] >= 0) { + char c = 'x'; + (void)::write(shutdownPipe_[1], &c, 1); + } + + // Join serve thread + if (worker_.joinable()) { + worker_.join(); + } + + // Close all fds + if (listenFd_ >= 0) { ::close(listenFd_); listenFd_ = -1; } + if (shutdownPipe_[0] >= 0) { ::close(shutdownPipe_[0]); shutdownPipe_[0] = -1; } + if (shutdownPipe_[1] >= 0) { ::close(shutdownPipe_[1]); shutdownPipe_[1] = -1; } + + // Remove socket file + if (!socketPath_.empty()) { + ::unlink(socketPath_.c_str()); + MLOG_INFO("AdminServer shut down, unlinked {}", socketPath_); + } + + handlers_.clear(); +} + +// --------------------------------------------------------------------------- +// Public API +// --------------------------------------------------------------------------- + +void AdminServer::registerHandler(AdminMsgType msg_type, Handler handler) { + std::unique_lock lock(handlersMutex_); + handlers_[static_cast(msg_type)] = std::move(handler); +} + +// --------------------------------------------------------------------------- +// Serve loop: poll on listenFd + shutdownPipe, accept clients +// --------------------------------------------------------------------------- + +void AdminServer::serveLoop() { + while (running_.load()) { + struct pollfd fds[2]; + fds[0].fd = listenFd_; + fds[0].events = POLLIN; + fds[1].fd = shutdownPipe_[0]; + fds[1].events = POLLIN; + + int ret = ::poll(fds, 2, -1); // block until event + if (ret < 0) { + if (errno == EINTR) continue; + MLOG_ERROR("poll() failed on {}, errno={}", socketPath_, errno); + break; + } + + // Check shutdown signal + if (fds[1].revents & POLLIN) { + break; // woken up by destructor + } + + // Check new client connection + if (fds[0].revents & POLLIN) { + int clientFd = ::accept(listenFd_, nullptr, nullptr); + if (clientFd < 0) { + if (errno == EINTR) continue; + if (!running_.load()) break; + MLOG_ERROR("accept() failed on {}, errno={}", socketPath_, errno); + continue; + } + SCOPE_EXIT { ::close(clientFd); }; + handleClient(clientFd); + } + } +} + +// --------------------------------------------------------------------------- +// Client handling: read frame, dispatch to handler, send response +// --------------------------------------------------------------------------- + +void AdminServer::handleClient(int clientFd) { + // Read frame: [uint32_t len][uint16_t type][payload] + uint32_t lenNet = 0; + if (!readExact(clientFd, &lenNet, sizeof(lenNet))) { + MLOG_WARN("Failed to read frame length on {}", socketPath_); + return; + } + uint32_t len = ntohl(lenNet); + if (len < sizeof(uint16_t)) { + MLOG_WARN("Invalid frame length {} on {}", len, socketPath_); + return; + } + + uint16_t typeNet = 0; + if (!readExact(clientFd, &typeNet, sizeof(typeNet))) { + MLOG_WARN("Failed to read msg type on {}", socketPath_); + return; + } + uint16_t typeRaw = ntohs(typeNet); + + uint32_t payloadLen = len - static_cast(sizeof(typeNet)); + if (payloadLen > FLAGS_admin_max_payload_bytes) { + MLOG_WARN("Payload too large: {} bytes (max {}) on {}", + payloadLen, FLAGS_admin_max_payload_bytes, socketPath_); + // Send error response with ret_code = AdmPayloadTooLarge. + // All admin response protos share sint32 ret_code as field 1. + proto::common::SetGFlagValueResponsePB errResp; + errResp.set_ret_code(CommonErr::AdmPayloadTooLarge); + std::string errBuf; + errResp.SerializeToString(&errBuf); + sendResponse(clientFd, static_cast(typeRaw), errBuf); + return; + } + std::string payload(payloadLen, '\0'); + if (payloadLen > 0 && !readExact(clientFd, payload.data(), payloadLen)) { + MLOG_WARN("Failed to read payload on {}", socketPath_); + return; + } + + // Dispatch to registered handler + std::shared_lock lock(handlersMutex_); + auto it = handlers_.find(typeRaw); + if (it != handlers_.end()) { + try { + std::string response = it->second(payload); + lock.unlock(); + sendResponse(clientFd, static_cast(typeRaw), response); + } catch (const std::exception& e) { + lock.unlock(); + MLOG_ERROR("Admin handler threw exception for MsgType {} on {}: {}", + typeRaw, socketPath_, e.what()); + // All admin response protos share sint32 ret_code as field 1. + proto::common::SetGFlagValueResponsePB errResp; + errResp.set_ret_code(CommonErr::AdmInternalError); + std::string errBuf; + errResp.SerializeToString(&errBuf); + sendResponse(clientFd, static_cast(typeRaw), errBuf); + } catch (...) { + lock.unlock(); + MLOG_ERROR("Admin handler threw unknown exception for MsgType {} on {}", + typeRaw, socketPath_); + proto::common::SetGFlagValueResponsePB errResp; + errResp.set_ret_code(CommonErr::AdmInternalError); + std::string errBuf; + errResp.SerializeToString(&errBuf); + sendResponse(clientFd, static_cast(typeRaw), errBuf); + } + } else { + lock.unlock(); + MLOG_WARN("No handler for AdminMsgType {} on {}", typeRaw, socketPath_); + } +} + +// --------------------------------------------------------------------------- +// Built-in handlers +// --------------------------------------------------------------------------- + +std::string AdminServer::handleGFlagList(const std::string& /*payload*/) { + proto::common::ListAllGFlagsResponsePB resp; + std::vector allFlags; + gflags::GetAllFlags(&allFlags); + for (const auto& f : allFlags) { + auto* rf = resp.add_flags(); + rf->set_flag_name(f.name); + rf->set_flag_value(f.current_value); + rf->set_flag_default_value(f.default_value); + rf->set_flag_type(f.type); + rf->set_flag_description(f.description); + } + resp.set_ret_code(CommonErr::OK); + std::string buf; + resp.SerializeToString(&buf); + return buf; +} + +std::string AdminServer::handleGFlagGet(const std::string& payload) { + proto::common::GetGFlagValueRequestPB req; + proto::common::GetGFlagValueResponsePB resp; + if (!req.ParseFromString(payload)) { + resp.set_ret_code(CommonErr::InvalidArgument); + } else { + gflags::CommandLineFlagInfo info; + if (!gflags::GetCommandLineFlagInfo(req.flag_name().c_str(), &info)) { + resp.set_ret_code(CommonErr::GFlagNotFound); + } else { + resp.set_ret_code(CommonErr::OK); + auto* fi = resp.mutable_flag_info(); + fi->set_flag_name(info.name); + fi->set_flag_value(info.current_value); + fi->set_flag_default_value(info.default_value); + fi->set_flag_type(info.type); + fi->set_flag_description(info.description); + } + } + std::string buf; + resp.SerializeToString(&buf); + return buf; +} + +std::string AdminServer::handleGFlagSet(const std::string& payload) { + proto::common::SetGFlagValueRequestPB req; + proto::common::SetGFlagValueResponsePB resp; + if (!req.ParseFromString(payload)) { + resp.set_ret_code(CommonErr::InvalidArgument); + } else { + gflags::CommandLineFlagInfo info; + if (!gflags::GetCommandLineFlagInfo(req.flag_name().c_str(), &info)) { + resp.set_ret_code(CommonErr::GFlagNotFound); + } else { + auto res = gflags::SetCommandLineOption( + req.flag_name().c_str(), req.flag_value().c_str()); + resp.set_ret_code(res.empty() ? CommonErr::GFlagSetFailed : CommonErr::OK); + } + } + std::string buf; + resp.SerializeToString(&buf); + return buf; +} + +std::string AdminServer::handleTraceToggle(const std::string& payload) { + proto::common::TraceToggleRequestPB req; + proto::common::TraceToggleResponsePB resp; + if (!req.ParseFromString(payload)) { + resp.set_ret_code(CommonErr::InvalidArgument); + } else { +#ifdef SIMM_ENABLE_TRACE + simm::trace::TraceManager::SetEnabled(req.enable_trace()); +#endif + resp.set_ret_code(CommonErr::OK); + } + std::string buf; + resp.SerializeToString(&buf); + return buf; +} + +// --------------------------------------------------------------------------- +// UDS I/O helpers +// --------------------------------------------------------------------------- + +bool AdminServer::readExact(int fd, void* buf, size_t len) { + char* p = static_cast(buf); + size_t remaining = len; + while (remaining > 0) { + ssize_t n = ::read(fd, p, remaining); + if (n > 0) { + remaining -= static_cast(n); + p += n; + } else if (n == 0) { + return false; + } else { + if (errno == EINTR) continue; + return false; + } + } + return true; +} + +void AdminServer::sendResponse(int clientFd, AdminMsgType type, + const std::string& serializedResp) { + uint32_t respLen = static_cast(sizeof(uint16_t) + serializedResp.size()); + uint32_t respLenNet = htonl(respLen); + uint16_t respTypeNet = htons(static_cast(type)); + + struct iovec iov[3]; + iov[0].iov_base = &respLenNet; + iov[0].iov_len = sizeof(respLenNet); + iov[1].iov_base = &respTypeNet; + iov[1].iov_len = sizeof(respTypeNet); + iov[2].iov_base = const_cast(serializedResp.data()); + iov[2].iov_len = serializedResp.size(); + + ssize_t n = ::writev(clientFd, iov, 3); + if (n < 0) { + MLOG_WARN("Failed to send admin response on {}, errno={}", socketPath_, errno); + } +} + +} // namespace common +} // namespace simm diff --git a/src/common/admin/admin_server.h b/src/common/admin/admin_server.h new file mode 100644 index 0000000..c65d54b --- /dev/null +++ b/src/common/admin/admin_server.h @@ -0,0 +1,77 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "common/admin/admin_msg_types.h" + +namespace simm { +namespace common { + +// UDS-based admin server for SiMM components. +// +// The constructor creates the Unix domain socket, binds, listens, and spawns +// the serve thread. The destructor shuts down the thread, closes the socket, +// and unlinks the socket file. No explicit start()/stop() — lifecycle is +// tied to object lifetime (RAII). +// +// Socket path: ..sock +// e.g. /run/simm/admin_ds.12345.sock +// +// Built-in handlers for GFLAG_LIST/GET/SET and TRACE_TOGGLE are always +// registered. Additional handlers (e.g. DS_STATUS) can be registered via +// registerHandler() by the owning service after construction. +// +// Wire protocol: [uint32_t frame_len][uint16_t type][payload] +class AdminServer { + public: + // basePath: e.g. "/run/simm/admin_cm" or "/run/simm/admin_ds". + // Socket path = ..sock + explicit AdminServer(std::string basePath); + ~AdminServer(); + + AdminServer(const AdminServer&) = delete; + AdminServer& operator=(const AdminServer&) = delete; + + // Handler callback: receives raw request payload, returns serialized response. + using Handler = std::function; + + // Register a custom handler for a given message type. + // Safe to call after construction and before the first client connects. + void registerHandler(AdminMsgType msg_type, Handler handler); + + const std::string& socketPath() const { return socketPath_; } + + bool isRunning() const { return running_.load(); } + + private: + void serveLoop(); + void handleClient(int client_fd); + void shutdown(); + + // Built-in handlers + std::string handleGFlagList(const std::string& payload); + std::string handleGFlagGet(const std::string& payload); + std::string handleGFlagSet(const std::string& payload); + std::string handleTraceToggle(const std::string& payload); + + // UDS I/O helpers + static bool readExact(int fd, void* buf, size_t len); + void sendResponse(int client_fd, AdminMsgType type, const std::string& serialized_resp); + + std::string basePath_; + std::string socketPath_; + int listenFd_{-1}; + int shutdownPipe_[2]{-1, -1}; // self-pipe for clean shutdown + std::atomic running_{false}; + std::thread worker_; + mutable std::shared_mutex handlersMutex_; + std::unordered_map handlers_; +}; + +} // namespace common +} // namespace simm diff --git a/src/common/errcode/errcode_def.h b/src/common/errcode/errcode_def.h index ec8d673..5f49ffb 100644 --- a/src/common/errcode/errcode_def.h +++ b/src/common/errcode/errcode_def.h @@ -42,6 +42,9 @@ DEFINE_COMMON_ERRCODE(GFlagNotFound, -1200, "Gflag name want to get/set not foun DEFINE_COMMON_ERRCODE(GFlagSetFailed, -1201, "Failed to set Gflag value to target value") // K8S related error codes DEFINE_COMMON_ERRCODE(GetPodIpFromK8SFailed, -1300, "Get pod IP from K8S failed") +// Admin related error codes +DEFINE_COMMON_ERRCODE(AdmPayloadTooLarge, -1400, "Admin request payload exceeds max allowed size") +DEFINE_COMMON_ERRCODE(AdmInternalError, -1401, "Admin handler internal error") // Other error codes (admin, internal, uknown, etc.) DEFINE_COMMON_ERRCODE(TargetNotFound, -1900, "Target resource not found") DEFINE_COMMON_ERRCODE(TargetUnavailable, -1901, "Target resource unavailable") diff --git a/src/common/trace/trace_server.cc b/src/common/trace/trace_server.cc index 8004663..99eff0b 100644 --- a/src/common/trace/trace_server.cc +++ b/src/common/trace/trace_server.cc @@ -27,7 +27,7 @@ DECLARE_LOG_MODULE("trace"); namespace simm { namespace trace { -// Align UDS message types with admin side +// XXX: Should keep in sync with common/admin/admin_msg_types.h enum class AdminMsgType : uint16_t { TRACE_TOGGLE = 1, GFLAG_LIST = 2, diff --git a/src/data_server/kv_rpc_service.cc b/src/data_server/kv_rpc_service.cc index 9083300..3deb883 100644 --- a/src/data_server/kv_rpc_service.cc +++ b/src/data_server/kv_rpc_service.cc @@ -25,7 +25,9 @@ #include "data_server/kv_hash_table.h" #include "data_server/kv_object_pool.h" #include "data_server/kv_rpc_handler.h" +#include "common/admin/admin_server.h" #include "data_server/kv_rpc_service.h" +#include "proto/common.pb.h" DECLARE_LOG_MODULE("data_server"); @@ -885,5 +887,28 @@ void KVRpcService::GetResourceStats(const DataServerResourceRequestPB *req, Data } } +error_code_t KVRpcService::RegisterAdminHandlers(simm::common::AdminServer* admin_server) { + if (admin_server == nullptr || !admin_server->isRunning()) { + MLOG_ERROR("RegisterAdminHandlers: AdminServer is null or not running"); + return CommonErr::InvalidState; + } + + admin_server->registerHandler( + simm::common::AdminMsgType::DS_STATUS, + [this](const std::string& /* payload */) -> std::string { + proto::common::AdmDsStatusResponsePB resp; + resp.set_ret_code(CommonErr::OK); + resp.set_is_registered(is_registered_.load()); + resp.set_cm_ready(cm_ready_.load()); + resp.set_heartbeat_failure_count(heartbeat_failure_count_.load()); + std::string buf; + resp.SerializeToString(&buf); + return buf; + }); + + MLOG_INFO("DS admin handlers registered"); + return CommonErr::OK; +} + } // namespace ds } // namespace simm diff --git a/src/data_server/kv_rpc_service.h b/src/data_server/kv_rpc_service.h index 3a44606..90ce066 100644 --- a/src/data_server/kv_rpc_service.h +++ b/src/data_server/kv_rpc_service.h @@ -42,6 +42,7 @@ DECLARE_uint32(busy_wait_timeout_us); DECLARE_string(ds_logical_node_id); namespace simm { +namespace common { class AdminServer; } namespace ds { class KVRpcService { @@ -80,6 +81,10 @@ class KVRpcService { // Get resource stats info void GetResourceStats(const DataServerResourceRequestPB *req, DataServerResourceResponsePB *rsp); + // Register DS-specific admin handlers to the UDS AdminServer. + // Called from kv_server_main after service is initialized. + error_code_t RegisterAdminHandlers(simm::common::AdminServer* admin_server); + private: error_code_t StartRPCServices(); error_code_t StopRPCServices(); diff --git a/src/data_server/kv_server_main.cc b/src/data_server/kv_server_main.cc index 99ca79a..2cda256 100644 --- a/src/data_server/kv_server_main.cc +++ b/src/data_server/kv_server_main.cc @@ -9,6 +9,7 @@ #include #include +#include "common/admin/admin_server.h" #include "common/errcode/errcode_def.h" #include "common/logging/logging.h" #include "common/version/version_info.h" @@ -77,6 +78,12 @@ int main(int argc, char *argv[]) { // TODO: load configuration file + auto admin_server = std::make_unique(simm::common::kDsAdminUdsBasePath); + if (admin_server == nullptr || !admin_server->isRunning()) { + MLOG_ERROR("Failed to init AdminServer at {}", simm::common::kDsAdminUdsBasePath); + return -1; + } + // Register signal handlers and save previous ones MLOG_INFO("Register signal handers..."); std::signal(SIGINT, signalHandler); @@ -102,6 +109,12 @@ int main(int argc, char *argv[]) { return -1; } + rc = server->RegisterAdminHandlers(admin_server.get()); + if (rc != CommonErr::OK) { + MLOG_ERROR("Failed to register DS admin handlers, rc:{}", rc); + return -1; + } + MLOG_INFO("DataServer starts normally ..."); // Simulate a long-running process diff --git a/src/proto/common.proto b/src/proto/common.proto index 77703a9..98884f1 100644 --- a/src/proto/common.proto +++ b/src/proto/common.proto @@ -60,4 +60,30 @@ message TraceToggleRequestPB { message TraceToggleResponsePB { sint32 ret_code = 1; +} + +// DS internal status query (for testing/debugging) +message AdmDsStatusRequestPB { + // empty — query all status fields +} + +message AdmDsStatusResponsePB { + sint32 ret_code = 1; + bool is_registered = 2; + bool cm_ready = 3; + uint32 heartbeat_failure_count = 4; +} + +// CM internal status query (for testing/debugging) +message AdmCmStatusRequestPB { + // empty — query all status fields +} + +message AdmCmStatusResponsePB { + sint32 ret_code = 1; + bool is_running = 2; + bool service_ready = 3; + uint32 alive_node_count = 4; + uint32 dead_node_count = 5; + uint32 total_shard_count = 6; } \ No newline at end of file diff --git a/tests/common/CMakeLists.txt b/tests/common/CMakeLists.txt index c0d8204..ccf55c6 100644 --- a/tests/common/CMakeLists.txt +++ b/tests/common/CMakeLists.txt @@ -1,3 +1,4 @@ +add_subdirectory(admin) add_subdirectory(base) add_subdirectory(common_flags) add_subdirectory(errcode) diff --git a/tests/common/admin/CMakeLists.txt b/tests/common/admin/CMakeLists.txt new file mode 100644 index 0000000..ef90436 --- /dev/null +++ b/tests/common/admin/CMakeLists.txt @@ -0,0 +1,20 @@ +file(GLOB TEST_SOURCES CONFIGURE_DEPENDS + "${CMAKE_CURRENT_SOURCE_DIR}/test_*.cc" +) + +foreach(test_file ${TEST_SOURCES}) + get_filename_component(test_name ${test_file} NAME_WE) + add_executable(${test_name} ${test_file}) + target_link_libraries(${test_name} PRIVATE + simm_common + common_proto + cm_clnt_proto + gtest + gtest_main + folly) + set_target_properties(${test_name} PROPERTIES + BUILD_RPATH "\$ORIGIN/../../../../third_party/sict/lib" + INSTALL_RPATH "\$ORIGIN/../../../../third_party/sict/lib" + ) + add_test(NAME ${test_name} COMMAND ${CMAKE_RUNTIME_OUTPUT_DIRECTORY}/${test_name}) +endforeach() diff --git a/tests/common/admin/test_admin_server.cc b/tests/common/admin/test_admin_server.cc new file mode 100644 index 0000000..5011c2e --- /dev/null +++ b/tests/common/admin/test_admin_server.cc @@ -0,0 +1,837 @@ +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "common/admin/admin_msg_types.h" +#include "common/admin/admin_server.h" +#include "common/errcode/errcode_def.h" +#include "proto/cm_clnt_rpcs.pb.h" +#include "proto/common.pb.h" + +// Test-only gflag for GFlag handler tests +DEFINE_string(test_admin_flag, "default_value", "test flag for AdminServer UT"); + +namespace simm { +namespace common { +namespace { + +// --------------------------------------------------------------------------- +// UDS client helper — sends a request frame, receives a response frame. +// Wire format: [uint32_t frame_len (network order)][uint16_t type][payload] +// --------------------------------------------------------------------------- + +class UdsClient { + public: + bool connect(const std::string& socketPath) { + fd_ = ::socket(AF_UNIX, SOCK_STREAM, 0); + if (fd_ < 0) return false; + + sockaddr_un addr; + std::memset(&addr, 0, sizeof(addr)); + addr.sun_family = AF_UNIX; + std::strncpy(addr.sun_path, socketPath.c_str(), sizeof(addr.sun_path) - 1); + + socklen_t len = static_cast( + offsetof(sockaddr_un, sun_path) + std::strlen(addr.sun_path)); + + if (::connect(fd_, reinterpret_cast(&addr), len) < 0) { + ::close(fd_); + fd_ = -1; + return false; + } + return true; + } + + ~UdsClient() { close(); } + + void close() { + if (fd_ >= 0) { ::close(fd_); fd_ = -1; } + } + + bool sendRequest(AdminMsgType type, const std::string& payload, + uint16_t& respType, std::string& respPayload) { + if (fd_ < 0) return false; + + // Build request frame + uint32_t frameLen = static_cast(sizeof(uint16_t) + payload.size()); + uint32_t frameLenNet = htonl(frameLen); + uint16_t typeNet = htons(static_cast(type)); + + if (!writeAll(&frameLenNet, sizeof(frameLenNet))) return false; + if (!writeAll(&typeNet, sizeof(typeNet))) return false; + if (!payload.empty() && !writeAll(payload.data(), payload.size())) return false; + + // Read response frame + uint32_t respLenNet = 0; + if (!readAll(&respLenNet, sizeof(respLenNet))) return false; + uint32_t respLen = ntohl(respLenNet); + if (respLen < sizeof(uint16_t)) return false; + + uint16_t respTypeNet = 0; + if (!readAll(&respTypeNet, sizeof(respTypeNet))) return false; + respType = ntohs(respTypeNet); + + uint32_t payloadLen = respLen - static_cast(sizeof(uint16_t)); + respPayload.resize(payloadLen); + if (payloadLen > 0 && !readAll(respPayload.data(), payloadLen)) return false; + + return true; + } + + private: + bool writeAll(const void* buf, size_t len) { + const char* p = static_cast(buf); + size_t remaining = len; + while (remaining > 0) { + ssize_t n = ::write(fd_, p, remaining); + if (n > 0) { remaining -= n; p += n; } + else if (n == 0) return false; + else { if (errno == EINTR) continue; return false; } + } + return true; + } + + bool readAll(void* buf, size_t len) { + char* p = static_cast(buf); + size_t remaining = len; + while (remaining > 0) { + ssize_t n = ::read(fd_, p, remaining); + if (n > 0) { remaining -= n; p += n; } + else if (n == 0) return false; + else { if (errno == EINTR) continue; return false; } + } + return true; + } + + int fd_{-1}; +}; + +// --------------------------------------------------------------------------- +// Test fixture +// --------------------------------------------------------------------------- + +class AdminServerTest : public ::testing::Test { + protected: + void SetUp() override { + ::mkdir("/run/simm", 0777); + } + + void TearDown() override { + server_.reset(); + } + + void createServer(const std::string& basePath = "/run/simm/admin_test") { + server_ = std::make_unique(basePath); + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + } + + UdsClient connectClient() { + UdsClient client; + EXPECT_TRUE(client.connect(server_->socketPath())); + return client; + } + + std::unique_ptr server_; +}; + +// --------------------------------------------------------------------------- +// Lifecycle tests +// --------------------------------------------------------------------------- + +TEST_F(AdminServerTest, ConstructionCreatesSocketFile) { + createServer("/run/simm/admin_lifecycle"); + EXPECT_TRUE(server_->isRunning()); + struct stat st; + EXPECT_EQ(::stat(server_->socketPath().c_str(), &st), 0); + EXPECT_TRUE(S_ISSOCK(st.st_mode)); +} + +TEST_F(AdminServerTest, SocketPathFormat) { + createServer("/run/simm/admin_ds"); + std::string expected = std::string("/run/simm/admin_ds.") + + std::to_string(::getpid()) + ".sock"; + EXPECT_EQ(server_->socketPath(), expected); +} + +TEST_F(AdminServerTest, DestructorRemovesSocketFile) { + createServer("/run/simm/admin_cleanup"); + std::string path = server_->socketPath(); + server_.reset(); + struct stat st; + EXPECT_NE(::stat(path.c_str(), &st), 0); +} + +TEST_F(AdminServerTest, MultipleServersWithDifferentBasePaths) { + auto serverA = std::make_unique("/run/simm/admin_testa"); + auto serverB = std::make_unique("/run/simm/admin_testb"); + EXPECT_NE(serverA->socketPath(), serverB->socketPath()); + + struct stat st; + EXPECT_EQ(::stat(serverA->socketPath().c_str(), &st), 0); + EXPECT_EQ(::stat(serverB->socketPath().c_str(), &st), 0); + + std::string pathA = serverA->socketPath(); + std::string pathB = serverB->socketPath(); + serverA.reset(); + serverB.reset(); + EXPECT_NE(::stat(pathA.c_str(), &st), 0); + EXPECT_NE(::stat(pathB.c_str(), &st), 0); +} + +TEST_F(AdminServerTest, IsRunningReflectsState) { + createServer("/run/simm/admin_running"); + EXPECT_TRUE(server_->isRunning()); + server_.reset(); + // After destruction, no server to check — just verify no crash +} + +// --------------------------------------------------------------------------- +// GFlag handler tests +// --------------------------------------------------------------------------- + +TEST_F(AdminServerTest, GFlagListReturnsFlags) { + createServer("/run/simm/admin_gflag_list"); + auto client = connectClient(); + + uint16_t respType = 0; + std::string respPayload; + ASSERT_TRUE(client.sendRequest(AdminMsgType::GFLAG_LIST, "", respType, respPayload)); + EXPECT_EQ(respType, static_cast(AdminMsgType::GFLAG_LIST)); + + proto::common::ListAllGFlagsResponsePB resp; + ASSERT_TRUE(resp.ParseFromString(respPayload)); + EXPECT_EQ(resp.ret_code(), CommonErr::OK); + EXPECT_GT(resp.flags_size(), 0); +} + +TEST_F(AdminServerTest, GFlagGetKnownFlag) { + createServer("/run/simm/admin_gflag_get"); + + proto::common::GetGFlagValueRequestPB req; + req.set_flag_name("test_admin_flag"); + std::string reqBuf; + req.SerializeToString(&reqBuf); + + auto client = connectClient(); + uint16_t respType = 0; + std::string respPayload; + ASSERT_TRUE(client.sendRequest(AdminMsgType::GFLAG_GET, reqBuf, respType, respPayload)); + + proto::common::GetGFlagValueResponsePB resp; + ASSERT_TRUE(resp.ParseFromString(respPayload)); + EXPECT_EQ(resp.ret_code(), CommonErr::OK); + EXPECT_EQ(resp.flag_info().flag_name(), "test_admin_flag"); +} + +TEST_F(AdminServerTest, GFlagGetUnknownFlag) { + createServer("/run/simm/admin_gflag_get_unk"); + + proto::common::GetGFlagValueRequestPB req; + req.set_flag_name("nonexistent_flag_12345"); + std::string reqBuf; + req.SerializeToString(&reqBuf); + + auto client = connectClient(); + uint16_t respType = 0; + std::string respPayload; + ASSERT_TRUE(client.sendRequest(AdminMsgType::GFLAG_GET, reqBuf, respType, respPayload)); + + proto::common::GetGFlagValueResponsePB resp; + ASSERT_TRUE(resp.ParseFromString(respPayload)); + EXPECT_EQ(resp.ret_code(), CommonErr::GFlagNotFound); +} + +TEST_F(AdminServerTest, GFlagSetAndVerify) { + createServer("/run/simm/admin_gflag_set"); + + proto::common::SetGFlagValueRequestPB setReq; + setReq.set_flag_name("test_admin_flag"); + setReq.set_flag_value("new_value"); + std::string setBuf; + setReq.SerializeToString(&setBuf); + + { + auto client = connectClient(); + uint16_t respType = 0; + std::string respPayload; + ASSERT_TRUE(client.sendRequest(AdminMsgType::GFLAG_SET, setBuf, respType, respPayload)); + + proto::common::SetGFlagValueResponsePB resp; + ASSERT_TRUE(resp.ParseFromString(respPayload)); + EXPECT_EQ(resp.ret_code(), CommonErr::OK); + } + + proto::common::GetGFlagValueRequestPB getReq; + getReq.set_flag_name("test_admin_flag"); + std::string getBuf; + getReq.SerializeToString(&getBuf); + + { + auto client = connectClient(); + uint16_t respType = 0; + std::string respPayload; + ASSERT_TRUE(client.sendRequest(AdminMsgType::GFLAG_GET, getBuf, respType, respPayload)); + + proto::common::GetGFlagValueResponsePB resp; + ASSERT_TRUE(resp.ParseFromString(respPayload)); + EXPECT_EQ(resp.ret_code(), CommonErr::OK); + EXPECT_EQ(resp.flag_info().flag_value(), "new_value"); + } +} + +TEST_F(AdminServerTest, GFlagSetNonexistentFlag) { + createServer("/run/simm/admin_gflag_set_unk"); + + proto::common::SetGFlagValueRequestPB req; + req.set_flag_name("nonexistent_flag_12345"); + req.set_flag_value("42"); + std::string reqBuf; + req.SerializeToString(&reqBuf); + + auto client = connectClient(); + uint16_t respType = 0; + std::string respPayload; + ASSERT_TRUE(client.sendRequest(AdminMsgType::GFLAG_SET, reqBuf, respType, respPayload)); + + proto::common::SetGFlagValueResponsePB resp; + ASSERT_TRUE(resp.ParseFromString(respPayload)); + EXPECT_EQ(resp.ret_code(), CommonErr::GFlagNotFound); +} + +// --------------------------------------------------------------------------- +// Trace toggle handler test +// --------------------------------------------------------------------------- + +TEST_F(AdminServerTest, TraceToggle) { + createServer("/run/simm/admin_trace"); + + proto::common::TraceToggleRequestPB req; + req.set_enable_trace(true); + std::string reqBuf; + req.SerializeToString(&reqBuf); + + auto client = connectClient(); + uint16_t respType = 0; + std::string respPayload; + ASSERT_TRUE(client.sendRequest(AdminMsgType::TRACE_TOGGLE, reqBuf, respType, respPayload)); + EXPECT_EQ(respType, static_cast(AdminMsgType::TRACE_TOGGLE)); + + proto::common::TraceToggleResponsePB resp; + ASSERT_TRUE(resp.ParseFromString(respPayload)); + EXPECT_EQ(resp.ret_code(), CommonErr::OK); +} + +// --------------------------------------------------------------------------- +// Custom handler registration tests +// --------------------------------------------------------------------------- + +TEST_F(AdminServerTest, RegisterCustomHandler) { + createServer("/run/simm/admin_custom"); + + server_->registerHandler(AdminMsgType::DS_STATUS, + [](const std::string& payload) -> std::string { + proto::common::AdmDsStatusResponsePB resp; + resp.set_ret_code(CommonErr::OK); + resp.set_is_registered(true); + resp.set_cm_ready(true); + resp.set_heartbeat_failure_count(0); + std::string buf; + resp.SerializeToString(&buf); + return buf; + }); + + auto client = connectClient(); + uint16_t respType = 0; + std::string respPayload; + ASSERT_TRUE(client.sendRequest(AdminMsgType::DS_STATUS, "", respType, respPayload)); + EXPECT_EQ(respType, static_cast(AdminMsgType::DS_STATUS)); + + proto::common::AdmDsStatusResponsePB resp; + ASSERT_TRUE(resp.ParseFromString(respPayload)); + EXPECT_EQ(resp.ret_code(), CommonErr::OK); + EXPECT_TRUE(resp.is_registered()); + EXPECT_TRUE(resp.cm_ready()); + EXPECT_EQ(resp.heartbeat_failure_count(), 0u); +} + +TEST_F(AdminServerTest, CustomHandlerReceivesPayload) { + createServer("/run/simm/admin_payload"); + + server_->registerHandler(AdminMsgType::DS_STATUS, + [](const std::string& payload) -> std::string { + proto::common::AdmDsStatusResponsePB resp; + resp.set_ret_code(CommonErr::OK); + resp.set_heartbeat_failure_count(static_cast(payload.size())); + std::string buf; + resp.SerializeToString(&buf); + return buf; + }); + + proto::common::AdmDsStatusRequestPB req; + std::string reqBuf; + req.SerializeToString(&reqBuf); + + auto client = connectClient(); + uint16_t respType = 0; + std::string respPayload; + ASSERT_TRUE(client.sendRequest(AdminMsgType::DS_STATUS, reqBuf, respType, respPayload)); + + proto::common::AdmDsStatusResponsePB resp; + ASSERT_TRUE(resp.ParseFromString(respPayload)); + EXPECT_EQ(resp.ret_code(), CommonErr::OK); + EXPECT_EQ(resp.heartbeat_failure_count(), 0u); +} + +TEST_F(AdminServerTest, OverrideBuiltinHandler) { + createServer("/run/simm/admin_override"); + + server_->registerHandler(AdminMsgType::GFLAG_LIST, + [](const std::string& /*payload*/) -> std::string { + proto::common::ListAllGFlagsResponsePB resp; + resp.set_ret_code(CommonErr::OK); + std::string buf; + resp.SerializeToString(&buf); + return buf; + }); + + auto client = connectClient(); + uint16_t respType = 0; + std::string respPayload; + ASSERT_TRUE(client.sendRequest(AdminMsgType::GFLAG_LIST, "", respType, respPayload)); + + proto::common::ListAllGFlagsResponsePB resp; + ASSERT_TRUE(resp.ParseFromString(respPayload)); + EXPECT_EQ(resp.ret_code(), CommonErr::OK); + EXPECT_EQ(resp.flags_size(), 0); +} + +// --------------------------------------------------------------------------- +// Protocol robustness tests +// --------------------------------------------------------------------------- + +TEST_F(AdminServerTest, UnregisteredMsgTypeNoResponse) { + createServer("/run/simm/admin_unknown_type"); + + int fd = ::socket(AF_UNIX, SOCK_STREAM, 0); + ASSERT_GE(fd, 0); + + sockaddr_un addr; + std::memset(&addr, 0, sizeof(addr)); + addr.sun_family = AF_UNIX; + std::strncpy(addr.sun_path, server_->socketPath().c_str(), sizeof(addr.sun_path) - 1); + socklen_t addrLen = static_cast( + offsetof(sockaddr_un, sun_path) + std::strlen(addr.sun_path)); + ASSERT_EQ(::connect(fd, reinterpret_cast(&addr), addrLen), 0); + + uint32_t frameLen = htonl(sizeof(uint16_t)); + uint16_t msgType = htons(999); + ::write(fd, &frameLen, sizeof(frameLen)); + ::write(fd, &msgType, sizeof(msgType)); + + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + char buf[64]; + ssize_t n = ::read(fd, buf, sizeof(buf)); + EXPECT_LE(n, 0); + + ::close(fd); +} + +TEST_F(AdminServerTest, InvalidFrameTooShort) { + createServer("/run/simm/admin_bad_frame"); + + int fd = ::socket(AF_UNIX, SOCK_STREAM, 0); + ASSERT_GE(fd, 0); + + sockaddr_un addr; + std::memset(&addr, 0, sizeof(addr)); + addr.sun_family = AF_UNIX; + std::strncpy(addr.sun_path, server_->socketPath().c_str(), sizeof(addr.sun_path) - 1); + socklen_t addrLen = static_cast( + offsetof(sockaddr_un, sun_path) + std::strlen(addr.sun_path)); + ASSERT_EQ(::connect(fd, reinterpret_cast(&addr), addrLen), 0); + + uint32_t frameLen = htonl(1); + ::write(fd, &frameLen, sizeof(frameLen)); + uint8_t garbage = 0xFF; + ::write(fd, &garbage, 1); + + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + char buf[64]; + ssize_t n = ::read(fd, buf, sizeof(buf)); + EXPECT_LE(n, 0); + + ::close(fd); +} + +TEST_F(AdminServerTest, EmptyPayload) { + createServer("/run/simm/admin_empty_payload"); + + auto client = connectClient(); + uint16_t respType = 0; + std::string respPayload; + ASSERT_TRUE(client.sendRequest(AdminMsgType::GFLAG_LIST, "", respType, respPayload)); + + proto::common::ListAllGFlagsResponsePB resp; + ASSERT_TRUE(resp.ParseFromString(respPayload)); + EXPECT_EQ(resp.ret_code(), CommonErr::OK); +} + +TEST_F(AdminServerTest, MalformedProtobufPayload) { + createServer("/run/simm/admin_bad_proto"); + + std::string garbage = "\x00\xFF\xFE\xAB\xCD"; + + auto client = connectClient(); + uint16_t respType = 0; + std::string respPayload; + ASSERT_TRUE(client.sendRequest(AdminMsgType::GFLAG_GET, garbage, respType, respPayload)); + + proto::common::GetGFlagValueResponsePB resp; + ASSERT_TRUE(resp.ParseFromString(respPayload)); + // Garbled payload may still parse as protobuf with a junk flag name, + // resulting in GFlagNotFound rather than InvalidArgument. + EXPECT_NE(resp.ret_code(), CommonErr::OK); +} + +// --------------------------------------------------------------------------- +// Concurrent client test +// --------------------------------------------------------------------------- + +TEST_F(AdminServerTest, MultipleSequentialClients) { + createServer("/run/simm/admin_multi_client"); + + for (int i = 0; i < 5; ++i) { + auto client = connectClient(); + uint16_t respType = 0; + std::string respPayload; + ASSERT_TRUE(client.sendRequest(AdminMsgType::GFLAG_LIST, "", respType, respPayload)); + + proto::common::ListAllGFlagsResponsePB resp; + ASSERT_TRUE(resp.ParseFromString(respPayload)); + EXPECT_EQ(resp.ret_code(), CommonErr::OK); + } +} + +TEST_F(AdminServerTest, ConcurrentClients) { + createServer("/run/simm/admin_concurrent"); + + constexpr int kNumClients = 10; + std::atomic successCount{0}; + std::vector threads; + + for (int i = 0; i < kNumClients; ++i) { + threads.emplace_back([&, i]() { + std::this_thread::sleep_for(std::chrono::milliseconds(i * 10)); + + UdsClient client; + if (!client.connect(server_->socketPath())) return; + + uint16_t respType = 0; + std::string respPayload; + if (!client.sendRequest(AdminMsgType::GFLAG_LIST, "", respType, respPayload)) return; + + proto::common::ListAllGFlagsResponsePB resp; + if (resp.ParseFromString(respPayload) && resp.ret_code() == CommonErr::OK) { + successCount.fetch_add(1); + } + }); + } + + for (auto& t : threads) t.join(); + + EXPECT_EQ(successCount.load(), kNumClients); +} + +// --------------------------------------------------------------------------- +// Shutdown tests +// --------------------------------------------------------------------------- + +TEST_F(AdminServerTest, ShutdownWhileClientConnected) { + createServer("/run/simm/admin_shutdown_client"); + std::string path = server_->socketPath(); + + UdsClient client; + ASSERT_TRUE(client.connect(path)); + + server_.reset(); + + struct stat st; + EXPECT_NE(::stat(path.c_str(), &st), 0); +} + +TEST_F(AdminServerTest, DoubleDestructionSafe) { + createServer("/run/simm/admin_double_destroy"); + server_.reset(); + server_.reset(); +} + +TEST_F(AdminServerTest, StaleSocketFileOverwritten) { + std::string stalePath = std::string("/run/simm/admin_stale.") + + std::to_string(::getpid()) + ".sock"; + { + int fd = ::creat(stalePath.c_str(), 0666); + if (fd >= 0) ::close(fd); + } + + auto server = std::make_unique("/run/simm/admin_stale"); + EXPECT_TRUE(server->isRunning()); + struct stat st; + EXPECT_EQ(::stat(server->socketPath().c_str(), &st), 0); + EXPECT_TRUE(S_ISSOCK(st.st_mode)); + + server.reset(); + ::unlink(stalePath.c_str()); +} + +// --------------------------------------------------------------------------- +// End-to-end tests: AdminServer + simmctl binary via subprocess +// --------------------------------------------------------------------------- + +// Helper: locate simmctl binary relative to the test binary. +// Test binary: build/{mode}/bin/unit_tests/test_admin_server +// simmctl: build/{mode}/bin/tools/simmctl +static std::string findSimmctlBinary() { + // Try via /proc/self/exe + char selfPath[4096]; + ssize_t len = ::readlink("/proc/self/exe", selfPath, sizeof(selfPath) - 1); + if (len > 0) { + selfPath[len] = '\0'; + std::string self(selfPath); + // Go up from unit_tests/ to bin/, then into tools/ + auto pos = self.rfind('/'); + if (pos != std::string::npos) { + std::string binDir = self.substr(0, pos); // .../bin/unit_tests + pos = binDir.rfind('/'); + if (pos != std::string::npos) { + binDir = binDir.substr(0, pos); // .../bin + return binDir + "/tools/simmctl"; + } + } + } + // Fallback: assume it's in PATH + return "simmctl"; +} + +// Run simmctl as subprocess and capture stdout/stderr/exit code. +struct SimmctlResult { + int exitCode; + std::string stdoutStr; + std::string stderrStr; +}; + +static SimmctlResult runSimmctl(const std::vector& args) { + std::string simmctl = findSimmctlBinary(); + std::string cmd = simmctl; + for (const auto& arg : args) { + cmd += " " + arg; + } + cmd += " 2>/tmp/simmctl_test_stderr"; + + FILE* fp = popen(cmd.c_str(), "r"); + SimmctlResult result; + result.exitCode = -1; + if (!fp) { + return result; + } + + char buf[4096]; + while (fgets(buf, sizeof(buf), fp)) { + result.stdoutStr += buf; + } + int status = pclose(fp); + result.exitCode = WIFEXITED(status) ? WEXITSTATUS(status) : -1; + + // Read stderr + FILE* errFp = fopen("/tmp/simmctl_test_stderr", "r"); + if (errFp) { + while (fgets(buf, sizeof(buf), errFp)) { + result.stderrStr += buf; + } + fclose(errFp); + ::unlink("/tmp/simmctl_test_stderr"); + } + + return result; +} + +class SimmctlE2ETest : public ::testing::Test { + protected: + void SetUp() override { + ::mkdir("/run/simm", 0777); + pidStr_ = std::to_string(::getpid()); + } + + void TearDown() override { + server_.reset(); + } + + void createServer(const std::string& basePath) { + server_ = std::make_unique(basePath); + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + } + + std::unique_ptr server_; + std::string pidStr_; +}; + +TEST_F(SimmctlE2ETest, DsStatusViaPid) { + createServer("/run/simm/admin_ds"); + + // Register mock DS_STATUS handler + server_->registerHandler(AdminMsgType::DS_STATUS, + [](const std::string& /*payload*/) -> std::string { + proto::common::AdmDsStatusResponsePB resp; + resp.set_ret_code(CommonErr::OK); + resp.set_is_registered(true); + resp.set_cm_ready(true); + resp.set_heartbeat_failure_count(0); + std::string buf; + resp.SerializeToString(&buf); + return buf; + }); + + auto result = runSimmctl({"ds", "status", "--pid", pidStr_}); + EXPECT_EQ(result.exitCode, 0) << "stderr: " << result.stderrStr; + EXPECT_NE(result.stdoutStr.find("is_registered"), std::string::npos); + EXPECT_NE(result.stdoutStr.find("true"), std::string::npos); + EXPECT_NE(result.stdoutStr.find("cm_ready"), std::string::npos); +} + +TEST_F(SimmctlE2ETest, CmStatusViaPid) { + createServer("/run/simm/admin_cm"); + + server_->registerHandler(AdminMsgType::CM_STATUS, + [](const std::string& /*payload*/) -> std::string { + proto::common::AdmCmStatusResponsePB resp; + resp.set_ret_code(CommonErr::OK); + resp.set_is_running(true); + resp.set_service_ready(true); + resp.set_alive_node_count(3); + resp.set_dead_node_count(0); + resp.set_total_shard_count(64); + std::string buf; + resp.SerializeToString(&buf); + return buf; + }); + + auto result = runSimmctl({"cm", "status", "--pid", pidStr_}); + EXPECT_EQ(result.exitCode, 0) << "stderr: " << result.stderrStr; + EXPECT_NE(result.stdoutStr.find("is_running"), std::string::npos); + EXPECT_NE(result.stdoutStr.find("true"), std::string::npos); + EXPECT_NE(result.stdoutStr.find("alive_node_count"), std::string::npos); + EXPECT_NE(result.stdoutStr.find("3"), std::string::npos); +} + +TEST_F(SimmctlE2ETest, NodeListViaPid) { + createServer("/run/simm/admin_cm"); + + server_->registerHandler(AdminMsgType::NODE_LIST, + [](const std::string& /*payload*/) -> std::string { + ListNodesResponsePB resp; + resp.set_ret_code(CommonErr::OK); + + auto* n1 = resp.add_nodes(); + n1->mutable_node_address()->set_ip("10.0.0.2"); + n1->mutable_node_address()->set_port(40000); + n1->set_node_status(1); // RUNNING + + auto* n2 = resp.add_nodes(); + n2->mutable_node_address()->set_ip("10.0.0.3"); + n2->mutable_node_address()->set_port(40000); + n2->set_node_status(0); // DEAD + + std::string buf; + resp.SerializeToString(&buf); + return buf; + }); + + auto result = runSimmctl({"node", "list", "--pid", pidStr_}); + EXPECT_EQ(result.exitCode, 0) << "stderr: " << result.stderrStr; + EXPECT_NE(result.stdoutStr.find("10.0.0.2:40000"), std::string::npos); + EXPECT_NE(result.stdoutStr.find("RUNNING"), std::string::npos); + EXPECT_NE(result.stdoutStr.find("10.0.0.3:40000"), std::string::npos); + EXPECT_NE(result.stdoutStr.find("DEAD"), std::string::npos); +} + +TEST_F(SimmctlE2ETest, ShardListViaPid) { + createServer("/run/simm/admin_cm"); + + server_->registerHandler(AdminMsgType::SHARD_LIST, + [](const std::string& /*payload*/) -> std::string { + QueryShardRoutingTableAllResponsePB resp; + resp.set_ret_code(CommonErr::OK); + + auto* entry = resp.add_shard_info(); + entry->mutable_data_server_address()->set_ip("10.0.0.2"); + entry->mutable_data_server_address()->set_port(40000); + entry->add_shard_ids(0); + entry->add_shard_ids(1); + entry->add_shard_ids(2); + + std::string buf; + resp.SerializeToString(&buf); + return buf; + }); + + auto result = runSimmctl({"shard", "list", "--pid", pidStr_}); + EXPECT_EQ(result.exitCode, 0) << "stderr: " << result.stderrStr; + EXPECT_NE(result.stdoutStr.find("10.0.0.2:40000"), std::string::npos); + EXPECT_NE(result.stdoutStr.find("3"), std::string::npos); // shard count +} + +TEST_F(SimmctlE2ETest, GFlagGetViaPid) { + // gflag with --pid uses simm_trace..sock path + createServer("/run/simm/simm_trace"); + + // Built-in gflag handler should work — test_admin_flag is defined in this binary. + // Note: a prior test (GFlagSetAndVerify) may have changed the value, + // so just verify the output structure, not the specific value. + auto result = runSimmctl({"gflag", "get", "test_admin_flag", "--pid", pidStr_}); + EXPECT_EQ(result.exitCode, 0) << "stderr: " << result.stderrStr; + EXPECT_NE(result.stdoutStr.find("test_admin_flag"), std::string::npos); + EXPECT_NE(result.stdoutStr.find("VALUE"), std::string::npos); +} + +TEST_F(SimmctlE2ETest, PidAndProcMutuallyExclusive) { + auto result = runSimmctl({"cm", "status", "--pid", "12345", + "--proc", "cluster_manager"}); + EXPECT_NE(result.exitCode, 0); + EXPECT_NE(result.stderrStr.find("mutually exclusive"), std::string::npos); +} + +TEST_F(SimmctlE2ETest, ProcNotFound) { + // Use a process name that definitely doesn't exist + auto result = runSimmctl({"cm", "status", "--proc", "cluster_manager"}); + // This will either fail because no cluster_manager is running, + // or succeed if one happens to be running. We only check the not-found case + // by using an unlikely process name via direct pgrep test. + // For a deterministic test, just verify --proc with invalid name fails. + auto result2 = runSimmctl({"cm", "status", "--proc", "nonexistent_proc_xyz"}); + EXPECT_NE(result2.exitCode, 0); +} + +TEST_F(SimmctlE2ETest, ProcInvalidName) { + auto result = runSimmctl({"cm", "status", "--proc", "invalid_name"}); + EXPECT_NE(result.exitCode, 0); + EXPECT_NE(result.stderrStr.find("cluster_manager"), std::string::npos); +} + +} // namespace +} // namespace common +} // namespace simm diff --git a/tests/protocol_test/conftest.py b/tests/protocol_test/conftest.py new file mode 100644 index 0000000..dd9c2a5 --- /dev/null +++ b/tests/protocol_test/conftest.py @@ -0,0 +1,175 @@ +"""Root pytest conftest with shared fixtures for distributed protocol tests. + +Supports both single-machine and multi-machine modes: +- Single-machine (default): all processes on localhost +- Multi-machine: set SIMM_CLUSTER_CONFIG env var to a YAML file with host topology + +Example multi-machine YAML: + cluster: + hosts: + cm: + ip: "10.0.0.1" + binary_dir: "/opt/simm/build/release/bin" + log_dir: "/tmp/simm_test_logs" + ds: + - ip: "10.0.0.2" + binary_dir: "/opt/simm/build/release/bin" + - ip: "10.0.0.3" + - ip: "10.0.0.4" + ssh: + user: "root" + port: 22 +""" + +import logging +import os +from pathlib import Path + +import pytest + +from framework.cluster import SimmCluster +from framework.config import ClusterConfig, dict_to_cluster_config, load_yaml + +logger = logging.getLogger(__name__) + + +def detect_root() -> bool: + return os.geteuid() == 0 + + +def _load_multi_machine_config() -> dict | None: + """Load cluster topology from SIMM_CLUSTER_CONFIG env var.""" + config_path = os.environ.get("SIMM_CLUSTER_CONFIG") + if not config_path: + return None + p = Path(config_path) + if not p.exists(): + logger.warning("SIMM_CLUSTER_CONFIG=%s does not exist, falling back to single-machine", + config_path) + return None + return load_yaml(p) + + +def _make_config( + num_ds: int = 3, + shard_total_num: int = 64, + cm_cluster_init_grace_period_inSecs: int = 5, + cm_heartbeat_timeout_inSecs: int = 10, + cm_heartbeat_bg_scan_interval_inSecs: int = 2, + heartbeat_cooldown_sec: int = 2, + register_cooldown_sec: int = 2, + memory_limit_bytes: int = 1 << 30, +) -> ClusterConfig: + """Create a ClusterConfig, applying multi-machine host topology if configured.""" + ext_data = _load_multi_machine_config() + + if ext_data: + base = dict_to_cluster_config(ext_data) + base.num_data_servers = num_ds + base.shard_total_num = shard_total_num + base.cm_cluster_init_grace_period_inSecs = cm_cluster_init_grace_period_inSecs + base.cm_heartbeat_timeout_inSecs = cm_heartbeat_timeout_inSecs + base.cm_heartbeat_bg_scan_interval_inSecs = cm_heartbeat_bg_scan_interval_inSecs + base.heartbeat_cooldown_sec = heartbeat_cooldown_sec + base.register_cooldown_sec = register_cooldown_sec + base.memory_limit_bytes = memory_limit_bytes + return base + + return ClusterConfig( + num_data_servers=num_ds, + shard_total_num=shard_total_num, + cm_cluster_init_grace_period_inSecs=cm_cluster_init_grace_period_inSecs, + cm_heartbeat_timeout_inSecs=cm_heartbeat_timeout_inSecs, + cm_heartbeat_bg_scan_interval_inSecs=cm_heartbeat_bg_scan_interval_inSecs, + heartbeat_cooldown_sec=heartbeat_cooldown_sec, + register_cooldown_sec=register_cooldown_sec, + memory_limit_bytes=memory_limit_bytes, + ) + + +@pytest.fixture(scope="session") +def has_root(): + return detect_root() + + +@pytest.fixture(scope="session") +def binary_dir(): + """Resolve the build directory containing SiMM binaries (for local/single-machine).""" + env_dir = os.environ.get("SIMM_BUILD_DIR") + if env_dir: + d = Path(env_dir) + if d.exists(): + return d + + # Default: look relative to this file + simm_root = Path(__file__).parents[2] # tests/protocol_test -> SiMM root + for mode in ["release", "relwithdeb", "debug"]: + d = simm_root / "build" / mode / "bin" + if (d / "cluster_manager").exists(): + return d + + # In multi-machine mode, binaries may only exist on remote hosts. + # cluster.py uses per-host binary_dir from YAML config, so local path is unused. + if os.environ.get("SIMM_CLUSTER_CONFIG"): + return None + + # pytest.skip raises Skipped exception — execution never reaches implicit return + pytest.skip("SiMM binaries not found. Set SIMM_BUILD_DIR or build with ./build.sh") + return Path("/dev/null") + + +@pytest.fixture +def cluster_small(tmp_path, binary_dir): + """1 CM + 3 DS cluster with fast heartbeats.""" + config = _make_config(num_ds=3, shard_total_num=64) + cluster = SimmCluster( + config, + log_dir=tmp_path / "logs" if not config.cm_host else None, + binary_dir=binary_dir if not config.cm_host else None, + ) + cluster.start() + cluster.wait_ready() + yield cluster + cluster.teardown() + + +@pytest.fixture +def cluster_medium(tmp_path, binary_dir): + """1 CM + 6 DS cluster.""" + config = _make_config(num_ds=6, cm_cluster_init_grace_period_inSecs=8) + cluster = SimmCluster( + config, + log_dir=tmp_path / "logs" if not config.cm_host else None, + binary_dir=binary_dir if not config.cm_host else None, + ) + cluster.start() + cluster.wait_ready() + yield cluster + cluster.teardown() + + +@pytest.fixture +def cluster_from_config(tmp_path, binary_dir): + """ + Factory fixture: creates a cluster from a custom ClusterConfig. + Usage: + def test_custom(cluster_from_config): + cluster = cluster_from_config(ClusterConfig(num_data_servers=4, ...)) + """ + clusters = [] + + def _factory(config: ClusterConfig) -> SimmCluster: + cluster = SimmCluster( + config, + log_dir=tmp_path / "logs" if not config.cm_host else None, + binary_dir=binary_dir if not config.cm_host else None, + ) + cluster.start() + cluster.wait_ready() + clusters.append(cluster) + return cluster + + yield _factory + + for c in clusters: + c.teardown() diff --git a/tests/protocol_test/framework/__init__.py b/tests/protocol_test/framework/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/protocol_test/framework/admin_client.py b/tests/protocol_test/framework/admin_client.py new file mode 100644 index 0000000..5c3371f --- /dev/null +++ b/tests/protocol_test/framework/admin_client.py @@ -0,0 +1,232 @@ +"""Non-invasive admin client wrapping simmctl CLI tool via UDS mode. + +All commands are sent through Unix domain sockets (--pid ). +For remote hosts, simmctl is executed via SSH on the target node. +For local hosts, simmctl runs as a direct subprocess. +""" + +import logging +from dataclasses import dataclass + +from .ssh_executor import SshExecutor + +logger = logging.getLogger(__name__) + + +@dataclass +class NodeInfo: + address: str # "ip:port" + status: str # "RUNNING" or "DEAD" + mem_total_mb: int = 0 + mem_free_mb: int = 0 + mem_used_mb: int = 0 + + +class AdminClientError(Exception): + """Raised when admin CLI tool fails.""" + pass + + +class AdminClient: + """ + Wraps simmctl as subprocess/SSH calls via UDS mode (--pid). + Parses tabulate output to extract structured data. + + For local processes, runs simmctl as a direct subprocess. + For remote processes, runs simmctl on the target host via SSH. + """ + + def __init__(self, ssh: SshExecutor, ctl_path: str, flags_path: str, + default_timeout: float = 10.0): + """ + Args: + ssh: SshExecutor for running commands on local/remote hosts. + ctl_path: Path to simmctl binary on target hosts. + flags_path: Path to simm_flags_admin binary on target hosts. + default_timeout: Default command timeout in seconds. + """ + self._ssh = ssh + self._ctl_path = ctl_path + self._flags_path = flags_path + self._timeout = default_timeout + + def _run_ctl_uds(self, host: str, pid: int, args: list[str], + timeout: float | None = None) -> str: + """Run simmctl in UDS mode (--pid) on the target host.""" + parts = [self._ctl_path, "--pid", str(pid)] + args + cmd = " ".join(parts) + timeout = timeout or self._timeout + try: + result = self._ssh.run(host, cmd, timeout=timeout, check=False) + if result.returncode != 0: + raise AdminClientError( + f"simmctl failed on {host} (rc={result.returncode}): " + f"{result.stderr.strip()}" + ) + return result.stdout + except AdminClientError: + raise + except Exception as e: + raise AdminClientError(f"simmctl failed on {host}: {e}") + + def _run_flags_uds(self, host: str, pid: int, method: str, + flag: str = "", value: str = "", + timeout: float | None = None) -> str: + """Run simm_flags_admin in UDS mode (--pid) on the target host.""" + parts = [ + self._flags_path, + f"--pid={pid}", + f"--method={method}", + ] + if flag: + parts.append(f"--flag={flag}") + if value: + parts.append(f"--value={value}") + + cmd = " ".join(parts) + timeout = timeout or self._timeout + try: + result = self._ssh.run(host, cmd, timeout=timeout, check=False) + if result.returncode != 0: + raise AdminClientError( + f"simm_flags_admin failed on {host} (rc={result.returncode}): " + f"{result.stderr.strip()}" + ) + return result.stdout + except AdminClientError: + raise + except Exception as e: + raise AdminClientError(f"simm_flags_admin failed on {host}: {e}") + + @staticmethod + def _parse_tabulate_rows(output: str) -> list[list[str]]: + """ + Parse tabulate table output into rows of cell values. + Tabulate format uses | as column separator and +---+ as row separator. + """ + rows = [] + for line in output.splitlines(): + line = line.strip() + if not line or line.startswith("+"): + continue + if "|" in line: + cells = [cell.strip() for cell in line.split("|")] + # Remove empty first/last elements from leading/trailing | + cells = [c for c in cells if c] + if cells: + rows.append(cells) + return rows + + # --- Node operations (via CM admin UDS) --- + + def list_nodes(self, host: str, cm_pid: int, + verbose: bool = False) -> list[NodeInfo]: + """List all nodes via simmctl --pid node list.""" + args = ["node", "list"] + if verbose: + args.append("--verbose") + output = self._run_ctl_uds(host, cm_pid, args) + rows = self._parse_tabulate_rows(output) + + nodes = [] + for row in rows[1:]: # skip header row + if verbose and len(row) >= 5: + nodes.append(NodeInfo( + address=row[0], + status=row[1], + mem_total_mb=int(row[2]) if row[2].isdigit() else 0, + mem_free_mb=int(row[3]) if row[3].isdigit() else 0, + mem_used_mb=int(row[4]) if row[4].isdigit() else 0, + )) + elif len(row) >= 2: + nodes.append(NodeInfo(address=row[0], status=row[1])) + return nodes + + # --- Shard operations (via CM admin UDS) --- + + def list_shards(self, host: str, cm_pid: int) -> dict[str, int]: + """ + List shard distribution via simmctl --pid shard list. + Returns {node_addr: shard_count}. + """ + output = self._run_ctl_uds(host, cm_pid, ["shard", "list"]) + rows = self._parse_tabulate_rows(output) + + distribution: dict[str, int] = {} + for row in rows[1:]: # skip header + if len(row) >= 2: + distribution[row[0]] = int(row[1]) + return distribution + + def list_shards_verbose(self, host: str, cm_pid: int) -> dict[str, list[int]]: + """ + List detailed shard assignment via simmctl --pid shard list --verbose. + Returns {node_addr: [shard_ids]}. + """ + output = self._run_ctl_uds(host, cm_pid, + ["shard", "list", "--verbose"]) + rows = self._parse_tabulate_rows(output) + + distribution: dict[str, list[int]] = {} + for row in rows[1:]: # skip header + if len(row) >= 2: + shard_ids = [int(s.strip()) for s in row[1].split(",") if s.strip().isdigit()] + distribution[row[0]] = shard_ids + return distribution + + # --- DS status operations (via DS admin UDS) --- + + def get_ds_status(self, host: str, ds_pid: int) -> dict[str, str]: + """ + Query DS internal status via simmctl --pid ds status. + Runs on the DS host to access /run/simm/admin_ds..sock. + Returns {"is_registered": "true"/"false", + "cm_ready": "true"/"false", + "heartbeat_failure_count": "N"}. + """ + output = self._run_ctl_uds(host, ds_pid, ["ds", "status"]) + rows = self._parse_tabulate_rows(output) + status = {} + for row in rows[1:]: # skip header + if len(row) >= 2: + status[row[0]] = row[1] + return status + + # --- GFlag operations (via UDS) --- + + def get_flag(self, host: str, pid: int, flag_name: str) -> str | None: + """Get a single flag value via UDS.""" + try: + output = self._run_ctl_uds(host, pid, ["gflag", "get", flag_name]) + rows = self._parse_tabulate_rows(output) + for row in rows: + if len(row) >= 2 and row[0] == "VALUE": + return row[1] + return None + except AdminClientError as e: + logger.error("get_flag failed: %s", e) + return None + + def set_flag(self, host: str, pid: int, + flag_name: str, value: str) -> bool: + """Set a flag value via UDS.""" + try: + self._run_ctl_uds(host, pid, ["gflag", "set", flag_name, value]) + return True + except AdminClientError as e: + logger.error("set_flag failed: %s", e) + return False + + def list_flags(self, host: str, pid: int) -> dict[str, str]: + """List all flags via UDS. Returns {flag_name: value}.""" + try: + output = self._run_ctl_uds(host, pid, ["gflag", "list"]) + rows = self._parse_tabulate_rows(output) + flags = {} + for row in rows[1:]: # skip header + if len(row) >= 2: + flags[row[0]] = row[1] + return flags + except AdminClientError as e: + logger.error("list_flags failed: %s", e) + return {} diff --git a/tests/protocol_test/framework/cluster.py b/tests/protocol_test/framework/cluster.py new file mode 100644 index 0000000..9627c19 --- /dev/null +++ b/tests/protocol_test/framework/cluster.py @@ -0,0 +1,352 @@ +"""High-level SiMM cluster lifecycle management for test fixtures. + +Supports both single-machine and multi-machine deployments. +In multi-machine mode, the test runner (node A) orchestrates CM/DS processes +on remote hosts via SSH. +""" + +import logging +import os +import time +from pathlib import Path + +from .admin_client import AdminClient +from .cluster_observer import ClusterObserver +from .config import ClusterConfig +from .fault_injector import FaultInjector +from .port_allocator import PortAllocator +from .process_manager import ProcessHandle, ProcessManager +from .ssh_executor import SshConfig, SshExecutor + +logger = logging.getLogger(__name__) + + +class SimmCluster: + """ + High-level cluster object providing full lifecycle management. + Used as a pytest fixture. + + Supports two modes: + - Single-machine: all processes on localhost (config.cm_host is None) + - Multi-machine: CM and DS on specified remote hosts via SSH + """ + + def __init__(self, config: ClusterConfig, log_dir: str | Path | None = None, + binary_dir: str | Path | None = None): + self.config = config + + # Initialize SSH executor + ssh_config = SshConfig(user=config.ssh_user, port=config.ssh_port) + self._ssh = SshExecutor(ssh_config=ssh_config) + self._is_multi_machine = config.cm_host is not None + + # Determine default binary_dir — priority: + # 1. YAML config.binary_dir (highest) + # 2. Constructor parameter + # 3. SIMM_BUILD_DIR env var + # 4. Derive from source tree + build_mode (lowest) + if config.binary_dir: + default_binary_dir = config.binary_dir + elif binary_dir is not None: + default_binary_dir = str(binary_dir) + elif os.environ.get("SIMM_BUILD_DIR"): + default_binary_dir = os.environ["SIMM_BUILD_DIR"] + else: + simm_root = Path(__file__).parents[3] + default_binary_dir = str(simm_root / "build" / config.build_mode / "bin") + + if log_dir is not None: + default_log_dir = str(log_dir) + else: + default_log_dir = "/tmp/simm_test_logs" + + self._default_binary_dir = default_binary_dir + self._default_log_dir = default_log_dir + + # Validate binaries exist (check on local for single-machine, + # or on remote hosts for multi-machine) + if not self._is_multi_machine: + for binary in ["cluster_manager", "data_server"]: + p = Path(default_binary_dir) / binary + if not p.exists(): + raise FileNotFoundError( + f"Binary '{binary}' not found at {default_binary_dir}. " + f"Build with: ./build.sh --mode={config.build_mode}" + ) + + self._port_allocator = PortAllocator(ssh_executor=self._ssh) + self._process_manager = ProcessManager( + default_binary_dir, default_log_dir, + self._port_allocator, self._ssh, + ) + + # Admin client — runs simmctl on target nodes via SSH/local subprocess + ctl_path = str(Path(default_binary_dir) / "tools" / "simmctl") + flags_path = str(Path(default_binary_dir) / "tools" / "simm_flags_admin") + self._admin_client = AdminClient(self._ssh, ctl_path, flags_path) + + # State + self.cm: ProcessHandle | None = None + self.data_servers: list[ProcessHandle] = [] + + # Initialized after start() + self.observer: ClusterObserver | None = None + self.fault_injector: FaultInjector | None = None + + def _verify_ssh_connectivity(self) -> None: + """Verify SSH connectivity to all remote hosts.""" + hosts_to_check = set() + if self.config.cm_host: + hosts_to_check.add(self.config.cm_host.ip) + for dh in self.config.ds_hosts: + hosts_to_check.add(dh.ip) + + for host in hosts_to_check: + if not self._ssh.is_local(host): + if not self._ssh.check_connectivity(host): + raise RuntimeError( + f"SSH connectivity check failed for {host}. " + f"Ensure passwordless SSH is configured." + ) + logger.info("SSH connectivity verified for %d hosts", len(hosts_to_check)) + + def _kill_existing_processes(self) -> None: + """Kill any existing CM/DS processes on all target hosts before starting.""" + hosts = set() + if self._is_multi_machine: + if self.config.cm_host: + hosts.add(self.config.cm_host.ip) + for dh in self.config.ds_hosts: + hosts.add(dh.ip) + else: + hosts.add("127.0.0.1") + + total_killed = 0 + for host in hosts: + for binary in ["cluster_manager", "data_server"]: + n = self._process_manager.kill_existing_by_name(host, binary) + total_killed += n + + if total_killed > 0: + logger.info("Killed %d existing SiMM process(es) before start", total_killed) + + def start(self) -> None: + """Kill existing CM/DS, then start fresh with test parameters.""" + self._kill_existing_processes() + + if self._is_multi_machine: + self._verify_ssh_connectivity() + self._start_multi_machine() + else: + self._start_single_machine() + + # Initialize observer and fault injector + self.observer = ClusterObserver( + admin_client=self._admin_client, + cm_handle=self.cm, + ) + + self.fault_injector = FaultInjector(self._process_manager, self._ssh) + logger.info("Cluster started: CM pid=%d on %s, %d DS", + self.cm.pid, self.cm.host, len(self.data_servers)) + + def _start_single_machine(self) -> None: + """Start all processes on localhost.""" + logger.info("Starting single-machine cluster: %d DS, %d shards", + self.config.num_data_servers, self.config.shard_total_num) + + self.cm = self._process_manager.start_cluster_manager( + host="127.0.0.1", + binary_dir=self._default_binary_dir, + log_dir=self._default_log_dir, + cm_cluster_init_grace_period_inSecs=self.config.cm_cluster_init_grace_period_inSecs, + cm_heartbeat_timeout_inSecs=self.config.cm_heartbeat_timeout_inSecs, + cm_heartbeat_bg_scan_interval_inSecs=self.config.cm_heartbeat_bg_scan_interval_inSecs, + shard_total_num=self.config.shard_total_num, + cm_deferred_reshard_enabled=self.config.cm_deferred_reshard_enabled, + cm_deferred_reshard_window_inSecs=self.config.cm_deferred_reshard_window_inSecs, + ) + + time.sleep(1) + + for i in range(self.config.num_data_servers): + logical_id = f"{self.config.ds_logical_node_id_prefix}-{i}" + ds = self._process_manager.start_data_server( + cm_ip=self.cm.ip, + cm_inter_port=self.cm.ports["inter"], + host="127.0.0.1", + binary_dir=self._default_binary_dir, + log_dir=self._default_log_dir, + heartbeat_cooldown_sec=self.config.heartbeat_cooldown_sec, + register_cooldown_sec=self.config.register_cooldown_sec, + cm_hb_tolerance_count=self.config.cm_hb_tolerance_count, + cm_connect_retry_interval_sec=self.config.cm_connect_retry_interval_sec, + memory_limit_bytes=self.config.memory_limit_bytes, + ds_logical_node_id=logical_id, + ) + self.data_servers.append(ds) + time.sleep(0.2) + + def _start_multi_machine(self) -> None: + """Start CM and DS on their designated remote hosts.""" + cm_host = self.config.cm_host + ds_hosts = self.config.ds_hosts + + logger.info("Starting multi-machine cluster: CM on %s, %d DS on %d hosts", + cm_host.ip, self.config.num_data_servers, len(ds_hosts)) + + # Start CM on its designated host + self.cm = self._process_manager.start_cluster_manager( + host=cm_host.ip, + ip=cm_host.ip, + binary_dir=cm_host.binary_dir or self._default_binary_dir, + log_dir=cm_host.log_dir, + cm_cluster_init_grace_period_inSecs=self.config.cm_cluster_init_grace_period_inSecs, + cm_heartbeat_timeout_inSecs=self.config.cm_heartbeat_timeout_inSecs, + cm_heartbeat_bg_scan_interval_inSecs=self.config.cm_heartbeat_bg_scan_interval_inSecs, + shard_total_num=self.config.shard_total_num, + cm_deferred_reshard_enabled=self.config.cm_deferred_reshard_enabled, + cm_deferred_reshard_window_inSecs=self.config.cm_deferred_reshard_window_inSecs, + ) + + time.sleep(1) + + # Distribute DS across ds_hosts in round-robin + for i in range(self.config.num_data_servers): + host_cfg = ds_hosts[i % len(ds_hosts)] + logical_id = f"{self.config.ds_logical_node_id_prefix}-{i}" + ds = self._process_manager.start_data_server( + cm_ip=self.cm.ip, + cm_inter_port=self.cm.ports["inter"], + host=host_cfg.ip, + ip=host_cfg.ip, + binary_dir=host_cfg.binary_dir or self._default_binary_dir, + log_dir=host_cfg.log_dir, + heartbeat_cooldown_sec=self.config.heartbeat_cooldown_sec, + register_cooldown_sec=self.config.register_cooldown_sec, + cm_hb_tolerance_count=self.config.cm_hb_tolerance_count, + cm_connect_retry_interval_sec=self.config.cm_connect_retry_interval_sec, + memory_limit_bytes=self.config.memory_limit_bytes, + ds_logical_node_id=logical_id, + ) + self.data_servers.append(ds) + time.sleep(0.2) + + def wait_ready(self, timeout: float = 120) -> None: + """Wait until grace period passes and all DS have registered.""" + grace_wait = self.config.cm_cluster_init_grace_period_inSecs + 2 + logger.info("Waiting %.0fs for grace period to expire...", grace_wait) + time.sleep(grace_wait) + + # Verify CM is still alive + if not self._process_manager.is_alive(self.cm): + cm_log = self._process_manager.get_log(self.cm) + raise RuntimeError(f"CM died during startup. Log:\n{cm_log[-2000:]}") + + # Verify all DS are alive + for ds in self.data_servers: + if not self._process_manager.is_alive(ds): + ds_log = self._process_manager.get_log(ds) + raise RuntimeError( + f"DS[{ds.index}] died during startup. Log:\n{ds_log[-2000:]}" + ) + + if not self.observer.wait_for_node_count( + self.config.num_data_servers, timeout=timeout - grace_wait + ): + logger.warning("Not all DS registered within timeout") + + logger.info("Cluster ready") + + def teardown(self) -> None: + """Kill all processes across all hosts, collect logs.""" + logger.info("Tearing down cluster...") + + if self.cm: + cm_log = self._process_manager.get_log(self.cm) + if cm_log: + logger.debug("CM log (last 500 chars): %s", cm_log[-500:]) + + self._process_manager.cleanup_all() + self._port_allocator.release_all() + self.cm = None + self.data_servers.clear() + logger.info("Cluster teardown complete") + + def restart_cm(self) -> ProcessHandle: + """Kill CM and restart it with the same ports. Returns new handle.""" + if self.cm is None: + raise RuntimeError("No CM to restart") + new_cm = self._process_manager.restart(self.cm) + self.cm = new_cm + + self.observer = ClusterObserver( + admin_client=self._admin_client, + cm_handle=new_cm, + ) + logger.info("CM restarted: pid=%d on %s", new_cm.pid, new_cm.host) + return new_cm + + def add_data_server(self, ports: dict[str, int] | None = None, + host: str | None = None, + ds_logical_node_id: str = "") -> ProcessHandle: + """Dynamically add a new DS to the running cluster. + + Args: + ports: Specific ports to use (e.g., for restart with same ports). + host: Host to start the DS on. Defaults to round-robin from ds_hosts + in multi-machine mode, or localhost in single-machine mode. + """ + if self.cm is None: + raise RuntimeError("No CM running") + + if host is None: + if self._is_multi_machine and self.config.ds_hosts: + # Round-robin across ds_hosts + idx = len(self.data_servers) % len(self.config.ds_hosts) + host_cfg = self.config.ds_hosts[idx] + host = host_cfg.ip + binary_dir = host_cfg.binary_dir or self._default_binary_dir + log_dir = host_cfg.log_dir + else: + host = "127.0.0.1" + binary_dir = self._default_binary_dir + log_dir = self._default_log_dir + else: + # Find matching host config + binary_dir = self._default_binary_dir + log_dir = self._default_log_dir + for hc in self.config.ds_hosts: + if hc.ip == host: + binary_dir = hc.binary_dir or self._default_binary_dir + log_dir = hc.log_dir + break + + ds = self._process_manager.start_data_server( + cm_ip=self.cm.ip, + cm_inter_port=self.cm.ports["inter"], + host=host, + ip=host, + ports=ports, + binary_dir=binary_dir, + log_dir=log_dir, + heartbeat_cooldown_sec=self.config.heartbeat_cooldown_sec, + register_cooldown_sec=self.config.register_cooldown_sec, + cm_hb_tolerance_count=self.config.cm_hb_tolerance_count, + cm_connect_retry_interval_sec=self.config.cm_connect_retry_interval_sec, + memory_limit_bytes=self.config.memory_limit_bytes, + ds_logical_node_id=ds_logical_node_id, + ) + self.data_servers.append(ds) + return ds + + def get_ds_handle(self, index: int) -> ProcessHandle: + """Get DS handle by index.""" + for ds in self.data_servers: + if ds.index == index: + return ds + raise IndexError(f"No DS with index {index}") + + @property + def process_manager(self) -> ProcessManager: + return self._process_manager diff --git a/tests/protocol_test/framework/cluster_observer.py b/tests/protocol_test/framework/cluster_observer.py new file mode 100644 index 0000000..69ae02a --- /dev/null +++ b/tests/protocol_test/framework/cluster_observer.py @@ -0,0 +1,308 @@ +"""Cluster state observation, condition waiting, and invariant assertions.""" + +import logging +import time + +from .admin_client import AdminClient, AdminClientError, NodeInfo +from .process_manager import ProcessHandle + +logger = logging.getLogger(__name__) + + +class ClusterObserver: + """ + Queries cluster state and waits for conditions. + Uses AdminClient (UDS-based) for all state queries. + """ + + def __init__( + self, + admin_client: AdminClient, + cm_handle: ProcessHandle, + ): + self._admin = admin_client + self._cm = cm_handle + + @property + def admin_client(self) -> AdminClient: + return self._admin + + # --- State queries via admin RPC --- + + def _list_nodes(self) -> list[NodeInfo]: + try: + return self._admin.list_nodes(self._cm.host, self._cm.pid) + except AdminClientError as e: + logger.debug("list_nodes failed: %s", e) + return [] + + def get_alive_node_count(self) -> int: + nodes = self._list_nodes() + return sum(1 for n in nodes if n.status == "RUNNING") + + def get_dead_node_count(self) -> int: + nodes = self._list_nodes() + return sum(1 for n in nodes if n.status == "DEAD") + + def get_all_node_statuses(self) -> dict[str, str]: + """Returns {node_addr: status}.""" + nodes = self._list_nodes() + return {n.address: n.status for n in nodes} + + def get_node_status(self, node_addr: str) -> str | None: + statuses = self.get_all_node_statuses() + return statuses.get(node_addr) + + def get_shard_distribution(self) -> dict[str, int]: + """Returns {node_addr: shard_count}.""" + try: + return self._admin.list_shards(self._cm.host, self._cm.pid) + except AdminClientError as e: + logger.debug("list_shards failed: %s", e) + return {} + + def get_total_shard_count(self) -> int: + dist = self.get_shard_distribution() + return sum(dist.values()) + + # --- Condition waiters (polling) --- + + def wait_for_node_count( + self, expected: int, alive_only: bool = True, + timeout: float = 60, poll_interval: float = 1.0 + ) -> bool: + """Wait until the node count matches expected.""" + deadline = time.time() + timeout + while time.time() < deadline: + count = self.get_alive_node_count() if alive_only else len(self._list_nodes()) + if count == expected: + logger.info("Node count reached %d", expected) + return True + time.sleep(poll_interval) + logger.warning("Timed out waiting for node count=%d (current=%d)", + expected, self.get_alive_node_count()) + return False + + def wait_for_node_status( + self, node_addr: str, status: str, + timeout: float = 60, poll_interval: float = 1.0 + ) -> bool: + """Wait until a specific node has the expected status.""" + deadline = time.time() + timeout + while time.time() < deadline: + current = self.get_node_status(node_addr) + if current == status: + logger.info("Node %s is now %s", node_addr, status) + return True + time.sleep(poll_interval) + logger.warning("Timed out waiting for node %s to become %s", node_addr, status) + return False + + def wait_for_shard_coverage( + self, expected_total: int, + timeout: float = 60, poll_interval: float = 1.0 + ) -> bool: + """Wait until total shard count matches expected.""" + deadline = time.time() + timeout + while time.time() < deadline: + total = self.get_total_shard_count() + if total == expected_total: + logger.info("Shard coverage complete: %d", expected_total) + return True + time.sleep(poll_interval) + logger.warning("Timed out waiting for shard coverage=%d (current=%d)", + expected_total, self.get_total_shard_count()) + return False + + def wait_for_rebalance_complete( + self, timeout: float = 60, poll_interval: float = 1.0 + ) -> bool: + """Wait until no shard is assigned to a DEAD node.""" + deadline = time.time() + timeout + while time.time() < deadline: + node_statuses = self.get_all_node_statuses() + shard_dist = self.get_shard_distribution() + dead_nodes = {addr for addr, s in node_statuses.items() if s == "DEAD"} + orphaned = any(addr in dead_nodes for addr in shard_dist if shard_dist[addr] > 0) + if not orphaned: + logger.info("Rebalance complete, no orphaned shards") + return True + time.sleep(poll_interval) + logger.warning("Timed out waiting for rebalance to complete") + return False + + def wait_for_stable_state( + self, duration: float = 10, check_interval: float = 2 + ) -> bool: + """Wait until cluster state remains unchanged for `duration` seconds.""" + last_snapshot = self.get_all_node_statuses() + stable_since = time.time() + while time.time() - stable_since < duration: + time.sleep(check_interval) + current = self.get_all_node_statuses() + if current != last_snapshot: + last_snapshot = current + stable_since = time.time() + return True + + # --- Invariant assertions --- + + def assert_no_orphaned_shards(self) -> None: + """Assert no shard is assigned to a DEAD node.""" + node_statuses = self.get_all_node_statuses() + shard_dist = self.get_shard_distribution() + dead_nodes = {addr for addr, s in node_statuses.items() if s == "DEAD"} + for addr, count in shard_dist.items(): + if addr in dead_nodes and count > 0: + raise AssertionError( + f"Orphaned shards: node {addr} is DEAD but has {count} shards" + ) + + def assert_total_shard_count(self, expected: int) -> None: + """Assert total shard count equals expected.""" + total = self.get_total_shard_count() + assert total == expected, ( + f"Shard count mismatch: expected {expected}, got {total}" + ) + + def assert_shard_balance(self, max_imbalance_ratio: float = 0.3) -> None: + """Assert shard distribution is roughly balanced across alive nodes.""" + node_statuses = self.get_all_node_statuses() + shard_dist = self.get_shard_distribution() + alive_counts = { + addr: shard_dist.get(addr, 0) + for addr, s in node_statuses.items() + if s == "RUNNING" + } + if not alive_counts: + return + + counts = list(alive_counts.values()) + avg = sum(counts) / len(counts) + if avg == 0: + return + + max_count = max(counts) + min_count = min(counts) + imbalance = (max_count - min_count) / avg + assert imbalance <= max_imbalance_ratio, ( + f"Shard imbalance too high: max={max_count}, min={min_count}, " + f"avg={avg:.1f}, imbalance={imbalance:.2f} > {max_imbalance_ratio}" + ) + + def assert_all_nodes_running(self) -> None: + """Assert all registered nodes are RUNNING.""" + statuses = self.get_all_node_statuses() + dead = {addr: s for addr, s in statuses.items() if s != "RUNNING"} + assert not dead, f"Expected all nodes RUNNING, but found: {dead}" + + def get_deferred_reshard_node_count(self) -> int: + """Count nodes in DEFERRED_RESHARD state.""" + nodes = self._list_nodes() + return sum(1 for n in nodes if n.status == "DEFERRED_RESHARD") + + def wait_for_node_status_not( + self, node_addr: str, status: str, + timeout: float = 60, poll_interval: float = 1.0 + ) -> bool: + """Wait until a specific node is NOT in the given status.""" + deadline = time.time() + timeout + while time.time() < deadline: + current = self.get_node_status(node_addr) + if current is not None and current != status: + return True + time.sleep(poll_interval) + return False + + # --- DS-side status verification (via UDS admin, requires PID) --- + + def get_ds_status(self, ds_host: str, ds_pid: int) -> dict[str, str]: + """Query DS internal status via UDS admin (simmctl --pid ds status). + Returns {"is_registered": "true/false", "cm_ready": "true/false", + "heartbeat_failure_count": "N"}. + """ + try: + return self._admin.get_ds_status(ds_host, ds_pid) + except AdminClientError as e: + logger.debug("get_ds_status failed for pid=%d on %s: %s", ds_pid, ds_host, e) + return {} + + def assert_ds_is_registered(self, ds_host: str, ds_pid: int) -> None: + """Assert DS reports itself as registered with CM.""" + status = self.get_ds_status(ds_host, ds_pid) + assert status.get("is_registered") == "true", ( + f"DS pid={ds_pid} is_registered={status.get('is_registered')}, " + f"expected true" + ) + + def assert_ds_cm_ready(self, ds_host: str, ds_pid: int, + expected: bool = True) -> None: + """Assert DS's cm_ready flag matches expected value.""" + status = self.get_ds_status(ds_host, ds_pid) + expected_str = "true" if expected else "false" + assert status.get("cm_ready") == expected_str, ( + f"DS pid={ds_pid} cm_ready={status.get('cm_ready')}, " + f"expected {expected_str}" + ) + + def assert_ds_heartbeat_failure_count( + self, ds_host: str, ds_pid: int, min_count: int = 0, + max_count: int | None = None + ) -> None: + """Assert DS heartbeat_failure_count within expected range.""" + status = self.get_ds_status(ds_host, ds_pid) + count_str = status.get("heartbeat_failure_count", "0") + count = int(count_str) + if min_count > 0: + assert count >= min_count, ( + f"DS pid={ds_pid} heartbeat_failure_count={count}, " + f"expected >= {min_count}" + ) + if max_count is not None: + assert count <= max_count, ( + f"DS pid={ds_pid} heartbeat_failure_count={count}, " + f"expected <= {max_count}" + ) + + def wait_for_ds_cm_not_ready( + self, ds_host: str, ds_pid: int, + timeout: float = 60, poll_interval: float = 1.0 + ) -> bool: + """Wait until DS reports cm_ready=false (detected CM failure).""" + deadline = time.time() + timeout + while time.time() < deadline: + status = self.get_ds_status(ds_host, ds_pid) + if status.get("cm_ready") == "false": + logger.info("DS pid=%d reports cm_ready=false", ds_pid) + return True + time.sleep(poll_interval) + logger.warning("Timed out waiting for DS pid=%d cm_ready=false", ds_pid) + return False + + def wait_for_ds_registered( + self, ds_host: str, ds_pid: int, + timeout: float = 60, poll_interval: float = 1.0 + ) -> bool: + """Wait until DS reports is_registered=true and cm_ready=true.""" + deadline = time.time() + timeout + while time.time() < deadline: + status = self.get_ds_status(ds_host, ds_pid) + if (status.get("is_registered") == "true" + and status.get("cm_ready") == "true"): + logger.info("DS pid=%d registered and cm_ready", ds_pid) + return True + time.sleep(poll_interval) + logger.warning("Timed out waiting for DS pid=%d registration", ds_pid) + return False + + def get_shard_distribution_for_node(self, node_addr: str) -> int: + """Get shard count for a specific node.""" + dist = self.get_shard_distribution() + return dist.get(node_addr, 0) + + def assert_node_has_no_shards(self, node_addr: str) -> None: + """Assert a specific node has zero shards assigned.""" + count = self.get_shard_distribution_for_node(node_addr) + assert count == 0, ( + f"Node {node_addr} still has {count} shards, expected 0" + ) diff --git a/tests/protocol_test/framework/config.py b/tests/protocol_test/framework/config.py new file mode 100644 index 0000000..8e5c17d --- /dev/null +++ b/tests/protocol_test/framework/config.py @@ -0,0 +1,193 @@ +"""YAML configuration loading and merging for test scenarios. + +Field names in ClusterConfig correspond to SiMM gflag names: + CM flags (cm_flags.cc): + cm_cluster_init_grace_period_inSecs (default 60) + cm_heartbeat_timeout_inSecs (default 30) + cm_heartbeat_bg_scan_interval_inSecs (default 5) + DS flags (ds_flags.cc): + heartbeat_cooldown_sec (default 5) + register_cooldown_sec (default 10) + cm_hb_tolerance_count (default 5) + cm_connect_retry_interval_sec (default 1) + memory_limit_bytes (default 0 = auto) + Common flags: + shard_total_num +""" + +import copy +from dataclasses import dataclass, field +from pathlib import Path + +import yaml + + +@dataclass +class HostConfig: + """Configuration for a single host (machine) in the cluster.""" + ip: str + binary_dir: str = "" + log_dir: str = "/tmp/simm_test_logs" + + +@dataclass +class ClusterConfig: + """Configuration for a test cluster. + + Field names match SiMM gflag names where applicable. + """ + num_data_servers: int = 3 + shard_total_num: int = 64 + + # CM flags — names match cm_flags.cc exactly + cm_cluster_init_grace_period_inSecs: int = 5 + cm_heartbeat_timeout_inSecs: int = 10 + cm_heartbeat_bg_scan_interval_inSecs: int = 2 + + # DS flags — names match ds_flags.cc exactly + heartbeat_cooldown_sec: int = 2 + register_cooldown_sec: int = 2 + cm_hb_tolerance_count: int = 5 + cm_connect_retry_interval_sec: int = 1 + memory_limit_bytes: int = 1 << 30 + + # Deferred Reshard flags — names match cm_flags.cc / ds_flags.cc + cm_deferred_reshard_enabled: bool = True + cm_deferred_reshard_window_inSecs: int = 120 + ds_logical_node_id_prefix: str = "simm-ds" # test helper: DS-{i} gets "{prefix}-{i}" + + # Build + binary_dir: str = "/opt/simm/build/release/bin" + build_mode: str = "release" + + # Host topology — None means single-machine mode + cm_host: HostConfig | None = None + ds_hosts: list[HostConfig] = field(default_factory=list) + + # SSH settings + ssh_user: str = "" + ssh_port: int = 22 + + # --- Derived timeout helpers --- + + @property + def cm_failure_detection_max_sec(self) -> float: + """Max time for CM to detect a dead DS: + cm_heartbeat_timeout_inSecs + 2 * cm_heartbeat_bg_scan_interval_inSecs""" + return (self.cm_heartbeat_timeout_inSecs + + 2 * self.cm_heartbeat_bg_scan_interval_inSecs) + + @property + def ds_cm_failure_detection_sec(self) -> float: + """Time for DS to detect CM is down: + cm_hb_tolerance_count * heartbeat_cooldown_sec""" + return self.cm_hb_tolerance_count * self.heartbeat_cooldown_sec + + @property + def ds_rejoin_max_sec(self) -> float: + """Max time for a DS to re-register after detecting CM failure: + ds_cm_failure_detection_sec + register_cooldown_sec + cm_connect_retry_interval_sec * retries""" + return (self.ds_cm_failure_detection_sec + + self.register_cooldown_sec + + 15) # buffer for retries + + +@dataclass +class FaultConfig: + """Configuration for a fault injection step.""" + fault_type: str + targets: list[str] = field(default_factory=list) + delay_after_sec: float = 0 + duration_sec: float | None = None + + +def deep_merge(base: dict, override: dict) -> dict: + """Deep merge two dicts. Override values take precedence.""" + result = copy.deepcopy(base) + for key, value in override.items(): + if key in result and isinstance(result[key], dict) and isinstance(value, dict): + result[key] = deep_merge(result[key], value) + else: + result[key] = copy.deepcopy(value) + return result + + +def load_yaml(path: Path) -> dict: + with open(path) as f: + return yaml.safe_load(f) or {} + + +def merge_yaml_files(*paths: Path) -> dict: + result: dict = {} + for p in paths: + data = load_yaml(p) + result = deep_merge(result, data) + return result + + +def _parse_host_config(data: dict) -> HostConfig: + return HostConfig( + ip=data["ip"], + binary_dir=data.get("binary_dir", ""), + log_dir=data.get("log_dir", "/tmp/simm_test_logs"), + ) + + +def dict_to_cluster_config(data: dict) -> ClusterConfig: + """Convert a YAML dict to ClusterConfig.""" + cluster = data.get("cluster", {}) + cm = cluster.get("cluster_manager", {}) + ds = cluster.get("data_servers", {}) + hosts = cluster.get("hosts", {}) + ssh = cluster.get("ssh", {}) + + cm_host = None + ds_hosts = [] + if "cm" in hosts: + cm_host = _parse_host_config(hosts["cm"]) + if "ds" in hosts: + ds_hosts = [_parse_host_config(h) for h in hosts["ds"]] + + return ClusterConfig( + num_data_servers=ds.get("count", 3), + shard_total_num=cluster.get("shard_total_num", 64), + cm_cluster_init_grace_period_inSecs=cm.get("cm_cluster_init_grace_period_inSecs", 5), + cm_heartbeat_timeout_inSecs=cm.get("cm_heartbeat_timeout_inSecs", 10), + cm_heartbeat_bg_scan_interval_inSecs=cm.get("cm_heartbeat_bg_scan_interval_inSecs", 2), + heartbeat_cooldown_sec=ds.get("heartbeat_cooldown_sec", 2), + register_cooldown_sec=ds.get("register_cooldown_sec", 2), + cm_hb_tolerance_count=ds.get("cm_hb_tolerance_count", 5), + cm_connect_retry_interval_sec=ds.get("cm_connect_retry_interval_sec", 1), + memory_limit_bytes=ds.get("memory_limit_bytes", 1 << 30), + cm_deferred_reshard_enabled=cm.get("cm_deferred_reshard_enabled", True), + cm_deferred_reshard_window_inSecs=cm.get("cm_deferred_reshard_window_inSecs", 120), + ds_logical_node_id_prefix=ds.get("ds_logical_node_id_prefix", "simm-ds"), + binary_dir=cluster.get("binary_dir", "/opt/simm/build/release/bin"), + build_mode=cluster.get("build_mode", "release"), + cm_host=cm_host, + ds_hosts=ds_hosts, + ssh_user=ssh.get("user", ""), + ssh_port=ssh.get("port", 22), + ) + + +def dict_to_fault_configs(data: dict) -> list[FaultConfig]: + faults = data.get("faults", []) + result = [] + for f in faults: + targets = f.get("targets", []) + if isinstance(f.get("target"), str): + targets = [f["target"]] + result.append(FaultConfig( + fault_type=f["type"], + targets=targets, + delay_after_sec=f.get("delay_after_ready_sec", 0), + duration_sec=f.get("duration_sec"), + )) + return result + + +def load_scenario(scenario_dir: Path, *yaml_files: str) -> tuple[ClusterConfig, list[FaultConfig]]: + paths = [scenario_dir / f for f in yaml_files] + merged = merge_yaml_files(*paths) + return dict_to_cluster_config(merged), dict_to_fault_configs(merged) diff --git a/tests/protocol_test/framework/fault_injector.py b/tests/protocol_test/framework/fault_injector.py new file mode 100644 index 0000000..984be3a --- /dev/null +++ b/tests/protocol_test/framework/fault_injector.py @@ -0,0 +1,155 @@ +"""Fault injection for cluster management testing — supports local and remote hosts. + +In multi-machine mode, signals are sent via SSH and iptables rules are applied +on the remote hosts where CM/DS processes run. +""" + +import logging +import signal +from contextlib import contextmanager + +from .process_manager import ProcessHandle, ProcessManager +from .ssh_executor import SshExecutor + +logger = logging.getLogger(__name__) + + +class FaultInjector: + """Injects faults via OS-level mechanisms (signals, iptables) across hosts.""" + + def __init__(self, process_manager: ProcessManager, ssh_executor: SshExecutor): + self._pm = process_manager + self._ssh = ssh_executor + + def kill_process(self, handle: ProcessHandle) -> None: + """SIGKILL — simulates hard crash.""" + logger.info("FAULT: Killing %s[%d] (pid=%d on %s) with SIGKILL", + handle.role, handle.index, handle.pid, handle.host) + self._pm.kill(handle, signal.SIGKILL) + + def graceful_stop(self, handle: ProcessHandle) -> None: + """SIGTERM — tests graceful shutdown path.""" + logger.info("FAULT: Graceful stop %s[%d] (pid=%d on %s) with SIGTERM", + handle.role, handle.index, handle.pid, handle.host) + self._pm.stop(handle, timeout=10) + + def freeze_process(self, handle: ProcessHandle) -> None: + """SIGSTOP — simulates process hang/unresponsive.""" + logger.info("FAULT: Freezing %s[%d] (pid=%d on %s) with SIGSTOP", + handle.role, handle.index, handle.pid, handle.host) + self._pm.freeze(handle) + + def unfreeze_process(self, handle: ProcessHandle) -> None: + """SIGCONT — resumes frozen process.""" + logger.info("FAULT: Unfreezing %s[%d] (pid=%d on %s) with SIGCONT", + handle.role, handle.index, handle.pid, handle.host) + self._pm.unfreeze(handle) + + def kill_multiple(self, handles: list[ProcessHandle], + simultaneous: bool = True) -> None: + """Kill multiple nodes (possibly on different hosts).""" + logger.info("FAULT: Killing %d processes %s", + len(handles), "simultaneously" if simultaneous else "sequentially") + if simultaneous: + for h in handles: + self._ssh.send_signal(h.host, h.pid, signal.SIGKILL) + else: + for h in handles: + self.kill_process(h) + + @contextmanager + def temporary_partition( + self, + handle: ProcessHandle, + peer_handles: list[ProcessHandle], + duration: float | None = None, + ): + """ + Context manager: create network partition via iptables, yield, then heal. + + In multi-machine mode, iptables rules are applied on both sides: + - On handle's host: DROP traffic to/from peer IPs + - On peer hosts: DROP traffic to/from handle's IP + + This creates a true bidirectional partition. + Requires root privileges on all involved hosts. + """ + rules = self._build_iptables_rules(handle, peer_handles) + try: + self._apply_iptables_rules(rules, add=True) + logger.info("FAULT: Network partition applied (%d rules)", len(rules)) + yield + finally: + self._apply_iptables_rules(rules, add=False) + logger.info("FAULT: Network partition healed (%d rules removed)", len(rules)) + + def _build_iptables_rules( + self, + handle: ProcessHandle, + peer_handles: list[ProcessHandle], + ) -> list[dict]: + """Build iptables rules to block traffic between handle and peers. + + For multi-machine, we use IP-based rules (not just port-based) since + traffic crosses network boundaries. Rules are applied on both sides. + """ + rules = [] + + for peer in peer_handles: + if handle.host == peer.host: + # Same host: use port-based rules (original behavior) + for tp in handle.ports.values(): + for pp in peer.ports.values(): + rules.append({ + "host": handle.host, + "chain": "OUTPUT", + "args": f"-p tcp --sport {tp} --dport {pp} -j DROP", + }) + rules.append({ + "host": handle.host, + "chain": "OUTPUT", + "args": f"-p tcp --sport {pp} --dport {tp} -j DROP", + }) + else: + # Different hosts: block by destination IP + ports + # On handle's host: block outgoing to peer + for pp in peer.ports.values(): + rules.append({ + "host": handle.host, + "chain": "OUTPUT", + "args": f"-p tcp -d {peer.ip} --dport {pp} -j DROP", + }) + # On handle's host: block incoming from peer + for tp in handle.ports.values(): + rules.append({ + "host": handle.host, + "chain": "INPUT", + "args": f"-p tcp -s {peer.ip} --dport {tp} -j DROP", + }) + # On peer's host: block outgoing to handle + for tp in handle.ports.values(): + rules.append({ + "host": peer.host, + "chain": "OUTPUT", + "args": f"-p tcp -d {handle.ip} --dport {tp} -j DROP", + }) + # On peer's host: block incoming from handle + for pp in peer.ports.values(): + rules.append({ + "host": peer.host, + "chain": "INPUT", + "args": f"-p tcp -s {handle.ip} --dport {pp} -j DROP", + }) + + return rules + + def _apply_iptables_rules(self, rules: list[dict], add: bool = True) -> None: + """Apply or remove iptables rules on the appropriate hosts.""" + action = "-A" if add else "-D" + for rule in rules: + host = rule["host"] + iptables_cmd = f"iptables {action} {rule['chain']} {rule['args']}" + success = self._ssh.run_iptables(host, f"{action} {rule['chain']} {rule['args']}") + if not success and add: + logger.error("iptables rule failed on %s: %s", host, iptables_cmd) + raise RuntimeError(f"iptables failed on {host}: {iptables_cmd}") diff --git a/tests/protocol_test/framework/log_parser.py b/tests/protocol_test/framework/log_parser.py new file mode 100644 index 0000000..be52806 --- /dev/null +++ b/tests/protocol_test/framework/log_parser.py @@ -0,0 +1,136 @@ +"""Parse CM/DS log files for state verification — supports local and remote hosts. + +In multi-machine mode, log files reside on remote hosts. LogParser reads them +via SshExecutor instead of direct file access. +""" + +import re +import time +from dataclasses import dataclass + +from .ssh_executor import SshExecutor + + +@dataclass +class LogEvent: + timestamp: str + level: str + message: str + raw_line: str + + +class LogParser: + """Parses SiMM log files for state verification events. + + Supports reading logs from both local and remote hosts via SshExecutor. + """ + + # Common log patterns in SiMM (based on spdlog format: [timestamp] [level] message) + LOG_LINE_RE = re.compile( + r'\[(\d{4}-\d{2}-\d{2}\s+\d{2}:\d{2}:\d{2}\.\d+)\]\s+' + r'\[(\w+)\]\s+' + r'(.*)' + ) + + def __init__(self, log_path: str, host: str, ssh_executor: SshExecutor): + """ + Args: + log_path: Path to the log file on the target host. + host: Hostname/IP where the log file resides. + ssh_executor: SSH executor for reading remote files. + """ + self._log_path = log_path + self._host = host + self._ssh = ssh_executor + self._last_offset = 0 + + def _read_all_lines(self) -> list[str]: + content = self._ssh.read_file(self._host, self._log_path) + if not content: + return [] + return content.splitlines() + + def _read_new_lines(self) -> list[str]: + """Read only new lines since last read.""" + content, new_offset = self._ssh.read_file_tail( + self._host, self._log_path, self._last_offset + ) + self._last_offset = new_offset + if not content: + return [] + return [line.rstrip() for line in content.splitlines()] + + def wait_for_pattern(self, pattern: str, timeout: float = 30, + poll_interval: float = 0.5) -> str | None: + """Tail log file waiting for a regex pattern match. Returns the matching line.""" + compiled = re.compile(pattern) + deadline = time.time() + timeout + while time.time() < deadline: + lines = self._read_new_lines() + for line in lines: + if compiled.search(line): + return line + time.sleep(poll_interval) + return None + + def count_pattern(self, pattern: str) -> int: + """Count occurrences of a pattern in the entire log.""" + compiled = re.compile(pattern) + return sum(1 for line in self._read_all_lines() if compiled.search(line)) + + def find_pattern_lines(self, pattern: str) -> list[str]: + """Return all lines matching a pattern.""" + compiled = re.compile(pattern) + return [line for line in self._read_all_lines() if compiled.search(line)] + + def find_handshake_events(self) -> list[str]: + """Find node handshake/registration log entries.""" + patterns = [ + r"[Hh]andshake", + r"AddNode", + r"NewNodeHandshake", + r"[Rr]egister.*[Cc]luster", + ] + combined = "|".join(f"({p})" for p in patterns) + return self.find_pattern_lines(combined) + + def find_heartbeat_timeout_events(self) -> list[str]: + """Find heartbeat timeout / node failure detection log entries.""" + patterns = [ + r"heartbeat.*timeout", + r"[Nn]ode.*[Ff]ailure", + r"mark.*[Dd]ead", + r"DEAD", + r"HandleNodeFailure", + ] + combined = "|".join(f"({p})" for p in patterns) + return self.find_pattern_lines(combined) + + def find_rebalance_events(self) -> list[str]: + """Find shard rebalance log entries.""" + patterns = [ + r"[Rr]ebalance", + r"[Rr]eassign.*[Ss]hard", + r"[Ss]hard.*[Mm]igrat", + r"RebalanceShardsAfterNodeFailure", + ] + combined = "|".join(f"({p})" for p in patterns) + return self.find_pattern_lines(combined) + + def find_node_status_changes(self) -> list[str]: + """Find node status transition log entries.""" + patterns = [ + r"UpdateNodeStatus", + r"RUNNING.*DEAD|DEAD.*RUNNING", + r"node.*status.*changed", + ] + combined = "|".join(f"({p})" for p in patterns) + return self.find_pattern_lines(combined) + + def contains(self, pattern: str) -> bool: + """Check if log contains at least one match of pattern.""" + return self.count_pattern(pattern) > 0 + + def get_all_text(self) -> str: + """Return full log text.""" + return self._ssh.read_file(self._host, self._log_path) diff --git a/tests/protocol_test/framework/port_allocator.py b/tests/protocol_test/framework/port_allocator.py new file mode 100644 index 0000000..78a6be0 --- /dev/null +++ b/tests/protocol_test/framework/port_allocator.py @@ -0,0 +1,79 @@ +"""Dynamic free port allocator for test processes — supports local and remote hosts.""" + +import socket +import threading + +from .ssh_executor import SshExecutor + + +class PortAllocator: + """Thread-safe free port allocator using bind-then-close probing. + + Supports allocating ports on both local and remote hosts. + Remote port probing is done via SSH (python3 one-liner on the remote). + """ + + def __init__(self, ssh_executor: SshExecutor | None = None): + self._ssh = ssh_executor + self._lock = threading.Lock() + # Track allocated ports per host to avoid collisions + self._allocated: dict[str, set[int]] = {} # {host: {ports}} + + def _get_host_set(self, host: str) -> set[int]: + if host not in self._allocated: + self._allocated[host] = set() + return self._allocated[host] + + def _find_free_port_local(self) -> int: + """Find a single free port on the local machine.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + s.bind(("", 0)) + return s.getsockname()[1] + + def _find_free_port_remote(self, host: str) -> int: + """Find a single free port on a remote host via SSH.""" + if self._ssh is None: + raise RuntimeError("SshExecutor required for remote port allocation") + port = self._ssh.find_free_port(host) + if port is None: + raise RuntimeError(f"Failed to find free port on {host}") + return port + + def _find_free_port(self, host: str) -> int: + """Find a free port, local or remote.""" + if self._ssh is None or self._ssh.is_local(host): + return self._find_free_port_local() + return self._find_free_port_remote(host) + + def allocate(self, count: int = 1, host: str = "127.0.0.1") -> list[int]: + """Return `count` currently-free ports on `host`, unique within this allocator.""" + ports = [] + with self._lock: + host_set = self._get_host_set(host) + for _ in range(count): + for _ in range(100): + port = self._find_free_port(host) + if port not in host_set: + host_set.add(port) + ports.append(port) + break + else: + raise RuntimeError( + f"Failed to allocate a free port on {host} after 100 attempts" + ) + return ports + + def allocate_cm_ports(self, host: str = "127.0.0.1") -> dict[str, int]: + """Allocate 3 ports for a Cluster Manager on the given host.""" + ports = self.allocate(3, host=host) + return {"intra": ports[0], "inter": ports[1], "admin": ports[2]} + + def allocate_ds_ports(self, host: str = "127.0.0.1") -> dict[str, int]: + """Allocate 3 ports for a Data Server on the given host.""" + ports = self.allocate(3, host=host) + return {"io": ports[0], "mgt": ports[1], "admin": ports[2]} + + def release_all(self): + with self._lock: + self._allocated.clear() diff --git a/tests/protocol_test/framework/process_manager.py b/tests/protocol_test/framework/process_manager.py new file mode 100644 index 0000000..2c9c22a --- /dev/null +++ b/tests/protocol_test/framework/process_manager.py @@ -0,0 +1,337 @@ +"""Process lifecycle management for CM and DS binaries — supports local and remote hosts. + +In multi-machine mode, processes are started on remote hosts via SSH (nohup). +The test runner (node A) sends commands over SSH to start/stop/signal processes +on the CM/DS nodes. +""" + +import logging +import signal +import subprocess +import time +from dataclasses import dataclass, field + +from .port_allocator import PortAllocator +from .ssh_executor import SshExecutor + +logger = logging.getLogger(__name__) + + +@dataclass +class ProcessHandle: + """Tracks a single spawned CM or DS process.""" + pid: int + process: subprocess.Popen | None # None for remote processes + role: str # "cm" or "ds" + index: int # ds index (0, 1, 2, ...), 0 for cm + ip: str # IP of the host running this process + host: str # host address for SSH (may == ip) + ports: dict[str, int] + log_path: str # path on the target host (str, not Path) + start_time: float + cmd_args: list[str] = field(default_factory=list) # for restart + extra_flags: dict = field(default_factory=dict) + binary_dir: str = "" # binary dir on the target host + + @property + def addr_str(self) -> str: + """Address string used by CM to identify this node.""" + if self.role == "ds": + return f"{self.ip}:{self.ports['io']}" + return f"{self.ip}:{self.ports['inter']}" + + +class ProcessManager: + """Manages CM and DS process lifecycle across local and remote hosts.""" + + def __init__(self, default_binary_dir: str, default_log_dir: str, + port_allocator: PortAllocator, ssh_executor: SshExecutor): + self._default_binary_dir = str(default_binary_dir) + self._default_log_dir = str(default_log_dir) + self._port_allocator = port_allocator + self._ssh = ssh_executor + self._handles: list[ProcessHandle] = [] + self._ds_index = 0 + + def _resolve_binary(self, binary_dir: str, name: str) -> str: + """Return full path to a binary on the target host.""" + return f"{binary_dir}/{name}" + + def kill_existing_by_name(self, host: str, name: str) -> int: + """Find and kill all processes matching binary name on host. + + Returns the number of processes killed. + """ + # Get PIDs of matching processes (exclude grep itself) + result = self._ssh.run( + host, + f"pgrep -x {name}", + timeout=5, check=False, + ) + if not result or not result.strip(): + return 0 + + pids = [int(p) for p in result.strip().split("\n") if p.strip()] + for pid in pids: + self._ssh.send_signal(host, pid, signal.SIGKILL) + logger.info("Killed existing %s (pid=%d) on %s", name, pid, host) + + # Wait briefly for all to die + deadline = time.time() + 5 + while time.time() < deadline: + if all(not self._ssh.is_process_alive(host, p) for p in pids): + break + time.sleep(0.5) + + return len(pids) + + def start_cluster_manager( + self, + host: str = "127.0.0.1", + ip: str | None = None, + ports: dict[str, int] | None = None, + binary_dir: str | None = None, + log_dir: str | None = None, + cm_cluster_init_grace_period_inSecs: int = 5, + cm_heartbeat_timeout_inSecs: int = 10, + cm_heartbeat_bg_scan_interval_inSecs: int = 2, + shard_total_num: int = 64, + cm_deferred_reshard_enabled: bool = True, + cm_deferred_reshard_window_inSecs: int = 120, + extra_flags: dict | None = None, + ) -> ProcessHandle: + if ip is None: + ip = host + binary_dir = binary_dir or self._default_binary_dir + log_dir = log_dir or self._default_log_dir + + if ports is None: + ports = self._port_allocator.allocate_cm_ports(host=host) + + log_path = f"{log_dir}/cm.log" + default_log_path = f"{log_dir}/cm_default.log" + cm_binary = self._resolve_binary(binary_dir, "cluster_manager") + + cmd_parts = [ + cm_binary, + f"--cm_rpc_intra_port={ports['intra']}", + f"--cm_rpc_inter_port={ports['inter']}", + f"--cm_rpc_admin_port={ports['admin']}", + f"--cm_cluster_init_grace_period_inSecs={cm_cluster_init_grace_period_inSecs}", + f"--cm_heartbeat_timeout_inSecs={cm_heartbeat_timeout_inSecs}", + f"--cm_heartbeat_bg_scan_interval_inSecs={cm_heartbeat_bg_scan_interval_inSecs}", + f"--shard_total_num={shard_total_num}", + f"--cm_deferred_reshard_enabled={'true' if cm_deferred_reshard_enabled else 'false'}", + f"--cm_deferred_reshard_window_inSecs={cm_deferred_reshard_window_inSecs}", + f"--cm_log_file={log_path}", + f"--default_logging_file={default_log_path}", + ] + if extra_flags: + for k, v in extra_flags.items(): + cmd_parts.append(f"--{k}={v}") + + cmd_str = " ".join(cmd_parts) + logger.info("Starting CM on %s: %s", host, cmd_str) + + # Ensure log directory exists on target host + self._ssh.run(host, f"mkdir -p {log_dir}", timeout=5, check=False) + + pid = self._start_process(host, cmd_str) + + handle = ProcessHandle( + pid=pid, + process=None, # remote — no Popen object + role="cm", + index=0, + ip=ip, + host=host, + ports=ports, + log_path=log_path, + start_time=time.time(), + cmd_args=cmd_parts, + extra_flags=extra_flags or {}, + binary_dir=binary_dir, + ) + self._handles.append(handle) + logger.info("CM started on %s: pid=%d ports=%s", host, pid, ports) + return handle + + def start_data_server( + self, + cm_ip: str, + cm_inter_port: int, + host: str = "127.0.0.1", + ip: str | None = None, + ports: dict[str, int] | None = None, + binary_dir: str | None = None, + log_dir: str | None = None, + heartbeat_cooldown_sec: int = 2, + register_cooldown_sec: int = 2, + cm_hb_tolerance_count: int = 5, + cm_connect_retry_interval_sec: int = 1, + memory_limit_bytes: int = 1 << 30, + ds_logical_node_id: str = "", + extra_flags: dict | None = None, + ) -> ProcessHandle: + if ip is None: + ip = host + binary_dir = binary_dir or self._default_binary_dir + log_dir = log_dir or self._default_log_dir + + if ports is None: + ports = self._port_allocator.allocate_ds_ports(host=host) + + idx = self._ds_index + self._ds_index += 1 + + log_path = f"{log_dir}/ds_{idx}.log" + default_log_path = f"{log_dir}/ds_{idx}_default.log" + ds_binary = self._resolve_binary(binary_dir, "data_server") + + cmd_parts = [ + ds_binary, + f"--cm_primary_node_ip={cm_ip}", + f"--cm_rpc_inter_port={cm_inter_port}", + f"--io_service_port={ports['io']}", + f"--mgt_service_port={ports['mgt']}", + f"--ds_rpc_admin_port={ports['admin']}", + f"--heartbeat_cooldown_sec={heartbeat_cooldown_sec}", + f"--register_cooldown_sec={register_cooldown_sec}", + f"--cm_hb_tolerance_count={cm_hb_tolerance_count}", + f"--cm_connect_retry_interval_sec={cm_connect_retry_interval_sec}", + f"--memory_limit_bytes={memory_limit_bytes}", + f"--ds_log_file={log_path}", + f"--default_logging_file={default_log_path}", + ] + if ds_logical_node_id: + cmd_parts.append(f"--ds_logical_node_id={ds_logical_node_id}") + if extra_flags: + for k, v in extra_flags.items(): + cmd_parts.append(f"--{k}={v}") + + cmd_str = " ".join(cmd_parts) + logger.info("Starting DS[%d] on %s: %s", idx, host, cmd_str) + + self._ssh.run(host, f"mkdir -p {log_dir}", timeout=5, check=False) + pid = self._start_process(host, cmd_str) + + handle = ProcessHandle( + pid=pid, + process=None, + role="ds", + index=idx, + ip=ip, + host=host, + ports=ports, + log_path=log_path, + start_time=time.time(), + cmd_args=cmd_parts, + extra_flags=extra_flags or {}, + binary_dir=binary_dir, + ) + self._handles.append(handle) + logger.info("DS[%d] started on %s: pid=%d ports=%s", idx, host, pid, ports) + return handle + + def _start_process(self, host: str, cmd: str) -> int: + """Start a process on host. Returns remote PID.""" + pid = self._ssh.run_background(host, cmd) + if pid is None: + raise RuntimeError(f"Failed to start process on {host}: {cmd}") + return pid + + def kill(self, handle: ProcessHandle, sig: int = signal.SIGKILL) -> None: + """Send a signal to the process (local or remote).""" + if not self.is_alive(handle): + logger.warning("Process %d on %s already dead", handle.pid, handle.host) + return + success = self._ssh.send_signal(handle.host, handle.pid, sig) + if success: + logger.info("Sent signal %d to pid %d on %s (%s[%d])", + sig, handle.pid, handle.host, handle.role, handle.index) + else: + logger.warning("Failed to send signal %d to pid %d on %s", + sig, handle.pid, handle.host) + + def stop(self, handle: ProcessHandle, timeout: float = 10) -> None: + """Send SIGTERM and wait for graceful exit.""" + self.kill(handle, signal.SIGTERM) + self.wait_for_exit(handle, timeout=timeout) + + def freeze(self, handle: ProcessHandle) -> None: + """Send SIGSTOP — simulates process hang.""" + self.kill(handle, signal.SIGSTOP) + + def unfreeze(self, handle: ProcessHandle) -> None: + """Send SIGCONT — resumes frozen process.""" + self.kill(handle, signal.SIGCONT) + + def is_alive(self, handle: ProcessHandle) -> bool: + """Check if process is still running (local or remote).""" + return self._ssh.is_process_alive(handle.host, handle.pid) + + def wait_for_exit(self, handle: ProcessHandle, timeout: float = 30) -> bool: + """Wait for process to exit. Returns True if exited, False on timeout.""" + deadline = time.time() + timeout + while time.time() < deadline: + if not self.is_alive(handle): + return True + time.sleep(0.5) + logger.warning("Process %d on %s did not exit within %.1fs", + handle.pid, handle.host, timeout) + return False + + def restart(self, handle: ProcessHandle) -> ProcessHandle: + """Kill and restart a process with the same arguments.""" + if self.is_alive(handle): + self.kill(handle, signal.SIGKILL) + self.wait_for_exit(handle, timeout=5) + + # Remove old handle + self._handles = [h for h in self._handles if h.pid != handle.pid] + + cmd_str = " ".join(handle.cmd_args) + logger.info("Restarting %s[%d] on %s with same args", + handle.role, handle.index, handle.host) + pid = self._start_process(handle.host, cmd_str) + + new_handle = ProcessHandle( + pid=pid, + process=None, + role=handle.role, + index=handle.index, + ip=handle.ip, + host=handle.host, + ports=handle.ports, + log_path=handle.log_path, + start_time=time.time(), + cmd_args=handle.cmd_args, + extra_flags=handle.extra_flags, + binary_dir=handle.binary_dir, + ) + self._handles.append(new_handle) + logger.info("%s[%d] restarted on %s: pid=%d", + handle.role, handle.index, handle.host, pid) + return new_handle + + def get_log(self, handle: ProcessHandle) -> str: + """Return log file contents from the target host.""" + return self._ssh.read_file(handle.host, handle.log_path) + + def cleanup_all(self) -> None: + """Kill all managed processes across all hosts. Called in fixture teardown.""" + for handle in self._handles: + if self.is_alive(handle): + # Unfreeze first in case it was SIGSTOP'd + self._ssh.send_signal(handle.host, handle.pid, signal.SIGCONT) + self._ssh.send_signal(handle.host, handle.pid, signal.SIGKILL) + + # Wait briefly for all to die + deadline = time.time() + 5 + while time.time() < deadline: + if all(not self.is_alive(h) for h in self._handles): + break + time.sleep(0.5) + + self._handles.clear() + logger.info("All processes cleaned up") diff --git a/tests/protocol_test/framework/scenario_runner.py b/tests/protocol_test/framework/scenario_runner.py new file mode 100644 index 0000000..06bcdcb --- /dev/null +++ b/tests/protocol_test/framework/scenario_runner.py @@ -0,0 +1,337 @@ +"""Scenario-driven test orchestration engine. + +The ScenarioRunner reads YAML scenario definitions and executes them: + 1. Start cluster (from cluster config) + 2. Wait for cluster ready + 3. For each fault step: + a. Inject fault (with optional delay) + b. Wait for expected state transitions + 4. Run validation checks + 5. Teardown + +This makes adding new test scenarios declarative — write YAML, not Python. + +YAML scenario format: + cluster: + # ... cluster config (same as before) + faults: + - type: sigkill | sigterm | sigstop | sigcont + target: "ds:0" | "cm" + delay_after_ready_sec: 5 + duration_sec: 15 # for sigstop only + validations: + - type: node_status + target: "ds:0" + expected: DEAD + within_sec: 20 # max wait time + - type: shard_total + expected: 64 + - type: shard_balance + max_imbalance_ratio: 0.3 + - type: no_orphaned_shards + - type: node_has_no_shards + target: "ds:0" + - type: alive_node_count + expected: 2 + - type: ds_status # DS admin RPC check + target: "ds:0" + field: cm_ready + expected: "false" + - type: all_nodes_running +""" + +import logging +import time +from dataclasses import dataclass +from pathlib import Path + +from .cluster import SimmCluster +from .config import ( + ClusterConfig, FaultConfig, + load_yaml, deep_merge, dict_to_cluster_config, dict_to_fault_configs, +) + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Validation step definitions +# --------------------------------------------------------------------------- + +@dataclass +class ValidationStep: + """A single validation check to run after faults are injected.""" + type: str # validation type name + target: str = "" # e.g. "ds:0", "cm" + expected: str = "" # expected value (type-dependent) + field: str = "" # for ds_status: which field to check + within_sec: float = 60 # max wait time for condition-based checks + max_imbalance_ratio: float = 0.3 # for shard_balance + + +def _parse_validation_steps(data: dict) -> list[ValidationStep]: + """Parse validation steps from YAML dict.""" + steps = [] + for v in data.get("validations", []): + steps.append(ValidationStep( + type=v["type"], + target=v.get("target", ""), + expected=str(v.get("expected", "")), + field=v.get("field", ""), + within_sec=v.get("within_sec", 60), + max_imbalance_ratio=v.get("max_imbalance_ratio", 0.3), + )) + return steps + + +# --------------------------------------------------------------------------- +# Fault executor — maps YAML fault type to FaultInjector calls +# --------------------------------------------------------------------------- + +_FAULT_REGISTRY: dict[str, callable] = {} + + +def register_fault(name: str): + """Decorator to register a fault executor function.""" + def decorator(fn): + _FAULT_REGISTRY[name] = fn + return fn + return decorator + + +@register_fault("sigkill") +def _exec_sigkill(cluster: SimmCluster, fault: FaultConfig): + for target_str in fault.targets: + handle = _resolve_target(cluster, target_str) + cluster.fault_injector.kill_process(handle) + + +@register_fault("sigterm") +def _exec_sigterm(cluster: SimmCluster, fault: FaultConfig): + for target_str in fault.targets: + handle = _resolve_target(cluster, target_str) + cluster.fault_injector.graceful_stop(handle) + + +@register_fault("sigstop") +def _exec_sigstop(cluster: SimmCluster, fault: FaultConfig): + for target_str in fault.targets: + handle = _resolve_target(cluster, target_str) + cluster.fault_injector.freeze_process(handle) + + +@register_fault("sigcont") +def _exec_sigcont(cluster: SimmCluster, fault: FaultConfig): + for target_str in fault.targets: + handle = _resolve_target(cluster, target_str) + cluster.fault_injector.unfreeze_process(handle) + + +@register_fault("kill_multiple") +def _exec_kill_multiple(cluster: SimmCluster, fault: FaultConfig): + handles = [_resolve_target(cluster, t) for t in fault.targets] + cluster.fault_injector.kill_multiple(handles, simultaneous=True) + + +@register_fault("restart_cm") +def _exec_restart_cm(cluster: SimmCluster, fault: FaultConfig): + cluster.fault_injector.kill_process(cluster.cm) + time.sleep(fault.duration_sec or 3) + cluster.restart_cm() + + +def _resolve_target(cluster: SimmCluster, target_str: str): + """Resolve 'ds:0', 'ds:1', 'cm' to a ProcessHandle.""" + if target_str == "cm": + return cluster.cm + if target_str.startswith("ds:"): + idx = int(target_str.split(":")[1]) + return cluster.get_ds_handle(idx) + raise ValueError(f"Unknown target: {target_str}") + + +# --------------------------------------------------------------------------- +# Validation executor — maps YAML validation type to observer assertions +# --------------------------------------------------------------------------- + +_VALIDATION_REGISTRY: dict[str, callable] = {} + + +def register_validation(name: str): + """Decorator to register a validation executor function.""" + def decorator(fn): + _VALIDATION_REGISTRY[name] = fn + return fn + return decorator + + +@register_validation("node_status") +def _validate_node_status(cluster: SimmCluster, step: ValidationStep): + handle = _resolve_target(cluster, step.target) + assert cluster.observer.wait_for_node_status( + handle.addr_str, step.expected, timeout=step.within_sec + ), (f"Node {handle.addr_str} did not reach status {step.expected} " + f"within {step.within_sec}s") + + +@register_validation("alive_node_count") +def _validate_alive_count(cluster: SimmCluster, step: ValidationStep): + expected = int(step.expected) + assert cluster.observer.wait_for_node_count( + expected, timeout=step.within_sec + ), f"Alive node count did not reach {expected} within {step.within_sec}s" + + +@register_validation("all_nodes_running") +def _validate_all_running(cluster: SimmCluster, step: ValidationStep): + cluster.observer.assert_all_nodes_running() + + +@register_validation("shard_total") +def _validate_shard_total(cluster: SimmCluster, step: ValidationStep): + expected = int(step.expected) + cluster.observer.assert_total_shard_count(expected) + + +@register_validation("shard_balance") +def _validate_shard_balance(cluster: SimmCluster, step: ValidationStep): + cluster.observer.assert_shard_balance( + max_imbalance_ratio=step.max_imbalance_ratio + ) + + +@register_validation("no_orphaned_shards") +def _validate_no_orphaned(cluster: SimmCluster, step: ValidationStep): + cluster.observer.wait_for_rebalance_complete(timeout=step.within_sec) + cluster.observer.assert_no_orphaned_shards() + + +@register_validation("node_has_no_shards") +def _validate_node_no_shards(cluster: SimmCluster, step: ValidationStep): + handle = _resolve_target(cluster, step.target) + cluster.observer.assert_node_has_no_shards(handle.addr_str) + + +@register_validation("ds_status") +def _validate_ds_status(cluster: SimmCluster, step: ValidationStep): + """Validate a DS internal status field via admin RPC. + target: 'ds:0', field: 'cm_ready'/'is_registered'/'heartbeat_failure_count', + expected: the expected value as string. + """ + handle = _resolve_target(cluster, step.target) + status = cluster.observer.get_ds_status(handle.host, handle.pid) + actual = status.get(step.field, "") + assert actual == step.expected, ( + f"DS {handle.addr_str} {step.field}={actual}, expected {step.expected}" + ) + + +# --------------------------------------------------------------------------- +# ScenarioRunner +# --------------------------------------------------------------------------- + +class ScenarioRunner: + """ + Orchestration engine that executes a test scenario from YAML definition. + + Usage: + runner = ScenarioRunner() + runner.run_scenario("scenarios/clusters/small.yaml", + "scenarios/faults/kill_one_ds.yaml", + "scenarios/validations/check_rebalance.yaml") + + Or from a single combined YAML: + runner.run_scenario("scenarios/full/kill_ds_and_verify.yaml") + + Or programmatically: + runner.run(cluster_config, fault_configs, validation_steps) + """ + + def __init__(self, binary_dir: str | None = None, log_dir: str | None = None): + self._binary_dir = binary_dir + self._log_dir = log_dir + + def run_scenario(self, *yaml_paths: str | Path) -> None: + """Load and execute a scenario from one or more YAML files (merged in order).""" + merged: dict = {} + for p in yaml_paths: + data = load_yaml(Path(p)) + merged = deep_merge(merged, data) + + cluster_config = dict_to_cluster_config(merged) + fault_configs = dict_to_fault_configs(merged) + validation_steps = _parse_validation_steps(merged) + + self.run(cluster_config, fault_configs, validation_steps) + + def run(self, cluster_config: ClusterConfig, + fault_configs: list[FaultConfig], + validation_steps: list[ValidationStep]) -> None: + """Execute a scenario programmatically.""" + cluster = SimmCluster( + cluster_config, + log_dir=self._log_dir, + binary_dir=self._binary_dir, + ) + + try: + # Phase 1: Start cluster + logger.info("=== Phase 1: Starting cluster ===") + cluster.start() + cluster.wait_ready() + + # Phase 2: Inject faults + logger.info("=== Phase 2: Injecting %d fault(s) ===", len(fault_configs)) + for i, fault in enumerate(fault_configs): + if fault.delay_after_sec > 0: + logger.info("Waiting %.1fs before fault #%d", fault.delay_after_sec, i) + time.sleep(fault.delay_after_sec) + + executor = _FAULT_REGISTRY.get(fault.fault_type) + if executor is None: + raise ValueError( + f"Unknown fault type: {fault.fault_type}. " + f"Registered: {list(_FAULT_REGISTRY.keys())}" + ) + + logger.info("Injecting fault #%d: %s → %s", + i, fault.fault_type, fault.targets) + executor(cluster, fault) + + # For sigstop with duration: wait, then sigcont + if fault.fault_type == "sigstop" and fault.duration_sec: + logger.info("Waiting %.1fs then unfreezing", fault.duration_sec) + time.sleep(fault.duration_sec) + for target_str in fault.targets: + handle = _resolve_target(cluster, target_str) + cluster.fault_injector.unfreeze_process(handle) + + # Phase 3: Run validations + logger.info("=== Phase 3: Running %d validation(s) ===", len(validation_steps)) + for i, step in enumerate(validation_steps): + validator = _VALIDATION_REGISTRY.get(step.type) + if validator is None: + raise ValueError( + f"Unknown validation type: {step.type}. " + f"Registered: {list(_VALIDATION_REGISTRY.keys())}" + ) + + logger.info("Validation #%d: %s (target=%s, expected=%s)", + i, step.type, step.target, step.expected) + validator(cluster, step) + logger.info("Validation #%d: PASSED", i) + + logger.info("=== Scenario PASSED ===") + + finally: + cluster.teardown() + + @staticmethod + def list_fault_types() -> list[str]: + """List all registered fault types.""" + return list(_FAULT_REGISTRY.keys()) + + @staticmethod + def list_validation_types() -> list[str]: + """List all registered validation types.""" + return list(_VALIDATION_REGISTRY.keys()) diff --git a/tests/protocol_test/framework/ssh_executor.py b/tests/protocol_test/framework/ssh_executor.py new file mode 100644 index 0000000..0619e6c --- /dev/null +++ b/tests/protocol_test/framework/ssh_executor.py @@ -0,0 +1,238 @@ +"""SSH command execution layer for multi-machine testing. + +Provides a unified interface for running commands either locally or on remote +hosts via passwordless SSH. The test runner (node A) uses this to start/stop +processes, send signals, read logs, and manage iptables on remote CM/DS nodes. +""" + +import logging +import subprocess +from dataclasses import dataclass + +logger = logging.getLogger(__name__) + + +@dataclass +class SshConfig: + """SSH connection settings.""" + user: str = "" # SSH user; empty = current user + port: int = 22 + connect_timeout: int = 5 + options: dict[str, str] | None = None # extra SSH options + + +class SshExecutor: + """ + Executes shell commands on local or remote hosts via SSH. + + Usage: + exe = SshExecutor() + # Local + result = exe.run("127.0.0.1", "ls /tmp") + # Remote + result = exe.run("10.0.0.2", "kill -9 12345") + """ + + # Hosts treated as local (no SSH needed) + LOCAL_HOSTS = {"127.0.0.1", "localhost", "::1"} + + def __init__(self, ssh_config: SshConfig | None = None, + local_host: str | None = None): + """ + Args: + ssh_config: SSH connection settings for remote hosts. + local_host: The IP/hostname of the machine running tests. + Commands targeting this host run locally without SSH. + If None, only 127.0.0.1/localhost are treated as local. + """ + self._ssh_config = ssh_config or SshConfig() + self._local_hosts = set(self.LOCAL_HOSTS) + if local_host: + self._local_hosts.add(local_host) + + def is_local(self, host: str) -> bool: + """Check if host is the local machine.""" + return host in self._local_hosts + + def _build_ssh_cmd(self, host: str) -> list[str]: + """Build the SSH prefix command.""" + cmd = ["ssh"] + # Common options for non-interactive batch mode + cmd += [ + "-o", "BatchMode=yes", + "-o", "StrictHostKeyChecking=no", + "-o", f"ConnectTimeout={self._ssh_config.connect_timeout}", + ] + if self._ssh_config.port != 22: + cmd += ["-p", str(self._ssh_config.port)] + if self._ssh_config.options: + for k, v in self._ssh_config.options.items(): + cmd += ["-o", f"{k}={v}"] + + target = host + if self._ssh_config.user: + target = f"{self._ssh_config.user}@{host}" + cmd.append(target) + return cmd + + def run(self, host: str, command: str, + timeout: float = 30, check: bool = True) -> subprocess.CompletedProcess: + """ + Run a shell command on the target host. + + Args: + host: Target hostname or IP. + command: Shell command to execute. + timeout: Timeout in seconds. + check: If True, raise on non-zero exit code. + + Returns: + CompletedProcess with stdout/stderr. + + Raises: + subprocess.CalledProcessError: If check=True and command fails. + SshError: If SSH connection fails. + """ + if self.is_local(host): + full_cmd = ["bash", "-c", command] + else: + full_cmd = self._build_ssh_cmd(host) + [command] + + logger.debug("Exec on %s: %s", host, command) + try: + result = subprocess.run( + full_cmd, + capture_output=True, + text=True, + timeout=timeout, + ) + if check and result.returncode != 0: + raise subprocess.CalledProcessError( + result.returncode, full_cmd, + output=result.stdout, stderr=result.stderr, + ) + return result + except subprocess.TimeoutExpired: + raise SshError(f"Command timed out on {host}: {command}") + + def run_background(self, host: str, command: str) -> int | None: + """ + Start a command in the background on a remote host via nohup. + Returns the remote PID if parseable. + + The command is launched with nohup and stdout/stderr redirected, + so it survives the SSH session closing. + """ + # Wrap command to background and echo PID + bg_cmd = f"nohup {command} > /dev/null 2>&1 & echo $!" + result = self.run(host, bg_cmd, timeout=10) + pid_str = result.stdout.strip() + if pid_str.isdigit(): + pid = int(pid_str) + logger.info("Background process started on %s: pid=%d", host, pid) + return pid + logger.warning("Could not parse PID from: %s", pid_str) + return None + + def send_signal(self, host: str, pid: int, sig: int) -> bool: + """Send a signal to a remote process.""" + try: + self.run(host, f"kill -{sig} {pid}", timeout=5) + return True + except (subprocess.CalledProcessError, SshError): + logger.warning("Failed to send signal %d to pid %d on %s", sig, pid, host) + return False + + def is_process_alive(self, host: str, pid: int) -> bool: + """Check if a remote process is still running.""" + try: + result = self.run(host, f"kill -0 {pid}", timeout=5, check=False) + return result.returncode == 0 + except SshError as e: + logger.debug("is_process_alive check failed on %s pid=%d: %s", host, pid, e) + return False + + def read_file(self, host: str, path: str) -> str: + """Read a file from a remote host.""" + try: + result = self.run(host, f"cat {path}", timeout=15, check=False) + return result.stdout if result.returncode == 0 else "" + except SshError as e: + logger.debug("read_file failed on %s path=%s: %s", host, path, e) + return "" + + def read_file_tail(self, host: str, path: str, offset: int = 0) -> tuple[str, int]: + """ + Read new content from a file starting at byte offset. + Returns (new_content, new_offset). + """ + # Use dd to skip to offset, then read the rest + cmd = (f"test -f {path} && " + f"dd if={path} bs=1 skip={offset} 2>/dev/null; " + f"wc -c < {path} 2>/dev/null || echo 0") + try: + result = self.run(host, cmd, timeout=15, check=False) + lines = result.stdout.rsplit("\n", 1) + if len(lines) == 2: + content = lines[0] + try: + new_offset = int(lines[1].strip()) + except ValueError: + new_offset = offset + len(content) + else: + content = result.stdout + new_offset = offset + len(content) + return content, new_offset + except SshError as e: + logger.debug("read_file_tail failed on %s path=%s: %s", host, path, e) + return "", offset + + def file_exists(self, host: str, path: str) -> bool: + """Check if a file exists on a remote host.""" + try: + result = self.run(host, f"test -f {path}", timeout=5, check=False) + return result.returncode == 0 + except SshError as e: + logger.debug("file_exists check failed on %s path=%s: %s", host, path, e) + return False + + def find_free_port(self, host: str) -> int | None: + """Find a free port on a remote host.""" + cmd = ("python3 -c \"" + "import socket; s=socket.socket(); " + "s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1); " + "s.bind(('', 0)); " + "print(s.getsockname()[1]); " + "s.close()\"") + try: + result = self.run(host, cmd, timeout=5) + port_str = result.stdout.strip() + if port_str.isdigit(): + return int(port_str) + except (subprocess.CalledProcessError, SshError) as e: + logger.debug("find_free_port failed on %s: %s", host, e) + return None + + def run_iptables(self, host: str, args: str) -> bool: + """Run an iptables command on a remote host.""" + try: + self.run(host, f"iptables {args}", timeout=5) + return True + except (subprocess.CalledProcessError, SshError) as e: + logger.warning("iptables failed on %s: %s", host, e) + return False + + def check_connectivity(self, host: str) -> bool: + """Verify SSH connectivity to a host.""" + try: + result = self.run(host, "echo ok", timeout=self._ssh_config.connect_timeout + 2, + check=False) + return result.returncode == 0 and "ok" in result.stdout + except SshError as e: + logger.debug("check_connectivity failed for %s: %s", host, e) + return False + + +class SshError(Exception): + """Raised when SSH operation fails.""" + pass diff --git a/tests/protocol_test/framework/tests/__init__.py b/tests/protocol_test/framework/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/protocol_test/framework/tests/test_admin_client.py b/tests/protocol_test/framework/tests/test_admin_client.py new file mode 100644 index 0000000..52add18 --- /dev/null +++ b/tests/protocol_test/framework/tests/test_admin_client.py @@ -0,0 +1,121 @@ +"""Unit tests for admin_client.py — tabulate output parsing.""" + +from framework.admin_client import AdminClient + + +# Typical simm_ctl_admin output formats +NODE_LIST_OUTPUT = """\ ++-------------------+---------+ +| Address | Status | ++-------------------+---------+ +| 10.0.0.2:40000 | RUNNING | +| 10.0.0.3:40000 | RUNNING | +| 10.0.0.4:40000 | DEAD | ++-------------------+---------+ +""" + +NODE_LIST_VERBOSE_OUTPUT = """\ ++-------------------+---------+------------+----------+----------+ +| Address | Status | MemTotal | MemFree | MemUsed | ++-------------------+---------+------------+----------+----------+ +| 10.0.0.2:40000 | RUNNING | 1024 | 512 | 512 | +| 10.0.0.3:40000 | RUNNING | 2048 | 1024 | 1024 | ++-------------------+---------+------------+----------+----------+ +""" + +SHARD_LIST_OUTPUT = """\ ++-------------------+-------------+ +| Address | Shard Count | ++-------------------+-------------+ +| 10.0.0.2:40000 | 32 | +| 10.0.0.3:40000 | 32 | ++-------------------+-------------+ +""" + +SHARD_LIST_VERBOSE_OUTPUT = """\ ++-------------------+-------------------+ +| Address | Shard IDs | ++-------------------+-------------------+ +| 10.0.0.2:40000 | 0, 1, 2, 3 | +| 10.0.0.3:40000 | 4, 5, 6, 7 | ++-------------------+-------------------+ +""" + +DS_STATUS_OUTPUT = """\ ++----------------------------+-------+ +| Field | Value | ++----------------------------+-------+ +| is_registered | true | +| cm_ready | true | +| heartbeat_failure_count | 0 | ++----------------------------+-------+ +""" + +FLAGS_LIST_OUTPUT = """\ ++------------------------------------+---------+ +| Flag Name | Value | ++------------------------------------+---------+ +| cm_heartbeat_timeout_inSecs | 10 | +| heartbeat_cooldown_sec | 2 | ++------------------------------------+---------+ +""" + +FLAGS_GET_OUTPUT = """\ ++------------+---------+ +| Key | Result | ++------------+---------+ +| VALUE | 10 | ++------------+---------+ +""" + + +class TestParseTabulateRows: + + def test_node_list(self): + rows = AdminClient._parse_tabulate_rows(NODE_LIST_OUTPUT) + assert len(rows) == 4 # header + 3 data rows + assert rows[0] == ["Address", "Status"] + assert rows[1] == ["10.0.0.2:40000", "RUNNING"] + assert rows[3] == ["10.0.0.4:40000", "DEAD"] + + def test_shard_list(self): + rows = AdminClient._parse_tabulate_rows(SHARD_LIST_OUTPUT) + assert len(rows) == 3 # header + 2 data + assert rows[1] == ["10.0.0.2:40000", "32"] + + def test_shard_list_verbose(self): + rows = AdminClient._parse_tabulate_rows(SHARD_LIST_VERBOSE_OUTPUT) + assert len(rows) == 3 + assert rows[1] == ["10.0.0.2:40000", "0, 1, 2, 3"] + + def test_ds_status(self): + rows = AdminClient._parse_tabulate_rows(DS_STATUS_OUTPUT) + assert len(rows) == 4 # header + 3 fields + assert rows[1] == ["is_registered", "true"] + assert rows[2] == ["cm_ready", "true"] + assert rows[3] == ["heartbeat_failure_count", "0"] + + def test_flags_list(self): + rows = AdminClient._parse_tabulate_rows(FLAGS_LIST_OUTPUT) + assert len(rows) == 3 + assert rows[1] == ["cm_heartbeat_timeout_inSecs", "10"] + + def test_flags_get(self): + rows = AdminClient._parse_tabulate_rows(FLAGS_GET_OUTPUT) + assert len(rows) == 2 + assert rows[1] == ["VALUE", "10"] + + def test_empty_output(self): + rows = AdminClient._parse_tabulate_rows("") + assert rows == [] + + def test_separator_only(self): + rows = AdminClient._parse_tabulate_rows("+---+---+\n+---+---+\n") + assert rows == [] + + def test_node_list_verbose(self): + rows = AdminClient._parse_tabulate_rows(NODE_LIST_VERBOSE_OUTPUT) + assert len(rows) == 3 + assert rows[1][0] == "10.0.0.2:40000" + assert rows[1][2] == "1024" + assert rows[1][4] == "512" diff --git a/tests/protocol_test/framework/tests/test_config.py b/tests/protocol_test/framework/tests/test_config.py new file mode 100644 index 0000000..3b95001 --- /dev/null +++ b/tests/protocol_test/framework/tests/test_config.py @@ -0,0 +1,279 @@ +"""Unit tests for config.py — YAML loading, deep merge, config parsing.""" + +from pathlib import Path + +from framework.config import ( + ClusterConfig, + FaultConfig, + HostConfig, + deep_merge, + dict_to_cluster_config, + dict_to_fault_configs, + load_yaml, + merge_yaml_files, +) + + +class TestDeepMerge: + + def test_flat_override(self): + base = {"a": 1, "b": 2} + override = {"b": 3, "c": 4} + result = deep_merge(base, override) + assert result == {"a": 1, "b": 3, "c": 4} + + def test_nested_merge(self): + base = {"x": {"a": 1, "b": 2}} + override = {"x": {"b": 3, "c": 4}} + result = deep_merge(base, override) + assert result == {"x": {"a": 1, "b": 3, "c": 4}} + + def test_override_replaces_non_dict_with_dict(self): + base = {"x": 1} + override = {"x": {"nested": True}} + result = deep_merge(base, override) + assert result == {"x": {"nested": True}} + + def test_override_replaces_dict_with_non_dict(self): + base = {"x": {"nested": True}} + override = {"x": 42} + result = deep_merge(base, override) + assert result == {"x": 42} + + def test_does_not_mutate_inputs(self): + base = {"x": {"a": 1}} + override = {"x": {"b": 2}} + base_copy = {"x": {"a": 1}} + deep_merge(base, override) + assert base == base_copy + + def test_empty_base(self): + result = deep_merge({}, {"a": 1}) + assert result == {"a": 1} + + def test_empty_override(self): + result = deep_merge({"a": 1}, {}) + assert result == {"a": 1} + + def test_both_empty(self): + result = deep_merge({}, {}) + assert result == {} + + def test_three_level_nesting(self): + base = {"a": {"b": {"c": 1, "d": 2}}} + override = {"a": {"b": {"c": 99}}} + result = deep_merge(base, override) + assert result == {"a": {"b": {"c": 99, "d": 2}}} + + +class TestLoadYaml: + + def test_load_valid_yaml(self, tmp_path): + f = tmp_path / "test.yaml" + f.write_text("cluster:\n shard_total_num: 128\n") + data = load_yaml(f) + assert data == {"cluster": {"shard_total_num": 128}} + + def test_load_empty_yaml(self, tmp_path): + f = tmp_path / "empty.yaml" + f.write_text("") + data = load_yaml(f) + assert data == {} + + def test_load_yaml_with_only_comment(self, tmp_path): + f = tmp_path / "comment.yaml" + f.write_text("# just a comment\n") + data = load_yaml(f) + assert data == {} + + +class TestMergeYamlFiles: + + def test_merge_two_files(self, tmp_path): + f1 = tmp_path / "base.yaml" + f1.write_text("cluster:\n shard_total_num: 64\n cluster_manager:\n cm_heartbeat_timeout_inSecs: 30\n") + f2 = tmp_path / "override.yaml" + f2.write_text("cluster:\n cluster_manager:\n cm_heartbeat_timeout_inSecs: 6\n") + result = merge_yaml_files(f1, f2) + assert result["cluster"]["shard_total_num"] == 64 + assert result["cluster"]["cluster_manager"]["cm_heartbeat_timeout_inSecs"] == 6 + + def test_merge_single_file(self, tmp_path): + f = tmp_path / "only.yaml" + f.write_text("key: value\n") + result = merge_yaml_files(f) + assert result == {"key": "value"} + + +class TestDictToClusterConfig: + + def test_minimal_dict(self): + config = dict_to_cluster_config({}) + assert config.num_data_servers == 3 + assert config.shard_total_num == 64 + assert config.cm_heartbeat_timeout_inSecs == 10 + assert config.cm_host is None + assert config.ds_hosts == [] + + def test_full_dict(self): + data = { + "cluster": { + "shard_total_num": 128, + "build_mode": "debug", + "cluster_manager": { + "cm_cluster_init_grace_period_inSecs": 15, + "cm_heartbeat_timeout_inSecs": 20, + "cm_heartbeat_bg_scan_interval_inSecs": 3, + "cm_deferred_reshard_enabled": False, + "cm_deferred_reshard_window_inSecs": 60, + }, + "data_servers": { + "count": 6, + "heartbeat_cooldown_sec": 3, + "register_cooldown_sec": 5, + "cm_hb_tolerance_count": 10, + "cm_connect_retry_interval_sec": 2, + "memory_limit_bytes": 2 << 30, + "ds_logical_node_id_prefix": "test-ds", + }, + "ssh": { + "user": "admin", + "port": 2222, + }, + } + } + config = dict_to_cluster_config(data) + assert config.num_data_servers == 6 + assert config.shard_total_num == 128 + assert config.build_mode == "debug" + assert config.cm_cluster_init_grace_period_inSecs == 15 + assert config.cm_heartbeat_timeout_inSecs == 20 + assert config.cm_heartbeat_bg_scan_interval_inSecs == 3 + assert config.cm_deferred_reshard_enabled is False + assert config.cm_deferred_reshard_window_inSecs == 60 + assert config.heartbeat_cooldown_sec == 3 + assert config.register_cooldown_sec == 5 + assert config.cm_hb_tolerance_count == 10 + assert config.cm_connect_retry_interval_sec == 2 + assert config.memory_limit_bytes == 2 << 30 + assert config.ds_logical_node_id_prefix == "test-ds" + assert config.ssh_user == "admin" + assert config.ssh_port == 2222 + + def test_with_host_topology(self): + data = { + "cluster": { + "hosts": { + "cm": {"ip": "10.0.0.1", "binary_dir": "/opt/bin"}, + "ds": [ + {"ip": "10.0.0.2"}, + {"ip": "10.0.0.3", "log_dir": "/var/log"}, + ], + }, + } + } + config = dict_to_cluster_config(data) + assert config.cm_host is not None + assert config.cm_host.ip == "10.0.0.1" + assert config.cm_host.binary_dir == "/opt/bin" + assert len(config.ds_hosts) == 2 + assert config.ds_hosts[0].ip == "10.0.0.2" + assert config.ds_hosts[0].log_dir == "/tmp/simm_test_logs" # default + assert config.ds_hosts[1].log_dir == "/var/log" + + +class TestDictToFaultConfigs: + + def test_empty(self): + assert dict_to_fault_configs({}) == [] + + def test_single_fault_with_target(self): + data = { + "faults": [ + {"type": "sigkill", "target": "ds:0", "delay_after_ready_sec": 5} + ] + } + faults = dict_to_fault_configs(data) + assert len(faults) == 1 + assert faults[0].fault_type == "sigkill" + assert faults[0].targets == ["ds:0"] + assert faults[0].delay_after_sec == 5 + assert faults[0].duration_sec is None + + def test_fault_with_targets_list(self): + data = { + "faults": [ + {"type": "kill_multiple", "targets": ["ds:0", "ds:1"]} + ] + } + faults = dict_to_fault_configs(data) + assert faults[0].targets == ["ds:0", "ds:1"] + + def test_fault_with_duration(self): + data = { + "faults": [ + {"type": "sigstop", "target": "ds:0", "duration_sec": 15} + ] + } + faults = dict_to_fault_configs(data) + assert faults[0].duration_sec == 15 + + def test_multiple_faults(self): + data = { + "faults": [ + {"type": "sigkill", "target": "ds:0"}, + {"type": "restart_cm", "target": "cm", "duration_sec": 3}, + ] + } + faults = dict_to_fault_configs(data) + assert len(faults) == 2 + assert faults[0].fault_type == "sigkill" + assert faults[1].fault_type == "restart_cm" + + +class TestConftestMakeConfig: + """Verify conftest._make_config produces valid ClusterConfig.""" + + def test_make_config_default(self): + # Import from conftest to catch parameter name typos at import time + import sys + sys.path.insert(0, str(Path(__file__).parents[2])) + from conftest import _make_config + config = _make_config() + assert config.num_data_servers == 3 + assert config.shard_total_num == 64 + assert config.cm_cluster_init_grace_period_inSecs == 5 + + def test_make_config_medium(self): + import sys + sys.path.insert(0, str(Path(__file__).parents[2])) + from conftest import _make_config + config = _make_config(num_ds=6, cm_cluster_init_grace_period_inSecs=8) + assert config.num_data_servers == 6 + assert config.cm_cluster_init_grace_period_inSecs == 8 + + +class TestClusterConfigProperties: + + def test_cm_failure_detection_max_sec(self): + config = ClusterConfig( + cm_heartbeat_timeout_inSecs=10, + cm_heartbeat_bg_scan_interval_inSecs=2, + ) + assert config.cm_failure_detection_max_sec == 14 # 10 + 2*2 + + def test_ds_cm_failure_detection_sec(self): + config = ClusterConfig( + cm_hb_tolerance_count=5, + heartbeat_cooldown_sec=2, + ) + assert config.ds_cm_failure_detection_sec == 10 # 5 * 2 + + def test_ds_rejoin_max_sec(self): + config = ClusterConfig( + cm_hb_tolerance_count=5, + heartbeat_cooldown_sec=2, + register_cooldown_sec=2, + ) + # ds_cm_failure_detection_sec(10) + register_cooldown_sec(2) + 15 + assert config.ds_rejoin_max_sec == 27 diff --git a/tests/protocol_test/framework/tests/test_log_parser.py b/tests/protocol_test/framework/tests/test_log_parser.py new file mode 100644 index 0000000..42eeae6 --- /dev/null +++ b/tests/protocol_test/framework/tests/test_log_parser.py @@ -0,0 +1,90 @@ +"""Unit tests for log_parser.py — regex matching on log content. + +Uses a mock SshExecutor to feed log content without actual file I/O. +""" + +from unittest.mock import MagicMock + +from framework.log_parser import LogParser + +SAMPLE_LOG = """\ +[2025-01-15 10:00:01.123] [info] ClusterManager service starts successfully! +[2025-01-15 10:00:02.456] [info] NewNodeHandshake from 10.0.0.2:40000 +[2025-01-15 10:00:02.789] [info] AddNode: 10.0.0.2:40000 registered +[2025-01-15 10:00:03.100] [info] NewNodeHandshake from 10.0.0.3:40000 +[2025-01-15 10:00:03.200] [info] AddNode: 10.0.0.3:40000 registered +[2025-01-15 10:00:15.000] [warn] heartbeat timeout for 10.0.0.2:40000 +[2025-01-15 10:00:15.001] [warn] HandleNodeFailure: mark 10.0.0.2:40000 as DEAD +[2025-01-15 10:00:15.002] [info] RebalanceShardsAfterNodeFailure: reassign 22 shards +[2025-01-15 10:00:15.003] [info] UpdateNodeStatus: 10.0.0.2:40000 RUNNING -> DEAD +""" + + +def _make_parser(log_content: str) -> LogParser: + """Create a LogParser with mocked SshExecutor.""" + ssh = MagicMock() + ssh.read_file.return_value = log_content + ssh.read_file_tail.return_value = (log_content, len(log_content)) + return LogParser("/fake/log.log", "127.0.0.1", ssh) + + +class TestLogParserPatterns: + + def test_find_handshake_events(self): + parser = _make_parser(SAMPLE_LOG) + events = parser.find_handshake_events() + assert len(events) >= 4 # 2 NewNodeHandshake + 2 AddNode + + def test_find_heartbeat_timeout_events(self): + parser = _make_parser(SAMPLE_LOG) + events = parser.find_heartbeat_timeout_events() + assert len(events) >= 2 # timeout + HandleNodeFailure + + def test_find_rebalance_events(self): + parser = _make_parser(SAMPLE_LOG) + events = parser.find_rebalance_events() + assert len(events) >= 1 + + def test_find_node_status_changes(self): + parser = _make_parser(SAMPLE_LOG) + events = parser.find_node_status_changes() + assert len(events) >= 1 + assert any("DEAD" in e for e in events) + + def test_contains_existing_pattern(self): + parser = _make_parser(SAMPLE_LOG) + assert parser.contains("RebalanceShardsAfterNodeFailure") + + def test_contains_missing_pattern(self): + parser = _make_parser(SAMPLE_LOG) + assert not parser.contains("ThisPatternDoesNotExist") + + def test_count_pattern(self): + parser = _make_parser(SAMPLE_LOG) + count = parser.count_pattern("NewNodeHandshake") + assert count == 2 + + def test_find_pattern_lines(self): + parser = _make_parser(SAMPLE_LOG) + lines = parser.find_pattern_lines(r"AddNode.*registered") + assert len(lines) == 2 + assert all("registered" in line for line in lines) + + +class TestLogParserEmptyLog: + + def test_empty_log(self): + parser = _make_parser("") + assert parser.find_handshake_events() == [] + assert parser.find_heartbeat_timeout_events() == [] + assert parser.find_rebalance_events() == [] + assert parser.count_pattern("anything") == 0 + assert not parser.contains("anything") + + +class TestLogParserGetAllText: + + def test_get_all_text(self): + parser = _make_parser(SAMPLE_LOG) + text = parser.get_all_text() + assert "ClusterManager service starts" in text diff --git a/tests/protocol_test/framework/tests/test_port_allocator.py b/tests/protocol_test/framework/tests/test_port_allocator.py new file mode 100644 index 0000000..1ad6e73 --- /dev/null +++ b/tests/protocol_test/framework/tests/test_port_allocator.py @@ -0,0 +1,64 @@ +"""Unit tests for port_allocator.py — local port allocation.""" + +import socket + +from framework.port_allocator import PortAllocator + + +class TestPortAllocator: + + def test_allocate_single_port(self): + pa = PortAllocator() + ports = pa.allocate(1) + assert len(ports) == 1 + assert 1024 < ports[0] < 65536 + + def test_allocate_multiple_ports_unique(self): + pa = PortAllocator() + ports = pa.allocate(10) + assert len(ports) == 10 + assert len(set(ports)) == 10 # all unique + + def test_no_collision_across_calls(self): + pa = PortAllocator() + ports1 = pa.allocate(5) + ports2 = pa.allocate(5) + all_ports = ports1 + ports2 + assert len(set(all_ports)) == 10 + + def test_allocate_cm_ports(self): + pa = PortAllocator() + ports = pa.allocate_cm_ports() + assert set(ports.keys()) == {"intra", "inter", "admin"} + assert len(set(ports.values())) == 3 # all different + + def test_allocate_ds_ports(self): + pa = PortAllocator() + ports = pa.allocate_ds_ports() + assert set(ports.keys()) == {"io", "mgt", "admin"} + assert len(set(ports.values())) == 3 + + def test_cm_and_ds_ports_no_overlap(self): + pa = PortAllocator() + cm = pa.allocate_cm_ports() + ds = pa.allocate_ds_ports() + cm_set = set(cm.values()) + ds_set = set(ds.values()) + assert cm_set.isdisjoint(ds_set) + + def test_release_all(self): + pa = PortAllocator() + pa.allocate(5) + pa.release_all() + # After release, internal tracking is cleared + assert pa._allocated == {} + + def test_ports_are_bindable(self): + """Allocated ports should be currently free (best-effort check).""" + pa = PortAllocator() + ports = pa.allocate(3) + for port in ports: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + # Should not raise — port was free when allocated + s.bind(("", port)) diff --git a/tests/protocol_test/framework/tests/test_scenario_runner.py b/tests/protocol_test/framework/tests/test_scenario_runner.py new file mode 100644 index 0000000..3004af4 --- /dev/null +++ b/tests/protocol_test/framework/tests/test_scenario_runner.py @@ -0,0 +1,121 @@ +"""Unit tests for scenario_runner.py — validation step parsing and registries.""" + +from framework.scenario_runner import ( + ScenarioRunner, + ValidationStep, + _FAULT_REGISTRY, + _VALIDATION_REGISTRY, + _parse_validation_steps, + _resolve_target, +) + + +class TestParseValidationSteps: + + def test_empty(self): + steps = _parse_validation_steps({}) + assert steps == [] + + def test_single_step(self): + data = { + "validations": [ + {"type": "node_status", "target": "ds:0", "expected": "DEAD", "within_sec": 20} + ] + } + steps = _parse_validation_steps(data) + assert len(steps) == 1 + assert steps[0].type == "node_status" + assert steps[0].target == "ds:0" + assert steps[0].expected == "DEAD" + assert steps[0].within_sec == 20 + + def test_defaults(self): + data = { + "validations": [ + {"type": "shard_total", "expected": 64} + ] + } + steps = _parse_validation_steps(data) + assert steps[0].target == "" + assert steps[0].field == "" + assert steps[0].within_sec == 60 + assert steps[0].max_imbalance_ratio == 0.3 + + def test_expected_converted_to_string(self): + data = { + "validations": [ + {"type": "shard_total", "expected": 64} + ] + } + steps = _parse_validation_steps(data) + assert steps[0].expected == "64" + assert isinstance(steps[0].expected, str) + + def test_ds_status_with_field(self): + data = { + "validations": [ + {"type": "ds_status", "target": "ds:1", "field": "cm_ready", "expected": "true"} + ] + } + steps = _parse_validation_steps(data) + assert steps[0].field == "cm_ready" + + def test_shard_balance_with_ratio(self): + data = { + "validations": [ + {"type": "shard_balance", "max_imbalance_ratio": 0.5} + ] + } + steps = _parse_validation_steps(data) + assert steps[0].max_imbalance_ratio == 0.5 + + def test_multiple_steps(self): + data = { + "validations": [ + {"type": "node_status", "target": "ds:0", "expected": "DEAD"}, + {"type": "shard_total", "expected": 64}, + {"type": "no_orphaned_shards"}, + ] + } + steps = _parse_validation_steps(data) + assert len(steps) == 3 + assert [s.type for s in steps] == ["node_status", "shard_total", "no_orphaned_shards"] + + +class TestFaultRegistry: + + def test_builtin_fault_types(self): + expected = {"sigkill", "sigterm", "sigstop", "sigcont", "kill_multiple", "restart_cm"} + assert expected.issubset(set(_FAULT_REGISTRY.keys())) + + def test_all_entries_callable(self): + for name, fn in _FAULT_REGISTRY.items(): + assert callable(fn), f"Fault '{name}' is not callable" + + +class TestValidationRegistry: + + def test_builtin_validation_types(self): + expected = { + "node_status", "alive_node_count", "all_nodes_running", + "shard_total", "shard_balance", "no_orphaned_shards", + "node_has_no_shards", "ds_status", + } + assert expected.issubset(set(_VALIDATION_REGISTRY.keys())) + + def test_all_entries_callable(self): + for name, fn in _VALIDATION_REGISTRY.items(): + assert callable(fn), f"Validation '{name}' is not callable" + + +class TestScenarioRunnerLists: + + def test_list_fault_types(self): + types = ScenarioRunner.list_fault_types() + assert "sigkill" in types + assert "restart_cm" in types + + def test_list_validation_types(self): + types = ScenarioRunner.list_validation_types() + assert "node_status" in types + assert "shard_total" in types diff --git a/tests/protocol_test/pytest.ini b/tests/protocol_test/pytest.ini new file mode 100644 index 0000000..7a8c31f --- /dev/null +++ b/tests/protocol_test/pytest.ini @@ -0,0 +1,10 @@ +[pytest] +testpaths = tests +markers = + requires_rdma: test requires RDMA hardware for CM/DS processes + requires_root: test requires root privileges (iptables) + slow: test takes more than 60 seconds +timeout = 120 +log_cli = true +log_cli_level = INFO +log_cli_format = %(asctime)s [%(levelname)s] %(name)s: %(message)s diff --git a/tests/protocol_test/requirements.txt b/tests/protocol_test/requirements.txt new file mode 100644 index 0000000..38eeb45 --- /dev/null +++ b/tests/protocol_test/requirements.txt @@ -0,0 +1,3 @@ +pytest>=7.0 +pytest-timeout>=2.0 +pyyaml>=6.0 diff --git a/tests/protocol_test/scenarios/clusters/medium.yaml b/tests/protocol_test/scenarios/clusters/medium.yaml new file mode 100644 index 0000000..1860744 --- /dev/null +++ b/tests/protocol_test/scenarios/clusters/medium.yaml @@ -0,0 +1,16 @@ +cluster: + build_mode: release + shard_total_num: 64 + + cluster_manager: + cm_cluster_init_grace_period_inSecs: 8 + cm_heartbeat_timeout_inSecs: 10 + cm_heartbeat_bg_scan_interval_inSecs: 2 + + data_servers: + count: 6 + heartbeat_cooldown_sec: 2 + register_cooldown_sec: 2 + cm_hb_tolerance_count: 5 + cm_connect_retry_interval_sec: 1 + memory_limit_bytes: 1073741824 diff --git a/tests/protocol_test/scenarios/clusters/multi_machine.yaml b/tests/protocol_test/scenarios/clusters/multi_machine.yaml new file mode 100644 index 0000000..94476c5 --- /dev/null +++ b/tests/protocol_test/scenarios/clusters/multi_machine.yaml @@ -0,0 +1,45 @@ +# Multi-machine cluster topology example. +# +# Usage: +# export SIMM_CLUSTER_CONFIG=scenarios/clusters/multi_machine.yaml +# pytest tests/ -v --timeout=120 +# +# The test runner (node A) must have passwordless SSH access to all listed hosts. +# SiMM binaries must be pre-built at the specified binary_dir on each host. + +cluster: + build_mode: release + shard_total_num: 64 + + hosts: + cm: + ip: "10.0.0.1" + binary_dir: "/opt/simm/build/release/bin" + log_dir: "/tmp/simm_test_logs" + ds: + - ip: "10.0.0.2" + binary_dir: "/opt/simm/build/release/bin" + log_dir: "/tmp/simm_test_logs" + - ip: "10.0.0.3" + binary_dir: "/opt/simm/build/release/bin" + log_dir: "/tmp/simm_test_logs" + - ip: "10.0.0.4" + binary_dir: "/opt/simm/build/release/bin" + log_dir: "/tmp/simm_test_logs" + + ssh: + user: "root" + port: 22 + + cluster_manager: + cm_cluster_init_grace_period_inSecs: 5 + cm_heartbeat_timeout_inSecs: 10 + cm_heartbeat_bg_scan_interval_inSecs: 2 + + data_servers: + count: 3 + heartbeat_cooldown_sec: 2 + register_cooldown_sec: 2 + cm_hb_tolerance_count: 5 + cm_connect_retry_interval_sec: 1 + memory_limit_bytes: 1073741824 diff --git a/tests/protocol_test/scenarios/clusters/small.yaml b/tests/protocol_test/scenarios/clusters/small.yaml new file mode 100644 index 0000000..a44af3a --- /dev/null +++ b/tests/protocol_test/scenarios/clusters/small.yaml @@ -0,0 +1,16 @@ +cluster: + build_mode: release + shard_total_num: 64 + + cluster_manager: + cm_cluster_init_grace_period_inSecs: 5 + cm_heartbeat_timeout_inSecs: 10 + cm_heartbeat_bg_scan_interval_inSecs: 2 + + data_servers: + count: 3 + heartbeat_cooldown_sec: 2 + register_cooldown_sec: 2 + cm_hb_tolerance_count: 5 + cm_connect_retry_interval_sec: 1 + memory_limit_bytes: 1073741824 # 1GB diff --git a/tests/protocol_test/scenarios/faults/freeze_ds.yaml b/tests/protocol_test/scenarios/faults/freeze_ds.yaml new file mode 100644 index 0000000..4ce9f5f --- /dev/null +++ b/tests/protocol_test/scenarios/faults/freeze_ds.yaml @@ -0,0 +1,5 @@ +faults: + - type: sigstop + target: "ds:0" + delay_after_ready_sec: 5 + duration_sec: 15 diff --git a/tests/protocol_test/scenarios/faults/kill_cm.yaml b/tests/protocol_test/scenarios/faults/kill_cm.yaml new file mode 100644 index 0000000..198bbe4 --- /dev/null +++ b/tests/protocol_test/scenarios/faults/kill_cm.yaml @@ -0,0 +1,4 @@ +faults: + - type: sigkill + target: "cm" + delay_after_ready_sec: 5 diff --git a/tests/protocol_test/scenarios/faults/kill_one_ds.yaml b/tests/protocol_test/scenarios/faults/kill_one_ds.yaml new file mode 100644 index 0000000..fd777e1 --- /dev/null +++ b/tests/protocol_test/scenarios/faults/kill_one_ds.yaml @@ -0,0 +1,4 @@ +faults: + - type: sigkill + target: "ds:0" + delay_after_ready_sec: 5 diff --git a/tests/protocol_test/scenarios/faults/kill_two_ds.yaml b/tests/protocol_test/scenarios/faults/kill_two_ds.yaml new file mode 100644 index 0000000..458744c --- /dev/null +++ b/tests/protocol_test/scenarios/faults/kill_two_ds.yaml @@ -0,0 +1,6 @@ +faults: + - type: sigkill + targets: + - "ds:0" + - "ds:1" + delay_after_ready_sec: 5 diff --git a/tests/protocol_test/scenarios/full/cm_restart_ds_rejoin.yaml b/tests/protocol_test/scenarios/full/cm_restart_ds_rejoin.yaml new file mode 100644 index 0000000..d5267dd --- /dev/null +++ b/tests/protocol_test/scenarios/full/cm_restart_ds_rejoin.yaml @@ -0,0 +1,31 @@ +# Complete scenario: kill CM, restart, verify all DS re-register. + +cluster: + shard_total_num: 64 + cluster_manager: + cm_cluster_init_grace_period_inSecs: 5 + cm_heartbeat_timeout_inSecs: 10 + cm_heartbeat_bg_scan_interval_inSecs: 2 + data_servers: + count: 3 + heartbeat_cooldown_sec: 2 + register_cooldown_sec: 2 + cm_hb_tolerance_count: 5 + +faults: + - type: restart_cm + target: "cm" + delay_after_ready_sec: 3 + # duration_sec: time to wait before restarting CM + # (allow DS to detect CM failure: cm_hb_tolerance_count * heartbeat_cooldown_sec) + duration_sec: 12 + +validations: + - type: alive_node_count + expected: 3 + within_sec: 30 # cm_cluster_init_grace_period_inSecs + ds_rejoin time + + - type: all_nodes_running + + - type: shard_total + expected: 64 diff --git a/tests/protocol_test/scenarios/full/deferred_reshard_window_timeout.yaml b/tests/protocol_test/scenarios/full/deferred_reshard_window_timeout.yaml new file mode 100644 index 0000000..0cb3ad0 --- /dev/null +++ b/tests/protocol_test/scenarios/full/deferred_reshard_window_timeout.yaml @@ -0,0 +1,45 @@ +# Scenario: Deferred reshard window timeout → degrade to standard reshard. +# +# DS killed, no replacement within window → DEFERRED_RESHARD → DEAD → reshard. + +cluster: + shard_total_num: 64 + cluster_manager: + cm_cluster_init_grace_period_inSecs: 5 + cm_heartbeat_timeout_inSecs: 6 + cm_heartbeat_bg_scan_interval_inSecs: 1 + cm_deferred_reshard_enabled: true + cm_deferred_reshard_window_inSecs: 10 + data_servers: + count: 3 + heartbeat_cooldown_sec: 2 + register_cooldown_sec: 2 + ds_logical_node_id_prefix: "simm-ds" + +faults: + - type: sigkill + target: "ds:0" + delay_after_ready_sec: 2 + +validations: + # First: DEFERRED_RESHARD (not immediate DEAD) + - type: node_status + target: "ds:0" + expected: DEFERRED_RESHARD + within_sec: 15 + + # Then: window expires → DEAD + - type: node_status + target: "ds:0" + expected: DEAD + within_sec: 20 + + # Standard reshard happens + - type: no_orphaned_shards + within_sec: 15 + + - type: node_has_no_shards + target: "ds:0" + + - type: shard_total + expected: 64 diff --git a/tests/protocol_test/scenarios/full/kill_ds_and_verify_rebalance.yaml b/tests/protocol_test/scenarios/full/kill_ds_and_verify_rebalance.yaml new file mode 100644 index 0000000..8ff6ee5 --- /dev/null +++ b/tests/protocol_test/scenarios/full/kill_ds_and_verify_rebalance.yaml @@ -0,0 +1,40 @@ +# Complete scenario: kill one DS, verify failure detection + shard rebalance. +# +# Usage: +# pytest tests/test_scenario.py -k "kill_ds_and_verify_rebalance" -v +# or programmatically: +# ScenarioRunner().run_scenario("scenarios/full/kill_ds_and_verify_rebalance.yaml") + +cluster: + shard_total_num: 64 + cluster_manager: + cm_cluster_init_grace_period_inSecs: 5 + cm_heartbeat_timeout_inSecs: 10 + cm_heartbeat_bg_scan_interval_inSecs: 2 + data_servers: + count: 3 + heartbeat_cooldown_sec: 2 + register_cooldown_sec: 2 + +faults: + - type: sigkill + target: "ds:0" + delay_after_ready_sec: 2 + +validations: + - type: node_status + target: "ds:0" + expected: DEAD + within_sec: 20 # cm_heartbeat_timeout_inSecs + 2 * cm_heartbeat_bg_scan_interval_inSecs + buffer + + - type: no_orphaned_shards + within_sec: 15 + + - type: node_has_no_shards + target: "ds:0" + + - type: shard_total + expected: 64 + + - type: shard_balance + max_imbalance_ratio: 0.3 diff --git a/tests/protocol_test/scenarios/overrides/fast_heartbeat.yaml b/tests/protocol_test/scenarios/overrides/fast_heartbeat.yaml new file mode 100644 index 0000000..32e1612 --- /dev/null +++ b/tests/protocol_test/scenarios/overrides/fast_heartbeat.yaml @@ -0,0 +1,8 @@ +# Override: reduce heartbeat timeouts for faster test execution +cluster: + cluster_manager: + cm_heartbeat_timeout_inSecs: 6 + cm_heartbeat_bg_scan_interval_inSecs: 1 + + data_servers: + heartbeat_cooldown_sec: 1 diff --git a/tests/protocol_test/tests/__init__.py b/tests/protocol_test/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/protocol_test/tests/test_cm_restart.py b/tests/protocol_test/tests/test_cm_restart.py new file mode 100644 index 0000000..111aa9e --- /dev/null +++ b/tests/protocol_test/tests/test_cm_restart.py @@ -0,0 +1,136 @@ +"""Tests for Cluster Manager crash and restart recovery. + +Verifies the full CM restart protocol from BOTH sides: + CM killed → DS detects HB failures (heartbeat_failure_count reaches cm_hb_tolerance_count) + → DS sets cm_ready=false → CM restarts → DS calls RegisterToRestartedManager + → DS sets cm_ready=true, is_registered=true, heartbeat_failure_count=0 + → CM rebuilds routing table from DS-reported shard lists +""" + +import time + +import pytest + + +@pytest.mark.requires_rdma +class TestCMRestart: + """Verify DS re-register after CM crash, with DS-side and shard verification.""" + + def test_cm_crash_all_ds_rejoin(self, cluster_small): + """Kill CM → DS detect failure → restart CM → all DS re-register.""" + + # Kill CM + cluster_small.fault_injector.kill_process(cluster_small.cm) + assert not cluster_small.process_manager.is_alive(cluster_small.cm) + + # All DS should still be alive (DS doesn't exit when CM dies) + for ds in cluster_small.data_servers: + assert cluster_small.process_manager.is_alive(ds), ( + f"DS[{ds.index}] should survive CM crash" + ) + + # Wait for DS to detect CM failure via heartbeat_failure_count + hb_failure_wait = ( + 5 * cluster_small.config.heartbeat_cooldown_sec + 2 + ) + time.sleep(hb_failure_wait) + + # Restart CM with same ports + new_cm = cluster_small.restart_cm() + assert cluster_small.process_manager.is_alive(new_cm) + + # Wait for DS re-registration + timeout = cluster_small.config.cm_cluster_init_grace_period_inSecs + 20 + assert cluster_small.observer.wait_for_node_count( + cluster_small.config.num_data_servers, timeout=timeout + ), "Not all DS re-registered after CM restart" + + # CM side: all nodes RUNNING + cluster_small.observer.assert_all_nodes_running() + + # Shard routing table fully rebuilt + cluster_small.observer.assert_total_shard_count(shard_total_before) + + def test_ds_detects_cm_failure_via_status(self, cluster_small): + """After CM crash, DS admin RPC should show cm_ready=false, failure_count >= tolerance.""" + cluster_small.fault_injector.kill_process(cluster_small.cm) + + # Wait for DS to accumulate heartbeat failures + # DS: cm_hb_tolerance_count(5) * heartbeat_cooldown_sec(2) = 10s + hb_failure_wait = ( + 5 * cluster_small.config.heartbeat_cooldown_sec + 5 + ) + + # Each DS should detect CM failure: cm_ready becomes false + for ds in cluster_small.data_servers: + assert cluster_small.observer.wait_for_ds_cm_not_ready( + ds.host, ds.pid, timeout=hb_failure_wait + ), f"DS[{ds.index}] did not detect CM failure (cm_ready still true)" + + # Verify heartbeat_failure_count is non-zero on each DS + for ds in cluster_small.data_servers: + cluster_small.observer.assert_ds_heartbeat_failure_count( + ds.host, ds.pid, min_count=5 + ) + + # Restart CM so teardown works + cluster_small.restart_cm() + timeout = cluster_small.config.cm_cluster_init_grace_period_inSecs + 20 + cluster_small.observer.wait_for_node_count( + cluster_small.config.num_data_servers, timeout=timeout + ) + + def test_ds_re_registration_resets_status(self, cluster_small): + """After CM restart, DS status should show is_registered=true, cm_ready=true, failure_count=0.""" + cluster_small.fault_injector.kill_process(cluster_small.cm) + + hb_failure_wait = ( + 5 * cluster_small.config.heartbeat_cooldown_sec + 2 + ) + time.sleep(hb_failure_wait) + + cluster_small.restart_cm() + + timeout = cluster_small.config.cm_cluster_init_grace_period_inSecs + 20 + assert cluster_small.observer.wait_for_node_count( + cluster_small.config.num_data_servers, timeout=timeout + ) + + # Each DS should now report: registered=true, cm_ready=true, failure_count=0 + for ds in cluster_small.data_servers: + assert cluster_small.observer.wait_for_ds_registered( + ds.host, ds.pid, timeout=10 + ), f"DS[{ds.index}] did not re-register properly" + + cluster_small.observer.assert_ds_is_registered(ds.host, ds.pid) + cluster_small.observer.assert_ds_cm_ready(ds.host, ds.pid, expected=True) + cluster_small.observer.assert_ds_heartbeat_failure_count( + ds.host, ds.pid, min_count=0, max_count=0 + ) + + def test_shard_table_rebuilt_after_cm_restart(self, cluster_small): + """After CM restart, shard routing table rebuilt from DS-reported shard lists.""" + shard_total_before = cluster_small.observer.get_total_shard_count() + dist_before = cluster_small.observer.get_shard_distribution() + + cluster_small.fault_injector.kill_process(cluster_small.cm) + time.sleep(5 * cluster_small.config.heartbeat_cooldown_sec + 2) + cluster_small.restart_cm() + + timeout = cluster_small.config.cm_cluster_init_grace_period_inSecs + 20 + assert cluster_small.observer.wait_for_node_count( + cluster_small.config.num_data_servers, timeout=timeout + ) + + # Shard table must be complete + cluster_small.observer.assert_total_shard_count(shard_total_before) + + # Distribution should match pre-crash state + dist_after = cluster_small.observer.get_shard_distribution() + for addr, count_before in dist_before.items(): + count_after = dist_after.get(addr, 0) + assert count_after == count_before, ( + f"Node {addr} shard count changed after CM restart: " + f"before={count_before}, after={count_after}" + ) + diff --git a/tests/protocol_test/tests/test_deferred_reshard.py b/tests/protocol_test/tests/test_deferred_reshard.py new file mode 100644 index 0000000..8e88b61 --- /dev/null +++ b/tests/protocol_test/tests/test_deferred_reshard.py @@ -0,0 +1,457 @@ +"""Tests for Deferred Reshard feature. + +Design doc: simm-deferred-reshard-design.docx v2 + +Core behavior: + DS down → CM enters DEFERRED_RESHARD (not DEAD, no reshard yet) + → Within cm_deferred_reshard_window_inSecs, new DS with same logical_node_id registers + → CM does in-place IP update in routing table, assigns original shards, no reshard + → If window expires, CM degrades to standard reshard (DEAD) + +Key flags: + cm_deferred_reshard_enabled (default true) + cm_deferred_reshard_window_inSecs (default 120s, shortened in tests) + cm_heartbeat_timeout_inSecs (default 30s, shortened in tests) + ds_logical_node_id (per-DS, e.g. "simm-ds-0") +""" + +import pytest + +from framework.cluster import SimmCluster +from framework.config import ClusterConfig + + +def _make_deferred_reshard_config( + cm_deferred_reshard_window_inSecs: int = 30, +) -> ClusterConfig: + """Create a ClusterConfig with deferred reshard enabled and short timeouts.""" + return ClusterConfig( + num_data_servers=3, + shard_total_num=64, + cm_cluster_init_grace_period_inSecs=5, + cm_heartbeat_timeout_inSecs=6, + cm_heartbeat_bg_scan_interval_inSecs=1, + heartbeat_cooldown_sec=2, + register_cooldown_sec=2, + cm_hb_tolerance_count=5, + cm_deferred_reshard_enabled=True, + cm_deferred_reshard_window_inSecs=cm_deferred_reshard_window_inSecs, + ds_logical_node_id_prefix="simm-ds", + ) + + +@pytest.fixture +def cluster_deferred(tmp_path, binary_dir): + """Cluster with deferred reshard enabled, short timeouts for fast testing.""" + config = _make_deferred_reshard_config(cm_deferred_reshard_window_inSecs=30) + cluster = SimmCluster(config, log_dir=tmp_path / "logs", binary_dir=binary_dir) + cluster.start() + cluster.wait_ready() + yield cluster + cluster.teardown() + + +@pytest.fixture +def cluster_deferred_short_window(tmp_path, binary_dir): + """Cluster with very short deferred reshard window (10s) to test timeout degradation.""" + config = _make_deferred_reshard_config(cm_deferred_reshard_window_inSecs=10) + cluster = SimmCluster(config, log_dir=tmp_path / "logs", binary_dir=binary_dir) + cluster.start() + cluster.wait_ready() + yield cluster + cluster.teardown() + + +@pytest.fixture +def cluster_deferred_disabled(tmp_path, binary_dir): + """Cluster with deferred reshard disabled — should behave like original.""" + config = _make_deferred_reshard_config() + config.cm_deferred_reshard_enabled = False + cluster = SimmCluster(config, log_dir=tmp_path / "logs", binary_dir=binary_dir) + cluster.start() + cluster.wait_ready() + yield cluster + cluster.teardown() + + +@pytest.mark.requires_rdma +class TestDeferredReshardReplace: + """DS killed within window → replacement DS with same logical_node_id → no reshard.""" + + def test_kill_ds_enters_deferred_reshard_not_dead(self, cluster_deferred): + """After DS killed, CM should enter DEFERRED_RESHARD, NOT DEAD.""" + c = cluster_deferred + ds0 = c.get_ds_handle(0) + addr = ds0.addr_str + + # Record shard distribution before kill + dist_before = c.observer.get_shard_distribution() + ds0_shards_before = dist_before.get(addr, 0) + assert ds0_shards_before > 0 + + c.fault_injector.kill_process(ds0) + + # Wait for CM to detect heartbeat timeout + detect_timeout = ( + c.config.cm_heartbeat_timeout_inSecs + + 2 * c.config.cm_heartbeat_bg_scan_interval_inSecs + + 3 + ) + assert c.observer.wait_for_node_status( + addr, "DEFERRED_RESHARD", timeout=detect_timeout + ), f"CM should mark {addr} as DEFERRED_RESHARD, not DEAD" + + # Shards should NOT have been migrated (that's the whole point) + dist_during = c.observer.get_shard_distribution() + ds0_shards_during = dist_during.get(addr, 0) + assert ds0_shards_during == ds0_shards_before, ( + f"Shards should stay assigned during DEFERRED_RESHARD: " + f"before={ds0_shards_before}, during={ds0_shards_during}" + ) + + # Total shard count unchanged + c.observer.assert_total_shard_count(c.config.shard_total_num) + + def test_replacement_ds_same_logical_id_recovers_shards(self, cluster_deferred): + """Kill DS0, start replacement with same logical_node_id → shards recovered, no reshard.""" + c = cluster_deferred + ds0 = c.get_ds_handle(0) + addr = ds0.addr_str + logical_id = f"{c.config.ds_logical_node_id_prefix}-0" # "simm-ds-0" + + # Record shard distribution before kill + dist_before = c.observer.get_shard_distribution() + ds0_shards_before = dist_before.get(addr, 0) + total_before = sum(dist_before.values()) + + # Kill DS0 + c.fault_injector.kill_process(ds0) + + # Wait for DEFERRED_RESHARD + detect_timeout = c.config.cm_heartbeat_timeout_inSecs + 5 + assert c.observer.wait_for_node_status( + addr, "DEFERRED_RESHARD", timeout=detect_timeout + ) + + # Start replacement DS with same logical_node_id (may use different ports) + new_ds = c.add_data_server(ds_logical_node_id=logical_id) + + # Wait for CM to recognize replacement and set node back to RUNNING + assert c.observer.wait_for_node_status( + new_ds.addr_str, "RUNNING", timeout=15 + ), "Replacement DS did not become RUNNING" + + # Shard total unchanged + c.observer.assert_total_shard_count(total_before) + + # The replacement DS should have received the SAME shards + # (CM does in-place IP update, assigns original shards via assigned_shard_ids) + dist_after = c.observer.get_shard_distribution() + new_ds_shards = dist_after.get(new_ds.addr_str, 0) + assert new_ds_shards == ds0_shards_before, ( + f"Replacement DS should inherit {ds0_shards_before} shards, " + f"got {new_ds_shards}" + ) + + # Other DS shard counts unchanged (no reshard happened) + for ds in c.data_servers: + if ds.index == 0 or ds.addr_str == new_ds.addr_str: + continue + before = dist_before.get(ds.addr_str, 0) + after = dist_after.get(ds.addr_str, 0) + assert after == before, ( + f"DS[{ds.index}] shards changed during deferred reshard: " + f"before={before}, after={after}. Reshard should NOT have happened." + ) + + def test_replacement_ds_status_healthy_after_recovery(self, cluster_deferred): + """After replacement, DS admin RPC should show is_registered=true, cm_ready=true.""" + c = cluster_deferred + ds0 = c.get_ds_handle(0) + logical_id = f"{c.config.ds_logical_node_id_prefix}-0" + + c.fault_injector.kill_process(ds0) + + detect_timeout = c.config.cm_heartbeat_timeout_inSecs + 5 + c.observer.wait_for_node_status(ds0.addr_str, "DEFERRED_RESHARD", timeout=detect_timeout) + + new_ds = c.add_data_server(ds_logical_node_id=logical_id) + + c.observer.wait_for_node_status(new_ds.addr_str, "RUNNING", timeout=15) + + # DS-side: healthy state + assert c.observer.wait_for_ds_registered( + new_ds.host, new_ds.pid, timeout=10 + ) + c.observer.assert_ds_cm_ready(new_ds.host, new_ds.pid, expected=True) + c.observer.assert_ds_heartbeat_failure_count( + new_ds.host, new_ds.pid, min_count=0, max_count=0 + ) + + +@pytest.mark.requires_rdma +class TestDeferredReshardWindowTimeout: + """Window expires without replacement → degrade to standard reshard.""" + + def test_window_timeout_triggers_reshard(self, cluster_deferred_short_window): + """If no replacement within cm_deferred_reshard_window_inSecs, CM degrades to DEAD + reshard.""" + c = cluster_deferred_short_window + ds0 = c.get_ds_handle(0) + addr = ds0.addr_str + + dist_before = c.observer.get_shard_distribution() + total_before = sum(dist_before.values()) + ds0_shards_before = dist_before.get(addr, 0) + assert ds0_shards_before > 0 + + c.fault_injector.kill_process(ds0) + + # First: should enter DEFERRED_RESHARD + detect_timeout = c.config.cm_heartbeat_timeout_inSecs + 5 + assert c.observer.wait_for_node_status( + addr, "DEFERRED_RESHARD", timeout=detect_timeout + ) + + # Wait for window to expire: cm_deferred_reshard_window_inSecs = 10s + window_timeout = ( + c.config.cm_deferred_reshard_window_inSecs + + c.config.cm_heartbeat_bg_scan_interval_inSecs + + 5 + ) + assert c.observer.wait_for_node_status( + addr, "DEAD", timeout=window_timeout + ), "CM should degrade to DEAD after deferred reshard window expires" + + # Now standard reshard should have happened + c.observer.wait_for_rebalance_complete(timeout=15) + + # Dead node: zero shards + c.observer.assert_node_has_no_shards(addr) + + # Total preserved + c.observer.assert_total_shard_count(total_before) + + # Alive nodes absorbed the shards + c.observer.assert_no_orphaned_shards() + + # Alive nodes gained ds0's shards + dist_after = c.observer.get_shard_distribution() + for ds in c.data_servers: + if ds.index == 0: + continue + before = dist_before.get(ds.addr_str, 0) + after = dist_after.get(ds.addr_str, 0) + assert after > before, ( + f"DS[{ds.index}] should have gained shards: before={before}, after={after}" + ) + + def test_deferred_then_late_replacement_treated_as_new_node( + self, cluster_deferred_short_window + ): + """If replacement arrives AFTER window expires, it's treated as a new node, not replacement.""" + c = cluster_deferred_short_window + ds0 = c.get_ds_handle(0) + addr = ds0.addr_str + logical_id = f"{c.config.ds_logical_node_id_prefix}-0" + + dist_before = c.observer.get_shard_distribution() + ds0_shards_before = dist_before.get(addr, 0) + assert ds0_shards_before > 0 + + c.fault_injector.kill_process(ds0) + + # Wait for DEFERRED_RESHARD then DEAD (window expired) + total_timeout = ( + c.config.cm_heartbeat_timeout_inSecs + + c.config.cm_deferred_reshard_window_inSecs + + 10 + ) + assert c.observer.wait_for_node_status( + addr, "DEAD", timeout=total_timeout + ) + c.observer.wait_for_rebalance_complete(timeout=15) + + # Now start "replacement" DS — but it's too late, reshard already happened + new_ds = c.add_data_server(ds_logical_node_id=logical_id) + + # Should register as new node (RUNNING), but gets new shard assignment + assert c.observer.wait_for_node_status( + new_ds.addr_str, "RUNNING", timeout=15 + ) + + c.observer.assert_total_shard_count(c.config.shard_total_num) + + # Late DS is treated as new node — it should NOT inherit ds0's original shard count + # (reshard already redistributed ds0's shards to other alive nodes) + dist_after = c.observer.get_shard_distribution() + new_ds_shards = dist_after.get(new_ds.addr_str, 0) + assert new_ds_shards != ds0_shards_before or new_ds_shards > 0, ( + "Late replacement should be treated as new node with fresh shard assignment" + ) + + +@pytest.mark.requires_rdma +class TestDeferredReshardDisabled: + """With cm_deferred_reshard_enabled=false, DS kill goes straight to DEAD + reshard.""" + + def test_disabled_goes_straight_to_dead(self, cluster_deferred_disabled): + """With deferred reshard disabled, DS kill → DEAD immediately (no DEFERRED_RESHARD).""" + c = cluster_deferred_disabled + ds0 = c.get_ds_handle(0) + addr = ds0.addr_str + + dist_before = c.observer.get_shard_distribution() + total_before = sum(dist_before.values()) + + c.fault_injector.kill_process(ds0) + + detect_timeout = ( + c.config.cm_heartbeat_timeout_inSecs + + 2 * c.config.cm_heartbeat_bg_scan_interval_inSecs + + 5 + ) + # Should go straight to DEAD, skip DEFERRED_RESHARD + assert c.observer.wait_for_node_status( + addr, "DEAD", timeout=detect_timeout + ), "With deferred reshard disabled, should go straight to DEAD" + + # Standard reshard should happen + c.observer.wait_for_rebalance_complete(timeout=15) + c.observer.assert_node_has_no_shards(addr) + c.observer.assert_total_shard_count(total_before) + + +@pytest.mark.requires_rdma +class TestDeferredReshardEdgeCases: + """Edge cases for the deferred reshard protocol.""" + + def test_replacement_with_different_logical_id_is_new_node(self, cluster_deferred): + """A DS with a DIFFERENT logical_node_id during DEFERRED_RESHARD is treated as new, not replacement.""" + c = cluster_deferred + ds0 = c.get_ds_handle(0) + addr = ds0.addr_str + + dist_before = c.observer.get_shard_distribution() + ds0_shards = dist_before.get(addr, 0) + + c.fault_injector.kill_process(ds0) + + detect_timeout = c.config.cm_heartbeat_timeout_inSecs + 5 + c.observer.wait_for_node_status(addr, "DEFERRED_RESHARD", timeout=detect_timeout) + + # Start DS with DIFFERENT logical_node_id — should NOT replace ds0 + new_ds = c.add_data_server(ds_logical_node_id="simm-ds-99-not-a-replacement") + + # New DS should register as new node + c.observer.wait_for_node_status(new_ds.addr_str, "RUNNING", timeout=15) + + # ds0 should STILL be in DEFERRED_RESHARD (not replaced) + status = c.observer.get_node_status(addr) + assert status == "DEFERRED_RESHARD", ( + f"DS0 should still be DEFERRED_RESHARD, got {status}" + ) + + # ds0's shards should still be assigned to ds0 (not migrated) + dist_after = c.observer.get_shard_distribution() + assert dist_after.get(addr, 0) == ds0_shards, ( + f"DS0 shards should be unchanged during DEFERRED_RESHARD" + ) + + def test_fast_restart_before_heartbeat_timeout(self, cluster_deferred): + """DS killed and restarted with same logical_id BEFORE heartbeat timeout + → IP update, never enters DEFERRED_RESHARD at all.""" + c = cluster_deferred + ds0 = c.get_ds_handle(0) + logical_id = f"{c.config.ds_logical_node_id_prefix}-0" + + dist_before = c.observer.get_shard_distribution() + total_before = sum(dist_before.values()) + + # Kill and immediately restart (within heartbeat_timeout) + c.fault_injector.kill_process(ds0) + new_ds = c.add_data_server(ds_logical_node_id=logical_id) + + # Should register quickly — never hits DEFERRED_RESHARD + assert c.observer.wait_for_node_status( + new_ds.addr_str, "RUNNING", timeout=10 + ) + + # Shards unchanged, total correct + c.observer.assert_total_shard_count(total_before) + + def test_insufficient_nodes_keeps_deferred_state(self, cluster_deferred_short_window): + """Kill 2/3 DS → window expires → CM keeps DEFERRED_RESHARD (not DEAD) + because only 1 alive node cannot absorb all orphaned shards safely.""" + c = cluster_deferred_short_window + ds0 = c.get_ds_handle(0) + ds1 = c.get_ds_handle(1) + + total_before = c.observer.get_total_shard_count() + + # Kill 2 of 3 DS + c.fault_injector.kill_multiple([ds0, ds1]) + + # Both should enter DEFERRED_RESHARD + detect_timeout = c.config.cm_heartbeat_timeout_inSecs + 5 + c.observer.wait_for_node_status(ds0.addr_str, "DEFERRED_RESHARD", timeout=detect_timeout) + c.observer.wait_for_node_status(ds1.addr_str, "DEFERRED_RESHARD", timeout=detect_timeout) + + # Wait for window to expire + window_timeout = ( + c.config.cm_deferred_reshard_window_inSecs + + c.config.cm_heartbeat_bg_scan_interval_inSecs + + 5 + ) + time.sleep(window_timeout) + + # Both should STILL be DEFERRED_RESHARD — not DEAD + # CM cannot reshard safely with only 1 alive node + status0 = c.observer.get_node_status(ds0.addr_str) + status1 = c.observer.get_node_status(ds1.addr_str) + assert status0 == "DEFERRED_RESHARD", ( + f"DS0 should stay DEFERRED_RESHARD with insufficient alive nodes, got {status0}" + ) + assert status1 == "DEFERRED_RESHARD", ( + f"DS1 should stay DEFERRED_RESHARD with insufficient alive nodes, got {status1}" + ) + + # Shards should still be assigned (no reshard happened) + c.observer.assert_total_shard_count(total_before) + + # Remaining DS should still be RUNNING + ds2 = c.get_ds_handle(2) + assert c.observer.get_node_status(ds2.addr_str) == "RUNNING" + + def test_multiple_ds_down_deferred_reshard(self, cluster_deferred): + """Kill 2/3 DS → both enter DEFERRED_RESHARD → replace both → all shards recovered.""" + c = cluster_deferred + ds0 = c.get_ds_handle(0) + ds1 = c.get_ds_handle(1) + logical_id_0 = f"{c.config.ds_logical_node_id_prefix}-0" + logical_id_1 = f"{c.config.ds_logical_node_id_prefix}-1" + + dist_before = c.observer.get_shard_distribution() + total_before = sum(dist_before.values()) + ds0_shards = dist_before.get(ds0.addr_str, 0) + ds1_shards = dist_before.get(ds1.addr_str, 0) + + # Kill both + c.fault_injector.kill_multiple([ds0, ds1]) + + detect_timeout = c.config.cm_heartbeat_timeout_inSecs + 5 + c.observer.wait_for_node_status(ds0.addr_str, "DEFERRED_RESHARD", timeout=detect_timeout) + c.observer.wait_for_node_status(ds1.addr_str, "DEFERRED_RESHARD", timeout=detect_timeout) + + # Replace both + new_ds0 = c.add_data_server(ds_logical_node_id=logical_id_0) + new_ds1 = c.add_data_server(ds_logical_node_id=logical_id_1) + + c.observer.wait_for_node_status(new_ds0.addr_str, "RUNNING", timeout=15) + c.observer.wait_for_node_status(new_ds1.addr_str, "RUNNING", timeout=15) + + # All shards recovered + c.observer.assert_total_shard_count(total_before) + + dist_after = c.observer.get_shard_distribution() + assert dist_after.get(new_ds0.addr_str, 0) == ds0_shards + assert dist_after.get(new_ds1.addr_str, 0) == ds1_shards diff --git a/tests/protocol_test/tests/test_failure_detection.py b/tests/protocol_test/tests/test_failure_detection.py new file mode 100644 index 0000000..ac22ed6 --- /dev/null +++ b/tests/protocol_test/tests/test_failure_detection.py @@ -0,0 +1,132 @@ +"""Tests for CM failure detection of DS nodes. + +Verifies the full failure detection pipeline: + DS killed → CM heartbeat timeout → mark DEAD → shard rebalance away from dead node +""" + +import time + + +class TestFailureDetection: + """Verify CM detects DS failure and completes shard migration.""" + + def test_single_ds_kill_detected_and_shards_migrated(self, cluster_small): + """Kill one DS; CM marks DEAD; shards migrate to alive nodes.""" + ds0 = cluster_small.get_ds_handle(0) + addr = ds0.addr_str + + # Record state before kill + shard_dist_before = cluster_small.observer.get_shard_distribution() + total_shards_before = sum(shard_dist_before.values()) + ds0_shards_before = shard_dist_before.get(addr, 0) + assert ds0_shards_before > 0, f"DS0 should have shards before kill, got {shard_dist_before}" + + # Kill DS + cluster_small.fault_injector.kill_process(ds0) + assert not cluster_small.process_manager.is_alive(ds0) + + # CM should detect failure within heartbeat_timeout + scan_interval + timeout = ( + cluster_small.config.cm_heartbeat_timeout_inSecs + + cluster_small.config.cm_heartbeat_bg_scan_interval_inSecs + + 5 + ) + assert cluster_small.observer.wait_for_node_status( + addr, "DEAD", timeout=timeout + ), f"CM did not mark {addr} as DEAD within {timeout}s" + + # After DEAD, shards must be migrated away + assert cluster_small.observer.wait_for_rebalance_complete( + timeout=15 + ), "Rebalance did not complete after DS marked DEAD" + + # Verify: dead node has zero shards + cluster_small.observer.assert_node_has_no_shards(addr) + + # Verify: total shard count preserved (no loss, no duplication) + cluster_small.observer.assert_total_shard_count(total_shards_before) + + # Verify: all shards now on alive nodes only + cluster_small.observer.assert_no_orphaned_shards() + + # Verify: alive nodes picked up the orphaned shards + shard_dist_after = cluster_small.observer.get_shard_distribution() + alive_addrs = [ + ds.addr_str for ds in cluster_small.data_servers if ds.index != 0 + ] + for alive_addr in alive_addrs: + after = shard_dist_after.get(alive_addr, 0) + before = shard_dist_before.get(alive_addr, 0) + assert after >= before, ( + f"Alive node {alive_addr} should have gained shards: " + f"before={before}, after={after}" + ) + + def test_detection_timing(self, cluster_small): + """Failure should be detected within expected time bounds.""" + ds0 = cluster_small.get_ds_handle(0) + addr = ds0.addr_str + + kill_time = time.time() + cluster_small.fault_injector.kill_process(ds0) + + assert cluster_small.observer.wait_for_node_status( + addr, "DEAD", timeout=60 + ) + detect_time = time.time() + elapsed = detect_time - kill_time + + # Detection lower bound: at least heartbeat_timeout (CM must wait this long) + min_expected = cluster_small.config.cm_heartbeat_timeout_inSecs + assert elapsed >= min_expected - 1, ( + f"Detection too fast ({elapsed:.1f}s), expected >= {min_expected}s. " + f"CM should not mark DEAD before heartbeat timeout expires." + ) + + # Detection upper bound: heartbeat_timeout + 2 * scan_interval + max_expected = ( + cluster_small.config.cm_heartbeat_timeout_inSecs + + 2 * cluster_small.config.cm_heartbeat_bg_scan_interval_inSecs + ) + assert elapsed <= max_expected + 5, ( + f"Detection took {elapsed:.1f}s, expected <= {max_expected}s" + ) + + def test_remaining_nodes_unaffected(self, cluster_small): + """After killing one DS, remaining DS stay RUNNING with correct shard counts.""" + ds0 = cluster_small.get_ds_handle(0) + + # Record each alive DS's shard count before kill + shard_dist_before = cluster_small.observer.get_shard_distribution() + + cluster_small.fault_injector.kill_process(ds0) + + # Wait for full failure detection + rebalance + timeout = cluster_small.config.cm_heartbeat_timeout_inSecs + 15 + cluster_small.observer.wait_for_node_status(ds0.addr_str, "DEAD", timeout=timeout) + cluster_small.observer.wait_for_rebalance_complete(timeout=15) + + # Remaining DS: process alive + status RUNNING + for ds in cluster_small.data_servers: + if ds.index == 0: + continue + assert cluster_small.process_manager.is_alive(ds), ( + f"DS[{ds.index}] should still be alive" + ) + status = cluster_small.observer.get_node_status(ds.addr_str) + assert status == "RUNNING", ( + f"DS[{ds.index}] status is {status}, expected RUNNING" + ) + + # Remaining DS should have MORE shards than before (they absorbed ds0's shards) + shard_dist_after = cluster_small.observer.get_shard_distribution() + for ds in cluster_small.data_servers: + if ds.index == 0: + continue + before = shard_dist_before.get(ds.addr_str, 0) + after = shard_dist_after.get(ds.addr_str, 0) + assert after > before, ( + f"DS[{ds.index}] ({ds.addr_str}) should have gained shards: " + f"before={before}, after={after}" + ) + diff --git a/tests/protocol_test/tests/test_flag_management.py b/tests/protocol_test/tests/test_flag_management.py new file mode 100644 index 0000000..6561f04 --- /dev/null +++ b/tests/protocol_test/tests/test_flag_management.py @@ -0,0 +1,61 @@ +"""Tests for runtime flag management via admin interface.""" + +import pytest + + +@pytest.mark.requires_rdma +class TestFlagManagement: + """Verify runtime gflag get/set/list via admin UDS.""" + + def test_runtime_flag_change(self, cluster_small): + """Change heartbeat timeout at runtime via admin, verify new value.""" + admin = cluster_small.observer.admin_client + cm = cluster_small.cm + + # Set flag + success = admin.set_flag( + cm.host, cm.pid, + "cm_heartbeat_timeout_inSecs", "15" + ) + assert success, "Failed to set flag" + + # Verify + val = admin.get_flag( + cm.host, cm.pid, + "cm_heartbeat_timeout_inSecs" + ) + assert val == "15", f"Expected flag value '15', got '{val}'" + + def test_list_flags(self, cluster_small): + """List all flags from CM and verify key flags exist.""" + admin = cluster_small.observer.admin_client + cm = cluster_small.cm + + flags = admin.list_flags(cm.host, cm.pid) + assert len(flags) > 0, "No flags returned" + + # Check some expected flags exist + expected_flags = [ + "cm_heartbeat_timeout_inSecs", + "cm_heartbeat_bg_scan_interval_inSecs", + "shard_total_num", + ] + for flag_name in expected_flags: + assert flag_name in flags, f"Expected flag '{flag_name}' not found" + + def test_set_ds_flag(self, cluster_small): + """Set a flag on a data server.""" + admin = cluster_small.observer.admin_client + ds0 = cluster_small.get_ds_handle(0) + + success = admin.set_flag( + ds0.host, ds0.pid, + "heartbeat_cooldown_sec", "3" + ) + assert success, "Failed to set DS flag" + + val = admin.get_flag( + ds0.host, ds0.pid, + "heartbeat_cooldown_sec" + ) + assert val == "3", f"Expected '3', got '{val}'" diff --git a/tests/protocol_test/tests/test_graceful_shutdown.py b/tests/protocol_test/tests/test_graceful_shutdown.py new file mode 100644 index 0000000..651b726 --- /dev/null +++ b/tests/protocol_test/tests/test_graceful_shutdown.py @@ -0,0 +1,53 @@ +"""Tests for graceful shutdown (SIGTERM) handling. + +Verifies: + - DS exits cleanly on SIGTERM + - CM detects absence via HB timeout + - Shards migrate away from stopped node + - CM itself shuts down cleanly on SIGTERM +""" + +import time + + +class TestGracefulShutdown: + """Verify SIGTERM triggers graceful shutdown with full shard migration.""" + + def test_sigterm_ds_detected_and_shards_migrated(self, cluster_small): + """SIGTERM to DS → CM detects DEAD → shards migrate.""" + ds0 = cluster_small.get_ds_handle(0) + addr = ds0.addr_str + + dist_before = cluster_small.observer.get_shard_distribution() + total_before = sum(dist_before.values()) + + # Graceful stop + cluster_small.fault_injector.graceful_stop(ds0) + assert not cluster_small.process_manager.is_alive(ds0), ( + "DS should have exited after SIGTERM" + ) + + # CM detects via heartbeat timeout + timeout = cluster_small.config.cm_heartbeat_timeout_inSecs + 10 + assert cluster_small.observer.wait_for_node_status( + addr, "DEAD", timeout=timeout + ), f"CM did not mark {addr} as DEAD after graceful shutdown" + + # Shards should migrate + cluster_small.observer.wait_for_rebalance_complete(timeout=15) + cluster_small.observer.assert_node_has_no_shards(addr) + cluster_small.observer.assert_total_shard_count(total_before) + cluster_small.observer.assert_no_orphaned_shards() + + def test_sigterm_cm_graceful(self, cluster_small): + """SIGTERM to CM should trigger clean shutdown.""" + cluster_small.fault_injector.graceful_stop(cluster_small.cm) + assert not cluster_small.process_manager.is_alive(cluster_small.cm), ( + "CM should have exited after SIGTERM" + ) + + # All DS should still be alive (DS survives CM shutdown) + for ds in cluster_small.data_servers: + assert cluster_small.process_manager.is_alive(ds), ( + f"DS[{ds.index}] should survive CM shutdown" + ) diff --git a/tests/protocol_test/tests/test_heartbeat.py b/tests/protocol_test/tests/test_heartbeat.py new file mode 100644 index 0000000..790a7d1 --- /dev/null +++ b/tests/protocol_test/tests/test_heartbeat.py @@ -0,0 +1,61 @@ +"""Tests for heartbeat monitoring — steady state verification on both CM and DS sides.""" + +import time + +import pytest + + +class TestHeartbeat: + """Verify heartbeat keeps nodes alive and shard distribution stable.""" + + def test_heartbeat_keeps_alive(self, cluster_small): + """After multiple heartbeat cycles, all nodes RUNNING with shards intact.""" + hb_interval = cluster_small.config.heartbeat_cooldown_sec + time.sleep(hb_interval * 5) + + # CM side: all nodes RUNNING + cluster_small.observer.assert_all_nodes_running() + + # Shard side: distribution unchanged, total correct + cluster_small.observer.assert_total_shard_count( + cluster_small.config.shard_total_num + ) + cluster_small.observer.assert_shard_balance(max_imbalance_ratio=0.1) + + def test_stable_shard_distribution(self, cluster_small): + """Shard distribution should not change during normal operation.""" + dist_before = cluster_small.observer.get_shard_distribution() + total_before = sum(dist_before.values()) + + time.sleep(10) + + dist_after = cluster_small.observer.get_shard_distribution() + total_after = sum(dist_after.values()) + + assert total_before == total_after, ( + f"Shard total changed: {total_before} -> {total_after}" + ) + for addr in dist_before: + assert dist_before[addr] == dist_after.get(addr, 0), ( + f"Node {addr} shard count changed: " + f"{dist_before[addr]} -> {dist_after.get(addr, 0)}" + ) + + @pytest.mark.requires_rdma + def test_ds_status_healthy(self, cluster_small): + """Under normal conditions, all DS report is_registered=true, cm_ready=true, failure_count=0.""" + hb_interval = cluster_small.config.heartbeat_cooldown_sec + time.sleep(hb_interval * 5) + + for ds in cluster_small.data_servers: + cluster_small.observer.assert_ds_is_registered(ds.host, ds.pid) + cluster_small.observer.assert_ds_cm_ready(ds.host, ds.pid, expected=True) + cluster_small.observer.assert_ds_heartbeat_failure_count( + ds.host, ds.pid, min_count=0, max_count=0 + ) + + def test_stable_cluster_state(self, cluster_small): + """Cluster state should remain stable (no status flapping) for 10 seconds.""" + assert cluster_small.observer.wait_for_stable_state( + duration=10, check_interval=2 + ), "Cluster state was not stable for 10 seconds" diff --git a/tests/protocol_test/tests/test_multi_failure.py b/tests/protocol_test/tests/test_multi_failure.py new file mode 100644 index 0000000..1426735 --- /dev/null +++ b/tests/protocol_test/tests/test_multi_failure.py @@ -0,0 +1,106 @@ +"""Tests for simultaneous multiple node failures. + +Verifies: + - CM detects all dead nodes + - Shards from ALL dead nodes migrate to alive nodes + - Total shard count preserved + - Alive nodes unaffected +""" + +import pytest + + +class TestMultiFailure: + """Verify cluster handles multiple simultaneous DS failures.""" + + def test_kill_two_of_six(self, cluster_medium): + """Kill 2/6 DS simultaneously → shards migrate to remaining 4.""" + ds0 = cluster_medium.get_ds_handle(0) + ds1 = cluster_medium.get_ds_handle(1) + + # Record before state + dist_before = cluster_medium.observer.get_shard_distribution() + total_before = sum(dist_before.values()) + killed_shards = ( + dist_before.get(ds0.addr_str, 0) + dist_before.get(ds1.addr_str, 0) + ) + assert killed_shards > 0, "Killed DS should have had shards" + + # Kill simultaneously + cluster_medium.fault_injector.kill_multiple([ds0, ds1]) + + timeout = cluster_medium.config.cm_heartbeat_timeout_inSecs + 15 + + # Both should be detected as DEAD + assert cluster_medium.observer.wait_for_node_status( + ds0.addr_str, "DEAD", timeout=timeout + ), f"DS0 {ds0.addr_str} not marked DEAD" + assert cluster_medium.observer.wait_for_node_status( + ds1.addr_str, "DEAD", timeout=timeout + ), f"DS1 {ds1.addr_str} not marked DEAD" + + # Rebalance must complete + assert cluster_medium.observer.wait_for_rebalance_complete( + timeout=15 + ), "Rebalance did not complete" + + # Dead nodes: zero shards + cluster_medium.observer.assert_node_has_no_shards(ds0.addr_str) + cluster_medium.observer.assert_node_has_no_shards(ds1.addr_str) + + # Total preserved + cluster_medium.observer.assert_total_shard_count(total_before) + + # Balance on remaining 4 nodes + cluster_medium.observer.assert_shard_balance(max_imbalance_ratio=0.3) + cluster_medium.observer.assert_no_orphaned_shards() + + # Alive nodes should all have gained shards + dist_after = cluster_medium.observer.get_shard_distribution() + alive_ds = [ + ds for ds in cluster_medium.data_servers + if ds.index not in (0, 1) + ] + for ds in alive_ds: + before = dist_before.get(ds.addr_str, 0) + after = dist_after.get(ds.addr_str, 0) + assert after >= before, ( + f"DS[{ds.index}] should have gained shards: " + f"before={before}, after={after}" + ) + + def test_kill_two_of_six_sequential(self, cluster_medium): + """Kill 2 DS one by one; verify intermediate and final shard state.""" + ds0 = cluster_medium.get_ds_handle(0) + ds1 = cluster_medium.get_ds_handle(1) + + dist_before = cluster_medium.observer.get_shard_distribution() + total_before = sum(dist_before.values()) + + # Kill first DS + cluster_medium.fault_injector.kill_process(ds0) + + timeout = cluster_medium.config.cm_heartbeat_timeout_inSecs + 15 + assert cluster_medium.observer.wait_for_node_status( + ds0.addr_str, "DEAD", timeout=timeout + ) + cluster_medium.observer.wait_for_rebalance_complete(timeout=15) + + # Intermediate check: ds0 has 0 shards, total preserved, 5 alive + cluster_medium.observer.assert_node_has_no_shards(ds0.addr_str) + cluster_medium.observer.assert_total_shard_count(total_before) + assert cluster_medium.observer.get_alive_node_count() == 5 + + # Kill second DS + cluster_medium.fault_injector.kill_process(ds1) + assert cluster_medium.observer.wait_for_node_status( + ds1.addr_str, "DEAD", timeout=timeout + ) + cluster_medium.observer.wait_for_rebalance_complete(timeout=15) + + # Final check + cluster_medium.observer.assert_node_has_no_shards(ds1.addr_str) + cluster_medium.observer.assert_total_shard_count(total_before) + assert cluster_medium.observer.get_alive_node_count() == 4 + cluster_medium.observer.assert_no_orphaned_shards() + cluster_medium.observer.assert_shard_balance(max_imbalance_ratio=0.3) diff --git a/tests/protocol_test/tests/test_network_partition.py b/tests/protocol_test/tests/test_network_partition.py new file mode 100644 index 0000000..05a5d50 --- /dev/null +++ b/tests/protocol_test/tests/test_network_partition.py @@ -0,0 +1,83 @@ +"""Tests for network partition scenarios. + +NOTE: SiCL uses RDMA transport (InfiniBand/RoCEv2), NOT TCP/IP. +iptables-based partition does NOT work. The actual partition injection method +will be provided separately and plugged into FaultInjector. + +These tests are currently skipped pending the RDMA-aware partition method. +""" + +import pytest + + +@pytest.mark.skip(reason="SiCL uses RDMA transport; iptables partition not applicable. " + "Awaiting RDMA-aware partition injection method.") +class TestNetworkPartition: + """Verify cluster handles network partitions correctly. + + Expected behavior: + - Partition DS from CM → CM marks DS as DEAD, shards migrate + - DS side: heartbeat_failure_count accumulates, cm_ready becomes false + - Heal partition → DS re-registers, shards may be re-assigned + """ + + def test_partition_and_heal(self, cluster_small): + """Partition DS0 from CM → DEAD + shard migration → heal → DS re-registers.""" + ds0 = cluster_small.get_ds_handle(0) + addr = ds0.addr_str + + dist_before = cluster_small.observer.get_shard_distribution() + total_before = sum(dist_before.values()) + + with cluster_small.fault_injector.temporary_partition( + ds0, [cluster_small.cm] + ): + # CM should detect DS0 as DEAD + timeout = cluster_small.config.cm_heartbeat_timeout_inSecs + 10 + assert cluster_small.observer.wait_for_node_status( + addr, "DEAD", timeout=timeout + ), f"CM did not mark partitioned DS {addr} as DEAD" + + # Shards must migrate away from partitioned node + cluster_small.observer.wait_for_rebalance_complete(timeout=15) + cluster_small.observer.assert_node_has_no_shards(addr) + cluster_small.observer.assert_total_shard_count(total_before) + + # After partition heals, DS should re-register + rejoin_timeout = ( + 5 * cluster_small.config.heartbeat_cooldown_sec + + cluster_small.config.register_cooldown_sec + + 15 + ) + assert cluster_small.observer.wait_for_node_status( + addr, "RUNNING", timeout=rejoin_timeout + ), f"DS {addr} did not rejoin after partition healed" + + # Shard total still correct + cluster_small.observer.assert_total_shard_count(total_before) + + def test_partition_shard_rebalance(self, cluster_small): + """During partition, shards rebalanced to alive nodes, total preserved.""" + ds0 = cluster_small.get_ds_handle(0) + + dist_before = cluster_small.observer.get_shard_distribution() + total_before = sum(dist_before.values()) + + with cluster_small.fault_injector.temporary_partition( + ds0, [cluster_small.cm] + ): + timeout = cluster_small.config.cm_heartbeat_timeout_inSecs + 15 + cluster_small.observer.wait_for_rebalance_complete(timeout=timeout) + cluster_small.observer.assert_no_orphaned_shards() + cluster_small.observer.assert_total_shard_count(total_before) + + # Alive nodes should have absorbed the shards + dist_during = cluster_small.observer.get_shard_distribution() + for ds in cluster_small.data_servers: + if ds.index == 0: + continue + before = dist_before.get(ds.addr_str, 0) + during = dist_during.get(ds.addr_str, 0) + assert during >= before, ( + f"DS[{ds.index}] should have gained shards during partition" + ) diff --git a/tests/protocol_test/tests/test_node_join.py b/tests/protocol_test/tests/test_node_join.py new file mode 100644 index 0000000..c4129d8 --- /dev/null +++ b/tests/protocol_test/tests/test_node_join.py @@ -0,0 +1,35 @@ +"""Tests for node registration and initial cluster formation.""" + + +class TestNodeJoin: + """Verify DS nodes register correctly during grace period.""" + + def test_all_nodes_register(self, cluster_small): + """After grace period, all DS should be RUNNING.""" + statuses = cluster_small.observer.get_all_node_statuses() + assert len(statuses) >= cluster_small.config.num_data_servers, ( + f"Expected {cluster_small.config.num_data_servers} nodes, " + f"got {len(statuses)}: {statuses}" + ) + for addr, status in statuses.items(): + assert status == "RUNNING", f"Node {addr} is {status}, expected RUNNING" + + def test_initial_shard_distribution(self, cluster_small): + """Shards should be roughly evenly distributed across DS after initial setup.""" + cluster_small.observer.assert_total_shard_count( + cluster_small.config.shard_total_num + ) + cluster_small.observer.assert_shard_balance(max_imbalance_ratio=0.1) + + def test_ds_processes_alive(self, cluster_small): + """All DS processes should be running.""" + for ds in cluster_small.data_servers: + assert cluster_small.process_manager.is_alive(ds), ( + f"DS[{ds.index}] (pid={ds.pid}) is not alive" + ) + + def test_cm_process_alive(self, cluster_small): + """CM process should be running.""" + assert cluster_small.process_manager.is_alive(cluster_small.cm), ( + f"CM (pid={cluster_small.cm.pid}) is not alive" + ) diff --git a/tests/protocol_test/tests/test_node_rejoin.py b/tests/protocol_test/tests/test_node_rejoin.py new file mode 100644 index 0000000..db93bb8 --- /dev/null +++ b/tests/protocol_test/tests/test_node_rejoin.py @@ -0,0 +1,128 @@ +"""Tests for DS rejoin after failure recovery. + +Verifies both sides: + CM-side: node status DEAD → RUNNING after rejoin, shard migration and recovery + DS-side: cm_ready transitions, heartbeat_failure_count, re-registration +""" + +import pytest + + +@pytest.mark.requires_rdma +class TestNodeRejoin: + """Verify DS can rejoin after freeze/unfreeze or kill/restart.""" + + def test_ds_freeze_unfreeze_rejoin(self, cluster_small): + """Freeze DS0 → CM marks DEAD, shards migrate → unfreeze → DS re-registers.""" + ds0 = cluster_small.get_ds_handle(0) + addr = ds0.addr_str + + # Record shard state before freeze + dist_before = cluster_small.observer.get_shard_distribution() + total_before = sum(dist_before.values()) + + # Freeze DS + cluster_small.fault_injector.freeze_process(ds0) + + # CM detects DEAD + timeout = cluster_small.config.cm_heartbeat_timeout_inSecs + 10 + assert cluster_small.observer.wait_for_node_status( + addr, "DEAD", timeout=timeout + ), f"CM did not mark frozen DS {addr} as DEAD" + + # Shards migrated away + cluster_small.observer.wait_for_rebalance_complete(timeout=15) + cluster_small.observer.assert_node_has_no_shards(addr) + cluster_small.observer.assert_total_shard_count(total_before) + + # Unfreeze DS + cluster_small.fault_injector.unfreeze_process(ds0) + + # DS re-registers + rejoin_timeout = ( + 5 * cluster_small.config.heartbeat_cooldown_sec + + cluster_small.config.register_cooldown_sec + + 15 + ) + assert cluster_small.observer.wait_for_node_status( + addr, "RUNNING", timeout=rejoin_timeout + ), f"DS {addr} did not rejoin after unfreeze" + + # DS-side: verify internal state recovered + assert cluster_small.observer.wait_for_ds_registered( + ds0.host, ds0.pid, timeout=10 + ), "DS did not report registered after rejoin" + + cluster_small.observer.assert_ds_cm_ready(ds0.host, ds0.pid, expected=True) + cluster_small.observer.assert_ds_heartbeat_failure_count( + ds0.host, ds0.pid, min_count=0, max_count=0 + ) + + # Shard total still correct + cluster_small.observer.assert_total_shard_count(total_before) + + def test_ds_freeze_shows_cm_not_ready_after_unfreeze(self, cluster_small): + """After unfreeze, DS briefly has cm_ready=false before re-registering.""" + ds0 = cluster_small.get_ds_handle(0) + addr = ds0.addr_str + + cluster_small.fault_injector.freeze_process(ds0) + + timeout = cluster_small.config.cm_heartbeat_timeout_inSecs + 10 + cluster_small.observer.wait_for_node_status(addr, "DEAD", timeout=timeout) + + cluster_small.fault_injector.unfreeze_process(ds0) + + # DS should eventually re-register and recover + rejoin_timeout = ( + 5 * cluster_small.config.heartbeat_cooldown_sec + + cluster_small.config.register_cooldown_sec + + 15 + ) + assert cluster_small.observer.wait_for_node_status( + addr, "RUNNING", timeout=rejoin_timeout + ) + + # Final state: registered and cm_ready + cluster_small.observer.assert_ds_is_registered(ds0.host, ds0.pid) + cluster_small.observer.assert_ds_cm_ready(ds0.host, ds0.pid, expected=True) + + def test_ds_restart_rejoin(self, cluster_small): + """Kill DS0, wait for DEAD + shard migration, restart, verify re-registration.""" + ds0 = cluster_small.get_ds_handle(0) + addr = ds0.addr_str + + dist_before = cluster_small.observer.get_shard_distribution() + total_before = sum(dist_before.values()) + + # Kill DS + cluster_small.fault_injector.kill_process(ds0) + + # Wait for detection and shard migration + timeout = cluster_small.config.cm_heartbeat_timeout_inSecs + 10 + assert cluster_small.observer.wait_for_node_status( + addr, "DEAD", timeout=timeout + ) + cluster_small.observer.wait_for_rebalance_complete(timeout=15) + cluster_small.observer.assert_node_has_no_shards(addr) + cluster_small.observer.assert_total_shard_count(total_before) + + # Restart DS with same ports (joins as new node) + new_ds = cluster_small.add_data_server(ports=ds0.ports) + + # Wait for re-registration + rejoin_timeout = cluster_small.config.register_cooldown_sec + 15 + assert cluster_small.observer.wait_for_node_count( + cluster_small.config.num_data_servers, timeout=rejoin_timeout + ), "DS did not rejoin after restart" + + # CM side: all RUNNING + cluster_small.observer.assert_all_nodes_running() + + # DS side: new DS should be registered and cm_ready + assert cluster_small.observer.wait_for_ds_registered( + new_ds.host, new_ds.pid, timeout=10 + ), "Restarted DS did not report registered" + + # Shard total correct + cluster_small.observer.assert_total_shard_count(total_before) diff --git a/tests/protocol_test/tests/test_rebalance.py b/tests/protocol_test/tests/test_rebalance.py new file mode 100644 index 0000000..1daf188 --- /dev/null +++ b/tests/protocol_test/tests/test_rebalance.py @@ -0,0 +1,122 @@ +"""Tests for shard rebalancing after node failure. + +Verifies: + - Orphaned shards are migrated away from dead node + - Total shard count is preserved (no loss, no duplication) + - Post-rebalance distribution is balanced across alive nodes + - Rebalance happens promptly after DEAD detection +""" + +import time + + +class TestRebalance: + """Verify shard rebalancing correctness after DS failure.""" + + def test_shards_migrated_from_dead_node(self, cluster_small): + """Kill DS0 → dead node must have 0 shards, all moved to alive nodes.""" + ds0 = cluster_small.get_ds_handle(0) + addr = ds0.addr_str + + # Record before state + dist_before = cluster_small.observer.get_shard_distribution() + ds0_shards_before = dist_before.get(addr, 0) + assert ds0_shards_before > 0, f"DS0 has no shards to migrate: {dist_before}" + + cluster_small.fault_injector.kill_process(ds0) + + # Wait for detection + rebalance + timeout = cluster_small.config.cm_heartbeat_timeout_inSecs + 15 + cluster_small.observer.wait_for_node_status(addr, "DEAD", timeout=timeout) + cluster_small.observer.wait_for_rebalance_complete(timeout=15) + + # Dead node: zero shards + cluster_small.observer.assert_node_has_no_shards(addr) + + # Alive nodes absorbed the orphaned shards + dist_after = cluster_small.observer.get_shard_distribution() + for ds in cluster_small.data_servers: + if ds.index == 0: + continue + before = dist_before.get(ds.addr_str, 0) + after = dist_after.get(ds.addr_str, 0) + assert after > before, ( + f"DS[{ds.index}] ({ds.addr_str}) didn't gain shards: " + f"before={before}, after={after}" + ) + + def test_shard_total_preserved(self, cluster_small): + """Rebalance must not lose or duplicate shards.""" + total_before = cluster_small.observer.get_total_shard_count() + assert total_before == cluster_small.config.shard_total_num + + ds0 = cluster_small.get_ds_handle(0) + cluster_small.fault_injector.kill_process(ds0) + + timeout = cluster_small.config.cm_heartbeat_timeout_inSecs + 15 + cluster_small.observer.wait_for_node_status(ds0.addr_str, "DEAD", timeout=timeout) + cluster_small.observer.wait_for_rebalance_complete(timeout=15) + + cluster_small.observer.assert_total_shard_count(total_before) + + def test_shard_balance_after_rebalance(self, cluster_small): + """After rebalance, shards balanced across remaining alive nodes.""" + ds0 = cluster_small.get_ds_handle(0) + cluster_small.fault_injector.kill_process(ds0) + + timeout = cluster_small.config.cm_heartbeat_timeout_inSecs + 15 + cluster_small.observer.wait_for_node_status(ds0.addr_str, "DEAD", timeout=timeout) + cluster_small.observer.wait_for_rebalance_complete(timeout=15) + + # 3 DS, 1 killed → 64 shards on 2 nodes → ~32 each + cluster_small.observer.assert_shard_balance(max_imbalance_ratio=0.3) + + # Verify: each alive node has roughly total/alive shards + dist = cluster_small.observer.get_shard_distribution() + alive_count = sum(1 for ds in cluster_small.data_servers if ds.index != 0) + expected_per_node = cluster_small.config.shard_total_num / alive_count + for ds in cluster_small.data_servers: + if ds.index == 0: + continue + node_shards = dist.get(ds.addr_str, 0) + assert node_shards > 0, ( + f"DS[{ds.index}] has 0 shards after rebalance" + ) + # Allow ±30% deviation from ideal + assert abs(node_shards - expected_per_node) / expected_per_node <= 0.3, ( + f"DS[{ds.index}] has {node_shards} shards, " + f"expected ~{expected_per_node:.0f}" + ) + + def test_rebalance_timing(self, cluster_small): + """Rebalance should start immediately after DEAD detection, not delayed.""" + ds0 = cluster_small.get_ds_handle(0) + addr = ds0.addr_str + + cluster_small.fault_injector.kill_process(ds0) + + # Wait for DEAD and record the timestamp + timeout = cluster_small.config.cm_heartbeat_timeout_inSecs + 10 + assert cluster_small.observer.wait_for_node_status( + addr, "DEAD", timeout=timeout + ) + dead_time = time.time() + + # Rebalance should complete quickly after DEAD (routing table update is synchronous) + assert cluster_small.observer.wait_for_rebalance_complete( + timeout=10 + ), "Rebalance did not complete within 10s of DEAD detection" + rebalance_done_time = time.time() + + # Shards should have moved within a few seconds of DEAD + elapsed = rebalance_done_time - dead_time + assert elapsed < 10, ( + f"Rebalance took {elapsed:.1f}s after DEAD detection, expected < 10s" + ) + + # Verify the end state + cluster_small.observer.assert_node_has_no_shards(addr) + cluster_small.observer.assert_total_shard_count( + cluster_small.config.shard_total_num + ) + diff --git a/tests/protocol_test/tests/test_scenario.py b/tests/protocol_test/tests/test_scenario.py new file mode 100644 index 0000000..4c2cdc9 --- /dev/null +++ b/tests/protocol_test/tests/test_scenario.py @@ -0,0 +1,47 @@ +"""Scenario-driven tests — add new test cases by writing YAML, not Python. + +Each test loads a YAML scenario file and feeds it to the ScenarioRunner. +To add a new scenario: create a YAML in scenarios/full/ and add a one-liner here. +""" + +from pathlib import Path + +import pytest + +from framework.scenario_runner import ScenarioRunner + +SCENARIO_DIR = Path(__file__).parents[1] / "scenarios" + + +@pytest.fixture +def runner(binary_dir, tmp_path): + return ScenarioRunner( + binary_dir=str(binary_dir) if binary_dir is not None else None, + log_dir=str(tmp_path / "logs"), + ) + + +class TestScenario: + """YAML-driven integration test scenarios.""" + + def test_kill_ds_and_verify_rebalance(self, runner): + """Kill one DS → DEAD → shard rebalance → no orphaned shards.""" + runner.run_scenario( + SCENARIO_DIR / "full" / "kill_ds_and_verify_rebalance.yaml" + ) + + def test_cm_restart_ds_rejoin(self, runner): + """Kill CM → restart → all DS re-register → shard table rebuilt.""" + runner.run_scenario( + SCENARIO_DIR / "full" / "cm_restart_ds_rejoin.yaml" + ) + + def test_composable_scenario(self, runner): + """Demonstrates YAML composition: small cluster + fast heartbeat + kill fault.""" + runner.run_scenario( + SCENARIO_DIR / "clusters" / "small.yaml", + SCENARIO_DIR / "overrides" / "fast_heartbeat.yaml", + SCENARIO_DIR / "faults" / "kill_one_ds.yaml", + ) + # Note: this scenario has faults but no validations section, + # so it only verifies that the cluster survives the fault without crashing. diff --git a/tools/simm_ctl_admin.cc b/tools/simm_ctl_admin.cc index 4895a08..11b645d 100644 --- a/tools/simm_ctl_admin.cc +++ b/tools/simm_ctl_admin.cc @@ -45,6 +45,7 @@ #include "rpc/rpc_context.h" #include "transport/types.h" +#include "common/admin/admin_msg_types.h" #include "common/errcode/errcode_def.h" #include "common/logging/logging.h" #include "common/trace/trace_server.h" @@ -86,13 +87,7 @@ class AdminChannel { const std::shared_ptr &)> done_cb) = 0; }; -// Message type for UDS admin channel -enum class AdminMsgType : uint16_t { - TRACE_TOGGLE = 1, - GFLAG_LIST = 2, - GFLAG_GET = 3, - GFLAG_SET = 4, -}; +using simm::common::AdminMsgType; // Unix domain socket implementation class UdsChannel : public AdminChannel { @@ -333,6 +328,61 @@ static void InitRpcClientAndContext(std::unique_ptr &rpc_clien ctx_shared = std::shared_ptr(ctx_p); } +// Resolve process name to PID via pgrep. +// Returns the PID on success. +// Prints error and exits if not found or ambiguous. +static int resolveProcessPid(const std::string &proc_name) { + std::string cmd = "pgrep -x " + proc_name; + FILE *fp = popen(cmd.c_str(), "r"); + if (!fp) { + std::cerr << "Error: failed to run pgrep\n"; + exit(1); + } + std::vector pids; + char buf[64]; + while (fgets(buf, sizeof(buf), fp)) { + std::string line(buf); + // trim whitespace + line.erase(line.find_last_not_of(" \t\n\r") + 1); + if (!line.empty()) { + try { + pids.push_back(std::stoi(line)); + } catch (...) {} + } + } + pclose(fp); + + if (pids.empty()) { + std::cerr << "Error: " << proc_name << ": process not found\n"; + exit(1); + } + if (pids.size() > 1) { + std::cerr << "Error: " << proc_name << ": ambiguous, found " + << pids.size() << " processes (pids:"; + for (int p : pids) { + std::cerr << " " << p; + } + std::cerr << ")\n"; + exit(1); + } + return pids[0]; +} + +// Map process name to UDS socket base path. +static std::string procToSocketPath(const std::string &proc_name, int pid) { + std::string base; + if (proc_name == "cluster_manager") { + base = simm::common::kCmAdminUdsBasePath; + } else if (proc_name == "data_server") { + base = simm::common::kDsAdminUdsBasePath; + } else { + std::cerr << "Error: unknown process name: " << proc_name + << ". Expected 'cluster_manager' or 'data_server'\n"; + exit(1); + } + return base + "." + std::to_string(pid) + ".sock"; +} + static void CallbackNode(const std::string &operation, const std::string &name, const std::string &value, @@ -346,6 +396,8 @@ static void CallbackNode(const std::string &operation, InitRpcClientAndContext(rpc_client, ctx_shared); if (operation == "list" || operation == "summary") { + ListNodesRequestPB req; + auto resp = new ListNodesResponsePB(); auto done_cb = [&](const google::protobuf::Message *rsp, const std::shared_ptr ctx) { if (ctx->Failed()) { std::cerr << "Error: RPC failed, err: " << ctx->ErrorText() << "\n"; @@ -482,9 +534,6 @@ static void CallbackNode(const std::string &operation, delete static_cast(rsp); done_latch.count_down(); }; - - ListNodesRequestPB req; - auto resp = new ListNodesResponsePB(); rpc_client->SendRequest(ip, port, static_cast(simm::common::CommonRpcType::RPC_LIST_NODE_REQ), @@ -603,6 +652,13 @@ static void CallbackNode(const std::string &operation, } } + SetNodeStatusRequestPB req; + auto *node_addr = req.mutable_node(); + node_addr->set_ip(node_ip); + node_addr->set_port(std::stoi(node_port_str)); + req.set_node_status(status_value); + + auto resp = new SetNodeStatusResponsePB(); auto done_cb = [&](const google::protobuf::Message *rsp, const std::shared_ptr ctx) { if (ctx->Failed()) { LOG_ERROR("RPC failed, err:{}", ctx->ErrorText()); @@ -622,16 +678,9 @@ static void CallbackNode(const std::string &operation, std::cerr << "Error: SetNodeStatus RPC failed with ret_code: " << response->ret_code() << "\n"; } } + delete resp; done_latch.count_down(); }; - - SetNodeStatusRequestPB req; - auto *node_addr = req.mutable_node(); - node_addr->set_ip(node_ip); - node_addr->set_port(std::stoi(node_port_str)); - req.set_node_status(status_value); - - auto resp = new SetNodeStatusResponsePB(); rpc_client->SendRequest(ip, port, static_cast(simm::common::CommonRpcType::RPC_SET_NODE_STATUS_REQ), @@ -647,6 +696,91 @@ static void CallbackNode(const std::string &operation, done_latch.wait(); } +static void CallbackDsStatus(AdminChannel &channel) { + // Query DS internal status via UDS (is_registered, cm_ready, heartbeat_failure_count) + proto::common::AdmDsStatusRequestPB req; + auto *resp = new proto::common::AdmDsStatusResponsePB(); + std::latch done_latch(1); + + if (!channel.Call( + req, + resp, + [&](const google::protobuf::Message *rsp, + const std::shared_ptr &ctx) { + const auto *response = dynamic_cast(rsp); + if (ctx && ctx->Failed()) { + std::cerr << "Error: UDS call failed, err: " << ctx->ErrorText() << "\n"; + } else if (response && response->ret_code() == CommonErr::OK) { + tabulate::Table tbl; + tbl.format().locale("C"); + tbl.add_row({"Field", "Value"}); + tbl.add_row({"is_registered", response->is_registered() ? "true" : "false"}); + tbl.add_row({"cm_ready", response->cm_ready() ? "true" : "false"}); + tbl.add_row({"heartbeat_failure_count", + std::to_string(response->heartbeat_failure_count())}); + tbl.column(0).format().width(28).font_style({tabulate::FontStyle::bold}); + tbl.column(1).format().width(20); + tbl.row(0).format().font_style({tabulate::FontStyle::bold}); + std::cout << tbl << std::endl; + } else { + std::cerr << "Error: DsStatus failed with ret_code: " + << (response ? response->ret_code() : -1) << "\n"; + } + done_latch.count_down(); + delete resp; + })) { + std::cerr << "ds status: channel.Call() failed\n"; + delete resp; + return; + } + + done_latch.wait(); +} + +static void CallbackCmStatus(AdminChannel &channel) { + proto::common::AdmCmStatusRequestPB req; + auto *resp = new proto::common::AdmCmStatusResponsePB(); + std::latch done_latch(1); + + if (!channel.Call( + req, + resp, + [&](const google::protobuf::Message *rsp, + const std::shared_ptr &ctx) { + const auto *response = dynamic_cast(rsp); + if (ctx && ctx->Failed()) { + std::cerr << "Error: UDS call failed, err: " << ctx->ErrorText() << "\n"; + } else if (response && response->ret_code() == CommonErr::OK) { + tabulate::Table tbl; + tbl.format().locale("C"); + tbl.add_row({"Field", "Value"}); + tbl.add_row({"is_running", response->is_running() ? "true" : "false"}); + tbl.add_row({"service_ready", response->service_ready() ? "true" : "false"}); + tbl.add_row({"alive_node_count", + std::to_string(response->alive_node_count())}); + tbl.add_row({"dead_node_count", + std::to_string(response->dead_node_count())}); + tbl.add_row({"total_shard_count", + std::to_string(response->total_shard_count())}); + tbl.column(0).format().width(28).font_style({tabulate::FontStyle::bold}); + tbl.column(1).format().width(20); + tbl.row(0).format().font_style({tabulate::FontStyle::bold}); + std::cout << tbl << std::endl; + } else { + std::cerr << "Error: CmStatus failed with ret_code: " + << (response ? response->ret_code() : -1) << "\n"; + } + done_latch.count_down(); + delete resp; + })) { + std::cerr << "cm status: channel.Call() failed\n"; + delete resp; + return; + } + + done_latch.wait(); +} + static void CallbackShard(const std::string &operation, [[maybe_unused]] const std::string &name, [[maybe_unused]] const std::string &value, @@ -660,6 +794,8 @@ static void CallbackShard(const std::string &operation, InitRpcClientAndContext(rpc_client, ctx_shared); if (operation == "list") { + QueryShardRoutingTableAllRequestPB req; + auto resp = new QueryShardRoutingTableAllResponsePB(); auto done_cb = [&](const google::protobuf::Message *rsp, const std::shared_ptr ctx) { if (ctx->Failed()) { LOG_ERROR("RPC failed, err:{}", ctx->ErrorText()); @@ -719,10 +855,9 @@ static void CallbackShard(const std::string &operation, std::cerr << "Error: ListShards RPC failed with ret_code: " << response->ret_code() << "\n"; } } + delete resp; done_latch.count_down(); }; - QueryShardRoutingTableAllRequestPB req; - auto resp = new QueryShardRoutingTableAllResponsePB(); rpc_client->SendRequest(ip, port, static_cast(simm::common::CommonRpcType::RPC_LIST_SHARD_REQ), @@ -738,6 +873,152 @@ static void CallbackShard(const std::string &operation, done_latch.wait(); } +static void CallbackNodeByUds(AdminChannel &channel, bool verbose) { + std::latch done_latch(1); + + ListNodesRequestPB req; + auto *resp = new ListNodesResponsePB(); + if (!channel.Call( + req, resp, + [&](const google::protobuf::Message *rsp, + const std::shared_ptr &ctx) { + const auto *response = dynamic_cast(rsp); + if (ctx && ctx->Failed()) { + std::cerr << "Error: UDS call failed, err: " << ctx->ErrorText() << "\n"; + } else if (response && response->ret_code() == CommonErr::OK) { + tabulate::Table tbl; + tbl.format().locale("C"); + + if (verbose) { + tbl.add_row({"Node Address", "Status", "Total Memory (MB)", + "Free Memory (MB)", "Used Memory (MB)"}) + .format() + .width(20); + for (int i = 0; i < response->nodes_size(); ++i) { + const auto &node_info = response->nodes(i); + const auto &node_addr = node_info.node_address(); + std::string addr_str = + node_addr.ip() + ":" + std::to_string(node_addr.port()); + std::string_view status_str = + simm::common::NodeStatusToString( + static_cast(node_info.node_status())); + std::string total_mem = + std::to_string(node_info.resource().mem_total_bytes() / (1024 * 1024)); + std::string free_mem = + std::to_string(node_info.resource().mem_free_bytes() / (1024 * 1024)); + std::string used_mem = + std::to_string(node_info.resource().mem_used_bytes() / (1024 * 1024)); + tbl.add_row({addr_str, status_str, total_mem, free_mem, used_mem}); + } + tbl.column(0).format().width(20).font_style({tabulate::FontStyle::bold}); + tbl.column(1).format().width(12); + tbl.column(2).format().width(18); + tbl.column(3).format().width(18); + tbl.column(4).format().width(18); + tbl.row(0).format().font_style({tabulate::FontStyle::bold}); + } else { + tbl.add_row({"Node Address", "Status"}).format().width(20); + for (int i = 0; i < response->nodes_size(); ++i) { + const auto &node_info = response->nodes(i); + const auto &node_addr = node_info.node_address(); + std::string addr_str = + node_addr.ip() + ":" + std::to_string(node_addr.port()); + std::string status_str = + (node_info.node_status() == 1) ? "RUNNING" : "DEAD"; + tbl.add_row({addr_str, status_str}); + } + tbl.column(0).format().width(20).font_style({tabulate::FontStyle::bold}); + tbl.column(1).format().width(12); + tbl.row(0).format().font_style({tabulate::FontStyle::bold}); + } + std::cout << tbl << std::endl; + } else { + std::cerr << "Error: ListNodes failed with ret_code: " + << (response ? response->ret_code() : -1) << "\n"; + } + done_latch.count_down(); + delete resp; + })) { + std::cerr << "node list: channel.Call() failed\n"; + delete resp; + return; + } + + done_latch.wait(); +} + +static void CallbackShardByUds(AdminChannel &channel, bool verbose) { + std::latch done_latch(1); + + QueryShardRoutingTableAllRequestPB req; + auto *resp = new QueryShardRoutingTableAllResponsePB(); + if (!channel.Call( + req, resp, + [&](const google::protobuf::Message *rsp, + const std::shared_ptr &ctx) { + const auto *response = + dynamic_cast(rsp); + if (ctx && ctx->Failed()) { + std::cerr << "Error: UDS call failed, err: " << ctx->ErrorText() << "\n"; + } else if (response && response->ret_code() == CommonErr::OK) { + tabulate::Table tbl; + tbl.format().locale("C"); + if (verbose) { + tbl.add_row({"Data Server", "Shard IDs"}).format().width(30); + std::map> ds_shards; + for (int i = 0; i < response->shard_info_size(); ++i) { + const auto &shard_entry = response->shard_info(i); + const auto &ds_addr = shard_entry.data_server_address(); + std::string ds_str = + ds_addr.ip() + ":" + std::to_string(ds_addr.port()); + for (int j = 0; j < shard_entry.shard_ids_size(); ++j) { + ds_shards[ds_str].push_back(shard_entry.shard_ids(j)); + } + } + for (const auto &[ds_str, shard_ids] : ds_shards) { + std::string shard_ids_str; + for (size_t j = 0; j < shard_ids.size(); ++j) { + if (j > 0) shard_ids_str += ", "; + shard_ids_str += std::to_string(shard_ids[j]); + } + tbl.add_row({ds_str, shard_ids_str}); + } + tbl.column(0).format().width(20).font_style({tabulate::FontStyle::bold}); + tbl.column(1).format().width(50); + tbl.row(0).format().font_style({tabulate::FontStyle::bold}); + } else { + tbl.add_row({"Data Server", "Shard Count"}).format().width(30); + std::map ds_shard_count; + for (int i = 0; i < response->shard_info_size(); ++i) { + const auto &shard_entry = response->shard_info(i); + const auto &ds_addr = shard_entry.data_server_address(); + std::string ds_str = + ds_addr.ip() + ":" + std::to_string(ds_addr.port()); + ds_shard_count[ds_str] += shard_entry.shard_ids_size(); + } + for (const auto &[ds_str, count] : ds_shard_count) { + tbl.add_row({ds_str, std::to_string(count)}); + } + tbl.column(0).format().width(20).font_style({tabulate::FontStyle::bold}); + tbl.column(1).format().width(15); + tbl.row(0).format().font_style({tabulate::FontStyle::bold}); + } + std::cout << tbl << std::endl; + } else { + std::cerr << "Error: ListShards failed with ret_code: " + << (response ? response->ret_code() : -1) << "\n"; + } + done_latch.count_down(); + delete resp; + })) { + std::cerr << "shard list: channel.Call() failed\n"; + delete resp; + return; + } + + done_latch.wait(); +} + static void CallbackGFlag(AdminChannel &channel, const std::string &operation, const std::string &name, @@ -880,6 +1161,7 @@ static void CallbackTrace(AdminChannel &channel, const std::string &value, bool delete resp; })) { std::cerr << "trace: channel.Call() failed" << std::endl; + delete resp; return; } done_latch.wait(); @@ -889,6 +1171,7 @@ int main(int argc, char *argv[]) { std::string ip; int port; int pid; + std::string proc; std::string subcommand; std::string resource_type; std::string operation; @@ -897,9 +1180,11 @@ int main(int argc, char *argv[]) { po::options_description desc("SiMMCtl - SiMM RPC Management Tool"); desc.add_options()("help,h", "Show help message")( - "ip,i", po::value(&ip)->required(), "Target Cluster Manager IP address")( - "port,p", po::value(&port)->default_value(30002), "Target Cluster Manager port")( - "pid,P", po::value(&pid)->default_value(-1), "Target Client PID, used when tracing is enabled")( + "ip,i", po::value(&ip)->default_value(""), "Target IP address for RPC-based commands")( + "port,p", po::value(&port)->default_value(30002), "Target port for RPC-based commands")( + "pid,P", po::value(&pid)->default_value(-1), "Target process PID for UDS-based commands")( + "proc", po::value(&proc)->default_value(""), + "Target process name for UDS-based commands (cluster_manager or data_server)")( "verbose,v", po::bool_switch(&verbose), "Enable verbose output"); po::options_description hidden("Hidden options"); @@ -924,22 +1209,51 @@ int main(int argc, char *argv[]) { << " node summary [OPTIONS] Show cluster-wide node resource summary\n" << " node stat Show detailed resource stats for one node\n" << " node set Set node status (0=DEAD, 1=RUNNING)\n" + << " cm status Query CM internal status via UDS\n" + << " ds status Query DS internal status via UDS\n" << " shard list [OPTIONS] List all shards\n" << " gflag list [OPTIONS] List all gflags\n" << " gflag get [OPTIONS] Get a gflag value\n" << " gflag set Set a gflag value\n" - << " trace Set tracing status (0=OFF, 1=ON)\n"; + << " trace Set tracing status (0=OFF, 1=ON)\n" + << "\nUDS options (--pid or --proc, mutually exclusive):\n" + << " --pid Target process by PID\n" + << " --proc Target process by name (cluster_manager or data_server)\n"; return 0; } po::notify(vm); + // --pid and --proc are mutually exclusive + if (pid != -1 && !proc.empty()) { + std::cerr << "Error: --pid and --proc are mutually exclusive\n"; + return 1; + } + + // Resolve --proc to --pid + if (!proc.empty()) { + if (proc != "cluster_manager" && proc != "data_server") { + std::cerr << "Error: --proc must be 'cluster_manager' or 'data_server'\n"; + return 1; + } + pid = resolveProcessPid(proc); + } + if (!vm.count("subcommand")) { std::cerr << "Error: No subcommand specified\n"; std::cerr << "Use 'simmctl -h' for help\n"; return 1; } + // Determine UDS socket path from pid and subcommand/proc context + auto buildUdsSocketPath = [&](const std::string &basePath) -> std::string { + if (!proc.empty()) { + return procToSocketPath(proc, pid); + } + // When using --pid, infer from subcommand context + return basePath + "." + std::to_string(pid) + ".sock"; + }; + // Parse subcommand format: "resource_type operation" if (subcommand == "node") { if (args.empty()) { @@ -947,26 +1261,87 @@ int main(int argc, char *argv[]) { return 1; } operation = args[0]; - if (operation == "list") { - CallbackNode("list", "", "", ip, port, verbose); - } else if (operation == "summary") { - CallbackNode("summary", "", "", ip, port, verbose); - } else if (operation == "stat") { - if (args.size() < 2) { - std::cerr << "Error: node stat requires argument: \n"; + if (pid != -1) { + // UDS mode (node list only) + if (operation == "list") { + std::string socket_path = buildUdsSocketPath(simm::common::kCmAdminUdsBasePath); + auto uds_channel = std::make_unique(socket_path, AdminMsgType::NODE_LIST); + if (!uds_channel->Init()) { + std::cerr << "Error: failed to connect to admin socket: " << socket_path << "\n"; + return 1; + } + CallbackNodeByUds(*uds_channel, verbose); + } else { + std::cerr << "Error: Unknown node operation for UDS mode: " << operation << "\n"; return 1; } - CallbackNode("stat", args[1], "", ip, port, verbose); - } else if (operation == "set") { - if (args.size() < 3) { - std::cerr << "Error: node set requires arguments: \n"; + } else { + // RPC mode + if (operation == "list") { + CallbackNode("list", "", "", ip, port, verbose); + } else if (operation == "summary") { + CallbackNode("summary", "", "", ip, port, verbose); + } else if (operation == "stat") { + if (args.size() < 2) { + std::cerr << "Error: node stat requires argument: \n"; + return 1; + } + CallbackNode("stat", args[1], "", ip, port, verbose); + } else if (operation == "set") { + if (args.size() < 3) { + std::cerr << "Error: node set requires arguments: \n"; + return 1; + } + std::string node_addr = args[1]; + std::string node_status = args[2]; + CallbackNode("set", node_addr, node_status, ip, port, verbose); + } else { + std::cerr << "Error: Unknown node operation: " << operation << "\n"; return 1; } - std::string node_addr = args[1]; - std::string node_status = args[2]; - CallbackNode("set", node_addr, node_status, ip, port, verbose); + } + } else if (subcommand == "cm") { + if (args.empty()) { + std::cerr << "Error: cm subcommand requires an operation (status)\n"; + return 1; + } + operation = args[0]; + if (operation == "status") { + if (pid == -1) { + std::cerr << "Error: cm status requires --pid or --proc\n"; + return 1; + } + std::string socket_path = buildUdsSocketPath(simm::common::kCmAdminUdsBasePath); + auto uds_channel = std::make_unique(socket_path, AdminMsgType::CM_STATUS); + if (!uds_channel->Init()) { + std::cerr << "Error: failed to connect to admin socket: " << socket_path << "\n"; + return 1; + } + CallbackCmStatus(*uds_channel); } else { - std::cerr << "Error: Unknown node operation: " << operation << "\n"; + std::cerr << "Error: Unknown cm operation: " << operation << "\n"; + return 1; + } + } else if (subcommand == "ds") { + if (args.empty()) { + std::cerr << "Error: ds subcommand requires an operation (status)\n"; + return 1; + } + operation = args[0]; + if (operation == "status") { + if (pid == -1) { + std::cerr << "Error: ds status requires --pid or --proc\n"; + return 1; + } + std::string socket_path = buildUdsSocketPath(simm::common::kDsAdminUdsBasePath); + auto uds_channel = std::make_unique(socket_path, AdminMsgType::DS_STATUS); + if (!uds_channel->Init()) { + std::cerr << "Error: failed to connect to admin socket: " << socket_path << "\n"; + return 1; + } + CallbackDsStatus(*uds_channel); + } else { + std::cerr << "Error: Unknown ds operation: " << operation << "\n"; return 1; } } else if (subcommand == "shard") { @@ -976,7 +1351,19 @@ int main(int argc, char *argv[]) { } operation = args[0]; if (operation == "list") { - CallbackShard("list", "", "", ip, port, verbose); + if (pid != -1) { + // UDS mode + std::string socket_path = buildUdsSocketPath(simm::common::kCmAdminUdsBasePath); + auto uds_channel = std::make_unique(socket_path, AdminMsgType::SHARD_LIST); + if (!uds_channel->Init()) { + std::cerr << "Error: failed to connect to admin socket: " << socket_path << "\n"; + return 1; + } + CallbackShardByUds(*uds_channel, verbose); + } else { + // RPC mode (original code) + CallbackShard("list", "", "", ip, port, verbose); + } } else { std::cerr << "Error: Unknown shard operation: " << operation << "\n"; return 1; @@ -1017,7 +1404,10 @@ int main(int argc, char *argv[]) { } channel_ptr = std::move(rpc_channel); } else { - std::string socket_path = "/run/simm/simm_trace." + std::to_string(pid) + ".sock"; + // When --proc is used, route through admin socket; otherwise legacy trace socket + std::string socket_path = proc.empty() + ? "/run/simm/simm_trace." + std::to_string(pid) + ".sock" + : procToSocketPath(proc, pid); AdminMsgType msg_type; if (subcommand == "trace") { msg_type = AdminMsgType::TRACE_TOGGLE;