On Wed, 6 Aug 2025 at 16:46, Mykyta Yatsenko <mykyta.yatsenko5@xxxxxxxxx> wrote:
>
> From: Mykyta Yatsenko <yatsenko@xxxxxxxx>
>
> Implementation of the bpf_task_work_schedule kfuncs.
>
> Main components:
> * struct bpf_task_work_context – Metadata and state management per task
> work.
> * enum bpf_task_work_state – A state machine to serialize work
> scheduling and execution.
> * bpf_task_work_schedule() – The central helper that initiates
> scheduling.
> * bpf_task_work_callback() – Invoked when the actual task_work runs.
> * bpf_task_work_irq() – An intermediate step (runs in softirq context)
> to enqueue task work.
> * bpf_task_work_cancel_and_free() – Cleanup for deleted BPF map entries.
>
> Flow of task work scheduling
> 1) bpf_task_work_schedule_* is called from BPF code.
> 2) Transition state from STANDBY to PENDING.
> 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.
> 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, state set back to
> STANDBY.
>
> Map value deletion
> If map value that contains bpf_task_work_context is deleted, BPF map
> implementation calls bpf_task_work_cancel_and_free().
> Deletion is handled by atomically setting state to FREED and
> releasing references or letting scheduler do that, depending on the
> last state before the deletion:
> * SCHEDULING: release references in bpf_task_work_cancel_and_free(),
> expect bpf_task_work_irq() to cancel task work.
> * SCHEDULED: release references and try to cancel task work in
> bpf_task_work_cancel_and_free().
> * other states: one of bpf_task_work_irq(), bpf_task_work_schedule(),
> bpf_task_work_callback() should cleanup upon detecting the state
> switching to FREED.
>
> The state transitions are controlled with atomic_cmpxchg, ensuring:
> * Only one thread can successfully enqueue work.
> * Proper handling of concurrent deletes (BPF_TW_FREED).
> * Safe rollback if task_work_add() fails.
>
In general I am not sure why we need so many acquire/release pairs.
Why not use test_and_set_bit etc.? Or simply cmpxchg?
What ordering of stores are we depending on that merits
acquire/release ordering?
We should probably document explicitly.
> Signed-off-by: Mykyta Yatsenko <yatsenko@xxxxxxxx>
> ---
> kernel/bpf/helpers.c | 188 ++++++++++++++++++++++++++++++++++++++++++-
> 1 file changed, 186 insertions(+), 2 deletions(-)
>
> diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
> index 516286f67f0d..4c8b1c9be7aa 100644
> --- a/kernel/bpf/helpers.c
> +++ b/kernel/bpf/helpers.c
> @@ -25,6 +25,8 @@
> #include <linux/kasan.h>
> #include <linux/bpf_verifier.h>
> #include <linux/uaccess.h>
> +#include <linux/task_work.h>
> +#include <linux/irq_work.h>
>
> #include "../../lib/kstrtox.h"
>
> @@ -3702,6 +3704,160 @@ __bpf_kfunc int bpf_strstr(const char *s1__ign, const char *s2__ign)
>
> typedef void (*bpf_task_work_callback_t)(struct bpf_map *, void *, void *);
>
> +enum bpf_task_work_state {
> + /* bpf_task_work is ready to be used */
> + BPF_TW_STANDBY = 0,
> + /* bpf_task_work is getting scheduled into irq_work */
> + BPF_TW_PENDING,
> + /* bpf_task_work is in irq_work and getting scheduled into task_work */
> + BPF_TW_SCHEDULING,
> + /* bpf_task_work is scheduled into task_work successfully */
> + BPF_TW_SCHEDULED,
> + /* callback is running */
> + BPF_TW_RUNNING,
> + /* BPF map value storing this bpf_task_work is deleted */
> + BPF_TW_FREED,
> +};
> +
> +struct bpf_task_work_context {
> + /* map that contains this structure in a value */
> + struct bpf_map *map;
> + /* bpf_task_work_state value, representing the state */
> + atomic_t state;
> + /* bpf_prog that schedules task work */
> + struct bpf_prog *prog;
> + /* task for which callback is scheduled */
> + struct task_struct *task;
> + /* notification mode for task work scheduling */
> + enum task_work_notify_mode mode;
> + /* callback to call from task work */
> + bpf_task_work_callback_t callback_fn;
> + struct callback_head work;
> + struct irq_work irq_work;
> +} __aligned(8);
I will echo Alexei's comments about the layout. We cannot inline all
this in map value.
Allocation using an init function or in some control function is
probably the only way.
> +
> +static bool task_work_match(struct callback_head *head, void *data)
> +{
> + struct bpf_task_work_context *ctx = container_of(head, struct bpf_task_work_context, work);
> +
> + return ctx == data;
> +}
> +
> +static void bpf_reset_task_work_context(struct bpf_task_work_context *ctx)
> +{
> + bpf_prog_put(ctx->prog);
> + bpf_task_release(ctx->task);
> + rcu_assign_pointer(ctx->map, NULL);
> +}
> +
> +static void bpf_task_work_callback(struct callback_head *cb)
> +{
> + enum bpf_task_work_state state;
> + struct bpf_task_work_context *ctx;
> + struct bpf_map *map;
> + u32 idx;
> + void *key;
> + void *value;
> +
> + rcu_read_lock_trace();
> + ctx = container_of(cb, struct bpf_task_work_context, work);
> +
> + state = atomic_cmpxchg_acquire(&ctx->state, BPF_TW_SCHEDULING, BPF_TW_RUNNING);
> + if (state == BPF_TW_SCHEDULED)
> + state = atomic_cmpxchg_acquire(&ctx->state, BPF_TW_SCHEDULED, BPF_TW_RUNNING);
> + if (state == BPF_TW_FREED)
> + goto out;
I am leaving out commenting on this, since I expect it to change per
later comments.
> +
> + map = rcu_dereference(ctx->map);
> + if (!map)
> + goto out;
> +
> + value = (void *)ctx - map->record->task_work_off;
> + key = (void *)map_key_from_value(map, value, &idx);
> +
> + migrate_disable();
> + ctx->callback_fn(map, key, value);
> + migrate_enable();
> +
> + /* State is running or freed, either way reset. */
> + bpf_reset_task_work_context(ctx);
> + atomic_cmpxchg_release(&ctx->state, BPF_TW_RUNNING, BPF_TW_STANDBY);
> +out:
> + rcu_read_unlock_trace();
> +}
> +
> +static void bpf_task_work_irq(struct irq_work *irq_work)
> +{
> + struct bpf_task_work_context *ctx;
> + enum bpf_task_work_state state;
> + int err;
> +
> + ctx = container_of(irq_work, struct bpf_task_work_context, irq_work);
> +
> + rcu_read_lock_trace();
What's the idea behind rcu_read_lock_trace? Let's add a comment.
> + state = atomic_cmpxchg_release(&ctx->state, BPF_TW_PENDING, BPF_TW_SCHEDULING);
> + if (state == BPF_TW_FREED) {
> + bpf_reset_task_work_context(ctx);
> + goto out;
> + }
> +
> + err = task_work_add(ctx->task, &ctx->work, ctx->mode);
Racy, SCHEDULING->FREE state claim from cancel_and_free will release ctx->task.
> + if (err) {
> + state = atomic_cmpxchg_acquire(&ctx->state, BPF_TW_SCHEDULING, BPF_TW_PENDING);
Races here look fine, since we don't act on FREED (for this block
atleast), cancel_and_free doesn't act on seeing PENDING,
so there is interlocking.
> + if (state == BPF_TW_SCHEDULING) {
> + bpf_reset_task_work_context(ctx);
> + atomic_cmpxchg_release(&ctx->state, BPF_TW_PENDING, BPF_TW_STANDBY);
> + }
> + goto out;
> + }
> + state = atomic_cmpxchg_release(&ctx->state, BPF_TW_SCHEDULING, BPF_TW_SCHEDULED);
> + if (state == BPF_TW_FREED)
> + task_work_cancel_match(ctx->task, task_work_match, ctx);
It looks like there is a similar race condition here.
If BPF_TW_SCHEDULING is set, cancel_and_free may invoke and attempt
bpf_task_release() from bpf_reset_task_work_context().
Meanwhile, we will access ctx->task here directly after seeing BPF_TW_FREED.
> +out:
> + rcu_read_unlock_trace();
> +}
> +
> +static int bpf_task_work_schedule(struct task_struct *task, struct bpf_task_work_context *ctx,
> + 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;
> +
> + BTF_TYPE_EMIT(struct bpf_task_work);
> +
> + prog = bpf_prog_inc_not_zero(aux->prog);
> + if (IS_ERR(prog))
> + return -EPERM;
> +
> + if (!atomic64_read(&map->usercnt)) {
> + bpf_prog_put(prog);
> + return -EPERM;
> + }
> + task = bpf_task_acquire(task);
> + if (!task) {
> + bpf_prog_put(prog);
> + return -EPERM;
> + }
> +
> + if (atomic_cmpxchg_acquire(&ctx->state, BPF_TW_STANDBY, BPF_TW_PENDING) != BPF_TW_STANDBY) {
If we are reusing map values, wouldn't a freed state stay perpetually freed?
I.e. after the first delete of array elements etc. it becomes useless.
Every array map update would invoke a cancel_and_free.
Who resets it?
> + bpf_task_release(task);
> + bpf_prog_put(prog);
> + return -EBUSY;
> + }
> +
> + ctx->task = task;
> + ctx->callback_fn = callback_fn;
> + ctx->prog = prog;
> + ctx->mode = mode;
> + init_irq_work(&ctx->irq_work, bpf_task_work_irq);
> + init_task_work(&ctx->work, bpf_task_work_callback);
> + rcu_assign_pointer(ctx->map, map);
> +
> + irq_work_queue(&ctx->irq_work);
> +
> + return 0;
> +}
> +
> /**
> * bpf_task_work_schedule_signal - Schedule BPF callback using task_work_add with TWA_SIGNAL mode
> * @task: Task struct for which callback should be scheduled
> @@ -3718,7 +3874,8 @@ __bpf_kfunc int bpf_task_work_schedule_signal(struct task_struct *task,
> bpf_task_work_callback_t callback,
> void *aux__prog)
> {
> - return 0;
> + return bpf_task_work_schedule(task, (struct bpf_task_work_context *)tw, map__map,
> + callback, aux__prog, TWA_SIGNAL);
> }
>
> /**
> @@ -3738,13 +3895,38 @@ __bpf_kfunc int bpf_task_work_schedule_resume(struct task_struct *task,
> bpf_task_work_callback_t callback,
> void *aux__prog)
> {
> - return 0;
> + enum task_work_notify_mode mode;
> +
> + mode = task == current && in_nmi() ? TWA_NMI_CURRENT : TWA_RESUME;
> + return bpf_task_work_schedule(task, (struct bpf_task_work_context *)tw, map__map,
> + callback, aux__prog, mode);
> }
>
> __bpf_kfunc_end_defs();
>
> void bpf_task_work_cancel_and_free(void *val)
> {
> + struct bpf_task_work_context *ctx = val;
> + enum bpf_task_work_state state;
> +
> + state = atomic_xchg(&ctx->state, BPF_TW_FREED);
> + switch (state) {
> + case BPF_TW_SCHEDULED:
> + task_work_cancel_match(ctx->task, task_work_match, ctx);
> + fallthrough;
> + /* Scheduling codepath is trying to schedule task work, reset context here. */
> + case BPF_TW_SCHEDULING:
> + bpf_reset_task_work_context(ctx);
> + break;
> + /* work is not initialized, mark as freed and exit */
> + case BPF_TW_STANDBY:
> + /* The context is in interim state, scheduling logic should cleanup. */
> + case BPF_TW_PENDING:
> + /* Callback is already running, it should reset context upon finishing. */
> + case BPF_TW_RUNNING:
> + default:
> + break;
> + }
> }
>
> BTF_KFUNCS_START(generic_btf_ids)
> @@ -3770,6 +3952,8 @@ BTF_ID_FLAGS(func, bpf_rbtree_first, KF_RET_NULL)
> BTF_ID_FLAGS(func, bpf_rbtree_root, KF_RET_NULL)
> BTF_ID_FLAGS(func, bpf_rbtree_left, KF_RET_NULL)
> BTF_ID_FLAGS(func, bpf_rbtree_right, KF_RET_NULL)
> +BTF_ID_FLAGS(func, bpf_task_work_schedule_signal, KF_TRUSTED_ARGS)
> +BTF_ID_FLAGS(func, bpf_task_work_schedule_resume, KF_TRUSTED_ARGS)
>
> #ifdef CONFIG_CGROUPS
> BTF_ID_FLAGS(func, bpf_cgroup_acquire, KF_ACQUIRE | KF_RCU | KF_RET_NULL)
> --
> 2.50.1
>
>
