scx_chaos: Add backwards sticky scheduling

Signed-off-by: Daniel Hodges <[email protected]>

Author: hodgesds

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

@@ -7,6 +7,9 @@
 
 #ifndef __KERNEL__
 typedef unsigned long long u64;
+#ifndef __kptr
+#define __kptr
+#endif
 #endif
 
 enum chaos_consts {
@@ -30,6 +33,7 @@ enum chaos_trait_kind {
 	CHAOS_TRAIT_RANDOM_DELAYS,
 	CHAOS_TRAIT_CPU_FREQ,
 	CHAOS_TRAIT_DEGRADATION,
+	CHAOS_TRAIT_BACKWARDS_STICKY,
 	CHAOS_TRAIT_MAX,
 };
 
@@ -40,6 +44,8 @@ struct chaos_task_ctx {
 	enum chaos_trait_kind	next_trait;
 	u64			enq_flags;
 	u64			p2dq_vtime;
+
+	struct bpf_cpumask __kptr	*bs_cpumask;
 };
 
 #endif /* __CHAOS_INTF_H */

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

@@ -50,6 +50,8 @@ const volatile u32 cpu_freq_max = SCX_CPUPERF_ONE;
 const volatile u32 degradation_freq_frac32 = 1;
 const volatile u64 degradation_frac7 = 0;
 
+const volatile u32 backwards_sticky_freq_frac32 = 1;
+
 #define MIN(x, y) ((x) < (y) ? (x) : (y))
 #define MAX(x, y) ((x) > (y) ? (x) : (y))
 
@@ -77,11 +79,21 @@ struct chaos_task_ctx *lookup_create_chaos_task_ctx(struct task_struct *p)
 	return bpf_task_storage_get(&chaos_task_ctxs, p, NULL, BPF_LOCAL_STORAGE_GET_F_CREATE);
 }
 
+static __always_inline bool keep_backwards_sticky(struct chaos_task_ctx *taskc)
+{
+	return taskc->next_trait == CHAOS_TRAIT_BACKWARDS_STICKY &&
+		taskc->bs_cpumask &&
+		bpf_cpumask_weight(cast_mask(taskc->bs_cpumask)) != nr_cpus;
+}
+
 static __always_inline enum chaos_trait_kind choose_chaos(struct chaos_task_ctx *taskc)
 {
 	if (taskc->match & CHAOS_MATCH_EXCLUDED)
 		return CHAOS_TRAIT_NONE;
 
+	if (keep_backwards_sticky(taskc))
+		return -1;
+
 	u32 roll = bpf_get_prandom_u32();
 
 	#pragma unroll
@@ -113,6 +125,23 @@ static __always_inline u64 get_cpu_delay_dsq(int cpu_idx)
 	return CHAOS_DSQ_BASE | cpu_idx;
 }
 
+static __always_inline s32 backwards_sticky_cpu(struct chaos_task_ctx *taskc)
+{
+	s32 cpu;
+	if (!taskc->bs_cpumask || bpf_cpumask_weight(cast_mask(taskc->bs_cpumask)) == nr_cpus)
+		return -ENOENT;
+
+	if (!taskc->bs_cpumask)
+		return -ENOENT;
+
+	cpu = bpf_cpumask_first_zero(cast_mask(taskc->bs_cpumask));
+	if (!taskc->bs_cpumask)
+		return -ENOENT;
+	bpf_cpumask_set_cpu(cpu, taskc->bs_cpumask);
+
+	return cpu;
+}
+
 static __always_inline s32 calculate_chaos_match(struct task_struct *p)
 {
 	struct chaos_task_ctx *taskc;
@@ -259,6 +288,15 @@ __weak s32 enqueue_chaotic(struct task_struct *p __arg_trusted, u64 enq_flags,
 		out = enqueue_random_delay(p, enq_flags, taskc);
 		break;
 
+	case CHAOS_TRAIT_BACKWARDS_STICKY:
+		out = false;
+		if (taskc->bs_cpumask &&
+		    bpf_cpumask_weight(cast_mask(taskc->bs_cpumask)) == nr_cpus) {
+			if (taskc->bs_cpumask)
+				bpf_cpumask_clear(taskc->bs_cpumask);
+			taskc->next_trait = CHAOS_TRAIT_NONE;
+		}
+		return out;
 	case CHAOS_TRAIT_NONE:
 	case CHAOS_TRAIT_CPU_FREQ:
 	case CHAOS_TRAIT_DEGRADATION:
@@ -474,6 +512,27 @@ void BPF_STRUCT_OPS(chaos_enqueue, struct task_struct *p __arg_trusted, u64 enq_
 		}
 	}
 
+	// Backwards sticky means the task will cycle through all CPUs before
+	// becoming sticky again. Do a direct dispatch to ensure the task lands
+	// on the correct CPU and trashes the caches.
+	if (taskc->next_trait == CHAOS_TRAIT_BACKWARDS_STICKY) {
+		s32 cpu = backwards_sticky_cpu(taskc);
+		if (cpu >= 0) {
+			u64 slice_ns;
+			if (promise.kind == P2DQ_ENQUEUE_PROMISE_VTIME)
+				slice_ns = promise.vtime.slice_ns;
+			if (promise.kind == P2DQ_ENQUEUE_PROMISE_FIFO)
+				slice_ns = promise.fifo.slice_ns;
+
+			scx_bpf_dsq_insert(p,
+					   SCX_DSQ_LOCAL_ON | cpu,
+					   slice_ns,
+					   enq_flags);
+			promise.kind = P2DQ_ENQUEUE_PROMISE_COMPLETE;
+			return;
+		}
+	}
+
 	complete_p2dq_enqueue(&promise, p);
 }
 
@@ -532,6 +591,26 @@ s32 BPF_STRUCT_OPS(chaos_select_cpu, struct task_struct *p, s32 prev_cpu, u64 wa
 s32 BPF_STRUCT_OPS_SLEEPABLE(chaos_init_task, struct task_struct *p,
 			     struct scx_init_task_args *args)
 {
+	struct bpf_cpumask *bs_cpumask;
+	struct chaos_task_ctx *taskc;
+
+	if (!(taskc = lookup_create_chaos_task_ctx(p))) {
+		scx_bpf_error("failed to lookup task context in init_task");
+		return -ENOMEM;
+	}
+
+	if (!(bs_cpumask = bpf_cpumask_create())) {
+		scx_bpf_error("chaos_task_ctx bs_cpumask allocation failure");
+		return -ENOMEM;
+	}
+
+	if ((bs_cpumask = bpf_kptr_xchg(taskc->bs_cpumask, bs_cpumask))) {
+		bpf_cpumask_release(bs_cpumask);
+		scx_bpf_error("chaos_task_ctx bs_cpumask allocation failure");
+		return -EINVAL;
+	}
+	bpf_cpumask_clear(taskc->bs_cpumask);
+
 	s32 ret = p2dq_init_task_impl(p, args);
 	if (ret)
 		return ret;

scheds/rust/scx_chaos/src/lib.rs

@@ -95,6 +95,9 @@ pub enum Trait {
         frequency: f64,
         degradation_frac7: u64,
     },
+    BackwardsSticky {
+        frequency: f64,
+    },
 }
 
 impl Trait {
@@ -103,6 +106,7 @@ impl Trait {
             Self::RandomDelays { .. } => bpf_intf::chaos_trait_kind_CHAOS_TRAIT_RANDOM_DELAYS,
             Self::CpuFreq { .. } => bpf_intf::chaos_trait_kind_CHAOS_TRAIT_CPU_FREQ,
             Self::PerfDegradation { .. } => bpf_intf::chaos_trait_kind_CHAOS_TRAIT_DEGRADATION,
+            Self::BackwardsSticky { .. } => bpf_intf::chaos_trait_kind_CHAOS_TRAIT_BACKWARDS_STICKY,
         }
     }
 
@@ -111,6 +115,7 @@ impl Trait {
             Self::RandomDelays { frequency, .. } => *frequency,
             Self::CpuFreq { frequency, .. } => *frequency,
             Self::PerfDegradation { frequency, .. } => *frequency,
+            Self::BackwardsSticky { frequency, .. } => *frequency,
         }
     }
 }
@@ -409,6 +414,10 @@ impl Builder<'_> {
                         (frequency * 2_f64.powf(32_f64)) as u32;
                     open_skel.maps.rodata_data.degradation_frac7 = *degradation_frac7;
                 }
+                Trait::BackwardsSticky { frequency } => {
+                    open_skel.maps.rodata_data.backwards_sticky_freq_frac32 =
+                        (frequency * 2_f64.powf(32_f64)) as u32;
+                }
             }
         }
 
@@ -499,6 +508,15 @@ pub struct PerfDegradationArgs {
     pub degradation_frac7: u64,
 }
 
+/// Backwards Sticky scheduling is where a task is scheduled on all CPUs iteratively to break
+/// caches and then returned to sticky scheduling.
+#[derive(Debug, Parser)]
+pub struct BackwardsStickyArgs {
+    /// Chance of applying backward sticky to a process.
+    #[clap(long)]
+    pub degradation_frequency: Option<f64>,
+}
+
 /// scx_chaos: A general purpose sched_ext scheduler designed to amplify race conditions
 ///
 /// WARNING: This scheduler is a very early alpha, and hasn't been production tested yet. The CLI
@@ -540,6 +558,9 @@ pub struct Args {
     #[command(flatten, next_help_heading = "Perf Degradation")]
     pub perf_degradation: PerfDegradationArgs,
 
+    #[command(flatten, next_help_heading = "Backward Sticky")]
+    pub backward_sticky: BackwardsStickyArgs,
+
     #[command(flatten, next_help_heading = "CPU Frequency")]
     pub cpu_freq: CpuFreqArgs,