[PATCH 8/8] io_uring: enable IORING_SETUP_ATTACH_WQ to attach to SQPOLL thread too

[Jens Axboe <axboe@xxxxxxxxx>]

We support using IORING_SETUP_ATTACH_WQ to share async backends between
rings created by the same process, this now also allows the same to
happen with SQPOLL. The setup procedure remains the same, the caller
sets io_uring_params->wq_fd to the 'parent' context, and then the newly
created ring will attach to that async backend.

Signed-off-by: Jens Axboe <axboe@xxxxxxxxx>
---
 fs/io_uring.c | 43 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/fs/io_uring.c b/fs/io_uring.c
index 5bafc7a2c65c..07e16049e62d 100644
--- a/fs/io_uring.c
+++ b/fs/io_uring.c
@@ -232,6 +232,10 @@ struct io_restriction {
 struct io_sq_data {
 	refcount_t		refs;
 
+	/* global sqd lookup */
+	struct list_head	all_sqd_list;
+	int			attach_fd;
+
 	/* ctx's that are using this sqd */
 	struct list_head	ctx_list;
 	struct list_head	ctx_new_list;
@@ -241,6 +245,9 @@ struct io_sq_data {
 	struct wait_queue_head	wait;
 };
 
+static LIST_HEAD(sqd_list);
+static DEFINE_MUTEX(sqd_lock);
+
 struct io_ring_ctx {
 	struct {
 		struct percpu_ref	refs;
@@ -6975,14 +6982,38 @@ static void io_put_sq_data(struct io_sq_data *sqd)
 			kthread_stop(sqd->thread);
 		}
 
+		mutex_lock(&sqd_lock);
+		list_del(&sqd->all_sqd_list);
+		mutex_unlock(&sqd_lock);
+
 		kfree(sqd);
 	}
 }
 
+static struct io_sq_data *io_attach_sq_data(struct io_uring_params *p)
+{
+	struct io_sq_data *sqd, *ret = ERR_PTR(-ENXIO);
+
+	mutex_lock(&sqd_lock);
+	list_for_each_entry(sqd, &sqd_list, all_sqd_list) {
+		if (sqd->attach_fd == p->wq_fd) {
+			refcount_inc(&sqd->refs);
+			ret = sqd;
+			break;
+		}
+	}
+	mutex_unlock(&sqd_lock);
+
+	return ret;
+}
+
 static struct io_sq_data *io_get_sq_data(struct io_uring_params *p)
 {
 	struct io_sq_data *sqd;
 
+	if (p->flags & IORING_SETUP_ATTACH_WQ)
+		return io_attach_sq_data(p);
+
 	sqd = kzalloc(sizeof(*sqd), GFP_KERNEL);
 	if (!sqd)
 		return ERR_PTR(-ENOMEM);
@@ -6992,6 +7023,10 @@ static struct io_sq_data *io_get_sq_data(struct io_uring_params *p)
 	INIT_LIST_HEAD(&sqd->ctx_new_list);
 	mutex_init(&sqd->ctx_lock);
 	init_waitqueue_head(&sqd->wait);
+
+	mutex_lock(&sqd_lock);
+	list_add_tail(&sqd->all_sqd_list, &sqd_list);
+	mutex_unlock(&sqd_lock);
 	return sqd;
 }
 
@@ -7675,6 +7710,9 @@ static int io_sq_offload_create(struct io_ring_ctx *ctx,
 		if (!ctx->sq_thread_idle)
 			ctx->sq_thread_idle = HZ;
 
+		if (sqd->thread)
+			goto done;
+
 		if (p->flags & IORING_SETUP_SQ_AFF) {
 			int cpu = p->sq_thread_cpu;
 
@@ -7701,6 +7739,7 @@ static int io_sq_offload_create(struct io_ring_ctx *ctx,
 		goto err;
 	}
 
+done:
 	ret = io_init_wq_offload(ctx, p);
 	if (ret)
 		goto err;
@@ -8831,6 +8870,10 @@ static int io_uring_create(unsigned entries, struct io_uring_params *p,
 	if (ret < 0)
 		goto err;
 
+	if ((ctx->flags & (IORING_SETUP_SQPOLL | IORING_SETUP_ATTACH_WQ)) ==
+	    IORING_SETUP_SQPOLL)
+		ctx->sq_data->attach_fd = ret;
+
 	trace_io_uring_create(ret, ctx, p->sq_entries, p->cq_entries, p->flags);
 	return ret;
 err:
-- 
2.28.0