scheduler: Fix stacked borrows violation in scheduler (#56850)

Release Notes:

- N/A or Added/Fixed/Improved ...

.github/workflows/run_tests.yml

@@ -448,6 +448,34 @@ jobs:
       run: |
         rm -rf ./../.cargo
     timeout-minutes: 60
+  miri_scheduler:
+    needs:
+    - orchestrate
+    if: needs.orchestrate.outputs.run_tests == 'true' && github.event_name != 'merge_group'
+    runs-on: namespace-profile-16x32-ubuntu-2204
+    steps:
+    - name: steps::checkout_repo
+      uses: actions/checkout@93cb6efe18208431cddfb8368fd83d5badbf9bfd
+      with:
+        clean: false
+    - name: steps::setup_cargo_config
+      run: |
+        mkdir -p ./../.cargo
+        cp ./.cargo/ci-config.toml ./../.cargo/config.toml
+    - name: steps::cache_rust_dependencies_namespace
+      uses: namespacelabs/nscloud-cache-action@a90bb5d4b27522ce881c6e98eebd7d7e6d1653f9
+      with:
+        cache: rust
+        path: ~/.rustup
+    - name: run_tests::miri_scheduler::install_miri
+      run: rustup toolchain install nightly --profile minimal --component miri --component rust-src
+    - name: run_tests::miri_scheduler::run_scheduler_tests_under_miri
+      run: cargo +nightly -q miri test -p scheduler
+    - name: steps::cleanup_cargo_config
+      if: always()
+      run: |
+        rm -rf ./../.cargo
+    timeout-minutes: 60
   doctests:
     needs:
     - orchestrate
@@ -804,6 +832,7 @@ jobs:
     - run_tests_windows
     - run_tests_linux
     - run_tests_mac
+    - miri_scheduler
     - doctests
     - check_workspace_binaries
     - build_visual_tests_binary
@@ -835,6 +864,7 @@ jobs:
         check_result "run_tests_windows" "$RESULT_RUN_TESTS_WINDOWS"
         check_result "run_tests_linux" "$RESULT_RUN_TESTS_LINUX"
         check_result "run_tests_mac" "$RESULT_RUN_TESTS_MAC"
+        check_result "miri_scheduler" "$RESULT_MIRI_SCHEDULER"
         check_result "doctests" "$RESULT_DOCTESTS"
         check_result "check_workspace_binaries" "$RESULT_CHECK_WORKSPACE_BINARIES"
         check_result "build_visual_tests_binary" "$RESULT_BUILD_VISUAL_TESTS_BINARY"
@@ -856,6 +886,7 @@ jobs:
         RESULT_RUN_TESTS_WINDOWS: ${{ needs.run_tests_windows.result }}
         RESULT_RUN_TESTS_LINUX: ${{ needs.run_tests_linux.result }}
         RESULT_RUN_TESTS_MAC: ${{ needs.run_tests_mac.result }}
+        RESULT_MIRI_SCHEDULER: ${{ needs.miri_scheduler.result }}
         RESULT_DOCTESTS: ${{ needs.doctests.result }}
         RESULT_CHECK_WORKSPACE_BINARIES: ${{ needs.check_workspace_binaries.result }}
         RESULT_BUILD_VISUAL_TESTS_BINARY: ${{ needs.build_visual_tests_binary.result }}

crates/scheduler/src/executor.rs

@@ -343,9 +343,10 @@ where
                 "local task dropped by a thread that didn't spawn it. Task spawned at {}",
                 self.location
             );
-            unsafe {
-                ManuallyDrop::drop(&mut self.inner);
-            }
+            // SAFETY: `inner` is wrapped in `ManuallyDrop`, so this is the only
+            // place it is dropped. The thread check above ensures local futures
+            // are dropped on the thread that created them.
+            unsafe { ManuallyDrop::drop(&mut self.inner) };
         }
     }
 
@@ -353,27 +354,34 @@ where
         type Output = F::Output;
 
         fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+            // SAFETY: We don't move any fields out of `self`; this mutable
+            // reference is only used to check metadata and to project the pin to
+            // `inner` below.
+            let this = unsafe { self.get_unchecked_mut() };
             assert!(
-                self.id == thread_id(),
+                this.id == thread_id(),
                 "local task polled by a thread that didn't spawn it. Task spawned at {}",
-                self.location
+                this.location
             );
-            unsafe { self.map_unchecked_mut(|c| &mut *c.inner).poll(cx) }
+            // SAFETY: `inner` is structurally pinned by `Checked`; after
+            // `Checked` is pinned, `inner` is never moved. The thread check
+            // above ensures the local future is only polled by its spawning
+            // thread.
+            unsafe { Pin::new_unchecked(&mut *this.inner).poll(cx) }
         }
     }
 
     let location = metadata.location;
 
-    unsafe {
-        async_task::Builder::new()
-            .metadata(metadata)
-            .spawn_unchecked(
-                move |_| Checked {
-                    id: thread_id(),
-                    inner: ManuallyDrop::new(future),
-                    location,
-                },
-                schedule,
-            )
-    }
+    let future = move |_| Checked {
+        id: thread_id(),
+        inner: ManuallyDrop::new(future),
+        location,
+    };
+
+    let builder = async_task::Builder::new().metadata(metadata);
+    // SAFETY: `Checked` enforces the invariants required by `spawn_unchecked`:
+    // the non-`Send` future is only polled and dropped on the thread that
+    // spawned it.
+    unsafe { builder.spawn_unchecked(future, schedule) }
 }

crates/scheduler/src/test_scheduler.rs

@@ -108,7 +108,12 @@ impl TestScheduler {
     pub fn end_test(&self) {
         let mut state = self.state.lock();
         if let Some((message, backtrace)) = &state.non_determinism_error {
-            panic!("{}\n{:?}", message, backtrace)
+            if cfg!(miri) {
+                // miri cannot debug print backtraces with `miri-disable-isolation` enabled
+                panic!("{}", message)
+            } else {
+                panic!("{}\n{:?}", message, backtrace)
+            }
         }
         state.finished = true;
     }
@@ -458,6 +463,9 @@ impl TestScheduler {
             }
         } else if deadline.is_some() {
             false
+        } else if cfg!(miri) {
+            // miri cannot debug print backtraces with `miri-disable-isolation` enabled
+            panic!("Parking forbidden.");
         } else if self.state.lock().capture_pending_traces {
             let mut pending_traces = String::new();
             for (_, trace) in mem::take(&mut self.state.lock().pending_traces) {

crates/scheduler/src/tests.rs

@@ -38,7 +38,7 @@ fn test_background_executor_spawn() {
 fn test_foreground_ordering() {
     let mut traces = HashSet::new();
 
-    TestScheduler::many(100, async |scheduler| {
+    TestScheduler::many(if cfg!(miri) { 5 } else { 100 }, async |scheduler| {
         #[derive(Hash, PartialEq, Eq)]
         struct TraceEntry {
             session: usize,
@@ -128,6 +128,24 @@ fn test_timer_ordering() {
     });
 }
 
+#[test]
+fn test_foreground_task_can_hold_mut_borrow_across_await() {
+    TestScheduler::once(async |scheduler| {
+        let foreground = scheduler.foreground();
+        let (sender, mut receiver) = mpsc::unbounded::<()>();
+
+        foreground
+            .spawn(async move {
+                receiver.next().await;
+            })
+            .detach();
+
+        scheduler.run();
+        sender.unbounded_send(()).unwrap();
+        scheduler.run();
+    });
+}
+
 #[test]
 fn test_send_from_bg_to_fg() {
     TestScheduler::once(async |scheduler| {
@@ -238,7 +256,7 @@ fn test_block() {
 }
 
 #[test]
-#[should_panic(expected = "Parking forbidden. Pending traces:")]
+#[should_panic(expected = "Parking forbidden.")]
 fn test_parking_panics() {
     let config = TestSchedulerConfig {
         capture_pending_traces: true,
@@ -340,7 +358,7 @@ fn test_block_with_timeout() {
 
     // Test case: future makes progress via timer but still times out
     let mut results = BTreeSet::new();
-    TestScheduler::many(100, async |scheduler| {
+    TestScheduler::many(if cfg!(miri) { 5 } else { 100 }, async |scheduler| {
         // Keep the existing probabilistic behavior here (do not force 0 ticks), since this subtest
         // is explicitly checking that some seeds/timeouts can complete while others can time out.
         let task = scheduler.background().spawn(async move {
@@ -354,7 +372,11 @@ fn test_block_with_timeout() {
     });
     assert_eq!(
         results.into_iter().collect::<Vec<_>>(),
-        vec![None, Some(42)]
+        if cfg!(miri) {
+            vec![Some(42)]
+        } else {
+            vec![None, Some(42)]
+        }
     );
 
     // Regression test:
@@ -400,7 +422,7 @@ fn test_block_with_timeout() {
 #[test]
 fn test_block_does_not_progress_same_session_foreground() {
     let mut task2_made_progress_once = false;
-    TestScheduler::many(1000, async |scheduler| {
+    TestScheduler::many(if cfg!(miri) { 5 } else { 1000 }, async |scheduler| {
         let foreground1 = scheduler.foreground();
         let foreground2 = scheduler.foreground();
 
@@ -614,7 +636,7 @@ fn test_background_priority_scheduling() {
 
     // Run many iterations to get statistical significance
     let mut high_before_low_count = 0;
-    let iterations = 100;
+    let iterations = if cfg!(miri) { 5 } else { 100 };
 
     for seed in 0..iterations {
         let config = TestSchedulerConfig::with_seed(seed);

tooling/xtask/src/tasks/workflows/run_tests.rs

@@ -70,6 +70,9 @@ pub(crate) fn run_tests() -> Workflow {
         should_run_tests
             .and_not_in_merge_queue()
             .then(run_platform_tests(Platform::Mac)),
+        should_run_tests
+            .and_not_in_merge_queue()
+            .then(miri_scheduler()),
         should_run_tests.and_not_in_merge_queue().then(doctests()),
         should_run_tests
             .and_not_in_merge_queue()
@@ -690,6 +693,29 @@ pub(crate) fn check_postgres_and_protobuf_migrations() -> NamedJob {
     )
 }
 
+fn miri_scheduler() -> NamedJob {
+    fn install_miri() -> Step<Run> {
+        named::bash(
+            "rustup toolchain install nightly --profile minimal --component miri --component rust-src",
+        )
+    }
+
+    fn run_scheduler_tests_under_miri() -> Step<Run> {
+        named::bash("cargo +nightly -q miri test -p scheduler")
+    }
+
+    named::job(
+        release_job(&[])
+            .runs_on(runners::LINUX_DEFAULT)
+            .add_step(steps::checkout_repo())
+            .add_step(steps::setup_cargo_config(Platform::Linux))
+            .add_step(steps::cache_rust_dependencies_namespace())
+            .add_step(install_miri())
+            .add_step(run_scheduler_tests_under_miri())
+            .add_step(steps::cleanup_cargo_config(Platform::Linux)),
+    )
+}
+
 fn doctests() -> NamedJob {
     fn run_doctests() -> Step<Run> {
         named::bash(indoc::indoc! {r#"