bpf: task work scheduling kfuncs

Implementation of the new bpf_task_work_schedule kfuncs, that let a BPF
program schedule task_work callbacks for a target task:
 * bpf_task_work_schedule_signal() - schedules with TWA_SIGNAL
 * bpf_task_work_schedule_resume() - schedules with TWA_RESUME

Each map value should embed a struct bpf_task_work, which the kernel
side pairs with struct bpf_task_work_kern, containing a pointer to
struct bpf_task_work_ctx, that maintains metadata relevant for the
concrete callback scheduling.

A small state machine and refcounting scheme ensures safe reuse and
teardown. State transitions:
    _______________________________
    |                             |
    v                             |
[standby] ---> [pending] --> [scheduling] --> [scheduled]
    ^                             |________________|_________
    |                                                       |
    |                                                       v
    |                                                   [running]
    |_______________________________________________________|

All states may transition into FREED state:
[pending] [scheduling] [scheduled] [running] [standby] -> [freed]

A FREED terminal state coordinates with map-value
deletion (bpf_task_work_cancel_and_free()).

Scheduling itself is deferred via irq_work to keep the kfunc callable
from NMI context.

Lifetime is guarded with refcount_t + RCU Tasks Trace.

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_acquire_ctx() - Attempts to take ownership of the context,
 pointed by passed struct bpf_task_work, allocates new context if none
 exists yet.
 * 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 successful task work scheduling
 1) bpf_task_work_schedule_* is called from BPF code.
 2) Transition state from STANDBY to PENDING, mark context as owned by
 this task work scheduler
 3) irq_work_queue() schedules bpf_task_work_irq().
 4) Transition state from PENDING to SCHEDULING (noop if transition
 successful)
 5) bpf_task_work_irq() attempts task_work_add(). If successful, state
 transitions to SCHEDULED.
 6) Task work calls bpf_task_work_callback(), which transition state to
 RUNNING.
 7) BPF callback is executed
 8) Context is cleaned up, refcounts released, context state set back to
 STANDBY.

Signed-off-by: Mykyta Yatsenko <yatsenko@meta.com>
Reviewed-by: Andrii Nakryiko <andrii@kernel.org>
Reviewed-by: Eduard Zingerman <eddyz87@gmail.com>
Acked-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
Link: https://lore.kernel.org/r/20250923112404.668720-8-mykyta.yatsenko5@gmail.com
Signed-off-by: Alexei Starovoitov <ast@kernel.org>

diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
index c28f48a310e353cd85eb1e2df48099bd09f243a0..c9fab9a356dfc18b0eb7fcfd435c9c9372add77d 100644 (file)
--- a/kernel/bpf/helpers.c
+++ b/kernel/bpf/helpers.c
@@ -26,6 +26,8 @@
 #include <linux/bpf_verifier.h>
 #include <linux/uaccess.h>
 #include <linux/verification.h>
+#include <linux/task_work.h>
+#include <linux/irq_work.h>
 
 #include "../../lib/kstrtox.h"
 
@@ -3904,6 +3906,265 @@ __bpf_kfunc int bpf_verify_pkcs7_signature(struct bpf_dynptr *data_p,
 
 typedef int (*bpf_task_work_callback_t)(struct bpf_map *map, void *key, void *value);
 
+enum bpf_task_work_state {
+       /* bpf_task_work is ready to be used */
+       BPF_TW_STANDBY = 0,
+       /* irq work scheduling in progress */
+       BPF_TW_PENDING,
+       /* task work scheduling in progress */
+       BPF_TW_SCHEDULING,
+       /* task work is scheduled successfully */
+       BPF_TW_SCHEDULED,
+       /* callback is running */
+       BPF_TW_RUNNING,
+       /* associated BPF map value is deleted */
+       BPF_TW_FREED,
+};
+
+struct bpf_task_work_ctx {
+       enum bpf_task_work_state state;
+       refcount_t refcnt;
+       struct callback_head work;
+       struct irq_work irq_work;
+       /* bpf_prog that schedules task work */
+       struct bpf_prog *prog;
+       /* task for which callback is scheduled */
+       struct task_struct *task;
+       /* the map and map value associated with this context */
+       struct bpf_map *map;
+       void *map_val;
+       enum task_work_notify_mode mode;
+       bpf_task_work_callback_t callback_fn;
+       struct rcu_head rcu;
+} __aligned(8);
+
+/* Actual type for struct bpf_task_work */
+struct bpf_task_work_kern {
+       struct bpf_task_work_ctx *ctx;
+};
+
+static void bpf_task_work_ctx_reset(struct bpf_task_work_ctx *ctx)
+{
+       if (ctx->prog) {
+               bpf_prog_put(ctx->prog);
+               ctx->prog = NULL;
+       }
+       if (ctx->task) {
+               bpf_task_release(ctx->task);
+               ctx->task = NULL;
+       }
+}
+
+static bool bpf_task_work_ctx_tryget(struct bpf_task_work_ctx *ctx)
+{
+       return refcount_inc_not_zero(&ctx->refcnt);
+}
+
+static void bpf_task_work_ctx_put(struct bpf_task_work_ctx *ctx)
+{
+       if (!refcount_dec_and_test(&ctx->refcnt))
+               return;
+
+       bpf_task_work_ctx_reset(ctx);
+
+       /* bpf_mem_free expects migration to be disabled */
+       migrate_disable();
+       bpf_mem_free(&bpf_global_ma, ctx);
+       migrate_enable();
+}
+
+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 will inevitably run or has already completed
+        * running, and it would have taken care of its ctx ref itself.
+        */
+       if (task_work_cancel(ctx->task, &ctx->work))
+               bpf_task_work_ctx_put(ctx);
+}
+
+static void bpf_task_work_callback(struct callback_head *cb)
+{
+       struct bpf_task_work_ctx *ctx = container_of(cb, struct bpf_task_work_ctx, work);
+       enum bpf_task_work_state state;
+       u32 idx;
+       void *key;
+
+       /* Read lock is needed to protect ctx and map key/value access */
+       guard(rcu_tasks_trace)();
+       /*
+        * This callback may start running before bpf_task_work_irq() switched to
+        * SCHEDULED state, so handle both transition variants SCHEDULING|SCHEDULED -> RUNNING.
+        */
+       state = cmpxchg(&ctx->state, BPF_TW_SCHEDULING, BPF_TW_RUNNING);
+       if (state == BPF_TW_SCHEDULED)
+               state = cmpxchg(&ctx->state, BPF_TW_SCHEDULED, BPF_TW_RUNNING);
+       if (state == BPF_TW_FREED) {
+               bpf_task_work_ctx_put(ctx);
+               return;
+       }
+
+       key = (void *)map_key_from_value(ctx->map, ctx->map_val, &idx);
+
+       migrate_disable();
+       ctx->callback_fn(ctx->map, key, ctx->map_val);
+       migrate_enable();
+
+       bpf_task_work_ctx_reset(ctx);
+       (void)cmpxchg(&ctx->state, BPF_TW_RUNNING, BPF_TW_STANDBY);
+
+       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;
+       }
+
+       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);
+               bpf_task_work_ctx_put(ctx);
+               return;
+       }
+
+       /*
+        * It's technically possible for just scheduled task_work callback to
+        * complete running by now, going SCHEDULING -> RUNNING and then
+        * dropping its ctx refcount. Instead of capturing extra ref just to
+        * protected below ctx->state access, we rely on RCU protection to
+        * perform below SCHEDULING -> SCHEDULED attempt.
+        */
+       state = cmpxchg(&ctx->state, BPF_TW_SCHEDULING, BPF_TW_SCHEDULED);
+       if (state == BPF_TW_FREED)
+               bpf_task_work_cancel(ctx); /* clean up if we switched into FREED state */
+}
+
+static struct bpf_task_work_ctx *bpf_task_work_fetch_ctx(struct bpf_task_work *tw,
+                                                        struct bpf_map *map)
+{
+       struct bpf_task_work_kern *twk = (void *)tw;
+       struct bpf_task_work_ctx *ctx, *old_ctx;
+
+       ctx = READ_ONCE(twk->ctx);
+       if (ctx)
+               return ctx;
+
+       ctx = bpf_mem_alloc(&bpf_global_ma, sizeof(struct bpf_task_work_ctx));
+       if (!ctx)
+               return ERR_PTR(-ENOMEM);
+
+       memset(ctx, 0, sizeof(*ctx));
+       refcount_set(&ctx->refcnt, 1); /* map's own ref */
+       ctx->state = BPF_TW_STANDBY;
+
+       old_ctx = cmpxchg(&twk->ctx, NULL, ctx);
+       if (old_ctx) {
+               /*
+                * tw->ctx is set by concurrent BPF program, release allocated
+                * memory and try to reuse already set context.
+                */
+               bpf_mem_free(&bpf_global_ma, ctx);
+               return old_ctx;
+       }
+
+       return ctx; /* Success */
+}
+
+static struct bpf_task_work_ctx *bpf_task_work_acquire_ctx(struct bpf_task_work *tw,
+                                                          struct bpf_map *map)
+{
+       struct bpf_task_work_ctx *ctx;
+
+       ctx = bpf_task_work_fetch_ctx(tw, map);
+       if (IS_ERR(ctx))
+               return ctx;
+
+       /* try to get ref for task_work callback to hold */
+       if (!bpf_task_work_ctx_tryget(ctx))
+               return ERR_PTR(-EBUSY);
+
+       if (cmpxchg(&ctx->state, BPF_TW_STANDBY, BPF_TW_PENDING) != BPF_TW_STANDBY) {
+               /* lost acquiring race or map_release_uref() stole it from us, put ref and bail */
+               bpf_task_work_ctx_put(ctx);
+               return ERR_PTR(-EBUSY);
+       }
+
+       /*
+        * If no process or bpffs is holding a reference to the map, no new callbacks should be
+        * scheduled. This does not address any race or correctness issue, but rather is a policy
+        * choice: dropping user references should stop everything.
+        */
+       if (!atomic64_read(&map->usercnt)) {
+               /* drop ref we just got for task_work callback itself */
+               bpf_task_work_ctx_put(ctx);
+               /* transfer map's ref into cancel_and_free() */
+               bpf_task_work_cancel_and_free(tw);
+               return ERR_PTR(-EBUSY);
+       }
+
+       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 = -EBADF;
+               goto release_prog;
+       }
+
+       ctx = bpf_task_work_acquire_ctx(tw, map);
+       if (IS_ERR(ctx)) {
+               err = PTR_ERR(ctx);
+               goto release_all;
+       }
+
+       ctx->task = task;
+       ctx->callback_fn = callback_fn;
+       ctx->prog = prog;
+       ctx->mode = mode;
+       ctx->map = map;
+       ctx->map_val = (void *)tw - map->record->task_work_off;
+       init_task_work(&ctx->work, bpf_task_work_callback);
+       init_irq_work(&ctx->irq_work, bpf_task_work_irq);
+
+       irq_work_queue(&ctx->irq_work);
+       return 0;
+
+release_all:
+       bpf_task_release(task);
+release_prog:
+       bpf_prog_put(prog);
+       return err;
+}
+
 /**
  * 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
@@ -3918,7 +4179,7 @@ __bpf_kfunc int bpf_task_work_schedule_signal(struct task_struct *task, struct b
                                              void *map__map, bpf_task_work_callback_t callback,
                                              void *aux__prog)
 {
-       return 0;
+       return bpf_task_work_schedule(task, tw, map__map, callback, aux__prog, TWA_SIGNAL);
 }
 
 /**
@@ -3935,13 +4196,38 @@ __bpf_kfunc int bpf_task_work_schedule_resume(struct task_struct *task, struct b
                                              void *map__map, bpf_task_work_callback_t callback,
                                              void *aux__prog)
 {
-       return 0;
+       return bpf_task_work_schedule(task, tw, map__map, callback, aux__prog, TWA_RESUME);
 }
 
 __bpf_kfunc_end_defs();
 
+static void bpf_task_work_cancel_scheduled(struct irq_work *irq_work)
+{
+       struct bpf_task_work_ctx *ctx = container_of(irq_work, struct bpf_task_work_ctx, irq_work);
+
+       bpf_task_work_cancel(ctx); /* this might put task_work callback's ref */
+       bpf_task_work_ctx_put(ctx); /* and here we put map's own ref that was transferred to us */
+}
+
 void bpf_task_work_cancel_and_free(void *val)
 {
+       struct bpf_task_work_kern *twk = val;
+       struct bpf_task_work_ctx *ctx;
+       enum bpf_task_work_state state;
+
+       ctx = xchg(&twk->ctx, NULL);
+       if (!ctx)
+               return;
+
+       state = xchg(&ctx->state, BPF_TW_FREED);
+       if (state == BPF_TW_SCHEDULED) {
+               /* run in irq_work to avoid locks in NMI */
+               init_irq_work(&ctx->irq_work, bpf_task_work_cancel_scheduled);
+               irq_work_queue(&ctx->irq_work);
+               return;
+       }
+
+       bpf_task_work_ctx_put(ctx); /* put bpf map's ref */
 }
 
 BTF_KFUNCS_START(generic_btf_ids)
@@ -4086,6 +4372,8 @@ BTF_ID_FLAGS(func, bpf_strnstr);
 BTF_ID_FLAGS(func, bpf_cgroup_read_xattr, KF_RCU)
 #endif
 BTF_ID_FLAGS(func, bpf_stream_vprintk, KF_TRUSTED_ARGS)
+BTF_ID_FLAGS(func, bpf_task_work_schedule_signal, KF_TRUSTED_ARGS)
+BTF_ID_FLAGS(func, bpf_task_work_schedule_resume, KF_TRUSTED_ARGS)
 BTF_KFUNCS_END(common_btf_ids)
 
 static const struct btf_kfunc_id_set common_kfunc_set = {