[TENT] Make UB batch submission all-or-nothing

mooncake-transfer-engine/tent/include/tent/transport/ub/workers.h

@@ -147,6 +147,8 @@ class UbWorkers final {
 
     Status start();
     Status stop();
+    Status submitBatch(const std::vector<UbTask::Ptr>& tasks,
+                       uint64_t device_mask = ~0ULL);
     Status submit(const UbTask::Ptr& task, uint64_t device_mask = ~0ULL);
     Status cancel(const UbTask::Ptr& task);
 
@@ -157,6 +159,11 @@ class UbWorkers final {
     [[nodiscard]] size_t inflightCount() const;
 
    private:
+    friend class UbWorkersTestPeer;
+    // Scoped to the submitting thread so transport tests can exercise the
+    // real workers without changing their lifecycle or racing posting lanes.
+    static thread_local bool fail_second_enqueue_for_test_;
+
     struct PendingSlice {
         // Keep the task alive until every queued/in-flight slice reaches a
         // terminal state. Callers may release the sub-batch handle before the

mooncake-transfer-engine/tent/src/transport/ub/ub_transport.cpp

@@ -8,6 +8,7 @@
 #include <chrono>
 #include <cstdint>
 #include <limits>
+#include <memory>
 #include <mutex>
 #include <new>
 #include <sstream>
@@ -705,6 +706,7 @@ Status UbTransport::submitTransferTasks(
 
     const auto notify_progress = ub_batch->notify_progress;
     const auto progress_batch_id = ub_batch->progress_batch_id;
+    auto progress_armed = std::make_shared<std::atomic<bool>>(false);
     std::vector<ub::UbTask::Ptr> new_tasks;
     new_tasks.reserve(request_list.size());
     for (const auto& request : request_list) {
@@ -717,9 +719,12 @@ Status UbTransport::submitTransferTasks(
                 "transports/ub/max_slices_per_task" LOC_MARK);
         }
         auto task = ub::UbTask::create(
-            request,
-            [notify_progress, progress_batch_id](const TransferStatus&) {
-                if (notify_progress) notify_progress(progress_batch_id);
+            request, [notify_progress, progress_batch_id,
+                      progress_armed](const TransferStatus&) {
+                if (progress_armed->load(std::memory_order_acquire) &&
+                    notify_progress) {
+                    notify_progress(progress_batch_id);
+                }
             });
         size_t offset = 0;
         while (offset < request.length) {
@@ -741,15 +746,15 @@ Status UbTransport::submitTransferTasks(
         new_tasks.push_back(std::move(task));
     }
 
-    for (auto& task : new_tasks) {
-        ub_batch->task_list.push_back(task);
-        auto status = impl_->workers->submit(task, ub_batch->device_mask);
-        if (!status.ok()) {
-            for (auto& queued : new_tasks) {
-                if (queued) queued->requestCancellation();
-            }
-            return status;
-        }
+    CHECK_STATUS(impl_->workers->submitBatch(new_tasks, ub_batch->device_mask));
+    ub_batch->task_list.insert(ub_batch->task_list.end(), new_tasks.begin(),
+                               new_tasks.end());
+    progress_armed->store(true, std::memory_order_release);
+    // A very small task can finish between worker publication and the batch
+    // append above. Its callback intentionally stays disarmed until the task
+    // is queryable; this notification closes that race.
+    if (!new_tasks.empty() && notify_progress) {
+        notify_progress(progress_batch_id);
     }
     return Status::OK();
 }

mooncake-transfer-engine/tent/src/transport/ub/workers.cpp

@@ -6,6 +6,8 @@
 #include <algorithm>
 #include <chrono>
 #include <limits>
+#include <new>
+#include <stdexcept>
 #include <unordered_set>
 #include <utility>
 
@@ -275,42 +277,106 @@ Status UbWorkers::submit(const UbTask::Ptr& task, uint64_t device_mask) {
     if (!task) {
         return Status::InvalidArgument("Cannot submit a null UB task" LOC_MARK);
     }
+    return submitBatch({task}, device_mask);
+}
+
+thread_local bool UbWorkers::fail_second_enqueue_for_test_ = false;
+
+Status UbWorkers::submitBatch(const std::vector<UbTask::Ptr>& tasks,
+                              uint64_t device_mask) {
+    const bool fail_second_enqueue =
+        std::exchange(fail_second_enqueue_for_test_, false);
     if (!accepting_.load(std::memory_order_acquire)) {
         return Status::InvalidArgument(
             "UB workers are not accepting work" LOC_MARK);
     }
-    const auto snapshot = task->snapshot();
-    if (!snapshot.sealed) {
-        return Status::InvalidArgument(
-            "UB task must be sealed before submission" LOC_MARK);
-    }
-    const int priority =
-        std::clamp(task->request().priority, static_cast<int>(PRIO_HIGH),
-                   static_cast<int>(PRIO_LOW));
 
     std::vector<PendingSlice> pending;
-    for (const auto& slice : task->slices()) {
-        if (slice && slice->markQueued()) {
+    std::unordered_set<const UbTask*> unique_tasks;
+    for (const auto& task : tasks) {
+        if (!task) {
+            return Status::InvalidArgument(
+                "Cannot submit a null UB task" LOC_MARK);
+        }
+        if (!unique_tasks.insert(task.get()).second) {
+            return Status::InvalidArgument(
+                "Cannot submit the same UB task twice in one batch" LOC_MARK);
+        }
+        const auto snapshot = task->snapshot();
+        if (!snapshot.sealed) {
+            return Status::InvalidArgument(
+                "UB task must be sealed before submission" LOC_MARK);
+        }
+        const int priority =
+            std::clamp(task->request().priority, static_cast<int>(PRIO_HIGH),
+                       static_cast<int>(PRIO_LOW));
+        for (const auto& slice : task->slices()) {
+            if (!slice || slice->snapshot().state != UbSliceState::kInitial) {
+                return Status::InvalidArgument(
+                    "UB task contains a slice that cannot be "
+                    "submitted" LOC_MARK);
+            }
             pending.push_back(PendingSlice{task, slice, device_mask, priority,
                                            task->request().target_id,
                                            task->request().opcode});
         }
     }
-    bool stopped_during_submit = false;
+
+    bool enqueue_failed = false;
+    size_t marked_count = 0;
     {
         std::lock_guard<std::mutex> lock(queue_mutex_);
         if (!accepting_.load(std::memory_order_relaxed)) {
-            stopped_during_submit = true;
-        } else {
-            for (auto& item : pending) {
-                queues_[item.priority].push_back(std::move(item));
+            return Status::InvalidArgument(
+                "UB workers stopped during submission" LOC_MARK);
+        }
+
+        for (const auto& item : pending) {
+            if (!item.slice->markQueued()) break;
+            ++marked_count;
+        }
+
+        if (marked_count == pending.size()) {
+            std::array<size_t, PRIO_LOW + 1> inserted_counts{};
+            try {
+                for (size_t index = 0; index < pending.size(); ++index) {
+                    if (fail_second_enqueue && index == 1) {
+                        throw std::length_error("Injected UB enqueue failure");
+                    }
+                    const auto& item = pending[index];
+                    queues_[item.priority].push_back(item);
+                    ++inserted_counts[item.priority];
+                }
+            } catch (const std::bad_alloc&) {
+                enqueue_failed = true;
+            } catch (const std::length_error&) {
+                enqueue_failed = true;
+            }
+            if (enqueue_failed) {
+                for (size_t priority = 0; priority < queues_.size();
+                     ++priority) {
+                    while (inserted_counts[priority] != 0) {
+                        queues_[priority].pop_back();
+                        --inserted_counts[priority];
+                    }
+                }
             }
         }
     }
-    if (stopped_during_submit) {
-        for (auto& item : pending) item.slice->requestCancellation();
+
+    if (enqueue_failed || marked_count != pending.size()) {
+        // A markQueued() failure or enqueue exception may occur after earlier
+        // slices changed state. None of them are visible to posting workers;
+        // cancel the private tasks so no slice remains orphaned in kQueued.
+        for (size_t index = 0; index < marked_count; ++index) {
+            (void)pending[index].task->requestCancellation();
+        }
+        if (enqueue_failed) {
+            return Status::InternalError(
+                "Unable to enqueue UB task batch" LOC_MARK);
+        }
         return Status::InvalidArgument(
-            "UB workers stopped during submission" LOC_MARK);
+            "UB task changed while it was being submitted" LOC_MARK);
     }
     if (!pending.empty()) queue_cv_.notify_all();
     return Status::OK();

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

@@ -27,6 +27,14 @@
 #include "tent/transport/ub/workers.h"
 
 namespace mooncake::tent::ub {
+
+class UbWorkersTestPeer {
+   public:
+    static void failSecondEnqueue() {
+        UbWorkers::fail_second_enqueue_for_test_ = true;
+    }
+};
+
 namespace {
 
 class FakeContext final : public Context {
@@ -547,6 +555,46 @@ TEST(UbNativeDataPathTest,
                           (void)endpoints.retire(endpoint);
                       });
     ASSERT_TRUE(workers.start().ok());
+
+    Request invalid_batch_request{};
+    invalid_batch_request.opcode = Request::WRITE;
+    invalid_batch_request.source = source.data();
+    invalid_batch_request.target_id = LOCAL_SEGMENT_ID;
+    invalid_batch_request.target_offset =
+        reinterpret_cast<uint64_t>(target.data());
+    invalid_batch_request.length = 16;
+    auto valid_task = UbTask::create(invalid_batch_request);
+    auto valid_slice = valid_task->addSlice(UbSliceSpec{
+        source.data(), reinterpret_cast<uint64_t>(target.data()), 16, 0, 1});
+    ASSERT_NE(valid_slice, nullptr);
+    ASSERT_TRUE(valid_task->seal());
+    auto unsealed_task = UbTask::create(invalid_batch_request);
+    ASSERT_NE(unsealed_task->addSlice(UbSliceSpec{
+                  source.data(), reinterpret_cast<uint64_t>(target.data()), 16,
+                  0, 1}),
+              nullptr);
+
+    EXPECT_FALSE(workers.submitBatch({valid_task, unsealed_task}).ok());
+    EXPECT_EQ(workers.queuedCount(), 0U);
+    EXPECT_EQ(workers.inflightCount(), 0U);
+    EXPECT_EQ(valid_slice->snapshot().state, UbSliceState::kInitial);
+
+    auto rollback_task = UbTask::create(invalid_batch_request);
+    auto rollback_slice = rollback_task->addSlice(UbSliceSpec{
+        source.data(), reinterpret_cast<uint64_t>(target.data()), 16, 0, 1});
+    ASSERT_NE(rollback_slice, nullptr);
+    ASSERT_TRUE(rollback_task->seal());
+
+    // Both slices are marked queued and the first is inserted before the
+    // second enqueue throws.
+    UbWorkersTestPeer::failSecondEnqueue();
+    EXPECT_TRUE(
+        workers.submitBatch({valid_task, rollback_task}).IsInternalError());
+    EXPECT_EQ(workers.queuedCount(), 0U);
+    EXPECT_EQ(workers.inflightCount(), 0U);
+    EXPECT_EQ(valid_slice->snapshot().state, UbSliceState::kCanceled);
+    EXPECT_EQ(rollback_slice->snapshot().state, UbSliceState::kCanceled);
+
     adapter->holdNextCompletion();
     adapter->dropNextQuiescedCompletion();
 
@@ -906,7 +954,7 @@ TEST(UbNativeDataPathTest, UbTransportRunsSelfReadWriteOverNativeControlPlane) {
     EXPECT_TRUE(transport.freeSubBatch(rejected_batch).ok());
 
     Transport::SubBatchRef batch = nullptr;
-    ASSERT_TRUE(transport.allocateSubBatch(batch, 2).ok());
+    ASSERT_TRUE(transport.allocateSubBatch(batch, 3).ok());
     Request request{};
     request.opcode = Request::WRITE;
     request.source = source.data();
@@ -928,10 +976,25 @@ TEST(UbNativeDataPathTest, UbTransportRunsSelfReadWriteOverNativeControlPlane) {
     EXPECT_EQ(source, target);
     EXPECT_GT(transport.getEstimatedBandwidth(), 0.0);
 
+    ASSERT_EQ(batch->size(), 1U);
+    auto single_slice_request = request;
+    single_slice_request.length = 16;
+    UbWorkersTestPeer::failSecondEnqueue();
+    EXPECT_TRUE(transport
+                    .submitTransferTasks(
+                        batch, {single_slice_request, single_slice_request})
+                    .IsInternalError());
+    EXPECT_EQ(batch->size(), 1U);
+    transfer = TransferStatus{PENDING, 0};
+    ASSERT_TRUE(transport.getTransferStatus(batch, 0, transfer).ok());
+    EXPECT_EQ(transfer.s, COMPLETED);
+    EXPECT_EQ(transfer.transferred_bytes, source.size());
+
     source.fill(0);
     Request read_request = request;
     read_request.opcode = Request::READ;
     ASSERT_TRUE(transport.submitTransferTasks(batch, {read_request}).ok());
+    EXPECT_EQ(batch->size(), 2U);
     transfer = TransferStatus{PENDING, 0};
     const auto read_deadline =
         std::chrono::steady_clock::now() + std::chrono::seconds(2);