Re: [PATCH bpf-next v3 6/7] bpf: task work scheduling kfuncs

[Eduard Zingerman <eddyz87@gmail.com>]

On Mon, 2025-09-08 at 23:33 -0400, Andrii Nakryiko wrote:
> On Mon, Sep 8, 2025 at 9:13 AM Mykyta Yatsenko
> <mykyta.yatsenko5@gmail.com> wrote:
> > 
> > On 9/6/25 21:22, Eduard Zingerman wrote:
> > > On Fri, 2025-09-05 at 17:45 +0100, Mykyta Yatsenko wrote:
> > > 
> > > [...]
> > > 
> > > > A small state machine and refcounting scheme ensures safe reuse and
> > > > teardown:
> > > >   STANDBY -> PENDING -> SCHEDULING -> SCHEDULED -> RUNNING -> STANDBY
> > > Nit: state machine is actually a bit more complex:
> > > 
> > >    digraph G {
> > >      scheduling  -> running    [label="callback 1"];
> > >      scheduled   -> running    [label="callback 2"];
> > >      running     -> standby    [label="callback 3"];
> > >      pending     -> scheduling [label="irq 1"];
> > >      scheduling  -> standby    [label="irq 2"];
> > >      scheduling  -> scheduled  [label="irq 3"];
> > >      standby     -> pending    [label="acquire_ctx"];
> > > 
> > >      freed      -> freed [label="cancel_and_free"];
> > >      pending    -> freed [label="cancel_and_free"];
> > >      running    -> freed [label="cancel_and_free"];
> > >      scheduled  -> freed [label="cancel_and_free"];
> > >      scheduling -> freed [label="cancel_and_free"];
> > >      standby    -> freed [label="cancel_and_free"];
> > >    }
> > > 
> > > [...]
> > > 
> > I'll update the description to contain proper graph.
> 
> Hm... I like the main linear chain of state transitions as is, tbh.
> It's fundamentally simple and helps get the general picture. Sure,
> there are important details, but I don't think we should overwhelm
> anyone reading with all of this upfront.
> 
> In the above, "callback 1" and so on is not really helpful for
> understanding, IMO.

I'm not suggesting to take the above graphviz spec as a description.
Just point out that current message is misleading.

> I'd just add the note that a) each state can transition to FREED and
> b) with tiny probability we might skip SCHEDULED and go SCHEDULING ->
> RUNNING (extremely unlikely if at all possible, tbh).
> 
> In short, let's not go too detailed here.

As you see fit. Transition graph is easier to read for me, but maybe
other people are better at reading text.

> 
> > > > Flow of successful task work scheduling
> > > >   1) bpf_task_work_schedule_* is called from BPF code.
> > > >   2) Transition state from STANDBY to PENDING, marks context is owned by
> > > >   this task work scheduler
> > > >   3) irq_work_queue() schedules bpf_task_work_irq().
> > > >   4) Transition state from PENDING to SCHEDULING.
> > > >   4) bpf_task_work_irq() attempts task_work_add(). If successful, state
> > > >   transitions to SCHEDULED.
> > > Nit: "4" repeated two times.
> > > 
> > > >   5) Task work calls bpf_task_work_callback(), which transition state to
> > > >   RUNNING.
> > > >   6) BPF callback is executed
> > > >   7) Context is cleaned up, refcounts released, context state set back to
> > > >   STANDBY.
> > > > 
> > > > Signed-off-by: Mykyta Yatsenko <yatsenko@meta.com>
> > > > ---
> > > >   kernel/bpf/helpers.c | 319 ++++++++++++++++++++++++++++++++++++++++++-
> > > >   1 file changed, 317 insertions(+), 2 deletions(-)
> > > > 
> > > > diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
> > > > index 109cb249e88c..418a0a211699 100644
> > > > --- a/kernel/bpf/helpers.c
> > > > +++ b/kernel/bpf/helpers.c
> > > [...]
> > > 
> > > > +static void bpf_task_work_cancel(struct bpf_task_work_ctx *ctx)
> > > > +{
> > > > +    /*
> > > > +     * Scheduled task_work callback holds ctx ref, so if we successfully
> > > > +     * cancelled, we put that ref on callback's behalf. If we couldn't
> > > > +     * cancel, callback is inevitably run or has already completed
> > > > +     * running, and it would have taken care of its ctx ref itself.
> > > > +     */
> > > > +    if (task_work_cancel_match(ctx->task, task_work_match, ctx))
> > > Will `task_work_cancel(ctx->task, ctx->work)` do the same thing here?
> > I think so, yes, thanks for checking.
> > > 
> > > > +            bpf_task_work_ctx_put(ctx);
> > > > +}
> > > [...]
> > > 
> > > > +static void bpf_task_work_irq(struct irq_work *irq_work)
> > > > +{
> > > > +    struct bpf_task_work_ctx *ctx = container_of(irq_work, struct bpf_task_work_ctx, irq_work);
> > > > +    enum bpf_task_work_state state;
> > > > +    int err;
> > > > +
> > > > +    guard(rcu_tasks_trace)();
> > > > +
> > > > +    if (cmpxchg(&ctx->state, BPF_TW_PENDING, BPF_TW_SCHEDULING) != BPF_TW_PENDING) {
> > > > +            bpf_task_work_ctx_put(ctx);
> > > > +            return;
> > > > +    }
> > > Why are separate PENDING and SCHEDULING states needed?
> > > Both indicate that the task had not been yet but is ready to be
> > > submitted to task_work_add(). So, on a first glance it seems that
> > > merging the two won't change the behaviour, what do I miss?
> > Yes, this is right, we may drop SCHEDULING state, it does not change any
> > behavior compared to PENDING.
> > The state check before task_work_add is needed anyway, so we won't
> > remove much code here.
> > I kept it just to be more consistent: with every state check we also
> > transition state machine forward.
> 
> Yeah, I like this property as well, I think it makes it easier to
> reason about all this. I'd keep the PENDING and SCHEDULING
> distinction, unless there is a strong reason not to.
> 
> It also gives us a natural point to check for FREED before doing
> unnecessary task_work scheduling + cancelling (if we were already in
> FREED). It doesn't seem like we'll simplify anything by SCHEDULING (or
> PENDING) state.

Again, people are probably different, but it took me some time trying
to figure out if I'm missing some details or SCHEDULING is there just
for the sake of it.

> > > 
> > > > +    err = task_work_add(ctx->task, &ctx->work, ctx->mode);
> > > > +    if (err) {
> > > > +            bpf_task_work_ctx_reset(ctx);
> > > > +            /*
> > > > +             * try to switch back to STANDBY for another task_work reuse, but we might have
> > > > +             * gone to FREED already, which is fine as we already cleaned up after ourselves
> > > > +             */
> > > > +            (void)cmpxchg(&ctx->state, BPF_TW_SCHEDULING, BPF_TW_STANDBY);
> > > > +
> > > > +            /* we don't have RCU protection, so put after switching state */
> > > > +            bpf_task_work_ctx_put(ctx);
> > > > +    }
> > > > +
> 
> [...]
> 
> > > > +
> > > > +    return ctx;
> > > > +}
> > > > +
> > > > +static int bpf_task_work_schedule(struct task_struct *task, struct bpf_task_work *tw,
> > > > +                              struct bpf_map *map, bpf_task_work_callback_t callback_fn,
> > > > +                              struct bpf_prog_aux *aux, enum task_work_notify_mode mode)
> > > > +{
> > > > +    struct bpf_prog *prog;
> > > > +    struct bpf_task_work_ctx *ctx;
> > > > +    int err;
> > > > +
> > > > +    BTF_TYPE_EMIT(struct bpf_task_work);
> > > > +
> > > > +    prog = bpf_prog_inc_not_zero(aux->prog);
> > > > +    if (IS_ERR(prog))
> > > > +            return -EBADF;
> > > > +    task = bpf_task_acquire(task);
> > > > +    if (!task) {
> > > > +            err = -EPERM;
> > > Nit: Why -EPERM? bpf_task_acquire() returns NULL if task->rcu_users
> > >       is zero, does not seem to be permission related.
> > Right, this probably should be -EBADF.
> 
> timer and wq (and now task_work) return -EPERM for map->usercnt==0
> check, but we have -EBADF for prog refcount being zero. It's a bit all
> over the place... I don't have a strong preference, but it would be
> nice to stay more or less consistent for all these "it's too late"
> conditions, IMO.

Ok, probably being consistently wrong is better option here,
let's stick with -EPERM.

[...]