Merge branch 'io_uring-exit-cancel.7' into for-next

* io_uring-exit-cancel.7:
  io_uring: wait for in-flight requests on ring release
  io_uring: drop registered files and buffers at release time
  io_uring: run cancelations synchronously on ring release
  io_uring/cancel: cancel and wait for all requests on process exit
  io_uring/notif: count pending zerocopy notifications per ring
  io_uring/uring_cmd: only cancel requests of the given task
  io_uring: put request files before posting the completions
  io_uring/rw: don't reap io-wq IOPOLL completions while io-wq has a reference
  io_uring: post io-wq completions from the last request reference
  io_uring/io-wq: put the request file before posting a completion

11 files changed, 319 insertions, 58 deletions

diff --git a/include/linux/io_uring_types.h b/include/linux/io_uring_types.h
index 39629ee77b91..90fea94ad202 100644
--- a/include/linux/io_uring_types.h
+++ b/include/linux/io_uring_types.h
@@ -522,6 +522,8 @@ struct io_ring_ctx {
 
 	/* protected by ->completion_lock */
 	unsigned			nr_req_allocated;
+	/* pending SEND_ZC notifications, protected by ->uring_lock */
+	unsigned			nr_notifs;
 
 #ifdef CONFIG_NET_RX_BUSY_POLL
 	struct list_head	napi_list;	/* track busy poll napi_id */
diff --git a/io_uring/cancel.c b/io_uring/cancel.c
index 7d7820eab878..911a47075064 100644
--- a/io_uring/cancel.c
+++ b/io_uring/cancel.c
@@ -514,8 +514,9 @@ static __cold bool io_uring_try_cancel_iowq(struct io_ring_ctx *ctx)
 
 __cold bool io_uring_try_cancel_requests(struct io_ring_ctx *ctx,
 					 struct io_uring_task *tctx,
-					 bool cancel_all, bool is_sqpoll_thread)
+					 unsigned int flags)
 {
+	bool cancel_all = flags & IO_CANCEL_ALL;
 	struct io_task_cancel cancel = { .tctx = tctx, .all = cancel_all, };
 	enum io_wq_cancel cret;
 	bool ret = false;
@@ -544,7 +545,7 @@ __cold bool io_uring_try_cancel_requests(struct io_ring_ctx *ctx,
 
 	/* SQPOLL thread does its own polling */
 	if ((!(ctx->flags & IORING_SETUP_SQPOLL) && cancel_all) ||
-	    is_sqpoll_thread) {
+	    (flags & IO_CANCEL_SQPOLL)) {
 		while (!list_empty(&ctx->iopoll_list)) {
 			io_iopoll_try_reap_events(ctx);
 			ret = true;
@@ -560,7 +561,9 @@ __cold bool io_uring_try_cancel_requests(struct io_ring_ctx *ctx,
 	ret |= io_poll_remove_all(ctx, tctx, cancel_all);
 	ret |= io_waitid_remove_all(ctx, tctx, cancel_all);
 	ret |= io_futex_remove_all(ctx, tctx, cancel_all);
-	ret |= io_uring_try_cancel_uring_cmd(ctx, tctx, cancel_all);
+	ret |= io_uring_try_cancel_uring_cmd(ctx, tctx);
+	if (flags & IO_CANCEL_KEEP_TIMEOUTS)
+		cancel_all = false;
 	ret |= io_kill_timeouts(ctx, tctx, cancel_all);
 	mutex_unlock(&ctx->uring_lock);
 	if (tctx)
@@ -576,6 +579,44 @@ static s64 tctx_inflight(struct io_uring_task *tctx, bool tracked)
 }
 
 /*
+ * If true, whole thread group is exiting, at which point no task is left that
+ * can reap completions and care about requests in-flight.
+ */
+static bool io_task_group_exiting(void)
+{
+	return current->signal->flags & SIGNAL_GROUP_EXIT;
+}
+
+/*
+ * Return a count of requests an exiting task should wait for. 
+ */
+static s64 tctx_inflight_exit(struct io_uring_task *tctx)
+{
+	struct io_tctx_node *node;
+	unsigned long index;
+	s64 inflight;
+
+	inflight = tctx_inflight(tctx, false);
+	xa_for_each(&tctx->xa, index, node) {
+		/* unlocked read is fine, the caller re-evaluates until done */
+		inflight -= data_race(node->ctx->nr_notifs);
+		/* takes ->uring_lock, we hold nothing on the group exit path */
+		inflight -= io_timeouts_armed(node->ctx, tctx);
+	}
+	return inflight;
+}
+
+static bool io_tctx_cancel_done(struct io_uring_task *tctx, bool cancel_all,
+				bool group_exit)
+{
+	if (cancel_all)
+		return !tctx_inflight(tctx, false);
+	if (tctx_inflight(tctx, true))
+		return false;
+	return !group_exit || tctx_inflight_exit(tctx) <= 0;
+}
+
+/*
  * Find any io_uring ctx that this task has registered or done IO on, and cancel
  * requests. @sqd should be not-null IFF it's an SQPOLL thread cancellation.
  */
@@ -584,7 +625,9 @@ __cold void io_uring_cancel_generic(bool cancel_all, struct io_sq_data *sqd)
 	struct io_uring_task *tctx = current->io_uring;
 	struct io_ring_ctx *ctx;
 	struct io_tctx_node *node;
+	unsigned int flags = 0;
 	unsigned long index;
+	bool group_exit;
 	s64 inflight;
 	DEFINE_WAIT(wait);
 
@@ -595,18 +638,30 @@ __cold void io_uring_cancel_generic(bool cancel_all, struct io_sq_data *sqd)
 	if (tctx->io_wq)
 		io_wq_exit_start(tctx->io_wq);
 
+	/*
+	 * If a whole thread group is exiting, nobody will look at completions.
+	 * If a single thread is exiting, cancel only those that belong to that
+	 * thread.
+	 */
+	group_exit = !cancel_all && io_task_group_exiting();
+	if (cancel_all)
+		flags = IO_CANCEL_ALL;
+	else if (group_exit)
+		flags = IO_CANCEL_ALL | IO_CANCEL_KEEP_TIMEOUTS;
+	if (sqd)
+		flags |= IO_CANCEL_SQPOLL;
+
 	atomic_inc(&tctx->in_cancel);
 	do {
 		bool loop = false;
+		unsigned int state;
 
 		io_uring_drop_tctx_refs(current);
-		if (!tctx_inflight(tctx, !cancel_all))
+		if (io_tctx_cancel_done(tctx, cancel_all, group_exit))
 			break;
 
 		/* read completions before cancelations */
 		inflight = tctx_inflight(tctx, false);
-		if (!inflight)
-			break;
 
 		if (!sqd) {
 			xa_for_each(&tctx->xa, index, node) {
@@ -615,15 +670,13 @@ __cold void io_uring_cancel_generic(bool cancel_all, struct io_sq_data *sqd)
 					continue;
 				loop |= io_uring_try_cancel_requests(node->ctx,
 							current->io_uring,
-							cancel_all,
-							false);
+							flags);
 			}
 		} else {
 			list_for_each_entry(ctx, &sqd->ctx_list, sqd_list)
 				loop |= io_uring_try_cancel_requests(ctx,
 								     current->io_uring,
-								     cancel_all,
-								     true);
+								     flags);
 		}
 
 		if (loop) {
@@ -631,7 +684,12 @@ __cold void io_uring_cancel_generic(bool cancel_all, struct io_sq_data *sqd)
 			continue;
 		}
 
-		prepare_to_wait(&tctx->wait, &wait, TASK_INTERRUPTIBLE);
+		state = TASK_INTERRUPTIBLE;
+		if (task_sigpending(current))
+			state = TASK_UNINTERRUPTIBLE;
+		if (!cancel_all)
+			state |= TASK_FREEZABLE;
+		prepare_to_wait(&tctx->wait, &wait, state);
 		io_run_task_work();
 		io_uring_drop_tctx_refs(current);
 		xa_for_each(&tctx->xa, index, node) {
@@ -646,8 +704,13 @@ __cold void io_uring_cancel_generic(bool cancel_all, struct io_sq_data *sqd)
 		 * avoids a race where a completion comes in before we did
 		 * prepare_to_wait().
 		 */
-		if (inflight == tctx_inflight(tctx, !cancel_all))
-			schedule();
+		if (inflight == tctx_inflight(tctx, false)) {
+			unsigned long timeout = 1;
+
+			if (state & TASK_INTERRUPTIBLE)
+				timeout = MAX_SCHEDULE_TIMEOUT;
+			schedule_timeout(timeout);
+		}
 end_wait:
 		finish_wait(&tctx->wait, &wait);
 	} while (1);
diff --git a/io_uring/cancel.h b/io_uring/cancel.h
index 1b201a094303..e49713a1bc66 100644
--- a/io_uring/cancel.h
+++ b/io_uring/cancel.h
@@ -30,9 +30,20 @@ bool io_cancel_remove_all(struct io_ring_ctx *ctx, struct io_uring_task *tctx,
 int io_cancel_remove(struct io_ring_ctx *ctx, struct io_cancel_data *cd,
 		     unsigned int issue_flags, struct hlist_head *list,
 		     bool (*cancel)(struct io_kiocb *));
+
+/* io_uring_try_cancel_requests() flags */
+enum {
+	/* match all requests, not just REQ_F_INFLIGHT */
+	IO_CANCEL_ALL		= 1,
+	/* ignore timeouts */
+	IO_CANCEL_KEEP_TIMEOUTS	= 2,
+	/* called by the SQPOLL thread */
+	IO_CANCEL_SQPOLL	= 4,
+};
+
 __cold bool io_uring_try_cancel_requests(struct io_ring_ctx *ctx,
 					 struct io_uring_task *tctx,
-					 bool cancel_all, bool is_sqpoll_thread);
+					 unsigned int flags);
 __cold void io_uring_cancel_generic(bool cancel_all, struct io_sq_data *sqd);
 __cold bool io_cancel_ctx_cb(struct io_wq_work *work, void *data);
 
diff --git a/io_uring/io_uring.c b/io_uring/io_uring.c
index 61053421d809..a67b2adeda36 100644
--- a/io_uring/io_uring.c
+++ b/io_uring/io_uring.c
@@ -907,6 +907,22 @@ bool io_req_post_cqe32(struct io_kiocb *req, struct io_uring_cqe cqe[2])
 	return posted;
 }
 
+/* drop the request file before the CQE is posted, not deferred after it */
+static void io_req_put_file(struct io_kiocb *req, bool sync)
+{
+	struct file *file = req->file;
+
+	if (!file || (req->flags & (REQ_F_FIXED_FILE | REQ_F_REISSUE)))
+		return;
+
+	WRITE_ONCE(req->file, NULL);
+	/* releasing a ring may wait on other rings, keep that deferred */
+	if (sync && !io_is_uring_fops(file))
+		__fput_sync(file);
+	else
+		fput(file);
+}
+
 static void io_req_complete_post(struct io_kiocb *req, unsigned issue_flags)
 {
 	struct io_ring_ctx *ctx = req->ctx;
@@ -919,6 +935,8 @@ static void io_req_complete_post(struct io_kiocb *req, unsigned issue_flags)
 	if (WARN_ON_ONCE(!(issue_flags & IO_URING_F_IOWQ)))
 		return;
 
+	io_req_put_file(req, true);
+
 	/*
 	 * Handle special CQ sync cases via task_work. DEFER_TASKRUN requires
 	 * the submitter task context, IOPOLL protects with uring_lock.
@@ -939,9 +957,11 @@ defer_complete:
 		goto defer_complete;
 
 	/*
-	 * We don't free the request here because we know it's called from
-	 * io-wq only, which holds a reference, so it cannot be the last put.
+	 * Request not freed here because we know it's called from io-wq only,
+	 * which holds a reference. Hence it can't be the last put. The CQE
+	 * has been posted, last put frees it.
 	 */
+	req->flags |= REQ_F_CQE_SKIP;
 	req_ref_put(req);
 }
 
@@ -995,8 +1015,6 @@ __cold void io_free_req(struct io_kiocb *req)
 {
 	/* refs were already put, restore them for io_req_task_complete() */
 	req->flags &= ~REQ_F_REFCOUNT;
-	/* we only want to free it, don't post CQEs */
-	req->flags |= REQ_F_CQE_SKIP;
 	req->io_task_work.func = io_req_task_complete;
 	io_req_task_work_add(req);
 }
@@ -1099,6 +1117,8 @@ static void io_free_batch_list(struct io_ring_ctx *ctx,
 			}
 			if (req->flags & REQ_F_REFCOUNT) {
 				node = req->comp_list.next;
+				/* CQE posted, the last put only frees */
+				req->flags |= REQ_F_CQE_SKIP;
 				if (!req_ref_put_and_test(req))
 					continue;
 			}
@@ -1130,6 +1150,19 @@ void __io_submit_flush_completions(struct io_ring_ctx *ctx)
 	struct io_submit_state *state = &ctx->submit_state;
 	struct io_wq_work_node *node;
 
+	/*
+	 * Drop the files before the CQEs are posted, so they're released by
+	 * the time the completions are visible. Not for requests that are
+	 * requeued or still referenced, those aren't freed below.
+	 */
+	__wq_list_for_each(node, &state->compl_reqs) {
+		struct io_kiocb *req = container_of(node, struct io_kiocb,
+					    comp_list);
+
+		if (!io_req_shared(req))
+			io_req_put_file(req, true);
+	}
+
 	__io_cq_lock(ctx);
 	__wq_list_for_each(node, &state->compl_reqs) {
 		struct io_kiocb *req = container_of(node, struct io_kiocb,
@@ -1262,7 +1295,10 @@ static int io_iopoll_check(struct io_ring_ctx *ctx, unsigned int min_events)
 
 void io_req_task_complete(struct io_tw_req tw_req, io_tw_token_t tw)
 {
-	io_req_complete_defer(tw_req.req);
+	struct io_kiocb *req = tw_req.req;
+
+	if (io_req_complete_ready(req))
+		io_req_complete_defer(req);
 }
 
 /*
@@ -1451,7 +1487,13 @@ struct io_wq_work *io_wq_free_work(struct io_wq_work *work)
 	struct io_kiocb *nxt = NULL;
 
 	if (req_ref_put_and_test_atomic(req)) {
-		if (req->flags & IO_REQ_LINK_FLAGS) {
+		/*
+		 * Only continue the link from here if the completion has been
+		 * posted. If it's still pending, io_free_req() completes the
+		 * request and the next link follows from there, after the CQE.
+		 */
+		if ((req->flags & IO_REQ_LINK_FLAGS) &&
+		    (req->flags & REQ_F_CQE_SKIP)) {
 			struct io_ring_ctx *ctx = req->ctx;
 
 			mutex_lock(&ctx->uring_lock);
@@ -1480,6 +1522,7 @@ void io_wq_submit_work(struct io_wq_work *work)
 	/* either cancelled or io-wq is dying, so don't touch tctx->iowq */
 	if (atomic_read(&work->flags) & IO_WQ_WORK_CANCEL) {
 fail:
+		io_req_put_file(req, false);
 		io_req_task_queue_fail(req, err);
 		return;
 	}
@@ -1555,8 +1598,10 @@ fail:
 	} while (1);
 
 	/* avoid locking problems by failing it from a clean context */
-	if (ret)
+	if (ret) {
+		io_req_put_file(req, true);
 		io_req_task_queue_fail(req, ret);
+	}
 }
 
 inline struct file *io_file_get_fixed(struct io_kiocb *req, int fd,
@@ -2157,8 +2202,6 @@ static __cold void io_ring_ctx_free(struct io_ring_ctx *ctx)
 	io_sq_thread_finish(ctx);
 
 	mutex_lock(&ctx->uring_lock);
-	io_sqe_buffers_unregister(ctx);
-	io_sqe_files_unregister(ctx);
 	io_unregister_zcrx(ctx);
 	io_cqring_overflow_kill(ctx);
 	io_eventfd_unregister(ctx);
@@ -2306,18 +2349,18 @@ static __cold void io_tctx_exit_cb(struct callback_head *cb)
 	complete(&work->completion);
 }
 
-static __cold void io_ring_exit_work(struct work_struct *work)
+/*
+ * Cancel what can be canceled on a dying ring and reap what has completed.
+ * Only waits on polled I/O, never anything else.
+ */
+static __cold void io_ring_ctx_cancel(struct io_ring_ctx *ctx)
 {
-	struct io_ring_ctx *ctx = container_of(work, struct io_ring_ctx, exit_work);
-	unsigned long timeout = jiffies + IO_URING_EXIT_WAIT_MAX;
-	unsigned long interval = HZ / 20;
-	struct io_tctx_exit exit;
-	struct io_tctx_node *node;
-	int ret;
+	struct io_sq_data *sqd = ctx->sq_data;
 
-	mutex_lock(&ctx->uring_lock);
-	io_terminate_zcrx(ctx);
-	mutex_unlock(&ctx->uring_lock);
+	if (test_bit(IO_CHECK_CQ_OVERFLOW_BIT, &ctx->check_cq)) {
+		scoped_guard(mutex, &ctx->uring_lock)
+			io_cqring_overflow_kill(ctx);
+	}
 
 	/*
 	 * If we're doing polled IO and end up having requests being
@@ -2326,32 +2369,106 @@ static __cold void io_ring_exit_work(struct work_struct *work)
 	 * as nobody else will be looking for them.
 	 */
 	do {
-		if (test_bit(IO_CHECK_CQ_OVERFLOW_BIT, &ctx->check_cq)) {
-			mutex_lock(&ctx->uring_lock);
-			io_cqring_overflow_kill(ctx);
-			mutex_unlock(&ctx->uring_lock);
+		if (ctx->flags & IORING_SETUP_DEFER_TASKRUN)
+			io_cancel_local_task_work(ctx);
+		cond_resched();
+	} while (io_uring_try_cancel_requests(ctx, NULL, IO_CANCEL_ALL));
+
+	if (sqd) {
+		struct task_struct *tsk;
+
+		io_sq_thread_park(sqd);
+		tsk = sqpoll_task_locked(sqd);
+		if (tsk && tsk->io_uring && tsk->io_uring->io_wq)
+			io_wq_cancel_cb(tsk->io_uring->io_wq,
+					io_cancel_ctx_cb, ctx, true);
+		io_sq_thread_unpark(sqd);
+	}
+
+	io_req_caches_free(ctx);
+}
+
+/* Number of requests that should be waited for */
+static __cold unsigned int io_ring_ctx_inflight(struct io_ring_ctx *ctx)
+{
+	guard(mutex)(&ctx->uring_lock);
+	__io_req_caches_free(ctx);
+	return ctx->nr_req_allocated - ctx->nr_notifs;
+}
+
+/*
+ * Run task_work completions for current. Only do so if the io_uring callback
+ * itself can get pruned first, otherwise we risk recursing.
+ */
+static __cold bool io_ring_run_own_completions(struct io_uring_task *tctx)
+{
+	unsigned int count = 0;
+
+	if (!tctx || mpscq_empty(&tctx->task_list))
+		return true;
+	if (!task_work_cancel(current, &tctx->task_work))
+		return false;
+	tctx_task_work_run(tctx, UINT_MAX, &count);
+	return true;
+}
+
+/*
+ * Requests may remain after cancelations have been run, as not all requests
+ * are cancelable. Storage I/O is an example. Wait for those so that once
+ * close(2) returns, files pinned by these requests have been released.
+ */
+static __cold void io_ring_ctx_wait_inflight(struct io_ring_ctx *ctx)
+{
+	struct io_uring_task *tctx = current->io_uring;
+	bool ran_own = true;
+
+	if (current->flags & (PF_KTHREAD | PF_EXITING))
+		return;
+	if (tctx && atomic_read(&tctx->in_cancel))
+		return;
+
+	while (io_ring_ctx_inflight(ctx) && !fatal_signal_pending(current)) {
+		unsigned int state;
+
+		if (test_thread_flag(TIF_NOTIFY_SIGNAL)) {
+			clear_notify_signal();
+			if (task_work_pending(current))
+				set_notify_resume(current);
 		}
+		state = TASK_INTERRUPTIBLE;
+		if (signal_pending(current))
+			state = TASK_KILLABLE;
+		set_current_state(state | TASK_FREEZABLE);
+		/* don't sleep on work that's already there and that we can run */
+		if (ran_own && ((tctx && !mpscq_empty(&tctx->task_list)) ||
+		    io_local_work_pending(ctx)))
+			__set_current_state(TASK_RUNNING);
+		else
+			schedule_timeout(1);
 
-		/* The SQPOLL thread never reaches this path */
-		do {
-			if (ctx->flags & IORING_SETUP_DEFER_TASKRUN)
-				io_cancel_local_task_work(ctx);
-			cond_resched();
-		} while (io_uring_try_cancel_requests(ctx, NULL, true, false));
-
-		if (ctx->sq_data) {
-			struct io_sq_data *sqd = ctx->sq_data;
-			struct task_struct *tsk;
-
-			io_sq_thread_park(sqd);
-			tsk = sqpoll_task_locked(sqd);
-			if (tsk && tsk->io_uring && tsk->io_uring->io_wq)
-				io_wq_cancel_cb(tsk->io_uring->io_wq,
-						io_cancel_ctx_cb, ctx, true);
-			io_sq_thread_unpark(sqd);
+		/* completions may be queued behind us */
+		if (!io_ring_run_own_completions(tctx)) {
+			if (!ran_own)
+				break;
+			ran_own = false;
+		} else {
+			ran_own = true;
 		}
+		io_ring_ctx_cancel(ctx);
+	}
+}
 
-		io_req_caches_free(ctx);
+static __cold void io_ring_exit_work(struct work_struct *work)
+{
+	struct io_ring_ctx *ctx = container_of(work, struct io_ring_ctx, exit_work);
+	unsigned long timeout = jiffies + IO_URING_EXIT_WAIT_MAX;
+	unsigned long interval = HZ / 20;
+	struct io_tctx_exit exit;
+	struct io_tctx_node *node;
+	int ret;
+
+	do {
+		io_ring_ctx_cancel(ctx);
 
 		if (WARN_ON_ONCE(time_after(jiffies, timeout))) {
 			/* there is little hope left, don't run it too often */
@@ -2419,8 +2536,23 @@ static __cold void io_ring_ctx_wait_and_kill(struct io_ring_ctx *ctx)
 	percpu_ref_kill(&ctx->refs);
 	xa_for_each(&ctx->personalities, index, creds)
 		io_unregister_personality(ctx, index);
+	io_terminate_zcrx(ctx);
+	/* Drop these now, rather than async, to unpin the files */
+	io_sqe_buffers_unregister(ctx);
+	io_sqe_files_unregister(ctx);
 	mutex_unlock(&ctx->uring_lock);
 
+	/*
+	 * Do the first round of cancelations upfront rather than leaving it
+	 * to to exit_work. Anything cancelable is then already on its way
+	 * out, and for requests owned by the task closing the ring, this
+	 * ensures any held files are put before close(2) returns.
+	 */
+	if (!(current->flags & PF_IO_WORKER)) {
+		io_ring_ctx_cancel(ctx);
+		io_ring_ctx_wait_inflight(ctx);
+	}
+
 	INIT_WORK(&ctx->exit_work, io_ring_exit_work);
 	/*
 	 * Use system_dfl_wq to avoid spawning tons of event kworkers
diff --git a/io_uring/notif.c b/io_uring/notif.c
index efce8ae12eaa..6c49aff515df 100644
--- a/io_uring/notif.c
+++ b/io_uring/notif.c
@@ -36,6 +36,7 @@ static void io_notif_tw_complete(struct io_tw_req tw_req, io_tw_token_t tw)
 		}
 
 		nd = nd->next;
+		ctx->nr_notifs--;
 		io_req_task_complete((struct io_tw_req){notif}, tw);
 	} while (nd);
 }
@@ -119,6 +120,7 @@ struct io_kiocb *io_alloc_notif(struct io_ring_ctx *ctx)
 
 	if (unlikely(!io_alloc_req(ctx, &notif)))
 		return NULL;
+	ctx->nr_notifs++;
 	notif->ctx = ctx;
 	notif->opcode = IORING_OP_NOP;
 	notif->flags = 0;
diff --git a/io_uring/refs.h b/io_uring/refs.h
index 0fe16b67c308..021bf0cd66af 100644
--- a/io_uring/refs.h
+++ b/io_uring/refs.h
@@ -56,6 +56,33 @@ static inline void __io_req_set_refcount(struct io_kiocb *req, int nr)
 	}
 }
 
+/*
+ * io-wq may still hold a reference if the issue completed async. If so, the
+ * last put completes the request, so that file drop and CQE visibility are
+ * ordered. The last reference stays for the free path to put, a linked
+ * timeout may still look at the request until then.
+ */
+static inline bool io_req_complete_ready(struct io_kiocb *req)
+{
+	if (!(req->flags & REQ_F_REFCOUNT) || (req->flags & REQ_F_REISSUE))
+		return true;
+	if (atomic_read(&req->refs) == 1)
+		return true;
+	if (!req_ref_put_and_test(req))
+		return false;
+	/* the other put raced us, ours was the last after all */
+	atomic_set(&req->refs, 1);
+	return true;
+}
+
+/* io-wq still holds a reference, the request can't be completed yet */
+static inline bool io_req_shared(struct io_kiocb *req)
+{
+	if (!(req->flags & REQ_F_REFCOUNT) || (req->flags & REQ_F_REISSUE))
+		return false;
+	return atomic_read(&req->refs) > 1;
+}
+
 static inline void io_req_set_refcount(struct io_kiocb *req)
 {
 	__io_req_set_refcount(req, 1);
diff --git a/io_uring/rw.c b/io_uring/rw.c
index 432820f86251..0c9494fd21be 100644
--- a/io_uring/rw.c
+++ b/io_uring/rw.c
@@ -16,6 +16,7 @@
 
 #include "filetable.h"
 #include "io_uring.h"
+#include "refs.h"
 #include "opdef.h"
 #include "kbuf.h"
 #include "alloc_cache.h"
@@ -1380,6 +1381,9 @@ int io_do_iopoll(struct io_ring_ctx *ctx, bool force_nonspin)
 		/* order with io_complete_rw_iopoll(), e.g. ->result updates */
 		if (!smp_load_acquire(&req->iopoll_completed))
 			continue;
+		/* io-wq still has a reference, reap it on the next pass */
+		if (io_req_shared(req))
+			continue;
 		list_del(&req->iopoll_node);
 		wq_list_add_tail(&req->comp_list, &ctx->submit_state.compl_reqs);
 		nr_events++;
diff --git a/io_uring/timeout.c b/io_uring/timeout.c
index c4dd26cf342d..9c239c0a1715 100644
--- a/io_uring/timeout.c
+++ b/io_uring/timeout.c
@@ -729,6 +729,24 @@ static bool io_match_task(struct io_kiocb *head, struct io_uring_task *tctx,
 	return false;
 }
 
+__cold unsigned int io_timeouts_armed(struct io_ring_ctx *ctx,
+				      struct io_uring_task *tctx)
+{
+	struct io_timeout *timeout;
+	unsigned int nr = 0;
+
+	guard(mutex)(&ctx->uring_lock);
+	raw_spin_lock_irq(&ctx->timeout_lock);
+	list_for_each_entry(timeout, &ctx->timeout_list, list) {
+		struct io_kiocb *req = cmd_to_io_kiocb(timeout);
+
+		if (req->tctx == tctx)
+			nr += io_linked_nr(req);
+	}
+	raw_spin_unlock_irq(&ctx->timeout_lock);
+	return nr;
+}
+
 /* Returns true if we found and killed one or more timeouts */
 __cold bool io_kill_timeouts(struct io_ring_ctx *ctx, struct io_uring_task *tctx,
 			     bool cancel_all)
diff --git a/io_uring/timeout.h b/io_uring/timeout.h
index 1620f94dd45a..b6cd608fb667 100644
--- a/io_uring/timeout.h
+++ b/io_uring/timeout.h
@@ -11,6 +11,8 @@ struct io_timeout_data {
 __cold void io_flush_timeouts(struct io_ring_ctx *ctx);
 struct io_cancel_data;
 int io_timeout_cancel(struct io_ring_ctx *ctx, struct io_cancel_data *cd);
+__cold unsigned int io_timeouts_armed(struct io_ring_ctx *ctx,
+				      struct io_uring_task *tctx);
 __cold bool io_kill_timeouts(struct io_ring_ctx *ctx, struct io_uring_task *tctx,
 			     bool cancel_all);
 void io_queue_linked_timeout(struct io_kiocb *req);
diff --git a/io_uring/uring_cmd.c b/io_uring/uring_cmd.c
index 3d5d8b5f4ebb..73b606899833 100644
--- a/io_uring/uring_cmd.c
+++ b/io_uring/uring_cmd.c
@@ -49,7 +49,7 @@ void io_uring_cmd_cleanup(struct io_kiocb *req)
 }
 
 bool io_uring_try_cancel_uring_cmd(struct io_ring_ctx *ctx,
-				   struct io_uring_task *tctx, bool cancel_all)
+				   struct io_uring_task *tctx)
 {
 	struct hlist_node *tmp;
 	struct io_kiocb *req;
@@ -63,7 +63,7 @@ bool io_uring_try_cancel_uring_cmd(struct io_ring_ctx *ctx,
 				struct io_uring_cmd);
 		struct file *file = req->file;
 
-		if (!cancel_all && req->tctx != tctx)
+		if (tctx && req->tctx != tctx)
 			continue;
 
 		if (cmd->flags & IORING_URING_CMD_CANCELABLE) {
diff --git a/io_uring/uring_cmd.h b/io_uring/uring_cmd.h
index 041aef8a8aa3..a5cb3f2ee1c5 100644
--- a/io_uring/uring_cmd.h
+++ b/io_uring/uring_cmd.h
@@ -14,7 +14,7 @@ void io_uring_cmd_sqe_copy(struct io_kiocb *req);
 void io_uring_cmd_cleanup(struct io_kiocb *req);
 
 bool io_uring_try_cancel_uring_cmd(struct io_ring_ctx *ctx,
-				   struct io_uring_task *tctx, bool cancel_all);
+				   struct io_uring_task *tctx);
 
 bool io_uring_cmd_post_mshot_cqe32(struct io_uring_cmd *cmd,
 				   unsigned int issue_flags,