UCP/RNDV: Fix rkey use-after-destroy on FAILOVER restart of rndv/put

[jyizheng]

What

ucp_proto_rndv_put_common_complete() destroys req->send.rndv.rkey (sets it to NULL) before ucp_proto_request_zcopy_complete() decides whether to complete or restart the request. In UCP_ERR_HANDLING_MODE_FAILOVER, an error completion restarts the request, and the restarted send dereferences the NULL rkey in ucp_proto_rndv_put_common_send() → ucp_rkey_get_tl_rkey() → segfault (release builds compile out the ucs_assert guards).

Fix: destroy the rkey only when the request is really completing (success, or non-FAILOVER EP). On the restart path the re-selected protocol instance destroys it on its own completion.

Context

Hit in production (boostrun/bst-use1) during a RoCE fabric congestion event: rail send failed with a hard error → PUT completed with error status → FAILOVER restart segfaulted the worker inside the UCXX progress thread. Backtrace:

ucp_rkey_get_tl_rkey            ucp_rkey.inl:55   <- NULL deref
ucp_proto_rndv_put_common_send  rndv_put.c:65
...
ucp_request_send                ucp_request.inl:330
ucp_proto_request_restart       proto_common.c:848
ucp_proto_multi_handle_send_error
...
ucp_proto_rndv_handle_rtr       proto_rndv.c:952

Same bug exists on upstream master; will report/port upstream separately.

Testing

Full release build (--with-verbs --with-cuda) passes with the patch on v1.21.x.

[wilsonliu-b10]

Independently traced the same chain and confirm the diagnosis — ucp_proto_rndv_rkey_destroy ends with req->send.rndv.rkey = NULL, put's complete calls it before ucp_proto_request_zcopy_complete_cb, and that hook restarts on error under FAILOVER, so the restarted ucp_proto_rndv_put_common_send reads a NULL rkey. Good find, and thanks for the backtrace.

Two things I'd raise before this lands.

1. The guard is wider than the restart condition, so it leaks the rkey.

The hook restarts only when:

status != UCS_OK && FAILOVER && !(req->send.ep->flags & UCP_EP_FLAG_FAILED)

This PR skips the destroy whenever status != UCS_OK && FAILOVER, without the UCP_EP_FLAG_FAILED term. In the case status != UCS_OK && FAILOVER && ep IS FAILED, the hook takes the else branch and completes the request — but the rkey was never destroyed. I checked what else runs on that path: ucp_request_rndv_flush_complete only decrements the flush-op count, and ucp_request_complete_send doesn't touch send.rndv.rkey, so nothing frees it.

That case isn't hypothetical — it's precisely the endpoint-death path (recovery retries exhausted → ucp_ep_set_lanes_failed_schedule sets UCP_EP_FLAG_FAILED, then in-flight put completions arrive with an error status). So the leak would fire in the same incidents as the segfault.

Suggestion: have the release site ask the same predicate the hook uses, rather than restating a subset of it. I have this locally as a small helper in proto_common.inl:

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);
}

used by the hook itself and by both schemes. The get path had this right already but open-coded the condition, which is exactly how the put path came to be written without it — one predicate keeps them from drifting again. Happy to hand that over or open it separately, whichever you prefer.

2. Base branch: is v1.21.x where this is reachable?

v1.21.x has no UCP_PROTO_COMMON_INIT_FLAG_FAILOVER in rndv_put.c at all, so put isn't admitted to failover there and the protocol shouldn't be selectable on that path. master does have the flag (from #1, merged 2026-08-08) and still has the unguarded destroy — so master looks like the branch carrying the reachable bug.

If production hit this, production is running something with the put admission, and a fix landing only on v1.21.x won't reach it. Worth confirming which lineage the affected build came from — I may be missing a cherry-pick that put the admission on v1.21.x.

[wilsonliu-b10]

Correction to my earlier comment — my first point was wrong, and I should have checked your base before raising it.

I read your guard against master's restart condition, which has three terms:

status != UCS_OK && FAILOVER && !(req->send.ep->flags & UCP_EP_FLAG_FAILED)

On v1.21.x it has two — there is no UCP_EP_FLAG_FAILED term (proto_common.inl, ucp_proto_request_zcopy_complete). So status == UCS_OK || !FAILOVER is the exact complement of the restart condition on this base, and there is no leak. Your patch is precisely correct for the branch it targets. Sorry for the noise.

My second point still stands as a question rather than a finding: v1.21.x has no UCP_PROTO_COMMON_INIT_FLAG_FAILOVER in rndv_put.c, so I can't see how put gets selected under failover on this base — which makes me unsure whether the build that crashed came from this lineage. If it did, I'm missing where the admission comes from and would like to know; if it didn't, the fix may need to go where that build's source actually lives.

[wilsonliu-b10]

Answering my own question about the lineage, since I found it: TRT-LLM builds UCX from this fork's master — docker/common/install_ucx.sh pinned 7840331732, which is the merge of the rendezvous failover work. That commit admits put to failover and still has the unguarded release, so the crashing production build came from master, not from v1.21.x.

That pin is now moved past the fix on our side. Flagging it here in case v1.21.x isn't actually feeding a build that can hit this — if nothing consumes that branch with the put admission, this PR may not be needed at all.