pulse-trader/src/engine/store/plugin.rs

use std::{
    marker::PhantomData,
    sync::{Arc, Weak},
};

use pulse_wire::{
    PulseWire,
    plugin::{
        RiskEngineMessage, RiskManifest, RiskMessage, StrategyEngineMessage, StrategyManifest,
        StrategyMessage,
    },
};

use tokio::{
    fs,
    io::{AsyncReadExt, AsyncWriteExt},
    process::{Child, ChildStdout, Command},
    sync::Mutex,
    task::JoinHandle,
};

use crate::{engine::Engine, store::pulse_plugin};

#[derive(Debug)]
pub struct Plugin<S: PulseWire, R: PulseWire> {
    pub stdout: Mutex<ChildStdout>,
    pub process: Mutex<Child>,
    pub _p: (PhantomData<S>, PhantomData<R>),
}

impl<S: PulseWire, R: PulseWire> Plugin<S, R> {
    pub fn new(mut child: Child) -> Self {
        Self {
            stdout: Mutex::new(child.stdout.take().expect("Failed to obtain child stdout")),
            process: Mutex::new(child),
            _p: (PhantomData, PhantomData),
        }
    }

    pub async fn recv(&self) -> tokio::io::Result<Option<R>> {
        let mut stdout = self.stdout.lock().await;

        let mut len_buf = [0u8; size_of::<usize>()];
        let size = stdout.read_exact(&mut len_buf).await?;

        let len = usize::from_le_bytes(len_buf);

        if size == 0 || len == 0 {
            return Ok(None);
        }

        let mut buffer = vec![0u8; len];

        stdout.read_exact(&mut buffer).await?;

        Ok(Some(R::from_com(&mut buffer)))
    }

    pub async fn send(&mut self, msg: &S) -> tokio::io::Result<()> {
        self.send_raw(&msg.to_com()).await
    }

    pub async fn send_raw(&mut self, msg: &[u8]) -> tokio::io::Result<()> {
        let mut process = self.process.lock().await;
        let stdin = process.stdin.as_mut().unwrap();

        stdin.write(&msg.len().to_le_bytes()).await?;
        stdin.write(msg).await?;
        stdin.flush().await?;

        Ok(())
    }
}

#[derive(Debug)]
pub struct StrategyPair {
    pub strategy: Plugin<StrategyEngineMessage, StrategyMessage>,
    pub risk: Plugin<RiskEngineMessage, RiskMessage>,

    pub strategy_manifest: Mutex<StrategyManifest>,
    pub risk_manifest: Mutex<RiskManifest>,
}

impl StrategyPair {
    pub async fn new(strategy_id: &str, risk_id: &str) -> tokio::io::Result<Arc<Self>> {
        let strategy = pulse_plugin(strategy_id)?;
        let risk = pulse_plugin(risk_id)?;

        Ok(Arc::new(Self {
            strategy: Plugin::new(
                Command::new("bash")
                    .arg(strategy.join("strategy.bash"))
                    .current_dir(&strategy)
                    .spawn()?,
            ),
            risk: Plugin::new(
                Command::new("bash")
                    .arg(strategy.join("risk.bash"))
                    .current_dir(&strategy)
                    .spawn()?,
            ),
            strategy_manifest: Mutex::new(
                toml::from_slice(&fs::read(strategy.join("strategy.toml")).await?).map_err(
                    |v| tokio::io::Error::new(std::io::ErrorKind::InvalidInput, v.to_string()),
                )?,
            ),
            risk_manifest: Mutex::new(
                toml::from_slice(&fs::read(risk.join("risk.toml")).await?).map_err(|v| {
                    tokio::io::Error::new(std::io::ErrorKind::InvalidInput, v.to_string())
                })?,
            ),
        }))
    }
}

#[derive(Debug)]
pub struct StrategyEngine {
    pub pair: Mutex<Arc<StrategyPair>>,

    pub strategy_handle: Mutex<Option<JoinHandle<tokio::io::Result<()>>>>,
    pub risk_handle: Mutex<Option<JoinHandle<tokio::io::Result<()>>>>,

    pub engine: Weak<Engine>,
}

impl StrategyEngine {
    pub async fn new(strategy_id: &str, risk_id: &str) -> tokio::io::Result<Self> {
        Ok(Self {
            pair: Mutex::new(StrategyPair::new(strategy_id, risk_id).await?),
            strategy_handle: Mutex::new(None),
            risk_handle: Mutex::new(None),
            engine: Weak::new(),
        })
    }

    pub fn initialize(mut self, engine: Weak<Engine>) -> Arc<Self> {
        self.engine = engine;

        Arc::new(self)
    }

    pub async fn run_strategy(&self) -> tokio::io::Result<()> {
        let engine = self
            .engine
            .upgrade()
            .expect("Failed to upgrade engine (StrategyEngine)");

        loop {}
    }

    pub async fn run_risk(&self) -> tokio::io::Result<()> {
        let engine = self
            .engine
            .upgrade()
            .expect("Failed to upgrade engine (StrategyEngine)");

        loop {}
    }

    pub async fn spawn(self: &Arc<Self>) {
        let engine = self.clone();

        *self.strategy_handle.lock().await =
            Some(tokio::spawn(async move { engine.run_strategy().await }));

        let engine = self.clone();

        *self.risk_handle.lock().await = Some(tokio::spawn(async move { engine.run_risk().await }));
    }

    pub async fn reload(
        self: &Arc<Self>,
        strategy_id: &str,
        risk_id: &str,
    ) -> tokio::io::Result<()> {
        {
            if let Some(handle) = &*self.strategy_handle.lock().await {
                handle.abort();
            }

            if let Some(handle) = &*self.risk_handle.lock().await {
                handle.abort();
            }

            {
                let pair = self.pair.lock().await.clone();

                pair.strategy.process.lock().await.kill().await?;
                pair.risk.process.lock().await.kill().await?;
            }

            *self.pair.lock().await = StrategyPair::new(strategy_id, risk_id).await?;
        }

        self.spawn().await;

        Ok(())
    }
}