[PATCH 0/2] objpool: fix nested pushes from NMI context

[PATCH 0/2] objpool: fix nested pushes from NMI context

[Shashank Mohan Jain]

objpool_push() runs with interrupts disabled, but a kretprobe that
returns in NMI context can push to the same per-CPU slot in the middle
of it. With rethook-based kretprobes (and, before v6.14, fprobe) this is
reachable since kretprobes moved to objpool in v6.7. The nested push
publishes slot->last past the unwritten entry of the interrupted push,
so a pop on another CPU can take a NULL or stale pointer, which hands
an object out twice and loses another, and slot->last can move
backwards, after which a pop on the owning CPU spins with interrupts
disabled.

Patch 1 publishes the entries in order with a cmpxchg() on slot->last.
Patch 2 adds a KUnit test in which a pinned hard hrtimer stands in for
the NMI. It fails without patch 1 and passes with it.

The race was found with a TLA+ model of __objpool_try_add_slot() and
__objpool_try_get_slot() on one slot: a task push that an NMI push can
interrupt at every step, pops on another CPU, and recycled objects.
Memory is sequentially consistent, with one level of nesting. TLC finds
the bug in current code, no violation with patch 1 (up to five objects,
four task pushes, three nested pushes, five remote pops), and shows that
a plain-store version of the fix lets last fall behind head. The model
(Objpool.tla and its .cfg files) can be posted if that is useful.

No earlier report of this was found: searches of patchwork (objpool,
rethook, kretprobe, objpool_push) and of the linux-kernel archive on
marc.info (objpool NMI, objpool_push, rethook NMI, kretprobe NMI) turned
up nothing related. lore was not searched.

Testing is described under "---" in each patch. In short: KUnit on UML
x86_64 (4 CPUs, 3 runs before and after, plus 1 CPU and CONFIG_SMP=n),
W=1 builds for x86_64 and i386, checkpatch --strict. Not tested: real
NMIs through a kretprobe on hardware, weakly ordered architectures, and
performance in the kernel.

This series was prepared with Claude Code (Anthropic), model Claude
Opus 5.5 (claude-opus-5-5). The code, the test, the changelogs and this
cover letter were written with the assistant; the TLA+ model checker
TLC found the race.

Shashank Mohan Jain (2):
  objpool: keep objpool_push() correct when a push from NMI nests in it
  lib/tests: add KUnit test for nested objpool pushes

 MAINTAINERS               |   1 +
 include/linux/objpool.h   |  37 ++++-
 lib/Kconfig.debug         |  13 ++
 lib/tests/Makefile        |   1 +
 lib/tests/objpool_kunit.c | 329 ++++++++++++++++++++++++++++++++++++++
 5 files changed, 373 insertions(+), 8 deletions(-)
 create mode 100644 lib/tests/objpool_kunit.c

-- 
2.43.0

[PATCH 1/2] objpool: keep objpool_push() correct when a push from NMI nests in it

[Shashank Mohan Jain]

objpool_push() adds an object to the slot of the local CPU with
interrupts disabled. It reserves an entry with a cmpxchg() on
slot->tail, writes the entry, and publishes it with
smp_store_release(&slot->last, tail + 1).

A push from NMI context can still interrupt it and push to the same
slot. kretprobes may run in NMI context since commit e03b4a084ea6
("kprobes: Remove NMI context check"). At the time, kretprobe instances
came from a CAS-based lockless freelist, which tolerates that. Commit
4bbd93455659 ("kprobes: kretprobe scalability improvement") moved
kretprobes and rethook to objpool.

With rethook, rethook_trampoline_handler() recycles instances with
objpool_push() after the user handler has run, when no kprobe is marked
running anymore; rethook_flush_task() does the same. An NMI that arrives
during such a push and runs a function probed by the same kretprobe takes
an instance in pre_handler_kretprobe(), and when the function returns
inside the NMI, pushes it back to the same slot. This needs a kretprobe
on a function that runs both in NMI context and outside it, for instance
one that perf calls from the PMU NMI handler on x86. Before v6.14, fprobe
also used rethook and could push from NMI the same way, with one pool
shared by all functions of an fprobe.

kretprobes without rethook are not affected: kretprobe_trampoline_handler()
and kprobe_flush_task() run under kprobe_busy_begin(), so a kprobe hit in
NMI context is counted as missed.

If the NMI lands between the reservation and the publication of the
interrupted push, it reserves the next entry, writes it and moves
slot->last past the entry of the interrupted push, which is not written
yet:

  CPU0 (irqs off)          CPU0 NMI              CPU1
  ---------------          --------              ----
  reserve entry t
                           reserve entry t+1
                           write entry t+1
                           last = t+2
                                                 pop: t < last,
                                                 returns entries[t]
  write entry t
  last = t+1

This has three effects:

- A pop on another CPU can take entry t before it is written. It gets
  whatever the ring held at that position: NULL, or a stale pointer to
  an object that was popped earlier and may still be in use, so the same
  object is handed out twice. The object pushed at t can never be popped
  again, so it is lost to the pool.
- The interrupted push moves slot->last backwards. If both entries were
  popped in between, head is now ahead of last, and
  __objpool_try_get_slot() spins with interrupts disabled until the next
  push to that slot. Only the owning CPU pushes there, so a pop on that
  CPU never comes back.
- Even without a concurrent pop, the entry of the NMI stays invisible
  until the next push. The WARN_ON_ONCE() sanity check can also fire
  spuriously, because it compares a snapshot of tail taken before the
  NMI with head read after it.

During the review of objpool, nested pushes were assumed not to happen
because the push runs with interrupts disabled. kretprobes push from NMI
context, though.

Publish the entries in order instead:

- A push advances slot->last only while last points at its own entry,
  meaning all earlier entries are written. It then also publishes the
  entries of pushes that interrupted it, which have completed by the
  time it resumes.
- A push that interrupted another one leaves slot->last alone. The
  interrupted push publishes both entries when it resumes.
- This needs a cmpxchg(). Masami sketched a plain-store version of
  this catch-up loop during the objpool review [1]. With a plain store,
  a nested push can take over as soon as last reaches its entry and
  publish it, and the store of an older tail by the interrupted push
  then moves last backwards, below head if a remote pop took the
  entries meanwhile. A pop from NMI on this CPU would then spin
  forever, as the interrupted push cannot resume to repair last.
- Read head for the sanity check only once the entry is reserved.

Without nesting, the push now costs one cmpxchg() and one extra read of
tail instead of one store.

Found with a TLA+ model of the push and pop paths (sequentially
consistent memory, one level of nesting), with a task push that an NMI
push can interrupt at every step and pops on another CPU. TLC finds pops
of unwritten entries, objects handed out twice and last moving backwards
with the current code. With this change it finds no violation for up to
five objects, four task pushes, three nested pushes and five remote
pops. In the same model, the plain-store variant hands out no bad
objects but lets last fall behind head.

On 4-CPU UML, the KUnit test added in the next patch lets an hrtimer
push while a task push is in progress. Before this change, it leaves
6960 to 7053 entries unpublished in 5 seconds. With a popper on another
CPU, it hands out objects twice and loses them until all 30 objects of
the task are gone. After this change, none of these happen (3 runs
each).

Link: https://lore.kernel.org/all/20231012230237.45726dfad125cf0e8d00ba53@kernel.org/ [1]
Fixes: b4edb8d2d464 ("lib: objpool added: ring-array based lockless MPMC")
Cc: stable@vger.kernel.org
Assisted-by: LLM TLC
Signed-off-by: Shashank Mohan Jain <jain.sm@gmail.com>
---
Notes (not part of the changelog):

Tested (master 72d3fcf802c4, v7.3-rc5):
- KUnit on UML x86_64 with ncpus=4 (CONFIG_SMP=y, HIGH_RES_TIMERS),
  with patch 2/2 applied, 3 runs each:
  - Without this patch, both cases fail every run. Local case: 6960-7053
    unpublished entries (8230-8391 timer pushes nested in a task push).
    Remote-pop case: 142-187 unpublished entries, 30 double handouts and
    30 objects lost (all of the task's objects). The spurious
    WARN_ON_ONCE() at objpool.h:202 fires in every run.
  - With this patch, both cases pass every run: 0 unpublished, 0 double
    handouts, 0 lost, with 5013-5237 (local) and 32994-35223 (remote
    pop) real nestings, and no WARN.
- KUnit on UML with ncpus=1: the local case fails without this patch
  (7112 unpublished entries) and passes with it (4972 nestings); the
  remote-pop case is skipped. With CONFIG_SMP=n it passes with this
  patch (72123 nestings).
- W=1 build of lib/objpool.o, lib/test_objpool.o,
  lib/tests/objpool_kunit.o, kernel/kprobes.o and kernel/trace/rethook.o
  for x86_64_defconfig and i386_defconfig (plus KPROBES, KRETPROBES,
  KRETPROBE_ON_RETHOOK, KUNIT, OBJPOOL_KUNIT_TEST=m, TEST_OBJPOOL=m):
  no warnings.
- checkpatch.pl --strict: clean apart from the missing Signed-off-by.
- git apply --check of this patch against include/linux/objpool.h of
  v6.12, v6.18 and v7.2: applies. It was not built or tested there.

Not tested:
- Real NMIs. The failure was not reproduced through a kretprobe on real
  hardware; an hrtimer on UML stands in for the NMI. The reachability
  argument (rethook recycles with no kprobe marked running) comes from
  reading the code. The hard lockup of a pop on the owning CPU follows
  from the code and the TLA+ model; it was not reproduced.
- Weakly ordered architectures. The pop side still reads slot->last
  without acquire; that is pre-existing and not addressed here.
- Performance in the kernel. A userspace model of the fast path on an
  i7-11700F goes from ~14 to ~22 ns per push+pop pair (one more locked
  cmpxchg); not measured in the kernel. lib/test_objpool.c was not run.
- The Link: message id was taken from patchwork; the lore URL was not
  opened.

An alternative would be to mark a kprobe busy around the recycle loops
in rethook_trampoline_handler() and rethook_flush_task(), as the
non-rethook kretprobe path does. That would not cover fprobe before
v6.14 and would keep objpool unsafe for pushes from NMI, so this fixes
objpool itself.

This patch was prepared with Claude Code (Anthropic), model Claude Opus
5.5 (claude-opus-5-5): the code and the changelog were written with the
assistant. The TLA+ model checker TLC found the race.

 include/linux/objpool.h | 37 +++++++++++++++++++++++++++++--------
 1 file changed, 29 insertions(+), 8 deletions(-)

diff --git a/include/linux/objpool.h b/include/linux/objpool.h
index b713a1fe7521..b86225e32ebe 100644
--- a/include/linux/objpool.h
+++ b/include/linux/objpool.h
@@ -193,19 +193,40 @@ __objpool_try_add_slot(void *obj, struct objpool_head *pool, int cpu)
 	struct objpool_slot *slot = pool->cpu_slots[cpu];
 	uint32_t head, tail;
 
-	/* loading tail and head as a local snapshot, tail first */
+	/*
+	 * Only the local CPU pushes to its slot, with irqs disabled, but a
+	 * push from NMI context (a kretprobe'd function returning in NMI)
+	 * can interrupt this one at any point.
+	 */
 	tail = READ_ONCE(slot->tail);
+	while (!try_cmpxchg_acquire(&slot->tail, &tail, tail + 1))
+		;
 
-	do {
-		head = READ_ONCE(slot->head);
-		/* fault caught: something must be wrong */
-		WARN_ON_ONCE(tail - head > pool->nr_objs);
-	} while (!try_cmpxchg_acquire(&slot->tail, &tail, tail + 1));
+	/*
+	 * fault caught: something must be wrong.  Read head only after the
+	 * reservation: a nested push and a pop on another CPU could have
+	 * moved head past an older snapshot of tail.
+	 */
+	head = READ_ONCE(slot->head);
+	WARN_ON_ONCE(tail - head > pool->nr_objs);
 
 	/* now the tail position is reserved for the given obj */
 	WRITE_ONCE(slot->entries[tail & slot->mask], obj);
-	/* update sequence to make this obj available for pop() */
-	smp_store_release(&slot->last, tail + 1);
+
+	/*
+	 * Update 'last' to make this obj available for pop(), in order:
+	 * 'last' must never move past an entry that is not written yet.
+	 * If this push interrupted another push to this slot, the entry of
+	 * the interrupted push comes first and is not written yet: leave
+	 * 'last' alone, the interrupted push publishes both when it resumes.
+	 * Pushes that interrupted this one have completed by now, so publish
+	 * their entries too.  A plain store is not enough: another nested
+	 * push can take over publishing as soon as 'last' reaches its entry.
+	 */
+	while (try_cmpxchg_release(&slot->last, &tail, tail + 1)) {
+		if (++tail == READ_ONCE(slot->tail))
+			break;
+	}
 
 	return 0;
 }
-- 
2.43.0

[PATCH 2/2] lib/tests: add KUnit test for nested objpool pushes

[Shashank Mohan Jain]

objpool_push() only pushes to the slot of the local CPU, but a push from
NMI context (a kretprobe'd function returning in NMI) can interrupt it
and push to the same slot.

UML has no NMIs. So the test pushes from task context with interrupts
enabled, through __objpool_try_add_slot(). A pinned hard hrtimer on the
same CPU pushes with objpool_push() every 10 us. That nests the same way
an NMI would. HRTIMER_MODE_REL_PINNED_HARD keeps the timer in hardirq
context on PREEMPT_RT as well.

- objpool_test_nested_push checks that no entry is left unpublished
  once the task push returns (slot->last == slot->tail).
- objpool_test_nested_push_remote_pop pops from another CPU with
  objpool_pop(). It checks that every popped object was in the pool,
  and that no object is lost from the pool. It is skipped when only one
  CPU is online.

Only a few timer pushes land inside a task push. The test counts those
and fails if there were none, as it would then have tested nothing.

Each case runs for 5 seconds.

Assisted-by: LLM
Signed-off-by: Shashank Mohan Jain <jain.sm@gmail.com>
---
Notes (not part of the changelog):

Tested (master 72d3fcf802c4, v7.3-rc5):
- UML x86_64 with ncpus=4, 3 runs without patch 1/2 (both cases fail)
  and 3 runs with it (both pass). Real nestings per case with patch 1/2:
  5013-5237 (local) and 32994-35223 (remote pop).
- UML with ncpus=1: the local case fails without patch 1/2 and passes
  with it; the remote-pop case is skipped. CONFIG_SMP=n: the local case
  passes with patch 1/2.
- W=1 build for x86_64 and i386, with OBJPOOL_KUNIT_TEST=m: no warnings.
- checkpatch.pl --strict: clean apart from the missing Signed-off-by.

Not tested:
- Architectures other than UML x86_64 (the test was only built for
  x86_64 and i386), and PREEMPT_RT.

It calls the internal __objpool_try_add_slot() with interrupts enabled
on purpose: that is what lets the hrtimer nest.

This patch was prepared with Claude Code (Anthropic), model Claude Opus
5.5 (claude-opus-5-5): the test and the changelog were written with the
assistant.

 MAINTAINERS               |   1 +
 lib/Kconfig.debug         |  13 ++
 lib/tests/Makefile        |   1 +
 lib/tests/objpool_kunit.c | 329 ++++++++++++++++++++++++++++++++++++++
 4 files changed, 344 insertions(+)
 create mode 100644 lib/tests/objpool_kunit.c

diff --git a/MAINTAINERS b/MAINTAINERS
index 72294ddfa5b7..0d319f2c2948 100644
--- a/MAINTAINERS
+++ b/MAINTAINERS
@@ -19872,6 +19872,7 @@ S:	Supported
 F:	include/linux/objpool.h
 F:	lib/objpool.c
 F:	lib/test_objpool.c
+F:	lib/tests/objpool_kunit.c
 
 OBJTOOL
 M:	Josh Poimboeuf <jpoimboe@kernel.org>
diff --git a/lib/Kconfig.debug b/lib/Kconfig.debug
index 134b15a44625..3d9c2783ac9b 100644
--- a/lib/Kconfig.debug
+++ b/lib/Kconfig.debug
@@ -3378,6 +3378,19 @@ config TEST_OBJPOOL
 
 	  If unsure, say N.
 
+config OBJPOOL_KUNIT_TEST
+	tristate "KUnit test for nested objpool pushes" if !KUNIT_ALL_TESTS
+	depends on KUNIT
+	default KUNIT_ALL_TESTS
+	help
+	  Stress test for an objpool push that is interrupted by another
+	  push to the same per-CPU slot, as happens when a kretprobe'd
+	  function returns in NMI context.  An hrtimer stands in for the
+	  NMI.  The test runs for about ten seconds.  The case that pops
+	  from another CPU is skipped when only one CPU is online.
+
+	  If unsure, say N.
+
 config TEST_KEXEC_HANDOVER
 	bool "Test for Kexec HandOver"
 	default n
diff --git a/lib/tests/Makefile b/lib/tests/Makefile
index 3cac3b63a752..f0238587cc8a 100644
--- a/lib/tests/Makefile
+++ b/lib/tests/Makefile
@@ -29,6 +29,7 @@ obj-$(CONFIG_IS_SIGNED_TYPE_KUNIT_TEST) += is_signed_type_kunit.o
 obj-$(CONFIG_KPROBES_SANITY_TEST) += test_kprobes.o
 obj-$(CONFIG_LIST_KUNIT_TEST) += list-test.o
 obj-$(CONFIG_LIST_PRIVATE_KUNIT_TEST) += list-private-test.o
+obj-$(CONFIG_OBJPOOL_KUNIT_TEST) += objpool_kunit.o
 obj-$(CONFIG_KFIFO_KUNIT_TEST) += kfifo_kunit.o
 obj-$(CONFIG_TEST_LIST_SORT) += test_list_sort.o
 obj-$(CONFIG_LINEAR_RANGES_TEST) += test_linear_ranges.o
diff --git a/lib/tests/objpool_kunit.c b/lib/tests/objpool_kunit.c
new file mode 100644
index 000000000000..d86ef00dd088
--- /dev/null
+++ b/lib/tests/objpool_kunit.c
@@ -0,0 +1,329 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * KUnit tests for objpool pushes that nest on the same per-CPU slot.
+ *
+ * objpool_push() only pushes to the slot of the local CPU and runs with
+ * interrupts disabled, but it can still be interrupted by an NMI that
+ * pushes to the same slot: kretprobes are allowed in NMI context, and a
+ * probed function that returns in NMI context recycles its instance with
+ * objpool_push().  UML has no NMIs, so these tests push from task context
+ * with interrupts enabled, and a pinned hrtimer pushes to the same slot
+ * from hardirq context in between.  Nesting is the same as with an NMI.
+ */
+
+#include <kunit/test.h>
+#include <linux/atomic.h>
+#include <linux/completion.h>
+#include <linux/cpumask.h>
+#include <linux/hrtimer.h>
+#include <linux/jiffies.h>
+#include <linux/kthread.h>
+#include <linux/objpool.h>
+#include <linux/sched.h>
+#include <linux/slab.h>
+#include <linux/spinlock.h>
+
+#define OP_NR_OBJS	64
+#define OP_RESERVE	4
+#define OP_TIMER_NS	(10 * NSEC_PER_USEC)
+#define OP_RUN_MS	5000
+
+enum { OP_TASK, OP_IRQ, OP_NR_STOCKS };
+
+struct op_obj {
+	atomic_t held;		/* 0 while the object is in the pool */
+	int owner;		/* stock the object goes back to */
+};
+
+struct op_ctx {
+	struct objpool_head pool;
+	struct op_obj *objs[OP_NR_OBJS];
+	struct op_obj *reserve[OP_RESERVE];
+	int nreserve;
+
+	raw_spinlock_t lock;	/* protects the stocks */
+	struct op_obj *stock[OP_NR_STOCKS][OP_NR_OBJS];
+	int nstock[OP_NR_STOCKS];
+
+	struct hrtimer timer;
+	int cpu;
+	bool remote_pop;
+	bool stop;
+	bool popper_done;
+
+	bool in_push;		/* the task is inside __objpool_try_add_slot() */
+
+	unsigned long pushes, irq_pushes, pops;
+	unsigned long nested;	/* timer pushes that interrupted a task push */
+	unsigned long hidden;	/* last != tail after the task's push returned */
+	unsigned long doubles;	/* pop returned an object that is not in the pool */
+	struct completion done;
+};
+
+static struct op_obj *op_take(struct op_ctx *ctx, int stock)
+{
+	struct op_obj *obj = NULL;
+	unsigned long flags;
+
+	raw_spin_lock_irqsave(&ctx->lock, flags);
+	if (ctx->nstock[stock])
+		obj = ctx->stock[stock][--ctx->nstock[stock]];
+	raw_spin_unlock_irqrestore(&ctx->lock, flags);
+	return obj;
+}
+
+static void op_give(struct op_ctx *ctx, struct op_obj *obj)
+{
+	unsigned long flags;
+
+	raw_spin_lock_irqsave(&ctx->lock, flags);
+	ctx->stock[obj->owner][ctx->nstock[obj->owner]++] = obj;
+	raw_spin_unlock_irqrestore(&ctx->lock, flags);
+}
+
+/* an object popped from the pool: it must have been in the pool */
+static void op_popped(struct op_ctx *ctx, struct op_obj *obj)
+{
+	ctx->pops++;
+	if (atomic_xchg(&obj->held, 1)) {
+		/* someone else holds it: don't put it in a stock twice */
+		ctx->doubles++;
+		return;
+	}
+	op_give(ctx, obj);
+}
+
+static enum hrtimer_restart op_timer_fn(struct hrtimer *timer)
+{
+	struct op_ctx *ctx = container_of(timer, struct op_ctx, timer);
+	struct op_obj *obj;
+
+	if (READ_ONCE(ctx->stop))
+		return HRTIMER_NORESTART;
+
+	/* the "NMI": push to the slot of the CPU it interrupted */
+	obj = op_take(ctx, OP_IRQ);
+	if (obj) {
+		if (READ_ONCE(ctx->in_push))
+			ctx->nested++;
+		atomic_set(&obj->held, 0);
+		objpool_push(obj, &ctx->pool);
+		ctx->irq_pushes++;
+	}
+
+	hrtimer_forward_now(timer, ns_to_ktime(OP_TIMER_NS));
+	return HRTIMER_RESTART;
+}
+
+/* push from task context, with interrupts enabled so that the timer can nest */
+static void op_task_push(struct op_ctx *ctx, struct op_obj *obj)
+{
+	struct objpool_slot *slot = ctx->pool.cpu_slots[ctx->cpu];
+	unsigned long flags;
+
+	atomic_set(&obj->held, 0);
+	WRITE_ONCE(ctx->in_push, true);
+	__objpool_try_add_slot(obj, &ctx->pool, ctx->cpu);
+	WRITE_ONCE(ctx->in_push, false);
+	ctx->pushes++;
+
+	/*
+	 * No push to this slot is in flight now: the timer only interrupts
+	 * us, and other CPUs never push to it.  Every entry must be visible.
+	 */
+	local_irq_save(flags);
+	if (READ_ONCE(slot->last) != READ_ONCE(slot->tail))
+		ctx->hidden++;
+	local_irq_restore(flags);
+}
+
+static int op_task_fn(void *data)
+{
+	struct op_ctx *ctx = data;
+	unsigned long end = jiffies + msecs_to_jiffies(OP_RUN_MS);
+	struct op_obj *obj;
+	unsigned long flags;
+
+	hrtimer_start(&ctx->timer, ns_to_ktime(OP_TIMER_NS),
+		      HRTIMER_MODE_REL_PINNED_HARD);
+
+	while (time_before(jiffies, end)) {
+		obj = op_take(ctx, OP_TASK);
+		if (obj)
+			op_task_push(ctx, obj);
+
+		if (!ctx->remote_pop) {
+			/* consume what is visible, like kretprobe entries would */
+			local_irq_save(flags);
+			while ((obj = __objpool_try_get_slot(&ctx->pool, ctx->cpu)))
+				op_popped(ctx, obj);
+			local_irq_restore(flags);
+		}
+		cond_resched();
+	}
+
+	WRITE_ONCE(ctx->stop, true);
+	hrtimer_cancel(&ctx->timer);
+
+	/*
+	 * A remote pop may be spinning because 'last' went backwards; any
+	 * later push to the slot releases it.  Use the reserve for that.
+	 */
+	while (ctx->remote_pop && !READ_ONCE(ctx->popper_done)) {
+		/* keep one reserve object for the teardown */
+		if (ctx->nreserve > 1) {
+			obj = ctx->reserve[--ctx->nreserve];
+			atomic_set(&obj->held, 0);
+			local_irq_save(flags);
+			__objpool_try_add_slot(obj, &ctx->pool, ctx->cpu);
+			local_irq_restore(flags);
+		}
+		schedule_timeout_uninterruptible(1);
+	}
+
+	complete(&ctx->done);
+	return 0;
+}
+
+static int op_popper_fn(void *data)
+{
+	struct op_ctx *ctx = data;
+	struct op_obj *obj;
+
+	while (!READ_ONCE(ctx->stop)) {
+		obj = objpool_pop(&ctx->pool);
+		if (obj)
+			op_popped(ctx, obj);
+		cond_resched();
+	}
+	WRITE_ONCE(ctx->popper_done, true);
+	return 0;
+}
+
+static int op_objinit(void *obj, void *context)
+{
+	atomic_set(&((struct op_obj *)obj)->held, 0);
+	return 0;
+}
+
+static void op_run(struct kunit *test, bool remote_pop)
+{
+	struct task_struct *task, *popper = NULL;
+	struct objpool_slot *slot;
+	struct op_ctx *ctx;
+	unsigned long flags;
+	int i, n = 0, cpu, pcpu, in_pool = 0, lost;
+	struct op_obj *obj;
+
+	if (remote_pop && num_online_cpus() < 2)
+		kunit_skip(test, "needs at least 2 CPUs");
+
+	ctx = kunit_kzalloc(test, sizeof(*ctx), GFP_KERNEL);
+	KUNIT_ASSERT_NOT_NULL(test, ctx);
+	ctx->remote_pop = remote_pop;
+	raw_spin_lock_init(&ctx->lock);
+	init_completion(&ctx->done);
+	hrtimer_setup(&ctx->timer, op_timer_fn, CLOCK_MONOTONIC,
+		      HRTIMER_MODE_REL_PINNED_HARD);
+
+	KUNIT_ASSERT_EQ(test, 0, objpool_init(&ctx->pool, OP_NR_OBJS,
+					      sizeof(struct op_obj), GFP_KERNEL,
+					      NULL, op_objinit, NULL));
+
+	/* take every object out of the pool; hand them to the stocks */
+	while (n < OP_NR_OBJS && (obj = objpool_pop(&ctx->pool)))
+		ctx->objs[n++] = obj;
+	KUNIT_ASSERT_EQ(test, n, OP_NR_OBJS);
+	for (i = 0; i < OP_NR_OBJS; i++) {
+		obj = ctx->objs[i];
+		atomic_set(&obj->held, 1);
+		if (i < OP_RESERVE) {
+			ctx->reserve[ctx->nreserve++] = obj;
+			obj->owner = OP_TASK;
+			continue;
+		}
+		obj->owner = i & 1 ? OP_IRQ : OP_TASK;
+		ctx->stock[obj->owner][ctx->nstock[obj->owner]++] = obj;
+	}
+
+	cpu = cpumask_first(cpu_online_mask);
+	pcpu = cpumask_next(cpu, cpu_online_mask);
+	ctx->cpu = cpu;
+	slot = ctx->pool.cpu_slots[cpu];
+
+	task = kthread_create(op_task_fn, ctx, "objpool_kunit_push");
+	KUNIT_ASSERT_FALSE(test, IS_ERR(task));
+	kthread_bind(task, cpu);
+	if (remote_pop) {
+		popper = kthread_create(op_popper_fn, ctx, "objpool_kunit_pop");
+		KUNIT_ASSERT_FALSE(test, IS_ERR(popper));
+		kthread_bind(popper, pcpu);
+		wake_up_process(popper);
+	}
+	wake_up_process(task);
+	wait_for_completion(&ctx->done);
+
+	/*
+	 * Quiescent now.  Objects that are in the pool but outside
+	 * [head, tail) can never be popped again: they are lost.
+	 */
+	for (i = 0; i < OP_NR_OBJS; i++)
+		if (!atomic_read(&ctx->objs[i]->held))
+			in_pool++;
+	lost = in_pool - (int)(READ_ONCE(slot->tail) - READ_ONCE(slot->head));
+
+	kunit_info(test, "%lu task pushes, %lu timer pushes (%lu nested), %lu pops: %lu hidden, %lu double handouts, %d lost\n",
+		   ctx->pushes, ctx->irq_pushes, ctx->nested, ctx->pops,
+		   ctx->hidden, ctx->doubles, lost);
+
+	/* no nesting means the test proved nothing */
+	KUNIT_EXPECT_GT(test, ctx->nested, 0UL);
+	KUNIT_EXPECT_EQ(test, ctx->hidden, 0UL);
+	KUNIT_EXPECT_EQ(test, ctx->doubles, 0UL);
+	KUNIT_EXPECT_EQ(test, lost, 0);
+
+	/*
+	 * Tear down: republish the slot (a push sets 'last' past 'head'),
+	 * drain it, then drop every object and the pool.
+	 */
+	if (ctx->nreserve) {
+		obj = ctx->reserve[--ctx->nreserve];
+		atomic_set(&obj->held, 0);
+		local_irq_save(flags);
+		__objpool_try_add_slot(obj, &ctx->pool, cpu);
+		local_irq_restore(flags);
+	}
+	if ((int)(READ_ONCE(slot->last) - READ_ONCE(slot->head)) >= 0) {
+		while (objpool_pop(&ctx->pool))
+			;
+	}
+	for (i = 0; i < OP_NR_OBJS; i++)
+		objpool_drop(ctx->objs[i], &ctx->pool);
+	objpool_fini(&ctx->pool);
+}
+
+static void objpool_test_nested_push(struct kunit *test)
+{
+	op_run(test, false);
+}
+
+static void objpool_test_nested_push_remote_pop(struct kunit *test)
+{
+	op_run(test, true);
+}
+
+static struct kunit_case objpool_test_cases[] = {
+	KUNIT_CASE_SLOW(objpool_test_nested_push),
+	KUNIT_CASE_SLOW(objpool_test_nested_push_remote_pop),
+	{}
+};
+
+static struct kunit_suite objpool_test_suite = {
+	.name = "objpool",
+	.test_cases = objpool_test_cases,
+};
+
+kunit_test_suite(objpool_test_suite);
+
+MODULE_DESCRIPTION("KUnit tests for objpool");
+MODULE_LICENSE("GPL");
-- 
2.43.0