Re: [PATCH 1/3] io_uring: add remote task_work execution helper

[Jens Axboe <axboe@kernel.dk>]

On 3/29/24 9:50 AM, Pavel Begunkov wrote:
> On 3/29/24 13:31, Jens Axboe wrote:
>> On 3/29/24 6:51 AM, Pavel Begunkov wrote:
>>> On 3/28/24 18:52, Jens Axboe wrote:
>>>> All our task_work handling is targeted at the state in the io_kiocb
>>>> itself, which is what it is being used for. However, MSG_RING rolls its
>>>> own task_work handling, ignoring how that is usually done.
>>>>
>>>> In preparation for switching MSG_RING to be able to use the normal
>>>> task_work handling, add io_req_task_work_add_remote() which allows the
>>>> caller to pass in the target io_ring_ctx and task.
>>>>
>>>> Signed-off-by: Jens Axboe <axboe@kernel.dk>
>>>> ---
>>>>    io_uring/io_uring.c | 27 +++++++++++++++++++--------
>>>>    io_uring/io_uring.h |  2 ++
>>>>    2 files changed, 21 insertions(+), 8 deletions(-)
>>>>
>>>> diff --git a/io_uring/io_uring.c b/io_uring/io_uring.c
>>>> index 9978dbe00027..609ff9ea5930 100644
>>>> --- a/io_uring/io_uring.c
>>>> +++ b/io_uring/io_uring.c
>>>> @@ -1241,9 +1241,10 @@ void tctx_task_work(struct callback_head *cb)
>>>>        WARN_ON_ONCE(ret);
>>>>    }
>>>>    -static inline void io_req_local_work_add(struct io_kiocb *req, unsigned tw_flags)
>>>> +static inline void io_req_local_work_add(struct io_kiocb *req,
>>>> +                     struct io_ring_ctx *ctx,
>>>> +                     unsigned tw_flags)
>>>>    {
>>>> -    struct io_ring_ctx *ctx = req->ctx;
>>>>        unsigned nr_wait, nr_tw, nr_tw_prev;
>>>>        unsigned long flags;
>>>>    @@ -1291,9 +1292,10 @@ static inline void io_req_local_work_add(struct io_kiocb *req, unsigned tw_flags
>>>>        wake_up_state(ctx->submitter_task, TASK_INTERRUPTIBLE);
>>>>    }
>>>>    -static void io_req_normal_work_add(struct io_kiocb *req)
>>>> +static void io_req_normal_work_add(struct io_kiocb *req,
>>>> +                   struct task_struct *task)
>>>>    {
>>>> -    struct io_uring_task *tctx = req->task->io_uring;
>>>> +    struct io_uring_task *tctx = task->io_uring;
>>>>        struct io_ring_ctx *ctx = req->ctx;
>>>>        unsigned long flags;
>>>>        bool was_empty;
>>>> @@ -1319,7 +1321,7 @@ static void io_req_normal_work_add(struct io_kiocb *req)
>>>>            return;
>>>>        }
>>>>    -    if (likely(!task_work_add(req->task, &tctx->task_work, ctx->notify_method)))
>>>> +    if (likely(!task_work_add(task, &tctx->task_work, ctx->notify_method)))
>>>>            return;
>>>>          io_fallback_tw(tctx, false);
>>>> @@ -1328,9 +1330,18 @@ static void io_req_normal_work_add(struct io_kiocb *req)
>>>>    void __io_req_task_work_add(struct io_kiocb *req, unsigned flags)
>>>>    {
>>>>        if (req->ctx->flags & IORING_SETUP_DEFER_TASKRUN)
>>>> -        io_req_local_work_add(req, flags);
>>>> +        io_req_local_work_add(req, req->ctx, flags);
>>>> +    else
>>>> +        io_req_normal_work_add(req, req->task);
>>>> +}
>>>> +
>>>> +void io_req_task_work_add_remote(struct io_kiocb *req, struct task_struct *task,
>>>> +                 struct io_ring_ctx *ctx, unsigned flags)
>>>
>>> Urgh, even the declration screams that there is something wrong
>>> considering it _either_ targets @ctx or @task.
>>>
>>> Just pass @ctx, so it either use ctx->submitter_task or
>>> req->task, hmm?
>>
>> I actually since changed the above to use a common helper, so was
>> scratching my head a bit over your comment as it can't really work in
>> that setup without needing to check for whether ->submitter_task is set
>> or not. But I do agree this would be nicer, so I'll just return to using
>> the separate helpers for this and it should fall out nicely. The only
>> odd caller is the MSG_RING side, so makes sense to have it a bit more
>> separate rather than try and fold it in with the regular side of using
>> task_work.
>>
>>> A side note, it's a dangerous game, I told it before. At least
>>> it would've been nice to abuse lockdep in a form of:
>>>
>>> io_req_task_complete(req, tw, ctx) {
>>>      lockdep_assert(req->ctx == ctx);
>>>      ...
>>> }
>>>
>>> but we don't have @ctx there, maybe we'll add it to tw later.
>>
>> Agree, but a separate thing imho.
> 
> It's not in a sense that condition couldn't have happened
> before and the patch opening all possibilities.
> 
> We actually have @ctx via struct io_tctx_node, so considering
> fallback it would probably be:
> 
> lockdep_assert(!current->io_uring ||
>                current->io_uring->ctx == req->ctx);

That's not a bad idea. I did run all the testing verifying the ctx, and
it all appears fine. But adding the check is a good idea in general.
Want to send it?

-- 
Jens Axboe