[PATCH v4 sched_ext/for-7.3 20/40] 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.

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

diff --git a/kernel/sched/ext/ext.c b/kernel/sched/ext/ext.c
index 46b80b31925e..b877f2459ff6 100644
--- a/kernel/sched/ext/ext.c
+++ b/kernel/sched/ext/ext.c
@@ -5696,6 +5696,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 programs (an armed timer, a tracing prog) can fire after
+	 * exit(). Now that exit()'s own kfuncs have resolved, 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 dce07b83c166..87582bca5cc4 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.
@@ -1968,14 +1969,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.
diff --git a/kernel/sched/ext/sub.c b/kernel/sched/ext/sub.c
index 17ed0d28f383..8d9825093649 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 programs (an armed timer, a tracing prog) can fire after
+	 * exit(). Now that exit()'s own kfuncs have resolved, 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.54.0