Add helper to convert PortRef to a Sink (#434)

Summary:
Pull Request resolved: #434

Makes it easier to combine `PortRef` with `SinkExt` helpers.

Reviewed By: mariusae, highker, shayne-fletcher

Differential Revision: D77694800

fbshipit-source-id: 7246f78e4a1a0f0610175664c187fc00b17b077f

Author: andrewjcg

hyperactor/src/mailbox.rs

@@ -88,6 +88,7 @@ use async_trait::async_trait;
 use dashmap::DashMap;
 use dashmap::DashSet;
 use dashmap::mapref::entry::Entry;
+use futures::Sink;
 use serde::Deserialize;
 use serde::Serialize;
 use serde::de::DeserializeOwned;
@@ -108,6 +109,7 @@ use crate::actor::Signal;
 use crate::actor::remote::USER_PORT_OFFSET;
 use crate::attrs::Attrs;
 use crate::cap;
+use crate::cap::CanSend;
 use crate::channel;
 use crate::channel::ChannelAddr;
 use crate::channel::ChannelError;
@@ -970,6 +972,39 @@ impl MailboxSender for MailboxClient {
     }
 }
 
+/// Wrapper to turn `PortRef` into a `Sink`.
+pub struct PortSink<'a, C: CanSend, M: RemoteMessage> {
+    caps: &'a C,
+    port: PortRef<M>,
+}
+
+impl<'a, C: CanSend, M: RemoteMessage> PortSink<'a, C, M> {
+    /// Create new PortSink
+    pub fn new(caps: &'a C, port: PortRef<M>) -> Self {
+        Self { caps, port }
+    }
+}
+
+impl<'a, C: CanSend, M: RemoteMessage> Sink<M> for PortSink<'a, C, M> {
+    type Error = MailboxSenderError;
+
+    fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        Poll::Ready(Ok(()))
+    }
+
+    fn start_send(self: Pin<&mut Self>, item: M) -> Result<(), Self::Error> {
+        self.port.send(self.caps, item)
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        Poll::Ready(Ok(()))
+    }
+
+    fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        Poll::Ready(Ok(()))
+    }
+}
+
 /// A mailbox coordinates message delivery to actors through typed
 /// [`Port`]s associated with the mailbox.
 #[derive(Clone, Debug)]

hyperactor/src/reference.rs

@@ -51,6 +51,7 @@ use crate::cap;
 use crate::data::Serialized;
 use crate::mailbox::MailboxSenderError;
 use crate::mailbox::MailboxSenderErrorKind;
+use crate::mailbox::PortSink;
 use crate::message::Bind;
 use crate::message::Bindings;
 use crate::message::Unbind;
@@ -893,6 +894,11 @@ impl<M: RemoteMessage> PortRef<M> {
     pub fn send_serialized(&self, caps: &impl cap::CanSend, message: Serialized, headers: Attrs) {
         caps.post(self.port_id.clone(), headers, message);
     }
+
+    /// Convert this port into a sink that can be used to send messages using the given capability.
+    pub fn into_sink<'a, C: cap::CanSend>(self, caps: &'a C) -> PortSink<'a, C, M> {
+        PortSink::new(caps, self)
+    }
 }
 
 impl<M: RemoteMessage> Clone for PortRef<M> {