[PATCH v3 09/10] sched/ext: Relinquish DL server reservations when not needed

[Joel Fernandes <joelagnelf@nvidia.com>]

I tested loading a test SCX program and verifying the bandwidth both
before and after applying the patch:

Without patch:
Before loading scx:
  .dl_bw->total_bw               : 1887408
After unloading scx:
  .dl_bw->total_bw               : 3774816

After patch:
Before loading scx:
  .dl_bw->total_bw               : 1887408
After unloading scx:
  .dl_bw->total_bw               : 1887408

Signed-off-by: Joel Fernandes <joelagnelf@nvidia.com>
---
 kernel/sched/deadline.c | 11 ++---------
 kernel/sched/ext.c      | 44 +++++++++++++++++++++++++++++++++++++----
 2 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/kernel/sched/deadline.c b/kernel/sched/deadline.c
index 0e73274d8c31..924dbbfb4b40 100644
--- a/kernel/sched/deadline.c
+++ b/kernel/sched/deadline.c
@@ -1786,18 +1786,11 @@ int dl_server_remove_params(struct sched_dl_entity *dl_se)
 	dl_b = dl_bw_of(cpu);
 	guard(raw_spinlock)(&dl_b->lock);
 	cpus = dl_bw_cpus(cpu);
-
+	if (dl_se->dl_non_contending)
+		sub_running_bw(dl_se, &rq->dl);
 	sub_rq_bw(dl_se, &rq->dl);
 	__dl_sub(dl_b, dl_se->dl_bw, cpus);
 
-	/*
-	 * If server was active and consuming bandwidth, remove it from
-	 * running bandwidth accounting. This should not happen because
-	 * we call this only after the last dl_server_stop().
-	 */
-	if (WARN_ON_ONCE(!dl_se->dl_non_contending))
-		sub_running_bw(dl_se, &rq->dl);
-
 	/*
 	 * Clear all server parameters. This will also clear ->dl_server so
 	 * the next dl_server_apply_params() will reconfigure the server.
diff --git a/kernel/sched/ext.c b/kernel/sched/ext.c
index 52f98c3944ed..2e77d9971c22 100644
--- a/kernel/sched/ext.c
+++ b/kernel/sched/ext.c
@@ -4784,13 +4784,28 @@ static void scx_ops_disable_workfn(struct kthread_work *work)
 	scx_task_iter_stop(&sti);
 	percpu_up_write(&scx_fork_rwsem);
 
-	/*
-	 * Invalidate all the rq clocks to prevent getting outdated
-	 * rq clocks from a previous scx scheduler.
-	 */
 	for_each_possible_cpu(cpu) {
 		struct rq *rq = cpu_rq(cpu);
+		struct rq_flags rf;
+
+		/*
+		 * Invalidate all the rq clocks to prevent getting outdated
+		 * rq clocks from a previous scx scheduler.
+		 */
 		scx_rq_clock_invalidate(rq);
+
+		/*
+		 * We are unloading the sched_ext scheduler, we do not need its
+		 * DL server bandwidth anymore, remove it for all CPUs. Whenever
+		 * the first SCX task is enqueued (when scx is re-loaded), its DL
+		 * server bandwidth will be re-initialized.
+		 */
+		rq_lock_irqsave(rq, &rf);
+		if (dl_server_active(&rq->ext_server)) {
+			dl_server_stop(&rq->ext_server);
+		}
+		dl_server_remove_params(&rq->ext_server);
+		rq_unlock_irqrestore(rq, &rf);
 	}
 
 	/* no task is on scx, turn off all the switches and flush in-progress calls */
@@ -5547,6 +5562,27 @@ static int scx_ops_enable(struct sched_ext_ops *ops, struct bpf_link *link)
 		check_class_changed(task_rq(p), p, old_class, p->prio);
 	}
 	scx_task_iter_stop(&sti);
+
+	if (scx_switching_all) {
+		for_each_possible_cpu(cpu) {
+			struct rq *rq = cpu_rq(cpu);
+			struct rq_flags rf;
+
+			/*
+			 * We are switching all fair tasks to the sched_ext scheduler,
+			 * we do not need fair server's DL bandwidth anymore, remove it
+			 * for all CPUs. Whenever the first CFS task is enqueued (when
+			 * scx is unloaded), the fair server's DL bandwidth will be
+			 * re-initialized.
+			 */
+			rq_lock_irqsave(rq, &rf);
+			if (dl_server_active(&rq->fair_server))
+				dl_server_stop(&rq->fair_server);
+			dl_server_remove_params(&rq->fair_server);
+			rq_unlock_irqrestore(rq, &rf);
+		}
+	}
+
 	percpu_up_write(&scx_fork_rwsem);
 
 	scx_ops_bypass(false);
-- 
2.34.1