[PATCH v5 sched_ext/for-7.3 12/33] sched_ext: Stop resolving a disabled scheduler's programs

[Tejun Heo <tj@kernel.org>]

A scheduler's BPF programs can outlive it. A timer it armed or a tracing
program it loaded can fire after ops.exit() has run, before the programs are
unloaded, and scx_prog_sched() still resolves the program to its scheduler
through ops->priv. Harmless while kfuncs touch only lifetime-stable state,
but a hazard once a kfunc reads global state a newly loaded scheduler can
change underneath it.

Add scx_sched->dead, set right after ops.exit() and drained with
synchronize_rcu(). It follows exit() rather than preceding it so exit()'s own
kfunc calls still resolve to @sch. scx_prog_sched() returns NULL for a dead
scheduler, so every kfunc's existing !sch bail rejects it at one choke
point.

v2: Check dead in the CONFIG_EXT_SUB_SCHED=n scx_prog_sched() too. (sashiko AI)

Signed-off-by: Tejun Heo <tj@kernel.org>
---
 kernel/sched/ext/ext.c      |  8 ++++++++
 kernel/sched/ext/internal.h | 19 +++++++++++++++----
 kernel/sched/ext/sub.c      |  9 +++++++++
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index ba83fe832343..fbf1c3477140 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -5733,6 +5733,14 @@ static void scx_root_disable(struct scx_sched *sch)
 	if (sch->ops.exit)
 		SCX_CALL_OP(sch, exit, NULL, sch->exit_info);
 
+	/*
+	 * @sch's non-ops programs such as timers and tracers can fire after
+	 * ops.exit(). Now that exit is complete, stop scx_prog_sched() from
+	 * resolving to @sch and drain in-flight resolvers.
+	 */
+	WRITE_ONCE(sch->dead, true);
+	synchronize_rcu();
+
 	scx_unlink_sched(sch);
 
 	/*
diff --git a/kernel/sched/ext/internal.h b/kernel/sched/ext/internal.h
index 5b18c4192c62..4452aac89b14 100644
--- a/kernel/sched/ext/internal.h
+++ b/kernel/sched/ext/internal.h
@@ -1202,6 +1202,7 @@ struct scx_sched {
 		struct sched_ext_ops_cid	ops_cid;
 	};
 	bool			is_cid_type;	/* true if registered via bpf_sched_ext_ops_cid */
+	bool			dead;		/* set after ops.exit(), gates scx_prog_sched() */
 
 	/*
 	 * Arena map auto-discovered from member progs at struct_ops attach.
@@ -1976,14 +1977,20 @@ static inline bool scx_task_on_sched(struct scx_sched *sch,
 static inline struct scx_sched *scx_prog_sched(const struct bpf_prog_aux *aux)
 {
 	struct sched_ext_ops *ops;
-	struct scx_sched *root;
+	struct scx_sched *sch, *root;
 
 	ops = bpf_prog_get_assoc_struct_ops(aux);
-	if (likely(ops))
-		return rcu_dereference_all(ops->priv);
+	if (likely(ops)) {
+		sch = rcu_dereference_all(ops->priv);
+		if (sch && unlikely(READ_ONCE(sch->dead)))
+			return NULL;
+		return sch;
+	}
 
 	root = rcu_dereference_all(scx_root);
 	if (root) {
+		if (unlikely(READ_ONCE(root->dead)))
+			return NULL;
 		/*
 		 * COMPAT-v6.19: Schedulers built before sub-sched support was
 		 * introduced may have unassociated non-struct_ops programs.
@@ -2035,7 +2042,11 @@ static inline bool scx_task_on_sched(struct scx_sched *sch,
 
 static inline struct scx_sched *scx_prog_sched(const struct bpf_prog_aux *aux)
 {
-	return rcu_dereference_all(scx_root);
+	struct scx_sched *root = rcu_dereference_all(scx_root);
+
+	if (root && unlikely(READ_ONCE(root->dead)))
+		return NULL;
+	return root;
 }
 
 static inline struct scx_sched *scx_parent(struct scx_sched *sch) { return NULL; }
diff --git a/kernel/sched/ext/sub.c b/kernel/sched/ext/sub.c
index 17ed0d28f383..017225ad492f 100644
--- a/kernel/sched/ext/sub.c
+++ b/kernel/sched/ext/sub.c
@@ -344,6 +344,15 @@ void scx_sub_disable(struct scx_sched *sch)
 
 	if (sch->ops.exit)
 		SCX_CALL_OP(sch, exit, NULL, sch->exit_info);
+
+	/*
+	 * @sch's non-ops programs such as timers and tracers can fire after
+	 * ops.exit(). Now that exit is complete, stop scx_prog_sched() from
+	 * resolving to @sch and drain in-flight resolvers.
+	 */
+	WRITE_ONCE(sch->dead, true);
+	synchronize_rcu();
+
 	if (sch->sub_kset)
 		kobject_del(&sch->sub_kset->kobj);
 	/* not added if enable failed before scx_sched_sysfs_add() */
-- 
2.55.0