actor system: Don't allow worker queues to become too deep, queue jobs in factory instead (risc0#3475)

This makes it so that we don't schedule more than a configurable amount
of work (default 3 tasks) on a worker. This is an improvement because it
delays commit a task running on a worker until a later point in time at
which we can make a potentially better scheduling decision.

This goes about implementing this by doing the following things:
- RPC system `ask` function takes async callback
- Actor code ReplySender `send` function becomes async function
- Add actor system `ask_callback` function that takes an async callback
that gets called with the reply
- This allows a reply to become a `tell` function call on another actor
- Factory no longer blocks messages waiting for a response to
`ScheduleTask`
- If the allocator was to only schedule a task much later in time, its
important to not block all factory messages until then
- Allocator only schedules a task once there is a worker available that
is under the task threshold.
- new allocator configuration for worker threshold (default to 3 for
now)
- A new trace span is added for tasks that measures the time it spends
in the factory waiting for a worker
- this keeps there being no major period of time unaccounted for with
tasks

Author: bobbobbio

risc0/r0vm/src/actors/actor.rs

@@ -129,6 +129,24 @@ impl<ActorT: Actor> ActorRunner<ActorT> {
     }
 }
 
+/// Called when an actor handles a message
+async fn actor_msg<ActorT, MessageT>(
+    actor: &mut ActorT,
+    actor_ref: ActorRef<ActorT>,
+    reply_sender: Option<ReplySender<<ActorT as Message<MessageT>>::Reply>>,
+    msg: MessageT,
+) where
+    ActorT: Actor + Message<MessageT>,
+{
+    let mut context = Context {
+        actor_ref,
+        reply_sender,
+    };
+    Message::handle(actor, msg, &mut context).await;
+
+    context.maybe_send_no_reply().await;
+}
+
 /// Spawn a new task which receives messages for the actor.
 pub fn spawn<ActorT: Actor>(actor: ActorT) -> ActorRef<ActorT> {
     let (actor_ref, actor_runner) = ActorController::run(actor);
@@ -253,6 +271,13 @@ impl<ActorT: Actor> ActorController<ActorT> {
     fn stop_waiter(&self) -> broadcast::Receiver<()> {
         self.stop.resubscribe()
     }
+
+    fn stop_reason(&self) -> Option<String> {
+        match &self.sender {
+            SenderState::Stopped { stop_reason } => Some(stop_reason.clone()),
+            _ => None,
+        }
+    }
 }
 
 pub struct ActorRef<ActorT: Actor> {
@@ -302,6 +327,9 @@ impl<ActorT: Actor> ActorRef<ActorT> {
     }
 
     /// Send a message to an actor and return an object which can be used to wait for a reply.
+    /// If the remote actor panic handling the message, the `PendingReply` will return
+    /// `Err(SendError::ActorNotRunning)`. If the remote actor doesn't reply, it will return
+    /// `Err(SendError::NoReply)`.
     pub async fn ask_enqueue<MessageT>(
         &self,
         msg: MessageT,
@@ -310,23 +338,39 @@ impl<ActorT: Actor> ActorRef<ActorT> {
         MessageT: Send + 'static,
         ActorT: Message<MessageT>,
     {
-        let actor_ref = self.clone();
         let (sender, reply) = oneshot::channel();
-        let reply_sender = ReplySender { sender };
+        self.ask_callback(msg, async |res| {
+            let _ = sender.send(res);
+        })
+        .await?;
+        Ok(PendingReply { reply })
+    }
+
+    /// Send a message to an actor and call a callback with the response.
+    /// If the remote actor panics handling the message the callback isn't called. If the remote
+    /// actor doesn't reply, the callback is called with `Err(SendError::NoReply)`.
+    pub async fn ask_callback<MessageT, FutT>(
+        &self,
+        msg: MessageT,
+        callback: impl FnOnce(Result<<ActorT as Message<MessageT>>::Reply, SendError>) -> FutT
+        + Send
+        + 'static,
+    ) -> Result<(), SendError>
+    where
+        MessageT: Send + 'static,
+        ActorT: Message<MessageT>,
+        FutT: Future<Output = ()> + Send + 'static,
+    {
+        let actor_ref = self.clone();
+        let reply_sender = ReplySender::new(callback);
 
         let sender = self.controller.lock().unwrap().sender()?;
         sender
             .send(Box::new(move |actor| {
-                Box::pin(async move {
-                    let mut context = Context {
-                        actor_ref,
-                        reply_sender: Some(reply_sender),
-                    };
-                    Message::handle(actor, msg, &mut context).await;
-                })
+                Box::pin(actor_msg(actor, actor_ref, Some(reply_sender), msg))
             }))
             .await?;
-        Ok(PendingReply { reply })
+        Ok(())
     }
 
     /// Send a message to an actor and don't wait for a reply.
@@ -339,13 +383,7 @@ impl<ActorT: Actor> ActorRef<ActorT> {
         let sender = self.controller.lock().unwrap().sender()?;
         sender
             .send(Box::new(move |actor| {
-                Box::pin(async move {
-                    let mut context = Context {
-                        actor_ref,
-                        reply_sender: None,
-                    };
-                    Message::handle(actor, msg, &mut context).await;
-                })
+                Box::pin(actor_msg(actor, actor_ref, None, msg))
             }))
             .await
     }
@@ -376,13 +414,7 @@ impl<ActorT: Actor> ActorRef<ActorT> {
         let actor_ref = self.clone();
         let sender = self.controller.lock().unwrap().sender()?;
         sender.blocking_send(Box::new(move |actor| {
-            Box::pin(async move {
-                let mut context = Context {
-                    actor_ref,
-                    reply_sender: None,
-                };
-                Message::handle(actor, msg, &mut context).await;
-            })
+            Box::pin(actor_msg(actor, actor_ref, None, msg))
         }))
     }
 
@@ -399,13 +431,7 @@ impl<ActorT: Actor> ActorRef<ActorT> {
         let sender = self.controller.lock().unwrap().sender()?;
         sender
             .send(Box::new(move |actor| {
-                Box::pin(async move {
-                    let mut context = Context {
-                        actor_ref,
-                        reply_sender,
-                    };
-                    Message::handle(actor, msg, &mut context).await;
-                })
+                Box::pin(actor_msg(actor, actor_ref, reply_sender, msg))
             }))
             .await
     }
@@ -435,6 +461,12 @@ impl<ActorT: Actor> ActorRef<ActorT> {
             controller: Arc::downgrade(&self.controller),
         }
     }
+
+    /// If the actor is stopped, get the reason.
+    #[allow(dead_code)]
+    pub fn stop_reason(&self) -> Option<String> {
+        self.controller.lock().unwrap().stop_reason()
+    }
 }
 
 /// A weak reference to an actor. Holding this reference will not stop an actor from being
@@ -461,18 +493,29 @@ impl<ActorT: Actor> Clone for WeakActorRef<ActorT> {
     }
 }
 
+type ReplySenderAsyncCb<ValueT> =
+    Box<dyn FnOnce(Result<ValueT, SendError>) -> Pin<Box<dyn Future<Output = ()> + Send>> + Send>;
+
 pub struct ReplySender<ValueT: Send + 'static> {
-    sender: oneshot::Sender<Result<ValueT, SendError>>,
+    sender: ReplySenderAsyncCb<ValueT>,
 }
 
 impl<ValueT: Send + 'static> ReplySender<ValueT> {
-    pub fn send(self, value: ValueT) {
-        // If the other side is no longer listening, we throw away the reply
-        let _ = self.sender.send(Ok(value));
+    fn new<FutT>(sender_fn: impl FnOnce(Result<ValueT, SendError>) -> FutT + Send + 'static) -> Self
+    where
+        FutT: Future<Output = ()> + Send + 'static,
+    {
+        Self {
+            sender: Box::new(|res| Box::pin(sender_fn(res))),
+        }
     }
 
-    pub fn no_reply(self) {
-        let _ = self.sender.send(Err(SendError::NoReply));
+    pub async fn send(self, value: ValueT) {
+        (self.sender)(Ok(value)).await;
+    }
+
+    pub async fn no_reply(self) {
+        (self.sender)(Err(SendError::NoReply)).await;
     }
 }
 
@@ -483,14 +526,6 @@ pub struct Context<ActorT: Actor, ReplyT: Send + 'static> {
     reply_sender: Option<ReplySender<ReplyT>>,
 }
 
-impl<ActorT: Actor, ReplyT: Send + 'static> Drop for Context<ActorT, ReplyT> {
-    fn drop(&mut self) {
-        if let Some(reply_sender) = self.reply_sender.take() {
-            reply_sender.no_reply()
-        }
-    }
-}
-
 impl<ActorT: Actor, ReplyT: Send + 'static> fmt::Debug for Context<ActorT, ReplyT> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("Context").finish_non_exhaustive()
@@ -507,9 +542,9 @@ impl<ActorT: Actor, ReplyT: Send + 'static> Context<ActorT, ReplyT> {
 
     /// If the sender of the message is expecting a reply, send the given message as a reply to
     /// them. Otherwise, if they aren't just ignore the given message.
-    pub fn reply(&mut self, msg: ReplyT) {
+    pub async fn reply(&mut self, msg: ReplyT) {
         if let Some(reply_sender) = self.reply_sender.take() {
-            reply_sender.send(msg);
+            reply_sender.send(msg).await;
         }
     }
 
@@ -530,6 +565,12 @@ impl<ActorT: Actor, ReplyT: Send + 'static> Context<ActorT, ReplyT> {
     pub fn actor_ref(&self) -> ActorRef<ActorT> {
         self.actor_ref.clone()
     }
+
+    async fn maybe_send_no_reply(&mut self) {
+        if let Some(reply_sender) = self.reply_sender.take() {
+            reply_sender.no_reply().await
+        }
+    }
 }
 
 /// The ability to receive and possible reply to a certain message.
@@ -613,7 +654,7 @@ mod tests {
             self.ping
                 .send(msg.clone())
                 .expect("messages should only be received when expected");
-            ctx.reply(Pong(msg.0));
+            ctx.reply(Pong(msg.0)).await;
         }
     }
 
@@ -661,6 +702,23 @@ mod tests {
         assert_eq!(spy.ping.recv().await.unwrap(), Ping(12));
     }
 
+    #[tokio::test]
+    async fn ask_callback() {
+        let (actor, mut spy) = TestActor::new();
+        let actor_ref = spawn(actor);
+
+        let (sender, reply) = oneshot::channel();
+        actor_ref
+            .ask_callback(Ping(12), async |res| {
+                sender.send(res.unwrap()).unwrap();
+            })
+            .await
+            .unwrap();
+        assert_eq!(reply.await.unwrap(), Pong(12));
+
+        assert_eq!(spy.ping.recv().await.unwrap(), Ping(12));
+    }
+
     #[tokio::test]
     async fn ask_with_runner() {
         let (actor, mut spy) = TestActor::new();