path: root/crates/core/tedge_api/examples/heartbeat.rs

use std::{collections::HashMap, time::Duration};

use async_trait::async_trait;
use futures::FutureExt;
use tedge_api::{
    address::ReplySenderFor,
    message::MessageType,
    plugin::{AcceptsReplies, BuiltPlugin, Handle, Message, PluginDeclaration, PluginExt},
    Address, CancellationToken, Plugin, PluginBuilder, PluginConfiguration, PluginDirectory,
    PluginError,
};

/// A message that represents a heartbeat that gets sent to plugins
#[derive(Debug)]
struct Heartbeat;
impl Message for Heartbeat {}
impl AcceptsReplies for Heartbeat {
    type Reply = HeartbeatStatus;
}

/// The reply for a heartbeat
#[derive(Debug)]
enum HeartbeatStatus {
    Alive,
    Degraded,
}
impl Message for HeartbeatStatus {}

/// A PluginBuilder that gets used to build a HeartbeatService plugin instance
#[derive(Debug)]
struct HeartbeatServiceBuilder;

#[derive(miette::Diagnostic, thiserror::Error, Debug)]
enum HeartbeatBuildError {
    #[error(transparent)]
    TomlParse(#[from] toml::de::Error),
}

#[async_trait]
impl<PD: PluginDirectory> PluginBuilder<PD> for HeartbeatServiceBuilder {
    fn kind_name() -> &'static str {
        todo!()
    }

    fn kind_message_types() -> tedge_api::plugin::HandleTypes
    where
        Self: Sized,
    {
        HeartbeatService::get_handled_types()
    }

    async fn verify_configuration(
        &self,
        _config: &PluginConfiguration,
    ) -> Result<(), tedge_api::error::PluginError> {
        Ok(())
    }

    async fn instantiate(
        &self,
        config: PluginConfiguration,
        cancellation_token: CancellationToken,
        plugin_dir: &PD,
    ) -> Result<BuiltPlugin, PluginError>
    where
        PD: 'async_trait,
    {
        let hb_config: HeartbeatConfig =
            toml::Value::try_into(config).map_err(HeartbeatBuildError::from)?;
        let monitored_services = hb_config
            .plugins
            .iter()
            .map(|name| {
                plugin_dir
                    .get_address_for::<HeartbeatMessages>(name)
                    .map(|addr| (name.clone(), addr))
            })
            .collect::<Result<Vec<_>, _>>()?;
        Ok(HeartbeatService::new(
            Duration::from_millis(hb_config.interval),
            monitored_services,
            cancellation_token,
        )
        .finish())
    }
}

/// The configuration a HeartbeatServices can receive is represented by this type
#[derive(serde::Deserialize, Debug)]
struct HeartbeatConfig {
    interval: u64,
    plugins: Vec<String>,
}

/// The HeartbeatService type represents the actual plugin
struct HeartbeatService {
    interval_duration: Duration,
    monitored_services: Vec<(String, Address<HeartbeatMessages>)>,
    cancel_token: CancellationToken,
}

impl PluginDeclaration for HeartbeatService {
    type HandledMessages = ();
}

#[async_trait]
impl Plugin for HeartbeatService {
    /// The setup function of the HeartbeatService can be used by the plugin author to setup for
    /// example a connection to an external service. In this example, it is simply used to send the
    /// heartbeat
    ///
    /// Because this example is _simple_, we do not spawn a background task that periodically sends
    /// the heartbeat. In a real world scenario, that background task would be started here.
    async fn start(&mut self) -> Result<(), PluginError> {
        println!(
            "HeartbeatService: Setting up heartbeat service with interval: {:?}!",
            self.interval_duration
        );

        for service in &self.monitored_services {
            let mut interval = tokio::time::interval(self.interval_duration);
            let service = service.clone();
            let cancel_token = self.cancel_token.child_token();
            tokio::spawn(async move {
                loop {
                    tokio::select! {
                        _ = interval.tick() => {}
                        _ = cancel_token.cancelled() => {
                            break
                        }
                    }
                    println!(
                        "HeartbeatService: Sending heartbeat to service: {:?}",
                        service
                    );
                    tokio::select! {
                        reply = service
                        .1
                        .send_and_wait(Heartbeat)
                        .then(|answer| {
                            answer.unwrap()
                            .wait_for_reply(Duration::from_millis(100))}
                        ) => {
                            match reply
                            {
                                Ok(HeartbeatStatus::Alive) => {
                                    println!("HeartbeatService: Received all is well!")
                                }
                                Ok(HeartbeatStatus::Degraded) => {
                                    println!(
                                        "HeartbeatService: Oh-oh! Plugin '{}' is not doing well",
                                        service.0
                                        )
                                }

                                Err(reply_error) => {
                                    println!(
                                        "HeartbeatService: Critical error for '{}'! {reply_error}",
                                        service.0
                                        )
                                }
                            }
                        }

                        _ = cancel_token.cancelled() => {
                            break
                        }
                    }
                }
            });
        }
        Ok(())
    }

    /// A plugin author can use this shutdown function to clean resources when thin-edge shuts down
    async fn shutdown(&mut self) -> Result<(), PluginError> {
        println!("HeartbeatService: Shutting down heartbeat service!");
        Ok(())
    }
}

impl HeartbeatService {
    fn new(
        interval_duration: Duration,
        monitored_services: Vec<(String, Address<HeartbeatMessages>)>,
        cancel_token: CancellationToken,
    ) -> Self {
        Self {
            interval_duration,
            monitored_services,
            cancel_token,
        }
    }
}

/// A plugin that receives heartbeats
struct CriticalServiceBuilder;

// declare a set of messages that the CriticalService can receive.
// In this example, it can only receive a Heartbeat.
tedge_api::make_receiver_bundle!(struct HeartbeatMessages(Heartbeat));

#[async_trait]
impl<PD: PluginDirectory> PluginBuilder<PD> for CriticalServiceBuilder {
    fn kind_name() -> &'static str {
        todo!()
    }

    fn kind_message_types() -> tedge_api::plugin::HandleTypes
    where
        Self: Sized,
    {
        CriticalService::get_handled_types()
    }

    async fn verify_configuration(
        &self,
        _config: &PluginConfiguration,
    ) -> Result<(), tedge_api::error::PluginError> {
        Ok(())
    }

    async fn instantiate(
        &self,
        _config: PluginConfiguration,
        _cancellation_token: CancellationToken,
        _plugin_dir: &PD,
    ) -> Result<BuiltPlugin, PluginError>
    where
        PD: 'async_trait,
    {
        Ok(CriticalService {
            status: tokio::sync::Mutex::new(true),
        }
        .finish())
    }
}

/// The actual "critical" plugin implementation
struct CriticalService {
    status: tokio::sync::Mutex<bool>,
}

/// The CriticalService can receive Heartbeat objects, thus it needs a Handle<Heartbeat>
/// implementation
#[async_trait]
impl Handle<Heartbeat> for CriticalService {
    async fn handle_message(
        &self,
        _message: Heartbeat,
        sender: ReplySenderFor<Heartbeat>,
    ) -> Result<(), PluginError> {
        println!("CriticalService: Received Heartbeat!");
        let mut status = self.status.lock().await;

        let _ = sender.