Raise timeout error if endpoint create callback does not complete

rapidsai · Oct 21, 2024 · 98a1191 · 98a1191
1 parent 3088fce
commit 98a1191
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/cpp/src/endpoint.cpp b/cpp/src/endpoint.cpp
@@ -97,6 +97,8 @@ void endpointErrorCallback(void* arg, ucp_ep_h ep, ucs_status_t status)
                ep,
                status,
                ucs_status_string(status));
+
+  std::swap(endpoint->_handle, endpoint->_originalHandle);
 }
 
 Endpoint::Endpoint(std::shared_ptr<Component> workerOrListener, bool endpointErrorHandling)
@@ -136,10 +138,12 @@ void Endpoint::create(ucp_ep_params_t* params)
             3000000000 /* 3s */))
         break;
 
-      if (i == maxAttempts - 1)
+      if (i == maxAttempts - 1) {
+        status = UCS_ERR_TIMED_OUT;
         ucxx_error("Timeout waiting for ucp_ep_create, all attempts failed");
-      else
+      } else {
         ucxx_warn("Timeout waiting for ucp_ep_create, retrying");
+      }
     }
     utils::ucsErrorThrow(status);
   } else {