diff --git a/src/bot.rs b/src/bot.rs index 0370fac..b8ccd13 100644 --- a/src/bot.rs +++ b/src/bot.rs @@ -3,7 +3,7 @@ use std::collections::HashMap; use anyhow::Context; use futures_util::{SinkExt, StreamExt}; use serde_json::json; -use tokio::net::TcpStream; +use tokio::{net::TcpStream, sync::Mutex}; use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async}; use crate::{ @@ -13,15 +13,15 @@ use crate::{ }; pub struct Bot { - pub ws: WebSocketStream>, - pub accounts: AccountManager, - pub executor: Box, - pub tokens: HashMap, + pub ws: Mutex>>, + pub accounts: Mutex, + pub executor: Mutex>, + pub tokens: Mutex>, } impl Bot { - pub async fn on_new_coin(&mut self, token: NewToken) -> anyhow::Result<()> { - self.tokens.insert( + pub async fn on_new_coin(&self, token: NewToken) -> anyhow::Result<()> { + self.tokens.lock().await.insert( token.mint.clone(), Token { mode: Mode::Observing, @@ -34,8 +34,12 @@ impl Bot { Ok(()) } - pub async fn on_trade(&mut self, trade: Trade) -> anyhow::Result<()> { - let Some(token) = self.tokens.get_mut(&trade.mint) else { + pub async fn on_trade(&self, trade: Trade) -> anyhow::Result<()> { + log::info!("Trade: {trade:?}"); + + let mut tokens = self.tokens.lock().await; + + let Some(token) = tokens.get_mut(&trade.mint) else { log::error!("Token not found on trade: {:?}", trade.mint); return Ok(()); }; @@ -51,8 +55,10 @@ impl Bot { Ok(()) } - pub async fn subscribe(&mut self, mint: &str) -> anyhow::Result<()> { + pub async fn subscribe(&self, mint: &str) -> anyhow::Result<()> { self.ws + .lock() + .await .send(tokio_tungstenite::tungstenite::Message::Text( json!({ "method": "subscribeTokenTrade", @@ -66,8 +72,10 @@ impl Bot { Ok(()) } - pub async fn unsubscribe(&mut self, mint: &str) -> anyhow::Result<()> { + pub async fn unsubscribe(&self, mint: &str) -> anyhow::Result<()> { self.ws + .lock() + .await .send(tokio_tungstenite::tungstenite::Message::Text( json!({ "method": "unsubscribeTokenTrade", @@ -93,28 +101,31 @@ impl Bot { .clone(); Ok(Self { - ws: connect_async("wss://pumpdev.io/ws").await?.0, - executor: Box::new(account.executor()), - accounts, - tokens: HashMap::new(), + ws: Mutex::new(connect_async("wss://pumpdev.io/ws").await?.0), + executor: Mutex::new(Box::new(account.executor())), + accounts: Mutex::new(accounts), + tokens: Mutex::new(HashMap::new()), }) } - pub async fn refresh_account(&mut self) -> anyhow::Result<()> { - let account = self + pub async fn refresh_account(&self) -> anyhow::Result<()> { + let accounts = self.accounts.lock().await; + + let account = accounts .accounts - .accounts - .get(&self.accounts.active) + .get(&accounts.active) .context("Failed to get account")? .clone(); - self.executor = Box::new(account.executor()); + *self.executor.lock().await = Box::new(account.executor()); Ok(()) } - pub async fn initialize_websocket_subscribe(&mut self) -> anyhow::Result<()> { + pub async fn initialize_websocket_subscribe(&self) -> anyhow::Result<()> { self.ws + .lock() + .await .send(tokio_tungstenite::tungstenite::Message::Text( json!({ "method": "subscribeNewToken" }).to_string().into(), )) @@ -123,19 +134,24 @@ impl Bot { Ok(()) } - pub async fn start(&mut self) -> anyhow::Result<()> { + pub async fn start(&self) -> anyhow::Result<()> { self.initialize_websocket_subscribe().await?; - while let Some(msg) = self.ws.next().await { - let msg = msg?; + let mut ws = self.ws.lock().await; + + while 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.on_new_coin(token).await?; + ws = self.ws.lock().await; } Ok(crate::types::PumpDevEvent::Trade(trade)) => { + drop(ws); self.on_trade(trade).await?; + ws = self.ws.lock().await; } Ok(event) => { @@ -143,7 +159,7 @@ impl Bot { } Err(err) => { - log::error!("{err}"); + log::error!("{err}, MSG -> {text}"); } } } diff --git a/src/executor/mod.rs b/src/executor/mod.rs index a7ecb40..61c449a 100644 --- a/src/executor/mod.rs +++ b/src/executor/mod.rs @@ -6,7 +6,7 @@ use crate::account::Account; #[allow(unused_variables)] #[async_trait::async_trait] -pub trait Executor { +pub trait Executor: Send + Sync { async fn buy( &self, mint: String, diff --git a/src/main.rs b/src/main.rs index d4ebe8d..73ba992 100644 --- a/src/main.rs +++ b/src/main.rs @@ -11,7 +11,7 @@ async fn main() -> anyhow::Result<()> { builder.filter_level(log::LevelFilter::Info); builder.init(); - let mut bot = Bot::new().await?; + let bot = Bot::new().await?; bot.start().await } diff --git a/src/types.rs b/src/types.rs index 1898e61..05243ec 100644 --- a/src/types.rs +++ b/src/types.rs @@ -16,11 +16,14 @@ impl<'de> Deserialize<'de> for PumpDevEvent { D: Deserializer<'de>, { let value = Value::deserialize(deserializer)?; + if value.get("txType").is_some() && value.get("name").is_none() { + let trade: Trade = serde_json::from_value(value).map_err(serde::de::Error::custom)?; + return Ok(PumpDevEvent::Trade(trade)); + } - if value.get("txType").is_some() { + if value.get("name").is_some() { let token: NewToken = serde_json::from_value(value).map_err(serde::de::Error::custom)?; - return Ok(PumpDevEvent::Create(token)); } @@ -33,10 +36,8 @@ impl<'de> Deserialize<'de> for PumpDevEvent { client_id: u64, message: String, }, - #[serde(rename = "connectionStatus")] ConnectionStatus { connected: bool, timestamp: u64 }, - #[serde(rename = "subscribed")] Subscribed { method: String }, } @@ -95,16 +96,16 @@ pub enum TradeType { pub struct Trade { pub signature: String, pub mint: String, + + #[serde(rename = "traderPublicKey")] pub trader: String, #[serde(rename = "txType")] pub tx_type: TradeType, - /// Value in lamports #[serde(rename = "solAmount")] pub sol_amount: f64, - /// Value in raw token units #[serde(rename = "tokenAmount")] pub token_amount: f64, @@ -116,6 +117,4 @@ pub struct Trade { #[serde(rename = "vSolInBondingCurve")] pub v_sol_in_bonding_curve: f64, - - pub pool: String, }