nit

Author: patrick-ogrady

.github/dylints/blocking_in_actor_loop/src/lib.rs

@@ -18,14 +18,17 @@ dylint_linting::declare_late_lint! {
     ///
     /// Detects blocking waits in actor event loops, such as awaiting an actor
     /// mailbox request/reply method or an async handler callback inside
-    /// `select_loop!`.
+    /// `select_loop!`. It also detects async mailbox request/reply helpers that
+    /// create a oneshot channel, enqueue a message with the sender, and then
+    /// await the receiver before returning.
     ///
     /// ### Why is this bad?
     ///
     /// Actor loops must keep polling their mailbox and other event sources.
     /// Awaiting another actor or an application callback from a branch body can
-    /// starve the loop. Spawn or pool the work and select on its completion
-    /// instead.
+    /// starve the loop. Hiding the wait behind an async mailbox helper makes it
+    /// too easy for actor loops to block accidentally. Spawn or pool the work
+    /// and select on its completion instead.
     ///
     /// ### Example
     ///
@@ -62,6 +65,20 @@ impl<'tcx> LateLintPass<'tcx> for BlockingInActorLoop {
             Some(kind) => BlockingAwaitVisitor { cx, kind }.visit_expr(body.value),
             None => {}
         }
+
+        if let Some(span) = request_reply_helper_wait(cx, body) {
+            cx.emit_span_lint(
+                BLOCKING_IN_ACTOR_LOOP,
+                span,
+                DiagDecorator(|diag| {
+                    diag.primary_message("async mailbox request/reply helper hides an actor wait");
+                    diag.span_help(
+                        span,
+                        "return the oneshot receiver and let callers select or await explicitly",
+                    );
+                }),
+            );
+        }
     }
 }
 
@@ -135,6 +152,50 @@ fn actor_loop_kind(cx: &LateContext<'_>, body: &Body<'_>) -> Option<ActorLoopKin
     None
 }
 
+fn request_reply_helper_wait(cx: &LateContext<'_>, body: &Body<'_>) -> Option<Span> {
+    let snippet = cx
+        .sess()
+        .source_map()
+        .span_to_snippet(body.value.span)
+        .ok()?;
+
+    if !snippet.contains("oneshot::channel")
+        || !snippet.contains(".enqueue(Message::")
+        || !snippet.contains(".await")
+    {
+        return None;
+    }
+
+    AwaitVisitor::default().await_span(body.value)
+}
+
+#[derive(Default)]
+struct AwaitVisitor {
+    span: Option<Span>,
+}
+
+impl AwaitVisitor {
+    fn await_span<'tcx>(mut self, expr: &'tcx Expr<'tcx>) -> Option<Span> {
+        self.visit_expr(expr);
+        self.span
+    }
+}
+
+impl<'tcx> Visitor<'tcx> for AwaitVisitor {
+    fn visit_expr(&mut self, expr: &'tcx Expr<'tcx>) {
+        if self.span.is_some() {
+            return;
+        }
+
+        if matches!(expr.kind, ExprKind::Match(_, _, MatchSource::AwaitDesugar)) {
+            self.span = Some(expr.span);
+            return;
+        }
+
+        walk_expr(self, expr);
+    }
+}
+
 fn awaited_call<'tcx>(expr: &'tcx Expr<'tcx>) -> Option<AwaitedCall<'tcx>> {
     let ExprKind::Match(awaited, _, MatchSource::AwaitDesugar) = expr.kind else {
         return None;

.github/dylints/blocking_in_actor_loop/ui/main.rs

@@ -14,6 +14,31 @@ macro_rules! select_loop {
     }};
 }
 
+mod oneshot {
+    use std::{
+        future::Future,
+        marker::PhantomData,
+        pin::Pin,
+        task::{Context, Poll},
+    };
+
+    pub struct Sender<T>(PhantomData<T>);
+
+    pub struct Receiver<T>(PhantomData<T>);
+
+    pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
+        (Sender(PhantomData), Receiver(PhantomData))
+    }
+
+    impl<T> Future for Receiver<T> {
+        type Output = Result<T, ()>;
+
+        fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Self::Output> {
+            Poll::Pending
+        }
+    }
+}
+
 trait Handler {
     async fn process(&mut self);
 }
@@ -50,6 +75,28 @@ impl Stream {
     async fn recv(&mut self) {}
 }
 
+enum Message {
+    Ready { responder: oneshot::Sender<()> },
+}
+
+struct RequestReplyMailbox;
+
+impl RequestReplyMailbox {
+    fn enqueue(&self, _: Message) {}
+
+    async fn ready(&self) -> Option<()> {
+        let (responder, receiver) = oneshot::channel();
+        self.enqueue(Message::Ready { responder });
+        receiver.await.ok()
+    }
+
+    fn ready_receiver(&self) -> oneshot::Receiver<()> {
+        let (responder, receiver) = oneshot::channel();
+        self.enqueue(Message::Ready { responder });
+        receiver
+    }
+}
+
 async fn bad_handler(mut handler: impl Handler) {
     let context = Context;
     select_loop!(context, {

.github/dylints/blocking_in_actor_loop/ui/main.stderr

@@ -1,67 +1,79 @@
+warning: async mailbox request/reply helper hides an actor wait
+  --> $DIR/main.rs:90:9
+   |
+LL |         receiver.await.ok()
+   |         ^^^^^^^^^^^^^^
+   |
+help: return the oneshot receiver and let callers select or await explicitly
+  --> $DIR/main.rs:90:9
+   |
+LL |         receiver.await.ok()
+   |         ^^^^^^^^^^^^^^
+note: the lint level is defined here
+  --> $DIR/main.rs:3:9
+   |
+LL | #![warn(blocking_in_actor_loop)]
+   |         ^^^^^^^^^^^^^^^^^^^^^^
+
 warning: blocking await in actor loop
-  --> $DIR/main.rs:56:9
+  --> $DIR/main.rs:103:9
    |
 LL |         handler.process().await;
    |         ^^^^^^^^^^^^^^^^^^^^^^^
    |
 help: queue this work and select on its completion instead of awaiting it in the actor loop
-  --> $DIR/main.rs:56:9
+  --> $DIR/main.rs:103:9
    |
 LL |         handler.process().await;
    |         ^^^^^^^^^^^^^^^^^
-note: the lint level is defined here
-  --> $DIR/main.rs:3:9
-   |
-LL | #![warn(blocking_in_actor_loop)]
-   |         ^^^^^^^^^^^^^^^^^^^^^^
 
 warning: blocking await in actor loop
-  --> $DIR/main.rs:63:9
+  --> $DIR/main.rs:110:9
    |
 LL |         monitor.collected().await;
    |         ^^^^^^^^^^^^^^^^^^^^^^^^^
    |
 help: queue this work and select on its completion instead of awaiting it in the actor loop
-  --> $DIR/main.rs:63:9
+  --> $DIR/main.rs:110:9
    |
 LL |         monitor.collected().await;
    |         ^^^^^^^^^^^^^^^^^^^
 
 warning: blocking await in actor loop
-  --> $DIR/main.rs:70:9
+  --> $DIR/main.rs:117:9
    |
 LL |         mailbox.dial().await;
    |         ^^^^^^^^^^^^^^^^^^^^
    |
 help: queue this work and select on its completion instead of awaiting it in the actor loop
-  --> $DIR/main.rs:70:9
+  --> $DIR/main.rs:117:9
    |
 LL |         mailbox.dial().await;
    |         ^^^^^^^^^^^^^^
 
 warning: blocking await in actor loop
-  --> $DIR/main.rs:83:9
+  --> $DIR/main.rs:130:9
    |
 LL |         stream.send().await;
    |         ^^^^^^^^^^^^^^^^^^^
    |
 help: queue this work and select on its completion instead of awaiting it in the actor loop
-  --> $DIR/main.rs:83:9
+  --> $DIR/main.rs:130:9
    |
 LL |         stream.send().await;
    |         ^^^^^^^^^^^^^
 
 warning: blocking await in actor loop
-  --> $DIR/main.rs:84:9
+  --> $DIR/main.rs:131:9
    |
 LL |         stream.recv().await;
    |         ^^^^^^^^^^^^^^^^^^^
    |
 help: queue this work and select on its completion instead of awaiting it in the actor loop
-  --> $DIR/main.rs:84:9
+  --> $DIR/main.rs:131:9
    |
 LL |         stream.recv().await;
    |         ^^^^^^^^^^^^^
 
-warning: 5 warnings emitted
+warning: 6 warnings emitted
 

broadcast/src/buffered/ingress.rs

@@ -127,11 +127,11 @@ impl<P: PublicKey, M: Digestible + Codec> Mailbox<P, M> {
 
     /// Get a message by digest.
     ///
-    /// If the engine has shut down, returns `None`.
-    pub async fn get(&self, digest: M::Digest) -> Option<M> {
+    /// The returned receiver is closed if the engine has shut down.
+    pub fn get(&self, digest: M::Digest) -> oneshot::Receiver<Option<M>> {
         let (responder, receiver) = oneshot::channel();
         let _ = self.sender.enqueue(Message::Get { digest, responder });
-        receiver.await.unwrap_or_default()
+        receiver
     }
 }