hyperactor/
endpoint.rs

1/*
2 * Copyright (c) Meta Platforms, Inc. and affiliates.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD-style license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9//! Generic send endpoints.
10
11use std::fmt;
12
13use hyperactor_config::Flattrs;
14use serde::Deserialize;
15use serde::Serialize;
16
17use crate::ActorAddr;
18use crate::PortAddr;
19use crate::context;
20use crate::mailbox::PortLocation;
21
22/// The logical location of an endpoint.
23#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq, typeuri::Named)]
24pub enum EndpointLocation {
25    /// An actor endpoint.
26    Actor(ActorAddr),
27    /// A port endpoint.
28    Port(PortAddr),
29    /// A local port handle that has not been bound to a routable port.
30    Local {
31        /// The actor that owns the local endpoint.
32        actor: ActorAddr,
33        /// The local endpoint's message type.
34        message_type: String,
35    },
36}
37
38impl EndpointLocation {
39    /// The actor address associated with this endpoint location.
40    pub fn actor_addr(&self) -> ActorAddr {
41        match self {
42            Self::Actor(actor) => actor.clone(),
43            Self::Port(port) => port.actor_addr(),
44            Self::Local { actor, .. } => actor.clone(),
45        }
46    }
47}
48
49impl From<PortLocation> for EndpointLocation {
50    fn from(location: PortLocation) -> Self {
51        match location {
52            PortLocation::Bound(port) => Self::Port(port),
53            PortLocation::Unbound(actor, message_type) => Self::Local {
54                actor,
55                message_type: message_type.to_string(),
56            },
57        }
58    }
59}
60
61impl fmt::Display for EndpointLocation {
62    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
63        match self {
64            Self::Actor(actor) => write!(f, "{}", actor),
65            Self::Port(port) => write!(f, "{}", port),
66            Self::Local {
67                actor,
68                message_type,
69            } => write!(f, "{}<{}>", actor, message_type),
70        }
71    }
72}
73
74/// A typed endpoint that can receive `M`.
75///
76/// This trait abstracts over local actor handles, local port handles, remote
77/// actor refs, remote port refs, and one-shot ports. It is sealed so that
78/// Hyperactor owns the post semantics for each endpoint kind.
79pub trait Endpoint<M>: crate::private::Sealed {
80    /// The logical location of this endpoint.
81    fn endpoint_location(&self) -> EndpointLocation;
82
83    /// Post `message` to this endpoint from `cx`.
84    fn post<C>(self, cx: &C, message: M)
85    where
86        C: context::Actor;
87}
88
89/// A typed endpoint that can receive `M` with message headers.
90///
91/// `RemoteEndpoint` is implemented only for endpoints whose post path preserves
92/// headers.
93pub trait RemoteEndpoint<M>: Endpoint<M> {
94    /// Post `message` and `headers` to this endpoint from `cx`.
95    fn post_with_headers<C>(self, cx: &C, headers: Flattrs, message: M)
96    where
97        C: context::Actor;
98}
99
100#[cfg(test)]
101mod tests {
102    use async_trait::async_trait;
103    use hyperactor_config::Flattrs;
104    use hyperactor_config::declare_attrs;
105    use tokio::sync::mpsc;
106    use typeuri::Named;
107
108    use super::*;
109    use crate::Actor;
110    use crate::Handler;
111    use crate::PortRef;
112    use crate::actor::Referable;
113    use crate::actor::RemoteHandles;
114    use crate::context::Mailbox as _;
115    use crate::proc::Context;
116    use crate::proc::Proc;
117
118    declare_attrs! {
119        attr ENDPOINT_TEST_HEADER: u64;
120    }
121
122    #[derive(Debug)]
123    struct EchoActor {
124        tx: PortRef<u64>,
125    }
126
127    #[async_trait]
128    impl Actor for EchoActor {}
129
130    #[async_trait]
131    impl Handler<u64> for EchoActor {
132        async fn handle(&mut self, cx: &Context<Self>, message: u64) -> anyhow::Result<()> {
133            Endpoint::post(&self.tx, cx, message);
134            Ok(())
135        }
136    }
137
138    struct TestBehavior;
139
140    impl Named for TestBehavior {
141        fn typename() -> &'static str {
142            "hyperactor::endpoint::tests::TestBehavior"
143        }
144    }
145
146    impl Referable for TestBehavior {}
147    impl RemoteHandles<u64> for TestBehavior {}
148
149    #[tokio::test]
150    async fn test_endpoint_actor_handle() {
151        let proc = Proc::isolated();
152        let client = proc.client("client");
153        let (tx, mut rx) = client.open_port();
154        let handle = proc.spawn(EchoActor { tx: tx.bind() });
155
156        Endpoint::post(&handle, &client, 123u64);
157
158        assert_eq!(rx.recv().await.expect("message should arrive"), 123);
159    }
160
161    #[tokio::test]
162    async fn test_endpoint_port_handle() {
163        let proc = Proc::isolated();
164        let client = proc.client("client");
165        let (tx, mut rx) = client.open_port();
166
167        Endpoint::post(&tx, &client, 123u64);
168
169        assert_eq!(rx.recv().await.expect("message should arrive"), 123);
170    }
171
172    #[tokio::test]
173    async fn test_endpoint_once_port_handle() {
174        let proc = Proc::isolated();
175        let client = proc.client("client");
176        let (tx, rx) = client.open_once_port();
177
178        Endpoint::post(tx, &client, 123u64);
179
180        assert_eq!(rx.recv().await.expect("message should arrive"), 123);
181    }
182
183    #[tokio::test]
184    async fn test_endpoint_actor_ref() {
185        let proc = Proc::isolated();
186        let (client, actor_ref, mut rx) = proc
187            .attach_actor::<TestBehavior, u64>("remote_actor")
188            .expect("attach actor should succeed");
189
190        Endpoint::post(&actor_ref, &client, 123u64);
191
192        assert_eq!(rx.recv().await.expect("message should arrive"), 123);
193    }
194
195    #[tokio::test]
196    async fn test_endpoint_port_ref() {
197        let proc = Proc::isolated();
198        let client = proc.client("client");
199        let (tx, mut rx) = client.open_port();
200        let port_ref = tx.bind();
201
202        Endpoint::post(&port_ref, &client, 123u64);
203
204        assert_eq!(rx.recv().await.expect("message should arrive"), 123);
205    }
206
207    #[tokio::test]
208    async fn test_endpoint_once_port_ref() {
209        let proc = Proc::isolated();
210        let client = proc.client("client");
211        let (tx, rx) = client.open_once_port();
212        let port_ref = tx.bind();
213
214        Endpoint::post(port_ref, &client, 123u64);
215
216        assert_eq!(rx.recv().await.expect("message should arrive"), 123);
217    }
218
219    #[tokio::test]
220    async fn test_remote_endpoint_headers() {
221        let proc = Proc::isolated();
222        let client = proc.client("client");
223        let (observed_tx, mut observed_rx) = mpsc::unbounded_channel();
224        let port =
225            client
226                .mailbox()
227                .open_handler_enqueue_port(move |headers: Flattrs, message: u64| {
228                    observed_tx
229                        .send((
230                            headers
231                                .get(ENDPOINT_TEST_HEADER)
232                                .expect("header should be present"),
233                            message,
234                        ))
235                        .expect("test receiver should be alive");
236                    Ok(())
237                });
238        let port_ref = port.bind();
239        let mut headers = Flattrs::new();
240        headers.set(ENDPOINT_TEST_HEADER, 456u64);
241
242        RemoteEndpoint::post_with_headers(&port_ref, &client, headers, 123u64);
243
244        assert_eq!(
245            observed_rx.recv().await.expect("message should arrive"),
246            (456, 123)
247        );
248    }
249}
hyperactor/endpoint.rs

hyperactor/
endpoint.rs
