mooncake-track
/
Mooncake
0
0
Fork 12
Code Issues Pull Requests 13 Actions Packages Projects Releases Wiki Activity

[CCF Archive] Store object type eviction policy submission #3

Closed
kancel wants to merge 382 commits from kancel:ccf-archive-pr2746 into main
pull from: kancel:ccf-archive-pr2746
merge into: mooncake-track:main
Conversation 0 Commits 382 Files Changed 654 +1069 -53
11 changed files with 1069 additions and 53 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit d6b661d5a4 - Show all commits

1
mooncake-transfer-engine/tent/config/transfer-engine.json
View File

@ -12,6 +12,7 @@
"log_level": "warning",
"max_failover_attempts": 3,
"enable_auto_failover_on_poll": true,
"enable_progress_worker": false,
"metrics": {
"enabled": true,
"http_port": 9100,

74
mooncake-transfer-engine/tent/include/tent/runtime/progress_worker.h Normal file
View File

@ -0,0 +1,74 @@
// Copyright 2024 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.
#ifndef PROGRESS_WORKER_H_
#define PROGRESS_WORKER_H_
#include <atomic>
#include <condition_variable>
#include <deque>
#include <mutex>
#include <thread>
#include <unordered_set>
#include "tent/common/types.h"
namespace mooncake {
namespace tent {
class TransferEngineImpl;
// Event-driven progress worker for issue #2116. When the engine is configured
// with enable_progress_worker=true, transports (or test hooks) call
// notifyBatchMaybeReady to wake this worker, which then drives one
// progressBatch step per notification. This decouples failover/resubmit from
// the caller polling loop, so integrators that turn off
// enable_auto_failover_on_poll do not need to spin a polling thread of their
// own to keep failover progressing.
class ProgressWorker {
public:
explicit ProgressWorker(TransferEngineImpl* impl);
~ProgressWorker();
ProgressWorker(const ProgressWorker&) = delete;
ProgressWorker& operator=(const ProgressWorker&) = delete;
void start();
// Idempotent. Signals the worker thread to exit and joins it. After stop()
// returns, notifyBatchMaybeReady becomes a no-op.
void stop();
// Safe from any thread. De-duplicates: enqueueing a batch that is already
// queued is a no-op. No-op if the worker has been stopped or never
// started.
void notifyBatchMaybeReady(BatchID batch_id);
private:
void runner();
TransferEngineImpl* impl_;
std::atomic<bool> running_{false};
std::thread thread_;
std::mutex mu_;
std::condition_variable cv_;
std::unordered_set<BatchID> queued_;
std::deque<BatchID> order_;
};
} // namespace tent
} // namespace mooncake
#endif // PROGRESS_WORKER_H_

28
mooncake-transfer-engine/tent/include/tent/runtime/transfer_engine_impl.h
View File

@ -46,6 +46,7 @@ class ControlService;
class SegmentTracker;
class Platform;
class ProxyManager;
class ProgressWorker;
struct TaskInfo {
TransportType type{UNSPEC};
@ -152,6 +153,8 @@ class TransferEngineImpl {
Status getTransferStatus(BatchID batch_id, TransferStatus& overall_status);
Status progressBatch(BatchID batch_id, TransferStatus& overall_status);
Status waitTransferCompletion(BatchID batch_id);
Status transferSync(const std::vector<Request>& request_list);
@ -172,6 +175,11 @@ class TransferEngineImpl {
}
}
// Wake the optional event-driven progress worker for `batch_id`. No-op if
// enable_progress_worker is false. Currently used by test/integration
// hooks; transports will be migrated to call this in a follow-up PR.
void notifyBatchMaybeReady(BatchID batch_id);
private:
Status construct();
@ -189,8 +197,15 @@ class TransferEngineImpl {
Status resubmitTransferTask(Batch* batch, size_t task_id);
void updateTaskStatusFromPoll(Batch* batch, size_t task_id,
TransferStatus& task_status);
Status pollTaskStatus(Batch* batch, size_t task_id,
TransferStatus& task_status);
void updateTaskStatusAfterPoll(Batch* batch, size_t task_id,
TransferStatus& task_status,
bool allow_failover);
Status getBatchStatus(BatchID batch_id, TransferStatus& overall_status,
bool allow_failover);
SelectionResult resolveTransport(const Request& req, int transport_index,
bool invalidate_on_fail = true);
@ -244,6 +259,15 @@ class TransferEngineImpl {
bool merge_requests_;
int max_failover_attempts_{3};
bool enable_auto_failover_on_poll_{true};
bool enable_progress_worker_{false};
// Guards alive_batches_ and serializes pollTaskStatus /
// updateTaskStatusAfterPoll / lazyFreeBatch against the optional
// ProgressWorker thread. Recursive because freeBatch -> lazyFreeBatch ->
// getTransferStatus can re-enter on the same thread. See issue #2116.
std::recursive_mutex progress_mutex_;
std::unordered_set<BatchID> alive_batches_;
std::unique_ptr<ProgressWorker> progress_worker_;
};
} // namespace tent
} // namespace mooncake

8
mooncake-transfer-engine/tent/include/tent/transfer_engine.h
View File

@ -308,6 +308,14 @@ class TransferEngine {
Status getTransferStatus(BatchID batch_id, TransferStatus& overall_status);
// Drive one progress step on a batch and return its aggregated status.
// Unlike getTransferStatus, this always allows internal failover/resubmit
// regardless of enable_auto_failover_on_poll. The call is non-blocking and
// performs at most one state-machine step per task; callers that want to
// wait for completion must invoke it in a loop. PENDING means "make
// progress later"; terminal states (COMPLETED/FAILED) will not be revived.
Status progressBatch(BatchID batch_id, TransferStatus& overall_status);
private:
std::unique_ptr<TransferEngineImpl> impl_;
};

81
mooncake-transfer-engine/tent/src/runtime/progress_worker.cpp Normal file
View File

@ -0,0 +1,81 @@
// Copyright 2024 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 "tent/runtime/progress_worker.h"
#include "tent/common/status.h"
#include "tent/runtime/transfer_engine_impl.h"
namespace mooncake {
namespace tent {
ProgressWorker::ProgressWorker(TransferEngineImpl* impl) : impl_(impl) {}
ProgressWorker::~ProgressWorker() { stop(); }
void ProgressWorker::start() {
if (running_.exchange(true, std::memory_order_acq_rel)) return;
thread_ = std::thread(&ProgressWorker::runner, this);
}
void ProgressWorker::stop() {
if (!running_.exchange(false, std::memory_order_acq_rel)) return;
{
std::lock_guard<std::mutex> lk(mu_);
// Drop pending work; outstanding batches will be reaped via the
// user thread's freeBatch path.
order_.clear();
queued_.clear();
}
cv_.notify_all();
if (thread_.joinable()) thread_.join();
}
void ProgressWorker::notifyBatchMaybeReady(BatchID batch_id) {
if (!batch_id) return;
if (!running_.load(std::memory_order_acquire)) return;
{
std::lock_guard<std::mutex> lk(mu_);
if (!queued_.insert(batch_id).second) return;
order_.push_back(batch_id);
}
cv_.notify_one();
}
void ProgressWorker::runner() {
while (true) {
BatchID batch_id = 0;
{
std::unique_lock<std::mutex> lk(mu_);
cv_.wait(lk, [&] {
return !running_.load(std::memory_order_acquire) ||
!order_.empty();
});
if (!running_.load(std::memory_order_acquire)) return;
batch_id = order_.front();
order_.pop_front();
queued_.erase(batch_id);
}
// progressBatch acquires the engine's progress_mutex_ and silently
// returns InvalidArgument if the batch was freed before we got here.
// PENDING means "kick again later"; the next notify wakes us up.
// Terminal states leave the batch alone — freeBatch on the user
// thread is responsible for reclamation.
TransferStatus s;
(void)impl_->progressBatch(batch_id, s);
}
}
} // namespace tent
} // namespace mooncake

3
mooncake-transfer-engine/tent/src/runtime/proxy_manager.cpp
View File

@ -385,8 +385,7 @@ Status ProxyManager::transferEventLoop(StagingTask& task,
case StageState::INFLIGHT: {
TransferStatus xfer_status;
CHECK_STATUS(
impl_->getTransferStatus(chunk.batch, xfer_status));
CHECK_STATUS(impl_->progressBatch(chunk.batch, xfer_status));
if (xfer_status.s == PENDING) {
event_queue.push(id);
break;

130
mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp
View File

@ -30,6 +30,7 @@
#include "tent/runtime/control_plane.h"
#include "tent/runtime/segment.h"
#include "tent/runtime/segment_tracker.h"
#include "tent/runtime/progress_worker.h"
#include "tent/runtime/proxy_manager.h"
#include "tent/runtime/transport.h"
#include "tent/runtime/topology.h"
@ -306,6 +307,7 @@ Status TransferEngineImpl::construct() {
max_failover_attempts_ = conf_->get("max_failover_attempts", 3);
enable_auto_failover_on_poll_ =
conf_->get("enable_auto_failover_on_poll", true);
enable_progress_worker_ = conf_->get("enable_progress_worker", false);
if (!hostname_.empty())
CHECK_STATUS(checkLocalIpAddress(hostname_, ipv6_));
else
@ -370,6 +372,11 @@ Status TransferEngineImpl::construct() {
staging_proxy_ = std::make_unique<ProxyManager>(this);
if (enable_progress_worker_) {
progress_worker_ = std::make_unique<ProgressWorker>(this);
progress_worker_->start();
}
// Initialize and start Metrics system
auto metrics_config = MetricsConfigLoader::loadWithDefaults(conf_.get());
if (metrics_config.enabled) {
@ -412,6 +419,13 @@ Status TransferEngineImpl::construct() {
Status TransferEngineImpl::deconstruct() {
// Metrics cleanup is handled automatically by TentMetrics destructor
// Stop the progress worker first so it cannot race with batch teardown
// below (it dereferences BatchID into Batch* via progressBatch).
if (progress_worker_) {
progress_worker_->stop();
progress_worker_.reset();
}
// Destroy staging_proxy_ first: its destructor calls back into
// unregisterLocalMemory/freeLocalMemory, which require
// local_segment_tracker_ and metadata_ to be alive.
@ -750,18 +764,25 @@ BatchID TransferEngineImpl::allocateBatch(size_t batch_size) {
if (!batch) return (BatchID)0;
batch->max_size = batch_size;
batch_set_.get().active.insert(batch);
return (BatchID)batch;
BatchID batch_id = (BatchID)batch;
{
std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
alive_batches_.insert(batch_id);
}
return batch_id;
}
Status TransferEngineImpl::freeBatch(BatchID batch_id) {
if (!batch_id) return Status::InvalidArgument("Invalid batch ID" LOC_MARK);
Batch* batch = (Batch*)(batch_id);
std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
batch_set_.get().freelist.push_back(batch);
lazyFreeBatch();
return Status::OK();
}
Status TransferEngineImpl::lazyFreeBatch() {
// Caller must hold progress_mutex_.
auto& batch_set = batch_set_.get();
for (auto it = batch_set.freelist.begin();
it != batch_set.freelist.end();) {
@ -778,6 +799,7 @@ Status TransferEngineImpl::lazyFreeBatch() {
if (transport && sub_batch) transport->freeSubBatch(sub_batch);
}
batch_set.active.erase(batch);
alive_batches_.erase((BatchID)batch);
Slab<Batch>::Get().deallocate(batch);
it = batch_set.freelist.erase(it);
}
@ -1379,11 +1401,35 @@ Status TransferEngineImpl::resubmitTransferTask(Batch* batch, size_t task_id) {
return transport->submitTransferTasks(sub_batch, {task.request});
}
void TransferEngineImpl::updateTaskStatusFromPoll(Batch* batch, size_t task_id,
TransferStatus& task_status) {
Status TransferEngineImpl::pollTaskStatus(Batch* batch, size_t task_id,
TransferStatus& task_status) {
auto& task = batch->task_list[task_id];
if (task.staging) {
return staging_proxy_->getStatus(&task, task_status);
}
if (task.type == UNSPEC) {
task_status.s = FAILED;
task_status.transferred_bytes = 0;
return Status::OK();
}
auto& transport = transport_list_[task.type];
auto& sub_batch = batch->sub_batch[task.type];
if (!transport || !sub_batch) {
return Status::InvalidArgument("Transport not available" LOC_MARK);
}
return transport->getTransferStatus(sub_batch, task.sub_task_id,
task_status);
}
void TransferEngineImpl::updateTaskStatusAfterPoll(Batch* batch, size_t task_id,
TransferStatus& task_status,
bool allow_failover) {
auto& task = batch->task_list[task_id];
task.status = task_status.s;
if (!enable_auto_failover_on_poll_ || task_status.s != FAILED) return;
if (!allow_failover || task_status.s != FAILED || task.type == UNSPEC)
return;
if (resubmitTransferTask(batch, task_id).ok()) {
task_status.s = PENDING;
@ -1433,29 +1479,17 @@ Status TransferEngineImpl::receiveNotification(
Status TransferEngineImpl::getTransferStatus(BatchID batch_id, size_t task_id,
TransferStatus& task_status) {
if (!batch_id) return Status::InvalidArgument("Invalid batch ID" LOC_MARK);
std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
if (!alive_batches_.count(batch_id))
return Status::InvalidArgument("Batch is not alive" LOC_MARK);
Batch* batch = (Batch*)(batch_id);
if (task_id >= batch->task_list.size())
return Status::InvalidArgument("Invalid task ID" LOC_MARK);
auto& task = batch->task_list[task_id];
auto prev_status = task.status;
if (task.staging) {
CHECK_STATUS(staging_proxy_->getStatus(&task, task_status));
} else {
if (task.type == UNSPEC) {
task_status.s = FAILED;
task_status.transferred_bytes = 0;
batch->task_list[task_id].status = task_status.s;
return Status::OK();
}
auto& transport = transport_list_[task.type];
auto& sub_batch = batch->sub_batch[task.type];
if (!transport || !sub_batch) {
return Status::InvalidArgument("Transport not available" LOC_MARK);
}
CHECK_STATUS(transport->getTransferStatus(sub_batch, task.sub_task_id,
task_status));
}
updateTaskStatusFromPoll(batch, task_id, task_status);
CHECK_STATUS(pollTaskStatus(batch, task_id, task_status));
updateTaskStatusAfterPoll(batch, task_id, task_status,
enable_auto_failover_on_poll_);
// Record metrics when task transitions to terminal state
recordTaskCompletionMetrics(batch->task_list[task_id], prev_status,
@ -1468,6 +1502,9 @@ Status TransferEngineImpl::getTransferStatus(BatchID batch_id, size_t task_id,
Status TransferEngineImpl::getTransferStatus(
BatchID batch_id, std::vector<TransferStatus>& status_list) {
if (!batch_id) return Status::InvalidArgument("Invalid batch ID" LOC_MARK);
std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
if (!alive_batches_.count(batch_id))
return Status::InvalidArgument("Batch is not alive" LOC_MARK);
Batch* batch = (Batch*)(batch_id);
status_list.clear();
for (size_t task_id = 0; task_id < batch->task_list.size(); ++task_id) {
@ -1478,9 +1515,13 @@ Status TransferEngineImpl::getTransferStatus(
return Status::OK();
}
Status TransferEngineImpl::getTransferStatus(BatchID batch_id,
TransferStatus& overall_status) {
Status TransferEngineImpl::getBatchStatus(BatchID batch_id,
TransferStatus& overall_status,
bool allow_failover) {
if (!batch_id) return Status::InvalidArgument("Invalid batch ID" LOC_MARK);
std::lock_guard<std::recursive_mutex> lk(progress_mutex_);
if (!alive_batches_.count(batch_id))
return Status::InvalidArgument("Batch is not alive" LOC_MARK);
Batch* batch = (Batch*)(batch_id);
overall_status.s = PENDING;
overall_status.transferred_bytes = 0;
@ -1512,26 +1553,8 @@ Status TransferEngineImpl::getTransferStatus(BatchID batch_id,
continue;
}
auto prev_status = task.status;
if (task.staging) {
CHECK_STATUS(staging_proxy_->getStatus(&task, task_status));
} else {
if (task.type == UNSPEC) {
task.status = FAILED;
failed_tasks++;
if (isWorse(FAILED, worst_failure)) worst_failure = FAILED;
continue;
}
auto& transport = transport_list_[task.type];
auto& sub_batch = batch->sub_batch[task.type];
if (!transport || !sub_batch) {
return Status::InvalidArgument(
"Transport not available" LOC_MARK);
}
CHECK_STATUS(transport->getTransferStatus(
sub_batch, task.sub_task_id, task_status));
}
// Preserve legacy auto-failover-on-poll before aggregating status.
updateTaskStatusFromPoll(batch, task_id, task_status);
CHECK_STATUS(pollTaskStatus(batch, task_id, task_status));
updateTaskStatusAfterPoll(batch, task_id, task_status, allow_failover);
if (task_status.s == COMPLETED) {
success_tasks++;
@ -1559,10 +1582,25 @@ Status TransferEngineImpl::getTransferStatus(BatchID batch_id,
return Status::OK();
}
Status TransferEngineImpl::getTransferStatus(BatchID batch_id,
TransferStatus& overall_status) {
return getBatchStatus(batch_id, overall_status,
enable_auto_failover_on_poll_);
}
Status TransferEngineImpl::progressBatch(BatchID batch_id,
TransferStatus& overall_status) {
return getBatchStatus(batch_id, overall_status, true);
}
void TransferEngineImpl::notifyBatchMaybeReady(BatchID batch_id) {
if (progress_worker_) progress_worker_->notifyBatchMaybeReady(batch_id);
}
Status TransferEngineImpl::waitTransferCompletion(BatchID batch_id) {
TransferStatus xfer_status;
while (true) {
CHECK_STATUS(getTransferStatus(batch_id, xfer_status));
CHECK_STATUS(progressBatch(batch_id, xfer_status));
if (xfer_status.s != PENDING) {
freeBatch(batch_id);
return xfer_status.s == COMPLETED
@ -1580,7 +1618,7 @@ Status TransferEngineImpl::transferSync(
CHECK_STATUS(submitTransfer(batch_id, request_list));
while (true) {
TransferStatus xfer_status;
CHECK_STATUS(getTransferStatus(batch_id, xfer_status));
CHECK_STATUS(progressBatch(batch_id, xfer_status));
if (xfer_status.s == COMPLETED) break;
if (xfer_status.s != PENDING) {
CHECK_STATUS(freeBatch(batch_id));

5
mooncake-transfer-engine/tent/src/transfer_engine.cpp
View File

@ -169,5 +169,10 @@ Status TransferEngine::getTransferStatus(BatchID batch_id,
return impl_->getTransferStatus(batch_id, overall_status);
}
Status TransferEngine::progressBatch(BatchID batch_id,
TransferStatus& overall_status) {
return impl_->progressBatch(batch_id, overall_status);
}
} // namespace tent
} // namespace mooncake

13
mooncake-transfer-engine/tent/tests/CMakeLists.txt
View File

@ -116,3 +116,16 @@ target_include_directories(tent_engine_failover_e2e_test
PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include)
add_test(NAME tent_engine_failover_e2e_test
COMMAND tent_engine_failover_e2e_test)
# ProgressWorker skeleton test: covers default-off behavior, event-driven
# progress without poll-failover, and freeBatch races (issue #2116).
add_executable(tent_progress_worker_test progress_worker_test.cpp)
target_link_libraries(tent_progress_worker_test
PRIVATE gtest gtest_main tent_link_group)
if(TARGET asio_shared)
target_link_libraries(tent_progress_worker_test PRIVATE asio_shared)
endif()
target_include_directories(tent_progress_worker_test
PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../include)
add_test(NAME tent_progress_worker_test
COMMAND tent_progress_worker_test)

292
mooncake-transfer-engine/tent/tests/engine_failover_e2e_test.cpp
View File

@ -65,12 +65,23 @@ class FakeSubBatch : public Transport::SubBatch {
public:
size_t size() const override { return task_count; }
size_t task_count = 0;
std::vector<Request> requests;
std::vector<TransferStatus> statuses;
std::vector<int> poll_counts;
};
class FakeTransport : public Transport {
public:
explicit FakeTransport(TransportType self_type) : self_type_(self_type) {
using StatusFactory = std::function<TransferStatus(const Request&)>;
using PollStatusFactory =
std::function<TransferStatus(const Request&, int)>;
explicit FakeTransport(TransportType self_type,
StatusFactory status_factory = {},
PollStatusFactory poll_status_factory = {})
: self_type_(self_type),
status_factory_(std::move(status_factory)),
poll_status_factory_(std::move(poll_status_factory)) {
caps.dram_to_dram = true; // so checkAvailability returns true
}
@ -103,7 +114,14 @@ class FakeTransport : public Transport {
++submit_calls;
auto* fb = static_cast<FakeSubBatch*>(batch);
for (const auto& req : request_list) {
fb->statuses.push_back({TransferStatusEnum::COMPLETED, req.length});
if (status_factory_) {
fb->statuses.push_back(status_factory_(req));
} else {
fb->statuses.push_back(
{TransferStatusEnum::COMPLETED, req.length});
}
fb->requests.push_back(req);
fb->poll_counts.push_back(0);
fb->task_count++;
}
return Status::OK();
@ -116,7 +134,13 @@ class FakeTransport : public Transport {
if (task_id < 0 || task_id >= (int)fb->statuses.size()) {
return Status::InvalidArgument("bad task_id" LOC_MARK);
}
status = fb->statuses[task_id];
++fb->poll_counts[task_id];
if (poll_status_factory_) {
status = poll_status_factory_(fb->requests[task_id],
fb->poll_counts[task_id]);
} else {
status = fb->statuses[task_id];
}
return Status::OK();
}
@ -163,6 +187,8 @@ class FakeTransport : public Transport {
private:
TransportType self_type_;
StatusFactory status_factory_;
PollStatusFactory poll_status_factory_;
};
// ---------------------------------------------------------------------------
@ -373,6 +399,266 @@ TEST(EngineFailoverE2E, AutoFailoverOnPollDisabledAppliesToOverallStatus) {
engine.unregisterLocalMemory(batch.buf.data(), batch.buf.size()).ok());
}
// ---------------------------------------------------------------------------
// P0c: Explicit progressBatch() drives one progress step and always allows
// failover/resubmit, regardless of enable_auto_failover_on_poll. Internal
// sync paths (waitTransferCompletion, transferSync) and the proxy event loop
// are wired through it so observation-only callers stay decoupled from
// progress-driving callers.
// ---------------------------------------------------------------------------
TEST(EngineFailoverE2E, ProgressBatchRetriesWhenPollAutoFailoverDisabled) {
auto cfg = makeMinimalP2PConfig();
cfg->set("enable_auto_failover_on_poll", false);
TransferEngineImpl engine(cfg);
ASSERT_TRUE(engine.available());
CorruptedRdmaBatch batch;
submitCorruptedRdmaBatch(engine, batch, 0xA4);
TransferStatus overall_status{};
ASSERT_TRUE(engine.progressBatch(batch.batch_id, overall_status).ok());
EXPECT_EQ(overall_status.s, TransferStatusEnum::PENDING);
EXPECT_EQ(batch.fake_rdma->submit_calls.load(), 1);
EXPECT_EQ(batch.fake_tcp->submit_calls.load(), 1);
EXPECT_EQ(batch.fake_tcp->status_calls.load(), 0)
<< "progressBatch should perform one progress step, not poll the "
"fallback submission immediately";
ASSERT_TRUE(engine.progressBatch(batch.batch_id, overall_status).ok());
EXPECT_EQ(overall_status.s, TransferStatusEnum::COMPLETED);
EXPECT_EQ(batch.fake_tcp->status_calls.load(), 1);
EXPECT_TRUE(engine.freeBatch(batch.batch_id).ok());
EXPECT_TRUE(
engine.unregisterLocalMemory(batch.buf.data(), batch.buf.size()).ok());
}
TEST(EngineFailoverE2E, ProgressBatchDoesNotReviveObservedFailedTask) {
auto cfg = makeMinimalP2PConfig();
cfg->set("enable_auto_failover_on_poll", false);
TransferEngineImpl engine(cfg);
ASSERT_TRUE(engine.available());
CorruptedRdmaBatch batch;
submitCorruptedRdmaBatch(engine, batch, 0xA5);
TransferStatus overall_status{};
ASSERT_TRUE(engine.getTransferStatus(batch.batch_id, overall_status).ok());
EXPECT_EQ(overall_status.s, TransferStatusEnum::FAILED);
EXPECT_EQ(batch.fake_tcp->submit_calls.load(), 0);
ASSERT_TRUE(engine.progressBatch(batch.batch_id, overall_status).ok());
EXPECT_EQ(overall_status.s, TransferStatusEnum::FAILED);
EXPECT_EQ(batch.fake_tcp->submit_calls.load(), 0);
EXPECT_TRUE(engine.freeBatch(batch.batch_id).ok());
EXPECT_TRUE(
engine.unregisterLocalMemory(batch.buf.data(), batch.buf.size()).ok());
}
TEST(EngineFailoverE2E, ProgressBatchHonorsMaxFailoverAttemptsZero) {
auto cfg = makeMinimalP2PConfig();
cfg->set("enable_auto_failover_on_poll", false);
cfg->set("max_failover_attempts", 0);
TransferEngineImpl engine(cfg);
ASSERT_TRUE(engine.available());
CorruptedRdmaBatch batch;
submitCorruptedRdmaBatch(engine, batch, 0xA6);
TransferStatus overall_status{};
ASSERT_TRUE(engine.progressBatch(batch.batch_id, overall_status).ok());
EXPECT_EQ(overall_status.s, TransferStatusEnum::FAILED);
EXPECT_EQ(batch.fake_rdma->submit_calls.load(), 1);
EXPECT_EQ(batch.fake_tcp->submit_calls.load(), 0);
EXPECT_TRUE(engine.freeBatch(batch.batch_id).ok());
EXPECT_TRUE(
engine.unregisterLocalMemory(batch.buf.data(), batch.buf.size()).ok());
}
TEST(EngineFailoverE2E, ProgressBatchKeepsOverallPendingWithMixedOutcomes) {
auto cfg = makeMinimalP2PConfig();
cfg->set("enable_auto_failover_on_poll", false);
cfg->set("max_failover_attempts", 0);
TransferEngineImpl engine(cfg);
ASSERT_TRUE(engine.available());
constexpr size_t kBufLen = 4096;
std::vector<uint8_t> failing_buf(kBufLen, 0xA7);
std::vector<uint8_t> pending_buf(kBufLen, 0xA8);
const uint64_t failing_addr =
reinterpret_cast<uint64_t>(failing_buf.data());
auto fake_rdma = std::make_shared<FakeTransport>(
RDMA, FakeTransport::StatusFactory{},
[failing_addr](const Request& req, int poll_count) {
if (req.target_offset == failing_addr) {
return TransferStatus{TransferStatusEnum::FAILED, 0};
}
if (poll_count == 1) {
return TransferStatus{TransferStatusEnum::PENDING, 0};
}
return TransferStatus{TransferStatusEnum::COMPLETED, req.length};
});
auto fake_tcp = std::make_shared<FakeTransport>(TCP);
std::string seg_name = engine.getSegmentName();
ASSERT_TRUE(fake_rdma->install(seg_name, nullptr, nullptr).ok());
ASSERT_TRUE(fake_tcp->install(seg_name, nullptr, nullptr).ok());
engine.swapTransportForTest(RDMA, fake_rdma);
engine.swapTransportForTest(TCP, fake_tcp);
ASSERT_TRUE(engine.registerLocalMemory(failing_buf.data(), kBufLen).ok());
ASSERT_TRUE(engine.registerLocalMemory(pending_buf.data(), kBufLen).ok());
BatchID batch_id = engine.allocateBatch(2);
ASSERT_NE(batch_id, (BatchID)0);
Request failing_req;
failing_req.opcode = Request::WRITE;
failing_req.source = failing_buf.data();
failing_req.target_id = LOCAL_SEGMENT_ID;
failing_req.target_offset = failing_addr;
failing_req.length = kBufLen;
Request pending_req;
pending_req.opcode = Request::WRITE;
pending_req.source = pending_buf.data();
pending_req.target_id = LOCAL_SEGMENT_ID;
pending_req.target_offset = reinterpret_cast<uint64_t>(pending_buf.data());
pending_req.length = kBufLen;
ASSERT_TRUE(
engine.submitTransfer(batch_id, {failing_req, pending_req}).ok());
TransferStatus overall_status{};
ASSERT_TRUE(engine.progressBatch(batch_id, overall_status).ok());
EXPECT_EQ(overall_status.s, TransferStatusEnum::PENDING);
EXPECT_EQ(fake_tcp->submit_calls.load(), 0);
ASSERT_TRUE(engine.progressBatch(batch_id, overall_status).ok());
EXPECT_EQ(overall_status.s, TransferStatusEnum::FAILED);
EXPECT_TRUE(engine.freeBatch(batch_id).ok());
EXPECT_TRUE(engine.unregisterLocalMemory(failing_buf.data(), kBufLen).ok());
EXPECT_TRUE(engine.unregisterLocalMemory(pending_buf.data(), kBufLen).ok());
}
TEST(EngineFailoverE2E,
WaitTransferCompletionUsesProgressBatchWhenPollDisabled) {
auto cfg = makeMinimalP2PConfig();
cfg->set("enable_auto_failover_on_poll", false);
TransferEngineImpl engine(cfg);
ASSERT_TRUE(engine.available());
CorruptedRdmaBatch batch;
submitCorruptedRdmaBatch(engine, batch, 0xA9);
EXPECT_TRUE(engine.waitTransferCompletion(batch.batch_id).ok());
EXPECT_EQ(batch.fake_tcp->submit_calls.load(), 1);
EXPECT_TRUE(
engine.unregisterLocalMemory(batch.buf.data(), batch.buf.size()).ok());
}
TEST(EngineFailoverE2E, TransferSyncUsesProgressBatchWhenPollDisabled) {
auto cfg = makeMinimalP2PConfig();
cfg->set("enable_auto_failover_on_poll", false);
TransferEngineImpl engine(cfg);
ASSERT_TRUE(engine.available());
auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
auto fake_tcp = std::make_shared<FakeTransport>(TCP);
FaultPolicy rdma_policy;
rdma_policy.status_corrupt_rate = 1.0;
auto proxied_rdma =
std::make_shared<FaultProxyTransport>(fake_rdma, rdma_policy);
std::string seg_name = engine.getSegmentName();
ASSERT_TRUE(proxied_rdma->install(seg_name, nullptr, nullptr).ok());
ASSERT_TRUE(fake_tcp->install(seg_name, nullptr, nullptr).ok());
engine.swapTransportForTest(RDMA, proxied_rdma);
engine.swapTransportForTest(TCP, fake_tcp);
constexpr size_t kBufLen = 4096;
std::vector<uint8_t> buf(kBufLen, 0xB0);
ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
Request req;
req.opcode = Request::WRITE;
req.source = buf.data();
req.target_id = LOCAL_SEGMENT_ID;
req.target_offset = reinterpret_cast<uint64_t>(buf.data());
req.length = kBufLen;
EXPECT_TRUE(engine.transferSync({req}).ok());
EXPECT_EQ(fake_rdma->submit_calls.load(), 1);
EXPECT_EQ(fake_tcp->submit_calls.load(), 1);
EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
}
TEST(EngineFailoverE2E, ProgressBatchAdvancesExactlyOneStepPerCall) {
auto cfg = makeMinimalP2PConfig();
cfg->set("enable_auto_failover_on_poll", false);
TransferEngineImpl engine(cfg);
ASSERT_TRUE(engine.available());
// Reach COMPLETED only on the third poll; earlier polls stay PENDING.
auto fake_rdma = std::make_shared<FakeTransport>(
RDMA, FakeTransport::StatusFactory{},
[](const Request& req, int poll_count) {
if (poll_count < 3) {
return TransferStatus{TransferStatusEnum::PENDING, 0};
}
return TransferStatus{TransferStatusEnum::COMPLETED, req.length};
});
auto fake_tcp = std::make_shared<FakeTransport>(TCP);
std::string seg_name = engine.getSegmentName();
ASSERT_TRUE(fake_rdma->install(seg_name, nullptr, nullptr).ok());
ASSERT_TRUE(fake_tcp->install(seg_name, nullptr, nullptr).ok());
engine.swapTransportForTest(RDMA, fake_rdma);
engine.swapTransportForTest(TCP, fake_tcp);
constexpr size_t kBufLen = 4096;
std::vector<uint8_t> buf(kBufLen, 0xC0);
ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
BatchID batch_id = engine.allocateBatch(1);
ASSERT_NE(batch_id, (BatchID)0);
Request req;
req.opcode = Request::WRITE;
req.source = buf.data();
req.target_id = LOCAL_SEGMENT_ID;
req.target_offset = reinterpret_cast<uint64_t>(buf.data());
req.length = kBufLen;
ASSERT_TRUE(engine.submitTransfer(batch_id, {req}).ok());
// Each progressBatch call must perform exactly one poll on the underlying
// transport — no internal loop until completion.
TransferStatus overall_status{};
ASSERT_TRUE(engine.progressBatch(batch_id, overall_status).ok());
EXPECT_EQ(overall_status.s, TransferStatusEnum::PENDING);
EXPECT_EQ(fake_rdma->status_calls.load(), 1);
ASSERT_TRUE(engine.progressBatch(batch_id, overall_status).ok());
EXPECT_EQ(overall_status.s, TransferStatusEnum::PENDING);
EXPECT_EQ(fake_rdma->status_calls.load(), 2);
ASSERT_TRUE(engine.progressBatch(batch_id, overall_status).ok());
EXPECT_EQ(overall_status.s, TransferStatusEnum::COMPLETED);
EXPECT_EQ(fake_rdma->status_calls.load(), 3);
EXPECT_EQ(fake_tcp->submit_calls.load(), 0);
EXPECT_TRUE(engine.freeBatch(batch_id).ok());
EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
}
// ---------------------------------------------------------------------------
// P1b: Both transports keep failing at status stage -> failover limit reached.
// ---------------------------------------------------------------------------

487
mooncake-transfer-engine/tent/tests/progress_worker_test.cpp Normal file
View File

@ -0,0 +1,487 @@
// 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.
//
// Tests the ProgressWorker skeleton (issue #2116, follow-up to PR #2160).
// Goals:
// * default-off behavior is byte-identical to the pre-worker world;
// * with the worker enabled and enable_auto_failover_on_poll=false, a
// caller that only submits + observes status (never calls
// progressBatch / waitTransferCompletion) still sees its batch
// progress through failover;
// * one notify advances the engine by exactly one progress step;
// * freeBatch racing the worker is safe (no UAF, no crash);
// * worker shuts down cleanly on engine destruction.
#include <gtest/gtest.h>
#include <atomic>
#include <chrono>
#include <cstdint>
#include <cstdlib>
#include <functional>
#include <memory>
#include <string>
#include <thread>
#include <utility>
#include <vector>
#include "tent/common/config.h"
#include "tent/common/types.h"
#include "tent/runtime/segment.h"
#include "tent/runtime/transfer_engine_impl.h"
#include "tent/runtime/transport.h"
#include "tent/transport/fault_proxy/fault_proxy_transport.h"
namespace mooncake {
namespace tent {
namespace {
// ---------------------------------------------------------------------------
// FakeTransport — same minimal shape used by engine_failover_e2e_test.cpp.
// Kept local to avoid cross-test linkage; sources of truth diverging is OK
// because we only exercise the "completes / status-can-be-overridden" surface.
// ---------------------------------------------------------------------------
class FakeSubBatch : public Transport::SubBatch {
public:
size_t size() const override { return task_count; }
size_t task_count = 0;
std::vector<Request> requests;
std::vector<TransferStatus> statuses;
std::vector<int> poll_counts;
};
class FakeTransport : public Transport {
public:
using StatusFactory = std::function<TransferStatus(const Request&)>;
using PollStatusFactory =
std::function<TransferStatus(const Request&, int)>;
explicit FakeTransport(TransportType self_type,
StatusFactory status_factory = {},
PollStatusFactory poll_status_factory = {})
: self_type_(self_type),
status_factory_(std::move(status_factory)),
poll_status_factory_(std::move(poll_status_factory)) {
caps.dram_to_dram = true;
}
std::atomic<int> install_calls{0};
std::atomic<int> submit_calls{0};
std::atomic<int> status_calls{0};
std::atomic<int> add_mem_calls{0};
Status install(std::string& /*local_segment_name*/,
std::shared_ptr<ControlService> /*metadata*/,
std::shared_ptr<Topology> /*local_topology*/,
std::shared_ptr<Config> /*conf*/ = nullptr) override {
++install_calls;
return Status::OK();
}
Status allocateSubBatch(SubBatchRef& batch, size_t /*max_size*/) override {
batch = new FakeSubBatch();
return Status::OK();
}
Status freeSubBatch(SubBatchRef& batch) override {
delete batch;
batch = nullptr;
return Status::OK();
}
Status submitTransferTasks(
SubBatchRef batch, const std::vector<Request>& request_list) override {
++submit_calls;
auto* fb = static_cast<FakeSubBatch*>(batch);
for (const auto& req : request_list) {
if (status_factory_) {
fb->statuses.push_back(status_factory_(req));
} else {
fb->statuses.push_back(
{TransferStatusEnum::COMPLETED, req.length});
}
fb->requests.push_back(req);
fb->poll_counts.push_back(0);
fb->task_count++;
}
return Status::OK();
}
Status getTransferStatus(SubBatchRef batch, int task_id,
TransferStatus& status) override {
++status_calls;
auto* fb = static_cast<FakeSubBatch*>(batch);
if (task_id < 0 || task_id >= (int)fb->statuses.size()) {
return Status::InvalidArgument("bad task_id" LOC_MARK);
}
++fb->poll_counts[task_id];
if (poll_status_factory_) {
status = poll_status_factory_(fb->requests[task_id],
fb->poll_counts[task_id]);
} else {
status = fb->statuses[task_id];
}
return Status::OK();
}
Status addMemoryBuffer(BufferDesc& desc,
const MemoryOptions& /*options*/) override {
++add_mem_calls;
desc.transports.push_back(self_type_);
return Status::OK();
}
Status addMemoryBuffer(std::vector<BufferDesc>& desc_list,
const MemoryOptions& options) override {
for (auto& d : desc_list) {
auto s = addMemoryBuffer(d, options);
if (!s.ok()) return s;
}
return Status::OK();
}
Status removeMemoryBuffer(BufferDesc& /*desc*/) override {
return Status::OK();
}
Status allocateLocalMemory(void** addr, size_t size,
MemoryOptions& /*options*/) override {
*addr = std::malloc(size);
if (!*addr) return Status::InternalError("malloc failed" LOC_MARK);
return Status::OK();
}
Status freeLocalMemory(void* addr, size_t /*size*/) override {
std::free(addr);
return Status::OK();
}
bool warmupMemory(void* /*addr*/, size_t /*length*/) override {
return false;
}
const char* getName() const override {
return self_type_ == RDMA ? "<fake-rdma>" : "<fake-tcp>";
}
private:
TransportType self_type_;
StatusFactory status_factory_;
PollStatusFactory poll_status_factory_;
};
std::shared_ptr<Config> makeMinimalP2PConfig() {
auto cfg = std::make_shared<Config>();
cfg->set("metadata_type", "p2p");
cfg->set("metadata_servers", "");
cfg->set("rpc_server_hostname", "127.0.0.1");
cfg->set("rpc_server_port", "0");
cfg->set("log_level", "warning");
cfg->set("merge_requests", false);
cfg->set("transports/tcp/enable", false);
cfg->set("transports/shm/enable", false);
cfg->set("transports/rdma/enable", false);
cfg->set("transports/io_uring/enable", false);
cfg->set("transports/nvlink/enable", false);
cfg->set("transports/mnnvl/enable", false);
cfg->set("transports/gds/enable", false);
cfg->set("transports/ascend_direct/enable", false);
cfg->set("max_failover_attempts", 3);
return cfg;
}
// ---------------------------------------------------------------------------
// 1. Default config: worker is not constructed, notifyBatchMaybeReady is a
// no-op, and behavior matches PR #2160 exactly.
// ---------------------------------------------------------------------------
TEST(ProgressWorker, DisabledByDefaultLeavesBehaviorUnchanged) {
auto cfg = makeMinimalP2PConfig();
TransferEngineImpl engine(cfg);
ASSERT_TRUE(engine.available());
auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
auto fake_tcp = std::make_shared<FakeTransport>(TCP);
std::string seg = engine.getSegmentName();
ASSERT_TRUE(fake_rdma->install(seg, nullptr, nullptr).ok());
ASSERT_TRUE(fake_tcp->install(seg, nullptr, nullptr).ok());
engine.swapTransportForTest(RDMA, fake_rdma);
engine.swapTransportForTest(TCP, fake_tcp);
constexpr size_t kBufLen = 4096;
std::vector<uint8_t> buf(kBufLen, 0x10);
ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
BatchID batch_id = engine.allocateBatch(1);
ASSERT_NE(batch_id, (BatchID)0);
Request req;
req.opcode = Request::WRITE;
req.source = buf.data();
req.target_id = LOCAL_SEGMENT_ID;
req.target_offset = reinterpret_cast<uint64_t>(buf.data());
req.length = kBufLen;
ASSERT_TRUE(engine.submitTransfer(batch_id, {req}).ok());
// No-op when the worker isn't constructed.
engine.notifyBatchMaybeReady(batch_id);
engine.notifyBatchMaybeReady((BatchID)0);
TransferStatus status{};
ASSERT_TRUE(engine.getTransferStatus(batch_id, status).ok());
EXPECT_EQ(status.s, TransferStatusEnum::COMPLETED);
EXPECT_EQ(fake_tcp->submit_calls.load(), 0);
EXPECT_TRUE(engine.freeBatch(batch_id).ok());
EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
}
// ---------------------------------------------------------------------------
// 2. Worker drives failover when the caller does not poll with
// allow_failover. This is the integration shape mooncake-pg needs.
// ---------------------------------------------------------------------------
TEST(ProgressWorker, ProgressesWithoutPollAutoFailover) {
auto cfg = makeMinimalP2PConfig();
cfg->set("enable_auto_failover_on_poll", false);
cfg->set("enable_progress_worker", true);
TransferEngineImpl engine(cfg);
ASSERT_TRUE(engine.available());
auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
auto fake_tcp = std::make_shared<FakeTransport>(TCP);
FaultPolicy rdma_policy;
rdma_policy.status_corrupt_rate = 1.0;
auto proxied_rdma =
std::make_shared<FaultProxyTransport>(fake_rdma, rdma_policy);
std::string seg = engine.getSegmentName();
ASSERT_TRUE(proxied_rdma->install(seg, nullptr, nullptr).ok());
ASSERT_TRUE(fake_tcp->install(seg, nullptr, nullptr).ok());
engine.swapTransportForTest(RDMA, proxied_rdma);
engine.swapTransportForTest(TCP, fake_tcp);
constexpr size_t kBufLen = 4096;
std::vector<uint8_t> buf(kBufLen, 0xC1);
ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
BatchID batch_id = engine.allocateBatch(1);
ASSERT_NE(batch_id, (BatchID)0);
Request req;
req.opcode = Request::WRITE;
req.source = buf.data();
req.target_id = LOCAL_SEGMENT_ID;
req.target_offset = reinterpret_cast<uint64_t>(buf.data());
req.length = kBufLen;
ASSERT_TRUE(engine.submitTransfer(batch_id, {req}).ok());
// Drive the worker until terminal. We deliberately never call
// progressBatch / waitTransferCompletion here — only
// notifyBatchMaybeReady + observation-only getTransferStatus
// (which, with auto-failover-on-poll disabled, will not advance failover
// by itself).
TransferStatus status{};
const auto deadline =
std::chrono::steady_clock::now() + std::chrono::milliseconds(2000);
while (std::chrono::steady_clock::now() < deadline) {
engine.notifyBatchMaybeReady(batch_id);
std::this_thread::sleep_for(std::chrono::milliseconds(5));
if (fake_tcp->submit_calls.load() == 0) continue;
status = {};
ASSERT_TRUE(engine.getTransferStatus(batch_id, status).ok());
if (status.s == TransferStatusEnum::COMPLETED) break;
}
EXPECT_EQ(status.s, TransferStatusEnum::COMPLETED)
<< "progress worker must drive failover when caller never calls "
"progressBatch";
EXPECT_EQ(fake_rdma->submit_calls.load(), 1);
EXPECT_GE(fake_tcp->submit_calls.load(), 1);
EXPECT_TRUE(engine.freeBatch(batch_id).ok());
EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
}
// ---------------------------------------------------------------------------
// 3. One notify == one progress step. The worker must not loop internally
// until completion; the next step requires another notify.
// ---------------------------------------------------------------------------
TEST(ProgressWorker, SingleNotifyAdvancesOneStep) {
auto cfg = makeMinimalP2PConfig();
cfg->set("enable_auto_failover_on_poll", false);
cfg->set("enable_progress_worker", true);
TransferEngineImpl engine(cfg);
ASSERT_TRUE(engine.available());
auto fake_rdma = std::make_shared<FakeTransport>(
RDMA, FakeTransport::StatusFactory{},
[](const Request& req, int poll_count) {
if (poll_count < 3) {
return TransferStatus{TransferStatusEnum::PENDING, 0};
}
return TransferStatus{TransferStatusEnum::COMPLETED, req.length};
});
auto fake_tcp = std::make_shared<FakeTransport>(TCP);
std::string seg = engine.getSegmentName();
ASSERT_TRUE(fake_rdma->install(seg, nullptr, nullptr).ok());
ASSERT_TRUE(fake_tcp->install(seg, nullptr, nullptr).ok());
engine.swapTransportForTest(RDMA, fake_rdma);
engine.swapTransportForTest(TCP, fake_tcp);
constexpr size_t kBufLen = 4096;
std::vector<uint8_t> buf(kBufLen, 0xC2);
ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
BatchID batch_id = engine.allocateBatch(1);
ASSERT_NE(batch_id, (BatchID)0);
Request req;
req.opcode = Request::WRITE;
req.source = buf.data();
req.target_id = LOCAL_SEGMENT_ID;
req.target_offset = reinterpret_cast<uint64_t>(buf.data());
req.length = kBufLen;
ASSERT_TRUE(engine.submitTransfer(batch_id, {req}).ok());
// Initial state: nothing polled yet.
EXPECT_EQ(fake_rdma->status_calls.load(), 0);
// Drive exactly one progress step via the worker. We can't observe the
// step instantly, but we can wait until status_calls increments by 1
// and then assert it does NOT keep climbing to 3 on its own.
engine.notifyBatchMaybeReady(batch_id);
const auto step_deadline =
std::chrono::steady_clock::now() + std::chrono::milliseconds(500);
while (std::chrono::steady_clock::now() < step_deadline &&
fake_rdma->status_calls.load() == 0) {
std::this_thread::sleep_for(std::chrono::milliseconds(2));
}
ASSERT_EQ(fake_rdma->status_calls.load(), 1)
<< "worker should issue exactly one poll for a single notify";
// Give the worker a generous window to misbehave. status_calls must
// stay at 1 because we did not notify again and the engine did not
// reach a terminal state.
std::this_thread::sleep_for(std::chrono::milliseconds(50));
EXPECT_EQ(fake_rdma->status_calls.load(), 1)
<< "worker must not loop internally until completion";
EXPECT_TRUE(engine.freeBatch(batch_id).ok());
EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
}
// ---------------------------------------------------------------------------
// 4. freeBatch races with worker notifications. With ASAN/UBSAN this
// catches missing alive_batches_ / progress_mutex_ coverage.
// ---------------------------------------------------------------------------
TEST(ProgressWorker, FreeBatchRacesWithWorker) {
auto cfg = makeMinimalP2PConfig();
cfg->set("enable_auto_failover_on_poll", false);
cfg->set("enable_progress_worker", true);
TransferEngineImpl engine(cfg);
ASSERT_TRUE(engine.available());
auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
auto fake_tcp = std::make_shared<FakeTransport>(TCP);
std::string seg = engine.getSegmentName();
ASSERT_TRUE(fake_rdma->install(seg, nullptr, nullptr).ok());
ASSERT_TRUE(fake_tcp->install(seg, nullptr, nullptr).ok());
engine.swapTransportForTest(RDMA, fake_rdma);
engine.swapTransportForTest(TCP, fake_tcp);
constexpr size_t kBufLen = 4096;
std::vector<uint8_t> buf(kBufLen, 0xC3);
ASSERT_TRUE(engine.registerLocalMemory(buf.data(), kBufLen).ok());
// Concurrently spam stale notifications from a second thread while the
// main thread submits, frees, and re-allocates batches.
std::atomic<bool> stop{false};
std::atomic<BatchID> latest{0};
std::thread spammer([&] {
while (!stop.load(std::memory_order_acquire)) {
BatchID bid = latest.load(std::memory_order_acquire);
engine.notifyBatchMaybeReady(bid);
std::this_thread::sleep_for(std::chrono::microseconds(50));
}
});
constexpr int kRounds = 100;
for (int i = 0; i < kRounds; ++i) {
BatchID batch_id = engine.allocateBatch(1);
ASSERT_NE(batch_id, (BatchID)0);
Request req;
req.opcode = Request::WRITE;
req.source = buf.data();
req.target_id = LOCAL_SEGMENT_ID;
req.target_offset = reinterpret_cast<uint64_t>(buf.data());
req.length = kBufLen;
ASSERT_TRUE(engine.submitTransfer(batch_id, {req}).ok());
latest.store(batch_id, std::memory_order_release);
engine.notifyBatchMaybeReady(batch_id);
// Free immediately; the worker may pick the notification up after
// free. The progress_mutex_ + alive_batches_ guard must keep this
// safe.
EXPECT_TRUE(engine.freeBatch(batch_id).ok());
}
stop.store(true, std::memory_order_release);
spammer.join();
EXPECT_TRUE(engine.unregisterLocalMemory(buf.data(), kBufLen).ok());
}
// ---------------------------------------------------------------------------
// 5. Engine teardown joins the worker cleanly even with pending notifies.
// ---------------------------------------------------------------------------
TEST(ProgressWorker, EngineDestructorJoinsWorker) {
auto cfg = makeMinimalP2PConfig();
cfg->set("enable_auto_failover_on_poll", false);
cfg->set("enable_progress_worker", true);
{
TransferEngineImpl engine(cfg);
ASSERT_TRUE(engine.available());
auto fake_rdma = std::make_shared<FakeTransport>(RDMA);
auto fake_tcp = std::make_shared<FakeTransport>(TCP);
std::string seg = engine.getSegmentName();
ASSERT_TRUE(fake_rdma->install(seg, nullptr, nullptr).ok());
ASSERT_TRUE(fake_tcp->install(seg, nullptr, nullptr).ok());
engine.swapTransportForTest(RDMA, fake_rdma);
engine.swapTransportForTest(TCP, fake_tcp);
// Push some notifies for non-existent batches; worker must reject
// them via alive_batches_ check and stay alive.
for (int i = 0; i < 8; ++i) {
engine.notifyBatchMaybeReady((BatchID)(uintptr_t)0xdeadbeef);
}
}
// If teardown hangs or crashes here, gtest fails this test on timeout
// / signal — no further assert needed.
SUCCEED();
}
} // namespace
} // namespace tent
} // namespace mooncake