io_uring: control ->async_data with a REQ_F flag

->async_data is a slow path, so it won't matter much if we do the clean
up inside io_clean_op(). Moreover, in many cases it's allocated together
with setting one or more of IO_REQ_CLEAN_FLAGS flags, so it'd go through
io_clean_op() anyway.

Control ->async_data allocation with a new flag REQ_F_ASYNC_DATA, so we
can do all the maintainence under io_req_needs_clean() fast check.

Signed-off-by: Pavel Begunkov <asml.silence@gmail.com>
Link: https://lore.kernel.org/r/6892cf5883c459f36bda26f30ceb16742b20b84b.1633373302.git.asml.silence@gmail.com
Signed-off-by: Jens Axboe <axboe@kernel.dk>

diff --git a/fs/io_uring.c b/fs/io_uring.c
index 1b007f77fe52a10c4e5225dd74f980ccd190028b..aee8ecc09168b8c7f6a37d3abcd93df47e498183 100644 (file)
--- a/fs/io_uring.c
+++ b/fs/io_uring.c
@@ -109,7 +109,8 @@
 #define SQE_VALID_FLAGS        (SQE_COMMON_FLAGS|IOSQE_BUFFER_SELECT|IOSQE_IO_DRAIN)
 
 #define IO_REQ_CLEAN_FLAGS (REQ_F_BUFFER_SELECTED | REQ_F_NEED_CLEANUP | \
-                               REQ_F_POLLED | REQ_F_INFLIGHT | REQ_F_CREDS)
+                               REQ_F_POLLED | REQ_F_INFLIGHT | REQ_F_CREDS | \
+                               REQ_F_ASYNC_DATA)
 
 #define IO_TCTX_REFS_CACHE_NR  (1U << 10)
 
@@ -732,6 +733,7 @@ enum {
        REQ_F_CREDS_BIT,
        REQ_F_REFCOUNT_BIT,
        REQ_F_ARM_LTIMEOUT_BIT,
+       REQ_F_ASYNC_DATA_BIT,
        /* keep async read/write and isreg together and in order */
        REQ_F_NOWAIT_READ_BIT,
        REQ_F_NOWAIT_WRITE_BIT,
@@ -787,6 +789,8 @@ enum {
        REQ_F_REFCOUNT          = BIT(REQ_F_REFCOUNT_BIT),
        /* there is a linked timeout that has to be armed */
        REQ_F_ARM_LTIMEOUT      = BIT(REQ_F_ARM_LTIMEOUT_BIT),
+       /* ->async_data allocated */
+       REQ_F_ASYNC_DATA        = BIT(REQ_F_ASYNC_DATA_BIT),
 };
 
 struct async_poll {
@@ -847,8 +851,6 @@ struct io_kiocb {
                struct io_completion    compl;
        };
 
-       /* opcode allocated if it needs to store data for async defer */
-       void                            *async_data;
        u8                              opcode;
        /* polled IO has completed */
        u8                              iopoll_completed;
@@ -863,6 +865,8 @@ struct io_kiocb {
        u64                             user_data;
 
        struct percpu_ref               *fixed_rsrc_refs;
+       /* store used ubuf, so we can prevent reloading */
+       struct io_mapped_ubuf           *imu;
 
        /* used by request caches, completion batching and iopoll */
        struct io_wq_work_node          comp_list;
@@ -872,8 +876,9 @@ struct io_kiocb {
        struct hlist_node               hash_node;
        /* internal polling, see IORING_FEAT_FAST_POLL */
        struct async_poll               *apoll;
-       /* store used ubuf, so we can prevent reloading */
-       struct io_mapped_ubuf           *imu;
+
+       /* opcode allocated if it needs to store data for async defer */
+       void                            *async_data;
        struct io_wq_work               work;
        /* custom credentials, valid IFF REQ_F_CREDS is set */
        const struct cred               *creds;
@@ -1219,6 +1224,11 @@ static bool io_match_task(struct io_kiocb *head, struct task_struct *task,
        return false;
 }
 
+static inline bool req_has_async_data(struct io_kiocb *req)
+{
+       return req->flags & REQ_F_ASYNC_DATA;
+}
+
 static inline void req_set_fail(struct io_kiocb *req)
 {
        req->flags |= REQ_F_FAIL;
@@ -1969,10 +1979,6 @@ static inline void io_dismantle_req(struct io_kiocb *req)
                io_put_file(req->file);
        if (req->fixed_rsrc_refs)
                percpu_ref_put(req->fixed_rsrc_refs);
-       if (req->async_data) {
-               kfree(req->async_data);
-               req->async_data = NULL;
-       }
 }
 
 static __cold void __io_free_req(struct io_kiocb *req)
@@ -2566,7 +2572,7 @@ static bool io_resubmit_prep(struct io_kiocb *req)
 {
        struct io_async_rw *rw = req->async_data;
 
-       if (!rw)
+       if (!req_has_async_data(req))
                return !io_req_prep_async(req);
        iov_iter_restore(&rw->iter, &rw->iter_state);
        return true;
@@ -2878,7 +2884,7 @@ static void kiocb_done(struct kiocb *kiocb, ssize_t ret,
        struct io_async_rw *io = req->async_data;
 
        /* add previously done IO, if any */
-       if (io && io->bytes_done > 0) {
+       if (req_has_async_data(req) && io->bytes_done > 0) {
                if (ret < 0)
                        ret = io->bytes_done;
                else
@@ -3257,7 +3263,11 @@ static inline bool io_alloc_async_data(struct io_kiocb *req)
 {
        WARN_ON_ONCE(!io_op_defs[req->opcode].async_size);
        req->async_data = kmalloc(io_op_defs[req->opcode].async_size, GFP_KERNEL);
-       return req->async_data == NULL;
+       if (req->async_data) {
+               req->flags |= REQ_F_ASYNC_DATA;
+               return false;
+       }
+       return true;
 }
 
 static int io_setup_async_rw(struct io_kiocb *req, const struct iovec *iovec,
@@ -3266,7 +3276,7 @@ static int io_setup_async_rw(struct io_kiocb *req, const struct iovec *iovec,
 {
        if (!force && !io_op_defs[req->opcode].needs_async_setup)
                return 0;
-       if (!req->async_data) {
+       if (!req_has_async_data(req)) {
                struct io_async_rw *iorw;
 
                if (io_alloc_async_data(req)) {
@@ -3399,12 +3409,13 @@ static int io_read(struct io_kiocb *req, unsigned int issue_flags)
        struct iovec inline_vecs[UIO_FASTIOV], *iovec = inline_vecs;
        struct kiocb *kiocb = &req->rw.kiocb;
        struct iov_iter __iter, *iter = &__iter;
-       struct io_async_rw *rw = req->async_data;
        bool force_nonblock = issue_flags & IO_URING_F_NONBLOCK;
        struct iov_iter_state __state, *state;
+       struct io_async_rw *rw;
        ssize_t ret, ret2;
 
-       if (rw) {
+       if (req_has_async_data(req)) {
+               rw = req->async_data;
                iter = &rw->iter;
                state = &rw->iter_state;
                /*
@@ -3534,12 +3545,13 @@ static int io_write(struct io_kiocb *req, unsigned int issue_flags)
        struct iovec inline_vecs[UIO_FASTIOV], *iovec = inline_vecs;
        struct kiocb *kiocb = &req->rw.kiocb;
        struct iov_iter __iter, *iter = &__iter;
-       struct io_async_rw *rw = req->async_data;
        bool force_nonblock = issue_flags & IO_URING_F_NONBLOCK;
        struct iov_iter_state __state, *state;
+       struct io_async_rw *rw;
        ssize_t ret, ret2;
 
-       if (rw) {
+       if (req_has_async_data(req)) {
+               rw = req->async_data;
                iter = &rw->iter;
                state = &rw->iter_state;
                iov_iter_restore(iter, state);
@@ -4708,8 +4720,9 @@ static int io_sendmsg(struct io_kiocb *req, unsigned int issue_flags)
        if (unlikely(!sock))
                return -ENOTSOCK;
 
-       kmsg = req->async_data;
-       if (!kmsg) {
+       if (req_has_async_data(req)) {
+               kmsg = req->async_data;
+       } else {
                ret = io_sendmsg_copy_hdr(req, &iomsg);
                if (ret)
                        return ret;
@@ -4925,8 +4938,9 @@ static int io_recvmsg(struct io_kiocb *req, unsigned int issue_flags)
        if (unlikely(!sock))
                return -ENOTSOCK;
 
-       kmsg = req->async_data;
-       if (!kmsg) {
+       if (req_has_async_data(req)) {
+               kmsg = req->async_data;
+       } else {
                ret = io_recvmsg_copy_hdr(req, &iomsg);
                if (ret)
                        return ret;
@@ -5117,7 +5131,7 @@ static int io_connect(struct io_kiocb *req, unsigned int issue_flags)
        int ret;
        bool force_nonblock = issue_flags & IO_URING_F_NONBLOCK;
 
-       if (req->async_data) {
+       if (req_has_async_data(req)) {
                io = req->async_data;
        } else {
                ret = move_addr_to_kernel(req->connect.addr,
@@ -5133,7 +5147,7 @@ static int io_connect(struct io_kiocb *req, unsigned int issue_flags)
        ret = __sys_connect_file(req->file, &io->address,
                                        req->connect.addr_len, file_flags);
        if ((ret == -EAGAIN || ret == -EINPROGRESS) && force_nonblock) {
-               if (req->async_data)
+               if (req_has_async_data(req))
                        return -EAGAIN;
                if (io_alloc_async_data(req)) {
                        ret = -ENOMEM;
@@ -5424,7 +5438,10 @@ static void __io_queue_proc(struct io_poll_iocb *poll, struct io_poll_table *pt,
                io_init_poll_iocb(poll, poll_one->events, io_poll_double_wake);
                req_ref_get(req);
                poll->wait.private = req;
+
                *poll_ptr = poll;
+               if (req->opcode == IORING_OP_POLL_ADD)
+                       req->flags |= REQ_F_ASYNC_DATA;
        }
 
        pt->nr_entries++;
@@ -6086,7 +6103,7 @@ static int io_timeout_prep(struct io_kiocb *req, const struct io_uring_sqe *sqe,
        if (unlikely(off && !req->ctx->off_timeout_used))
                req->ctx->off_timeout_used = true;
 
-       if (!req->async_data && io_alloc_async_data(req))
+       if (!req_has_async_data(req) && io_alloc_async_data(req))
                return -ENOMEM;
 
        data = req->async_data;
@@ -6395,7 +6412,7 @@ static int io_req_prep_async(struct io_kiocb *req)
 {
        if (!io_op_defs[req->opcode].needs_async_setup)
                return 0;
-       if (WARN_ON_ONCE(req->async_data))
+       if (WARN_ON_ONCE(req_has_async_data(req)))
                return -EFAULT;
        if (io_alloc_async_data(req))
                return -EAGAIN;
@@ -6538,7 +6555,10 @@ static void io_clean_op(struct io_kiocb *req)
        }
        if (req->flags & REQ_F_CREDS)
                put_cred(req->creds);
-
+       if (req->flags & REQ_F_ASYNC_DATA) {
+               kfree(req->async_data);
+               req->async_data = NULL;
+       }
        req->flags &= ~IO_REQ_CLEAN_FLAGS;
 }