Add a refcount to struct io_zcrx_ifq to track the number of rings that
share it. For now, this is only ever 1 i.e. not shared, but will be
larger once shared ifqs are added.

This ref is dec and tested in io_shutdown_zcrx_ifqs() to ensure that an
ifq is not cleaned up while there are still rings using it.

It's important to note that io_shutdown_zcrx_ifqs() may be called in a
loop in io_ring_exit_work() while waiting for ctx->refs to drop to 0.
Use XArray marks to ensure that the refcount dec only happens once.

The cleanup functions io_zcrx_scrub() and io_close_queue() only take ifq
locks and do not need anything from the ring ctx. Therefore it is safe
to call from any ring.

Opted for a bog standard refcount_t. The inc and dec ops are expected to
happen during the slow setup/teardown paths only, and a src ifq is only
expected to be shared a handful of times at most.

Signed-off-by: David Wei <dw@davidwei.uk>
---
 io_uring/zcrx.c | 18 ++++++++++++++++--
 io_uring/zcrx.h |  2 ++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/io_uring/zcrx.c b/io_uring/zcrx.c
index a816f5902091..569cc0338acb 100644
--- a/io_uring/zcrx.c
+++ b/io_uring/zcrx.c
@@ -587,6 +587,7 @@ int io_register_zcrx_ifq(struct io_ring_ctx *ctx,
 	if (!ifq)
 		return -ENOMEM;
 	ifq->rq_entries = reg.rq_entries;
+	refcount_set(&ifq->refs, 1);
 
 	scoped_guard(mutex, &ctx->mmap_lock) {
 		/* preallocate id */
@@ -730,8 +731,21 @@ void io_shutdown_zcrx_ifqs(struct io_ring_ctx *ctx)
 	lockdep_assert_held(&ctx->uring_lock);
 
 	xa_for_each(&ctx->zcrx_ctxs, index, ifq) {
-		io_zcrx_scrub(ifq);
-		io_close_queue(ifq);
+		if (xa_get_mark(&ctx->zcrx_ctxs, index, XA_MARK_0))
+			continue;
+
+		/* Safe to clean up from any ring. */
+		if (refcount_dec_and_test(&ifq->refs)) {
+			io_zcrx_scrub(ifq);
+			io_close_queue(ifq);
+		}
+
+		/*
+		 * This is called in a loop in io_ring_exit_work() until
+		 * ctx->refs drops to 0. Use marks to ensure refcounts are only
+		 * decremented once per ifq per ring.
+		 */
+		xa_set_mark(&ctx->zcrx_ctxs, index, XA_MARK_0);
 	}
 }
 
diff --git a/io_uring/zcrx.h b/io_uring/zcrx.h
index 33ef61503092..566d519cbaf6 100644
--- a/io_uring/zcrx.h
+++ b/io_uring/zcrx.h
@@ -60,6 +60,8 @@ struct io_zcrx_ifq {
 	 */
 	struct mutex			pp_lock;
 	struct io_mapped_region		region;
+
+	refcount_t			refs;
 };
 
 #if defined(CONFIG_IO_URING_ZCRX)
-- 
2.47.3