[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

[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 1/2] sched_ext: make DSQ iterator stay drained after NULL

[Tejun Heo]

Hello, Kang.

v1 is already applied to sched_ext/for-7.4-fixes.

The selftest doesn't catch the bug reliably. Every ops.dequeue() invocation
overwrites next_after_drain and the test reads it once at the end. If the
last invocation saw an empty DSQ, the restarted walk returns NULL too and
the test passes on a kernel without the fix. Here, it passed about half of
the runs on an idle VM without the fix.

Also, I don't think the selftest is all that useful.

Thanks.

--
tejun