From 20dfd86858e597495936bb4715463635254eab2b Mon Sep 17 00:00:00 2001 From: ShuweiShen772 Date: Fri, 13 Nov 2026 16:16:41 +0800 Subject: [PATCH 1/6] [Bugfix] Preserve deferred batch cleanup ownership after query errors Retain library-owned deferred descriptors and delete hooks on query errors. Back off recoverable queries and use sealed completion publication to drain errors safely after callback pins reach zero. Add focused ownership, persistent-error, pin, and shutdown regression coverage. Signed-off-by: ShuweiShen772 --- .../include/multi_transport.h | 18 +- .../src/multi_transport.cpp | 55 +++- mooncake-transfer-engine/tests/CMakeLists.txt | 7 + .../tests/deferred_batch_cleanup_test.cpp | 305 ++++++++++++++++++ 4 files changed, 372 insertions(+), 13 deletions(-) create mode 100644 mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp diff --git a/mooncake-transfer-engine/include/multi_transport.h b/mooncake-transfer-engine/include/multi_transport.h index e00581eca..d8130aebd 100644 --- a/mooncake-transfer-engine/include/multi_transport.h +++ b/mooncake-transfer-engine/include/multi_transport.h @@ -15,6 +15,7 @@ #ifndef MULTI_TRANSPORT_H_ #define MULTI_TRANSPORT_H_ +#include #include #include #include @@ -135,7 +136,9 @@ class MultiTransport { bool *allows_reuse = nullptr); Status tryFreeBatchID(BatchID batch_id, - const std::function &before_delete); + const std::function &before_delete, + bool allow_completion_fallback = false, + bool poll_transport = true); void deferredCleanupLoop(); @@ -154,8 +157,17 @@ class MultiTransport { std::mutex deferred_cleanup_mutex_; std::condition_variable deferred_cleanup_cv_; - std::unordered_map> - deferred_cleanup_batches_; + struct DeferredCleanupEntry { + explicit DeferredCleanupEntry(const std::function &callback) + : before_delete(callback) {} + + std::function before_delete; + Status last_error; + bool completion_only = false; + std::chrono::steady_clock::time_point next_poll{}; + std::chrono::milliseconds retry_delay{100}; + }; + std::unordered_map deferred_cleanup_batches_; bool stop_deferred_cleanup_ = false; std::thread deferred_cleanup_thread_; }; diff --git a/mooncake-transfer-engine/src/multi_transport.cpp b/mooncake-transfer-engine/src/multi_transport.cpp index 189d1dc60..51fcf51f3 100644 --- a/mooncake-transfer-engine/src/multi_transport.cpp +++ b/mooncake-transfer-engine/src/multi_transport.cpp @@ -155,7 +155,8 @@ Status MultiTransport::freeBatchID(BatchID batch_id, } Status MultiTransport::tryFreeBatchID( - BatchID batch_id, const std::function& before_delete) { + BatchID batch_id, const std::function& before_delete, + bool allow_completion_fallback, bool poll_transport) { auto& batch_desc = *((BatchDesc*)(batch_id)); const size_t task_count = batch_desc.task_list.size(); bool busy = false; @@ -172,6 +173,20 @@ Status MultiTransport::tryFreeBatchID( task.is_finished = true; continue; } + // A failed status query is not evidence that a task stopped using + // its buffers. After such an error, only the sealed last-slice + // publication can substitute for polling. CompletionCallbackPin still + // protects callback accesses until the batch-wide check below. + if (allow_completion_fallback && + __atomic_load_n(&task.completion_published, __ATOMIC_ACQUIRE)) { + __atomic_store_n(&task.is_finished, true, __ATOMIC_RELEASE); + if (metricsEnabled()) recordFinishedTask(task); + continue; + } + if (!poll_transport) { + busy = true; + continue; + } TransferStatus status; Status query_status = getTransferStatus(batch_id, task_id, status); if (!query_status.ok()) { @@ -230,18 +245,38 @@ void MultiTransport::deferredCleanupLoop() { for (auto it = deferred_cleanup_batches_.begin(); it != deferred_cleanup_batches_.end();) { - Status status = tryFreeBatchID(it->first, it->second); + auto& entry = it->second; + const auto now = std::chrono::steady_clock::now(); + const bool poll_transport = + !entry.completion_only && now >= entry.next_poll; + Status status = tryFreeBatchID( + it->first, entry.before_delete, + /*allow_completion_fallback=*/!entry.last_error.ok(), + poll_transport); if (status.ok()) { it = deferred_cleanup_batches_.erase(it); - } else if (!status.IsBatchBusy()) { - // A permanent error (e.g. a stale task reference) will never - // resolve on retry; drop the entry instead of polling forever. - LOG(ERROR) << "Dropping deferred cleanup for batch " - << it->first << ": " << status.ToString(); - it = deferred_cleanup_batches_.erase(it); - } else { - ++it; + continue; + } + if (!status.IsBatchBusy()) { + // The caller already invalidated this handle. Retain both + // the descriptor and its delete hook on every query error. + // InvalidArgument/NotImplemented cannot be repaired by + // polling; wait for callback-published completion instead. + if (!(entry.last_error == status)) { + LOG(ERROR) << "Retaining deferred cleanup for batch " + << it->first << ": " << status.ToString(); + } + entry.last_error = status; + entry.completion_only = + status.IsInvalidArgument() || status.IsNotImplemented(); + entry.next_poll = now + entry.retry_delay; + entry.retry_delay = std::min(entry.retry_delay * 2, + std::chrono::milliseconds(5000)); + } else if (poll_transport && !entry.completion_only) { + entry.last_error = Status::OK(); + entry.retry_delay = kCleanupInterval; } + ++it; } if (stop_deferred_cleanup_) { diff --git a/mooncake-transfer-engine/tests/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index 3686051af..9ed1ca4b0 100644 --- a/mooncake-transfer-engine/tests/CMakeLists.txt +++ b/mooncake-transfer-engine/tests/CMakeLists.txt @@ -1,5 +1,12 @@ set(WORKSPACE "${CMAKE_CURRENT_SOURCE_DIR}") +# Hardware-free ownership/error-path regression coverage for the batch reaper. +add_executable(deferred_batch_cleanup_test + ${WORKSPACE}/deferred_batch_cleanup_test.cpp) +target_link_libraries(deferred_batch_cleanup_test PRIVATE transfer_engine gtest + gtest_main) +add_test(NAME deferred_batch_cleanup_test COMMAND deferred_batch_cleanup_test) + if(MOONCAKE_ENABLE_ADAPTIVE_CONGESTION_CONTROL) add_executable(adaptive_congestion_control_test ${WORKSPACE}/adaptive_congestion_control_test.cpp) diff --git a/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp b/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp new file mode 100644 index 000000000..4d43ddc93 --- /dev/null +++ b/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp @@ -0,0 +1,305 @@ +// Copyright 2026 KVCache.AI +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "multi_transport.h" +#include "transport/transport.h" + +namespace mooncake { +class MultiTransportBatchTestPeer { + public: + static Status free(MultiTransport& transport, Transport::BatchID id, + const std::function& hook) { + return transport.freeBatchID(id, hook); + } + + static bool pending(MultiTransport& transport, Transport::BatchID id) { + std::lock_guard lock(transport.deferred_cleanup_mutex_); + return transport.deferred_cleanup_batches_.count(id) != 0; + } +}; + +namespace { +using Peer = MultiTransportBatchTestPeer; +using Clock = std::chrono::steady_clock; + +enum class QueryError { None, Context, InvalidArgument, NotImplemented }; + +class CounterTransport : public Transport { + public: + std::atomic error{QueryError::None}; + std::atomic error_queries{0}; + + Status submitTransfer(BatchID, + const std::vector&) override { + return Status::OK(); + } + + Status getTransferStatus(BatchID id, size_t index, + TransferStatus& status) override { + const auto result = error.load(); + if (result != QueryError::None) { + error_queries.fetch_add(1); + switch (result) { + case QueryError::Context: + return Status::Context("injected query failure"); + case QueryError::InvalidArgument: + return Status::InvalidArgument("injected query failure"); + case QueryError::NotImplemented: + return Status::NotImplemented("injected query failure"); + case QueryError::None: + break; + } + } + auto& task = toBatchDesc(id).task_list[index]; + const auto success = + __atomic_load_n(&task.success_slice_count, __ATOMIC_ACQUIRE); + const auto failed = + __atomic_load_n(&task.failed_slice_count, __ATOMIC_ACQUIRE); + status.transferred_bytes = success; + status.s = TransferStatusEnum::WAITING; + if (success + failed == task.slice_count) { + status.s = failed ? TransferStatusEnum::FAILED + : TransferStatusEnum::COMPLETED; + __atomic_store_n(&task.is_finished, true, __ATOMIC_RELEASE); + } + return Status::OK(); + } + + private: + int registerLocalMemory(void*, size_t, const std::string&, bool, + bool) override { + return 0; + } + int unregisterLocalMemory(void*, bool) override { return 0; } + int registerLocalMemoryBatch(const std::vector&, + const std::string&) override { + return 0; + } + int unregisterLocalMemoryBatch(const std::vector&) override { + return 0; + } + const char* getName() const override { return "counter-cleanup-test"; } +}; + +struct CleanupCounts { + std::atomic deleted{0}; + std::atomic slices{0}; + std::atomic hook_before_slice{true}; +}; + +struct Batch { + Transport::BatchID id; + Transport::BatchDesc* desc; + std::vector slices; +}; + +Batch make_batch(MultiTransport& transport, CounterTransport& backend, + CleanupCounts& counts, size_t task_count = 1) { + const auto id = transport.allocateBatchID(task_count); + auto& desc = Transport::toBatchDesc(id); + Batch batch{id, &desc, {}}; + batch.slices.reserve(task_count); + for (size_t i = 0; i < task_count; ++i) { + auto& task = desc.task_list.emplace_back(); + task.batch_id = id; + task.transport_ = &backend; + task.slice_count = 1; + auto* slice = new Transport::Slice{}; + slice->task = &task; + slice->length = 1; + slice->status = Transport::Slice::PENDING; + slice->ts = 0; + slice->source_addr = &counts; + slice->cleanup_callback = [](Transport::Slice* released) { + auto& counts = *static_cast(released->source_addr); + if (counts.deleted.load() != 1) counts.hook_before_slice = false; + counts.slices.fetch_add(1); + }; + task.slice_list.push_back(slice); + batch.slices.push_back(slice); + } + return batch; +} + +bool wait_until(const std::function& predicate) { + const auto deadline = Clock::now() + std::chrono::seconds(3); + while (!predicate() && Clock::now() < deadline) + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + return predicate(); +} + +void expect_reclaimed(MultiTransport& transport, const Batch& batch, + CleanupCounts& counts) { + ASSERT_TRUE( + wait_until([&] { return !Peer::pending(transport, batch.id); })); + EXPECT_EQ(counts.deleted.load(), 1); + EXPECT_EQ(counts.slices.load(), batch.slices.size()); + EXPECT_TRUE(counts.hook_before_slice.load()); +} + +TEST(DeferredBatchCleanupTest, InitialQueryErrorDoesNotTransferOwnership) { + std::string name = "initial-query-error"; + CounterTransport backend; + CleanupCounts counts; + MultiTransport transport(nullptr, name); + auto batch = make_batch(transport, backend, counts); + backend.error = QueryError::Context; + EXPECT_TRUE( + Peer::free(transport, batch.id, [&] { ++counts.deleted; }).IsContext()); + EXPECT_FALSE(Peer::pending(transport, batch.id)); + EXPECT_EQ(counts.deleted.load(), 0); + batch.slices.front()->markFailed(); + backend.error = QueryError::None; + EXPECT_TRUE( + Peer::free(transport, batch.id, [&] { ++counts.deleted; }).ok()); + EXPECT_EQ(counts.deleted.load(), 1); + EXPECT_EQ(counts.slices.load(), 1); +} + +TEST(DeferredBatchCleanupTest, TransientQueryErrorRetainsOwnerAndRecovers) { + std::string name = "transient-query-error"; + CounterTransport backend; + CleanupCounts counts; + MultiTransport transport(nullptr, name); + auto batch = make_batch(transport, backend, counts); + ASSERT_TRUE(Peer::free(transport, batch.id, [&] { + ++counts.deleted; + }).IsBatchCleanupDeferred()); + backend.error = QueryError::Context; + ASSERT_TRUE(wait_until([&] { return backend.error_queries.load() > 0; })); + EXPECT_TRUE(Peer::pending(transport, batch.id)); + EXPECT_EQ(counts.deleted.load(), 0); + backend.error = QueryError::None; + batch.slices.front()->markSuccess(); + expect_reclaimed(transport, batch, counts); +} + +TEST(DeferredBatchCleanupTest, PersistentQueryErrorDrainsWithoutQueryRecovery) { + std::string name = "persistent-query-error"; + CounterTransport backend; + CleanupCounts counts; + MultiTransport transport(nullptr, name); + auto batch = make_batch(transport, backend, counts); + ASSERT_TRUE(Peer::free(transport, batch.id, [&] { + ++counts.deleted; + }).IsBatchCleanupDeferred()); + backend.error = QueryError::Context; + ASSERT_TRUE(wait_until([&] { return backend.error_queries.load() > 0; })); + EXPECT_TRUE(Peer::pending(transport, batch.id)); + std::this_thread::sleep_for(std::chrono::milliseconds(1100)); + EXPECT_GE(backend.error_queries.load(), 2); + EXPECT_LE(backend.error_queries.load(), 5); + batch.slices.front()->markFailed(); + expect_reclaimed(transport, batch, counts); +} + +class PermanentCleanupErrorTest : public ::testing::TestWithParam { +}; + +TEST_P(PermanentCleanupErrorTest, StopsPollingButRetainsCompletionOwner) { + std::string name = "permanent-query-error"; + CounterTransport backend; + CleanupCounts counts; + MultiTransport transport(nullptr, name); + auto batch = make_batch(transport, backend, counts); + ASSERT_TRUE(Peer::free(transport, batch.id, [&] { + ++counts.deleted; + }).IsBatchCleanupDeferred()); + backend.error = GetParam(); + ASSERT_TRUE(wait_until([&] { return backend.error_queries.load() > 0; })); + std::this_thread::sleep_for(std::chrono::milliseconds(250)); + EXPECT_EQ(backend.error_queries.load(), 1); + EXPECT_TRUE(Peer::pending(transport, batch.id)); + EXPECT_EQ(counts.deleted.load(), 0); + batch.slices.front()->markFailed(); + expect_reclaimed(transport, batch, counts); +} + +INSTANTIATE_TEST_SUITE_P(QueryErrors, PermanentCleanupErrorTest, + ::testing::Values(QueryError::InvalidArgument, + QueryError::NotImplemented)); + +TEST(DeferredBatchCleanupTest, WaitsForEveryTaskAfterQueryError) { + std::string name = "partial-completion"; + CounterTransport backend; + CleanupCounts counts; + MultiTransport transport(nullptr, name); + auto batch = make_batch(transport, backend, counts, 2); + ASSERT_TRUE(Peer::free(transport, batch.id, [&] { + ++counts.deleted; + }).IsBatchCleanupDeferred()); + backend.error = QueryError::InvalidArgument; + ASSERT_TRUE(wait_until([&] { return backend.error_queries.load() > 0; })); + batch.slices.front()->markFailed(); + std::this_thread::sleep_for(std::chrono::milliseconds(150)); + EXPECT_TRUE(Peer::pending(transport, batch.id)); + EXPECT_EQ(counts.deleted.load(), 0); + batch.slices.back()->markSuccess(); + expect_reclaimed(transport, batch, counts); +} + +TEST(DeferredBatchCleanupTest, PublishedCompletionDoesNotBypassCallbackPin) { + std::string name = "completion-pin"; + CounterTransport backend; + CleanupCounts counts; + MultiTransport transport(nullptr, name); + auto batch = make_batch(transport, backend, counts); + ASSERT_TRUE(Peer::free(transport, batch.id, [&] { + ++counts.deleted; + }).IsBatchCleanupDeferred()); + backend.error = QueryError::Context; + ASSERT_TRUE(wait_until([&] { return backend.error_queries.load() > 0; })); + // Model an in-flight callback's internal pin, not a caller reusing its + // invalid public handle. The pin keeps this saved descriptor alive. + batch.desc->active_completion_callbacks.fetch_add(1); + batch.slices.front()->markFailed(); + std::this_thread::sleep_for(std::chrono::milliseconds(150)); + EXPECT_TRUE(Peer::pending(transport, batch.id)); + EXPECT_EQ(counts.deleted.load(), 0); + EXPECT_EQ(counts.slices.load(), 0); + batch.desc->active_completion_callbacks.fetch_sub(1); + expect_reclaimed(transport, batch, counts); +} + +TEST(DeferredBatchCleanupTest, ShutdownReclaimsAlreadyPublishedErrorBatch) { + std::string name = "shutdown-error-batch"; + CounterTransport backend; + CleanupCounts counts; + auto transport = std::make_unique(nullptr, name); + auto batch = make_batch(*transport, backend, counts); + ASSERT_TRUE(Peer::free(*transport, batch.id, [&] { + ++counts.deleted; + }).IsBatchCleanupDeferred()); + backend.error = QueryError::InvalidArgument; + ASSERT_TRUE(wait_until([&] { return backend.error_queries.load() > 0; })); + batch.slices.front()->markFailed(); + transport.reset(); + EXPECT_EQ(counts.deleted.load(), 1); + EXPECT_EQ(counts.slices.load(), 1); + EXPECT_TRUE(counts.hook_before_slice.load()); +} + +} // namespace +} // namespace mooncake From 6107763e5b0f3ddd800f9c06b9cc4b230421eee5 Mon Sep 17 00:00:00 2001 From: ShuweiShen772 Date: Fri, 13 Nov 2026 18:41:06 +0800 Subject: [PATCH 2/6] [TransferEngine] Expose caller-owned batch retirement status Keep borrowed buffers live until every task and completion callback has retired. Distinguish successful terminal queries from legacy finished flags, and preserve completion notification behavior. Signed-off-by: ShuweiShen772 --- .../include/multi_transport.h | 5 ++ .../include/transfer_engine.h | 12 +++ .../include/transfer_engine_impl.h | 13 +++ .../include/transport/transport.h | 3 + .../src/multi_transport.cpp | 83 +++++++++++++++++++ .../src/transfer_engine.cpp | 5 ++ .../tests/deferred_batch_cleanup_test.cpp | 83 +++++++++++++++++++ 7 files changed, 204 insertions(+) diff --git a/mooncake-transfer-engine/include/multi_transport.h b/mooncake-transfer-engine/include/multi_transport.h index d8130aebd..ae88d417e 100644 --- a/mooncake-transfer-engine/include/multi_transport.h +++ b/mooncake-transfer-engine/include/multi_transport.h @@ -81,6 +81,11 @@ class MultiTransport { Status getBatchTransferStatus(BatchID batch_id, TransferStatus &status); + // Observe retirement without releasing or transferring batch ownership. + // A non-OK result reports a query/transfer failure; only retired permits + // the caller to stop retaining borrowed requests and transfer buffers. + Status getBatchRetirementStatus(BatchID batch_id, bool &retired); + Transport *installTransport(const std::string &proto, std::shared_ptr topo); diff --git a/mooncake-transfer-engine/include/transfer_engine.h b/mooncake-transfer-engine/include/transfer_engine.h index 8767efac0..2c93cfbf0 100644 --- a/mooncake-transfer-engine/include/transfer_engine.h +++ b/mooncake-transfer-engine/include/transfer_engine.h @@ -289,6 +289,18 @@ class TransferEngine { Status getBatchTransferStatus(BatchID batch_id, TransferStatus& status); +#ifndef USE_TENT + // After submission returns, poll a caller-owned, valid batch until retired + // is true. All tasks and completion callbacks must have stopped accessing + // borrowed requests/buffers. This does not free or invalidate batch_id. + // A non-OK result reports an observed query/transfer failure independently + // of retirement. TIMEOUT/error alone never authorizes buffer reuse. + // Do not submit more work, free, or query this batch concurrently. As with + // transfer completion, transports must publish terminal state after I/O + // retirement (including failure paths). + Status getBatchRetirementStatus(BatchID batch_id, bool& retired); +#endif + Status getNicLoadStats(std::vector& stats) const; Transport* getTransport(const std::string& proto); diff --git a/mooncake-transfer-engine/include/transfer_engine_impl.h b/mooncake-transfer-engine/include/transfer_engine_impl.h index 234387993..43ffd25cf 100644 --- a/mooncake-transfer-engine/include/transfer_engine_impl.h +++ b/mooncake-transfer-engine/include/transfer_engine_impl.h @@ -342,6 +342,19 @@ class TransferEngineImpl { return result; } + Status getBatchRetirementStatus(BatchID batch_id, bool& retired) { + Status result = + multi_transports_->getBatchRetirementStatus(batch_id, retired); + if (retired && result.ok()) { + // Preserve completion metrics/notifications previously driven by + // the aggregate status polling path before the caller frees. + TransferStatus status; + Status completed = getBatchTransferStatus(batch_id, status); + if (result.ok() && !completed.ok()) return completed; + } + return result; + } + Transport* getTransport(const std::string& proto) { if (!multi_transports_) return nullptr; return multi_transports_->getTransport(proto); diff --git a/mooncake-transfer-engine/include/transport/transport.h b/mooncake-transfer-engine/include/transport/transport.h index 40ab57870..7e6a080b7 100644 --- a/mooncake-transfer-engine/include/transport/transport.h +++ b/mooncake-transfer-engine/include/transport/transport.h @@ -504,6 +504,9 @@ class Transport { volatile uint64_t completion_state = Slice::kCompletionSealBit; volatile uint64_t transferred_bytes = 0; volatile bool is_finished = false; + // Set by MultiTransport only after a successful terminal status query. + // Legacy is_finished alone is not proof after an API query error. + volatile bool retirement_observed = false; uint64_t total_bytes = 0; BatchID batch_id = 0; diff --git a/mooncake-transfer-engine/src/multi_transport.cpp b/mooncake-transfer-engine/src/multi_transport.cpp index 51fcf51f3..87e6aae37 100644 --- a/mooncake-transfer-engine/src/multi_transport.cpp +++ b/mooncake-transfer-engine/src/multi_transport.cpp @@ -589,6 +589,12 @@ Status MultiTransport::getTransferStatus(BatchID batch_id, size_t task_id, checkSliceTimeout(task)) { status.s = Transport::TransferStatusEnum::TIMEOUT; } + if (status.s == TransferStatusEnum::COMPLETED || + status.s == TransferStatusEnum::FAILED) { + __atomic_store_n(&task.retirement_observed, true, __ATOMIC_RELEASE); + __atomic_store_n(&task.transferred_bytes, status.transferred_bytes, + __ATOMIC_RELEASE); + } recordTaskEnd(task, status); return Status::OK(); } @@ -736,6 +742,83 @@ Status MultiTransport::getBatchTransferStatus(BatchID batch_id, return Status::OK(); } +Status MultiTransport::getBatchRetirementStatus(BatchID batch_id, + bool& retired) { + retired = false; + auto& batch = *((BatchDesc*)(batch_id)); + Status result = batch.has_failure.load(std::memory_order_acquire) + ? Status::Context("Mooncake transfer task failed") + : Status::OK(); + bool pending = false; + uint64_t completed_bytes = 0; + for (size_t task_id = 0; task_id < batch.task_list.size(); ++task_id) { + auto& task = batch.task_list[task_id]; + const auto slice_count = + __atomic_load_n(&task.slice_count, __ATOMIC_ACQUIRE); + const bool counter_completion = + task.transport_ && + task.transport_->supportsBatchCompletionCounter(); + bool done = + slice_count == 0 || + __atomic_load_n(&task.completion_published, __ATOMIC_ACQUIRE) || + (!counter_completion && + __atomic_load_n(&task.retirement_observed, __ATOMIC_ACQUIRE)); + if (!done && !counter_completion) { + // Counter transports publish completion without polling. Other + // transports need polling to drive completion. Visit every + // task even after a sibling failed or its query returned an error. + TransferStatus status{}; + const Status queried = getTransferStatus(batch_id, task_id, status); + if (!queried.ok()) { + if (result.ok()) result = queried; + // A sealed last-slice publication remains a retirement proof + // if a transport status query fails. Callback pins still gate + // the batch-wide result below. + done = __atomic_load_n(&task.completion_published, + __ATOMIC_ACQUIRE); + } else { + done = status.s == TransferStatusEnum::COMPLETED || + status.s == TransferStatusEnum::FAILED; + if (done) { + __atomic_store_n(&task.transferred_bytes, + status.transferred_bytes, + __ATOMIC_RELEASE); + } + if (status.s == TransferStatusEnum::FAILED) { + batch.has_failure.store(true, std::memory_order_release); + } + if (result.ok() && (status.s == TransferStatusEnum::FAILED || + status.s == TransferStatusEnum::TIMEOUT)) { + result = Status::Context("Mooncake transfer task failed"); + } + } + } + if (__atomic_load_n(&task.failed_slice_count, __ATOMIC_ACQUIRE) && + result.ok()) { + result = Status::Context("Mooncake transfer task failed"); + } + if (!done) { + pending = true; + continue; + } + __atomic_store_n(&task.is_finished, true, __ATOMIC_RELEASE); + completed_bytes += + __atomic_load_n(&task.transferred_bytes, __ATOMIC_ACQUIRE); + if (metricsEnabled()) recordFinishedTask(task); + } + retired = !pending && batch.active_completion_callbacks.load( + std::memory_order_acquire) == 0; + if (retired && result.ok()) { + // Let the wrapper publish completion metrics/notifications without + // querying a transport again after its terminal completion. + batch.finished_transfer_bytes.store(completed_bytes, + std::memory_order_relaxed); + batch.is_finished.store(true, std::memory_order_relaxed); + batch.status_cached.store(true, std::memory_order_release); + } + return result; +} + Transport* MultiTransport::installTransport(const std::string& proto, std::shared_ptr topo) { #ifdef USE_NCCL_HOST diff --git a/mooncake-transfer-engine/src/transfer_engine.cpp b/mooncake-transfer-engine/src/transfer_engine.cpp index 4a016afba..61959f0a4 100644 --- a/mooncake-transfer-engine/src/transfer_engine.cpp +++ b/mooncake-transfer-engine/src/transfer_engine.cpp @@ -303,6 +303,11 @@ Status TransferEngine::getBatchTransferStatus(BatchID batch_id, return impl_->getBatchTransferStatus(batch_id, status); } +Status TransferEngine::getBatchRetirementStatus(BatchID batch_id, + bool& retired) { + return impl_->getBatchRetirementStatus(batch_id, retired); +} + Status TransferEngine::getNicLoadStats(std::vector& stats) const { stats.clear(); return Status::OK(); diff --git a/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp b/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp index 4d43ddc93..c1aea8342 100644 --- a/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp +++ b/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp @@ -50,6 +50,10 @@ class CounterTransport : public Transport { public: std::atomic error{QueryError::None}; std::atomic error_queries{0}; + bool counter_completion = false; + bool supportsBatchCompletionCounter() const override { + return counter_completion; + } Status submitTransfer(BatchID, const std::vector&) override { @@ -301,5 +305,84 @@ TEST(DeferredBatchCleanupTest, ShutdownReclaimsAlreadyPublishedErrorBatch) { EXPECT_TRUE(counts.hook_before_slice.load()); } +TEST(DeferredBatchCleanupTest, RetirementKeepsHandleUntilAllTasksAndPinsDrain) { + std::string name = "retirement-failed-sibling"; + CounterTransport backend; + CleanupCounts counts; + MultiTransport transport(nullptr, name); + auto batch = make_batch(transport, backend, counts, 2); + bool retired = true; + batch.slices.front()->markFailed(); + EXPECT_TRUE( + transport.getBatchRetirementStatus(batch.id, retired).IsContext()); + EXPECT_FALSE(retired); + EXPECT_EQ(counts.deleted.load(), 0); + EXPECT_FALSE(Peer::pending(transport, batch.id)); + batch.desc->active_completion_callbacks.fetch_add(1); + batch.slices.back()->markSuccess(); + EXPECT_TRUE( + transport.getBatchRetirementStatus(batch.id, retired).IsContext()); + EXPECT_FALSE(retired); + batch.desc->active_completion_callbacks.fetch_sub(1); + EXPECT_TRUE( + transport.getBatchRetirementStatus(batch.id, retired).IsContext()); + EXPECT_TRUE(retired); + EXPECT_EQ(counts.deleted.load(), 0); + EXPECT_TRUE( + Peer::free(transport, batch.id, [&] { ++counts.deleted; }).ok()); + EXPECT_EQ(counts.deleted.load(), 1); +} + +TEST(DeferredBatchCleanupTest, RetirementRequiresSealedCounterPublication) { + std::string name = "retirement-submission-seal"; + CounterTransport backend; + backend.counter_completion = true; + CleanupCounts counts; + MultiTransport transport(nullptr, name); + auto batch = make_batch(transport, backend, counts); + auto& task = batch.desc->task_list.front(); + Transport::Slice::unsealTaskSubmission(&task); + batch.slices.front()->markSuccess(); + Transport::TransferStatus terminal{}; + EXPECT_TRUE(backend.getTransferStatus(batch.id, 0, terminal).ok()); + bool retired = true; + EXPECT_TRUE(transport.getBatchRetirementStatus(batch.id, retired).ok()); + EXPECT_FALSE(retired); + // A prior terminal getter cannot make is_finished bypass the seal. + EXPECT_TRUE(transport.getBatchRetirementStatus(batch.id, retired).ok()); + EXPECT_FALSE(retired); + Transport::Slice::sealTaskSubmission(&task); + EXPECT_TRUE(transport.getBatchRetirementStatus(batch.id, retired).ok()); + EXPECT_TRUE(retired); + EXPECT_TRUE( + Peer::free(transport, batch.id, [&] { ++counts.deleted; }).ok()); +} + +TEST(DeferredBatchCleanupTest, RetirementQueryFailureDoesNotInvalidateHandle) { + std::string name = "retirement-query-error"; + CounterTransport backend; + CleanupCounts counts; + MultiTransport transport(nullptr, name); + auto batch = make_batch(transport, backend, counts); + backend.error = QueryError::Context; + // A legacy finished flag does not make a failed API query a retirement + // proof; only a successful terminal query or publication can do that. + __atomic_store_n(&batch.desc->task_list.front().is_finished, true, + __ATOMIC_RELEASE); + bool retired = true; + EXPECT_TRUE( + transport.getBatchRetirementStatus(batch.id, retired).IsContext()); + EXPECT_FALSE(retired); + EXPECT_FALSE(Peer::pending(transport, batch.id)); + batch.slices.front()->markFailed(); + EXPECT_TRUE( + transport.getBatchRetirementStatus(batch.id, retired).IsContext()); + EXPECT_TRUE(retired); + EXPECT_EQ(counts.deleted.load(), 0); + EXPECT_TRUE( + Peer::free(transport, batch.id, [&] { ++counts.deleted; }).ok()); + EXPECT_EQ(counts.deleted.load(), 1); +} + } // namespace } // namespace mooncake From b2fccb83c3354d1f1e73aa3752cc116efe8f5154 Mon Sep 17 00:00:00 2001 From: ShuweiShen772 Date: Fri, 13 Nov 2026 23:40:05 +0800 Subject: [PATCH 3/6] [Bugfix] Make retirement byte accounting idempotent Signed-off-by: ShuweiShen772 --- .../include/transfer_engine_impl.h | 23 ++++- .../include/transport/transport.h | 3 + .../tests/transport_uint_test.cpp | 95 +++++++++++++++++++ 3 files changed, 116 insertions(+), 5 deletions(-) diff --git a/mooncake-transfer-engine/include/transfer_engine_impl.h b/mooncake-transfer-engine/include/transfer_engine_impl.h index 43ffd25cf..10e3deceb 100644 --- a/mooncake-transfer-engine/include/transfer_engine_impl.h +++ b/mooncake-transfer-engine/include/transfer_engine_impl.h @@ -276,9 +276,7 @@ class TransferEngineImpl { // Only record metrics for successful completions if (status.s == TransferStatusEnum::COMPLETED) { - if (status.transferred_bytes > 0) { - transferred_bytes_counter_.inc(status.transferred_bytes); - } + recordTransferredBytes(task, status.transferred_bytes); auto now = std::chrono::steady_clock::now(); auto duration = std::chrono::duration_cast( @@ -319,8 +317,14 @@ class TransferEngineImpl { #ifdef WITH_METRICS if (metrics_enabled_ && !skip_metrics && result.ok() && status.s == TransferStatusEnum::COMPLETED) { - if (status.transferred_bytes > 0) { - transferred_bytes_counter_.inc(status.transferred_bytes); + // The cached aggregate is built from these task byte counts. + // Share the per-task claim with getTransferStatus so mixed and + // repeated observers cannot count the same successful bytes twice. + auto& batch = Transport::toBatchDesc(batch_id); + for (auto& task : batch.task_list) { + recordTransferredBytes( + task, + __atomic_load_n(&task.transferred_bytes, __ATOMIC_ACQUIRE)); } } #endif @@ -540,6 +544,15 @@ class TransferEngineImpl { bool metrics_enabled_{false}; uint64_t metrics_interval_seconds_{5}; + void recordTransferredBytes(Transport::TransferTask& task, + uint64_t transferred_bytes) { + if (transferred_bytes > 0 && + !__atomic_exchange_n(&task.log_bytes_recorded, true, + __ATOMIC_ACQ_REL)) { + transferred_bytes_counter_.inc(transferred_bytes); + } + } + // Helper methods for metrics reporting thread management void InitializeMetricsConfig(); void StartMetricsReportingThread(); diff --git a/mooncake-transfer-engine/include/transport/transport.h b/mooncake-transfer-engine/include/transport/transport.h index 7e6a080b7..8d752d8f2 100644 --- a/mooncake-transfer-engine/include/transport/transport.h +++ b/mooncake-transfer-engine/include/transport/transport.h @@ -525,6 +525,9 @@ class Transport { // make the task non-movable). volatile bool export_metrics_recorded = false; volatile bool log_metrics_recorded = false; + // Task and aggregate/retirement observers share byte accounting. + // Keep latency publication independent from the first byte observer. + volatile bool log_bytes_recorded = false; // Set by the event-driven completion callback when it finishes the // task: steady_clock ticks (else 0) and COMPLETED or FAILED (else // WAITING). Atomic: a poll may read them while the callback writes. diff --git a/mooncake-transfer-engine/tests/transport_uint_test.cpp b/mooncake-transfer-engine/tests/transport_uint_test.cpp index bd75781d4..ed5d1f5c6 100644 --- a/mooncake-transfer-engine/tests/transport_uint_test.cpp +++ b/mooncake-transfer-engine/tests/transport_uint_test.cpp @@ -99,6 +99,42 @@ class TransferEngineImplTestPeer { return *engine.impl_; } +#ifdef WITH_METRICS + static void enableLogMetrics(TransferEngineImpl& engine) { + engine.metrics_enabled_ = true; + engine.multi_transports_ = std::make_shared( + nullptr, engine.local_server_name_); + } + + static uint64_t transferredBytes(TransferEngineImpl& engine) { + return engine.transferred_bytes_counter_.value(); + } + + static BatchID prepareCompletedBatch(TransferEngineImpl& engine, + size_t task_count, + uint64_t bytes_per_task) { + const BatchID id = engine.allocateBatchID(task_count); + auto& batch = Transport::toBatchDesc(id); + for (size_t i = 0; i < task_count; ++i) { + auto& task = batch.task_list.emplace_back(); + task.batch_id = id; + task.slice_count = 1; + task.success_slice_count = 1; + task.is_finished = true; + task.retirement_observed = true; + task.completion_published = true; + task.total_bytes = bytes_per_task; + task.transferred_bytes = bytes_per_task; + task.start_time = std::chrono::steady_clock::now(); + } + batch.finished_task_count = task_count; + batch.finished_transfer_bytes = task_count * bytes_per_task; + batch.is_finished = true; + batch.status_cached = true; + return id; + } +#endif + static AutoDiscoverConfig autoDiscoverConfig( const TransferEngineImpl& engine) { return engine.auto_discover_config_; @@ -611,6 +647,65 @@ TEST(TransferEngineTentCompatibilityTest, ForceTcpEnvForcesTcpTransport) { } #endif +#if defined(WITH_METRICS) && !defined(USE_TENT) +TEST(TransferEngineRetirementMetricsTest, RepeatedObservationsCountBytesOnce) { + ScopedEnvVar disable_reporter("MC_TE_METRIC", "0"); + TransferEngineImpl engine; + TransferEngineImplTestPeer::enableLogMetrics(engine); + const BatchID id = + TransferEngineImplTestPeer::prepareCompletedBatch(engine, 1, 64); + for (int i = 0; i < 2; ++i) { + bool retired = false; + EXPECT_TRUE(engine.getBatchRetirementStatus(id, retired).ok()); + EXPECT_TRUE(retired); + } + EXPECT_EQ(TransferEngineImplTestPeer::transferredBytes(engine), 64); + + // Concurrent read-only observations retain the caller-owned handle until + // every observer exits; the accounting claim itself must be atomic. + std::array observers; + for (auto& observer : observers) { + observer = std::thread([&] { + for (int i = 0; i < 10; ++i) { + bool retired = false; + EXPECT_TRUE(engine.getBatchRetirementStatus(id, retired).ok()); + EXPECT_TRUE(retired); + } + }); + } + for (auto& observer : observers) observer.join(); + EXPECT_EQ(TransferEngineImplTestPeer::transferredBytes(engine), 64); + EXPECT_TRUE(engine.freeBatchID(id).ok()); +} + +TEST(TransferEngineRetirementMetricsTest, TaskAndBatchObserversShareByteClaim) { + ScopedEnvVar disable_reporter("MC_TE_METRIC", "0"); + TransferEngineImpl engine; + TransferEngineImplTestPeer::enableLogMetrics(engine); + const BatchID id = + TransferEngineImplTestPeer::prepareCompletedBatch(engine, 2, 32); + TransferStatus status{}; + EXPECT_TRUE(engine.getTransferStatus(id, 0, status).ok()); + EXPECT_EQ(TransferEngineImplTestPeer::transferredBytes(engine), 32); + EXPECT_TRUE(engine.getBatchTransferStatus(id, status).ok()); + EXPECT_EQ(TransferEngineImplTestPeer::transferredBytes(engine), 64); + bool retired = false; + EXPECT_TRUE(engine.getBatchRetirementStatus(id, retired).ok()); + EXPECT_TRUE(retired); + EXPECT_TRUE(engine.getTransferStatus(id, 1, status).ok()); + EXPECT_EQ(TransferEngineImplTestPeer::transferredBytes(engine), 64); + EXPECT_TRUE(engine.freeBatchID(id).ok()); + + // Accounting belongs to tasks, not a pointer-derived historical BatchID. + const BatchID next = + TransferEngineImplTestPeer::prepareCompletedBatch(engine, 1, 64); + EXPECT_TRUE(engine.getBatchRetirementStatus(next, retired).ok()); + EXPECT_TRUE(retired); + EXPECT_EQ(TransferEngineImplTestPeer::transferredBytes(engine), 128); + EXPECT_TRUE(engine.freeBatchID(next).ok()); +} +#endif + class BatchResultTransport : public Transport { public: explicit BatchResultTransport(int unregister_result = 0) From ab4cc79e159217acb693f7c3d055893ca687b446 Mon Sep 17 00:00:00 2001 From: ShuweiShen772 Date: Sat, 14 Nov 2026 02:20:23 +0800 Subject: [PATCH 4/6] [Bugfix] Observe retained cleanup and clarify retirement contracts Signed-off-by: ShuweiShen772 --- .../include/multi_transport.h | 11 ++ .../include/transfer_engine.h | 10 +- .../include/transfer_engine_impl.h | 7 +- .../include/transport/transport.h | 7 + .../src/multi_transport.cpp | 47 +++++-- .../src/transfer_engine.cpp | 4 + .../tests/deferred_batch_cleanup_test.cpp | 131 +++++++++++++++++- 7 files changed, 206 insertions(+), 11 deletions(-) diff --git a/mooncake-transfer-engine/include/multi_transport.h b/mooncake-transfer-engine/include/multi_transport.h index ae88d417e..851407472 100644 --- a/mooncake-transfer-engine/include/multi_transport.h +++ b/mooncake-transfer-engine/include/multi_transport.h @@ -60,6 +60,8 @@ class MultiTransport { // cleanup means in-flight transport work may still own transfer buffers. Status freeBatchID(BatchID batch_id); + Transport::DeferredCleanupStats getDeferredCleanupStats(); + Status submitTransfer(BatchID batch_id, const std::vector &entries); @@ -140,6 +142,12 @@ class MultiTransport { Status selectTransport(const TransferRequest &entry, Transport *&transport, bool *allows_reuse = nullptr); + // The hook cleans descriptor-scoped state before deletion. It is not a + // user-buffer retirement callback; observe retirement before free when + // borrowed requests/buffers must be reused. + Transport::DeferredCleanupStats collectDeferredCleanupStats( + std::chrono::steady_clock::time_point now) const; + Status tryFreeBatchID(BatchID batch_id, const std::function &before_delete, bool allow_completion_fallback = false, @@ -167,12 +175,15 @@ class MultiTransport { : before_delete(callback) {} std::function before_delete; + std::chrono::steady_clock::time_point enqueued_at = + std::chrono::steady_clock::now(); Status last_error; bool completion_only = false; std::chrono::steady_clock::time_point next_poll{}; std::chrono::milliseconds retry_delay{100}; }; std::unordered_map deferred_cleanup_batches_; + std::chrono::steady_clock::time_point next_deferred_cleanup_warning_{}; bool stop_deferred_cleanup_ = false; std::thread deferred_cleanup_thread_; }; diff --git a/mooncake-transfer-engine/include/transfer_engine.h b/mooncake-transfer-engine/include/transfer_engine.h index 2c93cfbf0..ef0520e0b 100644 --- a/mooncake-transfer-engine/include/transfer_engine.h +++ b/mooncake-transfer-engine/include/transfer_engine.h @@ -297,8 +297,16 @@ class TransferEngine { // of retirement. TIMEOUT/error alone never authorizes buffer reuse. // Do not submit more work, free, or query this batch concurrently. As with // transfer completion, transports must publish terminal state after I/O - // retirement (including failure paths). + // retirement (including failure paths). TIMEOUT is not terminal: tasks may + // subsequently complete/fail, but the timeout itself never permits reuse. + // Never query after free returns OK or BatchCleanupDeferred: the handle is + // invalid, possibly already deleted or reused. This API does not validate + // expired pointer-derived IDs. Status getBatchRetirementStatus(BatchID batch_id, bool& retired); + + // Current reaper-owned backlog. This is not a count of all live batches, + // leaked allocations, historical IDs, or bytes in borrowed device buffers. + Transport::DeferredCleanupStats getDeferredCleanupStats(); #endif Status getNicLoadStats(std::vector& stats) const; diff --git a/mooncake-transfer-engine/include/transfer_engine_impl.h b/mooncake-transfer-engine/include/transfer_engine_impl.h index 10e3deceb..14fbbd323 100644 --- a/mooncake-transfer-engine/include/transfer_engine_impl.h +++ b/mooncake-transfer-engine/include/transfer_engine_impl.h @@ -354,11 +354,16 @@ class TransferEngineImpl { // the aggregate status polling path before the caller frees. TransferStatus status; Status completed = getBatchTransferStatus(batch_id, status); - if (result.ok() && !completed.ok()) return completed; + if (!completed.ok()) return completed; } return result; } + Transport::DeferredCleanupStats getDeferredCleanupStats() { + if (!multi_transports_) return {}; + return multi_transports_->getDeferredCleanupStats(); + } + Transport* getTransport(const std::string& proto) { if (!multi_transports_) return nullptr; return multi_transports_->getTransport(proto); diff --git a/mooncake-transfer-engine/include/transport/transport.h b/mooncake-transfer-engine/include/transport/transport.h index 8d752d8f2..b356abea2 100644 --- a/mooncake-transfer-engine/include/transport/transport.h +++ b/mooncake-transfer-engine/include/transport/transport.h @@ -552,6 +552,13 @@ class Transport { } }; + // A live snapshot of reaper-owned descriptors, not historical BatchIDs. + struct DeferredCleanupStats { + uint64_t pending_batches = 0; + uint64_t query_error_batches = 0; + uint64_t oldest_age_ms = 0; + }; + struct BatchDesc { BatchID id; size_t batch_size; diff --git a/mooncake-transfer-engine/src/multi_transport.cpp b/mooncake-transfer-engine/src/multi_transport.cpp index 87e6aae37..73b85b098 100644 --- a/mooncake-transfer-engine/src/multi_transport.cpp +++ b/mooncake-transfer-engine/src/multi_transport.cpp @@ -122,6 +122,27 @@ MultiTransport::BatchID MultiTransport::allocateBatchID(size_t batch_size) { return batch_desc->id; } +Transport::DeferredCleanupStats MultiTransport::collectDeferredCleanupStats( + std::chrono::steady_clock::time_point now) const { + Transport::DeferredCleanupStats stats; + stats.pending_batches = deferred_cleanup_batches_.size(); + for (const auto& [id, entry] : deferred_cleanup_batches_) { + if (!entry.last_error.ok()) ++stats.query_error_batches; + const auto age = std::chrono::duration_cast( + now - entry.enqueued_at) + .count(); + if (age > 0) + stats.oldest_age_ms = + std::max(stats.oldest_age_ms, static_cast(age)); + } + return stats; +} + +Transport::DeferredCleanupStats MultiTransport::getDeferredCleanupStats() { + std::lock_guard guard(deferred_cleanup_mutex_); + return collectDeferredCleanupStats(std::chrono::steady_clock::now()); +} + Status MultiTransport::freeBatchID(BatchID batch_id) { return freeBatchID(batch_id, std::function()); } @@ -279,6 +300,19 @@ void MultiTransport::deferredCleanupLoop() { ++it; } + const auto warning_now = std::chrono::steady_clock::now(); + if (warning_now >= next_deferred_cleanup_warning_) { + next_deferred_cleanup_warning_ = + warning_now + std::chrono::seconds(60); + const auto stats = collectDeferredCleanupStats(warning_now); + if (stats.oldest_age_ms >= 60000) { + LOG(WARNING) << "Deferred cleanup backlog: pending=" + << stats.pending_batches + << ", query_errors=" << stats.query_error_batches + << ", oldest_age_ms=" << stats.oldest_age_ms; + } + } + if (stop_deferred_cleanup_) { if (!deferred_cleanup_batches_.empty()) { LOG(WARNING) @@ -779,16 +813,13 @@ Status MultiTransport::getBatchRetirementStatus(BatchID batch_id, } else { done = status.s == TransferStatusEnum::COMPLETED || status.s == TransferStatusEnum::FAILED; - if (done) { - __atomic_store_n(&task.transferred_bytes, - status.transferred_bytes, - __ATOMIC_RELEASE); - } if (status.s == TransferStatusEnum::FAILED) { batch.has_failure.store(true, std::memory_order_release); } - if (result.ok() && (status.s == TransferStatusEnum::FAILED || - status.s == TransferStatusEnum::TIMEOUT)) { + if (result.ok() && status.s == TransferStatusEnum::TIMEOUT) { + result = Status::Context("Mooncake transfer task timeout"); + } else if (result.ok() && + status.s == TransferStatusEnum::FAILED) { result = Status::Context("Mooncake transfer task failed"); } } @@ -813,7 +844,7 @@ Status MultiTransport::getBatchRetirementStatus(BatchID batch_id, // querying a transport again after its terminal completion. batch.finished_transfer_bytes.store(completed_bytes, std::memory_order_relaxed); - batch.is_finished.store(true, std::memory_order_relaxed); + batch.is_finished.store(true, std::memory_order_release); batch.status_cached.store(true, std::memory_order_release); } return result; diff --git a/mooncake-transfer-engine/src/transfer_engine.cpp b/mooncake-transfer-engine/src/transfer_engine.cpp index 61959f0a4..1f1a02571 100644 --- a/mooncake-transfer-engine/src/transfer_engine.cpp +++ b/mooncake-transfer-engine/src/transfer_engine.cpp @@ -308,6 +308,10 @@ Status TransferEngine::getBatchRetirementStatus(BatchID batch_id, return impl_->getBatchRetirementStatus(batch_id, retired); } +Transport::DeferredCleanupStats TransferEngine::getDeferredCleanupStats() { + return impl_->getDeferredCleanupStats(); +} + Status TransferEngine::getNicLoadStats(std::vector& stats) const { stats.clear(); return Status::OK(); diff --git a/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp b/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp index c1aea8342..a5957313a 100644 --- a/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp +++ b/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp @@ -13,6 +13,7 @@ // limitations under the License. #include +#include #include #include @@ -20,6 +21,7 @@ #include #include #include +#include #include #include @@ -34,6 +36,15 @@ class MultiTransportBatchTestPeer { return transport.freeBatchID(id, hook); } + static void forceBacklogWarning(MultiTransport& transport) { + std::lock_guard lock(transport.deferred_cleanup_mutex_); + for (auto& [id, entry] : transport.deferred_cleanup_batches_) + entry.enqueued_at = + std::chrono::steady_clock::now() - std::chrono::seconds(61); + transport.next_deferred_cleanup_warning_ = {}; + transport.deferred_cleanup_cv_.notify_one(); + } + static bool pending(MultiTransport& transport, Transport::BatchID id) { std::lock_guard lock(transport.deferred_cleanup_mutex_); return transport.deferred_cleanup_batches_.count(id) != 0; @@ -50,7 +61,10 @@ class CounterTransport : public Transport { public: std::atomic error{QueryError::None}; std::atomic error_queries{0}; + std::atomic queries{0}; + std::atomic timed_out{false}; bool counter_completion = false; + bool publish_finished = true; bool supportsBatchCompletionCounter() const override { return counter_completion; } @@ -62,6 +76,7 @@ class CounterTransport : public Transport { Status getTransferStatus(BatchID id, size_t index, TransferStatus& status) override { + queries.fetch_add(1); const auto result = error.load(); if (result != QueryError::None) { error_queries.fetch_add(1); @@ -76,6 +91,11 @@ class CounterTransport : public Transport { break; } } + if (timed_out.load()) { + status.s = TransferStatusEnum::TIMEOUT; + status.transferred_bytes = 0; + return Status::OK(); + } auto& task = toBatchDesc(id).task_list[index]; const auto success = __atomic_load_n(&task.success_slice_count, __ATOMIC_ACQUIRE); @@ -86,7 +106,8 @@ class CounterTransport : public Transport { if (success + failed == task.slice_count) { status.s = failed ? TransferStatusEnum::FAILED : TransferStatusEnum::COMPLETED; - __atomic_store_n(&task.is_finished, true, __ATOMIC_RELEASE); + if (publish_finished) + __atomic_store_n(&task.is_finished, true, __ATOMIC_RELEASE); } return Status::OK(); } @@ -305,6 +326,114 @@ TEST(DeferredBatchCleanupTest, ShutdownReclaimsAlreadyPublishedErrorBatch) { EXPECT_TRUE(counts.hook_before_slice.load()); } +class BacklogLogSink final : public google::LogSink { + public: + BacklogLogSink() { google::AddLogSink(this); } + ~BacklogLogSink() override { google::RemoveLogSink(this); } + void send(google::LogSeverity, const char*, const char*, int, + const std::tm*, const char* message, size_t length) override { + if (std::string_view(message, length) + .find("Deferred cleanup backlog:") != std::string_view::npos) + warnings.fetch_add(1); + } + std::atomic warnings{0}; +}; + +TEST(DeferredBatchCleanupTest, RetainedBacklogHasLiveStatsAndPeriodicWarnings) { + BacklogLogSink logs; + std::string name = "retained-observability"; + CounterTransport backend; + CleanupCounts counts; + MultiTransport transport(nullptr, name); + auto batch = make_batch(transport, backend, counts); + ASSERT_TRUE(Peer::free(transport, batch.id, [&] { + ++counts.deleted; + }).IsBatchCleanupDeferred()); + backend.error = QueryError::Context; + ASSERT_TRUE(wait_until([&] { return backend.error_queries.load() > 0; })); + Peer::forceBacklogWarning(transport); + ASSERT_TRUE(wait_until([&] { return logs.warnings.load() >= 1; })); + auto stats = transport.getDeferredCleanupStats(); + EXPECT_EQ(stats.pending_batches, 1); + EXPECT_EQ(stats.query_error_batches, 1); + EXPECT_GE(stats.oldest_age_ms, 60000); + EXPECT_EQ(counts.deleted.load(), 0); + Peer::forceBacklogWarning(transport); + ASSERT_TRUE(wait_until([&] { return logs.warnings.load() >= 2; })); + batch.slices.front()->markFailed(); + expect_reclaimed(transport, batch, counts); + stats = transport.getDeferredCleanupStats(); + EXPECT_EQ(stats.pending_batches, 0); + EXPECT_EQ(stats.query_error_batches, 0); + EXPECT_EQ(stats.oldest_age_ms, 0); +} + +TEST(DeferredBatchCleanupTest, ObservedTerminalTaskSurvivesLaterQueryError) { + std::string name = "observed-terminal-proof"; + CounterTransport backend; + backend.publish_finished = false; + CleanupCounts counts; + MultiTransport transport(nullptr, name); + auto batch = make_batch(transport, backend, counts); + auto& task = batch.desc->task_list.front(); + // Model polling completion before the callback unwinds, with no sealed + // counter publication to substitute for the successful query proof. + task.success_slice_count = 1; + batch.desc->active_completion_callbacks = 1; + bool retired = true; + ASSERT_TRUE(transport.getBatchRetirementStatus(batch.id, retired).ok()); + EXPECT_FALSE(retired); + EXPECT_TRUE(__atomic_load_n(&task.is_finished, __ATOMIC_ACQUIRE)); + EXPECT_TRUE(__atomic_load_n(&task.retirement_observed, __ATOMIC_ACQUIRE)); + EXPECT_FALSE(__atomic_load_n(&task.completion_published, __ATOMIC_ACQUIRE)); + const int queries = backend.queries.load(); + backend.error = QueryError::Context; + ASSERT_TRUE(transport.getBatchRetirementStatus(batch.id, retired).ok()); + EXPECT_FALSE(retired); + batch.desc->active_completion_callbacks = 0; + ASSERT_TRUE(transport.getBatchRetirementStatus(batch.id, retired).ok()); + EXPECT_TRUE(retired); + EXPECT_EQ(backend.queries.load(), queries); + EXPECT_EQ(backend.error_queries.load(), 0); + EXPECT_TRUE(__atomic_load_n(&task.is_finished, __ATOMIC_ACQUIRE)); + const auto freed = + Peer::free(transport, batch.id, [&] { ++counts.deleted; }); + EXPECT_TRUE(freed.ok()); + EXPECT_EQ(backend.queries.load(), queries); + if (!freed.ok()) { + // Clean up the still-owned batch in the old-code negative control. + backend.error = QueryError::None; + EXPECT_TRUE( + Peer::free(transport, batch.id, [&] { ++counts.deleted; }).ok()); + } +} + +TEST(DeferredBatchCleanupTest, TimeoutIsNotRetirementButLaterFailureCanRetire) { + std::string name = "retirement-timeout"; + CounterTransport backend; + CleanupCounts counts; + MultiTransport transport(nullptr, name); + auto batch = make_batch(transport, backend, counts); + backend.timed_out = true; + bool retired = true; + const auto timed_out = + transport.getBatchRetirementStatus(batch.id, retired); + EXPECT_TRUE(timed_out.IsContext()); + EXPECT_NE(timed_out.ToString().find("timeout"), std::string::npos); + EXPECT_FALSE(retired); + EXPECT_FALSE(__atomic_load_n( + &batch.desc->task_list.front().retirement_observed, __ATOMIC_ACQUIRE)); + EXPECT_EQ(counts.deleted.load(), 0); + backend.timed_out = false; + batch.slices.front()->markFailed(); + const auto failed = transport.getBatchRetirementStatus(batch.id, retired); + EXPECT_TRUE(failed.IsContext()); + EXPECT_NE(failed.ToString().find("failed"), std::string::npos); + EXPECT_TRUE(retired); + EXPECT_TRUE( + Peer::free(transport, batch.id, [&] { ++counts.deleted; }).ok()); +} + TEST(DeferredBatchCleanupTest, RetirementKeepsHandleUntilAllTasksAndPinsDrain) { std::string name = "retirement-failed-sibling"; CounterTransport backend; From a9e588be00358b369ba2f734ce5fa6f9a6cd95a6 Mon Sep 17 00:00:00 2001 From: ShuweiShen772 Date: Sat, 14 Nov 2026 11:38:08 +0800 Subject: [PATCH 5/6] [Bugfix] Clarify successful deferred cleanup ownership Signed-off-by: ShuweiShen772 --- .../include/multi_transport.h | 8 ++++-- .../include/transfer_engine.h | 8 ++++-- .../tests/deferred_batch_cleanup_test.cpp | 26 +++++++++++++++++++ 3 files changed, 38 insertions(+), 4 deletions(-) diff --git a/mooncake-transfer-engine/include/multi_transport.h b/mooncake-transfer-engine/include/multi_transport.h index 851407472..a64481c9d 100644 --- a/mooncake-transfer-engine/include/multi_transport.h +++ b/mooncake-transfer-engine/include/multi_transport.h @@ -56,8 +56,12 @@ class MultiTransport { BatchID allocateBatchID(size_t batch_size); - // An OK or BatchCleanupDeferred return invalidates batch_id. A deferred - // cleanup means in-flight transport work may still own transfer buffers. + // An OK or BatchCleanupDeferred return invalidates batch_id. Deferred + // transfers batch cleanup responsibility, not ownership of user buffers. + // Cleanup status is independent of transfer success: do not turn an + // otherwise successful transfer into a failure solely due to deferral. + // In-flight work may still access borrowed buffers; keep them alive until + // retirement is independently established, even after handle handoff. Status freeBatchID(BatchID batch_id); Transport::DeferredCleanupStats getDeferredCleanupStats(); diff --git a/mooncake-transfer-engine/include/transfer_engine.h b/mooncake-transfer-engine/include/transfer_engine.h index ef0520e0b..085b7c6b3 100644 --- a/mooncake-transfer-engine/include/transfer_engine.h +++ b/mooncake-transfer-engine/include/transfer_engine.h @@ -268,8 +268,12 @@ class TransferEngine { BatchID allocateBatchID(size_t batch_size); - // An OK or BatchCleanupDeferred return invalidates batch_id. A deferred - // cleanup means in-flight transport work may still own transfer buffers. + // An OK or BatchCleanupDeferred return invalidates batch_id. Deferred + // transfers batch cleanup responsibility, not ownership of user buffers. + // Cleanup status is independent of transfer success: do not turn an + // otherwise successful transfer into a failure solely due to deferral. + // In-flight work may still access borrowed buffers; keep them alive until + // retirement is independently established, even after handle handoff. Status freeBatchID(BatchID batch_id); int getNotifies(std::vector& notifies); diff --git a/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp b/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp index a5957313a..9d775eafe 100644 --- a/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp +++ b/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp @@ -285,6 +285,32 @@ TEST(DeferredBatchCleanupTest, WaitsForEveryTaskAfterQueryError) { expect_reclaimed(transport, batch, counts); } +TEST(DeferredBatchCleanupTest, CompletedTransferCanDeferCleanupWithoutFailure) { + std::string name = "completed-cleanup-deferred"; + CounterTransport backend; + CleanupCounts counts; + MultiTransport transport(nullptr, name); + auto batch = make_batch(transport, backend, counts); + // Model a callback that published success but is still accessing state. + batch.desc->active_completion_callbacks.fetch_add(1); + batch.slices.front()->markSuccess(); + Transport::TransferStatus status{}; + ASSERT_TRUE(transport.getBatchTransferStatus(batch.id, status).ok()); + EXPECT_EQ(status.s, Transport::TransferStatusEnum::COMPLETED); + EXPECT_FALSE(batch.desc->has_failure.load()); + bool retired = true; + EXPECT_TRUE(transport.getBatchRetirementStatus(batch.id, retired).ok()); + EXPECT_FALSE(retired); + EXPECT_TRUE(Peer::free(transport, batch.id, [&] { + ++counts.deleted; + }).IsBatchCleanupDeferred()); + EXPECT_TRUE(Peer::pending(transport, batch.id)); + EXPECT_EQ(counts.deleted.load(), 0); + // Use only the callback's pinned descriptor after the public ID expires. + batch.desc->active_completion_callbacks.fetch_sub(1); + expect_reclaimed(transport, batch, counts); +} + TEST(DeferredBatchCleanupTest, PublishedCompletionDoesNotBypassCallbackPin) { std::string name = "completion-pin"; CounterTransport backend; From 69f6b0efdbb99da12e04f42349c3ea1582e39810 Mon Sep 17 00:00:00 2001 From: ShuweiShen772 Date: Sat, 14 Nov 2026 14:01:49 +0800 Subject: [PATCH 6/6] [Bugfix] Preserve terminal failure and retirement status Signed-off-by: ShuweiShen772 --- .../include/multi_transport.h | 2 + .../include/transfer_engine.h | 3 +- .../src/multi_transport.cpp | 35 +++++++++++---- .../tests/deferred_batch_cleanup_test.cpp | 44 +++++++++++++++++++ 4 files changed, 74 insertions(+), 10 deletions(-) diff --git a/mooncake-transfer-engine/include/multi_transport.h b/mooncake-transfer-engine/include/multi_transport.h index a64481c9d..902869dad 100644 --- a/mooncake-transfer-engine/include/multi_transport.h +++ b/mooncake-transfer-engine/include/multi_transport.h @@ -62,6 +62,8 @@ class MultiTransport { // otherwise successful transfer into a failure solely due to deferral. // In-flight work may still access borrowed buffers; keep them alive until // retirement is independently established, even after handle handoff. + // A pre-delete hook (when installed by TransferEngineImpl) releases + // descriptor side state; its invocation is not a buffer-retirement signal. Status freeBatchID(BatchID batch_id); Transport::DeferredCleanupStats getDeferredCleanupStats(); diff --git a/mooncake-transfer-engine/include/transfer_engine.h b/mooncake-transfer-engine/include/transfer_engine.h index 085b7c6b3..15a8deb6a 100644 --- a/mooncake-transfer-engine/include/transfer_engine.h +++ b/mooncake-transfer-engine/include/transfer_engine.h @@ -298,7 +298,8 @@ class TransferEngine { // is true. All tasks and completion callbacks must have stopped accessing // borrowed requests/buffers. This does not free or invalidate batch_id. // A non-OK result reports an observed query/transfer failure independently - // of retirement. TIMEOUT/error alone never authorizes buffer reuse. + // of retirement. TIMEOUT, CANCELED, INVALID, or another error alone never + // authorizes buffer reuse; rely on retired for the independent proof. // Do not submit more work, free, or query this batch concurrently. As with // transfer completion, transports must publish terminal state after I/O // retirement (including failure paths). TIMEOUT is not terminal: tasks may diff --git a/mooncake-transfer-engine/src/multi_transport.cpp b/mooncake-transfer-engine/src/multi_transport.cpp index 73b85b098..2c4a72622 100644 --- a/mooncake-transfer-engine/src/multi_transport.cpp +++ b/mooncake-transfer-engine/src/multi_transport.cpp @@ -131,9 +131,10 @@ Transport::DeferredCleanupStats MultiTransport::collectDeferredCleanupStats( const auto age = std::chrono::duration_cast( now - entry.enqueued_at) .count(); - if (age > 0) + if (age > 0) { stats.oldest_age_ms = std::max(stats.oldest_age_ms, static_cast(age)); + } } return stats; } @@ -251,6 +252,7 @@ Status MultiTransport::tryFreeBatchID( void MultiTransport::deferredCleanupLoop() { constexpr auto kCleanupInterval = std::chrono::milliseconds(100); + constexpr auto kDeferredCleanupWarningInterval = std::chrono::seconds(60); std::unique_lock lock(deferred_cleanup_mutex_); while (true) { if (deferred_cleanup_batches_.empty()) { @@ -303,9 +305,11 @@ void MultiTransport::deferredCleanupLoop() { const auto warning_now = std::chrono::steady_clock::now(); if (warning_now >= next_deferred_cleanup_warning_) { next_deferred_cleanup_warning_ = - warning_now + std::chrono::seconds(60); + warning_now + kDeferredCleanupWarningInterval; const auto stats = collectDeferredCleanupStats(warning_now); - if (stats.oldest_age_ms >= 60000) { + const auto oldest_age = + std::chrono::milliseconds(stats.oldest_age_ms); + if (oldest_age >= kDeferredCleanupWarningInterval) { LOG(WARNING) << "Deferred cleanup backlog: pending=" << stats.pending_batches << ", query_errors=" << stats.query_error_batches @@ -813,15 +817,28 @@ Status MultiTransport::getBatchRetirementStatus(BatchID batch_id, } else { done = status.s == TransferStatusEnum::COMPLETED || status.s == TransferStatusEnum::FAILED; - if (status.s == TransferStatusEnum::FAILED) { + if (status.s == TransferStatusEnum::FAILED || + status.s == TransferStatusEnum::CANCELED || + status.s == TransferStatusEnum::INVALID) { batch.has_failure.store(true, std::memory_order_release); } - if (result.ok() && status.s == TransferStatusEnum::TIMEOUT) { - result = Status::Context("Mooncake transfer task timeout"); - } else if (result.ok() && - status.s == TransferStatusEnum::FAILED) { - result = Status::Context("Mooncake transfer task failed"); + if (result.ok()) { + if (status.s == TransferStatusEnum::TIMEOUT) { + result = + Status::Context("Mooncake transfer task timeout"); + } else if (status.s == TransferStatusEnum::FAILED) { + result = + Status::Context("Mooncake transfer task failed"); + } else if (status.s == TransferStatusEnum::CANCELED) { + result = + Status::Context("Mooncake transfer task canceled"); + } else if (status.s == TransferStatusEnum::INVALID) { + result = + Status::Context("Mooncake transfer task invalid"); + } } + // INVALID/CANCELED expose failure without proving physical + // I/O retirement. Keep the batch pending for later proof. } } if (__atomic_load_n(&task.failed_slice_count, __ATOMIC_ACQUIRE) && diff --git a/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp b/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp index 9d775eafe..35c878a9e 100644 --- a/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp +++ b/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp @@ -63,6 +63,8 @@ class CounterTransport : public Transport { std::atomic error_queries{0}; std::atomic queries{0}; std::atomic timed_out{false}; + std::atomic override_terminal_status{false}; + std::atomic terminal_status{TransferStatusEnum::FAILED}; bool counter_completion = false; bool publish_finished = true; bool supportsBatchCompletionCounter() const override { @@ -96,6 +98,11 @@ class CounterTransport : public Transport { status.transferred_bytes = 0; return Status::OK(); } + if (override_terminal_status.load()) { + status.s = terminal_status.load(); + status.transferred_bytes = 0; + return Status::OK(); + } auto& task = toBatchDesc(id).task_list[index]; const auto success = __atomic_load_n(&task.success_slice_count, __ATOMIC_ACQUIRE); @@ -460,6 +467,43 @@ TEST(DeferredBatchCleanupTest, TimeoutIsNotRetirementButLaterFailureCanRetire) { Peer::free(transport, batch.id, [&] { ++counts.deleted; }).ok()); } +TEST(DeferredBatchCleanupTest, InvalidAndCanceledAreErrorsUntilRetired) { + for (const auto terminal_status : + {Transport::TransferStatusEnum::INVALID, + Transport::TransferStatusEnum::CANCELED}) { + std::string name = + terminal_status == Transport::TransferStatusEnum::INVALID + ? "retirement-invalid" + : "retirement-canceled"; + CounterTransport backend; + backend.override_terminal_status = true; + backend.terminal_status = terminal_status; + CleanupCounts counts; + MultiTransport transport(nullptr, name); + auto batch = make_batch(transport, backend, counts); + + bool retired = true; + const Status failure = + transport.getBatchRetirementStatus(batch.id, retired); + EXPECT_TRUE(failure.IsContext()); + EXPECT_FALSE(retired); + EXPECT_EQ(counts.deleted.load(), 0); + + // A later completion publication supplies the retirement proof; retain + // the original transport failure for the caller. + backend.override_terminal_status = false; + batch.slices.front()->markFailed(); + const Status completed = + transport.getBatchRetirementStatus(batch.id, retired); + EXPECT_TRUE(completed.IsContext()); + EXPECT_NE(completed.ToString().find("failed"), std::string::npos); + EXPECT_TRUE(retired); + EXPECT_TRUE( + Peer::free(transport, batch.id, [&] { ++counts.deleted; }).ok()); + EXPECT_EQ(counts.deleted.load(), 1); + } +} + TEST(DeferredBatchCleanupTest, RetirementKeepsHandleUntilAllTasksAndPinsDrain) { std::string name = "retirement-failed-sibling"; CounterTransport backend;