1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

//! A supervisor actor that catches panics in supervised actors and
//! allows for them to restart
use crate::actor::{Actor, ActorHandle};
use crate::context::Context;
use crate::handler::Handler;
use crate::mailbox::{Addr, Mailbox, Message};
use async_trait::async_trait;
use futures::channel::oneshot;
use std::any::Any;

/// Supervisor is an actor that can catch panics in other actors.
///
/// To supervisor an actor, start a supervisor and then pass the structure to the supervisor with
/// the message `SuperviseActor(...)`,
#[derive(Default, Debug)]
pub struct Supervisor;

impl Supervisor {
    pub fn new() -> Self {
        Self
    }

    pub async fn supervise<A>(&mut self, mut actor: A) -> (Addr<A>, ActorHandle)
    where
        A: 'static + Supervised,
    {
        let (addr_tx, addr_rx) = oneshot::channel();
        let task = tokio::spawn(async move {
            let mailbox = Mailbox::new(16);

            let addr = mailbox.new_addr();
            let (mut actor_ctx, handle) = Context::new(mailbox);

            addr_tx.send((addr, handle.clone())).unwrap();

            // Every time the actor panics restart it. If the actor gracefully stops don't restart it.
            while let Err(err) = actor_ctx.run(actor).await {
                actor = A::new(&mut actor_ctx).await;
                A::restarting(ActorStopReason::Panic(err), &mut actor_ctx).await;
            }
        });

        let (addr, ctrl_handle) = addr_rx.await.unwrap();
        (addr, ActorHandle::new(ctrl_handle, task))
    }
}

impl Actor for Supervisor {}

pub struct SuperviseActor<A>(pub A)
where
    A: Actor + Supervised;

impl<A> Message for SuperviseActor<A>
where
    A: 'static + Actor + Supervised,
{
    type Return = (Addr<A>, ActorHandle);
}

#[async_trait]
impl<A> Handler<SuperviseActor<A>> for Supervisor
where
    A: 'static + Actor + Supervised,
{
    async fn handle(
        &mut self,
        msg: SuperviseActor<A>,
        _: &mut Context<Self>,
    ) -> <SuperviseActor<A> as Message>::Return {
        let actor = msg.0;
        self.supervise(actor).await
    }
}

#[cfg(test)]
mod test {
    use crate::actors::supervisor::{Supervised, Supervisor};
    use crate::prelude::*;

    struct SupervisedActor;
    impl Actor for SupervisedActor {}

    #[async_trait]
    impl Supervised for SupervisedActor {
        async fn new(_: &mut Context<Self>) -> Self
        where
            Self: Sized,
        {
            SupervisedActor
        }
    }

    struct Panic;
    impl Message for Panic {
        type Return = ();
    }

    #[async_trait]
    impl Handler<Panic> for SupervisedActor {
        async fn handle(&mut self, _: Panic, _: &mut Context<Self>) -> <Panic as Message>::Return {
            panic!("Got panic message")
        }
    }

    #[tokio::test]
    async fn restart_panicking_supervised_actor() {
        let (mut addr, mut handle) = Supervisor::new().supervise(SupervisedActor).await;

        assert!(addr.send(Panic).await.is_err());

        // The actor should still be running
        assert!(handle.heartbeat().await)
    }

    #[tokio::test]
    async fn unsupervised_actor_should_panic_and_not_restart() {
        let (mut addr, mut handle) = SupervisedActor.start().await;
        assert!(addr.send(Panic).await.is_err());

        // The actor should not be running any longer
        assert!(!handle.heartbeat().await)
    }
}

/// In order for an Supervisor to supervise an actor, an actor needs to implement this trait
/// to provide new instances of the actor if it stops for whatever reason.
#[async_trait]
pub trait Supervised: Actor {
    /// Construct a new instance of the actor
    async fn new(context: &mut Context<Self>) -> Self
    where
        Self: Sized;

    #[allow(unused_variables)]
    /// This method is called whenever the actor has shutdown and the supervisor has begun
    /// to restart the actor with a new context
    async fn restarting(reason: ActorStopReason, ctx: &mut Context<Self>)
    where
        Self: Sized + Supervised,
    {
    }
}

#[derive(Debug)]
pub enum ActorStopReason {
    Panic(Box<dyn Any + Send>),
}