[PATCH bpf-next v7 0/4] bpf: Fix task work scheduling races

[PATCH bpf-next v7 0/4] bpf: Fix task work scheduling races

[Yun Lu]

From: Yun Lu <luyun@kylinos.cn>

This series fixes two races in the bpf_task_work scheduling kfuncs and
adds a regression test for each.  Each fix is followed by its test.

Patch 1 captures the task pointer before publishing the callback, which
can clear ctx->task before the scheduling handler attempts cancellation.
The existing RCU read-side section keeps the captured task alive because
bpf_task_release() is put_task_struct_rcu_user().

Patch 2 adds a 1500-round cross-CPU stress test for the cancellation
races: context reuse and map-value deletion racing with scheduling or
callback completion.  It does not force the exact faulting interleavings.

Patch 3 rejects context reuse while the previous scheduling irq_work is
still BUSY.  Otherwise a delayed handler can consume the next round's
SCHEDULING state and leave a false SCHEDULED state with a NULL task.

Patch 4 adds a separate test for that BUSY gate.  A bounded fexit wait on
task_work_add() holds the real scheduling handler after publication.
The target's next syscall attempts context reuse after its callback
wrapper has returned.  Observing its acknowledgment inside the fexit
wait proves the attempt occurred while the old irq_work was still BUSY.
Check rejection with -EBUSY, bounded recovery and exactly-once delivery.
Missed windows are retried within a bounded attempt budget.

Changes since v6 (patch 4 only):

- Skip on PREEMPT_RT through the CONFIG_PREEMPT_RT kconfig extern
  read from the skeleton before attaching, instead of the
  /sys/kernel/realtime probe and the irq_work thread-name detection.
- Share the status accessors, sequence polling, polling bounds and
  the userspace value layout with the cancellation race test.
- Drop the auxiliary stats map, the loop-count diagnostic and the
  wall-clock assertion; the elapsed time is read once after the
  worker cleanup, purely as a diagnostic.  Publish every first-round
  scheduling result, including -EBUSY, and drop the redundant
  userspace CB1 wait and the pacing-only inter-attempt sleep.

The core window protocol and its assertions are unchanged.

v6: https://lore.kernel.org/all/20261010042132.137430-1-luyun_611@163.com/

Changes since v5:

- Narrow the race test to the cancellation races it covers and add the
  task_work_busy test: a bounded fexit wait on task_work_add() holds
  the real window while the target thread attempts reuse after its
  callback wrapper returned, asserting -EBUSY, bounded recovery and
  exactly-once delivery.
- Give each attempt dedicated workers and a fresh map value; reset the
  shared state before the workers are enabled and reuse it only after
  their TIDs have left PID lookup.
- Skip on fewer than two CPUs or confirmed PREEMPT_RT.  Reorder the
  series so each fix is followed by its test.  Kernel patches
  unchanged.

v5: https://lore.kernel.org/all/20260930085417.186729-1-luyun_611@163.com/

Changes since v4 (selftest only; patches 1 and 2 are unchanged):

- Rework the commit message to map each variant to the races fixed by
  this series, and note that the stress test does not force the exact
  faulting interleavings.
- Merge the trigger and target threads into one worker function, name
  the polling bounds, and accept only -ENOENT in the deletion-race
  variant (-EBUSY is retried inside the BPF program and never
  published to userspace).

v4: https://lore.kernel.org/all/20260929072800.113129-1-luyun_611@163.com/

Changes since v3:

- Drop the get_task_struct()/put_task_struct() pair: capturing the
  pointer before task_work_add() is sufficient, since
  bpf_task_release() is put_task_struct_rcu_user() and the handler's
  RCU read-side section keeps the captured task alive.  Reword the
  RCU explanation.
- Split the BUSY check and the one-time irq_work initialization into
  a separate patch 2, with the full interleaving and the failure it
  prevents.
- Selftest code unchanged; commit message wording only.

v3: https://lore.kernel.org/all/20260922100042.136818-1-luyun_611@163.com/

Changes since v2:

- Kernel: drop irq_work_sync() from the callback and both
  WARN_ON_ONCE() checks, as requested in review.  Following the
  suggested direction, the scheduling handler now takes its own task
  reference around task_work_add() and the post-add cancellation; the
  temporary ctx refcount from v2 is gone as well.  The callback ends
  up unchanged from upstream.
- The BUSY check in acquire is kept: without it a handler delayed
  between task_work_add() and its state cmpxchg can flip the next
  round into a bogus SCHEDULED, which later routes map-value deletion
  into task_work_cancel() with an already-reset ctx->task.
- Selftest is unchanged apart from the rebase.

v2: https://lore.kernel.org/all/20260921101453.69273-1-luyun_611@163.com/

Changes since v1:

- Kernel: rework the fix per review feedback.  Replace the round_refs
  ownership counter and the extra irq_work objects with a completion
  edge: the callback claims RUNNING and waits for the scheduling
  irq_work via irq_work_sync(), plus one temporary reference from the
  existing ctx refcount and a BUSY check covering the add-failure
  path.  No context fields or states are added.
- Selftest: address review feedback (single packed result store,
  bounded thread-start wait, tolerate the expected deletion-race
  outcomes).

v1: https://lore.kernel.org/bpf/20260918080014.54012-1-luyun_611@163.com/

Testing:

- QEMU/KVM on x86-64 with KASAN, lockdep, PROVE_RCU and
  DEBUG_ATOMIC_SLEEP: task_work_busy passed three ordinary runs and
  three under target-CPU contention, with missed windows retried
  within the attempt budget; the cancellation race test passed as
  well.
- Without the BUSY fix the same binary fails the in-window rejection
  assertion as expected.  A PREEMPT_RT kernel skips before attaching,
  and a tracefs instance without the global function listing still
  attaches and passes.

Yun Lu (4):
  bpf: Fix NULL task dereference in task work cancellation
  selftests/bpf: Add task work cancellation race test
  bpf: Prevent task work context reuse while irq_work is busy
  selftests/bpf: Add task work BUSY window test

 kernel/bpf/helpers.c                          |  54 +-
 .../selftests/bpf/prog_tests/test_task_work.c | 694 ++++++++++++++++++
 .../selftests/bpf/progs/task_work_busy.c      | 325 ++++++++
 .../selftests/bpf/progs/task_work_race.c      | 126 ++++
 4 files changed, 1189 insertions(+), 10 deletions(-)
 create mode 100644 tools/testing/selftests/bpf/progs/task_work_busy.c
 create mode 100644 tools/testing/selftests/bpf/progs/task_work_race.c


base-commit: 5dd1818b15d98d4a20806cd00b1b40320b06004f
-- 
2.43.0

[PATCH bpf-next v7 1/4] bpf: Fix NULL task dereference in task work cancellation

[Yun Lu]

From: Yun Lu <luyun@kylinos.cn>

bpf_task_work_irq() publishes ctx->work with task_work_add() before
changing the context from SCHEDULING to SCHEDULED.  The target task can
run the callback on another CPU before the handler reaches its final
state transition.

The handler then reloads ctx->task for cancellation if its final cmpxchg
observes FREED.  By that point the callback or context destruction may
already have cleared ctx->task, causing task_work_cancel() to dereference
a NULL task pointer.

One concrete interleaving is:

  1. CPU 0 changes PENDING to SCHEDULING and successfully queues
     ctx->work with task_work_add(), but has not attempted the final
     state cmpxchg.

  2. The target task on CPU 1 runs the callback.  It changes SCHEDULING
     to RUNNING, executes the BPF subprogram, resets ctx->task,
     publishes STANDBY and drops the callback's context reference.

  3. CPU 2 deletes the map value.  It changes STANDBY to FREED and drops
     the map's context reference without queuing asynchronous
     cancellation.

  4. CPU 0 resumes.  Its SCHEDULING-to-SCHEDULED cmpxchg observes FREED,
     and the cancellation attempt passes NULL to task_work_cancel().

If deletion publishes FREED before the callback claims RUNNING, the
callback takes its FREED exit and drops its context reference.  Context
destruction can then clear ctx->task before CPU 0 attempts cancellation.

A controlled reproducer with a delay before the final cmpxchg triggered
the following fault on a v7.3-rc3 based x86-64 KVM guest:

  BUG: kernel NULL pointer dereference, address: 0000000000000890
  RIP: 0010:task_work_cancel+0xd/0xa0
  RDI: 0000000000000000
  CR2: 0000000000000890
  Call Trace:
   <IRQ>
   bpf_task_work_irq+0x95/0x100
   irq_work_run_list+0x4f/0x90
   irq_work_run+0x18/0x50
   </IRQ>

0x890 is the task_struct::task_works offset in that build, read through
the NULL task argument.

Capture ctx->task before task_work_add() publishes the callback and use
that pointer for the handler's cancellation attempt.  The round still
holds its task reference when the pointer is captured.  Releasing that
reference uses bpf_task_release(), which calls put_task_struct_rcu_user(),
so the handler's existing RCU read-side section keeps the captured task
alive even if ctx->task is subsequently cleared.

The asynchronous cancellation path reads ctx->task directly, which is
safe there: queuing it requires deletion to observe SCHEDULED, and that
xchg wins against the callback's transition to RUNNING, so the callback
takes its FREED exit without resetting ctx->task, while the map
reference transferred to the cancellation handler prevents context
destruction until cancellation completes.

The handler-tail and asynchronous cancellation attempts are mutually
exclusive: asynchronous cancellation is queued only if deletion observes
SCHEDULED after the scheduling handler's final cmpxchg has succeeded.
If that cmpxchg instead observes FREED, the scheduling handler attempts
cancellation with the captured pointer.

Fixes: 38aa7003e369 ("bpf: task work scheduling kfuncs")
Reported-by: Jackie Liu <liuyun01@kylinos.cn>
Suggested-by: Mykyta Yatsenko <mykyta.yatsenko5@gmail.com>
Signed-off-by: Yun Lu <luyun@kylinos.cn>
---
 kernel/bpf/helpers.c | 40 +++++++++++++++++++++++++++++++---------
 1 file changed, 31 insertions(+), 9 deletions(-)

diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
index b3cc5c8fc875..028fc6796b59 100644
--- a/kernel/bpf/helpers.c
+++ b/kernel/bpf/helpers.c
@@ -4437,15 +4437,21 @@ static void bpf_task_work_ctx_put(struct bpf_task_work_ctx *ctx)
 	}
 }
 
-static void bpf_task_work_cancel(struct bpf_task_work_ctx *ctx)
+static void bpf_task_work_cancel(struct bpf_task_work_ctx *ctx,
+				 struct task_struct *task)
 {
 	/*
 	 * 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.
+	 *
+	 * The caller passes the task explicitly: the scheduling handler must
+	 * not read ctx->task after the work is published (the callback may
+	 * have reset it), while in the asynchronous cancellation path
+	 * ctx->task is stable, see bpf_task_work_cancel_scheduled().
 	 */
-	if (task_work_cancel(ctx->task, &ctx->work))
+	if (task_work_cancel(task, &ctx->work))
 		bpf_task_work_ctx_put(ctx);
 }
 
@@ -4486,6 +4492,7 @@ static void bpf_task_work_callback(struct callback_head *cb)
 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);
+	struct task_struct *task;
 	enum bpf_task_work_state state;
 	int err;
 
@@ -4496,7 +4503,16 @@ static void bpf_task_work_irq(struct irq_work *irq_work)
 		return;
 	}
 
-	err = task_work_add(ctx->task, &ctx->work, ctx->mode);
+	/*
+	 * Capture the task pointer before the work is published: once the
+	 * callback runs, it may reset ctx->task. The capture happens while
+	 * this round's task reference is held, and every path that can
+	 * release it runs after this handler entered its rcu read-side
+	 * section. bpf_task_release() is put_task_struct_rcu_user(), so the
+	 * task_struct itself cannot be freed before this section ends.
+	 */
+	task = ctx->task;
+	err = task_work_add(task, &ctx->work, ctx->mode);
 	if (err) {
 		bpf_task_work_ctx_reset(ctx);
 		/*
@@ -4510,14 +4526,14 @@ static void bpf_task_work_irq(struct irq_work *irq_work)
 
 	/*
 	 * 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 an extra ref just
-	 * to protect below ctx->state access, we rely on rcu_read_lock
-	 * above to prevent kfree_rcu from freeing ctx before we return.
+	 * complete running by now, going SCHEDULING -> RUNNING, resetting
+	 * ctx->task and dropping its ctx refcount. The rcu read-side section
+	 * above keeps both the ctx memory and the captured task pointer
+	 * valid until this handler returns.
 	 */
 	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 */
+		bpf_task_work_cancel(ctx, task); /* clean up if we switched into FREED state */
 }
 
 static struct bpf_task_work_ctx *bpf_task_work_fetch_ctx(struct bpf_task_work *tw,
@@ -4777,7 +4793,13 @@ 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 */
+	/*
+	 * Deletion observed SCHEDULED and won against the callback's
+	 * transition to RUNNING. The callback takes its FREED exit without
+	 * resetting ctx->task, and the transferred map reference prevents
+	 * destruction until cancellation completes.
+	 */
+	bpf_task_work_cancel(ctx, ctx->task); /* 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 */
 }
 
-- 
2.43.0

[PATCH bpf-next v7 2/4] selftests/bpf: Add task work cancellation race test

[Yun Lu]

From: Yun Lu <luyun@kylinos.cn>

Add cross-CPU stress coverage for the captured-task cancellation fix.
Run scheduling and the target callback on different CPUs, rotating
through three map-value lifetime patterns for 1500 rounds.

Reuse rounds keep the map value and start the next generation after the
callback records its completion.  The next scheduling attempt can race
with the remaining callback cleanup, checking that the context becomes
reusable and subsequent generations make progress.

Late-deletion rounds remove the value after the callback records its
completion, racing deletion with the remaining callback cleanup and the
scheduling handler's final state transition.  Early-deletion rounds
remove it once the BPF program announces that it is about to schedule,
racing deletion with context acquisition and callback publication.
These deletion variants stress the cancellation paths fixed by the
captured-task change: on an unfixed kernel, a deletion observed by a
delayed scheduling handler can route cancellation through an already
cleared ctx->task.

Publish the scheduling result together with its generation in one
packed value, so a stale result cannot be mistaken for the current
round's.  The BPF program retries transient -EBUSY results without
publishing them; in early-deletion rounds, accept -ENOENT when
deletion removes the value before a retry looks it up.  Check callback
completion separately for each generation so that an earlier callback
cannot satisfy a later round.

This is stress coverage for task-work cancellation races; it does not
force the exact faulting interleavings or validate rejection of context
reuse while the scheduling irq_work is still BUSY.

Signed-off-by: Yun Lu <luyun@kylinos.cn>
---
 .../selftests/bpf/prog_tests/test_task_work.c | 281 ++++++++++++++++++
 .../selftests/bpf/progs/task_work_race.c      | 126 ++++++++
 2 files changed, 407 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/progs/task_work_race.c

diff --git a/tools/testing/selftests/bpf/prog_tests/test_task_work.c b/tools/testing/selftests/bpf/prog_tests/test_task_work.c
index 774b31a5f6ca..6b45210e0578 100644
--- a/tools/testing/selftests/bpf/prog_tests/test_task_work.c
+++ b/tools/testing/selftests/bpf/prog_tests/test_task_work.c
@@ -5,10 +5,14 @@
 #include <stdio.h>
 #include "task_work.skel.h"
 #include "task_work_fail.skel.h"
+#include "task_work_race.skel.h"
 #include <linux/bpf.h>
 #include <linux/perf_event.h>
 #include <sys/syscall.h>
 #include <time.h>
+#include <pthread.h>
+#include <sched.h>
+#include <unistd.h>
 
 static int perf_event_open(__u32 type, __u64 config, int pid)
 {
@@ -155,3 +159,280 @@ void test_task_work(void)
 
 	RUN_TESTS(task_work_fail);
 }
+
+#define TASK_WORK_RACE_ROUNDS		1500
+#define TASK_WORK_RACE_POLL_US		100
+#define TASK_WORK_RACE_POLL_ITERS	20000
+#define TASK_WORK_RACE_START_POLL_ITERS	50000
+
+/* Must match progs/task_work_race.c. */
+enum task_work_race_status {
+	RACE_ARM_SEQ,
+	RACE_READY_SEQ,
+	/* High 32 bits hold the sequence, low 32 bits the signed error. */
+	RACE_RESULT,
+};
+
+struct task_work_race_value {
+	__u32 seq;
+	char data[60];
+	struct bpf_task_work tw;
+};
+
+struct task_work_race_ctx {
+	int stop;
+	int setup_err;
+	int trigger_tid;
+	int target_tid;
+};
+
+struct task_work_race_thread_arg {
+	struct task_work_race_ctx *ctx;
+	int cpu;
+	int *tid;
+};
+
+static int task_work_race_pin_cpu(int cpu)
+{
+	cpu_set_t set;
+
+	CPU_ZERO(&set);
+	CPU_SET(cpu, &set);
+	return pthread_setaffinity_np(pthread_self(), sizeof(set), &set);
+}
+
+static void *task_work_race_worker(void *arg)
+{
+	struct task_work_race_thread_arg *thread = arg;
+	struct task_work_race_ctx *ctx = thread->ctx;
+	int err;
+
+	err = task_work_race_pin_cpu(thread->cpu);
+	if (err)
+		__atomic_store_n(&ctx->setup_err, err, __ATOMIC_RELEASE);
+	__atomic_store_n(thread->tid, syscall(__NR_gettid), __ATOMIC_RELEASE);
+	while (!__atomic_load_n(&ctx->stop, __ATOMIC_ACQUIRE))
+		getppid();
+	return NULL;
+}
+
+static int task_work_race_status(struct task_work_race *skel, int idx,
+				 __s64 *value)
+{
+	return bpf_map_lookup_elem(bpf_map__fd(skel->maps.status), &idx,
+				   value);
+}
+
+static int task_work_race_set_status(struct task_work_race *skel, int idx,
+				     __s64 value)
+{
+	return bpf_map_update_elem(bpf_map__fd(skel->maps.status), &idx,
+				   &value, BPF_ANY);
+}
+
+static int task_work_race_wait_ready(struct task_work_race *skel, __u32 seq)
+{
+	int i;
+
+	for (i = 0; i < TASK_WORK_RACE_POLL_ITERS; i++) {
+		__s64 value = 0;
+
+		if (!task_work_race_status(skel, RACE_READY_SEQ, &value) &&
+		    value == seq)
+			return 0;
+		if (!task_work_race_status(skel, RACE_RESULT, &value) &&
+		    (__u32)((__u64)value >> 32) == seq)
+			/*
+			 * The BPF program finished between the two reads:
+			 * RESULT makes waiting for READY unnecessary. Any error
+			 * that occurred before READY is checked by the caller.
+			 */
+			return 0;
+		usleep(TASK_WORK_RACE_POLL_US);
+	}
+	return -ETIMEDOUT;
+}
+
+static int task_work_race_prepare_elem(struct task_work_race *skel)
+{
+	struct task_work_race_value value = {};
+	int key = 0;
+
+	if (!bpf_map_lookup_elem(bpf_map__fd(skel->maps.hmap), &key, &value))
+		return 0;
+	if (errno != ENOENT)
+		return -errno;
+	if (bpf_map_update_elem(bpf_map__fd(skel->maps.hmap), &key, &value,
+				BPF_NOEXIST))
+		return -errno;
+	return 0;
+}
+
+static int task_work_race_arm(struct task_work_race *skel, __u32 seq)
+{
+	if (task_work_race_set_status(skel, RACE_ARM_SEQ, seq))
+		return -errno;
+	return task_work_race_wait_ready(skel, seq);
+}
+
+static int task_work_race_check_done(struct task_work_race *skel, __u32 seq)
+{
+	int i;
+
+	for (i = 0; i < TASK_WORK_RACE_POLL_ITERS; i++) {
+		__s64 result = 0;
+
+		if (!task_work_race_status(skel, RACE_RESULT, &result) &&
+		    (__u32)((__u64)result >> 32) == seq)
+			return (__s32)result;
+		usleep(TASK_WORK_RACE_POLL_US);
+	}
+	return -ETIMEDOUT;
+}
+
+static int task_work_race_wait_callback(struct task_work_race *skel,
+					__u32 seq)
+{
+	int i, key = seq;
+
+	for (i = 0; i < TASK_WORK_RACE_POLL_ITERS; i++) {
+		__u64 completed = 0;
+
+		if (!bpf_map_lookup_elem(bpf_map__fd(skel->maps.completed),
+					 &key, &completed) && completed)
+			return 0;
+		usleep(TASK_WORK_RACE_POLL_US);
+	}
+	return -ETIMEDOUT;
+}
+
+void serial_test_task_work_race(void)
+{
+	struct task_work_race_ctx ctx = {};
+	struct task_work_race_thread_arg trigger_arg = {
+		.ctx = &ctx,
+		.tid = &ctx.trigger_tid,
+	};
+	struct task_work_race_thread_arg target_arg = {
+		.ctx = &ctx,
+		.tid = &ctx.target_tid,
+	};
+	struct task_work_race *skel;
+	pthread_t trigger, target;
+	cpu_set_t allowed;
+	int cpu, cpu_count = 0;
+	bool target_started = false, trigger_started = false;
+	int err, i, key = 0;
+
+	if (sched_getaffinity(0, sizeof(allowed), &allowed)) {
+		ASSERT_OK(-errno, "sched_getaffinity");
+		return;
+	}
+	for (cpu = 0; cpu < CPU_SETSIZE && cpu_count < 2; cpu++) {
+		if (!CPU_ISSET(cpu, &allowed))
+			continue;
+		if (!cpu_count)
+			trigger_arg.cpu = cpu;
+		else
+			target_arg.cpu = cpu;
+		cpu_count++;
+	}
+	if (cpu_count < 2) {
+		printf("%s:SKIP:need two CPUs in the process affinity mask\n",
+		       __func__);
+		test__skip();
+		return;
+	}
+
+	skel = task_work_race__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "open_and_load"))
+		return;
+	if (!ASSERT_OK(task_work_race__attach(skel), "attach"))
+		goto cleanup;
+
+	err = pthread_create(&target, NULL, task_work_race_worker, &target_arg);
+	if (!ASSERT_OK(err, "pthread_create target"))
+		goto cleanup;
+	target_started = true;
+	err = pthread_create(&trigger, NULL, task_work_race_worker, &trigger_arg);
+	if (!ASSERT_OK(err, "pthread_create trigger"))
+		goto stop;
+	trigger_started = true;
+
+	for (i = 0; i < TASK_WORK_RACE_START_POLL_ITERS; i++) {
+		if (__atomic_load_n(&ctx.trigger_tid, __ATOMIC_ACQUIRE) &&
+		    __atomic_load_n(&ctx.target_tid, __ATOMIC_ACQUIRE))
+			break;
+		usleep(TASK_WORK_RACE_POLL_US);
+	}
+	if (!ASSERT_TRUE(__atomic_load_n(&ctx.trigger_tid, __ATOMIC_ACQUIRE) &&
+			 __atomic_load_n(&ctx.target_tid, __ATOMIC_ACQUIRE),
+			 "thread start"))
+		goto stop;
+	if (!ASSERT_OK(__atomic_load_n(&ctx.setup_err, __ATOMIC_ACQUIRE),
+		       "thread affinity"))
+		goto stop;
+
+	skel->bss->trigger_tid = __atomic_load_n(&ctx.trigger_tid, __ATOMIC_ACQUIRE);
+	skel->bss->target_tid = __atomic_load_n(&ctx.target_tid, __ATOMIC_ACQUIRE);
+
+	for (i = 0; i < TASK_WORK_RACE_ROUNDS; i++) {
+		__u32 seq = i + 1;
+		int variant = i % 3;
+
+		if (!ASSERT_OK(task_work_race_prepare_elem(skel), "prepare elem") ||
+		    !ASSERT_OK(task_work_race_arm(skel, seq), "round ready"))
+			goto stop;
+
+		if (variant == 2 &&
+		    !ASSERT_OK(bpf_map_delete_elem(bpf_map__fd(skel->maps.hmap),
+						   &key), "early delete"))
+			goto stop;
+
+		err = task_work_race_check_done(skel, seq);
+		/*
+		 * Variant 2 deletes the element after READY. A scheduling attempt
+		 * racing with deletion can return -EBUSY, which the BPF program
+		 * retries without publishing RESULT. The retry can then report
+		 * -ENOENT because the element is gone; accept that terminal result.
+		 * A later variant 0 round recreates the element.
+		 */
+		if (variant == 2 && err == -ENOENT)
+			err = 0;
+		if (!ASSERT_OK(err, "schedule")) {
+			fprintf(stderr, "round %d schedule failed: %d\n", i, err);
+			goto stop;
+		}
+
+		if (variant != 2 &&
+		    !ASSERT_OK(task_work_race_wait_callback(skel, seq),
+			       "callback"))
+			goto stop;
+
+		if (variant == 1 &&
+		    !ASSERT_OK(bpf_map_delete_elem(bpf_map__fd(skel->maps.hmap),
+						   &key), "late delete"))
+			goto stop;
+	}
+
+	/* A distinct completion generation prevents an old callback satisfying this. */
+	if (!ASSERT_OK(task_work_race_prepare_elem(skel), "final prepare") ||
+	    !ASSERT_OK(task_work_race_arm(skel, TASK_WORK_RACE_ROUNDS + 1),
+		       "final ready") ||
+	    !ASSERT_OK(task_work_race_check_done(skel,
+					TASK_WORK_RACE_ROUNDS + 1),
+		       "final schedule") ||
+	    !ASSERT_OK(task_work_race_wait_callback(skel,
+					     TASK_WORK_RACE_ROUNDS + 1),
+		       "final callback"))
+		goto stop;
+
+stop:
+	__atomic_store_n(&ctx.stop, 1, __ATOMIC_RELEASE);
+	if (trigger_started)
+		pthread_join(trigger, NULL);
+	if (target_started)
+		pthread_join(target, NULL);
+cleanup:
+	task_work_race__destroy(skel);
+}
diff --git a/tools/testing/selftests/bpf/progs/task_work_race.c b/tools/testing/selftests/bpf/progs/task_work_race.c
new file mode 100644
index 000000000000..c94eca510cb3
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/task_work_race.c
@@ -0,0 +1,126 @@
+// SPDX-License-Identifier: GPL-2.0
+/* Copyright (c) 2026 KylinSoft Corporation. */
+#include <vmlinux.h>
+#include <stdbool.h>
+#include <errno.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+#include "bpf_misc.h"
+
+char _license[] SEC("license") = "GPL";
+
+#define TASK_WORK_RACE_MAX_SEQ	2048
+
+enum task_work_race_status {
+	RACE_ARM_SEQ,
+	RACE_READY_SEQ,
+	/* High 32 bits hold the sequence, low 32 bits the signed error. */
+	RACE_RESULT,
+	RACE_STATUS_MAX,
+};
+
+int trigger_tid;
+int target_tid;
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__uint(max_entries, RACE_STATUS_MAX);
+	__type(key, int);
+	__type(value, __s64);
+} status SEC(".maps");
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__uint(max_entries, TASK_WORK_RACE_MAX_SEQ);
+	__type(key, int);
+	__type(value, __u64);
+} completed SEC(".maps");
+
+struct task_work_race_value {
+	__u32 seq;
+	char data[60];
+	struct bpf_task_work tw;
+};
+
+struct {
+	__uint(type, BPF_MAP_TYPE_HASH);
+	__uint(map_flags, BPF_F_NO_PREALLOC);
+	__uint(max_entries, 4);
+	__type(key, int);
+	__type(value, struct task_work_race_value);
+} hmap SEC(".maps");
+
+static __always_inline void set_status(int key, __s64 value)
+{
+	__s64 *slot;
+
+	slot = bpf_map_lookup_elem(&status, &key);
+	if (slot)
+		*slot = value;
+}
+
+static int process_work(struct bpf_map *map, void *key, void *value)
+{
+	struct task_work_race_value *work = value;
+	__u64 *done;
+	int seq = work->seq;
+
+	if (seq <= 0 || seq >= TASK_WORK_RACE_MAX_SEQ)
+		return 0;
+	done = bpf_map_lookup_elem(&completed, &seq);
+	if (done)
+		*done = 1;
+	return 0;
+}
+
+SEC("tracepoint/syscalls/sys_enter_getppid")
+int race_sched_work(void *ctx)
+{
+	struct task_work_race_value *work;
+	struct task_struct *task;
+	__s64 *arm, *result;
+	__u32 tid = (__u32)bpf_get_current_pid_tgid();
+	int key = 0, err;
+	__u32 seq;
+
+	if (tid != trigger_tid)
+		return 0;
+
+	key = RACE_ARM_SEQ;
+	arm = bpf_map_lookup_elem(&status, &key);
+	key = RACE_RESULT;
+	result = bpf_map_lookup_elem(&status, &key);
+	if (!arm || !result || *arm <= 0 ||
+	    (__u32)((__u64)*result >> 32) == *arm)
+		return 0;
+	seq = *arm;
+
+	task = bpf_task_from_pid(target_tid);
+	if (!task) {
+		err = -ESRCH;
+		goto out_done;
+	}
+
+	key = 0;
+	work = bpf_map_lookup_elem(&hmap, &key);
+	if (!work) {
+		err = -ENOENT;
+		goto out_task;
+	}
+
+	work->seq = seq;
+	set_status(RACE_READY_SEQ, seq);
+	err = bpf_task_work_schedule_signal(task, &work->tw, &hmap,
+					    process_work);
+	if (err == -EBUSY) {
+		bpf_task_release(task);
+		return 0;
+	}
+
+out_task:
+	bpf_task_release(task);
+out_done:
+	/* Publish the sequence and its error in one store. */
+	*result = ((__u64)seq << 32) | (__u32)err;
+	return 0;
+}
-- 
2.43.0

[PATCH bpf-next v7 3/4] bpf: Prevent task work context reuse while irq_work is busy

[Yun Lu]

From: Yun Lu <luyun@kylinos.cn>

A bpf_task_work context can return to STANDBY before its scheduling
irq_work handler returns.  Both callback completion and task_work_add()
failure can publish STANDBY while the irq_work is still BUSY.  A new
round can then reinitialize and queue the same irq_work on another CPU
while the old handler is still running.

One concrete interleaving is:

  1. CPU 0 successfully queues ctx->work with task_work_add() and pauses
     before its SCHEDULING-to-SCHEDULED cmpxchg.

  2. The target task on CPU 1 completes the callback, resets ctx->task
     and publishes STANDBY.

  3. CPU 2 starts a new round, changes STANDBY to PENDING and queues the
     same irq_work.  The new handler changes PENDING to SCHEDULING, but
     task_work_add() returns -ESRCH because its target task has exited.
     The handler has not yet reset the context or rolled back the state.

  4. CPU 0 resumes and changes the new round's SCHEDULING to SCHEDULED,
     even though no task work was queued for that round.

  5. The new handler resets ctx->task, but its SCHEDULING-to-STANDBY
     cmpxchg fails.  The context remains SCHEDULED with a NULL task, and
     subsequent attempts to acquire this context fail with -EBUSY.

  6. Map-value deletion observes SCHEDULED and queues asynchronous
     cancellation, which passes the NULL ctx->task to task_work_cancel().

Overlapping invocations also share one BUSY bit.  An invocation can clear
that bit while the other invocation is still running.

Initialize the scheduling irq_work when the context is created and stop
reinitializing it for each round.  After claiming the context with the
STANDBY-to-PENDING cmpxchg, check BUSY and reject reuse with -EBUSY if a
previous invocation is still in progress.  Checking after the cmpxchg
orders the BUSY read after the STANDBY publication.

No irq_work has been queued for the attempted round at this point, so
the rejection can roll PENDING back to STANDBY.  Use cmpxchg for that
rollback to preserve FREED if a concurrent deletion has published it.
Once BUSY clears, scheduling can be retried.

Fixes: 38aa7003e369 ("bpf: task work scheduling kfuncs")
Signed-off-by: Yun Lu <luyun@kylinos.cn>
Acked-by: Mykyta Yatsenko <yatsenko@meta.com>
---
 kernel/bpf/helpers.c | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
index 028fc6796b59..350269b8646e 100644
--- a/kernel/bpf/helpers.c
+++ b/kernel/bpf/helpers.c
@@ -4553,6 +4553,7 @@ static struct bpf_task_work_ctx *bpf_task_work_fetch_ctx(struct bpf_task_work *t
 	memset(ctx, 0, sizeof(*ctx));
 	refcount_set(&ctx->refcnt, 1); /* map's own ref */
 	ctx->state = BPF_TW_STANDBY;
+	init_irq_work(&ctx->irq_work, bpf_task_work_irq);
 
 	old_ctx = cmpxchg(&twk->ctx, NULL, ctx);
 	if (old_ctx) {
@@ -4595,6 +4596,18 @@ static struct bpf_task_work_ctx *bpf_task_work_acquire_ctx(struct bpf_task_work
 		bpf_task_work_ctx_put(ctx);
 		return ERR_PTR(-EBUSY);
 	}
+	/*
+	 * STANDBY can be published before the previous scheduling irq_work
+	 * returns. Claim PENDING before checking BUSY to order the check after
+	 * that publication. No irq_work has been queued for this attempt yet,
+	 * so reject reuse and roll back to STANDBY while BUSY is set. Preserve
+	 * FREED if deletion wins the rollback race.
+	 */
+	if (unlikely(irq_work_is_busy(&ctx->irq_work))) {
+		(void)cmpxchg(&ctx->state, BPF_TW_PENDING, BPF_TW_STANDBY);
+		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
@@ -4644,7 +4657,6 @@ static int bpf_task_work_schedule(struct task_struct *task, struct bpf_task_work
 	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;
-- 
2.43.0

Re: [PATCH bpf-next v7 3/4] bpf: Prevent task work context reuse while irq_work is busy

[sashiko-bot]

Thank you for your contribution! Sashiko AI review found 1 potential issue(s) to consider:
- [High] bpf: UAF in task_work/timer callbacks due to async context re-init on deleted map elements
--

commit 580f53f590a0a7077f231292b13de90213f646c5
Author: Yun Lu <luyun@kylinos.cn>

bpf: Prevent task work context reuse while irq_work is busy

This commit fixes a race condition where a task work context could return to
standby before its scheduling irq_work handler completes. It initializes the
irq_work during context creation and checks for the BUSY state before reusing
it to prevent overlapping invocations.

> diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
> index 028fc6796b597..350269b8646e4 100644
> --- a/kernel/bpf/helpers.c
> +++ b/kernel/bpf/helpers.c
> @@ -4553,6 +4553,7 @@ static struct bpf_task_work_ctx *bpf_task_work_fetch_ctx(struct bpf_task_work *t
>  	memset(ctx, 0, sizeof(*ctx));
>  	refcount_set(&ctx->refcnt, 1); /* map's own ref */
>  	ctx->state = BPF_TW_STANDBY;
> +	init_irq_work(&ctx->irq_work, bpf_task_work_irq);
>  
>  	old_ctx = cmpxchg(&twk->ctx, NULL, ctx);

[Severity: High]
Can a newly allocated context be installed on a deleted map element here?

If a concurrent map deletion calls bpf_task_work_cancel_and_free(), it
synchronously clears the context:

bpf_task_work_cancel_and_free(struct bpf_task_work *tw) {
    ...
    ctx = xchg(&twk->ctx, NULL);
    ...
}

Since it doesn't leave a poison value, a concurrent BPF program running under
RCU read lock that calls bpf_task_work_schedule() would see twk->ctx as NULL
and successfully cmpxchg the newly allocated context into the deleted element.

[ ... ]

> @@ -4644,7 +4657,6 @@ static int bpf_task_work_schedule(struct task_struct *task, struct bpf_task_work
>  	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;

[Severity: High]
Could this lead to a use-after-free if the element was concurrently deleted?

If bpf_task_work_fetch_ctx() re-initializes a context on a deleted element as
described above, ctx->map_val captures a pointer to the element's memory:

bpf_task_work_schedule(...) {
    ...
    ctx->map_val = (void *)tw - map->record->task_work_off;
    ...
    irq_work_queue(&ctx->irq_work);
}

Since the deleted map element will be freed after the RCU grace period, but
the newly initialized asynchronous irq_work remains armed, could the later
execution of this callback access the freed element memory through
ctx->map_val?

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20261012061915.231254-1-luyun_611@163.com?part=3

Re: [PATCH bpf-next v7 3/4] bpf: Prevent task work context reuse while irq_work is busy

[luyun]

在 2026/10/12 14:34, sashiko-bot@kernel.org 写道:
> Thank you for your contribution! Sashiko AI review found 1 potential issue(s) to consider:
> - [High] bpf: UAF in task_work/timer callbacks due to async context re-init on deleted map elements
> --
Hi,
I reproduced this issue with KASAN on qemu/kvm.

After the map value is deleted, a stale value pointer can still re-publish
a context and schedule a task work through it (the schedule succeeds, no
error). Callbacks that have not entered RUNNING yet are cancelled by the
final allocator destructor and are safe. The gap is a callback that is
already RUNNING when the freeing grace period completes: the destructor
only publishes FREED without waiting, and the callback keeps using
ctx->map_val after the element has been kfree:

   BUG: KASAN: slab-use-after-free in bpf_repro_probe_tw
   Read of size 1 at addr ffff88800d999558 by task uaf_repro/89

   Call Trace:
     bpf_repro_probe_tw          <- debug-only probe, scratch build
     process_work_cb             <- BPF task work callback
     bpf_task_work_callback
     task_work_run

   Freed by task 11:
     kfree
     free_all
     __free_rcu
     srcu_invoke_callbacks       <- bpf_mem_cache freeing grace period

This is not introduced by this series: fetch_ctx(), schedule(),
cancel_and_free() and the callback are identical to the pre-series base,
and patch 3 only moves init_irq_work() around. I hit the same reinstall
path on the submitted tree with no modifications.

The fix has to stop retired elements from re-publishing async work and
keep elements alive while a callback still runs on them, so I'd rather fix

it in a separate series than grow this patch set.

---

Thanks.
Yun Lu

> commit 580f53f590a0a7077f231292b13de90213f646c5
> Author: Yun Lu <luyun@kylinos.cn>
>
> bpf: Prevent task work context reuse while irq_work is busy
>
> This commit fixes a race condition where a task work context could return to
> standby before its scheduling irq_work handler completes. It initializes the
> irq_work during context creation and checks for the BUSY state before reusing
> it to prevent overlapping invocations.
>
>> diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
>> index 028fc6796b597..350269b8646e4 100644
>> --- a/kernel/bpf/helpers.c
>> +++ b/kernel/bpf/helpers.c
>> @@ -4553,6 +4553,7 @@ static struct bpf_task_work_ctx *bpf_task_work_fetch_ctx(struct bpf_task_work *t
>>   	memset(ctx, 0, sizeof(*ctx));
>>   	refcount_set(&ctx->refcnt, 1); /* map's own ref */
>>   	ctx->state = BPF_TW_STANDBY;
>> +	init_irq_work(&ctx->irq_work, bpf_task_work_irq);
>>   
>>   	old_ctx = cmpxchg(&twk->ctx, NULL, ctx);
> [Severity: High]
> Can a newly allocated context be installed on a deleted map element here?
>
> If a concurrent map deletion calls bpf_task_work_cancel_and_free(), it
> synchronously clears the context:
>
> bpf_task_work_cancel_and_free(struct bpf_task_work *tw) {
>      ...
>      ctx = xchg(&twk->ctx, NULL);
>      ...
> }
>
> Since it doesn't leave a poison value, a concurrent BPF program running under
> RCU read lock that calls bpf_task_work_schedule() would see twk->ctx as NULL
> and successfully cmpxchg the newly allocated context into the deleted element.
>
> [ ... ]
>
>> @@ -4644,7 +4657,6 @@ static int bpf_task_work_schedule(struct task_struct *task, struct bpf_task_work
>>   	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;
> [Severity: High]
> Could this lead to a use-after-free if the element was concurrently deleted?
>
> If bpf_task_work_fetch_ctx() re-initializes a context on a deleted element as
> described above, ctx->map_val captures a pointer to the element's memory:
>
> bpf_task_work_schedule(...) {
>      ...
>      ctx->map_val = (void *)tw - map->record->task_work_off;
>      ...
>      irq_work_queue(&ctx->irq_work);
> }
>
> Since the deleted map element will be freed after the RCU grace period, but
> the newly initialized asynchronous irq_work remains armed, could the later
> execution of this callback access the freed element memory through
> ctx->map_val?
>

[PATCH bpf-next v7 4/4] selftests/bpf: Add task work BUSY window test

[Yun Lu]

From: Yun Lu <luyun@kylinos.cn>

Add a regression test for reusing a bpf_task_work context while its
scheduling irq_work is still BUSY.  A test that schedules from a thread
pinned to one CPU never observes that window: irq_work_queue() self-IPIs
the local CPU and the invocation completes before the thread's next
syscall.

Use an fexit program on task_work_add() to delay the scheduling
handler's return with a time- and iteration-bounded busy-wait, after
the first-round callback was published and before the handler's tail
state transition.  The target task attempts to reuse the same context
from a later syscall, strictly after its own callback wrapper returned,
so the context is STANDBY at that point.  A valid window requires the
fexit program to observe the attempt's acknowledgment before returning,
which proves the attempt ran while the old handler, and thus the BUSY
irq_work, was still in flight.  With the map and program references
held and no competing scheduler, the only possible source of -EBUSY for
that attempt is the BUSY check in the acquire path.  Check that the
attempt is rejected with -EBUSY, that scheduling succeeds on a bounded
retry once BUSY clears, and that the second callback runs exactly once.

Each attempt uses dedicated workers and a fresh map value.  Reset the
shared slots before enabling the workers.  Stop and join them at the
end of the attempt, then wait for their TIDs to disappear from PID
lookup before checking the final callback count or starting another
attempt.  join alone can return before exit_task_work(); a signal-zero
tgkill probe still finds the worker until after its exit task_work has
run.  Bound that wait and fail on unexpected probe errors.  Callback
deliveries are also matched against the window's generation.

Missed windows are counted and retried within a bounded attempt budget;
hard failures fail the test.  Record the busy-wait elapsed time after
worker cleanup as a diagnostic, without asserting a wall-clock bound,
since host vCPU preemption can stretch wall-clock time without affecting
the gate behavior under test.

Skip when fewer than two CPUs are available or when the loaded kconfig
extern identifies CONFIG_PREEMPT_RT.  Check that extern before attaching
programs or starting workers.  Load and attach errors fail the test
instead of being inferred from tracefs file paths and reported as
unsupported.

Signed-off-by: Yun Lu <luyun@kylinos.cn>
---
 .../selftests/bpf/prog_tests/test_task_work.c | 477 ++++++++++++++++--
 .../selftests/bpf/progs/task_work_busy.c      | 325 ++++++++++++
 2 files changed, 770 insertions(+), 32 deletions(-)
 create mode 100644 tools/testing/selftests/bpf/progs/task_work_busy.c

diff --git a/tools/testing/selftests/bpf/prog_tests/test_task_work.c b/tools/testing/selftests/bpf/prog_tests/test_task_work.c
index 6b45210e0578..deabc524dd49 100644
--- a/tools/testing/selftests/bpf/prog_tests/test_task_work.c
+++ b/tools/testing/selftests/bpf/prog_tests/test_task_work.c
@@ -6,6 +6,7 @@
 #include "task_work.skel.h"
 #include "task_work_fail.skel.h"
 #include "task_work_race.skel.h"
+#include "task_work_busy.skel.h"
 #include <linux/bpf.h>
 #include <linux/perf_event.h>
 #include <sys/syscall.h>
@@ -161,9 +162,9 @@ void test_task_work(void)
 }
 
 #define TASK_WORK_RACE_ROUNDS		1500
-#define TASK_WORK_RACE_POLL_US		100
-#define TASK_WORK_RACE_POLL_ITERS	20000
-#define TASK_WORK_RACE_START_POLL_ITERS	50000
+#define TASK_WORK_POLL_US		100
+#define TASK_WORK_POLL_ITERS	20000
+#define TASK_WORK_START_POLL_ITERS	50000
 
 /* Must match progs/task_work_race.c. */
 enum task_work_race_status {
@@ -173,7 +174,7 @@ enum task_work_race_status {
 	RACE_RESULT,
 };
 
-struct task_work_race_value {
+struct task_work_value {
 	__u32 seq;
 	char data[60];
 	struct bpf_task_work tw;
@@ -216,31 +217,46 @@ static void *task_work_race_worker(void *arg)
 	return NULL;
 }
 
-static int task_work_race_status(struct task_work_race *skel, int idx,
-				 __s64 *value)
+static int task_work_status(struct bpf_map *map, int idx, __s64 *value)
 {
-	return bpf_map_lookup_elem(bpf_map__fd(skel->maps.status), &idx,
-				   value);
+	return bpf_map_lookup_elem(bpf_map__fd(map), &idx, value);
 }
 
-static int task_work_race_set_status(struct task_work_race *skel, int idx,
-				     __s64 value)
+static int task_work_set_status(struct bpf_map *map, int idx, __s64 value)
 {
-	return bpf_map_update_elem(bpf_map__fd(skel->maps.status), &idx,
-				   &value, BPF_ANY);
+	return bpf_map_update_elem(bpf_map__fd(map), &idx, &value, BPF_ANY);
+}
+
+/* Wait for seq, keeping polling errors separate from the scheduling result. */
+static int task_work_wait_seq(struct bpf_map *map, int idx, __u32 seq,
+			      __s32 *result)
+{
+	int i;
+
+	for (i = 0; i < TASK_WORK_POLL_ITERS; i++) {
+		__s64 value = 0;
+
+		if (!task_work_status(map, idx, &value) &&
+		    (__u32)((__u64)value >> 32) == seq) {
+			*result = (__s32)value;
+			return 0;
+		}
+		usleep(TASK_WORK_POLL_US);
+	}
+	return -ETIMEDOUT;
 }
 
 static int task_work_race_wait_ready(struct task_work_race *skel, __u32 seq)
 {
 	int i;
 
-	for (i = 0; i < TASK_WORK_RACE_POLL_ITERS; i++) {
+	for (i = 0; i < TASK_WORK_POLL_ITERS; i++) {
 		__s64 value = 0;
 
-		if (!task_work_race_status(skel, RACE_READY_SEQ, &value) &&
+		if (!task_work_status(skel->maps.status, RACE_READY_SEQ, &value) &&
 		    value == seq)
 			return 0;
-		if (!task_work_race_status(skel, RACE_RESULT, &value) &&
+		if (!task_work_status(skel->maps.status, RACE_RESULT, &value) &&
 		    (__u32)((__u64)value >> 32) == seq)
 			/*
 			 * The BPF program finished between the two reads:
@@ -248,14 +264,14 @@ static int task_work_race_wait_ready(struct task_work_race *skel, __u32 seq)
 			 * that occurred before READY is checked by the caller.
 			 */
 			return 0;
-		usleep(TASK_WORK_RACE_POLL_US);
+		usleep(TASK_WORK_POLL_US);
 	}
 	return -ETIMEDOUT;
 }
 
 static int task_work_race_prepare_elem(struct task_work_race *skel)
 {
-	struct task_work_race_value value = {};
+	struct task_work_value value = {};
 	int key = 0;
 
 	if (!bpf_map_lookup_elem(bpf_map__fd(skel->maps.hmap), &key, &value))
@@ -270,24 +286,18 @@ static int task_work_race_prepare_elem(struct task_work_race *skel)
 
 static int task_work_race_arm(struct task_work_race *skel, __u32 seq)
 {
-	if (task_work_race_set_status(skel, RACE_ARM_SEQ, seq))
+	if (task_work_set_status(skel->maps.status, RACE_ARM_SEQ, seq))
 		return -errno;
 	return task_work_race_wait_ready(skel, seq);
 }
 
 static int task_work_race_check_done(struct task_work_race *skel, __u32 seq)
 {
-	int i;
-
-	for (i = 0; i < TASK_WORK_RACE_POLL_ITERS; i++) {
-		__s64 result = 0;
+	__s32 result;
+	int err;
 
-		if (!task_work_race_status(skel, RACE_RESULT, &result) &&
-		    (__u32)((__u64)result >> 32) == seq)
-			return (__s32)result;
-		usleep(TASK_WORK_RACE_POLL_US);
-	}
-	return -ETIMEDOUT;
+	err = task_work_wait_seq(skel->maps.status, RACE_RESULT, seq, &result);
+	return err ?: result;
 }
 
 static int task_work_race_wait_callback(struct task_work_race *skel,
@@ -295,13 +305,13 @@ static int task_work_race_wait_callback(struct task_work_race *skel,
 {
 	int i, key = seq;
 
-	for (i = 0; i < TASK_WORK_RACE_POLL_ITERS; i++) {
+	for (i = 0; i < TASK_WORK_POLL_ITERS; i++) {
 		__u64 completed = 0;
 
 		if (!bpf_map_lookup_elem(bpf_map__fd(skel->maps.completed),
 					 &key, &completed) && completed)
 			return 0;
-		usleep(TASK_WORK_RACE_POLL_US);
+		usleep(TASK_WORK_POLL_US);
 	}
 	return -ETIMEDOUT;
 }
@@ -359,11 +369,11 @@ void serial_test_task_work_race(void)
 		goto stop;
 	trigger_started = true;
 
-	for (i = 0; i < TASK_WORK_RACE_START_POLL_ITERS; i++) {
+	for (i = 0; i < TASK_WORK_START_POLL_ITERS; i++) {
 		if (__atomic_load_n(&ctx.trigger_tid, __ATOMIC_ACQUIRE) &&
 		    __atomic_load_n(&ctx.target_tid, __ATOMIC_ACQUIRE))
 			break;
-		usleep(TASK_WORK_RACE_POLL_US);
+		usleep(TASK_WORK_POLL_US);
 	}
 	if (!ASSERT_TRUE(__atomic_load_n(&ctx.trigger_tid, __ATOMIC_ACQUIRE) &&
 			 __atomic_load_n(&ctx.target_tid, __ATOMIC_ACQUIRE),
@@ -436,3 +446,406 @@ void serial_test_task_work_race(void)
 cleanup:
 	task_work_race__destroy(skel);
 }
+
+#define TASK_WORK_BUSY_ROUNDS		32
+#define TASK_WORK_BUSY_MAX_ATTEMPTS	96
+
+/* Must match progs/task_work_busy.c. */
+enum task_work_busy_status {
+	BUSY_ARM_SEQ,
+	/*
+	 * High 32 bits hold the sequence.  BUSY_WINDOW low bits: 1 open,
+	 * 2 the in-window attempt was seen before the handler returned,
+	 * 3 the wait bound was reached without it.  Result slots: signed
+	 * error in the low 32 bits.
+	 */
+	BUSY_WINDOW,
+	BUSY_ACK,
+	BUSY_GEN1_RESULT,
+	BUSY_GEN2_RESULT,
+	BUSY_CB1_SEQ,
+	BUSY_CB2_COUNT,
+	BUSY_CB2_SEQ,
+	BUSY_RECOVER_SEQ,
+	BUSY_LAST_WAIT_NS,
+	BUSY_STATUS_MAX,
+};
+
+enum task_work_busy_attempt {
+	TWB_VALID,
+	TWB_MISSED,
+	TWB_HARD_FAIL,
+};
+
+static const char *twb_attempt_name(enum task_work_busy_attempt r)
+{
+	switch (r) {
+	case TWB_VALID:
+		return "valid";
+	case TWB_MISSED:
+		return "missed";
+	default:
+		return "fail";
+	}
+}
+
+static int task_work_busy_wait_value(struct task_work_busy *skel, int idx,
+				     __s64 want)
+{
+	int i;
+
+	for (i = 0; i < TASK_WORK_POLL_ITERS; i++) {
+		__s64 value = 0;
+
+		if (!task_work_status(skel->maps.status, idx, &value) && value == want)
+			return 0;
+		usleep(TASK_WORK_POLL_US);
+	}
+	return -ETIMEDOUT;
+}
+
+/* Wait until the window slot reports a final state (2 or 3) for seq. */
+static int task_work_busy_wait_window(struct task_work_busy *skel, __u32 seq,
+				      __s64 *window)
+{
+	int i;
+
+	for (i = 0; i < TASK_WORK_POLL_ITERS; i++) {
+		__s64 value = 0;
+
+		if (!task_work_status(skel->maps.status, BUSY_WINDOW, &value) &&
+		    (__u32)((__u64)value >> 32) == seq &&
+		    (__u32)value >= 2) {
+			*window = value;
+			return 0;
+		}
+		usleep(TASK_WORK_POLL_US);
+	}
+	return -ETIMEDOUT;
+}
+
+static int task_work_busy_prepare_elem(struct task_work_busy *skel, int key)
+{
+	struct task_work_value value = {};
+
+	if (!bpf_map_update_elem(bpf_map__fd(skel->maps.hmap), &key, &value,
+				 BPF_NOEXIST))
+		return 0;
+	return -errno;
+}
+
+static int task_work_busy_reset(struct task_work_busy *skel)
+{
+	static const int zeroed[] = {
+		BUSY_WINDOW, BUSY_ACK, BUSY_GEN1_RESULT, BUSY_GEN2_RESULT,
+		BUSY_CB1_SEQ, BUSY_CB2_COUNT, BUSY_CB2_SEQ, BUSY_RECOVER_SEQ,
+		BUSY_LAST_WAIT_NS,
+	};
+	int i;
+
+	for (i = 0; i < ARRAY_SIZE(zeroed); i++)
+		if (task_work_set_status(skel->maps.status, zeroed[i], 0))
+			return -errno;
+	return 0;
+}
+
+/*
+ * join waits for clear_child_tid, which is cleared before exit_task_work().
+ * PID lookup stops finding the worker only after its exit task_work has run.
+ * Probe with signal 0; no signal is delivered.  Do not create new workers
+ * until both old TIDs have disappeared.
+ */
+static int task_work_busy_wait_exit(int tid)
+{
+	int i;
+
+	for (i = 0; i < TASK_WORK_POLL_ITERS; i++) {
+		if (syscall(__NR_tgkill, getpid(), tid, 0)) {
+			if (errno == ESRCH)
+				return 0;
+			return -errno;
+		}
+		usleep(TASK_WORK_POLL_US);
+	}
+	return -ETIMEDOUT;
+}
+
+/*
+ * Run one attempt on dedicated workers and report its outcome.  The
+ * workers are stopped, joined and waited out of PID lookup before another
+ * attempt can reset the shared state.  join alone does not wait for their
+ * exit task_work to finish.
+ */
+static enum task_work_busy_attempt
+task_work_busy_run_attempt(struct task_work_busy *skel, int seq,
+			   int trigger_cpu, int target_cpu, __s64 *out_wait_ns)
+{
+	struct task_work_race_ctx ctx = {};
+	struct task_work_race_thread_arg trigger_arg = {
+		.ctx = &ctx,
+		.cpu = trigger_cpu,
+		.tid = &ctx.trigger_tid,
+	};
+	struct task_work_race_thread_arg target_arg = {
+		.ctx = &ctx,
+		.cpu = target_cpu,
+		.tid = &ctx.target_tid,
+	};
+	enum task_work_busy_attempt result = TWB_HARD_FAIL;
+	pthread_t trigger, target;
+	bool trigger_started = false, target_started = false;
+	bool cleanup_ok = true;
+	__s64 window = 0, wait_ns = 0, count = -1;
+	__s32 res = 0;
+	int err, i;
+
+	*out_wait_ns = 0;
+	/* Do not let a reused TID match the previous attempt's filter. */
+	skel->bss->trigger_tid = 0;
+	skel->bss->target_tid = 0;
+
+	/* Clear old state before enabling workers that can read it. */
+	if (!ASSERT_OK(task_work_busy_prepare_elem(skel, seq), "prepare elem"))
+		goto out;
+	if (!ASSERT_OK(task_work_busy_reset(skel), "reset slots"))
+		goto out;
+
+	err = pthread_create(&target, NULL, task_work_race_worker, &target_arg);
+	if (err) {
+		ASSERT_OK(err, "pthread_create target");
+		goto out;
+	}
+	target_started = true;
+	err = pthread_create(&trigger, NULL, task_work_race_worker, &trigger_arg);
+	if (err) {
+		ASSERT_OK(err, "pthread_create trigger");
+		goto out;
+	}
+	trigger_started = true;
+
+	for (i = 0; i < TASK_WORK_START_POLL_ITERS; i++) {
+		if (__atomic_load_n(&ctx.trigger_tid, __ATOMIC_ACQUIRE) &&
+		    __atomic_load_n(&ctx.target_tid, __ATOMIC_ACQUIRE))
+			break;
+		usleep(TASK_WORK_POLL_US);
+	}
+	if (!ASSERT_TRUE(__atomic_load_n(&ctx.trigger_tid, __ATOMIC_ACQUIRE) &&
+			 __atomic_load_n(&ctx.target_tid, __ATOMIC_ACQUIRE),
+			 "thread start"))
+		goto out;
+	if (!ASSERT_OK(__atomic_load_n(&ctx.setup_err, __ATOMIC_ACQUIRE),
+		       "thread affinity"))
+		goto out;
+
+	skel->bss->trigger_tid = __atomic_load_n(&ctx.trigger_tid, __ATOMIC_ACQUIRE);
+	skel->bss->target_tid = __atomic_load_n(&ctx.target_tid, __ATOMIC_ACQUIRE);
+
+	/* Arm last, once this attempt's value, slots and TIDs are in place. */
+	if (!ASSERT_OK(task_work_set_status(skel->maps.status, BUSY_ARM_SEQ, seq),
+		       "arm"))
+		goto out;
+
+	/* A fresh value has no expected BUSY source: check the first result. */
+	err = task_work_wait_seq(skel->maps.status, BUSY_GEN1_RESULT, seq, &res);
+	if (!ASSERT_OK(err, "gen1 result"))
+		goto out;
+	if (!ASSERT_OK(res, "gen1 schedule")) {
+		fprintf(stderr, "attempt %d: gen1 failed: %d\n", seq, res);
+		goto out;
+	}
+
+	if (!ASSERT_OK(task_work_busy_wait_window(skel, seq, &window),
+		       "window close"))
+		goto out;
+
+	if ((__u32)window == 3) {
+		/*
+		 * The attempt did not arrive inside the wait bound.  The
+		 * elapsed time is a diagnostic only; count
+		 * the miss and retry with a fresh attempt.
+		 */
+		result = TWB_MISSED;
+		goto out;
+	}
+
+	/*
+	 * Valid window: the handler saw this sequence's attempt before
+	 * returning, so the attempt ran while the scheduling irq_work was
+	 * still BUSY.  The elapsed time is a diagnostic only: host vCPU
+	 * preemption can stretch wall-clock time without affecting the
+	 * gate behavior under test.
+	 */
+	err = task_work_wait_seq(skel->maps.status, BUSY_ACK, seq, &res);
+	if (!ASSERT_OK(err, "in-window attempt"))
+		goto out;
+	if (res != -EBUSY) {
+		fprintf(stderr,
+			"attempt %d: in-window reuse %s (%d), expected -EBUSY\n",
+			seq, res == 0 ? "accepted" : "failed", res);
+		ASSERT_EQ(res, -EBUSY, "in-window BUSY gate rejection");
+		goto out;
+	}
+
+	/*
+	 * The window was released; retry until BUSY clears.  The second
+	 * callback must then run exactly once.
+	 */
+	if (!ASSERT_OK(task_work_set_status(skel->maps.status, BUSY_RECOVER_SEQ, seq),
+		       "enable recovery"))
+		goto out;
+	err = task_work_wait_seq(skel->maps.status, BUSY_GEN2_RESULT, seq, &res);
+	if (!ASSERT_OK(err, "recovery result"))
+		goto out;
+	if (!ASSERT_OK(res, "recovery schedule")) {
+		fprintf(stderr, "attempt %d: recovery failed: %d\n", seq, res);
+		goto out;
+	}
+	if (!ASSERT_OK(task_work_busy_wait_value(skel, BUSY_CB2_SEQ, seq),
+		       "gen2 callback"))
+		goto out;
+	result = TWB_VALID;
+out:
+	__atomic_store_n(&ctx.stop, 1, __ATOMIC_RELEASE);
+	if (trigger_started &&
+	    (!ASSERT_OK(pthread_join(trigger, NULL), "trigger join") ||
+	     !ASSERT_OK(task_work_busy_wait_exit(ctx.trigger_tid), "trigger exit"))) {
+		cleanup_ok = false;
+		result = TWB_HARD_FAIL;
+	}
+	if (target_started &&
+	    (!ASSERT_OK(pthread_join(target, NULL), "target join") ||
+	     !ASSERT_OK(task_work_busy_wait_exit(ctx.target_tid), "target exit"))) {
+		cleanup_ok = false;
+		result = TWB_HARD_FAIL;
+	}
+	/* Both workers have passed exit_task_work; the count is now stable. */
+	if (result == TWB_VALID) {
+		if (!ASSERT_OK(task_work_status(skel->maps.status, BUSY_CB2_COUNT, &count),
+			       "cb2 count")) {
+			result = TWB_HARD_FAIL;
+		} else if (!ASSERT_EQ(count, 1, "gen2 callback exactly once")) {
+			fprintf(stderr, "attempt %d: gen2 callback count %lld\n",
+				seq, (long long)count);
+			result = TWB_HARD_FAIL;
+		}
+	}
+	/*
+	 * Only a captured final window belongs to this attempt.  The fexit
+	 * runs in the trigger's interrupted context, so successful worker
+	 * cleanup also completes its diagnostic write.  Read once here;
+	 * the elapsed value does not decide whether the attempt passed.
+	 */
+	if (cleanup_ok && window) {
+		err = task_work_status(skel->maps.status, BUSY_LAST_WAIT_NS, &wait_ns);
+		if (err) {
+			err = -errno;
+			if (result != TWB_HARD_FAIL) {
+				ASSERT_OK(err, "wait diagnostic");
+				result = TWB_HARD_FAIL;
+			} else {
+				fprintf(stderr,
+					"attempt %d: elapsed diagnostic unavailable: %d\n",
+					seq, err);
+			}
+		}
+	}
+	*out_wait_ns = wait_ns;
+	fprintf(stderr, "attempt %d: %s (window %u wait_ns %lld)\n",
+		seq, twb_attempt_name(result), (__u32)window, (long long)wait_ns);
+	return result;
+}
+
+void serial_test_task_work_busy(void)
+{
+	struct task_work_busy *skel;
+	struct bpf_link *link_tp = NULL, *link_fexit = NULL;
+	cpu_set_t allowed;
+	int cpu, cpu_count = 0, trigger_cpu = -1, target_cpu = -1;
+	int attempt, attempts_started = 0, valid = 0, missed = 0;
+	bool hard_fail = false;
+	__s64 max_wait_ns = 0;
+
+	if (sched_getaffinity(0, sizeof(allowed), &allowed)) {
+		ASSERT_OK(-errno, "sched_getaffinity");
+		return;
+	}
+	for (cpu = 0; cpu < CPU_SETSIZE && cpu_count < 2; cpu++) {
+		if (!CPU_ISSET(cpu, &allowed))
+			continue;
+		if (!cpu_count)
+			trigger_cpu = cpu;
+		else
+			target_cpu = cpu;
+		cpu_count++;
+	}
+	if (cpu_count < 2) {
+		printf("%s:SKIP:need two CPUs in the process affinity mask\n",
+		       __func__);
+		test__skip();
+		return;
+	}
+
+	skel = task_work_busy__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "open_and_load"))
+		return;
+	if (skel->kconfig->CONFIG_PREEMPT_RT) {
+		printf("%s:SKIP:CONFIG_PREEMPT_RT=y runs irq_work in a kernel thread\n",
+		       __func__);
+		test__skip();
+		goto cleanup;
+	}
+
+	link_tp = bpf_program__attach_tracepoint(skel->progs.busy_sched_work,
+						 "syscalls",
+						 "sys_enter_getppid");
+	if (!ASSERT_OK_PTR(link_tp, "attach tracepoint"))
+		goto cleanup;
+
+	link_fexit = bpf_program__attach_trace(skel->progs.busy_fexit_task_work_add);
+	if (!ASSERT_OK_PTR(link_fexit, "attach fexit"))
+		goto cleanup;
+
+	for (attempt = 1;
+	     attempt <= TASK_WORK_BUSY_MAX_ATTEMPTS &&
+	     valid < TASK_WORK_BUSY_ROUNDS;
+	     attempt++) {
+		enum task_work_busy_attempt res;
+		__s64 wait_ns = 0;
+
+		attempts_started = attempt;
+		res = task_work_busy_run_attempt(skel, attempt, trigger_cpu,
+						 target_cpu, &wait_ns);
+		switch (res) {
+		case TWB_VALID:
+			valid++;
+			if (wait_ns > max_wait_ns)
+				max_wait_ns = wait_ns;
+			break;
+		case TWB_MISSED:
+			missed++;
+			break;
+		case TWB_HARD_FAIL:
+			hard_fail = true;
+			break;
+		}
+		if (hard_fail)
+			break;
+	}
+
+	fprintf(stderr,
+		"task_work_busy: %d/%d valid, %d missed, %d attempts, max_wait %lld ns\n",
+		valid, TASK_WORK_BUSY_ROUNDS, missed, attempts_started,
+		(long long)max_wait_ns);
+
+	/*
+	 * The coverage assertion only makes sense when no hard
+	 * failure aborted the loop; after a hard failure keep the primary
+	 * reason instead of deriving more assertion noise from it.
+	 */
+	if (!hard_fail)
+		ASSERT_EQ(valid, TASK_WORK_BUSY_ROUNDS, "valid busy windows");
+cleanup:
+	bpf_link__destroy(link_tp);
+	bpf_link__destroy(link_fexit);
+	task_work_busy__destroy(skel);
+}
diff --git a/tools/testing/selftests/bpf/progs/task_work_busy.c b/tools/testing/selftests/bpf/progs/task_work_busy.c
new file mode 100644
index 000000000000..9960011edadc
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/task_work_busy.c
@@ -0,0 +1,325 @@
+// SPDX-License-Identifier: GPL-2.0
+/* Copyright (c) 2026 KylinSoft Corporation. */
+#include <vmlinux.h>
+#include <stdbool.h>
+#include <errno.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+#include <bpf/bpf_core_read.h>
+
+char _license[] SEC("license") = "GPL";
+
+extern bool CONFIG_PREEMPT_RT __kconfig __weak;
+
+/*
+ * Bounds for the hardirq busy-wait that holds the scheduling handler
+ * inside the real window: task_work_add() has published the callback,
+ * but the handler's tail state transition has not run yet.
+ * irq_work_single() clears BUSY only after the handler returns, so
+ * while this wait runs, the context's scheduling irq_work is BUSY,
+ * even though the first callback may already have completed and
+ * returned the context to STANDBY.
+ *
+ * The wait is bounded both by a bpf_ktime_get_ns() check and by an
+ * iteration budget.  The time bound is checked against a monotonic
+ * clock: host vCPU preemption can stretch the wall-clock duration of
+ * the window, which only makes the window longer, never shorter.
+ */
+#define BUSY_WAIT_NS		4000000ULL
+#define BUSY_WAIT_ITERS		1000000ULL
+
+int trigger_tid;
+int target_tid;
+
+/*
+ * High 32 bits of the result slots hold the sequence, low 32 bits the
+ * signed error.  BUSY_WINDOW low bits: 1 open, 2 the in-window attempt
+ * was seen before the handler returned, 3 the wait bound was reached
+ * without it.  Must match the test.
+ */
+enum task_work_busy_status {
+	BUSY_ARM_SEQ,
+	BUSY_WINDOW,
+	BUSY_ACK,
+	BUSY_GEN1_RESULT,
+	BUSY_GEN2_RESULT,
+	BUSY_CB1_SEQ,
+	BUSY_CB2_COUNT,
+	BUSY_CB2_SEQ,
+	BUSY_RECOVER_SEQ,
+	BUSY_LAST_WAIT_NS,
+	BUSY_STATUS_MAX,
+};
+
+struct {
+	__uint(type, BPF_MAP_TYPE_ARRAY);
+	__uint(max_entries, BUSY_STATUS_MAX);
+	__type(key, int);
+	__type(value, __s64);
+} status SEC(".maps");
+
+struct task_work_busy_value {
+	__u32 seq;
+	char data[60];
+	struct bpf_task_work tw;
+};
+
+/*
+ * max_entries must cover TASK_WORK_BUSY_MAX_ATTEMPTS in the test: each
+ * attempt schedules into its own fresh value.
+ */
+struct {
+	__uint(type, BPF_MAP_TYPE_HASH);
+	__uint(map_flags, BPF_F_NO_PREALLOC);
+	__uint(max_entries, 128);
+	__type(key, int);
+	__type(value, struct task_work_busy_value);
+} hmap SEC(".maps");
+
+static __always_inline void set_status(int key, __s64 value)
+{
+	__s64 *slot;
+
+	slot = bpf_map_lookup_elem(&status, &key);
+	if (slot)
+		*slot = value;
+}
+
+static __always_inline __s64 get_status(int key)
+{
+	__s64 *slot = bpf_map_lookup_elem(&status, &key);
+
+	return slot ? *slot : 0;
+}
+
+struct busy_wait_ctx {
+	__u64 seq;
+	__u64 start;
+	bool acked;
+};
+
+/* bpf_loop() callback: return 0 continues the loop, return 1 stops it. */
+static long busy_wait_cb(__u32 i, void *pctx)
+{
+	struct busy_wait_ctx *w = pctx;
+
+	if (((__u64)get_status(BUSY_ACK) >> 32) == w->seq) {
+		w->acked = true;
+		return 1;
+	}
+	if (bpf_ktime_get_ns() - w->start > BUSY_WAIT_NS)
+		return 1;
+	return 0;
+}
+
+static int busy_cb_gen1(struct bpf_map *map, void *key, void *value)
+{
+	struct task_work_busy_value *work = value;
+	__u64 we = (__u64)get_status(BUSY_WINDOW) >> 32;
+
+	/*
+	 * Backstop for deliveries from abandoned attempts: accept only
+	 * while no window for a different sequence is published.  The
+	 * test waits for the old workers to disappear from PID lookup,
+	 * after their exit task_work, before resetting the shared state.
+	 * This callback may run before fexit publishes the window state.
+	 */
+	if (we && we != (__u64)work->seq)
+		return 0;
+	set_status(BUSY_CB1_SEQ, work->seq);
+	return 0;
+}
+
+static int busy_cb_gen2(struct bpf_map *map, void *key, void *value)
+{
+	struct task_work_busy_value *work = value;
+	__s64 *count;
+	int key2 = BUSY_CB2_COUNT;
+
+	/*
+	 * Count only deliveries belonging to the sequence whose window is
+	 * (or was) published: a late delivery from an abandoned attempt
+	 * must not satisfy a later exactly-once check.
+	 */
+	if ((__u64)work->seq != ((__u64)get_status(BUSY_WINDOW) >> 32))
+		return 0;
+	set_status(BUSY_CB2_SEQ, work->seq);
+	count = bpf_map_lookup_elem(&status, &key2);
+	if (count)
+		__sync_fetch_and_add(count, 1);
+	return 0;
+}
+
+SEC("fexit/task_work_add")
+int BPF_PROG(busy_fexit_task_work_add, struct task_struct *task,
+	     struct callback_head *work, int notify, int retval)
+{
+	struct busy_wait_ctx w = {};
+	__u64 seq, elapsed;
+
+	seq = (__u64)get_status(BUSY_ARM_SEQ);
+	if (!seq)
+		return 0;
+
+	/*
+	 * One-shot arm.  Only this test's own first-round add may consume
+	 * it and hold the window: a successful TWA_SIGNAL add of the
+	 * bpf_task_work callback, issued from the trigger thread's CPU by
+	 * the self-IPI'd scheduling handler and targeting the target
+	 * thread.  Everything else leaves the arm untouched.
+	 */
+	if (retval != 0 || notify != TWA_SIGNAL ||
+	    BPF_CORE_READ(task, pid) != target_tid)
+		return 0;
+	if ((__u32)bpf_get_current_pid_tgid() != (__u32)trigger_tid)
+		return 0;
+
+	set_status(BUSY_ARM_SEQ, 0);
+	set_status(BUSY_WINDOW, (seq << 32) | 1);
+
+	w.seq = seq;
+	w.start = bpf_ktime_get_ns();
+	bpf_loop(BUSY_WAIT_ITERS, busy_wait_cb, &w, 0);
+	elapsed = bpf_ktime_get_ns() - w.start;
+
+	/*
+	 * Publish elapsed time for logging after worker cleanup.  Window
+	 * validity depends on seeing ACK here, not on this diagnostic.
+	 */
+	set_status(BUSY_LAST_WAIT_NS, (__s64)elapsed);
+	if (w.acked)
+		set_status(BUSY_WINDOW, (seq << 32) | 2);
+	else
+		set_status(BUSY_WINDOW, (seq << 32) | 3);
+	return 0;
+}
+
+/*
+ * Attempt the second scheduling round on the same context, targeting the
+ * trigger thread.  The attempt uses the value keyed by its own sequence,
+ * so it can never touch another attempt's context.  Returns true if a
+ * terminal result was published into res_key; a transient -EBUSY is left
+ * unpublished for the caller's retry policy, mirroring the race test's
+ * convention.
+ */
+static __always_inline bool busy_try_gen2(__u64 seq, int res_key)
+{
+	struct task_work_busy_value *work;
+	struct task_struct *task;
+	int key = (int)seq, err;
+
+	task = bpf_task_from_pid(trigger_tid);
+	if (!task) {
+		set_status(res_key, ((__s64)seq << 32) | (__u32)-ESRCH);
+		return true;
+	}
+	work = bpf_map_lookup_elem(&hmap, &key);
+	if (!work) {
+		bpf_task_release(task);
+		set_status(res_key, ((__s64)seq << 32) | (__u32)-ENOENT);
+		return true;
+	}
+	work->seq = seq;
+	err = bpf_task_work_schedule_signal(task, &work->tw, &hmap,
+					    busy_cb_gen2);
+	bpf_task_release(task);
+	if (err == -EBUSY)
+		return false;
+	set_status(res_key, ((__s64)seq << 32) | (__u32)err);
+	return true;
+}
+
+SEC("tracepoint/syscalls/sys_enter_getppid")
+int busy_sched_work(void *ctx)
+{
+	__u32 tid = (__u32)bpf_get_current_pid_tgid();
+	struct task_work_busy_value *work;
+	struct task_struct *task;
+	__u64 window, seq;
+	int key, err;
+
+	/* Also keep the kconfig extern in object BTF for the userspace skip. */
+	if (CONFIG_PREEMPT_RT)
+		return 0;
+
+	if (tid == (__u32)trigger_tid) {
+		/*
+		 * Trigger thread: schedule the first round once per armed
+		 * sequence, targeting the target thread.  The kfunc queues
+		 * the scheduling irq_work and returns; the handler then
+		 * parks inside the fexit window above.  A fresh value has no
+		 * expected BUSY source, so publish every scheduling result.
+		 */
+		seq = (__u64)get_status(BUSY_ARM_SEQ);
+		if (!seq ||
+		    ((__u64)get_status(BUSY_GEN1_RESULT) >> 32) == seq)
+			return 0;
+		key = (int)seq;
+		task = bpf_task_from_pid(target_tid);
+		if (!task) {
+			set_status(BUSY_GEN1_RESULT,
+				   ((__s64)seq << 32) | (__u32)-ESRCH);
+			set_status(BUSY_ARM_SEQ, 0);
+			return 0;
+		}
+		work = bpf_map_lookup_elem(&hmap, &key);
+		if (!work) {
+			bpf_task_release(task);
+			set_status(BUSY_GEN1_RESULT,
+				   ((__s64)seq << 32) | (__u32)-ENOENT);
+			set_status(BUSY_ARM_SEQ, 0);
+			return 0;
+		}
+		work->seq = seq;
+		err = bpf_task_work_schedule_signal(task, &work->tw, &hmap,
+						    busy_cb_gen1);
+		bpf_task_release(task);
+		set_status(BUSY_GEN1_RESULT, ((__s64)seq << 32) | (__u32)err);
+		return 0;
+	}
+
+	if (tid != (__u32)target_tid)
+		return 0;
+
+	window = (__u64)get_status(BUSY_WINDOW);
+	seq = window >> 32;
+	if (!seq || (__u32)window == 0)
+		return 0;
+	/*
+	 * The first callback must have completed for this sequence.  Being
+	 * inside the target thread's own subsequent syscall proves the
+	 * callback wrapper has returned and the context is STANDBY, which
+	 * is stronger than observing the completion marker from another
+	 * thread.
+	 */
+	if ((__u64)get_status(BUSY_CB1_SEQ) != seq)
+		return 0;
+
+	if ((__u64)get_status(BUSY_ACK) >> 32 != seq) {
+		/*
+		 * The in-window attempt, exactly once per sequence.  On a
+		 * fixed kernel it must be rejected with -EBUSY: the context
+		 * is STANDBY, the map and program references are alive and
+		 * no other scheduler exists, so only the BUSY gate can
+		 * reject here.
+		 */
+		if (!busy_try_gen2(seq, BUSY_ACK))
+			set_status(BUSY_ACK,
+				   ((__s64)seq << 32) | (__u32)-EBUSY);
+		return 0;
+	}
+
+	/*
+	 * Recovery, gated by the test: the fexit program releasing the
+	 * window does not mean the whole old handler has returned and
+	 * BUSY has cleared yet, so retry until the schedule succeeds.
+	 */
+	if ((__u64)get_status(BUSY_RECOVER_SEQ) != seq)
+		return 0;
+	if ((__s32)get_status(BUSY_ACK) == 0)
+		return 0;	/* scheduled in-window; nothing to recover */
+	if (((__u64)get_status(BUSY_GEN2_RESULT) >> 32) == seq)
+		return 0;	/* terminal recovery result published */
+	busy_try_gen2(seq, BUSY_GEN2_RESULT);
+	return 0;
+}
-- 
2.43.0