PLUGINS/UCX: wait for in-flight transfers before release() returns

src/plugins/ucx/ucx_backend.cpp

@@ -26,6 +26,7 @@
 #include <optional>
 #include <limits>
 #include <future>
+#include <thread>
 #include <string.h>
 #include <unistd.h>
 #include "absl/strings/numbers.h"
@@ -133,20 +134,73 @@ class nixlUcxBackendReqH : public nixlBackendReqH {
 
     virtual void
     release() {
-        // TODO: Error log: uncompleted requests found! Cancelling ...
         for (nixlUcxReq req : requests_) {
             const nixl_status_t ret = nixl::ucx::ucsToNixlStatus(ucp_request_check_status(req));
             if (ret == NIXL_IN_PROG) {
-                // TODO: Need process this properly.
-                // it may not be enough to cancel UCX request
+                // A no-op for RMA - ucp_request_cancel() only acts on tag receives - but
+                // kept so tag-based operations still get cancelled before the drain.
                 worker_->reqCancel(req);
+                if (!drainRequest(req)) {
+                    NIXL_ERROR << "Request " << req
+                               << " is still in flight after the drain timeout. It is not "
+                                  "safe to deregister the memory backing this transfer yet: "
+                                  "UCX dereferences the ucp_mem_h when the operation "
+                                  "eventually completes.";
+                }
             }
             worker_->reqRelease(req);
         }
         requests_.clear();
         conn_.reset();
     }
 
+    // Progress the worker until req reaches a terminal state, so that UCX is done with the
+    // ucp_mem_h before the caller is free to deregister it. Returns false if the request is
+    // still in flight when the deadline expires.
+    //
+    // RMA operations are posted with UCP_OP_ATTR_FIELD_MEMH, and UCX keeps that ucp_mem_h in
+    // the request as a plain pointer (ucp_datatype_iter contig.memh): a user memh is not
+    // reference counted. On completion ucp_datatype_iter_cleanup() calls ucp_memh_put() on
+    // it, which dereferences memh->context and memh->parent. If the caller deregistered in
+    // the meantime, ucp_mem_unmap() has already ucs_free()d the memh and the completion
+    // faults inside ucp_memh_put().
+    //
+    // Releasing the request object itself is safe either way: ucp_request_free() on an
+    // uncompleted request only marks it released, and UCX returns it to the request pool
+    // when it completes. What the drain protects is the memory handle, not the request.
+    //
+    // The deadline keeps release() from blocking forever when a transfer is wedged on an
+    // endpoint that is alive but no longer making progress.
+    [[nodiscard]] bool
+    drainRequest(nixlUcxReq req) const {
+        using namespace std::chrono_literals;
+
+        const auto timeout =
+            nixl::config::getValueDefaulted("NIXL_UCX_REQUEST_DRAIN_TIMEOUT", 10'000ms);
+        const auto warning_interval =
+            nixl::config::getValueDefaulted("NIXL_UCX_WARNING_TIMEOUT", 5'000ms);
+        auto next_warning = warning_interval;
+
+        const auto start = std::chrono::steady_clock::now();
+        while (nixl::ucx::ucsToNixlStatus(ucp_request_check_status(req)) == NIXL_IN_PROG) {
+            const auto elapsed = std::chrono::steady_clock::now() - start;
+            if (elapsed > timeout) {
+                return false;
+            }
+
+            if (elapsed > next_warning) {
+                NIXL_WARN << "Still draining in-flight request " << req << " after "
+                          << next_warning.count() << " ms";
+                next_warning += warning_interval;
+            }
+
+            if (worker_->progress() == 0) {
+                std::this_thread::sleep_for(1ms);
+            }
+        }
+        return true;
+    }
+
     [[nodiscard]] virtual nixl_status_t
     status() {
         if (requests_.empty()) {
@@ -547,12 +601,59 @@ class nixlUcxCompositeBackendReqH : public nixlUcxBackendReqH {
         if (sharedState_) {
             // Set failed status to stop progress chunks
             sharedState_->status.store(NIXL_ERR_NOT_FOUND);
+
+            if (!waitForPendingChunks()) {
+                NIXL_ERROR << *this << " still has " << sharedState_->pendingReqs.load()
+                           << " chunk request(s) in flight after the drain timeout; "
+                              "deregistering the memory backing this transfer is not safe.";
+            }
+
             // Reset shared state - it will be effectively released when the last chunk
             // resets the shared state pointer
             sharedState_.reset();
         }
     }
 
+    // Wait for chunks already in flight to finish. Setting the failed status stops new
+    // chunks from starting, but it does not complete requests that are already posted,
+    // and resetting the shared state only drops our reference to it. Returns false if
+    // chunks are still in flight when the deadline expires.
+    //
+    // Returning while a chunk request is still in flight lets the caller deregister its
+    // memory: ucp_mem_unmap() then frees the memory handle while a zcopy completion is
+    // still pending, and that completion later dereferences the freed handle on a
+    // progress thread (use-after-free inside ucp_memh_put).
+    //
+    // The deadline keeps release() from blocking forever when a chunk is wedged on an
+    // endpoint that is alive but no longer making progress.
+    [[nodiscard]] bool
+    waitForPendingChunks() const {
+        using namespace std::chrono_literals;
+
+        const auto timeout =
+            nixl::config::getValueDefaulted("NIXL_UCX_REQUEST_DRAIN_TIMEOUT", 10'000ms);
+        const auto warning_interval =
+            nixl::config::getValueDefaulted("NIXL_UCX_WARNING_TIMEOUT", 5'000ms);
+        auto next_warning = warning_interval;
+
+        const auto start = std::chrono::steady_clock::now();
+        while (sharedState_->pendingReqs.load() > 0) {
+            const auto elapsed = std::chrono::steady_clock::now() - start;
+            if (elapsed > timeout) {
+                return false;
+            }
+
+            if (elapsed > next_warning) {
+                NIXL_WARN << *this << " still waiting for " << sharedState_->pendingReqs.load()
+                          << " in-flight chunk(s) after " << next_warning.count() << " ms";
+                next_warning += warning_interval;
+            }
+
+            std::this_thread::yield();
+        }
+        return true;
+    }
+
     [[nodiscard]] nixl_status_t
     status() override {
         getWorker()->progressLoop();

test/gtest/test_transfer.cpp

@@ -26,7 +26,10 @@
 #include <absl/strings/str_format.h>
 #include <absl/time/clock.h>
 #include <gtest/gtest.h>
+#include <algorithm>
+#include <atomic>
 #include <cstdlib>
+#include <cstring>
 #include <filesystem>
 #include <memory>
 #include <optional>
@@ -664,6 +667,117 @@ TEST_P(TestTransferTelemetry, GetXferTelemetryDisabled) {
     EXPECT_LE(lig.getIgnoredCount(), 1);
 }
 
+// Releasing a transfer handle while its requests are still in flight must not return until
+// UCX is done with the memory handles the transfer was posted with. NIXL posts RMA with
+// UCP_OP_ATTR_FIELD_MEMH, and UCX keeps that ucp_mem_h in the request as a plain pointer,
+// dereferencing it from ucp_memh_put() when the operation completes. A caller that
+// deregisters as soon as release() returns therefore frees the memh from under UCX.
+//
+// The data path is pinned to TCP so that a large transfer is genuinely asynchronous: over
+// shm/self UCX copies inline, the post completes immediately and there is nothing in flight
+// for release() to drain.
+class TestTransferRelease : public TestTransfer {
+protected:
+    void
+    SetUp() override {
+        env.addVar("NIXL_TELEMETRY_ENABLE", "n");
+        env.addVar("UCX_TLS", "tcp");
+        addAgent(0);
+        addAgent(1);
+    }
+
+    static void
+    fillBuffers(const std::vector<MemBuffer> &buffers, uint8_t value) {
+        for (const auto &buffer : buffers) {
+            std::memset(
+                reinterpret_cast<void *>(static_cast<uintptr_t>(buffer)), value, buffer.getSize());
+        }
+    }
+
+    static size_t
+    countBytes(const std::vector<MemBuffer> &buffers, uint8_t value) {
+        size_t found = 0;
+        for (const auto &buffer : buffers) {
+            const auto *data = reinterpret_cast<const uint8_t *>(static_cast<uintptr_t>(buffer));
+            found += static_cast<size_t>(std::count(data, data + buffer.getSize(), value));
+        }
+        return found;
+    }
+};
+
+TEST_P(TestTransferRelease, InFlightXferIsDrainedBeforeReleaseReturns) {
+    // The descriptor count is above the fixture's split_batch_size so that the threadpool
+    // parameterisations exercise the composite handle rather than the plain one.
+    constexpr size_t size = 1024 * 1024;
+    constexpr size_t count = 64;
+    constexpr uint8_t src_pattern = 0xab;
+    constexpr uint8_t dst_pattern = 0x00;
+
+    std::vector<MemBuffer> local_buffers, remote_buffers;
+    createRegisteredMem(getAgent(0), size, count, DRAM_SEG, local_buffers);
+    createRegisteredMem(getAgent(1), size, count, DRAM_SEG, remote_buffers);
+    fillBuffers(remote_buffers, src_pattern);
+    fillBuffers(local_buffers, dst_pattern);
+
+    exchangeMD(0, 1);
+
+    nixlXferReqH *xfer_req = nullptr;
+    ASSERT_EQ(getAgent(0).createXferReq(NIXL_READ,
+                                        makeDescList<nixlBasicDesc>(local_buffers, DRAM_SEG),
+                                        makeDescList<nixlBasicDesc>(remote_buffers, DRAM_SEG),
+                                        getAgentName(1),
+                                        xfer_req),
+              NIXL_SUCCESS);
+    ASSERT_NE(xfer_req, nullptr);
+
+    const nixl_status_t post_status = getAgent(0).postXferReq(xfer_req);
+    ASSERT_TRUE((post_status == NIXL_SUCCESS) || (post_status == NIXL_IN_PROG));
+
+    // release() only progresses the local worker. Over TCP the peer has to progress too, so
+    // stand in for a live remote process while the drain runs; without a progress thread
+    // nothing else would advance agent 1.
+    std::atomic<bool> pump_peer{true};
+    std::thread peer_thread([&]() {
+        nixl_notifs_t notifs;
+        while (pump_peer.load(std::memory_order_relaxed)) {
+            getAgent(1).getNotifs(notifs);
+        }
+    });
+
+    // A timed-out drain is a failure of this test, not of the run, so keep it out of the
+    // global problem counter and assert on it here instead.
+    const LogIgnoreGuard lig_drain("Still draining in-flight request");
+    const LogIgnoreGuard lig_timeout("is still in flight after the drain timeout");
+
+    const nixl_status_t release_status = getAgent(0).releaseXferReq(xfer_req);
+    pump_peer = false;
+    peer_thread.join();
+
+    EXPECT_EQ(release_status, NIXL_SUCCESS);
+    EXPECT_EQ(lig_timeout.getIgnoredCount(), 0u) << "release() gave up before the read finished";
+
+    // A read completes only once the data has landed locally, so a release() that drained
+    // implies the whole destination is written by the time it returns. Without the drain the
+    // read is still outstanding here and the destination is only partly filled.
+    if (post_status == NIXL_IN_PROG) {
+        EXPECT_EQ(countBytes(local_buffers, src_pattern), size * count)
+            << "releaseXferReq() returned while the read was still in flight";
+    } else {
+        Logger() << "transfer completed inline, nothing was in flight to drain";
+    }
+
+    // Deregistering is what would free the ucp_mem_h from under an in-flight request.
+    invalidateMD(0, 1);
+    deregisterMem(getAgent(0), local_buffers, DRAM_SEG);
+    deregisterMem(getAgent(1), remote_buffers, DRAM_SEG);
+}
+
+NIXL_INSTANTIATE_TEST(ucx, TestTransferRelease, "UCX", true, 2, 0, "");
+NIXL_INSTANTIATE_TEST(ucx_no_pt, TestTransferRelease, "UCX", false, 2, 0, "");
+// The threadpool parameterisations take the composite handle path.
+NIXL_INSTANTIATE_TEST(ucx_threadpool, TestTransferRelease, "UCX", true, 6, 4, "");
+NIXL_INSTANTIATE_TEST(ucx_threadpool_no_pt, TestTransferRelease, "UCX", false, 6, 4, "");
+
 NIXL_INSTANTIATE_TEST(ucx, TestTransfer, "UCX", true, 2, 0, "");
 NIXL_INSTANTIATE_TEST(ucx_no_pt, TestTransfer, "UCX", false, 2, 0, "");
 NIXL_INSTANTIATE_TEST(ucx_threadpool, TestTransfer, "UCX", true, 6, 4, "");