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

[wilsonliu-b10]

A production worker segfaulted inside the UCXX progress thread during a RoCE fabric event: a rail send failed, the rendezvous PUT completed with an error status, and the FAILOVER restart dereferenced a released remote key.

The bug

ucp_proto_rndv_put_common_complete()
  → ucp_proto_rndv_rkey_destroy(req)          // ends with req->send.rndv.rkey = NULL
  → ucp_proto_rndv_request_zcopy_complete()
      → ucp_proto_request_zcopy_complete_cb()
          status != OK && FAILOVER && !FAILED → ucp_proto_request_restart(req)
              → ucp_proto_rndv_put_common_send() → ucp_rkey_get_tl_rkey(NULL) → segfault

Release builds compile out the ucs_assert at the top of ucp_proto_rndv_rkey_destroy, so this surfaces as a clean NULL dereference rather than an assert.

The put scheme was admitted to failover on the reasoning that its completion already routed through the hook carrying the restart, so no new restart code was needed. That skipped checking what the caller does before invoking the hook. The get scheme had it right and documents why the rkey must survive a restart — but the condition was open-coded in each scheme, so the put path could be written without it.

The fix

Keep the rkey when the completion is going to restart rather than complete.

The guard has to match the restart condition exactly: status != OK && FAILOVER && !(ep->flags & UCP_EP_FLAG_FAILED). Guarding on only status != OK && FAILOVER would stop the crash but skip the release when the endpoint is failed — where the hook completes rather than restarts, and nothing else frees the key (ucp_request_rndv_flush_complete only decrements the flush count, ucp_request_complete_send does not touch it). That would leak the rkey on precisely the endpoint-death path that produces these incidents.

So the predicate is hoisted into ucp_proto_request_is_failover_restart() and asked by all three sites: the hook that performs the restart, and both schemes that must not release what a restart still needs. One predicate means they cannot drift apart again — which is the actual reason this bug existed.

A backstop, because the invariant is subtle

ucp_proto_request_restart re-enters protocol selection, so the request must still own whatever the newly selected protocol will use. The rule is not "never clean up before restarting" — the get scheme deliberately deregisters its datatype iterator first, and reset re-registers it. The rule is release only what the restart can re-acquire, and the peer's rkey, which only the RTS could supply, is the one thing it cannot.

Nothing stated or enforced that. So reset() may now return a status other than UCS_OK/UCS_ERR_CANCELED to mean "cannot resume", and restart aborts the request instead of tripping ucs_assertv_always. Both rndv zcopy protocols answer it by refusing to restart without their remote key.

That assert fired in every build, so this trades a crash for a failed transfer the caller already knows how to retry. With the guard above in place the refusal should never fire; it exists for the next protocol admitted to failover.

Adjacent sites, deliberately unchanged

• rndv_rkey_ptr.c:320 releases the rkey and then calls the restart-capable completion — safe only because the status is hardcoded UCS_OK. One refactor away from the identical crash; the new backstop would catch it.
• proto_rndv.c:995 releases the request ID and then passes a live status through. Restart allocates a fresh ID, so I believe it is fine, but it is the same shape and I have not proven it.
• The RTR paths NULL-guard their release and complete through ucp_proto_rndv_recv_complete rather than the restart hook, so they were never exposed.

Test

libucp builds clean on this base and on the b300 bed's tree. Not yet exercised against a live NIC kill on the put scheme: the failover gtest needs a separate --enable-gtest build, and kill_put_send/kill_put_recv on the two-pod harness is the real validation. Happy to run that before merge if you would rather gate on it than on review.