io_uring: ensure IORING_REGISTER_IOWQ_MAX_WORKERS works with SQPOLL

SQPOLL has a different thread doing submissions, we need to check for
that and use the right task context when updating the worker values.
Just hold the sqd->lock across the operation, this ensures that the
thread cannot go away while we poke at ->io_uring.

Link: https://github.com/axboe/liburing/issues/420
Fixes: 2e480058ddc2 ("io-wq: provide a way to limit max number of workers")
Reported-by: Johannes Lundberg <johalun0@gmail.com>
Tested-by: Johannes Lundberg <johalun0@gmail.com>
Signed-off-by: Jens Axboe <axboe@kernel.dk>

diff --git a/fs/io_uring.c b/fs/io_uring.c
index 4a5eb9e..4ad0d17 100644 (file)
--- a/fs/io_uring.c
+++ b/fs/io_uring.c
@@ -10323,26 +10323,46 @@ static int io_unregister_iowq_aff(struct io_ring_ctx *ctx)
 static int io_register_iowq_max_workers(struct io_ring_ctx *ctx,
                                        void __user *arg)
 {
-       struct io_uring_task *tctx = current->io_uring;
+       struct io_uring_task *tctx = NULL;
+       struct io_sq_data *sqd = NULL;
        __u32 new_count[2];
        int i, ret;
 
-       if (!tctx || !tctx->io_wq)
-               return -EINVAL;
        if (copy_from_user(new_count, arg, sizeof(new_count)))
                return -EFAULT;
        for (i = 0; i < ARRAY_SIZE(new_count); i++)
                if (new_count[i] > INT_MAX)
                        return -EINVAL;
 
+       if (ctx->flags & IORING_SETUP_SQPOLL) {
+               sqd = ctx->sq_data;
+               if (sqd) {
+                       mutex_lock(&sqd->lock);
+                       tctx = sqd->thread->io_uring;
+               }
+       } else {
+               tctx = current->io_uring;
+       }
+
+       ret = -EINVAL;
+       if (!tctx || !tctx->io_wq)
+               goto err;
+
        ret = io_wq_max_workers(tctx->io_wq, new_count);
        if (ret)
-               return ret;
+               goto err;
+
+       if (sqd)
+               mutex_unlock(&sqd->lock);
 
        if (copy_to_user(arg, new_count, sizeof(new_count)))
                return -EFAULT;
 
        return 0;
+err:
+       if (sqd)
+               mutex_unlock(&sqd->lock);
+       return ret;
 }
 
 static bool io_register_op_must_quiesce(int op)