Fixes a race in the UCX backend between synchronous rkey unpack and endpoint teardown.

src/plugins/ucx/ucx_backend.cpp

@@ -966,6 +966,12 @@ nixlUcxEngine::internalMDHelper (const nixl_blob_t &blob,
             // TODO: err: remote connection not found
             return NIXL_ERR_NOT_FOUND;
         }
+        for (size_t i = 0; i < workers_.size(); ++i) {
+            const nixl_status_t status = it->second->getEp(i)->checkTxState();
+            if (status != NIXL_SUCCESS) {
+                return status;
+            }
+        }
         // nixlSerDes::_stringToBytes() was used to "unpack" blob here.
         output = new nixlUcxPublicMetadata(
             it->second, makePublicMetadataRkeys(it->second, workers_.size(), blob.data()));

src/plugins/ucx/ucx_utils.cpp

@@ -86,15 +86,15 @@ err_cb_wrapper(void *arg, ucp_ep_h ucp_ep, ucs_status_t status) {
 
 void
 nixlUcxEp::err_cb(ucp_ep_h ucp_ep, ucs_status_t status) {
-    ucs_status_ptr_t request;
+    const auto current_state = state_.load(std::memory_order_relaxed);
 
-    NIXL_DEBUG << "ep " << eph << ": state " << state
+    NIXL_DEBUG << "ep " << eph << ": state " << current_state
                << ", UCX error handling callback was invoked with status " << status << " ("
                << ucs_status_string(status) << ")";
 
     NIXL_ASSERT(eph == ucp_ep);
 
-    switch (state) {
+    switch (current_state) {
     case nixl::ucx::ep_state_t::UNINITIALIZED:
     case nixl::ucx::ep_state_t::FAILED:
         // The error was already handled, nothing to do
@@ -103,40 +103,46 @@ nixlUcxEp::err_cb(ucp_ep_h ucp_ep, ucs_status_t status) {
         return;
     case nixl::ucx::ep_state_t::CONNECTED:
         setState(nixl::ucx::ep_state_t::FAILED);
-        request = ucp_ep_close_nb(ucp_ep, UCP_EP_CLOSE_MODE_FORCE);
-        if (UCS_PTR_IS_PTR(request)) {
-            ucp_request_free(request);
-        }
         return;
     }
-    NIXL_FATAL << "Invalid endpoint state: " << state;
+    NIXL_FATAL << "Invalid endpoint state: " << current_state;
     std::terminate();
 }
 
 void
 nixlUcxEp::setState(nixl::ucx::ep_state_t new_state) {
-    NIXL_ASSERT(new_state != state);
-    NIXL_DEBUG << "ep " << eph << ": state " << state << " -> " << new_state;
-    state = new_state;
+    const auto old_state = state_.load(std::memory_order_relaxed);
+    NIXL_ASSERT(new_state != old_state);
+    NIXL_DEBUG << "ep " << eph << ": state " << old_state << " -> " << new_state;
+    state_ = new_state;
 }
 
 nixl_status_t
 nixlUcxEp::closeImpl() {
     ucs_status_ptr_t request = nullptr;
+    const nixl::ucx::ep_state_t current_state = state_;
     const ucp_request_param_t req_param = {.op_attr_mask = UCP_OP_ATTR_FIELD_FLAGS,
                                            .flags = closeFlags_};
 
-    switch (state) {
+    switch (current_state) {
     case nixl::ucx::ep_state_t::UNINITIALIZED:
     case nixl::ucx::ep_state_t::DISCONNECTED:
         // The EP has not been connected, or already disconnected.
         // Nothing to do.
         NIXL_ASSERT(eph == nullptr);
         return NIXL_SUCCESS;
-    case nixl::ucx::ep_state_t::FAILED:
-        // The EP was closed in error callback, just return error.
+    case nixl::ucx::ep_state_t::FAILED: {
+        const ucp_request_param_t force_req_param = {
+            .op_attr_mask = UCP_OP_ATTR_FIELD_FLAGS,
+            .flags = UCP_EP_CLOSE_FLAG_FORCE,
+        };
+        request = ucp_ep_close_nbx(eph, &force_req_param);
+        if (UCS_PTR_IS_PTR(request)) {
+            ucp_request_free(request);
+        }
         eph = nullptr;
         return NIXL_ERR_REMOTE_DISCONNECT;
+    }
     case nixl::ucx::ep_state_t::CONNECTED:
         request = ucp_ep_close_nbx(eph, &req_param);
         if (request == nullptr) {
@@ -153,7 +159,7 @@ nixlUcxEp::closeImpl() {
         eph = nullptr;
         return NIXL_SUCCESS;
     }
-    NIXL_FATAL << "Invalid endpoint state: " << state;
+    NIXL_FATAL << "Invalid endpoint state: " << current_state;
     std::terminate();
 }
 

src/plugins/ucx/ucx_utils.h

@@ -17,6 +17,7 @@
 #ifndef NIXL_SRC_UTILS_UCX_UCX_UTILS_H
 #define NIXL_SRC_UTILS_UCX_UCX_UTILS_H
 
+#include <atomic>
 #include <memory>
 #include <type_traits>
 
@@ -44,7 +45,7 @@ class nixlUcxMem;
 class nixlUcxEp {
 private:
     ucp_ep_h eph{nullptr};
-    nixl::ucx::ep_state_t state = nixl::ucx::ep_state_t::UNINITIALIZED;
+    std::atomic<nixl::ucx::ep_state_t> state_{nixl::ucx::ep_state_t::UNINITIALIZED};
     const uint32_t closeFlags_;
 
     void
@@ -65,7 +66,7 @@ class nixlUcxEp {
 
     [[nodiscard]] nixl_status_t
     checkTxState() const noexcept {
-        return nixl::ucx::toNixlStatus(state);
+        return nixl::ucx::toNixlStatus(state_);
     }
 
     nixlUcxEp(ucp_worker_h worker,

test/gtest/error_handling.cpp

@@ -20,6 +20,9 @@
 #include <nixl_types.h>
 #include "common.h"
 #include "nixl.h"
+#ifdef HAVE_UCX_BACKEND
+#include "ucx_utils.h"
+#endif
 
 namespace gtest {
 namespace nixl {
@@ -450,6 +453,29 @@ TEST_P(TestErrorHandling, XferPostThenFail) {
     testXfer<TestType::FAIL_AFTER_POST, NIXL_READ>();
 }
 
+#ifdef HAVE_UCX_BACKEND
+TEST_P(TestErrorHandling, ErrorCallbackMarksEndpointFailedWithoutClosingIt) {
+    std::vector<std::string> devices;
+    const size_t num_workers = std::get<1>(GetParam());
+    const bool use_progress_thread = std::get<2>(GetParam()) > 0;
+    nixlUcxContext consumer_context(
+        devices, use_progress_thread, num_workers, nixl_thread_sync_t::NIXL_THREAD_SYNC_STRICT, 1);
+    nixlUcxContext producer_context(
+        devices, use_progress_thread, num_workers, nixl_thread_sync_t::NIXL_THREAD_SYNC_STRICT, 1);
+    nixlUcxWorker consumer(consumer_context, UCP_ERR_HANDLING_MODE_PEER);
+    nixlUcxWorker producer(producer_context, UCP_ERR_HANDLING_MODE_PEER);
+    std::string producer_address = producer.epAddr();
+    auto endpoint = consumer.connect(producer_address.data(), producer_address.size());
+    ASSERT_NE(endpoint, nullptr);
+
+    const ucp_ep_h native_endpoint = endpoint->getEp();
+    endpoint->err_cb(native_endpoint, UCS_ERR_CONNECTION_RESET);
+
+    EXPECT_EQ(endpoint->checkTxState(), NIXL_ERR_REMOTE_DISCONNECT);
+    EXPECT_EQ(endpoint->getEp(), native_endpoint);
+}
+#endif
+
 INSTANTIATE_TEST_SUITE_P(ucx, TestErrorHandling, testing::Values(std::make_tuple("UCX", 1, 0)));
 INSTANTIATE_TEST_SUITE_P(ucx_threadpool,
                          TestErrorHandling,

test/gtest/meson.build

@@ -106,6 +106,7 @@ endif
 
 if ucx_dep.found() and is_variable('ucx_backend_interface')
     gtest_sources += 'hw_warning_test.cpp'
+    cpp_flags += '-DHAVE_UCX_BACKEND'
     ucx_hw_warning_dep = [ucx_backend_interface, ucx_dep]
 else
     ucx_hw_warning_dep = []