[TransferEngine]Add grouped scatter tasks for TCP

mooncake-transfer-engine/include/transport/tcp_transport/tcp_transport.h

@@ -169,6 +169,8 @@ class TcpTransport : public Transport {
     Status submitTransferTaskGroup(
         const std::vector<TransferTask *> &task_list) override;
 
+    bool supportsGroupedScatter() const override { return true; }
+
     Status getTransferStatus(BatchID batch_id, size_t task_id,
                              TransferStatus &status) override;
 

mooncake-transfer-engine/src/transport/tcp_transport/tcp_transport.cpp

@@ -574,7 +574,8 @@ Status TcpTransport::submitTransferTask(
         do {
             task = task_list[i];
             assert(task && task->request);
-            slices.push_back(prepareTransfer(task, *task->request));
+            for (size_t j = 0; j < task->request_count; ++j)
+                slices.push_back(prepareTransfer(task, task->request[j]));
             ++i;
         } while (i < task_list.size() && task_list[i]->request &&
                  task_list[i]->request->task_group_id == group_id &&
@@ -591,15 +592,16 @@ Status TcpTransport::submitTransferTaskGroup(
     slices.reserve(task_list.size());
     for (auto* task : task_list) {
         assert(task && task->request);
-        slices.push_back(prepareTransfer(task, *task->request));
+        for (size_t j = 0; j < task->request_count; ++j)
+            slices.push_back(prepareTransfer(task, task->request[j]));
     }
     startTransferSequence(std::move(slices));
     return Status::OK();
 }
 
 Transport::Slice* TcpTransport::prepareTransfer(
     TransferTask* task, const TransferRequest& request) {
-    task->total_bytes = request.length;
+    task->total_bytes += request.length;
     Slice* slice = getSliceCache().allocate();
     slice->source_addr = static_cast<char*>(request.source);
     slice->length = request.length;

mooncake-transfer-engine/tests/tcp_write_visibility_test.cpp

@@ -1861,6 +1861,106 @@ TEST(TcpWriteVisibilityTest, V2ReadRoundTripAndRejectedRead) {
     EXPECT_EQ(runOne(h.engine.get(), r), TransferStatusEnum::COMPLETED);
 }
 
+TEST(TcpWriteVisibilityTest, GroupedScatterReadWritePreservesEveryFragment) {
+    EngineHandle h;
+    h.init(P2PHANDSHAKE, "127.0.0.1:17941", 4ull << 20);
+    ASSERT_TRUE(h.ok);
+    auto descriptor = h.engine->getMetadata()->getSegmentDescByID(h.segment_id);
+    auto* local = static_cast<char*>(h.pool) + (2ull << 20);
+    auto* remote = static_cast<char*>(h.pool);
+    constexpr size_t count = 256;
+    constexpr size_t length = 128;
+    std::vector<size_t> local_offsets(count), remote_offsets(count);
+    std::vector<size_t> lengths(count, length);
+    for (size_t i = 0; i < count; ++i) {
+        local_offsets[i] = i * 256;
+        remote_offsets[i] = i * 384;
+    }
+    for (auto opcode : {TransferRequest::READ, TransferRequest::WRITE}) {
+        memset(local, 0, 1ull << 20);
+        memset(remote, 0, 1ull << 20);
+        for (size_t i = 0; i < count; ++i) {
+            auto* source = opcode == TransferRequest::READ
+                               ? remote + remote_offsets[i]
+                               : local + local_offsets[i];
+            memset(source, static_cast<int>(i % 251 + 1), length);
+        }
+        std::vector<size_t> completions(count);
+        TransferEngine::ScatterTransferRange range{
+            .opcode = opcode,
+            .remote_segment = descriptor->name,
+            .remote_base_offset = h.remote_base,
+            .remote_size = 1ull << 20,
+            .local_buffer = local,
+            .local_capacity = 1ull << 20,
+            .local_offsets = local_offsets,
+            .remote_offsets = remote_offsets,
+            .lengths = lengths,
+            .on_fragment_complete =
+                [&](size_t i, const Status& status) {
+                    EXPECT_TRUE(status.ok());
+                    ++completions[i];
+                },
+        };
+        auto operation = h.engine->submitScatter({range});
+        ASSERT_TRUE(operation.wait().ok());
+        for (size_t i = 0; i < count; ++i) {
+            EXPECT_EQ(completions[i], 1u);
+            EXPECT_EQ(memcmp(local + local_offsets[i],
+                             remote + remote_offsets[i], length),
+                      0);
+            EXPECT_EQ(local[local_offsets[i] + length], 0);
+            EXPECT_EQ(remote[remote_offsets[i] + length], 0);
+        }
+    }
+}
+
+TEST(TcpWriteVisibilityTest, GroupedScatterPreservesPartialFailureResults) {
+    EngineHandle h;
+    constexpr size_t pool_size = 4ull << 20;
+    h.init(P2PHANDSHAKE, "127.0.0.1:17942", pool_size);
+    ASSERT_TRUE(h.ok);
+    auto descriptor = h.engine->getMetadata()->getSegmentDescByID(h.segment_id);
+    auto* local = static_cast<char*>(h.pool) + (2ull << 20);
+    std::vector<size_t> local_offsets{0, 256, 512};
+    // The middle fragment passes the supplied descriptor bounds but is
+    // rejected by the peer's actual registration, exercising transport failure.
+    std::vector<size_t> remote_offsets{0, pool_size + 4096, 512};
+    std::vector<size_t> lengths(3, 128);
+    for (auto opcode : {TransferRequest::READ, TransferRequest::WRITE}) {
+        memset(h.pool, 37, 1024);
+        memset(local, 73, 1024);
+        std::vector<size_t> completions(3);
+        std::vector<bool> succeeded(3);
+        TransferEngine::ScatterTransferRange range{
+            .opcode = opcode,
+            .remote_segment = descriptor->name,
+            .remote_base_offset = h.remote_base,
+            .remote_size = pool_size + 8192,
+            .local_buffer = local,
+            .local_capacity = 1024,
+            .local_offsets = local_offsets,
+            .remote_offsets = remote_offsets,
+            .lengths = lengths,
+            .on_fragment_complete =
+                [&](size_t i, const Status& status) {
+                    ++completions[i];
+                    succeeded[i] = status.ok();
+                },
+        };
+        auto operation = h.engine->submitScatter({range});
+        EXPECT_FALSE(operation.wait().ok());
+        EXPECT_EQ(completions, (std::vector<size_t>{1, 1, 1}));
+        EXPECT_EQ(succeeded, (std::vector<bool>{true, false, true}));
+        EXPECT_EQ(local[256], 73);
+        for (size_t i : {0u, 2u})
+            EXPECT_EQ(memcmp(local + local_offsets[i],
+                             static_cast<char*>(h.pool) + remote_offsets[i],
+                             lengths[i]),
+                      0);
+    }
+}
+
 // Mixed-version quadrant: a pooled legacy (v1) initiator against the current
 // server still transfers data correctly over repeated exchanges on one fixed
 // lane (with the old weaker WRITE completion semantics).

mooncake-transfer-engine/tests/transport_uint_test.cpp

@@ -22,14 +22,17 @@
 
 #include <algorithm>
 #include <array>
+#include <chrono>
 #include <condition_variable>
 #include <cstdlib>
+#include <cstring>
 #include <fstream>
 #include <future>
 #include <iomanip>
 #include <memory>
 #include <mutex>
 #include <optional>
+#include <thread>
 #include <utility>
 
 #include "multi_transport.h"
@@ -1085,6 +1088,95 @@ TEST_F(TransportTest, BatchCleanupRunsAfterBeforeDeleteCallback) {
     EXPECT_TRUE(cleanup_observed_callback);
 }
 
+#ifdef USE_TCP
+TEST_F(TransportTest, TcpScatterCombinesRequestsIntoOneTask) {
+    constexpr size_t kRequestCount = 256;
+    constexpr size_t kRequestBytes = 128;
+    constexpr size_t kBytes = kRequestCount * kRequestBytes;
+    // Keep request metadata and buffers alive until after engine shutdown,
+    // including when an unexpected transport failure aborts the test.
+    std::array<char, 2 * kBytes> buffer{};
+    std::vector<TransferRequest> requests;
+    requests.reserve(kRequestCount);
+    TransferEngineImpl engine(false);
+    ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0);
+    ASSERT_NE(engine.installTransport("tcp", nullptr), nullptr);
+    ASSERT_EQ(engine.registerLocalMemory(buffer.data(), buffer.size(), "cpu:0"),
+              0);
+    auto segment = engine.openSegment(engine.getLocalIpAndPort());
+    ASSERT_NE(engine.getMetadata()->getSegmentDescByID(segment), nullptr);
+
+    for (auto opcode : {TransferRequest::READ, TransferRequest::WRITE}) {
+        SCOPED_TRACE(opcode == TransferRequest::READ ? "READ" : "WRITE");
+        auto* source = opcode == TransferRequest::READ ? buffer.data() + kBytes
+                                                       : buffer.data();
+        auto* destination = opcode == TransferRequest::READ
+                                ? buffer.data()
+                                : buffer.data() + kBytes;
+        std::fill(buffer.begin(), buffer.end(), 0);
+        std::fill_n(source, kBytes, 37);
+        requests.clear();
+        for (size_t i = 0; i < kRequestCount; ++i) {
+            requests.push_back(TransferRequest{
+                .opcode = opcode,
+                .source = buffer.data() + i * kRequestBytes,
+                .target_id = segment,
+                .target_offset =
+                    reinterpret_cast<uint64_t>(buffer.data() + kBytes) +
+                    i * kRequestBytes,
+                .length = kRequestBytes,
+                .task_group_id = 1,
+            });
+        }
+
+        MultiTransport::ScatterSubmission submission;
+        const auto submitted = engine.submitScatter(requests, submission);
+        ASSERT_NE(submission.batch_id, INVALID_BATCH_ID);
+        const auto& tasks =
+            Transport::toBatchDesc(submission.batch_id).task_list;
+        const size_t task_count = tasks.size();
+        std::vector<size_t> request_counts;
+        for (const auto& task : tasks)
+            request_counts.push_back(task.request_count);
+
+        // Drain every actual task before asserting grouping, so the regression
+        // case (one task per request) also releases its batch normally.
+        std::vector<bool> done(task_count, false);
+        bool successful = true;
+        const auto deadline =
+            std::chrono::steady_clock::now() + std::chrono::seconds(15);
+        size_t remaining = task_count;
+        while (remaining != 0 && std::chrono::steady_clock::now() < deadline) {
+            for (size_t i = 0; i < task_count; ++i) {
+                if (done[i]) continue;
+                TransferStatus status;
+                auto result =
+                    engine.getTransferStatus(submission.batch_id, i, status);
+                if (!result.ok()) continue;
+                if (status.s == TransferStatusEnum::COMPLETED ||
+                    status.s == TransferStatusEnum::FAILED) {
+                    successful &= status.s == TransferStatusEnum::COMPLETED;
+                    done[i] = true;
+                    --remaining;
+                }
+            }
+            if (remaining != 0)
+                std::this_thread::sleep_for(std::chrono::milliseconds(1));
+        }
+        ASSERT_EQ(remaining, 0u)
+            << "TCP scatter did not reach physical completion";
+        EXPECT_TRUE(engine.freeBatchID(submission.batch_id).ok());
+        EXPECT_TRUE(submitted.ok());
+        EXPECT_TRUE(successful);
+        EXPECT_EQ(memcmp(source, destination, kBytes), 0);
+        EXPECT_EQ(task_count, 1u);
+        EXPECT_EQ(request_counts, (std::vector<size_t>{kRequestCount}));
+        EXPECT_EQ(submission.task_sizes, (std::vector<size_t>{kRequestCount}));
+    }
+    EXPECT_EQ(engine.closeSegment(segment), 0);
+}
+#endif
+
 TEST_F(TransportTest, ScatterSubmitFailurePreservesCompletedFragments) {
     TransferEngine engine(false);
     ASSERT_EQ(engine.init(P2PHANDSHAKE, "127.0.0.1:12345"), 0);