[PATCH] sched_ext: make DSQ iterator stay drained after NULL

[PATCH] sched_ext: make DSQ iterator stay drained after NULL

[Kang Chen]

nldsq_cursor_next_task() unlinks the iteration cursor when it reaches the
end of the DSQ.  A later call sees list_empty(&cursor->node), starts
from the DSQ head again, and returns the first task of the same
iteration.

The verifier relies on the iterator contract that once next() returns
NULL it keeps returning NULL: process_iter_next_call() marks the
iterator DRAINED and treats the result of later next() calls as the
constant zero.  A restarted DSQ iteration therefore hands the BPF
program a task pointer while the verifier believes R0 == 0.

Remember the drained state in the cursor flags and return NULL on every
later next() call.  The cursor has spare flag bits and the BPF-visible
iterator layout is unchanged.

Fixes: 650ba21b131e ("sched_ext: Implement DSQ iterator")
Signed-off-by: Kang Chen <k.chen@smail.nju.edu.cn>
---
 kernel/sched/ext/ext.c | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 9bae826d69f9..03e0fbf17011 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -555,12 +555,14 @@ enum scx_dsq_iter_flags {
 	/* iterate in the reverse dispatch order */
 	SCX_DSQ_ITER_REV		= 1U << 16,
 
+	__SCX_DSQ_ITER_DRAINED		= 1U << 28,
 	__SCX_DSQ_ITER_HAS_REENQ	= 1U << 29,
 	__SCX_DSQ_ITER_HAS_SLICE	= 1U << 30,
 	__SCX_DSQ_ITER_HAS_VTIME	= 1U << 31,
 
 	__SCX_DSQ_ITER_USER_FLAGS	= SCX_DSQ_ITER_REV,
 	__SCX_DSQ_ITER_ALL_FLAGS	= __SCX_DSQ_ITER_USER_FLAGS |
+					  __SCX_DSQ_ITER_DRAINED |
 					  __SCX_DSQ_ITER_HAS_REENQ |
 					  __SCX_DSQ_ITER_HAS_SLICE |
 					  __SCX_DSQ_ITER_HAS_VTIME,
@@ -10615,13 +10617,24 @@ __bpf_kfunc int bpf_iter_scx_dsq_new(struct bpf_iter_scx_dsq *it, u64 dsq_id,
 __bpf_kfunc struct task_struct *bpf_iter_scx_dsq_next(struct bpf_iter_scx_dsq *it)
 {
 	struct bpf_iter_scx_dsq_kern *kit = (void *)it;
+	struct task_struct *p;
 
-	if (!kit->dsq)
+	if (!kit->dsq || (kit->cursor.flags & __SCX_DSQ_ITER_DRAINED))
 		return NULL;
 
 	guard(raw_spinlock_irqsave)(&kit->dsq->lock);
 
-	return nldsq_cursor_next_task(&kit->cursor, kit->dsq);
+	p = nldsq_cursor_next_task(&kit->cursor, kit->dsq);
+	/*
+	 * nldsq_cursor_next_task() unlinks the cursor and starts from the DSQ
+	 * head when it is called again.  The verifier relies on an iterator
+	 * which returned NULL to keep returning NULL (see
+	 * process_iter_next_call()), so remember the drained state.
+	 */
+	if (!p)
+		kit->cursor.flags |= __SCX_DSQ_ITER_DRAINED;
+
+	return p;
 }
 
 /**
-- 
2.34.1

Re: [PATCH] sched_ext: make DSQ iterator stay drained after NULL

[sashiko-bot]

Thank you for your contribution! Sashiko AI review found no new issues
in this patch, but noticed 1 related pre-existing issue it does not
introduce:
- [High] sched_ext: recursive dsq->lock deadlock via
  bpf_iter_scx_dsq_next() in tracing [1]
--

commit 892356944fb79fb9f40ab2f24779b6976bd4ec25
Author: Kang Chen <k.chen@smail.nju.edu.cn>

sched_ext: make DSQ iterator stay drained after NULL

This commit records a drained flag in the DSQ iterator cursor once
nldsq_cursor_next_task() returns NULL so subsequent calls to
bpf_iter_scx_dsq_next() continue returning NULL.

> diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
> index ebdf7b91f9d37..f2bfd43d75526 100644
> --- a/kernel/sched/ext/ext.c
> +++ b/kernel/sched/ext/ext.c

[ ... ]

> @@ -11028,13 +11030,24 @@ __bpf_kfunc int bpf_iter_scx_dsq_new(struct bpf_iter_scx_dsq *it, u64 dsq_id,
>  __bpf_kfunc struct task_struct *bpf_iter_scx_dsq_next(struct bpf_iter_scx_dsq *it)
>  {
>  	struct bpf_iter_scx_dsq_kern *kit = (void *)it;
> +	struct task_struct *p;
>  
> -	if (!kit->dsq)
> +	if (!kit->dsq || (kit->cursor.flags & __SCX_DSQ_ITER_DRAINED))
>  		return NULL;
>  
>  	guard(raw_spinlock_irqsave)(&kit->dsq->lock);

[Severity: High]
This problem wasn't introduced by this patch, but can bpf_iter_scx_dsq_next()
and bpf_iter_scx_dsq_destroy() deadlock [1] on kit->dsq->lock when
called from a BPF tracing program?

The DSQ iterator kfuncs are registered in scx_kfunc_ids_any and are not
listed in scx_kfunc_ids_no_tracing, so scx_kfunc_context_filter() allows
BPF_PROG_TYPE_TRACING programs to call them.  If a tracing program runs
while the current CPU already holds the same user DSQ's lock, iterating that
DSQ attempts to acquire kit->dsq->lock recursively on the same CPU:

scx_dispatch_enqueue()
  raw_spin_lock(&dsq->lock)
  ...
  scx_error()
    trace_sched_ext_exit()
      <BPF tracing prog>
        bpf_iter_scx_dsq_next()
          raw_spin_lock_irqsave(&kit->dsq->lock)  <-- spins forever

Adding the DSQ iterator kfuncs to scx_kfunc_ids_no_tracing would prevent
tracing programs from calling them.

[1] https://sashiko.dev/#/bug/linux-243627bf-8f5e-401e-a2af-0b6c5d9b3088

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20261110125113.2935742-1-k.chen@smail.nju.edu.cn?part=1

Re: [PATCH] sched_ext: make DSQ iterator stay drained after NULL

[Tejun Heo]

Applied to sched_ext/for-7.4-fixes.

Thanks.

--
tejun

[PATCH sched_ext v2 1/2] sched_ext: make DSQ iterator stay drained after NULL

[Kang Chen]

nldsq_cursor_next_task() unlinks the cursor and returns NULL when the
DSQ walk is finished.  A later bpf_iter_scx_dsq_next() call sees an empty
cursor and nldsq_next_task(dsq, NULL, rev) starts again from the DSQ
head; the old tasks are still queued, so the iterator returns them
again.

The verifier relies on the iterator contract that once next() returns
NULL it keeps returning NULL: process_iter_next_call() marks the
iterator DRAINED and treats later next() results as the constant zero.
A restarted DSQ iteration therefore hands the BPF program a kernel
pointer while the verifier believes R0 == 0.

Remember the drained state in the cursor flags and return NULL on every
later next() call.  The cursor has spare flag bits and the BPF-visible
iterator layout is unchanged.

Fixes: 650ba21b131e ("sched_ext: Implement DSQ iterator")
Signed-off-by: Kang Chen <k.chen@smail.nju.edu.cn>
---
v2:
- Shorten the code comment.

 kernel/sched/ext/ext.c | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index f74ffb3b1aca..8a7a33e203c2 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -580,12 +580,14 @@ enum scx_dsq_iter_flags {
 	/* iterate in the reverse dispatch order */
 	SCX_DSQ_ITER_REV		= 1U << 16,
 
+	__SCX_DSQ_ITER_DRAINED		= 1U << 28,
 	__SCX_DSQ_ITER_HAS_REENQ	= 1U << 29,
 	__SCX_DSQ_ITER_HAS_SLICE	= 1U << 30,
 	__SCX_DSQ_ITER_HAS_VTIME	= 1U << 31,
 
 	__SCX_DSQ_ITER_USER_FLAGS	= SCX_DSQ_ITER_REV,
 	__SCX_DSQ_ITER_ALL_FLAGS	= __SCX_DSQ_ITER_USER_FLAGS |
+					  __SCX_DSQ_ITER_DRAINED |
 					  __SCX_DSQ_ITER_HAS_REENQ |
 					  __SCX_DSQ_ITER_HAS_SLICE |
 					  __SCX_DSQ_ITER_HAS_VTIME,
@@ -11028,13 +11030,18 @@ __bpf_kfunc int bpf_iter_scx_dsq_new(struct bpf_iter_scx_dsq *it, u64 dsq_id,
 __bpf_kfunc struct task_struct *bpf_iter_scx_dsq_next(struct bpf_iter_scx_dsq *it)
 {
 	struct bpf_iter_scx_dsq_kern *kit = (void *)it;
+	struct task_struct *p;
 
-	if (!kit->dsq)
+	if (!kit->dsq || (kit->cursor.flags & __SCX_DSQ_ITER_DRAINED))
 		return NULL;
 
 	guard(raw_spinlock_irqsave)(&kit->dsq->lock);
 
-	return nldsq_cursor_next_task(&kit->cursor, kit->dsq);
+	p = nldsq_cursor_next_task(&kit->cursor, kit->dsq);
+	if (!p)
+		kit->cursor.flags |= __SCX_DSQ_ITER_DRAINED;
+
+	return p;
 }
 
 /**
-- 
2.34.1

[PATCH sched_ext v2 2/2] selftests/sched_ext: check that a drained DSQ iterator stays NULL

[Kang Chen]

Extend the dequeue_iter selftest: after the DSQ iterator drains, call
bpf_iter_scx_dsq_next() once more and store the raw result.  The
user-space test checks that at least one task was iterated and that the
extra next() returned NULL.

The test fails on a kernel without the drained cursor state: the cursor
was unlinked and the extra next() starts again from the DSQ head.

Signed-off-by: Kang Chen <k.chen@smail.nju.edu.cn>
---
v2:
- New selftest in v2.

 tools/testing/selftests/sched_ext/dequeue_iter.bpf.c | 7 ++++++-
 tools/testing/selftests/sched_ext/dequeue_iter.c     | 4 ++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/tools/testing/selftests/sched_ext/dequeue_iter.bpf.c b/tools/testing/selftests/sched_ext/dequeue_iter.bpf.c
index 76c2c71a90b6..38eb2df54a6b 100644
--- a/tools/testing/selftests/sched_ext/dequeue_iter.bpf.c
+++ b/tools/testing/selftests/sched_ext/dequeue_iter.bpf.c
@@ -18,6 +18,8 @@ UEI_DEFINE(uei);
 #define TEST_DSQ_ID 1000
 
 u64 dq_count;
+u64 iter_count;
+u64 next_after_drain = -1;
 
 s32 BPF_STRUCT_OPS_SLEEPABLE(dequeue_iter_init)
 {
@@ -47,7 +49,10 @@ void BPF_STRUCT_OPS(dequeue_iter_dequeue, struct task_struct *p, u64 deq_flags)
 
 	if (!bpf_iter_scx_dsq_new(&it, TEST_DSQ_ID, 0)) {
 		while ((t = bpf_iter_scx_dsq_next(&it)))
-			;
+			iter_count++;
+
+		/* NULL must stay final after the iterator is drained. */
+		next_after_drain = (__u64)bpf_iter_scx_dsq_next(&it);
 	}
 	bpf_iter_scx_dsq_destroy(&it);
 
diff --git a/tools/testing/selftests/sched_ext/dequeue_iter.c b/tools/testing/selftests/sched_ext/dequeue_iter.c
index f60711bf3b39..48e8c94b7bf1 100644
--- a/tools/testing/selftests/sched_ext/dequeue_iter.c
+++ b/tools/testing/selftests/sched_ext/dequeue_iter.c
@@ -59,6 +59,10 @@ static enum scx_test_status run(void *ctx)
 		return SCX_TEST_FAIL;
 	}
 
+	SCX_FAIL_IF(!skel->bss->iter_count,
+		    "DSQ iterator drained no tasks");
+	SCX_EQ(skel->data->next_after_drain, 0);
+
 	return SCX_TEST_PASS;
 }
 
-- 
2.34.1

Re: [PATCH sched_ext v2 2/2] selftests/sched_ext: check that a drained DSQ iterator stays NULL

[sashiko-bot]

Thank you for your contribution! Sashiko AI review found 2 potential
issues to consider:
- [Medium] selftests/sched_ext: unconditional overwrite of
  next_after_drain in dequeue_iter_dequeue()
- [Medium] selftests/sched_ext: empty TEST_DSQ_ID in
  dequeue_iter_dequeue() leaves iter_count at 0
--

commit 2125e4ff3cb11bd1e11aff810c77c47f61782533
Author: Kang Chen <k.chen@smail.nju.edu.cn>

selftests/sched_ext: check that a drained DSQ iterator stays NULL

This commit extends the dequeue_iter selftest to verify that calling
bpf_iter_scx_dsq_next() again after a DSQ iterator is drained continues
to return NULL.

> diff --git a/tools/testing/selftests/sched_ext/dequeue_iter.bpf.c b/tools/testing/selftests/sched_ext/dequeue_iter.bpf.c
> index 76c2c71a90b61..38eb2df54a6b5 100644
> --- a/tools/testing/selftests/sched_ext/dequeue_iter.bpf.c
> +++ b/tools/testing/selftests/sched_ext/dequeue_iter.bpf.c
> @@ -18,6 +18,8 @@ UEI_DEFINE(uei);
>  #define TEST_DSQ_ID 1000
>  
>  u64 dq_count;
> +u64 iter_count;
> +u64 next_after_drain = -1;
>  
>  s32 BPF_STRUCT_OPS_SLEEPABLE(dequeue_iter_init)
>  {
> @@ -47,7 +49,10 @@ void BPF_STRUCT_OPS(dequeue_iter_dequeue, struct task_struct *p, u64 deq_flags)
>  
>  	if (!bpf_iter_scx_dsq_new(&it, TEST_DSQ_ID, 0)) {
>  		while ((t = bpf_iter_scx_dsq_next(&it)))
> -			;
> +			iter_count++;
> +
> +		/* NULL must stay final after the iterator is drained. */
> +		next_after_drain = (__u64)bpf_iter_scx_dsq_next(&it);
>  	}
>  	bpf_iter_scx_dsq_destroy(&it);

[Severity: Medium]
Can this assignment overwrite a non-NULL next_after_drain captured by an
earlier dequeue and cause a false pass on an unfixed kernel?

scx_consume_dispatch_q() unlinks the consumed task from TEST_DSQ_ID
before calling ops.dequeue().  If an earlier dequeue saw another task in
TEST_DSQ_ID and recorded a non-NULL task pointer in next_after_drain, any
subsequent dequeue that runs when TEST_DSQ_ID is empty will unconditionally
overwrite next_after_drain with 0.  In addition, iter_count and
next_after_drain are updated with plain stores even though ops.dequeue()
can run concurrently on multiple CPUs after dsq->lock is dropped.

> diff --git a/tools/testing/selftests/sched_ext/dequeue_iter.c b/tools/testing/selftests/sched_ext/dequeue_iter.c
> index f60711bf3b395..48e8c94b7bf14 100644
> --- a/tools/testing/selftests/sched_ext/dequeue_iter.c
> +++ b/tools/testing/selftests/sched_ext/dequeue_iter.c
> @@ -59,6 +59,10 @@ static enum scx_test_status run(void *ctx)
>  		return SCX_TEST_FAIL;
>  	}
>  
> +	SCX_FAIL_IF(!skel->bss->iter_count,
> +		    "DSQ iterator drained no tasks");

[Severity: Medium]
Will skel->bss->iter_count stay at 0 on an idle or single-CPU system?

Because scx_consume_dispatch_q() unlinks the task being consumed from
TEST_DSQ_ID and clears SCX_TASK_IN_CUSTODY before invoking
dequeue_iter_dequeue(), bpf_iter_scx_dsq_next() only yields a task if a
second task happens to be queued in TEST_DSQ_ID at the same instant.
run() only drives a single thread calling usleep(100) until dq_count
reaches DQ_TARGET, so if no other task is concurrently queued in
TEST_DSQ_ID during those 10 dequeues, the test fails with "DSQ iterator
drained no tasks".

-- 
Sashiko AI review · https://sashiko.dev/#/patchset/20261113024505.922854-1-k.chen@smail.nju.edu.cn?part=2