Skip to content

[TENT] Fix TaskInfo dangling pointers, data races, stage buffer leaks #3380

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

Open
RuiqingFeng wants to merge 2 commits into
base: main
Choose a base branch
from
+293 −14
Open

[TENT] Fix TaskInfo dangling pointers, data races, stage buffer leaks #3380

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
dc935a7
[TENT] Add unit tests for staging buffers and batch capacity limits
RuiqingFeng Aug 11, 2026
f9b524d
[TENT] Fix staging runtime bugs in ProxyManager and TransferEngineImpl
RuiqingFeng Aug 11, 2026
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: 3 additions & 1 deletion mooncake-transfer-engine/tent/include/tent/runtime/proxy_manager.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,8 @@ class ProxyManager {
private:
void runner(size_t id);

Status transferEventLoop(StagingTask& task, StageBufferCache* cache);
Status transferEventLoop(StagingTask& task, StageBufferCache* cache,
bool& buffers_safe_to_release);

Status transferSync(StagingTask& task, StageBufferCache* cache);

Expand Down Expand Up @@ -101,6 +102,7 @@ class ProxyManager {
const size_t chunk_count_;
TransferEngineImpl* impl_;
std::unordered_map<std::string, StageBuffers> stage_buffers_;
std::recursive_mutex stage_buffers_mu_;
std::atomic<bool> running_;
struct WorkerShard {
std::thread thread;
Expand Down
80 changes: 67 additions & 13 deletions mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,8 @@ Status ProxyManager::deconstruct() {
shards_[i].cv.notify_all();
shards_[i].thread.join();
}
for (auto entry : stage_buffers_) {
std::lock_guard<std::recursive_mutex> lock(stage_buffers_mu_);
for (auto& entry : stage_buffers_) {
impl_->unregisterLocalMemory(entry.second.chunks);
impl_->freeLocalMemory(entry.second.chunks);
delete[] entry.second.bitmap;
Expand Down Expand Up @@ -155,7 +156,9 @@ Status ProxyManager::submit(TaskInfo* task, BatchID batch,

Status ProxyManager::getStatus(TaskInfo* task, TransferStatus& task_status) {
if (!task || !task->staging) return Status::InvalidArgument("Invalid task");
task_status.s = task->staging_status;
TransferStatusEnum staging_status;
__atomic_load(&task->staging_status, &staging_status, __ATOMIC_ACQUIRE);

@alogfans alogfans Aug 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Recommend to use std::atomic for task->staging_status

task_status.s = staging_status;
if (task_status.s == COMPLETED) {
task_status.transferred_bytes = task->request.length;
}
Expand All @@ -173,7 +176,7 @@ struct StageBufferCache {
uint64_t addr = 0;
auto status = mgr.pinStageBuffer(location, addr);
if (!status.ok()) {
LOG(FATAL) << "Failed to pin local stage buffer: " << status
LOG(ERROR) << "Failed to pin local stage buffer: " << status
<< ", location " << location;
return 0;
}
Expand All @@ -191,7 +194,7 @@ struct StageBufferCache {
auto status =
ControlClient::pinStageBuffer(server_addr, location, addr);
if (!status.ok()) {
LOG(FATAL) << "Failed to pin remote stage buffer: " << status
LOG(ERROR) << "Failed to pin remote stage buffer: " << status
<< ", location " << location;
return 0;
}
Expand Down Expand Up @@ -220,7 +223,6 @@ struct StageBufferCache {
};

void ProxyManager::runner(size_t id) {
StageBufferCache cache(*this);
auto& shard = shards_[id];
while (running_) {
StagingTask task;
Expand All @@ -239,17 +241,26 @@ void ProxyManager::runner(size_t id) {
}

if (!task.native) continue;
auto status = transferEventLoop(task, &cache);
StageBufferCache cache(*this);
bool buffers_safe_to_release;
auto status = transferEventLoop(task, &cache, buffers_safe_to_release);
if (buffers_safe_to_release) {
cache.reset();
} else {
LOG(ERROR) << "Staging cleanup could not drain all in-flight "
"operations; keeping stage buffers pinned";
}
Comment on lines +250 to +252

@alogfans alogfans Aug 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to put cache in a deferred-cleanup collection?

auto staging_status = status.ok() ? COMPLETED : FAILED;
__atomic_store(&task.native->staging_status, &staging_status,
__ATOMIC_RELEASE);
impl_->notifyBatchMaybeReady(task.batch);
}
cache.reset();
}

Status ProxyManager::transferEventLoop(StagingTask& task,
StageBufferCache* cache) {
StageBufferCache* cache,
bool& buffers_safe_to_release) {
buffers_safe_to_release = true;
auto& request = task.native->request;
auto server_addr = task.params[0];
bool local_staging = !task.params[1].empty();
Expand Down Expand Up @@ -318,6 +329,44 @@ Status ProxyManager::transferEventLoop(StagingTask& task,

for (size_t i = 0; i < chunks.size(); ++i) event_queue.push(i);
std::vector<std::future<Status>> remote_futures(chunks.size());
auto drain_batch = [&](BatchID batch) {
while (true) {
TransferStatus xfer_status;
auto status = impl_->progressBatch(batch, xfer_status);
if (!status.ok()) {
LOG(ERROR) << "Failed to poll in-flight staging batch: "
<< status;
return false;
}
if (xfer_status.s == PENDING) continue;
auto free_status = impl_->freeBatch(batch);
if (!free_status.ok()) {
LOG(WARNING)
<< "Failed to free drained staging batch: " << free_status;
}
return true;
}
};
auto cleanup_inflight = [&]() {
Comment on lines +332 to +350

@alogfans alogfans Aug 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could this function run indefinitely? This may block the dedicated workers.

bool all_drained = true;
for (auto& chunk : chunks) {
if (!chunk.batch) continue;
if (drain_batch(chunk.batch)) {
chunk.batch = 0;
} else {
all_drained = false;
}
}
for (auto& future : remote_futures) {
if (!future.valid()) continue;
auto status = future.get();
if (!status.ok()) {
LOG(WARNING)
<< "Failed to drain remote staging request: " << status;
}
}
return all_drained;
};

while (!event_queue.empty()) {
auto id = event_queue.front();
Expand Down Expand Up @@ -416,7 +465,11 @@ Status ProxyManager::transferEventLoop(StagingTask& task,

case StageState::INFLIGHT: {
TransferStatus xfer_status;
CHECK_STATUS(impl_->progressBatch(chunk.batch, xfer_status));
auto status = impl_->progressBatch(chunk.batch, xfer_status);
if (!status.ok()) {
buffers_safe_to_release = cleanup_inflight();
return status;
}
if (xfer_status.s == PENDING) {
event_queue.push(id);
break;
Expand Down Expand Up @@ -453,10 +506,7 @@ Status ProxyManager::transferEventLoop(StagingTask& task,
}

case StageState::FAILED: {
// Drain the queue to avoid losing chunks
while (!event_queue.empty()) {
event_queue.pop();
}
buffers_safe_to_release = cleanup_inflight();
return Status::InternalError(
"Proxy event loop in failed state");
}
Expand Down Expand Up @@ -561,6 +611,7 @@ Status ProxyManager::transferSync(StagingTask& task, StageBufferCache* cache) {
}

Status ProxyManager::allocateStageBuffers(const std::string& location) {
std::lock_guard<std::recursive_mutex> lock(stage_buffers_mu_);
if (stage_buffers_.count(location)) return Status::OK();
StageBuffers buf;
auto total_size = chunk_size_ * chunk_count_;
Expand All @@ -575,6 +626,7 @@ Status ProxyManager::allocateStageBuffers(const std::string& location) {
}

Status ProxyManager::freeStageBuffers(const std::string& location) {
std::lock_guard<std::recursive_mutex> lock(stage_buffers_mu_);
auto it = stage_buffers_.find(location);
if (it == stage_buffers_.end())
return Status::InvalidArgument("Stage buffer not allocated" LOC_MARK);
Expand All @@ -587,6 +639,7 @@ Status ProxyManager::freeStageBuffers(const std::string& location) {

Status ProxyManager::pinStageBuffer(const std::string& location,
uint64_t& addr) {
std::lock_guard<std::recursive_mutex> lock(stage_buffers_mu_);
auto it = stage_buffers_.find(location);
if (it == stage_buffers_.end()) {
CHECK_STATUS(allocateStageBuffers(location));
Expand All @@ -605,6 +658,7 @@ Status ProxyManager::pinStageBuffer(const std::string& location,
}

Status ProxyManager::unpinStageBuffer(uint64_t addr) {
std::lock_guard<std::recursive_mutex> lock(stage_buffers_mu_);
for (auto& [location, buf] : stage_buffers_) {
auto base = reinterpret_cast<uint64_t>(buf.chunks);
auto end = base + chunk_size_ * chunk_count_;
Expand Down
6 changes: 6 additions & 0 deletions mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -795,6 +795,7 @@ BatchID TransferEngineImpl::allocateBatch(size_t batch_size) {
Batch* batch = Slab<Batch>::Get().allocate();
if (!batch) return (BatchID)0;
batch->max_size = batch_size;
batch->task_list.reserve(batch_size);
BatchID batch_id = (BatchID)batch;
std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
batch_set_.active.insert(batch);
Expand Down Expand Up @@ -1457,6 +1458,11 @@ void TransferEngineImpl::attachProgressNotifier(
Status TransferEngineImpl::commitPreparedSubmit(
Batch* batch, const PreparedSubmit& prepared) {
if (!batch) return Status::InvalidArgument("Invalid batch" LOC_MARK);
if (batch->task_list.size() > batch->max_size ||
prepared.tasks.size() > batch->max_size - batch->task_list.size()) {
return Status::TooManyRequests(
"batch public task capacity exceeded" LOC_MARK);
}

std::vector<Request> classified_request_list[kSupportedTransportTypes];
std::vector<size_t> task_id_list[kSupportedTransportTypes];
Expand Down
Loading
Loading