Replace &Instance<A> with &Context<'_, A> in the Handler trait (#368)

Summary:
Pull Request resolved: #368

In order to make the headers accessible to message handlers in a cleaner way, this diff introduces a new struct `Context<'_, A>` which contains a `&Instance<A>` and the headers associated with the current message. As per T228176204, the `Handler` trait now takes a `&Context<Self>` for the `this` parameter. `Context` has the same capabilities as the contained `Instance`. Moreover, I have implemented `Deref` for `Context` to return `&Instance` -- with deref coercion, this means that no code has to change except for inside the method signatures (this deref coercion is also why the `Handler` takes `&Context` instead of `&Instance`.

Disclaimer: many of these changes were made by devmate.

Reviewed By: shayne-fletcher, mariusae

Differential Revision: D77341169

fbshipit-source-id: 4b7309b04beca43249a881d4c8a435f007caea3f

Author: Sam Lurye

controller/src/lib.rs

@@ -22,6 +22,7 @@ use async_trait::async_trait;
 use hyperactor::Actor;
 use hyperactor::ActorId;
 use hyperactor::ActorRef;
+use hyperactor::Context;
 use hyperactor::GangId;
 use hyperactor::Handler;
 use hyperactor::Named;
@@ -212,7 +213,7 @@ impl ControllerActor {
     // M = self.worker_progress_check_interval
     async fn request_status_if_needed(
         &mut self,
-        this: &hyperactor::Instance<Self>,
+        this: &Context<'_, Self>,
     ) -> Result<(), anyhow::Error> {
         if let Some((expected_seq, ..)) = self.history.deadline(
             self.operations_per_worker_progress_request,
@@ -259,7 +260,7 @@ struct CheckWorkerProgress;
 impl Handler<CheckWorkerProgress> for ControllerActor {
     async fn handle(
         &mut self,
-        this: &hyperactor::Instance<Self>,
+        this: &Context<Self>,
         _check_worker_progress: CheckWorkerProgress,
     ) -> Result<(), anyhow::Error> {
         let client = self.client()?;
@@ -359,7 +360,7 @@ fn slice_to_selection(slice: Slice) -> Selection {
 impl ControllerMessageHandler for ControllerActor {
     async fn attach(
         &mut self,
-        this: &hyperactor::Instance<Self>,
+        this: &Context<Self>,
         client_actor: ActorRef<ClientActor>,
     ) -> Result<(), anyhow::Error> {
         tracing::debug!("attaching client actor {}", client_actor);
@@ -378,7 +379,7 @@ impl ControllerMessageHandler for ControllerActor {
 
     async fn node(
         &mut self,
-        this: &hyperactor::Instance<Self>,
+        this: &Context<Self>,
         seq: Seq,
         defs: Vec<Ref>,
         uses: Vec<Ref>,
@@ -395,7 +396,7 @@ impl ControllerMessageHandler for ControllerActor {
 
     async fn drop_refs(
         &mut self,
-        _this: &hyperactor::Instance<Self>,
+        _this: &Context<Self>,
         refs: Vec<Ref>,
     ) -> Result<(), anyhow::Error> {
         self.history.delete_invocations_for_refs(refs);
@@ -404,7 +405,7 @@ impl ControllerMessageHandler for ControllerActor {
 
     async fn send(
         &mut self,
-        this: &hyperactor::Instance<Self>,
+        this: &Context<Self>,
         ranks: Ranks,
         message: Serialized,
     ) -> Result<(), anyhow::Error> {
@@ -456,7 +457,7 @@ impl ControllerMessageHandler for ControllerActor {
 
     async fn remote_function_failed(
         &mut self,
-        this: &hyperactor::Instance<Self>,
+        this: &Context<Self>,
         seq: Seq,
         error: WorkerError,
     ) -> Result<(), anyhow::Error> {
@@ -469,7 +470,7 @@ impl ControllerMessageHandler for ControllerActor {
 
     async fn status(
         &mut self,
-        _this: &hyperactor::Instance<Self>,
+        _this: &Context<Self>,
         seq: Seq,
         worker_actor_id: ActorId,
         controller: bool,
@@ -486,18 +487,15 @@ impl ControllerMessageHandler for ControllerActor {
 
     async fn fetch_result(
         &mut self,
-        _this: &hyperactor::Instance<Self>,
+        _this: &Context<Self>,
         seq: Seq,
         result: Result<Serialized, WorkerError>,
     ) -> Result<(), anyhow::Error> {
         self.history.set_result(seq, result);
         Ok(())
     }
 
-    async fn check_supervision(
-        &mut self,
-        this: &hyperactor::Instance<Self>,
-    ) -> Result<(), anyhow::Error> {
+    async fn check_supervision(&mut self, this: &Context<Self>) -> Result<(), anyhow::Error> {
         let gang_id: GangId = self.worker_gang_ref.clone().into();
         let world_state = self
             .system_supervision_actor_ref
@@ -565,7 +563,7 @@ impl ControllerMessageHandler for ControllerActor {
 
     async fn debugger_message(
         &mut self,
-        this: &hyperactor::Instance<Self>,
+        this: &Context<Self>,
         debugger_actor_id: ActorId,
         action: DebuggerAction,
     ) -> Result<(), anyhow::Error> {
@@ -577,15 +575,15 @@ impl ControllerMessageHandler for ControllerActor {
     #[cfg(test)]
     async fn get_first_incomplete_seqs_unit_tests_only(
         &mut self,
-        _this: &hyperactor::Instance<Self>,
+        _this: &Context<Self>,
     ) -> Result<Vec<Seq>, anyhow::Error> {
         Ok(self.history.first_incomplete_seqs().to_vec())
     }
 
     #[cfg(not(test))]
     async fn get_first_incomplete_seqs_unit_tests_only(
         &mut self,
-        _this: &hyperactor::Instance<Self>,
+        _this: &Context<Self>,
     ) -> Result<Vec<Seq>, anyhow::Error> {
         unimplemented!("get_first_incomplete_seqs_unit_tests_only is only for unit tests")
     }
@@ -616,7 +614,6 @@ mod tests {
 
     use hyperactor::HandleClient;
     use hyperactor::Handler;
-    use hyperactor::Instance;
     use hyperactor::RefClient;
     use hyperactor::channel;
     use hyperactor::channel::ChannelTransport;
@@ -1801,7 +1798,7 @@ mod tests {
     impl PanickingMessageHandler for PanickingActor {
         async fn panic(
             &mut self,
-            _this: &Instance<Self>,
+            _this: &Context<Self>,
             err_msg: String,
         ) -> Result<(), anyhow::Error> {
             panic!("{}", err_msg);

hyper/src/commands/demo.rs

@@ -12,6 +12,7 @@ use std::time::Duration;
 use async_trait::async_trait;
 use clap::Subcommand;
 use hyperactor::Actor;
+use hyperactor::Context;
 use hyperactor::HandleClient;
 use hyperactor::Handler;
 use hyperactor::Named;
@@ -229,40 +230,29 @@ impl Actor for DemoActor {
 impl DemoMessageHandler for DemoActor {
     async fn echo(
         &mut self,
-        _this: &hyperactor::Instance<Self>,
+        _this: &Context<Self>,
         message: String,
     ) -> Result<String, anyhow::Error> {
         tracing::info!("demo: message: {}", message);
         Ok(message)
     }
 
-    async fn increment(
-        &mut self,
-        _this: &hyperactor::Instance<Self>,
-        num: u64,
-    ) -> Result<u64, anyhow::Error> {
+    async fn increment(&mut self, _this: &Context<Self>, num: u64) -> Result<u64, anyhow::Error> {
         tracing::info!("demo: increment: {}", num);
         Ok(num + 1)
     }
 
-    async fn panic(&mut self, _this: &hyperactor::Instance<Self>) -> Result<(), anyhow::Error> {
+    async fn panic(&mut self, _this: &Context<Self>) -> Result<(), anyhow::Error> {
         tracing::info!("demo: panic!");
         panic!()
     }
 
-    async fn spawn_child(
-        &mut self,
-        this: &hyperactor::Instance<Self>,
-    ) -> Result<ActorRef<Self>, anyhow::Error> {
+    async fn spawn_child(&mut self, this: &Context<Self>) -> Result<ActorRef<Self>, anyhow::Error> {
         tracing::info!("demo: spawn child");
         Ok(Self::spawn(this, ()).await?.bind())
     }
 
-    async fn error(
-        &mut self,
-        _this: &hyperactor::Instance<Self>,
-        message: String,
-    ) -> Result<(), anyhow::Error> {
+    async fn error(&mut self, _this: &Context<Self>, message: String) -> Result<(), anyhow::Error> {
         tracing::info!("demo: message: {}", message);
         anyhow::bail!("{}", message)
     }

hyperactor/example/derive.rs

@@ -11,9 +11,9 @@ use std::time::Duration;
 
 use async_trait::async_trait;
 use hyperactor::Actor;
+use hyperactor::Context;
 use hyperactor::HandleClient;
 use hyperactor::Handler;
-use hyperactor::Instance;
 use hyperactor::Named;
 use hyperactor::OncePortRef;
 use hyperactor::RefClient;
@@ -68,27 +68,23 @@ impl Actor for ShoppingListActor {
 #[async_trait]
 #[hyperactor::forward(ShoppingList)]
 impl ShoppingListHandler for ShoppingListActor {
-    async fn add(&mut self, _this: &Instance<Self>, item: String) -> Result<(), anyhow::Error> {
+    async fn add(&mut self, _this: &Context<Self>, item: String) -> Result<(), anyhow::Error> {
         eprintln!("insert {}", item);
         self.0.insert(item);
         Ok(())
     }
 
-    async fn remove(&mut self, _this: &Instance<Self>, item: String) -> Result<(), anyhow::Error> {
+    async fn remove(&mut self, _this: &Context<Self>, item: String) -> Result<(), anyhow::Error> {
         eprintln!("remove {}", item);
         self.0.remove(&item);
         Ok(())
     }
 
-    async fn exists(
-        &mut self,
-        _this: &Instance<Self>,
-        item: String,
-    ) -> Result<bool, anyhow::Error> {
+    async fn exists(&mut self, _this: &Context<Self>, item: String) -> Result<bool, anyhow::Error> {
         Ok(self.0.contains(&item))
     }
 
-    async fn list(&mut self, _this: &Instance<Self>) -> Result<Vec<String>, anyhow::Error> {
+    async fn list(&mut self, _this: &Context<Self>) -> Result<Vec<String>, anyhow::Error> {
         Ok(self.0.iter().cloned().collect())
     }
 }

hyperactor/src/actor.rs

@@ -49,6 +49,7 @@ use crate::mailbox::UndeliverableMessageError;
 use crate::mailbox::log::MessageLogError;
 use crate::message::Castable;
 use crate::message::IndexedErasedUnbound;
+use crate::proc::Context;
 use crate::proc::Instance;
 use crate::proc::InstanceCell;
 use crate::proc::Ports;
@@ -142,7 +143,7 @@ pub trait Actor: Sized + Send + Debug + 'static {
 #[async_trait]
 pub trait Handler<M>: Actor {
     /// Handle the next M-typed message.
-    async fn handle(&mut self, this: &Instance<Self>, message: M) -> Result<(), anyhow::Error>;
+    async fn handle(&mut self, this: &Context<Self>, message: M) -> Result<(), anyhow::Error>;
 }
 
 /// We provide this handler to indicate that actors can handle the [`Signal`] message.
@@ -151,7 +152,7 @@ pub trait Handler<M>: Actor {
 impl<A: Actor> Handler<Signal> for A {
     async fn handle(
         &mut self,
-        _this: &Instance<Self>,
+        _this: &Context<Self>,
         _message: Signal,
     ) -> Result<(), anyhow::Error> {
         unimplemented!("signal handler should not be called directly")
@@ -164,7 +165,7 @@ impl<A: Actor> Handler<Signal> for A {
 impl<A: Actor> Handler<Undeliverable<MessageEnvelope>> for A {
     async fn handle(
         &mut self,
-        this: &Instance<Self>,
+        this: &Context<Self>,
         message: Undeliverable<MessageEnvelope>,
     ) -> Result<(), anyhow::Error> {
         self.handle_undeliverable_message(this, message).await
@@ -181,7 +182,7 @@ where
 {
     async fn handle(
         &mut self,
-        this: &Instance<Self>,
+        this: &Context<Self>,
         msg: IndexedErasedUnbound<M>,
     ) -> anyhow::Result<()> {
         let message = msg.downcast()?.bind()?;
@@ -712,7 +713,7 @@ mod tests {
     impl Handler<u64> for EchoActor {
         async fn handle(
             &mut self,
-            this: &Instance<Self>,
+            this: &Context<Self>,
             message: u64,
         ) -> Result<(), anyhow::Error> {
             let Self(port) = self;
@@ -818,7 +819,7 @@ mod tests {
     impl Handler<OncePortHandle<bool>> for InitActor {
         async fn handle(
             &mut self,
-            _this: &Instance<Self>,
+            _this: &Context<Self>,
             port: OncePortHandle<bool>,
         ) -> Result<(), anyhow::Error> {
             port.send(self.0)?;
@@ -861,7 +862,7 @@ mod tests {
 
     #[async_trait]
     impl Handler<u64> for CheckpointActor {
-        async fn handle(&mut self, this: &Instance<Self>, value: u64) -> Result<(), anyhow::Error> {
+        async fn handle(&mut self, this: &Context<Self>, value: u64) -> Result<(), anyhow::Error> {
             self.sum += value;
             self.port.send(this, self.sum)?;
             Ok(())
@@ -946,7 +947,7 @@ mod tests {
     impl Handler<u64> for MultiActor {
         async fn handle(
             &mut self,
-            _this: &Instance<Self>,
+            _this: &Context<Self>,
             message: u64,
         ) -> Result<(), anyhow::Error> {
             let mut vals = self.0.lock().unwrap();
@@ -959,7 +960,7 @@ mod tests {
     impl Handler<String> for MultiActor {
         async fn handle(
             &mut self,
-            _this: &Instance<Self>,
+            _this: &Context<Self>,
             message: String,
         ) -> Result<(), anyhow::Error> {
             let mut vals = self.0.lock().unwrap();
@@ -972,7 +973,7 @@ mod tests {
     impl Handler<OncePortHandle<bool>> for MultiActor {
         async fn handle(
             &mut self,
-            _this: &Instance<Self>,
+            _this: &Context<Self>,
             message: OncePortHandle<bool>,
         ) -> Result<(), anyhow::Error> {
             message.send(true).unwrap();

hyperactor/src/actor/remote.rs

@@ -139,8 +139,8 @@ mod tests {
 
     use super::*;
     use crate as hyperactor; // for macros
+    use crate::Context;
     use crate::Handler;
-    use crate::Instance;
 
     #[derive(Debug)]
     #[hyperactor::export(handlers = [()])]
@@ -161,7 +161,7 @@ mod tests {
 
     #[async_trait]
     impl Handler<()> for MyActor {
-        async fn handle(&mut self, _this: &Instance<Self>, _message: ()) -> anyhow::Result<()> {
+        async fn handle(&mut self, _this: &Context<Self>, _message: ()) -> anyhow::Result<()> {
             unimplemented!()
         }
     }

hyperactor/src/lib.rs

@@ -149,6 +149,7 @@ pub use mailbox::RemoteMessage;
 pub use opentelemetry;
 #[doc(hidden)]
 pub use paste::paste;
+pub use proc::Context;
 pub use proc::Instance;
 pub use reference::ActorId;
 pub use reference::ActorRef;