Skip to content

[Fix]: ds will crash(with SIGTERM) without doing cleanup when repeatedly failed to connect to CM #25

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Merged
Sebastian-Keith merged 1 commit into from
Apr 7, 2026
Merged

[Fix]: ds will crash(with SIGTERM) without doing cleanup when repeatedly failed to connect to CM #25

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion src/data_server/kv_rpc_handler.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
#include "data_server/kv_hash_table.h"
#include "transport/types.h"

#include "common/context/context.h"`
#include "common/context/context.h"
#include "common/errcode/errcode_def.h"
#include "common/logging/logging.h"
#include "common/metrics/metrics.h"
Expand Down
53 changes: 41 additions & 12 deletions src/data_server/kv_rpc_service.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,12 @@ KVRpcService::~KVRpcService() {
all_tables_.clear();
}

void KVRpcService::SetClusterDisconnectHandler(std::function<void()> handler) {
if (handler) {
cluster_disconnect_handler_ = std::move(handler);
}
}

error_code_t KVRpcService::Init() {
if (!initialized_) {
sicl::rpc::SiRPC *io_service = nullptr;
Expand Down Expand Up @@ -164,8 +170,13 @@ error_code_t KVRpcService::Stop() {
return CommonErr::OK; // Already stopped
}
is_registered_ = true;
cm_ready_ = true;
register_condv_.post();
is_running_ = false;
{
std::unique_lock<std::mutex> heartbeat_lock(heartbeat_mutex_);
heartbeat_condv_.notify_all();
}
StopRPCServices();
return CommonErr::OK;
}
Expand Down Expand Up @@ -298,13 +309,13 @@ void KVRpcService::SetCMAddressFromK8S() {
}

void KVRpcService::KeepAlive() {
while (true) {
while (!is_registered_) {
while (is_running_) {
while (is_running_ && !is_registered_) {
RegisterOnCluster();
register_condv_.try_wait_for(std::chrono::seconds(FLAGS_register_cooldown_sec));
register_condv_.reset(); // folly::Bation should be reset before reuse
}
while (!cm_ready_) {
while (is_running_ && !cm_ready_) {
register_condv_.reset();

sicl::rpc::SiRPC *mgt_client = nullptr;
Expand All @@ -315,12 +326,16 @@ void KVRpcService::KeepAlive() {
RegisterToRestartedManager();
register_condv_.try_wait_for(std::chrono::seconds(FLAGS_cm_connect_retry_interval_sec));
}
if (!is_running_) {
break;
}
{
std::unique_lock<std::mutex> heartbeat_lock(heartbeat_mutex_);
heartbeat_condv_.wait_for(
heartbeat_lock, std::chrono::seconds(FLAGS_heartbeat_cooldown_sec), [this] { return !is_running_; });
if (!is_running_)
if (!is_running_) {
break;
}
}
// do heartbeat with cluster manager
HeartBeatToCluster();
Expand All @@ -347,11 +362,14 @@ void KVRpcService::RegisterOnCluster() {
is_registered_ = true;
register_condv_.post();
} else {
MLOG_ERROR("RegisterOnCluster return not OK");
MLOG_ERROR("RegisterOnCluster to CM returns not OK, ret_code:{}", response->ret_code());
}
} else {
MLOG_ERROR(
"Failed to RegisterOnCluster to Cluster Manager {}:{}", FLAGS_cm_primary_node_ip, FLAGS_cm_rpc_inter_port);
MLOG_ERROR("RegisterOnCluster RPC to CM {}:{} failed: {} ({})",
FLAGS_cm_primary_node_ip,
FLAGS_cm_rpc_inter_port,
ctx->ErrorText(),
ctx->ErrorCode());
}
};

Expand Down Expand Up @@ -391,10 +409,10 @@ void KVRpcService::RegisterToRestartedManager() {
cm_ready_ = true;
register_condv_.post();
} else {
MLOG_ERROR("RegisterToRestartedManager return not OK: {}", ret_code);
MLOG_ERROR("RegisterToRestartedManager to CM returns not OK, ret_code:{}", response->ret_code());
}
} else {
MLOG_ERROR("Failed to RegisterToRestartedManager to Cluster Manager {}:{}: {} ({})",
MLOG_ERROR("RegisterToRestartedManager RPC to CM {}:{} failed: {} ({})",
FLAGS_cm_primary_node_ip,
FLAGS_cm_rpc_inter_port,
ctx->ErrorText(),
Expand Down Expand Up @@ -429,9 +447,14 @@ void KVRpcService::HeartBeatToCluster() {
auto response = dynamic_cast<const DataServerHeartBeatResponsePB *>(rsp);
OnHeartbeatResult(false, response->ret_code());
if (response->ret_code() != CommonErr::OK) {
MLOG_ERROR("HeartBeatToCluster return not OK");
MLOG_ERROR("HeartBeatToCluster to CM {}:{} returns not OK, ret_code:{}", response->ret_code());
}
} else {
MLOG_ERROR("HeartBeatToCluster RPC to CM {}:{} failed: {} ({})",
FLAGS_cm_primary_node_ip,
FLAGS_cm_rpc_inter_port,
ctx->ErrorText(),
ctx->ErrorCode());
OnHeartbeatResult(true, CommonErr::InvalidState);
}
};
Expand All @@ -453,14 +476,16 @@ void KVRpcService::OnHeartbeatResult(bool rpc_failed, error_code_t ret_code) {
}

if (!rpc_failed) {
// RPC request succeed, but CM returns error code, do nothing
return;
}

const auto failure_count = heartbeat_failure_count_.fetch_add(1) + 1;
MLOG_ERROR("Failed to HeartbeatToCluster to Cluster Manager {}:{} (counter={})",
MLOG_ERROR("Failed to HeartbeatToCluster to Cluster Manager {}:{} (counter={}), ret_code:{}",
FLAGS_cm_primary_node_ip,
FLAGS_cm_rpc_inter_port,
failure_count);
failure_count,
ret_code);
if (failure_count >= FLAGS_cm_hb_tolerance_count) {
heartbeat_failure_count_.store(0);
if (FLAGS_ds_process_exit_cm_disconnection) {
Expand All @@ -475,6 +500,10 @@ void KVRpcService::OnHeartbeatResult(bool rpc_failed, error_code_t ret_code) {
}

void KVRpcService::HandleClusterManagerDisconnect() {
if (!cluster_disconnect_handler_) {
MLOG_ERROR("Cluster disconnect handler is not set, skip forced shutdown path");
return;
}
cluster_disconnect_handler_();
}

Expand Down
11 changes: 4 additions & 7 deletions src/data_server/kv_rpc_service.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ class KVRpcService {
error_code_t Init();
error_code_t Start();
error_code_t Stop();
void SetClusterDisconnectHandler(std::function<void()> handler);

sicl::rpc::SiRPC *GetIOService() { return io_service_.get(); }
sicl::rpc::SiRPC *GetMgtService() { return mgt_service_.get(); }
Expand Down Expand Up @@ -113,12 +114,7 @@ class KVRpcService {
std::condition_variable heartbeat_condv_;
std::atomic<uint32_t> heartbeat_failure_count_{0};
std::unique_ptr<std::thread> keepalive_thread_{nullptr};
std::function<void()> cluster_disconnect_handler_{[]() {
// raise SIGTERM to trigger handler to do data_server clean destruction
if (std::raise(SIGTERM) != 0) {
std::_Exit(EXIT_FAILURE);
}
}};
std::function<void()> cluster_disconnect_handler_{};

std::string local_ip_;
std::deque<std::atomic<size_t>> shard_used_bytes_;
Expand All @@ -129,7 +125,8 @@ class KVRpcService {
FRIEND_TEST(KVServiceLightTest, TestHeartbeatFailureCountResetOnSuccess);
FRIEND_TEST(KVServiceLightTest, TestClusterManagerDisconnectHandlerInvokedOnToleranceReached);
FRIEND_TEST(KVServiceLightTest, TestHeartbeatFailureToleranceTriggersReconnectWhenExitDisabled);
FRIEND_TEST(KVServiceLightTest, TestClusterManagerDisconnectSignalPathRaisesSigterm);
FRIEND_TEST(KVServiceLightTest, TestHandleClusterManagerDisconnectWithoutHandlerIsNoOp);
FRIEND_TEST(KVServiceLightTest, TestSetClusterDisconnectHandlerOverridesDefaultNoOp);
FRIEND_TEST(KVServiceLightTest, TestShmAllocatorDestructorReleasesSharedMemory);
#endif
};
Expand Down
36 changes: 12 additions & 24 deletions src/data_server/kv_server_main.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,27 +27,8 @@ DECLARE_string(ds_log_file);
DECLARE_LOG_MODULE("data_server");

static std::atomic<bool> quitPorcess{false};
static void (*prevSigIntHandler)(int) = nullptr;
static void (*prevSigTermHandler)(int) = nullptr;
static void (*prevSigSegvHandler)(int) = nullptr;

void signalHandler(int signal) {
// Call previous handler if exists
switch (signal) {
case SIGINT:
if (prevSigIntHandler)
prevSigIntHandler(signal);
break;
case SIGTERM:
if (prevSigTermHandler)
prevSigTermHandler(signal);
break;
case SIGSEGV:
if (prevSigSegvHandler)
prevSigSegvHandler(signal);
break;
}

// Our own handling logic
switch (signal) {
case SIGINT:
Expand All @@ -58,7 +39,9 @@ void signalHandler(int signal) {
break;
case SIGSEGV:
MLOG_CRITICAL("SIGSEGV received. Trigger codedump and exit...");
abort(); // trigger coredump
std::signal(SIGSEGV, SIG_DFL);
std::raise(SIGSEGV);
std::_Exit(EXIT_FAILURE);
break;
default:
MLOG_WARN("Unknown signal {} received. Exiting gracefully...", signal);
Expand Down Expand Up @@ -96,11 +79,16 @@ int main(int argc, char *argv[]) {

// Register signal handlers and save previous ones
MLOG_INFO("Register signal handers...");
prevSigIntHandler = std::signal(SIGINT, signalHandler);
prevSigTermHandler = std::signal(SIGTERM, signalHandler);
prevSigSegvHandler = std::signal(SIGSEGV, signalHandler);
std::signal(SIGINT, signalHandler);
std::signal(SIGTERM, signalHandler);
std::signal(SIGSEGV, signalHandler);

std::unique_ptr<simm::ds::KVRpcService> server = std::make_unique<simm::ds::KVRpcService>();
server->SetClusterDisconnectHandler([server_ptr = server.get()]() {
MLOG_CRITICAL("Cluster manager disconnected repeatedly, request clean data server shutdown");
server_ptr->Stop();
quitPorcess.store(true);
});
error_code_t rc;
rc = server->Init();
if (rc != CommonErr::OK) {
Expand All @@ -124,4 +112,4 @@ int main(int argc, char *argv[]) {
MLOG_INFO("Signal caught, cleanup done, exiting process ...");

return 0;
}
}
32 changes: 16 additions & 16 deletions tests/data_server/test_ds_kv_service.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,14 +33,6 @@ namespace ds {

namespace {

std::atomic<bool> g_sigterm_received{false};

void TestSigTermHandler(int signal) {
if (signal == SIGTERM) {
g_sigterm_received.store(true);
}
}

std::string MakeShmPath(const std::string &shm_name) {
std::ostringstream oss;
oss << "/dev/shm/" << shm_name;
Expand Down Expand Up @@ -266,7 +258,7 @@ TEST_F(KVServiceLightTest, TestClusterManagerDisconnectHandlerInvokedOnTolerance
FLAGS_ds_process_exit_cm_disconnection = true;

bool disconnect_handler_invoked = false;
rpcServicePtr->cluster_disconnect_handler_ = [&]() { disconnect_handler_invoked = true; };
rpcServicePtr->SetClusterDisconnectHandler([&]() { disconnect_handler_invoked = true; });
rpcServicePtr->heartbeat_failure_count_.store(FLAGS_cm_hb_tolerance_count - 1);

rpcServicePtr->OnHeartbeatResult(true, CommonErr::InvalidArgument);
Expand All @@ -281,7 +273,7 @@ TEST_F(KVServiceLightTest, TestHeartbeatFailureToleranceTriggersReconnectWhenExi
FLAGS_ds_process_exit_cm_disconnection = false;

bool disconnect_handler_invoked = false;
rpcServicePtr->cluster_disconnect_handler_ = [&]() { disconnect_handler_invoked = true; };
rpcServicePtr->SetClusterDisconnectHandler([&]() { disconnect_handler_invoked = true; });
rpcServicePtr->heartbeat_failure_count_.store(FLAGS_cm_hb_tolerance_count - 1);
rpcServicePtr->cm_ready_.store(true);

Expand All @@ -292,13 +284,20 @@ TEST_F(KVServiceLightTest, TestHeartbeatFailureToleranceTriggersReconnectWhenExi
EXPECT_EQ(rpcServicePtr->heartbeat_failure_count_.load(), 0);
}

TEST_F(KVServiceLightTest, TestClusterManagerDisconnectSignalPathRaisesSigterm) {
g_sigterm_received.store(false);
auto prev_handler = std::signal(SIGTERM, TestSigTermHandler);
auto restore_handler = folly::makeGuard([&]() { std::signal(SIGTERM, prev_handler); });
TEST_F(KVServiceLightTest, TestHandleClusterManagerDisconnectWithoutHandlerIsNoOp) {
rpcServicePtr->cluster_disconnect_handler_ = nullptr;

rpcServicePtr->cluster_disconnect_handler_();
EXPECT_TRUE(g_sigterm_received.load());
rpcServicePtr->HandleClusterManagerDisconnect();
SUCCEED();
}

TEST_F(KVServiceLightTest, TestSetClusterDisconnectHandlerOverridesDefaultNoOp) {
bool invoked = false;
rpcServicePtr->SetClusterDisconnectHandler([&]() { invoked = true; });

rpcServicePtr->HandleClusterManagerDisconnect();

EXPECT_TRUE(invoked);
}

TEST_F(KVServiceLightTest, TestShmAllocatorDestructorReleasesSharedMemory) {
Expand All @@ -314,6 +313,7 @@ TEST_F(KVServiceLightTest, TestShmAllocatorDestructorReleasesSharedMemory) {

EXPECT_NE(access(shm_path.c_str(), F_OK), 0) << "shared memory should be unlinked after allocator destruction";
}

} // namespace ds
} // namespace simm

Expand Down
Loading