[PATCH 5/6] tools/sched_ext: scx_flatcg: Don't strand a cgroup on the ENQ_RACE path

[Tao Cui <cui.tao@linux.dev>]

From: Tao Cui <cuitao@kylinos.cn>

cgrp_enqueued() bails out on the ENQ_RACE path when bpf_kptr_xchg()
returns NULL: the node is either back on the rbtree or in the hands
of try_pick_next_cgroup(), which found this cgroup empty and is
deciding whether to stash it back or re-add it to the rbtree.

The bailout relies on the assumption that a task being enqueued is
already visible to try_pick_next_cgroup()'s emptiness test. That
doesn't hold: inserts from ops.enqueue() are deferred till after
ops.enqueue() returns, so try_pick_next_cgroup() can see an empty DSQ,
clear ->queued and stash the node while ->queued stays at 1 from this
enqueue, leaving the cgroup off the rbtree with every subsequent
enqueue skipped; the whole group strands until it is destroyed or the
scheduler is reloaded, and the kernel watchdog eventually kills the
scheduler with SCX_EXIT_ERROR_STALL.

Instead of bailing out, wait for the picker to resolve and retry:
grabbing the node when it's stashed back, or bailing out when the
picker has re-added the cgroup to the rbtree. Exhausted retries mean
the cgroup is in the common cases either going away with its node
dropped or already back on the rbtree; a narrow window with the picker
still holding the node in flight can strand the cgroup the same way
the unpatched code does, at a rate orders of magnitude below stock.

Verified in qemu: driving ops.enqueue() races under CPU saturation strands
cgroups into scheduler-killing stalls 6/10 rounds on stock, 0/10 with the
change.

Fixes: a4103eacc2ab ("sched_ext: Add a cgroup scheduler which uses flattened hierarchy")
Signed-off-by: Tao Cui <cuitao@kylinos.cn>
---
 tools/sched_ext/scx_flatcg.bpf.c | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/tools/sched_ext/scx_flatcg.bpf.c b/tools/sched_ext/scx_flatcg.bpf.c
index 1d168d492203..5f2b9d0b867a 100644
--- a/tools/sched_ext/scx_flatcg.bpf.c
+++ b/tools/sched_ext/scx_flatcg.bpf.c
@@ -285,11 +285,24 @@ static void cgrp_enqueued(struct cgroup *cgrp, struct fcg_cgrp_ctx *cgc)
 		return;
 	}
 
-	/* NULL if the node is already on the rbtree */
+	/* NULL if the node is on the rbtree or in the picker's hands */
 	cgv_node = bpf_kptr_xchg(&stash->node, NULL);
 	if (!cgv_node) {
+		/*
+		 * The node is either back on the rbtree or in the
+		 * picker's hands; bailing out strands the cgroup. Wait
+		 * for the picker to resolve and retry.
+		 */
 		stat_inc(FCG_STAT_ENQ_RACE);
-		return;
+		bpf_repeat(CGROUP_MAX_RETRIES) {
+			cgv_node = bpf_kptr_xchg(&stash->node, NULL);
+			if (cgv_node)
+				break;
+			if (scx_bpf_dsq_nr_queued(cgid) > 0)
+				break;
+		}
+		if (!cgv_node)
+			return;
 	}
 
 	bpf_spin_lock(&cgv_tree_lock);
-- 
2.53.0

base-commit: 5b740a72b55c4267ba4d8010f17dde26b26f52d3