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

[Pavel Begunkov <asml.silence@gmail.com>]

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);

-- 
Pavel Begunkov