From 1f834e7f9420e12af023a5a7ebdd4f58654fde45 Mon Sep 17 00:00:00 2001
From: David Vernet <void@manifault.com>
Date: Wed, 21 Feb 2024 14:57:10 -0600
Subject: [PATCH] layered: Initialize layers before attaching scheduler

We currently have a bug in layered wherein we could fail to propagate
layer updates from user space to kernel space if a layer is never
adjusted after it's first initialized. For example, in the following
configuration:

[
	{
		"name": "workload.slice",
		"comment": "main workload slice",
		"matches": [
			[
				{
					"CgroupPrefix": "workload.slice/"
				}
			]
		],
		"kind": {
			"Grouped": {
				"cpus_range": [30, 30],
				"util_range": [
					0.0,
					1.0
				],
				"preempt": false
			}
		}
	},
	{
		"name": "normal",
		"comment": "the rest",
		"matches": [
			[]
		],
		"kind": {
			"Grouped": {
				"cpus_range": [2, 2],
				"util_range": [
					0.0,
					1.0
				],
				"preempt": false
			}
		}
	}
]

Both layers are static, and need only be resized a single time, so the
configuration would never be propagated to the kernel due to us never
calling update_bpf_layer_cpumask(). Let's instead have the
initialization propagate changes to the skeleton before we attach the
scheduler.

This has the advantage both of fixing the bug mentioned above where a
static configuration is never propagated to the kernel, and that we
don't have a short period when the scheduler is first attached where we
don't make optimal scheduling decisions due to the layer resizing not
being propagated.

Signed-off-by: David Vernet <void@manifault.com>
---
 scheds/rust/scx_layered/src/main.rs | 74 ++++++++++++-----------------
 1 file changed, 31 insertions(+), 43 deletions(-)

diff --git a/scheds/rust/scx_layered/src/main.rs b/scheds/rust/scx_layered/src/main.rs
index 1612724..580bc47 100644
--- a/scheds/rust/scx_layered/src/main.rs
+++ b/scheds/rust/scx_layered/src/main.rs
@@ -841,37 +841,13 @@ impl Layer {
 
         let nr_cpus = cpu_pool.nr_cpus;
 
-        let mut layer = Self {
+        Ok(Self {
             name: name.into(),
             kind,
 
             nr_cpus: 0,
             cpus: bitvec![0; nr_cpus],
-        };
-
-        match &layer.kind {
-            LayerKind::Confined {
-                cpus_range,
-                util_range,
-            }
-            | LayerKind::Grouped {
-                cpus_range,
-                util_range,
-                ..
-            } => {
-                layer.resize_confined_or_grouped(
-                    cpu_pool,
-                    *cpus_range,
-                    *util_range,
-                    (0.0, 0.0),
-                    (0.0, 0.0),
-                    false,
-                )?;
-            }
-            _ => {}
-        }
-
-        Ok(layer)
+        })
     }
 
     fn grow_confined_or_grouped(
@@ -1239,17 +1215,7 @@ impl<'a> Scheduler<'a> {
         }
         Self::init_layers(&mut skel, &layer_specs)?;
 
-        // Attach.
         let mut skel = skel.load().context("Failed to load BPF program")?;
-        skel.attach().context("Failed to attach BPF program")?;
-        let struct_ops = Some(
-            skel.maps_mut()
-                .layered()
-                .attach_struct_ops()
-                .context("Failed to attach layered struct ops")?,
-        );
-        info!("Layered Scheduler Attached");
-
         let mut layers = vec![];
         for spec in layer_specs.iter() {
             layers.push(Layer::new(&mut cpu_pool, &spec.name, spec.kind.clone())?);
@@ -1258,8 +1224,8 @@ impl<'a> Scheduler<'a> {
         // Other stuff.
         let proc_reader = procfs::ProcReader::new();
 
-        Ok(Self {
-            struct_ops, // should be held to keep it attached
+        let mut sched = Self {
+            struct_ops: None,
             layer_specs,
 
             sched_intv: Duration::from_secs_f64(opts.interval),
@@ -1281,7 +1247,22 @@ impl<'a> Scheduler<'a> {
 
             om_stats: OpenMetricsStats::new(),
             om_format: opts.open_metrics_format,
-        })
+        };
+
+        // Initialize layers before we attach the scheduler
+        sched.refresh_cpumasks()?;
+
+        // Attach.
+        sched.skel.attach().context("Failed to attach BPF program")?;
+        sched.struct_ops = Some(
+            sched.skel.maps_mut()
+                .layered()
+                .attach_struct_ops()
+                .context("Failed to attach layered struct ops")?,
+        );
+        info!("Layered Scheduler Attached");
+
+        Ok(sched)
     }
 
     fn update_bpf_layer_cpumask(layer: &Layer, bpf_layer: &mut bpf_bss_types::layer) {
@@ -1295,10 +1276,7 @@ impl<'a> Scheduler<'a> {
         bpf_layer.refresh_cpus = 1;
     }
 
-    fn step(&mut self) -> Result<()> {
-        let started_at = Instant::now();
-        self.sched_stats
-            .refresh(&mut self.skel, &self.proc_reader, started_at)?;
+    fn refresh_cpumasks(&mut self) -> Result<()> {
         let mut updated = false;
 
         for idx in 0..self.layers.len() {
@@ -1366,6 +1344,16 @@ impl<'a> Scheduler<'a> {
             }
         }
 
+        Ok(())
+    }
+
+    fn step(&mut self) -> Result<()> {
+        let started_at = Instant::now();
+        self.sched_stats
+            .refresh(&mut self.skel, &self.proc_reader, started_at)?;
+
+        self.refresh_cpumasks()?;
+
         self.processing_dur += Instant::now().duration_since(started_at);
         Ok(())
     }