UCP/RNDV: keep the rkey when a put completion restarts on failover

src/ucp/proto/proto_common.c

@@ -876,10 +876,22 @@ void ucp_proto_request_restart(ucp_request_t *req)
                   proto_config->proto->name, req->send.proto_stage);
 
     status = proto_config->proto->reset(req);
+    if (status == UCS_ERR_CANCELED) {
+        /* The request was already completed by cancellation. */
+        return;
+    }
+
     if (status != UCS_OK) {
-        ucs_assertv_always(status == UCS_ERR_CANCELED,
-                           "req %p, failed to reset: status %s", req,
-                           ucs_status_string(status));
+        /* The protocol cannot resume from this state - a rendezvous protocol
+         * whose remote key has already been released, for instance. Restart
+         * re-enters protocol selection, so the request must still own
+         * everything the selected protocol will use, and reset is the only
+         * place that can tell. Fail the request rather than re-selecting on
+         * top of resources it no longer holds: the caller sees a failed
+         * transfer instead of the worker dereferencing a released one. */
+        ucs_error("req %p: proto %s cannot restart (%s); aborting the request",
+                  req, proto_config->proto->name, ucs_status_string(status));
+        ucp_proto_request_abort(req, status);
         return;
     }
 

src/ucp/proto/proto_common.inl

@@ -93,6 +93,19 @@ ucp_proto_request_zcopy_clean(ucp_request_t *req, unsigned dt_mask)
     req->flags &= ~UCP_REQUEST_FLAG_PROTO_INITIALIZED;
 }
 
+/* Whether a completion with this status will restart the request instead of
+ * completing it. Callers that release per-request resources before completing
+ * must ask first: anything the restarted send still needs (a remote rkey, for
+ * one) has to outlive the completion. */
+static UCS_F_ALWAYS_INLINE int
+ucp_proto_request_is_failover_restart(const ucp_request_t *req,
+                                      ucs_status_t status)
+{
+    return ucs_unlikely(status != UCS_OK) &&
+           ucp_ep_err_mode_eq(req->send.ep, UCP_ERR_HANDLING_MODE_FAILOVER) &&
+           !(req->send.ep->flags & UCP_EP_FLAG_FAILED);
+}
+
 static UCS_F_ALWAYS_INLINE void
 ucp_proto_request_zcopy_complete_cb(ucp_request_t *req, ucs_status_t status,
                                     ucp_request_callback_t complete_cb)
@@ -105,9 +118,7 @@ ucp_proto_request_zcopy_complete_cb(ucp_request_t *req, ucs_status_t status,
         UCP_EP_STAT_TAG_OP(req->send.ep, EAGER)
     }
 
-    if (ucs_unlikely(status != UCS_OK) &&
-        ucp_ep_err_mode_eq(req->send.ep, UCP_ERR_HANDLING_MODE_FAILOVER) &&
-        !(req->send.ep->flags & UCP_EP_FLAG_FAILED)) {
+    if (ucp_proto_request_is_failover_restart(req, status)) {
         ucp_proto_request_restart(req);
     } else {
         if (complete_cb != NULL) {

src/ucp/rndv/rndv_get.c

@@ -115,8 +115,7 @@ ucp_proto_rndv_get_zcopy_fetch_completion(uct_completion_t *uct_comp)
     ucp_datatype_iter_mem_dereg(&req->send.state.dt_iter,
                                 UCS_BIT(UCP_DATATYPE_CONTIG));
     if (ucs_unlikely(uct_comp->status != UCS_OK)) {
-        if (ucp_ep_err_mode_eq(req->send.ep, UCP_ERR_HANDLING_MODE_FAILOVER) &&
-            !(req->send.ep->flags & UCP_EP_FLAG_FAILED)) {
+        if (ucp_proto_request_is_failover_restart(req, uct_comp->status)) {
             /* A lane failed but the endpoint is recovering: restart the
              * fetch on the surviving lanes instead of failing the receive.
              * Keep the rkey - the remote buffer is still registered until
@@ -228,6 +227,13 @@ static ucs_status_t ucp_rndv_get_zcopy_proto_reset(ucp_request_t *req)
         return UCS_OK;
     }
 
+    /* Same contract as the put scheme: the fetch resumes against the peer's
+     * remote key, so a restart is only possible while the request still owns
+     * it. */
+    if (req->send.rndv.rkey == NULL) {
+        return UCS_ERR_CONNECTION_RESET;
+    }
+
     req->flags &= ~UCP_REQUEST_FLAG_PROTO_INITIALIZED;
 
     switch (req->send.proto_stage) {

src/ucp/rndv/rndv_put.c

@@ -43,12 +43,22 @@ ucp_proto_rndv_put_common_complete(ucp_request_t *req)
 {
     const ucp_proto_rndv_put_priv_t UCS_V_UNUSED *rpriv =
                                                    req->send.proto_config->priv;
+    ucs_status_t status = req->send.state.uct_comp.status;
+
     ucp_trace_req(req, "rndv_put_common_complete");
     UCS_STATS_UPDATE_COUNTER(req->send.ep->worker->stats, rpriv->stat_counter,
                              +1);
-    ucp_proto_rndv_rkey_destroy(req);
-    ucp_proto_rndv_request_zcopy_complete(req,
-                                          req->send.state.uct_comp.status);
+
+    /* Keep the rkey when the completion below is going to restart this write
+     * on the surviving lanes rather than complete it: the remote buffer stays
+     * registered until ATP, and the restarted send reads the rkey again. The
+     * get scheme keeps it for the same reason. Destroying it here first left
+     * the restart dereferencing a NULL rkey. */
+    if (!ucp_proto_request_is_failover_restart(req, status)) {
+        ucp_proto_rndv_rkey_destroy(req);
+    }
+
+    ucp_proto_rndv_request_zcopy_complete(req, status);
 }
 
 static void ucp_proto_rndv_put_zcopy_completion(uct_completion_t *uct_comp)
@@ -455,6 +465,14 @@ static ucs_status_t ucp_proto_rndv_put_zcopy_reset(ucp_request_t *req)
 {
     const ucp_proto_rndv_put_priv_t *rpriv = req->send.proto_config->priv;
 
+    /* The write resumes against the peer's remote key, which came from the
+     * RTS and cannot be re-acquired here. Refuse the restart if a completion
+     * path released it, so the request fails instead of re-selecting a
+     * protocol that would dereference it. */
+    if (req->send.rndv.rkey == NULL) {
+        return UCS_ERR_CONNECTION_RESET;
+    }
+
     if (req->send.rndv.put.atp_count == rpriv->atp_num_lanes) {
         /* Sent all ATPs so the iterator should be at the end */
         ucs_assertv_always(ucp_datatype_iter_is_end(&req->send.state.dt_iter),