Skip to content

[Feat]: add dataserver resource query and summary mechanism in CM #28

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 10, 2026
Merged

[Feat]: add dataserver resource query and summary mechanism in CM #28

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
15 changes: 13 additions & 2 deletions src/cluster_manager/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,11 @@ list(REMOVE_ITEM MODULE_SRC_LIB
"${CMAKE_CURRENT_SOURCE_DIR}/cm_main.cc"
)

add_executable(${MODULE_NAME} ${MODULE_SRC} ${MODULE_HEADERS})
set(DS_SHARED_FLAG_SRC
"${PROJECT_SOURCE_DIR}/src/data_server/ds_flags.cc"
)

add_executable(${MODULE_NAME} ${MODULE_SRC} ${MODULE_HEADERS} ${DS_SHARED_FLAG_SRC})
target_link_libraries(${MODULE_NAME} PRIVATE
gflags
folly
Expand All @@ -32,7 +36,7 @@ set_target_properties(${MODULE_NAME} PROPERTIES
INSTALL_RPATH "\$ORIGIN/../../../third_party/sict/lib"
)

add_library(${MODULE_NAME}_static STATIC ${MODULE_SRC_LIB})
add_library(${MODULE_NAME}_static STATIC ${MODULE_SRC_LIB} ${DS_SHARED_FLAG_SRC})
# PRIVATE means that the libraries are only used by this target and
# not propagated to targets that link against this one.
#
Expand All @@ -46,3 +50,10 @@ target_link_libraries(${MODULE_NAME}_static PRIVATE
sict
simm_common
)

if(ENABLE_TESTS)
target_compile_definitions(${MODULE_NAME}_static PRIVATE SIMM_UNIT_TEST)
target_include_directories(${MODULE_NAME}_static PRIVATE
${PROJECT_SOURCE_DIR}/third_party/gtest/googletest/include
)
endif()
92 changes: 77 additions & 15 deletions src/cluster_manager/cm_node_manager.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
#include <algorithm>
#include <chrono>
#include <functional>
#include <optional>
Expand All @@ -12,13 +13,15 @@
#include "rpc/connection.h"
#include "rpc/rpc.h"

#include "cm_rpc_handler.h"
#include "common/logging/logging.h"
#include "common/utils/time_util.h"
#include "cm_node_manager.h"

DECLARE_string(dataserver_namespace);
DECLARE_string(dataserver_svc_name);
DECLARE_string(dataserver_port_name);
DECLARE_int32(mgt_service_port);
DECLARE_uint32(rpc_timeout_inSecs);
DECLARE_uint32(dataserver_resource_interval_inSecs);
DECLARE_bool(cm_deferred_reshard_enabled);
Expand All @@ -37,13 +40,17 @@ ClusterManagerNodeManager::ClusterManagerNodeManager() {

ClusterManagerNodeManager::~ClusterManagerNodeManager() {
MLOG_INFO("Start destruct Node Manager");
Stop();
if (rpc_client_ != nullptr) {
delete rpc_client_;
rpc_client_ = nullptr;
}
}

void ClusterManagerNodeManager::Init() {
if (resource_thread_ != nullptr) {
return;
}
start_timestamp_us_ = simm::utils::current_microseconds();
resource_thread_stop_.store(false);

Expand All @@ -52,7 +59,8 @@ void ClusterManagerNodeManager::Init() {
std::function<void()> resource_loop = [self]() {
while (!self->resource_thread_stop_.load()) {
self->updateAllNodeResource();
self->resource_thread_baton_.timed_wait(std::chrono::milliseconds(FLAGS_dataserver_resource_interval_inSecs * 1000));
self->resource_thread_baton_.timed_wait(
std::chrono::milliseconds(FLAGS_dataserver_resource_interval_inSecs * 1000));
self->resource_thread_baton_.reset();
}
};
Expand All @@ -67,6 +75,12 @@ void ClusterManagerNodeManager::Stop() {
if (resource_thread_ != nullptr && resource_thread_->joinable()) {
resource_thread_->join();
}
{
std::lock_guard<std::mutex> lock(resource_conn_mutex_);
resource_conn_map_.clear();
}
delete resource_thread_;
resource_thread_ = nullptr;
MLOG_INFO("Delete resource query thread in Node Manager succeed");
}

Expand All @@ -88,11 +102,7 @@ std::vector<std::shared_ptr<simm::common::NodeAddress>> ClusterManagerNodeManage
}

error_code_t ClusterManagerNodeManager::AddNode(const std::string & addr_str) {
auto result = node_status_map_.insert_or_assign(addr_str, NodeStatus::RUNNING);
if (!result.second) {
MLOG_ERROR("Add node {} in Node Manager status map failed", addr_str);
return CmErr::NodeManagerAddNodeFailed;
}
node_status_map_.insert_or_assign(addr_str, NodeStatus::RUNNING);

// FIXME(ytji): for test, just comment below codes
// // get node resource info, only print log if error
Expand All @@ -110,6 +120,10 @@ error_code_t ClusterManagerNodeManager::AddNode(const std::string & addr_str) {
error_code_t ClusterManagerNodeManager::DelNode(const std::string & addr_str) {
node_info_map_.erase(addr_str);
node_status_map_.erase(addr_str);
{
std::lock_guard<std::mutex> lock(resource_conn_mutex_);
resource_conn_map_.erase(addr_str);
}
MLOG_DEBUG("Delete node {} in Node Manager succeed", addr_str);
return CommonErr::OK;
}
Expand All @@ -122,8 +136,13 @@ error_code_t ClusterManagerNodeManager::UpdateNodeStatus(const std::string & add

std::shared_ptr<simm::common::NodeResource> ClusterManagerNodeManager::getNodeResource(
const std::string & addr_str) {
#if defined(SIMM_UNIT_TEST)
if (test_resource_query_hook_) {
return test_resource_query_hook_(addr_str);
}
#endif
DataServerResourceRequestPB req;
auto resp = new DataServerResourceResponsePB;
auto resp = std::make_unique<DataServerResourceResponsePB>();
sicl::rpc::RpcContext *ctx_p;
sicl::rpc::RpcContext::newInstance(ctx_p);
std::shared_ptr<sicl::rpc::RpcContext> ctx = std::shared_ptr<sicl::rpc::RpcContext>(ctx_p);
Expand All @@ -134,16 +153,54 @@ std::shared_ptr<simm::common::NodeResource> ClusterManagerNodeManager::getNodeRe
return nullptr;
}

rpc_client_->SendRequest(ds_addr->node_ip_, ds_addr->node_port_, static_cast<sicl::rpc::ReqType>(0), req, resp, ctx, nullptr);
std::shared_ptr<sicl::rpc::Connection> conn;
{
std::lock_guard<std::mutex> lock(resource_conn_mutex_);
auto it = resource_conn_map_.find(addr_str);
if (it != resource_conn_map_.end()) {
conn = it->second;
}
}
if (!conn) {
conn = rpc_client_->connect(ds_addr->node_ip_, FLAGS_mgt_service_port);
if (!conn) {
MLOG_ERROR("Get node {} resource failed, connect failed to management port {}", addr_str, FLAGS_mgt_service_port);
return nullptr;
}
std::lock_guard<std::mutex> lock(resource_conn_mutex_);
resource_conn_map_[addr_str] = conn;
}

ctx->set_timeout(sicl::transport::TimerTick::TIMER_3S);
rpc_client_->SendRequest(conn,
static_cast<sicl::rpc::ReqType>(cm::ClusterManagerRpcType::RPC_DATASERVER_RESOURCE_QUERY),
req,
resp.get(),
ctx);
if (ctx->Failed()) {
std::string errmsg = ctx->ErrorText();
std::lock_guard<std::mutex> lock(resource_conn_mutex_);
resource_conn_map_.erase(addr_str);
MLOG_ERROR("Get node {} resource failed, err:{}", addr_str, errmsg);
return nullptr;
}

std::vector<simm::common::NodeResource::ShardMemResource> shard_mem_infos;
shard_mem_infos.reserve(resp->shard_mem_infos_size());
for (const auto &shard_pb : resp->shard_mem_infos()) {
shard_mem_infos.push_back(
{static_cast<shard_id_t>(shard_pb.shard_id()), static_cast<int64_t>(shard_pb.shard_mem_used_bytes())});
}
std::sort(shard_mem_infos.begin(), shard_mem_infos.end(), [](const auto &lhs, const auto &rhs) {
return lhs.shard_id_ < rhs.shard_id_;
});

auto resource = std::make_shared<simm::common::NodeResource>(static_cast<int64_t>(resp->mem_total_bytes()),
static_cast<int64_t>(resp->mem_allocated_bytes()),
static_cast<int64_t>(resp->mem_used_bytes()),
static_cast<int64_t>(resp->mem_free_bytes()),
static_cast<int64_t>(resp->mem_used_bytes()));
resp->last_report_timestamp_us(),
std::move(shard_mem_infos));

MLOG_DEBUG("Get node {} resource info in Node Manager succeed", addr_str);
return resource;
Expand Down Expand Up @@ -191,6 +248,15 @@ std::shared_ptr<simm::common::NodeResource> ClusterManagerNodeManager::GetNodeRe
return it->second;
}

std::shared_ptr<simm::common::NodeResource> ClusterManagerNodeManager::RefreshNodeResource(
const std::string &addr_str) {
auto resource_ret = getNodeResource(addr_str);
if (resource_ret != nullptr) {
node_info_map_.insert_or_assign(addr_str, resource_ret);
}
return resource_ret;
}

std::unordered_map<std::string, NodeStatus> ClusterManagerNodeManager::GetAllNodeStatus() {
std::unordered_map<std::string, NodeStatus> status_map;
for (auto &pair : node_status_map_) {
Expand Down Expand Up @@ -225,7 +291,6 @@ HandshakeResult ClusterManagerNodeManager::ProcessHandshake(
const std::string& logical_id,
const std::string& new_ip_port,
const std::vector<shard_id_t>& reported_shards) {

HandshakeResult result;

auto it = logical_node_table_.find(logical_id);
Expand Down Expand Up @@ -259,8 +324,7 @@ HandshakeResult ClusterManagerNodeManager::ProcessHandshake(
entry.status = NodeStatus::RUNNING;
entry.deferred_reshard_since = {};
logical_node_table_.assign(logical_id, entry);
MLOG_INFO("Node replacement: logical_id={} old_ip={} new_ip={}",
logical_id, result.old_ip_port, new_ip_port);
MLOG_INFO("Node replacement: logical_id={} old_ip={} new_ip={}", logical_id, result.old_ip_port, new_ip_port);
return result;
}

Expand All @@ -276,8 +340,7 @@ HandshakeResult ClusterManagerNodeManager::ProcessHandshake(
// Case 4: IP changed while still RUNNING (rare IP drift)
migrateNodeIp(logical_id, entry, new_ip_port);
logical_node_table_.assign(logical_id, entry);
MLOG_INFO("Node IP update: logical_id={} old_ip={} new_ip={}",
logical_id, result.old_ip_port, new_ip_port);
MLOG_INFO("Node IP update: logical_id={} old_ip={} new_ip={}", logical_id, result.old_ip_port, new_ip_port);
}
return result;
}
Expand Down Expand Up @@ -323,7 +386,6 @@ error_code_t ClusterManagerNodeManager::SetNodeStatus(
NodeStatus status,
std::chrono::steady_clock::time_point ts,
std::optional<NodeStatus> expected_status) {

auto it = logical_node_table_.find(logical_id);
if (it == logical_node_table_.end()) {
MLOG_WARN("SetNodeStatus: logical_id={} not found", logical_id);
Expand Down
11 changes: 11 additions & 0 deletions src/cluster_manager/cm_node_manager.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
#include <optional>
#include <string>
#include <thread>
#include <mutex>
#include <unordered_map>
#include <vector>

Expand Down Expand Up @@ -88,6 +89,9 @@ class ClusterManagerNodeManager : public std::enable_shared_from_this<ClusterMan
// get resource info of a single data node
std::shared_ptr<simm::common::NodeResource> GetNodeResource(const std::string &addr_str);

// actively query one data node resource info and refresh cache
std::shared_ptr<simm::common::NodeResource> RefreshNodeResource(const std::string &addr_str);

// get status of all data nodes
// Returns an unordered_map of address string to node status
std::unordered_map<std::string, NodeStatus> GetAllNodeStatus();
Expand Down Expand Up @@ -149,13 +153,20 @@ class ClusterManagerNodeManager : public std::enable_shared_from_this<ClusterMan
folly::ConcurrentHashMap<std::string, std::string> addr_to_logical_;

sicl::rpc::SiRPC *rpc_client_{nullptr};
std::mutex resource_conn_mutex_;
std::unordered_map<std::string, std::shared_ptr<sicl::rpc::Connection>> resource_conn_map_;
std::thread *resource_thread_{nullptr};
std::atomic<bool> resource_thread_stop_{false};
folly::Baton<> resource_thread_baton_;
uint64_t start_timestamp_us_{0};

#if defined(SIMM_UNIT_TEST)
std::function<std::shared_ptr<simm::common::NodeResource>(const std::string &)> test_resource_query_hook_;
#endif

// only for UT test
#if defined(SIMM_UNIT_TEST)
friend class ClusterManagerNodeManagerTestPeer;
FRIEND_TEST(ClusterManagerHBMonitorTest, TestHBMonitor);
#endif
};
Expand Down
Loading
Loading