Skip to content

[Fix]: 1) add ds GetHandler check for case that req buf size less than kvmeta size ; 2) optimize stable test #26

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

[Fix]: 1) add ds GetHandler check for case that req buf size less than kvmeta size ; 2) optimize stable test #26

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
4 changes: 4 additions & 0 deletions src/client/clnt_messenger.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ DECLARE_uint32(clnt_thread_pool_size);
DECLARE_uint32(clnt_cm_addr_check_interval_inSecs);
DECLARE_uint32(clnt_syncreq_retry_count);
DECLARE_bool(clnt_syncreq_enable_retry);
DECLARE_bool(clnt_use_k8s);
DECLARE_bool(simm_enable_trace);

DECLARE_LOG_MODULE("simm_client");
Expand Down Expand Up @@ -177,6 +178,9 @@ std::string ClientMessenger::get_cm_address() {
}
#endif
const std::string default_cm_addr = FLAGS_cm_primary_node_ip + ":" + std::to_string(FLAGS_cm_rpc_inter_port);
if (!FLAGS_clnt_use_k8s) {
return default_cm_addr;
}
// get cluster manager pod info from K8S api, get vector of [pod_name, pod_ip]
// TODO: change to get cluster manager info from etcd
auto [ret, cm_ips] =
Expand Down
159 changes: 99 additions & 60 deletions src/data_server/kv_rpc_handler.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,44 +30,71 @@ void KVGetHandler::Work(const std::shared_ptr<sicl::rpc::RpcContext> ctx,
auto ret = service_->KVGet(simm_ctx, req, kv_entry);
if (ret == CommonErr::OK) {
auto [meta, data] = KVCachePool::GetBufferPair(&kv_entry->slab_info);
if (req->buf_len() < meta->value_len) {
MLOG_ERROR("KVGetHandler::Work client buffer too small for key {}, buf_len:{}, value_len:{}",
req->key(),
req->buf_len(),
meta->value_len);
service_->KVGetCallback(kv_entry);
rsp->set_ret_code(CommonErr::InvalidArgument);
simm::common::Metrics::Instance("data_server")
.ObserveRequestDuration("Get",
static_cast<double>(std::chrono::duration_cast<std::chrono::microseconds>(
std::chrono::steady_clock::now() - simm_ctx->GetReqStartTs())
.count()));
simm::common::Metrics::Instance("data_server").IncRequestsTotal("Get");
simm::common::Metrics::Instance("data_server").IncErrorsTotal("Get");
conn->SendResponse(*rsp, ctx, [rsp](std::shared_ptr<sicl::rpc::RpcContext> ctx) {
if (ctx->Failed()) {
MLOG_ERROR("{} response failed: {}({})", rsp->GetTypeName(), ctx->ErrorText(), ctx->ErrorCode());
}
});
return;
}
sicl::transport::RequestParam param{
.mem_desc = static_cast<sicl::transport::MemDesc *>(kv_entry->slab_info.block_addr->descr)};
std::vector<uint32_t> rkeys(req->buf_rkey().begin(), req->buf_rkey().end());
sicl::transport::WriteCallback done =
[this, kv_entry, rsp, conn, ctx, meta, simm_ctx](sicl::transport::Status status) mutable {
error_code_t ret = CommonErr::OK;
if (status.isOk()) {
// return actual value length to client
rsp->set_val_len(meta->value_len);
// record bytes written on successful write
simm::common::Metrics::Instance("data_server").IncWrittenTotal(static_cast<double>(meta->value_len));
} else {
MLOG_ERROR("KVGetHandler::Work connection write failed, err_code:{}, err_msg:{}",
std::to_string(status.errCode()),
status.errMsg());
ret = DsErr::DataRDMATransportFailed;
simm::common::Metrics::Instance("data_server").IncErrorsTotal("Get");
}
service_->KVGetCallback(kv_entry);
rsp->set_ret_code(ret);
simm::common::Metrics::Instance("data_server").ObserveRequestDuration("Get",
static_cast<double>(std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::steady_clock::now() - simm_ctx->GetReqStartTs()).count()));
simm::common::Metrics::Instance("data_server").IncRequestsTotal("Get");
conn->SendResponse(*rsp, ctx, [rsp](std::shared_ptr<sicl::rpc::RpcContext> ctx) {
if (ctx->Failed()) {
MLOG_ERROR("{} response failed: {}({})", rsp->GetTypeName(), ctx->ErrorText(), ctx->ErrorCode());
}
});
};
sicl::transport::WriteCallback done = [this, kv_entry, rsp, conn, ctx, meta, simm_ctx](
sicl::transport::Status status) mutable {
error_code_t ret = CommonErr::OK;
if (status.isOk()) {
// return actual value length to client
rsp->set_val_len(meta->value_len);
// record bytes written on successful write
simm::common::Metrics::Instance("data_server").IncWrittenTotal(static_cast<double>(meta->value_len));
} else {
MLOG_ERROR("KVGetHandler::Work connection write failed, err_code:{}, err_msg:{}",
std::to_string(status.errCode()),
status.errMsg());
ret = DsErr::DataRDMATransportFailed;
simm::common::Metrics::Instance("data_server").IncErrorsTotal("Get");
}
service_->KVGetCallback(kv_entry);
rsp->set_ret_code(ret);
simm::common::Metrics::Instance("data_server")
.ObserveRequestDuration("Get",
static_cast<double>(std::chrono::duration_cast<std::chrono::microseconds>(
std::chrono::steady_clock::now() - simm_ctx->GetReqStartTs())
.count()));
simm::common::Metrics::Instance("data_server").IncRequestsTotal("Get");
conn->SendResponse(*rsp, ctx, [rsp](std::shared_ptr<sicl::rpc::RpcContext> ctx) {
if (ctx->Failed()) {
MLOG_ERROR("{} response failed: {}({})", rsp->GetTypeName(), ctx->ErrorText(), ctx->ErrorCode());
}
});
};
auto res = conn->write(data, meta->value_len, req->buf_addr(), rkeys, done, param);
if (res != sicl::transport::Result::SICL_SUCCESS) {
done(sicl::transport::Status(res)); // synchronized return
}
return;
}
rsp->set_ret_code(ret);
simm::common::Metrics::Instance("data_server").ObserveRequestDuration("Get",
static_cast<double>(std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::steady_clock::now() - simm_ctx->GetReqStartTs()).count()));
simm::common::Metrics::Instance("data_server")
.ObserveRequestDuration("Get",
static_cast<double>(std::chrono::duration_cast<std::chrono::microseconds>(
std::chrono::steady_clock::now() - simm_ctx->GetReqStartTs())
.count()));
simm::common::Metrics::Instance("data_server").IncRequestsTotal("Get");
if (ret != CommonErr::OK) {
simm::common::Metrics::Instance("data_server").IncErrorsTotal("Get");
Expand Down Expand Up @@ -112,41 +139,47 @@ void KVPutHandler::Work(const std::shared_ptr<sicl::rpc::RpcContext> ctx,
sicl::transport::RequestParam param{
.mem_desc = static_cast<sicl::transport::MemDesc *>(kv_entry->slab_info.block_addr->descr)};
std::vector<uint32_t> rkeys(req->buf_rkey().begin(), req->buf_rkey().end());
sicl::transport::ReadCallback done =
[this, shard_id, kv_entry, rsp, conn, ctx, simm_ctx, meta](sicl::transport::Status status) mutable {
error_code_t ret;
if (status.isOk()) {
service_->KVPutSuccessHooks(kv_entry);
ret = CommonErr::OK;
// record bytes read when Put succeeded
simm::common::Metrics::Instance("data_server").IncReadTotal(static_cast<double>(meta->value_len));
} else {
service_->KVPutFailedRewind(shard_id, kv_entry);
MLOG_ERROR("KVPutHandler::Work connection read failed: err_code:{}, err_msg:{}",
std::to_string(status.errCode()),
status.errMsg());
ret = DsErr::DataRDMATransportFailed;
simm::common::Metrics::Instance("data_server").IncErrorsTotal("Put");
}
rsp->set_ret_code(ret);
simm::common::Metrics::Instance("data_server").ObserveRequestDuration("Put",
static_cast<double>(std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::steady_clock::now() - simm_ctx->GetReqStartTs()).count()));
simm::common::Metrics::Instance("data_server").IncRequestsTotal("Put");
conn->SendResponse(*rsp, ctx, [rsp](std::shared_ptr<sicl::rpc::RpcContext> ctx) {
if (ctx->Failed()) {
MLOG_ERROR("{} response failed: {}({})", rsp->GetTypeName(), ctx->ErrorText(), ctx->ErrorCode());
}
});
};
sicl::transport::ReadCallback done = [this, shard_id, kv_entry, rsp, conn, ctx, simm_ctx, meta](
sicl::transport::Status status) mutable {
error_code_t ret;
if (status.isOk()) {
service_->KVPutSuccessHooks(kv_entry);
ret = CommonErr::OK;
// record bytes read when Put succeeded
simm::common::Metrics::Instance("data_server").IncReadTotal(static_cast<double>(meta->value_len));
} else {
service_->KVPutFailedRewind(shard_id, kv_entry);
MLOG_ERROR("KVPutHandler::Work connection read failed: err_code:{}, err_msg:{}",
std::to_string(status.errCode()),
status.errMsg());
ret = DsErr::DataRDMATransportFailed;
simm::common::Metrics::Instance("data_server").IncErrorsTotal("Put");
}
rsp->set_ret_code(ret);
simm::common::Metrics::Instance("data_server")
.ObserveRequestDuration("Put",
static_cast<double>(std::chrono::duration_cast<std::chrono::microseconds>(
std::chrono::steady_clock::now() - simm_ctx->GetReqStartTs())
.count()));
simm::common::Metrics::Instance("data_server").IncRequestsTotal("Put");
conn->SendResponse(*rsp, ctx, [rsp](std::shared_ptr<sicl::rpc::RpcContext> ctx) {
if (ctx->Failed()) {
MLOG_ERROR("{} response failed: {}({})", rsp->GetTypeName(), ctx->ErrorText(), ctx->ErrorCode());
}
});
};
auto res = conn->read(data, meta->value_len, req->buf_addr(), rkeys, done, param);
if (res != sicl::transport::Result::SICL_SUCCESS) {
done(sicl::transport::Status(res)); // synchronized return
}
return;
}
rsp->set_ret_code(ret);
simm::common::Metrics::Instance("data_server").ObserveRequestDuration("Put",
static_cast<double>(std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::steady_clock::now() - simm_ctx->GetReqStartTs()).count()));
simm::common::Metrics::Instance("data_server")
.ObserveRequestDuration("Put",
static_cast<double>(std::chrono::duration_cast<std::chrono::microseconds>(
std::chrono::steady_clock::now() - simm_ctx->GetReqStartTs())
.count()));
simm::common::Metrics::Instance("data_server").IncRequestsTotal("Put");
if (ret != CommonErr::OK) {
simm::common::Metrics::Instance("data_server").IncErrorsTotal("Put");
Expand All @@ -167,8 +200,11 @@ void KVDelHandler::Work(const std::shared_ptr<sicl::rpc::RpcContext> ctx,
auto rsp = std::make_shared<KVDelResponsePB>();
auto ret = service_->KVDel(simm_ctx, req);
rsp->set_ret_code(ret);
simm::common::Metrics::Instance("data_server").ObserveRequestDuration("Delete",
static_cast<double>(std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::steady_clock::now() - simm_ctx->GetReqStartTs()).count()));
simm::common::Metrics::Instance("data_server")
.ObserveRequestDuration("Delete",
static_cast<double>(std::chrono::duration_cast<std::chrono::microseconds>(
std::chrono::steady_clock::now() - simm_ctx->GetReqStartTs())
.count()));
simm::common::Metrics::Instance("data_server").IncRequestsTotal("Delete");
if (ret != CommonErr::OK) {
simm::common::Metrics::Instance("data_server").IncErrorsTotal("Delete");
Expand All @@ -189,8 +225,11 @@ void KVLookupHandler::Work(const std::shared_ptr<sicl::rpc::RpcContext> ctx,
auto rsp = std::make_shared<KVLookupResponsePB>();
auto ret = service_->KVLookup(simm_ctx, req);
rsp->set_ret_code(ret);
simm::common::Metrics::Instance("data_server").ObserveRequestDuration("Lookup",
static_cast<double>(std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::steady_clock::now() - simm_ctx->GetReqStartTs()).count()));
simm::common::Metrics::Instance("data_server")
.ObserveRequestDuration("Lookup",
static_cast<double>(std::chrono::duration_cast<std::chrono::microseconds>(
std::chrono::steady_clock::now() - simm_ctx->GetReqStartTs())
.count()));
simm::common::Metrics::Instance("data_server").IncRequestsTotal("Lookup");
if (ret != CommonErr::OK) {
simm::common::Metrics::Instance("data_server").IncErrorsTotal("Lookup");
Expand Down
3 changes: 2 additions & 1 deletion tests/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,4 +10,5 @@ add_subdirectory(common)
add_subdirectory(client)
add_subdirectory(cluster_manager)
add_subdirectory(data_server)
add_subdirectory(correctness)
add_subdirectory(tools)
add_subdirectory(correctness)
21 changes: 21 additions & 0 deletions tests/client/test_clnt_messenger.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,9 @@

DECLARE_bool(clnt_syncreq_enable_retry);
DECLARE_uint32(clnt_syncreq_retry_count);
DECLARE_bool(clnt_use_k8s);
DECLARE_string(cm_primary_node_ip);
DECLARE_int32(cm_rpc_inter_port);
DECLARE_uint32(shard_total_num);

namespace simm {
Expand Down Expand Up @@ -297,6 +300,8 @@ class ClientMessengerTestPeer {

static std::string CmAddr() { return ClientMessenger::Instance().cm_addr_; }

static std::string GetCmAddress() { return ClientMessenger::Instance().get_cm_address(); }

static sicl::rpc::SiRPC *SwapRpcClient(sicl::rpc::SiRPC *replacement) {
auto &messenger = ClientMessenger::Instance();
auto *original = messenger.rpc_client_;
Expand Down Expand Up @@ -598,6 +603,22 @@ TEST_F(ClientMessengerUnitTest, ReInitRemovesStaleConnectionFromLocalContextMap)
EXPECT_EQ(ClientMessengerTestPeer::ShardOwner(1), "10.0.0.1:1001");
}

TEST_F(ClientMessengerUnitTest, GetCmAddressUsesFlagToSkipK8SLookup) {
const auto old_use_k8s = FLAGS_clnt_use_k8s;
const auto old_cm_ip = FLAGS_cm_primary_node_ip;
const auto old_cm_port = FLAGS_cm_rpc_inter_port;

FLAGS_clnt_use_k8s = false;
FLAGS_cm_primary_node_ip = "10.8.0.1";
FLAGS_cm_rpc_inter_port = 30001;

EXPECT_EQ(ClientMessengerTestPeer::GetCmAddress(), "10.8.0.1:30001");

FLAGS_clnt_use_k8s = old_use_k8s;
FLAGS_cm_primary_node_ip = old_cm_ip;
FLAGS_cm_rpc_inter_port = old_cm_port;
}

} // namespace clnt
} // namespace simm

Expand Down
114 changes: 114 additions & 0 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 @@ -240,6 +240,120 @@ TEST_F(KVServiceTest, TestClientHandlers) {
clientPool->join();
}

TEST_F(KVServiceTest, TestGetRejectsTooSmallClientBuffer) {
auto serverPool = std::make_unique<folly::IOThreadPoolExecutor>(1);
auto clientPool = std::make_unique<folly::IOThreadPoolExecutor>(1);
folly::Baton<> serverReady, serverExit;

folly::via(serverPool->getEventBase(), [&] {
auto ret = rpcServicePtr->Start();
EXPECT_EQ(ret, CommonErr::OK);
serverReady.post();
serverExit.wait();
});

serverReady.wait();

folly::Baton<> clientDone;
folly::via(clientPool->getEventBase(), [&] {
sicl::rpc::SiRPC *sirpc = nullptr;
sicl::rpc::SiRPC::newInstance(sirpc, false);
sicl::rpc::RpcContext *ctx_raw = nullptr;
sicl::rpc::RpcContext::newInstance(ctx_raw);
auto ctx = std::shared_ptr<sicl::rpc::RpcContext>(ctx_raw);
ctx->set_timeout(sicl::transport::TimerTick::TIMER_1S);

auto conn = sirpc->connect("127.0.0.1", FLAGS_io_service_port);
ASSERT_TRUE(conn != nullptr);

const std::string key = "test_get_too_small_buffer";
const uint32_t shard_id =
hashkit::HashkitBase::Instance().generate_16bit_hash_value(key.c_str(), key.length()) % FLAGS_shard_total_num;
constexpr size_t kValueLen = 4096;
constexpr size_t kSmallBufLen = 512;

auto *mempool = sirpc->GetMempool();
sicl::transport::MemDesc *put_descr = nullptr;
ASSERT_EQ(mempool->alloc(put_descr, kValueLen), sicl::transport::Result::SICL_SUCCESS);
std::string value;
get_random_string(kValueLen, &value);
memcpy(put_descr->getAddr(), value.data(), value.size());

auto put_req = std::make_shared<KVPutRequestPB>();
auto put_rsp = std::make_shared<KVPutResponsePB>();
put_req->set_shard_id(shard_id);
put_req->set_key(key);
put_req->set_val_len(kValueLen);
put_req->set_buf_addr(reinterpret_cast<uint64_t>(put_descr->getAddr()));
put_req->set_buf_ofs(0);
put_req->set_buf_len(kValueLen);
for (auto rkey : put_descr->getRemoteKeys()) {
put_req->add_buf_rkey(rkey);
}

std::atomic<bool> put_done{false};
sirpc->SendRequest(
conn,
static_cast<sicl::rpc::ReqType>(KVServerRpcType::RPC_CLIENT_KV_PUT),
*put_req,
put_rsp.get(),
ctx,
[&put_done](const google::protobuf::Message *rsp, const std::shared_ptr<sicl::rpc::RpcContext> ctx) {
ASSERT_FALSE(ctx->Failed());
auto *put_rsp = dynamic_cast<const KVPutResponsePB *>(rsp);
ASSERT_NE(put_rsp, nullptr);
EXPECT_EQ(put_rsp->ret_code(), CommonErr::OK);
put_done.store(true);
});
while (!put_done.load()) {
std::this_thread::sleep_for(std::chrono::milliseconds(10));
}

sicl::transport::MemDesc *get_descr = nullptr;
ASSERT_EQ(mempool->alloc(get_descr, kSmallBufLen), sicl::transport::Result::SICL_SUCCESS);
memset(get_descr->getAddr(), 0, kSmallBufLen);

auto get_req = std::make_shared<KVGetRequestPB>();
auto get_rsp = std::make_shared<KVGetResponsePB>();
get_req->set_shard_id(shard_id);
get_req->set_key(key);
get_req->set_buf_addr(reinterpret_cast<uint64_t>(get_descr->getAddr()));
get_req->set_buf_ofs(0);
get_req->set_buf_len(kSmallBufLen);
for (auto rkey : get_descr->getRemoteKeys()) {
get_req->add_buf_rkey(rkey);
}

std::atomic<bool> get_done{false};
sirpc->SendRequest(
conn,
static_cast<sicl::rpc::ReqType>(KVServerRpcType::RPC_CLIENT_KV_GET),
*get_req,
get_rsp.get(),
ctx,
[&get_done](const google::protobuf::Message *rsp, const std::shared_ptr<sicl::rpc::RpcContext> ctx) {
ASSERT_FALSE(ctx->Failed());
auto *get_rsp = dynamic_cast<const KVGetResponsePB *>(rsp);
ASSERT_NE(get_rsp, nullptr);
EXPECT_EQ(get_rsp->ret_code(), CommonErr::InvalidArgument);
get_done.store(true);
});
while (!get_done.load()) {
std::this_thread::sleep_for(std::chrono::milliseconds(10));
}

mempool->release(put_descr);
mempool->release(get_descr);
delete sirpc;
clientDone.post();
});

clientDone.wait();
serverExit.post();
serverPool->join();
clientPool->join();
}

TEST_F(KVServiceLightTest, TestHeartbeatFailureCountResetOnSuccess) {
rpcServicePtr->heartbeat_failure_count_.store(4);
rpcServicePtr->cm_ready_.store(true);
Expand Down
Loading
Loading