hyperactor_remote/
keepalive.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//! Message-based link implementation.
10//!
11//! The link has two actors. `KeepaliveWorker` runs on the worker side and owns
12//! the outbound keepalive stream. `KeepaliveSupervisor` runs on the supervisor
13//! side and acknowledges each keepalive. The worker sends [`Keepalive`] with a
14//! reply port for the worker's [`KeepaliveAck`] handler. The supervisor answers
15//! with [`KeepaliveAck`] on that reply port. Either side may fail the link: the
16//! worker fails if a keepalive is not acknowledged within `timeout`, and the
17//! supervisor fails if no newer keepalive is delivered within `timeout`.
18//!
19//! Keepalives carry monotonically increasing generation numbers. A generation
20//! identifies one worker-issued keepalive and the local timers derived from it.
21//! We use generations instead of canceling timers: if a later generation has
22//! already been observed, a delayed timer for an older generation is stale and
23//! is ignored. This relies on Hyperactor's in-order delivery guarantee for
24//! messages sent along a single actor-to-actor path. In particular, the
25//! supervisor observes keepalives from one worker in generation order, and each
26//! actor observes its own scheduled control messages in the order in which they
27//! become deliverable. The generation checks make late timers harmless, while
28//! ordered delivery preserves the meaning of "no newer generation arrived before
29//! this deadline."
30//!
31//! The worker uses two internal self messages. [`SendKeepalive`] advances the
32//! stream after `interval` if its generation is still current. [`AckDeadline`]
33//! checks after `timeout` whether the corresponding generation was acknowledged.
34//! The supervisor uses [`Deadline`] to check whether a newer generation arrived
35//! before its `timeout`.
36
37use std::time::Duration;
38
39use async_trait::async_trait;
40use hyperactor::Actor;
41use hyperactor::ActorEnvironment;
42use hyperactor::ActorHandle;
43use hyperactor::Context;
44use hyperactor::Endpoint as _;
45use hyperactor::HandleClient;
46use hyperactor::Handler;
47use hyperactor::Instance;
48use hyperactor::OncePortRef;
49use hyperactor::PortRef;
50use hyperactor::RefClient;
51use hyperactor::RemoteSpawn;
52use hyperactor::Uid;
53use hyperactor::context;
54use serde::Deserialize;
55use serde::Serialize;
56use typeuri::Named;
57
58use crate::LinkSpec;
59
60/// Keepalive timing parameters.
61#[derive(Clone, Debug, Serialize, Deserialize, Named, PartialEq, Eq)]
62pub struct KeepaliveParams {
63    /// Duration between keepalive messages sent by the worker side.
64    pub interval: Duration,
65    /// Maximum duration to wait for keepalive delivery and acknowledgment.
66    ///
67    /// The worker side waits this long for each acknowledgment, and the
68    /// supervisor side uses this as the maximum accepted gap between
69    /// keepalives.
70    pub timeout: Duration,
71}
72wirevalue::register_type!(KeepaliveParams);
73
74impl KeepaliveParams {
75    /// Create keepalive timing parameters.
76    ///
77    /// The worker sends keepalives every `interval`. Each keepalive must be
78    /// acknowledged within `timeout`, and the supervisor side also uses
79    /// `timeout` as the maximum accepted gap between keepalives.
80    pub fn new(interval: Duration, timeout: Duration) -> Self {
81        Self { interval, timeout }
82    }
83}
84
85/// Keepalive link configuration.
86#[derive(Clone, Debug, PartialEq, Eq)]
87pub struct KeepaliveLink(KeepaliveParams);
88
89impl Default for KeepaliveLink {
90    /// A 60-second keepalive interval and 60-second timeout.
91    fn default() -> Self {
92        Self::new(Duration::from_secs(60), Duration::from_secs(60))
93    }
94}
95
96impl KeepaliveLink {
97    /// Create a keepalive link configuration.
98    ///
99    /// The worker sends keepalives every `interval`. Each keepalive must be
100    /// acknowledged within `timeout`, and the supervisor side also uses
101    /// `timeout` as the maximum accepted gap between keepalives.
102    pub fn new(interval: Duration, timeout: Duration) -> Self {
103        Self(KeepaliveParams::new(interval, timeout))
104    }
105
106    /// Create a keepalive link configuration from timing parameters.
107    pub fn from_params(params: KeepaliveParams) -> Self {
108        Self(params)
109    }
110
111    /// Spawn the supervisor side and return the worker-side link spec.
112    pub fn spawn_supervisor<C: context::Actor>(
113        self,
114        this: &C,
115    ) -> anyhow::Result<(ActorHandle<KeepaliveSupervisor>, LinkSpec)> {
116        self.spawn_supervisor_uid(this, Uid::anonymous())
117    }
118
119    /// Spawn the supervisor side and return the worker-side link spec,
120    /// used to instantiate the worker side of the link.
121    ///
122    /// The passed uid determines the uid of the worker side implementation
123    /// actor; it is specified by the supervisor to enable efficient group-style
124    /// supervision.
125    pub fn spawn_supervisor_uid<C: context::Actor>(
126        self,
127        this: &C,
128        uid: Uid,
129    ) -> anyhow::Result<(ActorHandle<KeepaliveSupervisor>, LinkSpec)> {
130        let params = self.0;
131        let supervisor =
132            this.instance()
133                .spawn(KeepaliveSupervisor::new(KeepaliveSupervisorParams::new(
134                    params.timeout,
135                )));
136        let worker = KeepaliveWorkerParams::new(supervisor.port::<Keepalive>().bind(), params)
137            .link_spec_uid(uid)?;
138        Ok((supervisor, worker))
139    }
140}
141
142/// Parameters for [`KeepaliveWorker`].
143#[derive(Clone, Debug, Serialize, Deserialize, Named, PartialEq, Eq)]
144pub struct KeepaliveWorkerParams {
145    /// Supervisor-side keepalive handler port.
146    pub supervisor: PortRef<Keepalive>,
147    /// Keepalive timing parameters.
148    pub keepalive: KeepaliveParams,
149}
150wirevalue::register_type!(KeepaliveWorkerParams);
151
152impl KeepaliveWorkerParams {
153    /// Create keepalive worker parameters.
154    pub fn new(supervisor: PortRef<Keepalive>, keepalive: KeepaliveParams) -> Self {
155        Self {
156            supervisor,
157            keepalive,
158        }
159    }
160
161    /// Create a worker-side link spec for a keepalive worker actor with a fresh uid.
162    pub fn link_spec(self) -> anyhow::Result<LinkSpec> {
163        self.link_spec_uid(Uid::anonymous())
164    }
165
166    /// Create a worker-side link spec for a keepalive worker actor with an explicit uid.
167    pub fn link_spec_uid(self, uid: Uid) -> anyhow::Result<LinkSpec> {
168        let params = bincode::serde::encode_to_vec(self, bincode::config::legacy())?;
169        Ok(LinkSpec::for_actor_uid::<KeepaliveWorker>(uid, params))
170    }
171}
172
173/// Parameters for [`KeepaliveSupervisor`].
174#[derive(Clone, Debug, Serialize, Deserialize, Named, PartialEq, Eq)]
175pub struct KeepaliveSupervisorParams {
176    /// Maximum duration between accepted keepalive messages.
177    pub timeout: Duration,
178}
179wirevalue::register_type!(KeepaliveSupervisorParams);
180
181impl KeepaliveSupervisorParams {
182    /// Create keepalive supervisor parameters.
183    pub fn new(timeout: Duration) -> Self {
184        Self { timeout }
185    }
186}
187
188/// Keepalive request sent to a [`KeepaliveSupervisor`].
189#[derive(
190    Clone,
191    Debug,
192    Serialize,
193    Deserialize,
194    Named,
195    Handler,
196    HandleClient,
197    RefClient
198)]
199pub struct Keepalive {
200    /// Worker-issued keepalive generation.
201    pub generation: u64,
202    /// Reply port that receives the keepalive acknowledgment.
203    pub reply: OncePortRef<KeepaliveAck>,
204}
205wirevalue::register_type!(Keepalive);
206
207/// Acknowledgment for an accepted keepalive.
208#[derive(Clone, Debug, Serialize, Deserialize, Named, PartialEq, Eq)]
209pub struct KeepaliveAck {
210    /// Generation accepted by the link actor.
211    pub generation: u64,
212}
213wirevalue::register_type!(KeepaliveAck);
214
215/// Supervisor-side deadline for receiving a newer keepalive.
216///
217/// `generation` is the supervisor's current generation when the deadline is
218/// scheduled. If the supervisor still has the same current generation when this
219/// message is handled, no newer keepalive arrived before the deadline, and the
220/// supervisor fails the link. If the current generation is larger, this message
221/// is stale.
222#[derive(Clone, Debug, Serialize, Deserialize, Named)]
223struct Deadline {
224    generation: u64,
225}
226wirevalue::register_type!(Deadline);
227
228/// Worker-side timer for sending the next keepalive.
229///
230/// `generation` is the keepalive generation that scheduled this timer. The
231/// worker sends the next keepalive only if this generation is still current. If
232/// the worker has already advanced, this message is stale.
233#[derive(Clone, Debug, Serialize, Deserialize, Named)]
234struct SendKeepalive {
235    generation: u64,
236}
237wirevalue::register_type!(SendKeepalive);
238
239/// Worker-side deadline for receiving a keepalive acknowledgment.
240///
241/// `generation` is the keepalive generation that scheduled this deadline. If
242/// the worker has not recorded an acknowledgment for this generation, or any
243/// later generation, when the message is handled, the worker fails the link. If
244/// the acknowledged generation is equal or larger, this deadline is satisfied.
245#[derive(Clone, Debug, Serialize, Deserialize, Named)]
246struct AckDeadline {
247    generation: u64,
248}
249wirevalue::register_type!(AckDeadline);
250
251/// Worker-side keepalive link actor.
252///
253/// `KeepaliveWorker` owns the generation sequence and handles supervisor
254/// [`KeepaliveAck`] messages directly. On startup, the worker sends generation
255/// 1 and schedules [`SendKeepalive`] and [`AckDeadline`] for that generation.
256/// Each live [`SendKeepalive`] advances to the next generation, sends a new
257/// [`Keepalive`], and schedules the next pair of timers. Each [`KeepaliveAck`]
258/// records progress. Each [`AckDeadline`] proves that the matching generation
259/// failed to receive an acknowledgment before `timeout`, unless a same-or-newer
260/// acknowledgment has already been recorded.
261#[derive(Debug)]
262#[hyperactor::export]
263pub struct KeepaliveWorker {
264    supervisor: PortRef<Keepalive>,
265    interval: Duration,
266    timeout: Duration,
267    generation: u64,
268    acked_generation: u64,
269}
270
271#[async_trait]
272impl Actor for KeepaliveWorker {
273    async fn init(&mut self, this: &Instance<Self>) -> anyhow::Result<()> {
274        self.send_keepalive(this)
275    }
276}
277
278#[async_trait]
279impl RemoteSpawn for KeepaliveWorker {
280    type Params = KeepaliveWorkerParams;
281
282    async fn new(
283        params: KeepaliveWorkerParams,
284        _environment: &ActorEnvironment,
285    ) -> anyhow::Result<Self> {
286        Ok(Self {
287            supervisor: params.supervisor,
288            interval: params.keepalive.interval,
289            timeout: params.keepalive.timeout,
290            generation: 0,
291            acked_generation: 0,
292        })
293    }
294}
295
296#[async_trait]
297impl Handler<SendKeepalive> for KeepaliveWorker {
298    async fn handle(&mut self, cx: &Context<Self>, message: SendKeepalive) -> anyhow::Result<()> {
299        if message.generation == self.generation {
300            self.send_keepalive(cx)?;
301        }
302        Ok(())
303    }
304}
305
306#[async_trait]
307impl Handler<KeepaliveAck> for KeepaliveWorker {
308    async fn handle(&mut self, _cx: &Context<Self>, ack: KeepaliveAck) -> anyhow::Result<()> {
309        self.acked_generation = self.acked_generation.max(ack.generation);
310        Ok(())
311    }
312}
313
314#[async_trait]
315impl Handler<AckDeadline> for KeepaliveWorker {
316    async fn handle(&mut self, cx: &Context<Self>, message: AckDeadline) -> anyhow::Result<()> {
317        if self.acked_generation < message.generation {
318            cx.kill("keepalive acknowledgment missed")?;
319        }
320        Ok(())
321    }
322}
323
324impl KeepaliveWorker {
325    fn send_keepalive(&mut self, this: &Instance<Self>) -> anyhow::Result<()> {
326        self.generation += 1;
327        let generation = self.generation;
328        (&self.supervisor).post(
329            this,
330            Keepalive {
331                generation,
332                reply: this.port::<KeepaliveAck>().bind().into_once(),
333            },
334        );
335
336        this.post_after(this, SendKeepalive { generation }, self.interval);
337        this.post_after(this, AckDeadline { generation }, self.timeout);
338        Ok(())
339    }
340}
341
342/// Supervisor-side keepalive link actor.
343///
344/// `KeepaliveSupervisor` records the maximum keepalive generation it has
345/// accepted and acknowledges every [`Keepalive`] with the same generation. Each
346/// accepted keepalive schedules a [`Deadline`] for the current generation. A
347/// deadline is active only while its generation remains current; if a newer
348/// keepalive has arrived, the deadline is stale.
349#[derive(Debug)]
350pub struct KeepaliveSupervisor {
351    timeout: Duration,
352    generation: u64,
353}
354
355#[async_trait]
356impl Actor for KeepaliveSupervisor {
357    async fn init(&mut self, this: &Instance<Self>) -> anyhow::Result<()> {
358        self.schedule_deadline(this)
359    }
360}
361
362#[async_trait]
363impl Handler<Keepalive> for KeepaliveSupervisor {
364    async fn handle(&mut self, cx: &Context<Self>, message: Keepalive) -> anyhow::Result<()> {
365        self.generation = self.generation.max(message.generation);
366        message.reply.post(
367            cx,
368            KeepaliveAck {
369                generation: message.generation,
370            },
371        );
372        self.schedule_deadline(cx)
373    }
374}
375
376#[async_trait]
377impl Handler<Deadline> for KeepaliveSupervisor {
378    async fn handle(&mut self, cx: &Context<Self>, message: Deadline) -> anyhow::Result<()> {
379        // A deadline is scheduled for the generation that was current at the
380        // time. Equality means no newer keepalive arrived before the deadline.
381        // A smaller deadline generation is stale. A larger deadline generation
382        // cannot be produced by this actor before it observes that generation.
383        if message.generation == self.generation {
384            let reason = format!("keepalive missed for generation {}", message.generation);
385            cx.kill(&reason)?;
386        }
387        Ok(())
388    }
389}
390
391impl KeepaliveSupervisor {
392    /// Create a supervisor-side keepalive actor.
393    pub fn new(params: KeepaliveSupervisorParams) -> Self {
394        Self {
395            timeout: params.timeout,
396            generation: 0,
397        }
398    }
399
400    fn schedule_deadline(&self, this: &Instance<Self>) -> anyhow::Result<()> {
401        this.post_after(
402            this,
403            Deadline {
404                generation: self.generation,
405            },
406            self.timeout,
407        );
408        Ok(())
409    }
410}
411
412hyperactor::register_spawnable!(KeepaliveWorker);
413
414#[cfg(test)]
415mod tests {
416    use hyperactor::Actor;
417    use hyperactor::PortRef;
418    use hyperactor::Proc;
419    use hyperactor::actor::ActorErrorKind;
420    use hyperactor::actor::ActorStatus;
421    use hyperactor::supervision::ActorSupervisionEvent;
422
423    use super::*;
424
425    #[derive(Debug)]
426    enum ParentSpawn {
427        Link(LinkSpec),
428        Supervisor(KeepaliveSupervisor),
429    }
430
431    #[derive(Debug)]
432    struct ParentActor {
433        spawn: Option<ParentSpawn>,
434        events: PortRef<ActorSupervisionEvent>,
435    }
436
437    #[async_trait]
438    impl Actor for ParentActor {
439        async fn init(&mut self, this: &Instance<Self>) -> anyhow::Result<()> {
440            match self.spawn.take().unwrap() {
441                ParentSpawn::Link(link) => {
442                    link.spawn_worker(this).await?;
443                }
444                ParentSpawn::Supervisor(supervisor) => {
445                    this.spawn(supervisor);
446                }
447            }
448            Ok(())
449        }
450
451        async fn handle_supervision_event(
452            &mut self,
453            this: &Instance<Self>,
454            event: &ActorSupervisionEvent,
455        ) -> anyhow::Result<bool> {
456            self.events.post(this, event.clone());
457            Ok(true)
458        }
459    }
460
461    #[derive(Debug)]
462    #[hyperactor::export(Keepalive)]
463    struct SilentSupervisor;
464
465    #[async_trait]
466    impl Actor for SilentSupervisor {}
467
468    #[async_trait]
469    impl Handler<Keepalive> for SilentSupervisor {
470        async fn handle(&mut self, _cx: &Context<Self>, _message: Keepalive) -> anyhow::Result<()> {
471            Ok(())
472        }
473    }
474
475    #[tokio::test]
476    async fn test_keepalive_supervisor_replies_to_keepalive() {
477        let proc = Proc::isolated();
478        let parent = proc.client("parent");
479        let supervisor = parent.spawn(KeepaliveSupervisor::new(KeepaliveSupervisorParams::new(
480            Duration::from_secs(60),
481        )));
482        let (reply, ack_rx) = parent.open_once_port::<KeepaliveAck>();
483
484        supervisor.post(
485            &parent,
486            Keepalive {
487                generation: 41,
488                reply: reply.bind(),
489            },
490        );
491
492        let ack = ack_rx.recv().await.unwrap();
493        assert_eq!(ack.generation, 41);
494
495        supervisor.stop("test").unwrap();
496        supervisor.await;
497    }
498
499    #[tokio::test]
500    async fn test_keepalive_supervisor_failure_propagates_to_parent() {
501        let proc = Proc::isolated();
502        let client = proc.client("client");
503        let (events, mut event_rx) = client.open_port::<ActorSupervisionEvent>();
504        let supervisor =
505            KeepaliveSupervisor::new(KeepaliveSupervisorParams::new(Duration::from_millis(10)));
506        let parent: ActorHandle<ParentActor> = proc.spawn(ParentActor {
507            spawn: Some(ParentSpawn::Supervisor(supervisor)),
508            events: events.bind(),
509        });
510
511        let event = event_rx.recv().await.unwrap();
512
513        assert!(matches!(
514            event.actor_status,
515            ActorStatus::Failed(ActorErrorKind::Generic(ref reason))
516                if reason == "actor explicitly aborted due to: keepalive missed for generation 0"
517        ));
518        assert!(!event.actor_id.is_root());
519
520        parent.stop("test").unwrap();
521        parent.await;
522    }
523
524    #[tokio::test]
525    async fn test_keepalive_worker_sends_keepalives() {
526        let proc = Proc::isolated();
527        let parent = proc.client("parent");
528        let supervisor = parent.spawn(KeepaliveSupervisor::new(KeepaliveSupervisorParams::new(
529            Duration::from_secs(60),
530        )));
531        let worker = KeepaliveWorkerParams::new(
532            supervisor.port::<Keepalive>().bind(),
533            KeepaliveParams::new(Duration::from_millis(5), Duration::from_millis(50)),
534        )
535        .link_spec()
536        .unwrap()
537        .spawn_worker(&parent)
538        .await
539        .unwrap()
540        .downcast::<KeepaliveWorker>()
541        .unwrap();
542        let status = worker.status();
543
544        tokio::time::sleep(Duration::from_millis(30)).await;
545
546        assert!(!status.borrow().is_terminal());
547
548        worker.stop("test").unwrap();
549        worker.await;
550        supervisor.stop("test").unwrap();
551        supervisor.await;
552    }
553
554    #[tokio::test]
555    async fn test_keepalive_worker_failure_propagates_to_parent() {
556        let proc = Proc::isolated();
557        let client = proc.client("client");
558        let (events, mut event_rx) = client.open_port::<ActorSupervisionEvent>();
559        let supervisor = proc.spawn(SilentSupervisor);
560        let uid = Uid::anonymous();
561        let link = KeepaliveWorkerParams::new(
562            supervisor.port::<Keepalive>().bind(),
563            KeepaliveParams::new(Duration::from_millis(100), Duration::from_millis(10)),
564        )
565        .link_spec_uid(uid.clone())
566        .unwrap();
567        let parent: ActorHandle<ParentActor> = proc.spawn(ParentActor {
568            spawn: Some(ParentSpawn::Link(link)),
569            events: events.bind(),
570        });
571
572        let event = event_rx.recv().await.unwrap();
573
574        assert_eq!(event.actor_id.uid(), &uid);
575        assert!(matches!(
576            event.actor_status,
577            ActorStatus::Failed(ActorErrorKind::Generic(ref reason))
578                if reason == "actor explicitly aborted due to: keepalive acknowledgment missed"
579        ));
580
581        parent.stop("test").unwrap();
582        parent.await;
583        supervisor.stop("test").unwrap();
584        supervisor.await;
585    }
586
587    #[tokio::test]
588    async fn test_keepalive_link_spawn_mints_shared_worker_spec() {
589        let proc = Proc::isolated();
590        let parent = proc.client("parent");
591
592        let (supervisor, worker_link) =
593            KeepaliveLink::new(Duration::from_secs(60), Duration::from_secs(60))
594                .spawn_supervisor(&parent)
595                .unwrap();
596        let mut supervisor_status = supervisor.status();
597        supervisor_status
598            .wait_for(|status| matches!(status, ActorStatus::Idle))
599            .await
600            .unwrap();
601        let worker_uid = worker_link.uid().clone();
602        let remote_worker = worker_link.spawn_worker(&parent).await.unwrap();
603
604        assert_eq!(remote_worker.actor_id().uid(), &worker_uid);
605
606        supervisor_status.changed().await.unwrap();
607        supervisor_status
608            .wait_for(|status| matches!(status, ActorStatus::Idle))
609            .await
610            .unwrap();
611        remote_worker.drain_and_stop("test").unwrap();
612        remote_worker.await;
613        supervisor.stop("test").unwrap();
614        supervisor.await;
615    }
616}
hyperactor_remote/keepalive.rs

hyperactor_remote/
keepalive.rs
