scheduler: Fix scheduler retention cycle (#57789)

async_tasks can build a cycle with the scheduler where the runnable
itself keeps a reference to the scheduler and the scheduler keeping
references to the runnable. This effectively can cause scheduler / task
leaks, most notably in test environments, so we break this cycle by
using weak pointers instead.

Release Notes:

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

crates/scheduler/src/executor.rs

@@ -43,12 +43,14 @@ impl ForegroundExecutor {
         F::Output: 'static,
     {
         let session_id = self.session_id;
-        let scheduler = Arc::clone(&self.scheduler);
+        let scheduler = Arc::downgrade(&self.scheduler);
         let location = Location::caller();
         let (runnable, task) = spawn_local_with_source_location(
             future,
             move |runnable| {
-                scheduler.schedule_foreground(session_id, runnable);
+                if let Some(scheduler) = scheduler.upgrade() {
+                    scheduler.schedule_foreground(session_id, runnable);
+                }
             },
             RunnableMeta { location },
         );
@@ -135,14 +137,16 @@ impl BackgroundExecutor {
         F: Future + Send + 'static,
         F::Output: Send + 'static,
     {
-        let scheduler = Arc::clone(&self.scheduler);
+        let scheduler = Arc::downgrade(&self.scheduler);
         let location = Location::caller();
         let (runnable, task) = async_task::Builder::new()
             .metadata(RunnableMeta { location })
             .spawn(
                 move |_| future,
                 move |runnable| {
-                    scheduler.schedule_background_with_priority(runnable, priority);
+                    if let Some(scheduler) = scheduler.upgrade() {
+                        scheduler.schedule_background_with_priority(runnable, priority);
+                    }
                 },
             );
         runnable.schedule();

crates/scheduler/src/tests.rs

@@ -34,6 +34,44 @@ fn test_background_executor_spawn() {
     });
 }
 
+#[test]
+fn test_scheduler_drops_with_stalled_detached_foreground_task() {
+    let scheduler = Arc::new(TestScheduler::new(TestSchedulerConfig::default()));
+    let weak_scheduler = Arc::downgrade(&scheduler);
+    let (sender, receiver) = oneshot::channel::<()>();
+
+    scheduler
+        .foreground()
+        .spawn(async move {
+            receiver.await.ok();
+        })
+        .detach();
+    scheduler.run();
+
+    drop(scheduler);
+    assert!(weak_scheduler.upgrade().is_none());
+    drop(sender);
+}
+
+#[test]
+fn test_scheduler_drops_with_stalled_detached_background_task() {
+    let scheduler = Arc::new(TestScheduler::new(TestSchedulerConfig::default()));
+    let weak_scheduler = Arc::downgrade(&scheduler);
+    let (sender, receiver) = oneshot::channel::<()>();
+
+    scheduler
+        .background()
+        .spawn(async move {
+            receiver.await.ok();
+        })
+        .detach();
+    scheduler.run();
+
+    drop(scheduler);
+    assert!(weak_scheduler.upgrade().is_none());
+    drop(sender);
+}
+
 #[test]
 fn test_foreground_ordering() {
     let mut traces = HashSet::new();