rusty: Refactor + slightly improve wake_sync

Right now, the SCX_WAKE_SYNC logic in rusty is very primitive. We only check to
see if the waker CPU's runqueue is empty, and then migrate the wakee there if
so. We'll want to expand this to be more thorough, such as:

- Checking to see if prev_cpu and waker_cpu share the same LLC when determining
  where to migrate
- Check for whether SCX_WAKE_SYNC migration helps load imbalance between cores
- ...

Right now all of that code is just a big blob in the middle of
rusty_select_cpu(). Let's pull it into its own function to improve readability,
and also add some logic to stay on prev_cpu if it shares an LLC with the waker.

Signed-off-by: David Vernet <void@manifault.com>

scheds/rust/scx_rusty/src/bpf/intf.h

@@ -65,6 +65,7 @@ enum consts {
 enum stat_idx {
 	/* The following fields add up to all dispatched tasks */
 	RUSTY_STAT_WAKE_SYNC,
+	RUSTY_STAT_SYNC_PREV_IDLE,
 	RUSTY_STAT_PREV_IDLE,
 	RUSTY_STAT_GREEDY_IDLE,
 	RUSTY_STAT_PINNED,

scheds/rust/scx_rusty/src/bpf/main.bpf.c

@@ -782,6 +782,61 @@ static bool task_set_domain(struct task_ctx *taskc, struct task_struct *p,
 	return taskc->dom_id == new_dom_id;
 }
 
+
+static s32 try_sync_wakeup(struct task_struct *p, struct task_ctx *taskc,
+			   s32 prev_cpu)
+{
+	struct task_struct *current = (void *)bpf_get_current_task_btf();
+	s32 cpu;
+	const struct cpumask *idle_cpumask;
+	bool share_llc, has_idle;
+	struct dom_ctx *domc;
+	struct bpf_cpumask *d_cpumask;
+	struct pcpu_ctx *pcpuc;
+
+	cpu = bpf_get_smp_processor_id();
+	pcpuc = lookup_pcpu_ctx(cpu);
+	if (!pcpuc)
+		return -ENOENT;
+
+	domc = lookup_dom_ctx(pcpuc->dom_id);
+	if (!domc)
+		return -ENOENT;
+
+	d_cpumask = domc->cpumask;
+	if (!d_cpumask) {
+		scx_bpf_error("Failed to acquire dom%u cpumask kptr",
+				taskc->dom_id);
+		return -ENOENT;
+	}
+
+	idle_cpumask = scx_bpf_get_idle_cpumask();
+
+	share_llc = bpf_cpumask_test_cpu(prev_cpu, (const struct cpumask *)d_cpumask);
+	if (share_llc && scx_bpf_test_and_clear_cpu_idle(prev_cpu)) {
+		stat_add(RUSTY_STAT_SYNC_PREV_IDLE, 1);
+
+		cpu = prev_cpu;
+		goto err_out;
+	}
+
+	has_idle = bpf_cpumask_intersects((const struct cpumask *)d_cpumask,
+			idle_cpumask);
+
+	if (has_idle && bpf_cpumask_test_cpu(cpu, p->cpus_ptr) &&
+	    !(current->flags & PF_EXITING) && taskc->dom_id < MAX_DOMS &&
+	    scx_bpf_dsq_nr_queued(SCX_DSQ_LOCAL_ON | cpu) == 0) {
+		stat_add(RUSTY_STAT_WAKE_SYNC, 1);
+		goto err_out;
+	}
+
+	cpu = -ENOENT;
+
+err_out:
+	scx_bpf_put_idle_cpumask(idle_cpumask);
+	return cpu;
+}
+
 s32 BPF_STRUCT_OPS(rusty_select_cpu, struct task_struct *p, s32 prev_cpu,
 		   u64 wake_flags)
 {
@@ -811,41 +866,9 @@ s32 BPF_STRUCT_OPS(rusty_select_cpu, struct task_struct *p, s32 prev_cpu,
 	 * local dsq of the waker.
 	 */
 	if (wake_flags & SCX_WAKE_SYNC) {
-		struct task_struct *current = (void *)bpf_get_current_task_btf();
-
-		cpu = bpf_get_smp_processor_id();
-		if (!(current->flags & PF_EXITING) &&
-		    taskc->dom_id < MAX_DOMS &&
-		    scx_bpf_dsq_nr_queued(SCX_DSQ_LOCAL_ON | cpu) == 0) {
-			struct dom_ctx *domc;
-			struct bpf_cpumask *d_cpumask;
-			const struct cpumask *idle_cpumask;
-			bool has_idle;
-
-			domc = lookup_dom_ctx(taskc->dom_id);
-			if (!domc)
-				goto enoent;
-			d_cpumask = domc->cpumask;
-			if (!d_cpumask) {
-				scx_bpf_error("Failed to acquire dom%u cpumask kptr",
-					      taskc->dom_id);
-				goto enoent;
-			}
-
-			idle_cpumask = scx_bpf_get_idle_cpumask();
-
-			has_idle = bpf_cpumask_intersects((const struct cpumask *)d_cpumask,
-							  idle_cpumask);
-
-			scx_bpf_put_idle_cpumask(idle_cpumask);
-
-			if (has_idle) {
-				if (bpf_cpumask_test_cpu(cpu, p->cpus_ptr)) {
-					stat_add(RUSTY_STAT_WAKE_SYNC, 1);
-					goto direct;
-				}
-			}
-		}
+		cpu = try_sync_wakeup(p, taskc, prev_cpu);
+		if (cpu >= 0)
+			goto direct;
 	}
 
 	has_idle_cores = !bpf_cpumask_empty(idle_smtmask);

scheds/rust/scx_rusty/src/main.rs

@@ -448,6 +448,7 @@ impl<'a> Scheduler<'a> {
     ) {
         let stat = |idx| bpf_stats[idx as usize];
         let total = stat(bpf_intf::stat_idx_RUSTY_STAT_WAKE_SYNC)
+            + stat(bpf_intf::stat_idx_RUSTY_STAT_SYNC_PREV_IDLE)
             + stat(bpf_intf::stat_idx_RUSTY_STAT_PREV_IDLE)
             + stat(bpf_intf::stat_idx_RUSTY_STAT_GREEDY_IDLE)
             + stat(bpf_intf::stat_idx_RUSTY_STAT_PINNED)
@@ -473,9 +474,14 @@ impl<'a> Scheduler<'a> {
         let stat_pct = |idx| stat(idx) as f64 / total as f64 * 100.0;
 
         info!(
-            "tot={:7} wsync={:5.2} prev_idle={:5.2} greedy_idle={:5.2} pin={:5.2}",
+            "tot={:7} wsync_prev_idle={:5.2} wsync={:5.2}",
             total,
+            stat_pct(bpf_intf::stat_idx_RUSTY_STAT_SYNC_PREV_IDLE),
             stat_pct(bpf_intf::stat_idx_RUSTY_STAT_WAKE_SYNC),
+        );
+
+        info!(
+            "prev_idle={:5.2} greedy_idle={:5.2} pin={:5.2}",
             stat_pct(bpf_intf::stat_idx_RUSTY_STAT_PREV_IDLE),
             stat_pct(bpf_intf::stat_idx_RUSTY_STAT_GREEDY_IDLE),
             stat_pct(bpf_intf::stat_idx_RUSTY_STAT_PINNED),