graceful exit

Cargo.toml

@@ -10,12 +10,7 @@ futures-util = "0.3.33"
 serde = { version = "1.0.229", features = ["derive", "serde_derive"] }
 serde_json = "1.0.151"
 
-tokio = { version = "1.53.1", features = [
-    "macros",
-    "rt-multi-thread",
-    "sync",
-    "fs",
-] }
+tokio = { version = "1.53.1", features = ["macros", "rt-multi-thread", "sync", "fs", "io-std"] }
 tokio-tungstenite = { version = "0.30.0", features = ["native-tls"] }
 log = "0.4.33"
 env_logger = "0.11.11"

src/bot.rs

@@ -1,14 +1,17 @@
-use std::sync::Arc;
+use std::{collections::HashMap, sync::Arc};
 
 use anyhow::Context;
 use futures_util::{SinkExt, StreamExt};
 use serde_json::json;
-use tokio::{net::TcpStream, sync::Mutex};
+use tokio::{
+    net::TcpStream,
+    sync::{Mutex, watch},
+};
 use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
 
 use crate::{
     account::AccountManager,
-    executor::Executor,
+    executor::ExecutorWrapper,
     strategy::{Strategy, veloc::MomentumVelocityStrategy},
     tradelog::TradeLog,
 };
@@ -16,7 +19,7 @@ use crate::{
 pub struct Bot {
     pub ws: Mutex<WebSocketStream<MaybeTlsStream<TcpStream>>>,
     pub accounts: Mutex<AccountManager>,
-    pub executor: Mutex<Box<dyn Executor>>,
+    pub executor: Mutex<ExecutorWrapper>,
     pub strategy: Mutex<Box<dyn Strategy>>,
     pub trade_log: Mutex<Vec<TradeLog>>,
 }
@@ -69,14 +72,23 @@ impl Bot {
 
         Ok(Arc::new(Self {
             ws: Mutex::new(connect_async("wss://pumpdev.io/ws").await?.0),
-            executor: Mutex::new(Box::new(account.executor())),
+            executor: Mutex::new(ExecutorWrapper {
+                executor: Box::new(account.executor()),
+                positions: HashMap::new(),
+            }),
             accounts: Mutex::new(accounts),
             strategy: Mutex::new(Box::new(MomentumVelocityStrategy::new())),
             trade_log: Mutex::new(Vec::new()),
         }))
     }
 
-    pub async fn refresh_account(&self) -> anyhow::Result<()> {
+    pub async fn refresh_account(self: &Arc<Self>) -> anyhow::Result<()> {
+        self.strategy
+            .lock()
+            .await
+            .execute_sell_all(self.clone())
+            .await?;
+
         let accounts = self.accounts.lock().await;
 
         let account = accounts
@@ -85,7 +97,7 @@ impl Bot {
             .context("Failed to get account")?
             .clone();
 
-        *self.executor.lock().await = Box::new(account.executor());
+        self.executor.lock().await.executor = Box::new(account.executor());
 
         Ok(())
     }
@@ -102,32 +114,66 @@ impl Bot {
         Ok(())
     }
 
-    pub async fn start(self: &Arc<Self>) -> anyhow::Result<()> {
+    pub async fn start(
+        self: &Arc<Self>,
+        mut shutdown: watch::Receiver<bool>,
+    ) -> anyhow::Result<()> {
         self.initialize_websocket_subscribe().await?;
 
+        loop {
+            tokio::select! {
+                _ = shutdown.changed() => {
+                    if *shutdown.borrow() {
+                        log::info!("Shutdown signal received.");
+
+                        self.strategy
+                            .lock()
+                            .await
+                            .execute_sell_all(self.clone())
+                            .await?;
+
+                        log::info!("Bye!");
+
+                        break;
+                    }
+                }
+
+                result = self.tick() => {
+                    if result? {
+                        log::warn!("Websocket closed.");
+                        break;
+                    }
+                }
+            }
+        }
+
+        Ok(())
+    }
+
+    pub async fn tick(self: &Arc<Self>) -> anyhow::Result<bool> {
         let mut ws = self.ws.lock().await;
 
-        while let Some(msg) = ws.next().await.transpose()? {
+        if let Some(msg) = ws.next().await.transpose()? {
             if let tokio_tungstenite::tungstenite::Message::Text(text) = msg {
                 match serde_json::from_str::<crate::types::PumpDevEvent>(&text) {
                     Ok(crate::types::PumpDevEvent::Create(token)) => {
                         drop(ws);
+
                         self.strategy
                             .lock()
                             .await
                             .on_new_coin(self.clone(), token)
                             .await?;
-                        ws = self.ws.lock().await;
                     }
 
                     Ok(crate::types::PumpDevEvent::Trade(trade)) => {
                         drop(ws);
+
                         self.strategy
                             .lock()
                             .await
                             .on_trade(self.clone(), trade)
                             .await?;
-                        ws = self.ws.lock().await;
                     }
 
                     Ok(_event) => {
@@ -139,8 +185,10 @@ impl Bot {
                     }
                 }
             }
-        }
 
-        Ok(())
+            Ok(false)
+        } else {
+            Ok(true)
+        }
     }
 }

src/executor/mod.rs

@@ -1,9 +1,16 @@
 pub mod pump_fun;
 
+use std::collections::HashMap;
+
 use rust_decimal::Decimal;
 
 use crate::account::Account;
 
+pub struct ExecutorWrapper {
+    pub executor: Box<dyn Executor>,
+    pub positions: HashMap<String, Decimal>,
+}
+
 #[allow(unused_variables)]
 #[async_trait::async_trait]
 pub trait Executor: Send + Sync {
@@ -45,3 +52,95 @@ impl Account {
         }
     }
 }
+
+impl ExecutorWrapper {
+    pub async fn buy(
+        &mut self,
+        mint: &str,
+        amount: Decimal,
+        priority: Decimal,
+        slippage: u16,
+    ) -> anyhow::Result<()> {
+        self.executor.buy(mint, amount, priority, slippage).await?;
+
+        let position = self
+            .positions
+            .entry(mint.to_string())
+            .or_insert(Decimal::ZERO);
+
+        *position += amount;
+
+        Ok(())
+    }
+
+    pub async fn sell(
+        &mut self,
+        mint: &str,
+        amount: Decimal,
+        priority: Decimal,
+        slippage: u16,
+    ) -> anyhow::Result<()> {
+        self.executor.sell(mint, amount, priority, slippage).await?;
+
+        if let Some(position) = self.positions.get_mut(mint) {
+            *position -= amount;
+
+            if *position <= Decimal::ZERO {
+                self.positions.remove(mint);
+            }
+        }
+
+        Ok(())
+    }
+
+    pub async fn sell_percent(
+        &mut self,
+        mint: &str,
+        amount: u8,
+        priority: Decimal,
+        slippage: u16,
+    ) -> anyhow::Result<()> {
+        let sell_amount = match self.positions.get(mint) {
+            Some(position) => *position * Decimal::from(amount) / Decimal::from(100),
+            None => return Ok(()),
+        };
+
+        self.executor
+            .sell_percent(mint, amount, priority, slippage)
+            .await?;
+
+        if let Some(position) = self.positions.get_mut(mint) {
+            *position -= sell_amount;
+
+            if *position <= Decimal::ZERO {
+                self.positions.remove(mint);
+            }
+        }
+
+        Ok(())
+    }
+
+    pub async fn sell_all(&mut self, priority: Decimal, slippage: u16) -> anyhow::Result<()> {
+        let positions: Vec<(String, Decimal)> = self
+            .positions
+            .iter()
+            .map(|(mint, amount)| (mint.clone(), *amount))
+            .collect();
+
+        for (mint, amount) in positions {
+            self.executor
+                .sell(&mint, amount, priority, slippage)
+                .await?;
+
+            if let Some(position) = self.positions.get_mut(&mint) {
+                *position -= amount;
+
+                if *position <= Decimal::ZERO {
+                    self.positions.remove(&mint);
+                }
+            }
+        }
+
+        Ok(())
+    }
+}

src/main.rs

@@ -2,18 +2,38 @@ pub mod account;
 pub mod bot;
 pub mod executor;
 pub mod strategy;
-pub mod types;
 pub mod tradelog;
+pub mod types;
+
+use tokio::io::{self, AsyncBufReadExt, BufReader};
+use tokio::sync::watch;
 
 use crate::bot::Bot;
 
+async fn shutdown_listener(tx: watch::Sender<bool>) {
+    let stdin = BufReader::new(io::stdin());
+    let mut lines = stdin.lines();
+
+    while let Ok(Some(line)) = lines.next_line().await {
+        if line.trim().eq_ignore_ascii_case("exit") {
+            println!("Shutdown requested.");
+            let _ = tx.send(true);
+            break;
+        }
+    }
+}
+
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
     let mut builder = env_logger::Builder::from_default_env();
     builder.filter_level(log::LevelFilter::Info);
     builder.init();
 
+    let (shutdown_tx, shutdown_rx) = watch::channel(false);
+
+    tokio::spawn(shutdown_listener(shutdown_tx));
+
     let bot = Bot::new().await?;
 
-    bot.start().await
+    bot.start(shutdown_rx).await
 }

src/strategy/mod.rs

@@ -9,6 +9,8 @@ use crate::{
 
 #[async_trait::async_trait]
 pub trait Strategy: Send + Sync {
+    async fn execute_sell_all(&mut self, bot: Arc<Bot>) -> anyhow::Result<()>;
+
     async fn on_new_coin(&mut self, bot: Arc<Bot>, token: NewToken) -> anyhow::Result<()>;
 
     async fn on_trade(&mut self, bot: Arc<Bot>, trade: Trade) -> anyhow::Result<()>;

src/strategy/veloc.rs

@@ -74,6 +74,10 @@ impl MomentumVelocityStrategy {
 
 #[async_trait::async_trait]
 impl Strategy for MomentumVelocityStrategy {
+    async fn execute_sell_all(&mut self, bot: Arc<Bot>) -> anyhow::Result<()> {
+        bot.executor.lock().await.sell_all(PRIORITY, SLIPPAGE).await
+    }
+
     async fn on_new_coin(&mut self, bot: Arc<Bot>, token: NewToken) -> anyhow::Result<()> {
         trace!("[NEW COIN] Event received for token: {}", token.mint);
 
@@ -176,6 +180,8 @@ impl Strategy for MomentumVelocityStrategy {
                     pos.trade.close(current_price, reason);
 
                     info!("TRADE RESULT: {:?}", pos.trade);
+
+                    bot.trade_log.lock().await.push(pos.trade);
                 }
 
                 self.cleanup_and_unsubscribe(&bot, mint).await?;