[TENT] fix: Reduce unnecessary bandwidth consumption in TCP recvData operations#1513
Conversation
Summary of ChangesHello @00fish0, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request primarily focuses on optimizing the TCP transport within the TENT backend to reduce unnecessary bandwidth consumption during Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
00fish0
changed the title
00fish0
changed the title
There was a problem hiding this comment.
Pull request overview
This PR optimizes the TENT backend’s control-plane data path by shrinking recvData RPC requests (removing redundant preallocated payload) and tightening error signaling, along with benchmark transport-selection fixes and a TENT stub for getBaseAddr().
Changes:
- Reduce
recvDataRPC request size to only the descriptor header and adjust server/client behavior accordingly. - Add error propagation for
sendData/recvDatawhen the target address is not in a registered buffer. - Update TENT benchmark config to honor
--xport_typeby disabling/enabling transports; addTransferEngine::getBaseAddr()stub for TENT.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| mooncake-transfer-engine/tent/src/runtime/control_plane.cpp | Shrinks recvData request payload; adds/adjusts RPC error handling for Send/Recv. |
| mooncake-transfer-engine/src/transfer_engine.cpp | Adds getBaseAddr() behavior for TENT mode (returns nullptr). |
| mooncake-transfer-engine/benchmark/tent_backend.cpp | Changes --xport_type handling to disable all transports and enable only the chosen one; minor loop formatting tweak. |
Comments suppressed due to low confidence (1)
mooncake-transfer-engine/tent/src/runtime/control_plane.cpp:257
onRecvData()leavesresponseempty when the address range isn't registered, so the client can only infer failure via a size mismatch. For consistency withonSendData()and better diagnosability, set an explicit error message in theelsepath (and have the client surface it when present).
if (local_desc->findBuffer(peer_mem_addr, length)) {
response.resize(length);
Platform::getLoader().copy(response.data(), (void*)peer_mem_addr,
length);
}
}
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if (response.size() != length) | ||
| return Status::RpcServiceError( | ||
| "RecvData failed: target address not in registered buffer"); |
There was a problem hiding this comment.
recvData() treats any response size mismatch as "target address not in registered buffer" and discards the actual server response. This can misreport other failure modes (e.g., server-side error string, partial response) and makes debugging harder. Consider returning RpcServiceError(response) when !response.empty() and otherwise include expected vs actual sizes in the error message.
| if (response.size() != length) | |
| return Status::RpcServiceError( | |
| "RecvData failed: target address not in registered buffer"); | |
| if (response.size() != length) { | |
| if (!response.empty()) { | |
| // Server likely returned an explicit error message instead of raw data. | |
| return Status::RpcServiceError(response); | |
| } | |
| return Status::RpcServiceError( | |
| "RecvData failed: response size mismatch (expected " + | |
| std::to_string(length) + " bytes, got " + | |
| std::to_string(response.size()) + " bytes)"); | |
| } |
| // Enable only the specified transport | ||
| if (XferBenchConfig::xport_type == "rdma") { | ||
| config->set("transports/rdma/enable", true); | ||
| } else if (XferBenchConfig::xport_type == "tcp") { | ||
| config->set("transports/tcp/enable", true); | ||
| } else if (XferBenchConfig::xport_type == "shm") { | ||
| config->set("transports/shm/enable", true); | ||
| } else if (XferBenchConfig::xport_type == "iouring") { | ||
| config->set("transports/io_uring/enable", true); | ||
| } else if (XferBenchConfig::xport_type == "gds") { | ||
| config->set("transports/gds/enable", true); | ||
| } else if (XferBenchConfig::xport_type == "mnnvl") { | ||
| config->set("transports/mnnvl/enable", true); | ||
| } |
There was a problem hiding this comment.
If XferBenchConfig::xport_type is set to an unexpected value, this code disables all transports and enables none, which will make the benchmark fail in a non-obvious way. Add an else branch that logs/returns an error (or falls back to defaults) when the value is not one of the supported strings.
| // Disable all transports by default | ||
| config->set("transports/rdma/enable", false); | ||
| config->set("transports/tcp/enable", false); | ||
| config->set("transports/shm/enable", false); | ||
| config->set("transports/io_uring/enable", false); | ||
| config->set("transports/gds/enable", false); | ||
| config->set("transports/mnnvl/enable", false); |
There was a problem hiding this comment.
When xport_type is specified, this block disables only a subset of transports. Others that are enabled-by-default in the TENT transport loader (e.g. transports/nvlink/enable and transports/ascend_direct/enable) can still be loaded, so it may not actually "enable only the specified transport". Consider explicitly disabling all known transport enable keys before turning on the selected one.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces an important optimization to the recvData operation by reducing RPC request size, which significantly cuts down bandwidth consumption for READ operations. It also includes a fix for the --xport_type parameter handling and adds a getBaseAddr() stub for the TENT backend. However, critical vulnerabilities have been identified in the RPC handlers onSendData and onRecvData. These functions lack sufficient validation of user-supplied memory addresses and lengths, potentially leading to arbitrary memory read/write and Denial of Service (DoS) vulnerabilities due to integer overflows in memory range checks and missing buffer size validations. Robust input validation and overflow checks are highly recommended before merging these changes.
| if (XferBenchConfig::xport_type == "rdma") { | ||
| config->set("transports/rdma/enable", true); | ||
| } else if (XferBenchConfig::xport_type == "tcp") { | ||
| config->set("transports/tcp/enable", true); | ||
| } else if (XferBenchConfig::xport_type == "shm") { | ||
| config->set("transports/shm/enable", true); | ||
| } else if (XferBenchConfig::xport_type == "iouring") { | ||
| config->set("transports/io_uring/enable", true); | ||
| } else if (XferBenchConfig::xport_type == "gds") { | ||
| config->set("transports/gds/enable", true); | ||
| } else if (XferBenchConfig::xport_type == "mnnvl") { | ||
| config->set("transports/mnnvl/enable", true); |
There was a problem hiding this comment.
The series of if-else if statements for enabling transport types can become cumbersome as more transport types are added. Consider refactoring this into a std::map<std::string, bool> where keys are transport names and values indicate enablement, or using a switch statement if XferBenchConfig::xport_type can be mapped to an enum. This would improve maintainability and scalability of the configuration logic.
std::unordered_map<std::string, bool> transport_enable_map;
transport_enable_map["rdma"] = false;
transport_enable_map["tcp"] = false;
transport_enable_map["shm"] = false;
transport_enable_map["iouring"] = false;
transport_enable_map["gds"] = false;
transport_enable_map["mnnvl"] = false;
if (transport_enable_map.count(XferBenchConfig::xport_type)) {
transport_enable_map[XferBenchConfig::xport_type] = true;
}
for (const auto& entry : transport_enable_map) {
config->set("transports/" + entry.first + "/enable", entry.second);
}
There was a problem hiding this comment.
Following this comment is better if extending transports.
- Add request size validation in onSendData to prevent buffer over-read - Add length upper bound check in onRecvData to prevent DoS via memory exhaustion - Add integer overflow checks in findBuffer for memory range validation Fixes security vulnerabilities identified by Gemini Code Assist review
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
3 participants
Description
Optimizes TCP transport in TENT backend by reducing unnecessary bandwidth consumption in recvData operations.
--xport_typeparameter handling in TENT backend. Now properly disables unwanted transports and enables only the specified one (rdma/tcp/shm/iouring/gds/mnnvl).Type of Change
How Has This Been Tested?
Before:
After:
Checklist
./scripts/code_format.shbefore submitting.