Fix memory leak

crux_core/Cargo.toml

@@ -29,6 +29,7 @@ serde-reflection = { version = "0.4.0", optional = true }
 serde_json = "1.0.128"
 slab = "0.4.9"
 thiserror = "1.0.63"
+uuid = { version = "1.10.0", features = ["v4", "serde"] }
 
 [dev-dependencies]
 assert_fs = "1.0.13"
@@ -41,4 +42,3 @@ serde = { version = "1.0.210", features = ["derive"] }
 static_assertions = "1.1"
 rand = "0.8"
 url = "2.5.2"
-uuid = { version = "1.10.0", features = ["v4", "serde"] }

crux_core/src/capability/executor.rs

@@ -1,4 +1,5 @@
 use std::{
+ collections::HashMap,
  sync::{Arc, Mutex},
  task::Context,
 };
@@ -9,46 +10,76 @@ use futures::{
  task::{waker_ref, ArcWake},
  Future, FutureExt,
 };
+use uuid::Uuid;
 
 // used in docs/internals/runtime.md
 // ANCHOR: executor
 pub(crate) struct QueuingExecutor {
- ready_queue: Receiver<Arc<Task>>,
+ task_queue: Receiver<Task>,
+ ready_queue: Receiver<Uuid>,
+ tasks: Mutex<HashMap<Uuid, Task>>,
 }
 // ANCHOR_END: executor
 
 // used in docs/internals/runtime.md
 // ANCHOR: spawner
 #[derive(Clone)]
 pub struct Spawner {
- task_sender: Sender<Arc<Task>>,
+ task_sender: Sender<Task>,
+ ready_sender: Sender<Uuid>,
 }
 // ANCHOR_END: spawner
 
 // used in docs/internals/runtime.md
 // ANCHOR: task
 struct Task {
- future: Mutex<Option<future::BoxFuture<'static, ()>>>,
+ id: Uuid,
+ future: future::BoxFuture<'static, ()>,
+ ready_sender: Sender<Uuid>,
+}
+
+impl Task {
+ fn id(&self) -> Uuid {
+ self.id
+ }
 
- task_sender: Sender<Arc<Task>>,
+ fn notify(&self) -> NotifyTask {
+ NotifyTask {
+ task_id: self.id,
+ sender: self.ready_sender.clone(),
+ }
+ }
 }
+
 // ANCHOR_END: task
 
 pub(crate) fn executor_and_spawner() -> (QueuingExecutor, Spawner) {
- let (task_sender, ready_queue) = crossbeam_channel::unbounded();
+ let (task_sender, task_queue) = crossbeam_channel::unbounded();
+ let (ready_sender, ready_queue) = crossbeam_channel::unbounded();
 
- (QueuingExecutor { ready_queue }, Spawner { task_sender })
+ (
+ QueuingExecutor {
+ ready_queue,
+ task_queue,
+ tasks: Mutex::new(HashMap::new()),
+ },
+ Spawner {
+ task_sender,
+ ready_sender,
+ },
+ )
 }
 
 // used in docs/internals/runtime.md
 // ANCHOR: spawning
 impl Spawner {
  pub fn spawn(&self, future: impl Future<Output = ()> + 'static + Send) {
  let future = future.boxed();
- let task = Arc::new(Task {
- future: Mutex::new(Some(future)),
- task_sender: self.task_sender.clone(),
- });
+ let task = Task {
+ id: Uuid::new_v4(),
+ future,
+ ready_sender: self.ready_sender.clone(),
+ };
 
  self.task_sender
  .send(task)
@@ -59,38 +90,88 @@ impl Spawner {
 
 // used in docs/internals/runtime.md
 // ANCHOR: arc_wake
-impl ArcWake for Task {
+impl ArcWake for NotifyTask {
  fn wake_by_ref(arc_self: &Arc<Self>) {
- let cloned = arc_self.clone();
- arc_self
- .task_sender
- .send(cloned)
- .expect("unable to wake an async task, task sender channel is disconnected.")
+ let _ = arc_self.sender.send(arc_self.task_id);
+ // TODO should we report an error if send fails?
  }
 }
 // ANCHOR_END: arc_wake
 
+struct NotifyTask {
+ task_id: Uuid,
+ sender: Sender<Uuid>,
+}
+
 // used in docs/internals/runtime.md
 // ANCHOR: run_all
 impl QueuingExecutor {
  pub fn run_all(&self) {
- // While there are tasks to be processed
- while let Ok(task) = self.ready_queue.try_recv() {
- // Unlock the future in the Task
- let mut future_slot = task.future.lock().unwrap();
-
- // Take it, replace with None, ...
- if let Some(mut future) = future_slot.take() {
- let waker = waker_ref(&task);
- let context = &mut Context::from_waker(&waker);
-
- // ...and poll it
- if future.as_mut().poll(context).is_pending() {
- // If it's still pending, put it back
- *future_slot = Some(future)
- }
+ // we read off both queues and execute the tasks we receive.
+ // Since either queue can generate work for the other queue,
+ // we read from them in a loop until we are sure both queues
+ // are exhaused
+ let mut did_some_work = true;
+
+ while did_some_work {
+ did_some_work = false;
+ // While there are tasks to be processed
+ while let Ok(task) = self.task_queue.try_recv() {
+ let task_id = task.id();
+ self.tasks.lock().unwrap().insert(task_id, task);
+ self.run_task(task_id);
+ did_some_work = true;
  }
+ while let Ok(task_id) = self.ready_queue.try_recv() {
+ self.run_task(task_id);
+ did_some_work = true;
+ }
+ }
+ }
+
+ fn run_task(&self, task_id: Uuid) {
+ let mut tasks = self.tasks.lock().unwrap();
+ let mut task = tasks.remove(&task_id).unwrap();
+ drop(tasks);
+
+ let notify = Arc::new(task.notify());
+ let waker = waker_ref(&notify);
+ let context = &mut Context::from_waker(&waker);
+
+ // ...and poll it
+ if task.future.as_mut().poll(context).is_pending() {
+ // If it's still pending, put it back
+ self.tasks.lock().unwrap().insert(task.id, task);
  }
  }
 }
 // ANCHOR_END: run_all
+
+#[cfg(test)]
+mod tests {
+ use crate::capability::shell_request::ShellRequest;
+
+ use super::*;
+
+ #[test]
+ fn test_task_does_not_leak() {
+ let counter: Arc<()> = Arc::new(());
+ assert_eq!(Arc::strong_count(&counter), 1);
+
+ let (executor, spawner) = executor_and_spawner();
+
+ let future = {
+ let counter = counter.clone();
+ async move {
+ assert_eq!(Arc::strong_count(&counter), 2);
+ ShellRequest::<()>::new().await;
+ }
+ };
+
+ spawner.spawn(future);
+ executor.run_all();
+ drop(executor);
+ drop(spawner);
+ assert_eq!(Arc::strong_count(&counter), 1);
+ }
+}

crux_core/src/capability/shell_request.rs

@@ -13,6 +13,19 @@ pub struct ShellRequest<T> {
  shared_state: Arc<Mutex<SharedState<T>>>,
 }
 
+#[cfg(test)]
+impl ShellRequest<()> {
+ pub(crate) fn new() -> Self {
+ Self {
+ shared_state: Arc::new(Mutex::new(SharedState {
+ result: None,
+ waker: None,
+ send_request: None,
+ })),
+ }
+ }
+}
+
 struct SharedState<T> {
  result: Option<T>,
  waker: Option<Waker>,
@@ -38,7 +51,8 @@ impl<T> Future for ShellRequest<T> {
  match shared_state.result.take() {
  Some(result) => Poll::Ready(result),
  None => {
- shared_state.waker = Some(cx.waker().clone());
+ let cloned_waker = cx.waker().clone();
+ shared_state.waker = Some(cloned_waker);
  Poll::Pending
  }
  }