diff --git a/mooncake-transfer-engine/include/multi_transport.h b/mooncake-transfer-engine/include/multi_transport.h index e00581ec..902869da 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 @@ -55,10 +56,18 @@ 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. + // 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(); + Status submitTransfer(BatchID batch_id, const std::vector &entries); @@ -80,6 +89,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); @@ -134,8 +148,16 @@ 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); + const std::function &before_delete, + bool allow_completion_fallback = false, + bool poll_transport = true); void deferredCleanupLoop(); @@ -154,8 +176,20 @@ 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; + 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 8767efac..15a8deb6 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); @@ -289,6 +293,27 @@ 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, 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 + // 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; 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 23438799..14fbbd32 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 @@ -342,6 +346,24 @@ 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 (!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); @@ -527,6 +549,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 40ab5787..b356abea 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; @@ -522,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. @@ -546,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 189d1dc6..2c4a7262 100644 --- a/mooncake-transfer-engine/src/multi_transport.cpp +++ b/mooncake-transfer-engine/src/multi_transport.cpp @@ -122,6 +122,28 @@ 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()); } @@ -155,7 +177,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 +195,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()) { @@ -215,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()) { @@ -230,17 +268,52 @@ 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; + } + + const auto warning_now = std::chrono::steady_clock::now(); + if (warning_now >= next_deferred_cleanup_warning_) { + next_deferred_cleanup_warning_ = + warning_now + kDeferredCleanupWarningInterval; + const auto stats = collectDeferredCleanupStats(warning_now); + 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 + << ", oldest_age_ms=" << stats.oldest_age_ms; } } @@ -554,6 +627,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(); } @@ -701,6 +780,93 @@ 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 (status.s == TransferStatusEnum::FAILED || + status.s == TransferStatusEnum::CANCELED || + status.s == TransferStatusEnum::INVALID) { + batch.has_failure.store(true, std::memory_order_release); + } + 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) && + 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_release); + 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 4a016afb..1f1a0257 100644 --- a/mooncake-transfer-engine/src/transfer_engine.cpp +++ b/mooncake-transfer-engine/src/transfer_engine.cpp @@ -303,6 +303,15 @@ 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); +} + +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/CMakeLists.txt b/mooncake-transfer-engine/tests/CMakeLists.txt index 7986553a..1678aded 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 00000000..35c878a9 --- /dev/null +++ b/mooncake-transfer-engine/tests/deferred_batch_cleanup_test.cpp @@ -0,0 +1,587 @@ +// 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 +#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 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; + } +}; + +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}; + 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 { + return counter_completion; + } + + Status submitTransfer(BatchID, + const std::vector&) override { + return Status::OK(); + } + + 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); + 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; + } + } + if (timed_out.load()) { + status.s = TransferStatusEnum::TIMEOUT; + 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); + 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; + if (publish_finished) + __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, 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; + 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()); +} + +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, 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; + 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 diff --git a/mooncake-transfer-engine/tests/transport_uint_test.cpp b/mooncake-transfer-engine/tests/transport_uint_test.cpp index bd75781d..ed5d1f5c 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)