Merge pull request #678 from sched-ext/rustland-core-fix-mm-stall

scx_rustland_core: fix mm stall

rust/scx_rustland_core/assets/bpf/main.bpf.c

@@ -310,6 +310,46 @@ static bool in_setaffinity(pid_t pid)
 	return value != NULL;
 }
 
+/*
+ * Intercept when a task is executing __handle_mm_fault().
+ */
+struct {
+	__uint(type, BPF_MAP_TYPE_HASH);
+	__type(key, __u32);
+	__type(value, __u64);
+	__uint(max_entries, MAX_ENQUEUED_TASKS);
+} pid_mm_fault_map SEC(".maps");
+
+SEC("kprobe/__handle_mm_fault")
+int BPF_KPROBE(kprobe_handle_mm_fault, void *vma,
+			unsigned long address, unsigned int flags)
+{
+	pid_t pid = bpf_get_current_pid_tgid() >> 32;
+	u64 value = true;
+
+	bpf_map_update_elem(&pid_mm_fault_map, &pid, &value, BPF_ANY);
+
+	return 0;
+}
+
+SEC("kretprobe/__handle_mm_fault")
+int BPF_KRETPROBE(kretprobe_handle_mm_fault)
+{
+	pid_t pid = bpf_get_current_pid_tgid() >> 32;
+	bpf_map_delete_elem(&pid_mm_fault_map, &pid);
+
+	return 0;
+}
+
+/*
+ * Return true if a task is handling a page fault, false otherwise.
+ */
+static bool in_mm_fault(pid_t pid)
+{
+	u64 *value = bpf_map_lookup_elem(&pid_mm_fault_map, &pid);
+	return value != NULL;
+}
+
 /*
  * Heartbeat timer used to periodically trigger the check to run the user-space
  * scheduler.
@@ -848,12 +888,13 @@ void BPF_STRUCT_OPS(rustland_enqueue, struct task_struct *p, u64 enq_flags)
 		return;
 
 	/*
-	 * Always dispatch all kthreads directly on their target CPU.
+	 * Always dispatch per-CPU kthreads directly on their target CPU.
 	 *
-	 * This allows to prevent critical kernel threads from being stuck in
-	 * user-space causing system hangs.
+	 * This allows to prioritize critical kernel threads that may
+	 * potentially stall the entire system if they are blocked for too long
+	 * (i.e., ksoftirqd/N, rcuop/N, etc.).
 	 */
-	if (is_kthread(p)) {
+	if (is_kthread(p) && p->nr_cpus_allowed == 1) {
 		s32 cpu = scx_bpf_task_cpu(p);
 
 		scx_bpf_dispatch_vtime(p, cpu_to_dsq(cpu),
@@ -862,6 +903,16 @@ void BPF_STRUCT_OPS(rustland_enqueue, struct task_struct *p, u64 enq_flags)
 		return;
 	}
 
+	/*
+	 * Bypass user-space scheduling for faulting tasks to prevent potential
+	 * deadlock conditions. They can just be dispatched to the shared DSQ
+	 * using ith the highest priority.
+	 */
+	if (in_mm_fault(p->pid)) {
+		scx_bpf_dispatch_vtime(p, SHARED_DSQ, SCX_SLICE_DFL, 0, enq_flags);
+		return;
+	}
+
 	/*
 	 * Add tasks to the @queued list, they will be processed by the
 	 * user-space scheduler.