[Bugfix][TENT] Don't throw when aggregating transfer status during teardown

mooncake-transfer-engine/benchmark/bench_runner.h

@@ -63,6 +63,9 @@ class BenchRunner {
     virtual uint64_t getTargetBufferBase(int thread_id, uint64_t block_size,
                                          uint64_t batch_size) const = 0;
 
+    // Returns elapsed microseconds on success, or a negative value if the
+    // transfer failed. Implementations must not call exit() from worker
+    // threads.
     virtual double runSingleTransfer(uint64_t local_addr, uint64_t target_id,
                                      uint64_t target_addr, uint64_t block_size,
                                      uint64_t batch_size, OpCode opcode,

mooncake-transfer-engine/benchmark/main.cpp

@@ -143,11 +143,18 @@ int processBatchSizes(
                    deadline_us * 1000ull;
         };
 
+        auto failTask = [&]() {
+            measurement_started.store(true, std::memory_order_release);
+            return -1;
+        };
+
         XferBenchTimer timer;
         while (timer.lap_us(false) < 1000000ull) {
-            runner.runSingleTransfer(local_addr, target_id, target_addr,
-                                     thread_block_size, thread_batch_size,
-                                     opcode, deadlineNs(), intent_type);
+            if (runner.runSingleTransfer(
+                    local_addr, target_id, target_addr, thread_block_size,
+                    thread_batch_size, opcode, deadlineNs(), intent_type) < 0) {
+                return failTask();
+            }
         }
         if (measurement_ready.fetch_add(1, std::memory_order_acq_rel) + 1 ==
             num_threads) {
@@ -174,6 +181,7 @@ int processBatchSizes(
                 auto val = runner.runSingleTransfer(
                     local_addr, target_id, target_addr, thread_block_size,
                     thread_batch_size, WRITE, deadlineNs(), intent_type);
+                if (val < 0) return failTask();
                 thread_instant_bandwidth.push_back(
                     gbPerSecond(batch_bytes, val));
                 transfer_duration.push_back(val);
@@ -182,6 +190,7 @@ int processBatchSizes(
                 val = runner.runSingleTransfer(
                     local_addr, target_id, target_addr, thread_block_size,
                     thread_batch_size, READ, deadlineNs(), intent_type);
+                if (val < 0) return failTask();
                 thread_instant_bandwidth.push_back(
                     gbPerSecond(batch_bytes, val));
                 if (XferBenchConfig::check_consistency)
@@ -206,6 +215,7 @@ int processBatchSizes(
                 auto val = runner.runSingleTransfer(
                     local_addr, target_id, target_addr, thread_block_size,
                     thread_batch_size, opcode, deadlineNs(), intent_type);
+                if (val < 0) return failTask();
                 thread_instant_bandwidth.push_back(
                     gbPerSecond(batch_bytes, val));
                 if (read_verify) {
@@ -459,5 +469,5 @@ int main(int argc, char* argv[]) {
         }
         runner->stopInitiator();
     }
-    return 0;
+    return interrupted ? EXIT_FAILURE : EXIT_SUCCESS;
 }

mooncake-transfer-engine/benchmark/tent_backend.cpp

@@ -15,6 +15,8 @@
 #include "tent_backend.h"
 #include "utils.h"
 #include "char_util.h"
+
+#include <atomic>
 #include "tent/common/types.h"
 #include "tent/runtime/platform.h"
 #include "tent/runtime/topology.h"
@@ -31,18 +33,18 @@
 namespace mooncake {
 namespace tent {
 
-volatile bool g_tent_running = true;
-volatile bool g_tent_triggered_sig = false;
+std::atomic<bool> g_tent_running{true};
+std::atomic<bool> g_tent_triggered_sig{false};
 
 void signalHandlerV1(int signum) {
-    if (g_tent_triggered_sig) {
+    if (g_tent_triggered_sig.load()) {
         LOG(ERROR) << "Received signal " << signum
                    << " again, forcefully terminating...";
         std::exit(EXIT_FAILURE);
     }
     LOG(INFO) << "Received signal " << signum << ", stopping target server...";
-    g_tent_running = false;
-    g_tent_triggered_sig = true;
+    g_tent_running.store(false);
+    g_tent_triggered_sig.store(true);
 }
 
 std::shared_ptr<Config> loadConfig() {
@@ -302,7 +304,7 @@ TENTBenchRunner::TENTBenchRunner() {
 TENTBenchRunner::~TENTBenchRunner() { freeBuffers(); }
 
 int TENTBenchRunner::runTarget() {
-    while (g_tent_running) sleep(1);
+    while (g_tent_running.load()) sleep(1);
     return 0;
 }
 
@@ -350,7 +352,7 @@ int TENTBenchRunner::startInitiator(int num_threads) {
     LOG(INFO) << "Opened " << target_handles_.size() << " target segments";
     threads_.resize(num_threads);
     current_task_.resize(threads_.size());
-    g_tent_running = true;
+    g_tent_running.store(true);
     for (size_t i = 0; i < threads_.size(); ++i)
         threads_[i] = std::thread(&TENTBenchRunner::runner, this, i);
     return 0;
@@ -359,7 +361,7 @@ int TENTBenchRunner::startInitiator(int num_threads) {
 int TENTBenchRunner::stopInitiator() {
     {
         std::unique_lock<std::mutex> lk(mtx_);
-        g_tent_running = false;
+        g_tent_running.store(false);
         cv_task_.notify_all();
         cv_done_.notify_all();
     }
@@ -434,14 +436,14 @@ void TENTBenchRunner::pinThread(int thread_id) {
 }
 
 int TENTBenchRunner::runner(int thread_id) {
-    while (g_tent_running) {
+    while (g_tent_running.load()) {
         std::function<int(int)> task;
         {
             std::unique_lock<std::mutex> lk(mtx_);
             cv_task_.wait(lk, [&] {
-                return !g_tent_running || current_task_[thread_id];
+                return !g_tent_running.load() || current_task_[thread_id];
             });
-            if (!g_tent_running) break;
+            if (!g_tent_running.load()) break;
             std::swap(task, current_task_[thread_id]);
         }
         if (task) task(thread_id);
@@ -460,15 +462,37 @@ int TENTBenchRunner::runInitiatorTasks(
         current_task_[id] = func;
     pending_ = (int)threads_.size();
     cv_task_.notify_all();
-    cv_done_.wait(lk, [&] { return !g_tent_running || pending_ == 0; });
-    return g_tent_running ? 0 : -1;
+    cv_done_.wait(lk, [&] { return !g_tent_running.load() || pending_ == 0; });
+    return g_tent_running.load() ? 0 : -1;
+}
+
+void TENTBenchRunner::noteTransferFailed() {
+    std::lock_guard<std::mutex> lk(mtx_);
+    g_tent_running.store(false);
+    cv_task_.notify_all();
+    cv_done_.notify_all();
 }
 
 double TENTBenchRunner::runSingleTransfer(
     uint64_t local_addr, uint64_t target_id, uint64_t target_addr,
     uint64_t block_size, uint64_t batch_size, OpCode opcode,
     uint64_t deadline_ns, IntentType intent_type) {
+    if (!g_tent_running.load()) return -1.0;
+
+    auto abortTransfer = [&](const char* why) {
+        LOG(ERROR) << why;
+        noteTransferFailed();
+        return -1.0;
+    };
+
     auto batch_id = engine_->allocateBatch(batch_size);
+    if (!batch_id) return abortTransfer("Failed to allocate transfer batch");
+
+    auto finish = [&](double duration) {
+        (void)engine_->freeBatch(batch_id);
+        return duration;
+    };
+
     std::vector<Request> requests;
     for (uint64_t i = 0; i < batch_size; ++i) {
         Request entry;
@@ -485,26 +509,39 @@ double TENTBenchRunner::runSingleTransfer(
         requests.emplace_back(entry);
     }
     XferBenchTimer timer;
+    Status submitted;
     if (XferBenchConfig::notifi) {
-        // Use target_addr as msg for verification by peer
         Notification notifi{"benchmark", std::to_string(target_addr)};
-        CHECK_FAIL(engine_->submitTransfer(batch_id, requests, notifi));
+        submitted = engine_->submitTransfer(batch_id, requests, notifi);
     } else {
-        CHECK_FAIL(engine_->submitTransfer(batch_id, requests));
+        submitted = engine_->submitTransfer(batch_id, requests);
     }
-    while (true) {
+    if (!submitted.ok()) {
+        LOG(ERROR) << "Failed to submit transfer: " << submitted.ToString();
+        noteTransferFailed();
+        return finish(-1.0);
+    }
+    while (g_tent_running.load()) {
         TransferStatus overall_status;
-        CHECK_FAIL(engine_->getTransferStatus(batch_id, overall_status));
+        auto polled = engine_->getTransferStatus(batch_id, overall_status);
+        if (!polled.ok()) {
+            LOG(ERROR) << "Failed to poll transfer: " << polled.ToString();
+            noteTransferFailed();
+            return finish(-1.0);
+        }
         if (overall_status.s == TransferStatusEnum::COMPLETED) {
-            break;
-        } else if (overall_status.s == TransferStatusEnum::FAILED) {
+            return finish(timer.lap_us());
+        }
+        if (overall_status.s == TransferStatusEnum::FAILED ||
+            overall_status.s == TransferStatusEnum::TIMEOUT ||
+            overall_status.s == TransferStatusEnum::CANCELED ||
+            overall_status.s == TransferStatusEnum::INVALID) {
             LOG(ERROR) << "Failed transfer detected";
-            exit(EXIT_FAILURE);
+            noteTransferFailed();
+            return finish(-1.0);
         }
     }
-    auto duration = timer.lap_us();
-    CHECK_FAIL(engine_->freeBatch(batch_id));
-    return duration;
+    return finish(-1.0);
 }
 
 }  // namespace tent

mooncake-transfer-engine/benchmark/tent_backend.h

@@ -170,6 +170,9 @@ class TENTBenchRunner : public BenchRunner {
                           "target address");
     }
 
+    // Returns elapsed microseconds on success, or a negative value if the
+    // transfer failed. Never calls exit() — workers set g_tent_running so
+    // main can join threads and stop the engine first.
     double runSingleTransfer(uint64_t local_addr, uint64_t target_id,
                              uint64_t target_addr, uint64_t block_size,
                              uint64_t batch_size, OpCode opcode,
@@ -182,6 +185,8 @@ class TENTBenchRunner : public BenchRunner {
 
     int runner(int thread_id);
 
+    void noteTransferFailed();
+
     size_t targetIndex(int thread_id) const;
 
     int localTargetThreadId(int thread_id) const;

mooncake-transfer-engine/tent/include/tent/common/types.h

@@ -162,6 +162,27 @@ enum TransferStatusEnum {
     FAILED
 };
 
+// Rank for aggregating batch status. Unknown values rank with FAILED so
+// getBatchStatus never throws from unordered_map::at during teardown.
+inline int transferStatusSeverity(TransferStatusEnum s) {
+    switch (s) {
+        case INITIAL:
+        case PENDING:
+        case COMPLETED:
+            return 0;
+        case INVALID:
+            return 1;
+        case CANCELED:
+            return 2;
+        case TIMEOUT:
+            return 3;
+        case FAILED:
+            return 4;
+        default:
+            return 4;
+    }
+}
+
 struct TransferStatus {
     TransferStatusEnum s;
     size_t transferred_bytes;

mooncake-transfer-engine/tent/src/runtime/transfer_engine_impl.cpp

@@ -2659,11 +2659,7 @@ Status TransferEngineImpl::getBatchStatus(BatchID batch_id,
     size_t total_tasks = 0;
     TransferStatusEnum worst_failure = PENDING;
     auto isWorse = [](TransferStatusEnum cur, TransferStatusEnum best) {
-        static const std::unordered_map<TransferStatusEnum, int> severity = {
-            {INITIAL, 0},  {PENDING, 0}, {COMPLETED, 0}, {INVALID, 1},
-            {CANCELED, 2}, {TIMEOUT, 3}, {FAILED, 4},
-        };
-        return severity.at(cur) > severity.at(best);
+        return transferStatusSeverity(cur) > transferStatusSeverity(best);
     };
     for (size_t task_id = 0; task_id < batch->task_list.size(); ++task_id) {
         auto& task = batch->task_list[task_id];

mooncake-transfer-engine/tent/tests/engine_failover_e2e_test.cpp

@@ -622,6 +622,22 @@ TEST(EngineFailoverE2E, OverallStatusUsesWorstFailureNotGenericFailed) {
         engine.unregisterLocalMemory(completed_buf.data(), kBufLen).ok());
 }
 
+TEST(TransferStatusSeverityTest, KnownRanksMatchFormerMap) {
+    EXPECT_EQ(transferStatusSeverity(INITIAL), 0);
+    EXPECT_EQ(transferStatusSeverity(PENDING), 0);
+    EXPECT_EQ(transferStatusSeverity(COMPLETED), 0);
+    EXPECT_EQ(transferStatusSeverity(INVALID), 1);
+    EXPECT_EQ(transferStatusSeverity(CANCELED), 2);
+    EXPECT_EQ(transferStatusSeverity(TIMEOUT), 3);
+    EXPECT_EQ(transferStatusSeverity(FAILED), 4);
+}
+
+TEST(TransferStatusSeverityTest, UnknownValueRanksWithFailedAndDoesNotThrow) {
+    auto unknown = static_cast<TransferStatusEnum>(99);
+    EXPECT_EQ(transferStatusSeverity(unknown), transferStatusSeverity(FAILED));
+    EXPECT_GT(transferStatusSeverity(unknown), transferStatusSeverity(TIMEOUT));
+}
+
 TEST(EngineFailoverE2E,
      WaitTransferCompletionUsesProgressBatchWhenPollDisabled) {
     auto cfg = makeMinimalP2PConfig();