scx_rustsched: Let callbacks hold kernel objects the kernel changes

rust/scx_rustsched/src/cgroup.rs

@@ -31,17 +31,18 @@
 //! ancestor or one found by its id, is a `CgroupGuard`: a reference that
 //! is given back when the guard goes out of scope.
 
-use core::{ffi::c_void, ops::Deref, ptr::read_volatile};
+use core::{cell::UnsafeCell, ffi::c_void, ops::Deref, ptr::read_volatile};
 
 use crate::{
     config, helpers,
     kfuncs::{bpf_cgroup_ancestor, bpf_cgroup_from_id, bpf_cgroup_release, scx_bpf_cgroup_nr_cpus},
     ops::{cgroup, scx_cgroup_init_args},
 };
 
-/// A cgroup, valid for the callback it was passed to.
+/// A cgroup, valid for the callback it was passed to. The kernel changes
+/// it while the callback holds it, hence the `UnsafeCell`, as for `Task`.
 #[repr(transparent)]
-pub struct Cgroup(cgroup);
+pub struct Cgroup(UnsafeCell<cgroup>);
 
 impl Cgroup {
     /// # Safety
@@ -54,7 +55,7 @@ impl Cgroup {
     /// The cgroup as the kfuncs take it.
     #[inline(always)]
     pub fn as_ptr(&self) -> *mut cgroup {
-        (self as *const Cgroup).cast_mut().cast()
+        self.0.get()
     }
 
     /// How deep the cgroup is in the hierarchy: 0 for the root.

rust/scx_rustsched/src/task.rs

@@ -5,7 +5,7 @@
 //! resolved from the kernel's BTF: its policy and weight, whether it is
 //! still runnable, and where it may run.
 
-use core::ptr::read_volatile;
+use core::{cell::UnsafeCell, ptr::read_volatile};
 
 use crate::{
     config, helpers,
@@ -154,8 +154,13 @@ pub fn allowed_on(task: &Task, cid: crate::sched::Cid) -> bool {
 
 /// A task, as a callback of `#[scheduler]` takes it: a reference to the
 /// kernel's task_struct that is valid for the callback it was passed to.
+///
+/// The kernel and the kfuncs change the task while a callback holds it.
+/// The `UnsafeCell` tells the compiler so: without it `&Task` would be a
+/// read-only, non-aliased reference, and it could drop writes made through
+/// it or reuse reads made before a kfunc changed the task.
 #[repr(transparent)]
-pub struct Task(task_struct);
+pub struct Task(UnsafeCell<task_struct>);
 
 impl Task {
     /// # Safety
@@ -168,7 +173,7 @@ impl Task {
     /// The task as the kfuncs take it.
     #[inline(always)]
     pub fn as_ptr(&self) -> *mut task_struct {
-        (self as *const Task).cast_mut().cast()
+        self.0.get()
     }
 
     /// The task's weight, from its nice level: 1024 at nice 0.

rust/scx_rustsched/src/timer.rs

@@ -23,16 +23,19 @@
 //! the hrtimer from there: it queues an irq_work, a self-IPI per arming.
 //! From the callback itself, with interrupts on, it arms directly.
 
-use core::ffi::c_void;
+use core::{cell::UnsafeCell, ffi::c_void};
 
 use crate::helpers;
 
 const CLOCK_MONOTONIC: u64 = 1;
 
-/// struct bpf_timer.
+/// struct bpf_timer. The kernel changes it while a callback holds a
+/// `&Timer`, hence the `UnsafeCell`, as for `Task`. It wraps the words
+/// rather than the struct, since the verifier finds the timer in a map
+/// value by the name bpf_timer.
 #[repr(C)]
 pub struct bpf_timer {
-    opaque: [u64; 2],
+    opaque: UnsafeCell<[u64; 2]>,
 }
 
 /// A timer map's value.
@@ -46,7 +49,7 @@ impl Timer {
     /// armed. A timer that was never set up stays as it is.
     #[inline(always)]
     pub fn start(&self, delta: u64) {
-        unsafe { start((self as *const Timer).cast_mut(), delta) };
+        unsafe { start(self.timer.opaque.get().cast(), delta) };
     }
 }