Filter fill updates on subscription

service-mango-fills/src/main.rs

@@ -3,9 +3,9 @@ use anchor_lang::prelude::Pubkey;
 use bytemuck::cast_slice;
 use client::{Client, MangoGroupContext};
 use futures_channel::mpsc::{unbounded, UnboundedSender};
-use futures_util::{pin_mut, SinkExt, StreamExt};
+use futures_util::{pin_mut, SinkExt, StreamExt, future::{self, Ready}, TryStreamExt};
 use log::*;
-use std::{collections::HashMap, fs::File, io::Read, net::SocketAddr, sync::Arc, sync::Mutex, time::Duration, convert::identity, str::FromStr};
+use std::{collections::{HashMap, HashSet}, fs::File, io::Read, net::SocketAddr, sync::Arc, sync::Mutex, time::Duration, convert::identity, str::FromStr};
 use tokio::{
     net::{TcpListener, TcpStream},
     pin,
@@ -13,7 +13,7 @@ use tokio::{
 use tokio_tungstenite::tungstenite::{protocol::Message, Error};
 
 use serde::Deserialize;
-use solana_geyser_connector_lib::{metrics::{MetricType, MetricU64}, FilterConfig, fill_event_filter::SerumFillCheckpoint};
+use solana_geyser_connector_lib::{metrics::{MetricType, MetricU64}, FilterConfig, fill_event_filter::SerumFillCheckpoint, StatusResponse};
 use solana_geyser_connector_lib::{
     fill_event_filter::{self, FillCheckpoint, FillEventFilterMessage},
     grpc_plugin_source, metrics, websocket_source, MetricsConfig, SourceConfig,
@@ -21,12 +21,42 @@ use solana_geyser_connector_lib::{
 
 type CheckpointMap = Arc<Mutex<HashMap<String, FillCheckpoint>>>;
 type SerumCheckpointMap = Arc<Mutex<HashMap<String, SerumFillCheckpoint>>>;
-type PeerMap = Arc<Mutex<HashMap<SocketAddr, UnboundedSender<Message>>>>;
+type PeerMap = Arc<Mutex<HashMap<SocketAddr, Peer>>>;
+
+#[derive(Clone, Debug, Deserialize)]
+#[serde(tag = "command")]
+pub enum Command {
+    #[serde(rename = "subscribe")]
+    Subscribe(SubscribeCommand),
+    #[serde(rename = "unsubscribe")]
+    Unsubscribe(UnsubscribeCommand),
+    #[serde(rename = "getMarkets")]
+    GetMarkets,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct SubscribeCommand {
+    pub market_id: String,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct UnsubscribeCommand {
+    pub market_id: String,
+}
+
+#[derive(Clone, Debug)]
+pub struct Peer {
+    pub sender: UnboundedSender<Message>,
+    pub subscriptions: HashSet<String>,
+}
 
 async fn handle_connection_error(
     checkpoint_map: CheckpointMap,
     serum_checkpoint_map: SerumCheckpointMap,
     peer_map: PeerMap,
+    market_ids: HashMap<String, String>,
     raw_stream: TcpStream,
     addr: SocketAddr,
     metrics_opened_connections: MetricU64,
@@ -34,7 +64,7 @@ async fn handle_connection_error(
 ) {
     metrics_opened_connections.clone().increment();
 
-    let result = handle_connection(checkpoint_map, serum_checkpoint_map, peer_map.clone(), raw_stream, addr).await;
+    let result = handle_connection(checkpoint_map, serum_checkpoint_map, peer_map.clone(), market_ids, raw_stream, addr).await;
     if result.is_err() {
         error!("connection {} error {}", addr, result.unwrap_err());
     };
@@ -48,42 +78,156 @@ async fn handle_connection(
     checkpoint_map: CheckpointMap,
     serum_checkpoint_map: SerumCheckpointMap,
     peer_map: PeerMap,
+    market_ids: HashMap<String, String>,
     raw_stream: TcpStream,
     addr: SocketAddr,
 ) -> Result<(), Error> {
     info!("ws connected: {}", addr);
     let ws_stream = tokio_tungstenite::accept_async(raw_stream).await?;
-    let (mut ws_tx, _ws_rx) = ws_stream.split();
+    let (ws_tx, ws_rx) = ws_stream.split();
 
     // 1: publish channel in peer map
     let (chan_tx, chan_rx) = unbounded();
     {
-        peer_map.lock().unwrap().insert(addr, chan_tx);
-        info!("ws published: {}", addr);
+        peer_map.lock().unwrap().insert(
+            addr,
+            Peer {
+                sender: chan_tx,
+                subscriptions: HashSet::<String>::new(),
+            },
+        );
     }
 
-    // todo: add subscribe logic
-    // 2: send initial checkpoint
-    {
-        let checkpoint_map_copy = checkpoint_map.lock().unwrap().clone();
-        for (_, ckpt) in checkpoint_map_copy.iter() {
-            ws_tx
-                .feed(Message::Text(serde_json::to_string(ckpt).unwrap()))
-                .await?;
-        }
-    }
-    info!("ws ckpt sent: {}", addr);
-    ws_tx.flush().await?;
-
-    // 3: forward all events from channel to peer socket
+    let receive_commands = ws_rx.try_for_each(|msg| {
+        handle_commands(
+            addr,
+            msg,
+            peer_map.clone(),
+            checkpoint_map.clone(),
+            serum_checkpoint_map.clone(),
+            market_ids.clone(),
+        )
+    });
     let forward_updates = chan_rx.map(Ok).forward(ws_tx);
-    pin_mut!(forward_updates);
-    forward_updates.await?;
 
+    pin_mut!(receive_commands, forward_updates);
+    future::select(receive_commands, forward_updates).await;
+
+    peer_map.lock().unwrap().remove(&addr);
     info!("ws disconnected: {}", &addr);
     Ok(())
 }
 
+fn handle_commands(
+    addr: SocketAddr,
+    msg: Message,
+    peer_map: PeerMap,
+    checkpoint_map: CheckpointMap,
+    serum_checkpoint_map: SerumCheckpointMap,
+    market_ids: HashMap<String, String>,
+) -> Ready<Result<(), Error>> {
+    let msg_str = msg.clone().into_text().unwrap();
+    let command: Result<Command, serde_json::Error> = serde_json::from_str(&msg_str);
+    let mut peers = peer_map.lock().unwrap();
+    let peer = peers.get_mut(&addr).expect("peer should be in map");
+    match command {
+        Ok(Command::Subscribe(cmd)) => {
+            let market_id = cmd.clone().market_id;
+            match market_ids.get(&market_id) {
+                None => {
+                    let res = StatusResponse {
+                        success: false,
+                        message: "market not found",
+                    };
+                    peer.sender
+                        .unbounded_send(Message::Text(serde_json::to_string(&res).unwrap()))
+                        .unwrap();
+                    return future::ok(());
+                }
+                _ => {}
+            }
+            let subscribed = peer.subscriptions.insert(market_id.clone());
+
+            let res = if subscribed {
+                StatusResponse {
+                    success: true,
+                    message: "subscribed",
+                }
+            } else {
+                StatusResponse {
+                    success: false,
+                    message: "already subscribed",
+                }
+            };
+            peer.sender
+                .unbounded_send(Message::Text(serde_json::to_string(&res).unwrap()))
+                .unwrap();
+
+            if subscribed {
+                // todo: this is janky af
+                let checkpoint_map = checkpoint_map.lock().unwrap();
+                let serum_checkpoint_map = serum_checkpoint_map.lock().unwrap();
+                let checkpoint = checkpoint_map.get(&market_id);
+                match checkpoint {
+                    Some(checkpoint) => {
+                        peer.sender
+                            .unbounded_send(Message::Text(
+                                serde_json::to_string(&checkpoint).unwrap(),
+                            ))
+                            .unwrap();
+                    }
+                    None => match serum_checkpoint_map.get(&market_id) {
+                        Some(checkpoint) => {
+                            peer.sender
+                                .unbounded_send(Message::Text(
+                                    serde_json::to_string(&checkpoint).unwrap(),
+                                ))
+                                .unwrap();
+                        }
+                        None => info!("no checkpoint available on client subscription"), // todo: what to do here?
+                    }
+                }
+            }
+        }
+        Ok(Command::Unsubscribe(cmd)) => {
+            info!("unsubscribe {}", cmd.market_id);
+            let unsubscribed = peer.subscriptions.remove(&cmd.market_id);
+            let res = if unsubscribed {
+                StatusResponse {
+                    success: true,
+                    message: "unsubscribed",
+                }
+            } else {
+                StatusResponse {
+                    success: false,
+                    message: "not subscribed",
+                }
+            };
+            peer.sender
+                .unbounded_send(Message::Text(serde_json::to_string(&res).unwrap()))
+                .unwrap();
+        }
+        Ok(Command::GetMarkets) => {
+            info!("getMarkets");
+            peer.sender
+                .unbounded_send(Message::Text(serde_json::to_string(&market_ids).unwrap()))
+                .unwrap();
+        }
+        Err(err) => {
+            info!("error deserializing user input {:?}", err);
+            let res = StatusResponse {
+                success: false,
+                message: "invalid input",
+            };
+            peer.sender
+                .unbounded_send(Message::Text(serde_json::to_string(&res).unwrap()))
+                .unwrap();
+        }
+    };
+
+    future::ok(())
+}
+
 #[derive(Clone, Debug, Deserialize)]
 pub struct Config {
     pub source: SourceConfig,
@@ -170,6 +314,19 @@ async fn main() -> anyhow::Result<()> {
         })
         .collect();
 
+    let a: Vec<(String, String)> = group_context
+        .serum3_markets
+        .iter()
+        .map(|(_, context)| (context.market.name().to_owned(), context.market.serum_market_external.to_string())).collect();
+    let b: Vec<(String, String)> = group_context
+        .perp_markets
+        .iter()
+        .map(|(_, context)| (context.market.name().to_owned(), context.address.to_string())).collect();
+    let market_pubkey_strings: HashMap<String, String> = [a, b]
+        .concat()
+        .into_iter()
+        .collect();
+
     let (account_write_queue_sender, slot_queue_sender, fill_receiver) =
     fill_event_filter::init(perp_queue_pks.clone(), serum_queue_pks.clone(), metrics_tx.clone()).await?;
 
@@ -190,15 +347,18 @@ async fn main() -> anyhow::Result<()> {
                 FillEventFilterMessage::Update(update) => {
                     debug!("ws update {} {:?} fill", update.market, update.status);
                     let mut peer_copy = peers_ref_thread.lock().unwrap().clone();
-                    for (k, v) in peer_copy.iter_mut() {
-                        trace!("  > {}", k);
-                        let json = serde_json::to_string(&update);
-                        let result = v.send(Message::Text(json.unwrap())).await;
-                        if result.is_err() {
-                            error!(
-                                "ws update {} {:?} fill could not reach {}",
-                                update.market, update.status, k
-                            );
+                    for (addr, peer) in peer_copy.iter_mut() {
+                        let json = serde_json::to_string(&update).unwrap();
+                        
+                        // only send updates if the peer is subscribed
+                        if peer.subscriptions.contains(&update.market) {
+                            let result = peer.sender.send(Message::Text(json)).await;
+                            if result.is_err() {
+                                error!(
+                                    "ws update {} fill could not reach {}",
+                                    update.market, addr
+                                );
+                            }
                         }
                     }
                 }
@@ -211,15 +371,18 @@ async fn main() -> anyhow::Result<()> {
                 FillEventFilterMessage::SerumUpdate(update) => {
                     debug!("ws update {} {:?} serum fill", update.market, update.status);
                     let mut peer_copy = peers_ref_thread.lock().unwrap().clone();
-                    for (k, v) in peer_copy.iter_mut() {
-                        trace!("  > {}", k);
-                        let json = serde_json::to_string(&update);
-                        let result = v.send(Message::Text(json.unwrap())).await;
-                        if result.is_err() {
-                            error!(
-                                "ws update {} {:?} serum fill could not reach {}",
-                                update.market, update.status, k
-                            );
+                    for (addr, peer) in peer_copy.iter_mut() {
+                        let json = serde_json::to_string(&update).unwrap();
+                        
+                        // only send updates if the peer is subscribed
+                        if peer.subscriptions.contains(&update.market) {
+                            let result = peer.sender.send(Message::Text(json)).await;
+                            if result.is_err() {
+                                error!(
+                                    "ws update {} fill could not reach {}",
+                                    update.market, addr
+                                );
+                            }
                         }
                     }
                 }
@@ -243,6 +406,7 @@ async fn main() -> anyhow::Result<()> {
                 checkpoints.clone(),
                 serum_checkpoints.clone(),
                 peers.clone(),
+                market_pubkey_strings.clone(),
                 stream,
                 addr,
                 metrics_opened_connections.clone(),