Refactor websockets; add device stream events

Introduce a unified websocket helper and refactor websocket handling, plus add device stream event support across services.

Key changes:
- New apps/api/src/websocket.rs: common Redis setup (setup_ws_resources) and PushInfo decoding (decode_push_message).
- Refactor websocket_prices and websocket_stream: remove duplicated Redis setup, improve message handling, move configs out of Mutex where safe, better logging/error handling, and simplify stream loops.
- Split price handling into a dedicated PriceHandler for websocket_stream to encapsulate asset tracking, subscription logic, and payload building.
- Improve apps/api/src/main.rs: better CLI service parsing with strum iteration and unified Rocket launch path.
- Add device stream feature: DeviceStreamPayload/DeviceStreamEvent types (crates/streamer), queue name and producer API for DeviceStreamEvents, daemon consumer DeviceStreamConsumer to publish device-scoped stream events to cacher.
- Add CacherClient::publish helper to publish JSON messages to Redis channels.
- Various small API/serialization and testkit cleanups.

Why: removes duplication, isolates responsibilities (price handling, Redis setup), adds support for pushing device-targeted stream events, and improves robustness and logging for websocket components.

Author: gemcoder21

apps/api/src/main.rs

@@ -22,10 +22,12 @@ mod swap;
 mod transactions;
 mod wallets;
 mod webhooks;
+mod websocket;
 mod websocket_prices;
 mod websocket_stream;
 
 use std::{str::FromStr, sync::Arc};
+use strum::IntoEnumIterator;
 
 use ::fiat::FiatClient;
 use ::nft::{NFTClient, NFTProviderConfig};
@@ -268,7 +270,7 @@ async fn rocket_ws_prices(settings: Settings) -> Rocket<Build> {
     };
     rocket::build()
         .manage(Arc::new(Mutex::new(price_client)))
-        .manage(Arc::new(Mutex::new(price_observer_config)))
+        .manage(Arc::new(price_observer_config))
         .mount("/", routes![websocket_prices::ws_health])
         .mount("/v1/ws", routes![websocket_prices::ws_prices])
         .register("/", catchers![catchers::default_catcher])
@@ -292,7 +294,7 @@ async fn rocket_ws_stream(settings: Settings) -> Rocket<Build> {
         .manage(auth_config)
         .manage(database)
         .manage(Arc::new(Mutex::new(price_client)))
-        .manage(Arc::new(Mutex::new(stream_observer_config)))
+        .manage(Arc::new(stream_observer_config))
         .mount("/v2/devices", routes![websocket_stream::ws_stream])
         .mount("/", routes![websocket_stream::ws_health])
         .register("/", catchers![catchers::default_catcher])
@@ -302,25 +304,22 @@ async fn rocket_ws_stream(settings: Settings) -> Rocket<Build> {
 async fn main() {
     let settings = Settings::new().unwrap();
 
-    let service = std::env::args().nth(1).unwrap_or_default();
-    let service = APIService::from_str(service.as_str()).ok().unwrap_or(APIService::Api);
+    let service = match std::env::args().nth(1) {
+        Some(arg) => APIService::from_str(&arg).unwrap_or_else(|_| {
+            let services: Vec<_> = APIService::iter().map(|s| format!("api {}", s.as_ref())).collect();
+            panic!("unknown service: {arg}\nAvailable:\n {}", services.join("\n "))
+        }),
+        None => APIService::Api,
+    };
 
     println!("api start service: {}", service.as_ref());
 
-    match service {
-        APIService::WebsocketPrices => {
-            let rocket_api = rocket_ws_prices(settings.clone()).await;
-            rocket_api.launch().await.expect("Failed to launch Rocket");
-        }
-        APIService::WebsocketStream => {
-            let rocket_api = rocket_ws_stream(settings.clone()).await;
-            rocket_api.launch().await.expect("Failed to launch Rocket");
-        }
-        APIService::Api => {
-            let rocket_api = rocket_api(settings.clone()).await;
-            rocket_api.launch().await.expect("Failed to launch Rocket");
-        }
-    }
+    let rocket = match service {
+        APIService::Api => rocket_api(settings).await,
+        APIService::WebsocketPrices => rocket_ws_prices(settings).await,
+        APIService::WebsocketStream => rocket_ws_stream(settings).await,
+    };
+    rocket.launch().await.expect("Failed to launch Rocket");
 }
 
 #[cfg(test)]

apps/api/src/websocket.rs

@@ -0,0 +1,21 @@
+use std::error::Error;
+
+use redis::aio::MultiplexedConnection;
+use redis::{PushInfo, PushKind};
+use rocket_ws::stream::DuplexStream;
+use tokio::sync::mpsc::UnboundedReceiver;
+
+pub fn decode_push_message(message: &PushInfo) -> Option<(&str, &[u8])> {
+    match (&message.kind, message.data.as_slice()) {
+        (PushKind::Message, [redis::Value::BulkString(channel), redis::Value::BulkString(value)]) => Some((std::str::from_utf8(channel).ok()?, value)),
+        _ => None,
+    }
+}
+
+pub async fn setup_ws_resources(redis_url: &str, stream: DuplexStream) -> Result<(DuplexStream, MultiplexedConnection, UnboundedReceiver<PushInfo>), Box<dyn Error + Send + Sync>> {
+    let client = redis::Client::open(redis_url)?;
+    let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
+    let config = redis::AsyncConnectionConfig::new().set_push_sender(tx);
+    let redis_connection = client.get_multiplexed_async_connection_with_config(&config).await?;
+    Ok((stream, redis_connection, rx))
+}

apps/api/src/websocket_prices/client.rs

@@ -4,21 +4,23 @@ use std::sync::Arc;
 
 use pricer::PriceClient;
 use primitives::{AssetId, AssetPrice, AssetPriceInfo, WebSocketPriceAction, WebSocketPriceActionType, WebSocketPricePayload, asset::AssetHashSetExt};
+use redis::PushInfo;
 use redis::aio::MultiplexedConnection;
-use redis::{PushInfo, PushKind};
 use rocket::futures::SinkExt;
 use rocket::serde::json::serde_json;
 use rocket::tokio::sync::Mutex;
 use rocket_ws::Message;
 use rocket_ws::stream::DuplexStream;
 
+use crate::websocket::decode_push_message;
+
 pub struct PriceObserverConfig {
     pub redis_url: String,
 }
 
 pub struct PriceObserverClient {
-    pub price_client: Arc<Mutex<PriceClient>>,
-    pub assets: HashSet<AssetId>,
+    price_client: Arc<Mutex<PriceClient>>,
+    assets: HashSet<AssetId>,
     prices_to_publish: HashMap<String, AssetPrice>,
     interval: rocket::tokio::time::Interval,
 }
@@ -37,105 +39,82 @@ impl PriceObserverClient {
         self.interval.tick().await;
     }
 
-    pub fn get_asset_ids(&self) -> Vec<String> {
-        self.assets.iter().map(|id| id.to_string()).collect()
-    }
-
-    pub fn add_price_to_publish(&mut self, price: AssetPrice) {
-        self.prices_to_publish.insert(price.asset_id.to_string(), price);
-    }
-
-    pub fn clear_prices_to_publish(&mut self) {
-        self.prices_to_publish.clear();
+    pub fn take_prices(&mut self) -> Vec<AssetPrice> {
+        self.prices_to_publish.drain().map(|(_, v)| v).collect()
     }
 
-    pub fn get_prices_to_publish(&self) -> Vec<AssetPrice> {
-        self.prices_to_publish.values().cloned().collect()
+    fn get_channel_ids(&self) -> Vec<String> {
+        self.assets.iter().map(|id| id.to_string()).collect()
     }
 
-    pub async fn fetch_payload_data(&mut self, fetch_rates: bool) -> Result<WebSocketPricePayload, Box<dyn Error + Send + Sync>> {
-        let price_client_clone_prices = Arc::clone(&self.price_client);
-        let assets_clone_prices = self.assets.clone();
-        let prices = price_client_clone_prices
-            .lock()
-            .await
-            .get_cache_prices(assets_clone_prices.ids())
+    async fn fetch_payload(&self, fetch_rates: bool) -> Result<WebSocketPricePayload, Box<dyn Error + Send + Sync>> {
+        let client = self.price_client.lock().await;
+        let prices = client
+            .get_cache_prices(self.assets.ids())
             .await?
             .into_iter()
             .map(|x| x.as_asset_price_primitive())
             .collect();
-
-        if fetch_rates {
-            let rates = self.price_client.lock().await.get_cache_fiat_rates().await?;
-            Ok(WebSocketPricePayload { prices, rates })
-        } else {
-            Ok(WebSocketPricePayload { prices, rates: vec![] })
-        }
-    }
-
-    pub async fn build_and_send_payload(&mut self, stream: &mut DuplexStream, payload: WebSocketPricePayload) -> Result<(), Box<dyn Error + Send + Sync>> {
-        let text = serde_json::to_string(&payload)?;
-        let item = Message::Text(text);
-        Ok(stream.send(item).await?)
+        let rates = if fetch_rates { client.get_cache_fiat_rates().await? } else { vec![] };
+        Ok(WebSocketPricePayload { prices, rates })
     }
 
     pub async fn handle_ws_message(
         &mut self,
-        message: rocket_ws::Message,
+        message: Message,
         redis_connection: &mut MultiplexedConnection,
         stream: &mut DuplexStream,
     ) -> Result<(), Box<dyn Error + Send + Sync>> {
         match message {
             Message::Binary(data) => self.handle_message_payload(data, redis_connection, stream).await,
             Message::Text(text) => self.handle_message_payload(text.into_bytes(), redis_connection, stream).await.or(Ok(())),
             Message::Ping(data) => Ok(stream.send(Message::Pong(data)).await?),
-            Message::Close(_) => {
-                println!("Client closed connection gracefully");
-                Ok(())
-            }
+            Message::Close(_) => Ok(()),
             Message::Pong(_) | Message::Frame(_) => Ok(()),
         }
     }
 
-    pub async fn handle_message_payload(
-        &mut self,
-        data: Vec<u8>,
-        redis_connection: &mut MultiplexedConnection,
-        stream: &mut DuplexStream,
-    ) -> Result<(), Box<dyn Error + Send + Sync>> {
+    async fn handle_message_payload(&mut self, data: Vec<u8>, redis_connection: &mut MultiplexedConnection, stream: &mut DuplexStream) -> Result<(), Box<dyn Error + Send + Sync>> {
         let action = serde_json::from_slice::<WebSocketPriceAction>(&data)?;
         let new_assets: HashSet<AssetId> = action.assets.iter().cloned().collect();
 
-        match action.action {
+        let needs_rates = match action.action {
             WebSocketPriceActionType::Subscribe => {
+                let old_channels = self.get_channel_ids();
                 self.assets.clear();
+                self.prices_to_publish.clear();
+                if !old_channels.is_empty() {
+                    redis_connection.unsubscribe(old_channels).await?;
+                }
                 self.assets.extend(new_assets);
+                true
             }
             WebSocketPriceActionType::Add => {
                 self.assets.extend(new_assets);
+                false
             }
-        }
+        };
 
-        let asset_ids = self.assets.ids();
-        let _ = self.price_client.lock().await.track_observed_assets(&asset_ids).await;
+        let _ = self.price_client.lock().await.track_observed_assets(&self.assets.ids()).await;
 
-        let needs_rates = action.action == WebSocketPriceActionType::Subscribe;
-        let payload = self.fetch_payload_data(needs_rates).await?;
+        let payload = self.fetch_payload(needs_rates).await?;
+        self.send_payload(stream, payload).await?;
 
-        self.build_and_send_payload(stream, payload).await?;
+        redis_connection.subscribe(self.get_channel_ids()).await?;
+        Ok(())
+    }
 
-        Ok(redis_connection.subscribe(self.get_asset_ids()).await?)
+    pub async fn send_payload(&self, stream: &mut DuplexStream, payload: WebSocketPricePayload) -> Result<(), Box<dyn Error + Send + Sync>> {
+        let text = serde_json::to_string(&payload)?;
+        Ok(stream.send(Message::Text(text)).await?)
     }
 
-    pub fn handle_redis_message(&mut self, message: &PushInfo) -> Result<(), String> {
-        match (message.kind.clone(), message.data.last()) {
-            (PushKind::Message, Some(redis::Value::BulkString(value))) => {
-                let asset_price_info = serde_json::from_slice::<AssetPriceInfo>(value).map_err(|e| format!("Failed to deserialize AssetPrice: {e}"))?;
-                let asset_price = asset_price_info.as_asset_price_primitive();
-                self.add_price_to_publish(asset_price);
-                Ok(())
-            }
-            _ => Ok(()),
-        }
+    pub fn handle_redis_message(&mut self, message: &PushInfo) -> Result<(), Box<dyn Error + Send + Sync>> {
+        let Some((_, value)) = decode_push_message(message) else {
+            return Ok(());
+        };
+        let info = serde_json::from_slice::<AssetPriceInfo>(value)?;
+        self.prices_to_publish.insert(info.asset_id.to_string(), info.as_asset_price_primitive());
+        Ok(())
     }
 }

apps/api/src/websocket_prices/mod.rs

@@ -9,18 +9,17 @@ use rocket_ws::{Channel, WebSocket};
 mod client;
 mod stream;
 
-pub use client::{PriceObserverClient, PriceObserverConfig};
-use stream::Stream;
+pub use client::PriceObserverConfig;
 
 #[rocket::get("/prices")]
-pub async fn ws_prices(ws: WebSocket, price_client: &State<Arc<Mutex<PriceClient>>>, config: &State<Arc<Mutex<PriceObserverConfig>>>) -> Channel<'static> {
+pub async fn ws_prices(ws: WebSocket, price_client: &State<Arc<Mutex<PriceClient>>>, config: &State<Arc<PriceObserverConfig>>) -> Channel<'static> {
     let price_client = price_client.inner().clone();
-    let redis_url = config.lock().await.redis_url.clone();
+    let redis_url = config.redis_url.clone();
 
-    ws.channel(move |stream| {
+    ws.channel(move |ws_stream| {
         Box::pin(async move {
-            let mut observer = PriceObserverClient::new(price_client.clone());
-            Stream::new_stream(&redis_url, &mut observer, stream).await;
+            let mut observer = client::PriceObserverClient::new(price_client);
+            stream::new_stream(&redis_url, &mut observer, ws_stream).await;
             Ok::<(), rocket_ws::result::Error>(())
         })
     })