From 02cd4c6347d4ed391b90dcd6ededebf7565c903e Mon Sep 17 00:00:00 2001 From: Klesti Selimaj Date: Fri, 7 Aug 2026 05:30:15 +0200 Subject: [PATCH] graceful exit --- Cargo.toml | 7 +-- src/bot.rs | 74 ++++++++++++++++++++++++++------ src/executor/mod.rs | 99 +++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 24 ++++++++++- src/strategy/mod.rs | 2 + src/strategy/veloc.rs | 6 +++ 6 files changed, 191 insertions(+), 21 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 37de6f2..4aab1a1 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/bot.rs b/src/bot.rs index df2a95b..13ad9da 100644 --- a/src/bot.rs +++ b/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>>, pub accounts: Mutex, - pub executor: Mutex>, + pub executor: Mutex, pub strategy: Mutex>, pub trade_log: Mutex>, } @@ -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) -> 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) -> anyhow::Result<()> { + pub async fn start( + self: &Arc, + mut shutdown: watch::Receiver, + ) -> 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) -> anyhow::Result { 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::(&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) + } } } diff --git a/src/executor/mod.rs b/src/executor/mod.rs index f180176..0c915de 100644 --- a/src/executor/mod.rs +++ b/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, + pub positions: HashMap, +} + #[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(()) + } +} diff --git a/src/main.rs b/src/main.rs index 32b793e..a9acd57 100644 --- a/src/main.rs +++ b/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) { + 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 } diff --git a/src/strategy/mod.rs b/src/strategy/mod.rs index 6ca4731..d358b70 100644 --- a/src/strategy/mod.rs +++ b/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) -> anyhow::Result<()>; + async fn on_new_coin(&mut self, bot: Arc, token: NewToken) -> anyhow::Result<()>; async fn on_trade(&mut self, bot: Arc, trade: Trade) -> anyhow::Result<()>; diff --git a/src/strategy/veloc.rs b/src/strategy/veloc.rs index a8db349..229983f 100644 --- a/src/strategy/veloc.rs +++ b/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) -> anyhow::Result<()> { + bot.executor.lock().await.sell_all(PRIORITY, SLIPPAGE).await + } + async fn on_new_coin(&mut self, bot: Arc, 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?;