Blocktree/crates/btrun/tests/runtime_tests.rs

#![feature(impl_trait_in_assoc_type)]

use btrun::model::*;
use btrun::test_setup;
use btrun::*;

use btlib::Result;
use btproto::protocol;
use lazy_static::lazy_static;
use log;
use serde::{Deserialize, Serialize};
use std::{
    future::{ready, Future, Ready},
    sync::{
        atomic::{AtomicU8, Ordering},
        Arc,
    },
};

test_setup!();

mod ping_pong {
    use super::*;

    use btlib::bterr;

    // The following code is a proof-of-concept for what types should be generated for a
    // simple ping-pong protocol:
    protocol! {
        named PingProtocol;
        let server = [Server];
        let client = [Client];
        Client -> End, >service(Server)!Ping;
        Server?Ping -> End, >Client!Ping::Reply;
    }
    //
    // In words, the protocol is described as follows.
    // 1. When the Listening state receives the Ping message it returns the End state and a
    //    Ping::Reply message to be sent to the SentPing state.
    // 2. When the SentPing state receives the Ping::Reply message it returns the End state.
    //
    // The End state represents an end to the session described by the protocol. When an actor
    // transitions to the End state its function returns.
    // When a state is expecting a Reply message, an error occurs if the message is not received
    // in a timely manner.

    enum PingClientState<T: Client> {
        Client(T),
        End(End),
    }

    impl<T: Client> PingClientState<T> {
        const fn name(&self) -> &'static str {
            match self {
                Self::Client(_) => "Client",
                Self::End(_) => "End",
            }
        }
    }

    struct ClientHandle<T: Client> {
        state: Option<PingClientState<T>>,
        runtime: &'static Runtime,
    }

    impl<T: Client> ClientHandle<T> {
        async fn send_ping(&mut self, mut msg: Ping, service: ServiceAddr) -> Result<PingReply> {
            let state = self
                .state
                .take()
                .ok_or_else(|| bterr!("State was not returned."))?;
            let (new_state, result) = match state {
                PingClientState::Client(state) => match state.on_send_ping(&mut msg).await {
                    TransResult::Ok((new_state, _)) => {
                        let new_state = PingClientState::End(new_state);
                        let result = self
                            .runtime
                            .call_service(service, PingProtocolMsgs::Ping(msg))
                            .await;
                        (new_state, result)
                    }
                    TransResult::Abort { from, err } => {
                        let new_state = PingClientState::Client(from);
                        (new_state, Err(err))
                    }
                    TransResult::Fatal { err } => return Err(err),
                },
                state => {
                    let result = Err(bterr!("Can't send Ping in state {}.", state.name()));
                    (state, result)
                }
            };
            self.state = Some(new_state);
            let reply = result?;
            match reply {
                PingProtocolMsgs::PingReply(reply) => Ok(reply),
                msg => Err(bterr!(
                    "Unexpected message type sent in reply: {}",
                    msg.name()
                )),
            }
        }
    }

    async fn spawn_client<T: Client>(init: T, runtime: &'static Runtime) -> ClientHandle<T> {
        let state = Some(PingClientState::Client(init));
        ClientHandle { state, runtime }
    }

    async fn register_server<Init, F>(
        make_init: F,
        rt: &'static Runtime,
        id: ServiceId,
    ) -> Result<ServiceName>
    where
        Init: 'static + Server,
        F: 'static + Send + Sync + Clone + Fn() -> Init,
    {
        enum ServerState<S> {
            Server(S),
            End(End),
        }

        async fn server_loop<Init, F>(
            _runtime: &'static Runtime,
            make_init: F,
            mut mailbox: Mailbox<PingProtocolMsgs>,
            _act_id: ActorId,
        ) where
            Init: 'static + Server,
            F: 'static + Send + Sync + FnOnce() -> Init,
        {
            let mut state = ServerState::Server(make_init());
            while let Some(envelope) = mailbox.recv().await {
                let (msg, msg_kind) = envelope.split();
                state = match (state, msg) {
                    (ServerState::Server(listening_state), PingProtocolMsgs::Ping(msg)) => {
                        match listening_state.handle_ping(msg).await {
                            TransResult::Ok((new_state, reply)) => match msg_kind {
                                EnvelopeKind::Call { reply: replier } => {
                                    let replier =
                                        replier.expect("The reply has already been sent.");
                                    if let Err(_) = replier.send(PingProtocolMsgs::PingReply(reply))
                                    {
                                        panic!("Failed to send Ping reply.");
                                    }
                                    ServerState::End(new_state)
                                }
                                _ => panic!("'Ping' was expected to be a Call message."),
                            },
                            TransResult::Abort { from, err } => {
                                log::warn!("Aborted transition from the {} while handling the {} message: {}", "Server", "Ping", err);
                                ServerState::Server(from)
                            }
                            TransResult::Fatal { err } => {
                                panic!("Fatal error while handling Ping message in Server state: {err}");
                            }
                        }
                    }
                    (state, _) => state,
                };

                if let ServerState::End(_) = state {
                    break;
                }
            }
        }

        rt.register::<PingProtocolMsgs, _>(id, move |runtime| {
            let make_init = make_init.clone();
            let fut = async move {
                let actor_impl = runtime
                    .spawn(move |_, mailbox, act_id| {
                        server_loop(runtime, make_init, mailbox, act_id)
                    })
                    .await;
                Ok(actor_impl)
            };
            Box::pin(fut)
        })
        .await
    }

    #[derive(Serialize, Deserialize)]
    pub struct Ping;
    impl CallMsg for Ping {
        type Reply = PingReply;
    }

    #[derive(Serialize, Deserialize)]
    pub struct PingReply;

    struct ClientState {
        counter: Arc<AtomicU8>,
    }

    impl ClientState {
        fn new(counter: Arc<AtomicU8>) -> Self {
            counter.fetch_add(1, Ordering::SeqCst);
            Self { counter }
        }
    }

    impl Client for ClientState {
        fn actor_impl() -> String {
            "client".into()
        }

        type OnSendPingFut = impl Future<Output = TransResult<Self, (End, PingReply)>>;
        fn on_send_ping(self, _msg: &mut Ping) -> Self::OnSendPingFut {
            self.counter.fetch_sub(1, Ordering::SeqCst);
            ready(TransResult::Ok((End, PingReply)))
        }
    }

    struct ServerState {
        counter: Arc<AtomicU8>,
    }

    impl ServerState {
        fn new(counter: Arc<AtomicU8>) -> Self {
            counter.fetch_add(1, Ordering::SeqCst);
            Self { counter }
        }
    }

    impl Server for ServerState {
        fn actor_impl() -> String {
            "server".into()
        }

        type HandlePingFut = impl Future<Output = TransResult<Self, (End, PingReply)>>;
        fn handle_ping(self, _msg: Ping) -> Self::HandlePingFut {
            self.counter.fetch_sub(1, Ordering::SeqCst);
            ready(TransResult::Ok((End, PingReply)))
        }
    }

    #[test]
    fn ping_pong_test() {
        ASYNC_RT.block_on(async {
            const SERVICE_ID: &str = "PingPongProtocolServer";
            let service_id = ServiceId::from(SERVICE_ID);
            let counter = Arc::new(AtomicU8::new(0));
            let service_name = {
                let service_counter = counter.clone();
                let make_init = move || {
                    let server_counter = service_counter.clone();
                    ServerState::new(server_counter)
                };
                register_server(make_init, &RUNTIME, service_id.clone())
                    .await
                    .unwrap()
            };
            let mut client_handle = spawn_client(ClientState::new(counter.clone()), &RUNTIME).await;
            let service_addr = ServiceAddr::new(service_name, true);
            client_handle.send_ping(Ping, service_addr).await.unwrap();

            assert_eq!(0, counter.load(Ordering::SeqCst));

            RUNTIME.deregister(&service_id, None).await.unwrap();
        });
    }
}

mod travel_agency {
    use super::*;

    // Here's another protocol example. This is the Customer and Travel Agency protocol used as an
    // example in the survey paper "Behavioral Types in Programming Languages."
    // Note that the Choosing state can send messages at any time, not just in response to another
    // message because there is a transition from Choosing that doesn't use the receive operator
    // (`?`).
    protocol! {
        named TravelAgency;
        let agency = [Listening];
        let customer = [Choosing];
        Choosing -> Choosing, >service(Listening)!Query;
        Choosing -> Choosing, >service(Listening)!Accept;
        Choosing -> Choosing, >service(Listening)!Reject;
        Listening?Query -> Listening, >Choosing!Query::Reply;
        Choosing?Query::Reply -> Choosing;
        Listening?Accept -> End, >Choosing!Accept::Reply;
        Choosing?Accept::Reply -> End;
        Listening?Reject -> End, >Choosing!Reject::Reply;
        Choosing?Reject::Reply -> End;
    }

    #[derive(Serialize, Deserialize)]
    pub struct Query;

    impl CallMsg for Query {
        type Reply = ();
    }

    #[derive(Serialize, Deserialize)]
    pub struct Reject;

    impl CallMsg for Reject {
        type Reply = ();
    }

    #[derive(Serialize, Deserialize)]
    pub struct Accept;

    impl CallMsg for Accept {
        type Reply = ();
    }
}

#[allow(dead_code)]
mod client_callback {

    use super::*;

    use std::{panic::panic_any, time::Duration};
    use tokio::{sync::oneshot, time::timeout};

    #[derive(Serialize, Deserialize)]
    pub struct Register {
        factor: usize,
    }

    #[derive(Serialize, Deserialize)]
    pub struct Completed {
        value: usize,
    }

    protocol! {
        named ClientCallback;
        let server = [Listening];
        let worker = [Working];
        let client = [Unregistered, Registered];
        Unregistered -> Registered, >service(Listening)!Register[Registered];
        Listening?Register[Registered] -> Listening, Working[Registered];
        Working[Registered] -> End, >Registered!Completed;
        Registered?Completed -> End;
    }

    struct UnregisteredState {
        sender: oneshot::Sender<usize>,
    }

    impl Unregistered for UnregisteredState {
        fn actor_impl() -> String {
            "client".into()
        }

        type OnSendRegisterRegistered = RegisteredState;
        type OnSendRegisterFut = Ready<TransResult<Self, Self::OnSendRegisterRegistered>>;
        fn on_send_register(self, _arg: &mut Register) -> Self::OnSendRegisterFut {
            ready(TransResult::Ok(RegisteredState {
                sender: self.sender,
            }))
        }
    }

    struct RegisteredState {
        sender: oneshot::Sender<usize>,
    }

    impl Registered for RegisteredState {
        type HandleCompletedFut = Ready<TransResult<Self, End>>;
        fn handle_completed(self, arg: Completed) -> Self::HandleCompletedFut {
            self.sender.send(arg.value).unwrap();
            ready(TransResult::Ok(End))
        }
    }

    struct ListeningState {
        multiple: usize,
    }

    impl Listening for ListeningState {
        fn actor_impl() -> String {
            "server".into()
        }

        type HandleRegisterListening = ListeningState;
        type HandleRegisterWorking = WorkingState;
        type HandleRegisterFut = Ready<TransResult<Self, (ListeningState, WorkingState)>>;
        fn handle_register(self, arg: Register) -> Self::HandleRegisterFut {
            let multiple = self.multiple;
            ready(TransResult::Ok((
                self,
                WorkingState {
                    factor: arg.factor,
                    multiple,
                },
            )))
        }
    }

    struct WorkingState {
        factor: usize,
        multiple: usize,
    }

    impl Working for WorkingState {
        fn actor_impl() -> String {
            "worker".into()
        }

        type OnSendCompletedFut = Ready<TransResult<Self, (End, Completed)>>;
        fn on_send_completed(self) -> Self::OnSendCompletedFut {
            let value = self.multiple * self.factor;
            ready(TransResult::Ok((End, Completed { value })))
        }
    }

    use ::tokio::sync::Mutex;

    enum ClientState<Init: Unregistered> {
        Unregistered(Init),
        Registered(Init::OnSendRegisterRegistered),
        End(End),
    }

    impl<Init: Unregistered> ClientState<Init> {
        pub fn name(&self) -> &'static str {
            match self {
                Self::Unregistered(_) => "Unregistered",
                Self::Registered(_) => "Registered",
                Self::End(_) => "End",
            }
        }
    }

    struct ClientHandle<Init: Unregistered> {
        runtime: &'static Runtime,
        state: Arc<Mutex<Option<ClientState<Init>>>>,
        name: ActorName,
    }

    impl<Init: Unregistered> ClientHandle<Init> {
        async fn send_register(&self, to: ServiceAddr, mut msg: Register) -> Result<()> {
            let mut guard = self.state.lock().await;
            let state = guard
                .take()
                .unwrap_or_else(|| panic!("Logic error. The state was not returned."));
            let new_state = match state {
                ClientState::Unregistered(state) => match state.on_send_register(&mut msg).await {
                    TransResult::Ok(new_state) => {
                        let msg = ClientCallbackMsgs::Register(msg);
                        self.runtime
                            .send_service(to, self.name.clone(), msg)
                            .await?;
                        ClientState::Registered(new_state)
                    }
                    TransResult::Abort { from, err } => {
                        log::warn!(
                            "Aborted transition from the {} state: {}",
                            "Unregistered",
                            err
                        );
                        ClientState::Unregistered(from)
                    }
                    TransResult::Fatal { err } => {
                        return Err(err);
                    }
                },
                state => state,
            };
            *guard = Some(new_state);
            Ok(())
        }
    }

    async fn spawn_client<Init>(init: Init, runtime: &'static Runtime) -> ClientHandle<Init>
    where
        Init: 'static + Unregistered,
    {
        let state = Arc::new(Mutex::new(Some(ClientState::Unregistered(init))));
        let name = {
            let state = state.clone();
            runtime.spawn(move |_, mut mailbox, _act_id| async move {
                while let Some(envelope) = mailbox.recv().await {
                    let mut guard = state.lock().await;
                    let state = guard.take()
                        .unwrap_or_else(|| panic!("Logic error. The state was not returned."));
                    let (msg, _kind) = envelope.split();
                    let new_state = match (state, msg) {
                        (ClientState::Registered(curr_state), ClientCallbackMsgs::Completed(msg)) => {
                            match curr_state.handle_completed(msg).await {
                                TransResult::Ok(next) =>  ClientState::<Init>::End(next),
                                TransResult::Abort { from, err } => {
                                    log::warn!("Aborted transition from the {} state while handling the {} message: {}", "Registered", "Completed", err);
                                    ClientState::Registered(from)
                                }
                                TransResult::Fatal { err } => {
                                    panic_any(ActorPanic {
                                        actor_impl: Init::actor_impl(),
                                        state: "Registered",
                                        message: "Completed",
                                        kind: TransKind::Receive,
                                        err
                                    });
                                }
                            }
                        }
                        (state, msg) => {
                            log::error!("Unexpected message {} in state {}.", msg.name(), state.name());
                            state
                        }
                    };
                    *guard = Some(new_state);
                }
            }).await
        };
        ClientHandle {
            runtime,
            state,
            name,
        }
    }

    async fn register_server<Init, F>(
        make_init: F,
        runtime: &'static Runtime,
        service_id: ServiceId,
    ) -> Result<ServiceName>
    where
        Init: 'static + Listening<HandleRegisterListening = Init>,
        F: 'static + Send + Sync + Clone + Fn() -> Init,
    {
        enum ServerState<S: Listening> {
            Listening(S),
        }

        impl<S: Listening> ServerState<S> {
            fn name(&self) -> &'static str {
                match self {
                    Self::Listening(_) => "Listening",
                }
            }
        }

        async fn server_loop<Init, F>(
            runtime: &'static Runtime,
            make_init: F,
            mut mailbox: Mailbox<ClientCallbackMsgs>,
            _act_id: ActorId,
        ) where
            Init: 'static + Listening<HandleRegisterListening = Init>,
            F: 'static + Send + Sync + Fn() -> Init,
        {
            let mut state = ServerState::Listening(make_init());
            while let Some(envelope) = mailbox.recv().await {
                let (msg, msg_kind) = envelope.split();
                let new_state = match (state, msg) {
                    (ServerState::Listening(curr_state), ClientCallbackMsgs::Register(msg)) => {
                        match curr_state.handle_register(msg).await {
                            TransResult::Ok((new_state, working_state)) => {
                                if let EnvelopeKind::Send { from, .. } = msg_kind {
                                    start_worker(working_state, from, runtime).await;
                                } else {
                                    log::error!("Expected Register to be a Send message.");
                                }
                                ServerState::Listening(new_state)
                            }
                            TransResult::Abort { from, err } => {
                                log::warn!("Aborted transition from the {} state while handling the {} message: {}", "Listening", "Register", err);
                                ServerState::Listening(from)
                            }
                            TransResult::Fatal { err } => {
                                panic_any(ActorPanic {
                                    actor_impl: Init::actor_impl(),
                                    state: "Listening",
                                    message: "Register",
                                    kind: TransKind::Receive,
                                    err,
                                });
                            }
                        }
                    }
                    (state, msg) => {
                        log::error!(
                            "Unexpected message {} in state {}.",
                            msg.name(),
                            state.name()
                        );
                        state
                    }
                };
                state = new_state;
            }
        }

        runtime
            .register::<ClientCallbackMsgs, _>(service_id, move |runtime: &'static Runtime| {
                let make_init = make_init.clone();
                let fut = async move {
                    let make_init = make_init.clone();
                    let actor_impl = runtime
                        .spawn(move |_, mailbox, act_id| {
                            server_loop(runtime, make_init, mailbox, act_id)
                        })
                        .await;
                    Ok(actor_impl)
                };
                Box::pin(fut)
            })
            .await
    }

    async fn start_worker<Init>(
        init: Init,
        owned: ActorName,
        runtime: &'static Runtime,
    ) -> ActorName
    where
        Init: 'static + Working,
    {
        enum WorkerState<S: Working> {
            Working(S),
        }

        runtime
            .spawn::<ClientCallbackMsgs, _, _>(move |_, _, act_id| async move {
                let msg = match init.on_send_completed().await {
                    TransResult::Ok((End, msg)) => msg,
                    TransResult::Abort { err, .. } | TransResult::Fatal { err } => {
                        panic_any(ActorPanic {
                            actor_impl: Init::actor_impl(),
                            state: "Working",
                            message: "Completed",
                            kind: TransKind::Send,
                            err,
                        })
                    }
                };
                let from = runtime.actor_name(act_id);
                let msg = ClientCallbackMsgs::Completed(msg);
                runtime.send(owned, from, msg).await.unwrap_or_else(|err| {
                    panic_any(ActorPanic {
                        actor_impl: Init::actor_impl(),
                        state: "Working",
                        message: "Completed",
                        kind: TransKind::Send,
                        err,
                    });
                });
            })
            .await
    }

    #[test]
    fn client_callback_protocol() {
        ASYNC_RT.block_on(async {
            const SERVICE_ID: &str = "ClientCallbackProtocolListening";
            let service_id = ServiceId::from(SERVICE_ID);
            let service_name = {
                let make_init = move || ListeningState { multiple: 2 };
                register_server(make_init, &RUNTIME, service_id.clone())
                    .await
                    .unwrap()
            };
            let (sender, receiver) = oneshot::channel();
            let client_handle = spawn_client(UnregisteredState { sender }, &RUNTIME).await;
            let service_addr = ServiceAddr::new(service_name, false);
            client_handle
                .send_register(service_addr, Register { factor: 21 })
                .await
                .unwrap();
            let value = timeout(Duration::from_millis(500), receiver)
                .await
                .unwrap()
                .unwrap();

            assert_eq!(42, value);
        });
    }
}