io_uring/poll: avoid retry wakeup loss when dropping ownership

If a wakeup sets IO_POLL_RETRY_FLAG between the owning task consuming the existing instance of IO_POLL_RETRY_FLAG and the loop check of io_poll_check_events(), then we stall on processing the request until an unrelated wakeup comes along.

Using atomic_cmpxchg instead of atomic_sub_return to guarantee both the count v and poll_refs are 0 and no wake-ups are missed.

---
Validated that the stalls disappeared after applying the fix to a small zcrx testing script + bpftrace using nodev on 7.1 with a VM dropping from ~50 stalls to 0.

https://gist.github.com/jlhe97/cf5affffe75e709ba1627e5e0e6405a1/

Similarly, A/B tested on hw using 2 machines containing bnxt NICs for zcrx and set the needed ethtool flags/rules. Over 10 minute runs the patch dropped from 127 stalls to 0.

---
 io_uring/poll.c | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

index 5447a7c24dce6..44765eafa00dc 100644
@@ -330,7 +330,8 @@ static int io_poll_check_events(struct io_kiocb *req, io_tw_token_t tw)
 		 * Release all references, retry if someone tried to restart
 		 * task_work while we were executing it.
 		 */
-	} while (atomic_sub_return(v, &req->poll_refs) & IO_POLL_REF_MASK);
+		v &= IO_POLL_REF_MASK;
+	} while (atomic_cmpxchg(&req->poll_refs, v, 0) != v);
 
 	io_napi_add(req);
 	return IOU_POLL_NO_ACTION;

---
base-commit: e889b0e4cafddb3d68c3595dc01b7128fd35fbef
change-id: 20261013-pu-cmpxchg-9a70b8d5cd9f

Best regards,

Juanlu Herrero

Applied, thanks!

[1/1] io_uring/poll: avoid retry wakeup loss when dropping ownership

      commit: 937d38dcacd675dc0572d6c9b6743960df3128ed

Best regards,

Jens Axboe