bpf: task work scheduling kfuncs

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.

Signed-off-by: Mykyta Yatsenko <[email protected]>

Author: mykyta5

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,200 @@ __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,
+	/* task work scheduling in progress */
+	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;
+	/* map value that contains pointer to this context */
+	void *map_val;
+	/* 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;
+	/* bpf_task_work_state value, representing the state */
+	enum bpf_task_work_state state;
+	/* callback to call from task work */
+	bpf_task_work_callback_t callback_fn;
+	struct callback_head work;
+	struct irq_work irq_work;
+} __aligned(8);
+
+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);
+}
+
+static void bpf_task_work_callback(struct callback_head *cb)
+{
+	enum bpf_task_work_state state;
+	struct bpf_task_work_context *ctx;
+	u32 idx;
+	void *key;
+
+	ctx = container_of(cb, struct bpf_task_work_context, work);
+
+	/*
+	 * If map element is freed, context state is set to BPF_TW_FREED
+	 * Otherwise it is safe to access map key value under the rcu_read_lock
+	 */
+	rcu_read_lock_trace();
+	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) {
+		rcu_read_unlock_trace();
+		bpf_reset_task_work_context(ctx);
+		goto free;
+	}
+
+	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();
+
+	rcu_read_unlock_trace();
+	/* State is running or freed, either way reset. */
+	bpf_reset_task_work_context(ctx);
+	state = cmpxchg(&ctx->state, BPF_TW_RUNNING, BPF_TW_STANDBY);
+	if (state == BPF_TW_RUNNING)
+		return; /* Work is done, context is reset */
+
+free:
+	kfree(ctx);
+}
+
+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);
+
+	state = cmpxchg(&ctx->state, BPF_TW_PENDING, BPF_TW_SCHEDULING);
+	if (state == BPF_TW_FREED)
+		goto free_context;
+
+	err = task_work_add(ctx->task, &ctx->work, ctx->mode);
+	if (err) {
+		bpf_reset_task_work_context(ctx);
+		state = cmpxchg(&ctx->state, BPF_TW_SCHEDULING, BPF_TW_STANDBY);
+		if (state == BPF_TW_FREED)
+			kfree(ctx);
+		return;
+	}
+	state = cmpxchg(&ctx->state, BPF_TW_SCHEDULING, BPF_TW_SCHEDULED);
+	if (state == BPF_TW_FREED &&
+	    task_work_cancel_match(ctx->task, task_work_match, ctx))
+		goto free_context;
+	/* otherwise task work is not cancelled, expect callback to cleanup */
+	return;
+
+free_context:
+	bpf_reset_task_work_context(ctx);
+	kfree(ctx);
+}
+
+static struct bpf_task_work_context *
+bpf_task_work_aquire_ctx(struct bpf_task_work *tw, struct bpf_map *map)
+{
+	struct bpf_task_work_context *ctx;
+	enum bpf_task_work_state state;
+	atomic_long_t *ctx_ptr = (atomic_long_t *)&tw->ctx;
+
+	ctx = (void *)atomic_long_read(ctx_ptr);
+	if (!ctx) {
+		ctx = bpf_map_kmalloc_node(map,
+					   sizeof(struct bpf_task_work_context),
+					   GFP_ATOMIC, map->numa_node);
+		if (!ctx)
+			return ERR_PTR(-ENOMEM);
+		memset(ctx, 0, sizeof(*ctx));
+		if (atomic_long_cmpxchg(ctx_ptr, 0, (unsigned long)ctx) != 0) {
+			kfree(ctx);
+			return ERR_PTR(-EBUSY);
+		}
+	}
+	state = cmpxchg(&ctx->state, BPF_TW_STANDBY, BPF_TW_PENDING);
+	if (state != BPF_TW_STANDBY)
+		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_context *ctx = NULL;
+	int err;
+
+	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)) {
+		err = -EPERM;
+		goto release_prog;
+	}
+	task = bpf_task_acquire(task);
+	if (!task) {
+		err = -EPERM;
+		goto release_prog;
+	}
+	ctx = bpf_task_work_aquire_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_irq_work(&ctx->irq_work, bpf_task_work_irq);
+	init_task_work(&ctx->work, bpf_task_work_callback);
+
+	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
@@ -3718,7 +3914,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, tw, map__map,
+				      callback, aux__prog, TWA_SIGNAL);
 }
 
 /**
@@ -3738,13 +3935,44 @@ __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, tw, map__map,
+				      callback, aux__prog, mode);
 }
 
 __bpf_kfunc_end_defs();
 
 void bpf_task_work_cancel_and_free(void *val)
 {
+	struct bpf_task_work *tw = val;
+	struct bpf_task_work_context *ctx;
+	enum bpf_task_work_state state;
+
+	ctx = (void *)atomic_long_xchg((atomic_long_t *)&tw->ctx, 0);
+	if (!ctx)
+		return;
+
+	state = xchg(&ctx->state, BPF_TW_FREED);
+
+	switch (state) {
+	case BPF_TW_SCHEDULED:
+		/* If we can't cancel task work, rely on task work callback to free the context */
+		if (!task_work_cancel_match(ctx->task, task_work_match, ctx))
+			break;
+		bpf_reset_task_work_context(ctx);
+		fallthrough;
+	case BPF_TW_STANDBY:
+		kfree(ctx);
+		break;
+	/* In all below cases scheduling logic should detect context state change and cleanup */
+	case BPF_TW_SCHEDULING:
+	case BPF_TW_PENDING:
+	case BPF_TW_RUNNING:
+	default:
+		break;
+	}
 }
 
 BTF_KFUNCS_START(generic_btf_ids)
@@ -3770,6 +3998,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)