Hopefully fixed deadlocks in strategy engine

src/engine/store/plugin.rs

@@ -14,7 +14,7 @@ use pulse_wire::{
 use tokio::{
     fs,
     io::{AsyncReadExt, AsyncWriteExt},
-    process::Command,
+    process::{Child, ChildStdout, Command},
     sync::Mutex,
     task::JoinHandle,
 };
@@ -23,23 +23,25 @@ use crate::{engine::Engine, store::pulse_plugin};
 
 #[derive(Debug)]
 pub struct Plugin<S: PulseWire, R: PulseWire> {
-    pub process: tokio::process::Child,
+    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(child: tokio::process::Child) -> Self {
+    pub fn new(mut child: Child) -> Self {
         Self {
-            process: child,
+            stdout: Mutex::new(child.stdout.take().expect("Failed to obtain child stdout")),
+            process: Mutex::new(child),
             _p: (PhantomData, PhantomData),
         }
     }
 
-    pub async fn recv(&mut self) -> tokio::io::Result<Option<R>> {
-        let stderr = self.process.stderr.as_mut().unwrap();
+    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 = stderr.read_exact(&mut len_buf).await?;
+        let size = stdout.read_exact(&mut len_buf).await?;
 
         let len = usize::from_le_bytes(len_buf);
 
@@ -49,7 +51,7 @@ impl<S: PulseWire, R: PulseWire> Plugin<S, R> {
 
         let mut buffer = vec![0u8; len];
 
-        stderr.read_exact(&mut buffer).await?;
+        stdout.read_exact(&mut buffer).await?;
 
         Ok(Some(R::from_com(&mut buffer)))
     }
@@ -59,7 +61,8 @@ impl<S: PulseWire, R: PulseWire> Plugin<S, R> {
     }
 
     pub async fn send_raw(&mut self, msg: &[u8]) -> tokio::io::Result<()> {
-        let stdin = self.process.stdin.as_mut().unwrap();
+        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?;
@@ -71,31 +74,31 @@ impl<S: PulseWire, R: PulseWire> Plugin<S, R> {
 
 #[derive(Debug)]
 pub struct StrategyPair {
-    pub strategy: Mutex<Plugin<StrategyEngineMessage, StrategyMessage>>,
-    pub risk: Mutex<Plugin<RiskEngineMessage, RiskMessage>>,
+    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<Self> {
+    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(Self {
-            strategy: Mutex::new(Plugin::new(
+        Ok(Arc::new(Self {
+            strategy: Plugin::new(
                 Command::new("bash")
                     .arg(strategy.join("strategy.bash"))
                     .current_dir(&strategy)
                     .spawn()?,
-            )),
-            risk: Mutex::new(Plugin::new(
+            ),
+            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()),
@@ -106,13 +109,13 @@ impl StrategyPair {
                     tokio::io::Error::new(std::io::ErrorKind::InvalidInput, v.to_string())
                 })?,
             ),
-        })
+        }))
     }
 }
 
 #[derive(Debug)]
 pub struct StrategyEngine {
-    pub pair: StrategyPair,
+    pub pair: Mutex<Arc<StrategyPair>>,
 
     pub strategy_handle: Mutex<Option<JoinHandle<tokio::io::Result<()>>>>,
     pub risk_handle: Mutex<Option<JoinHandle<tokio::io::Result<()>>>>,
@@ -123,7 +126,7 @@ pub struct StrategyEngine {
 impl StrategyEngine {
     pub async fn new(strategy_id: &str, risk_id: &str) -> tokio::io::Result<Self> {
         Ok(Self {
-            pair: StrategyPair::new(strategy_id, risk_id).await?,
+            pair: Mutex::new(StrategyPair::new(strategy_id, risk_id).await?),
             strategy_handle: Mutex::new(None),
             risk_handle: Mutex::new(None),
             engine: Weak::new(),
@@ -143,8 +146,6 @@ impl StrategyEngine {
             .expect("Failed to upgrade engine (StrategyEngine)");
 
         loop {}
-
-        Ok(())
     }
 
     pub async fn run_risk(&self) -> tokio::io::Result<()> {
@@ -154,8 +155,6 @@ impl StrategyEngine {
             .expect("Failed to upgrade engine (StrategyEngine)");
 
         loop {}
-
-        Ok(())
     }
 
     pub async fn spawn(self: &Arc<Self>) {
@@ -183,32 +182,14 @@ impl StrategyEngine {
                 handle.abort();
             }
 
-            self.pair.strategy.lock().await.process.kill().await?;
-            self.pair.risk.lock().await.process.kill().await?;
+            {
+                let pair = self.pair.lock().await.clone();
+
+                pair.strategy.process.lock().await.kill().await?;
+                pair.risk.process.lock().await.kill().await?;
             }
 
-        {
-            let strategy = StrategyPair::new(strategy_id, risk_id).await?;
-
-            std::mem::swap(
-                &mut *self.pair.strategy.lock().await,
-                &mut *strategy.strategy.lock().await,
-            );
-
-            std::mem::swap(
-                &mut *self.pair.strategy_manifest.lock().await,
-                &mut *strategy.strategy_manifest.lock().await,
-            );
-
-            std::mem::swap(
-                &mut *self.pair.risk.lock().await,
-                &mut *strategy.risk.lock().await,
-            );
-
-            std::mem::swap(
-                &mut *self.pair.risk_manifest.lock().await,
-                &mut *strategy.risk_manifest.lock().await,
-            );
+            *self.pair.lock().await = StrategyPair::new(strategy_id, risk_id).await?;
         }
 
         self.spawn().await;