rusty_common/websocket/bridge/

channel_bridge.rs

//! Channel bridge handler implementation
//!
//! Bridges between WebSocketClient's callback pattern and channel-based consumers.

use super::super::{ExchangeHandler, Message, MessageHandler, WebSocketError, WebSocketResult};
use super::message_router::{MessageRouter, SubscriptionType};
use crate::collections::FxHashMap;
use async_trait::async_trait;
use parking_lot::RwLock;
use std::any::Any;
use std::sync::Arc;
use tokio::sync::mpsc;

/// Channel information for a subscription
pub struct SubscriptionChannel {
    /// Channel sender
    pub sender: mpsc::UnboundedSender<Box<dyn Any + Send>>,
    /// Type of subscription
    pub subscription_type: SubscriptionType,
}

/// Channel bridge handler
pub struct ChannelBridgeHandler<H: ExchangeHandler, R: MessageRouter> {
    /// Inner exchange handler
    inner: H,
    /// Message router
    router: R,
    /// Active subscriptions
    subscriptions: Arc<RwLock<FxHashMap<String, SubscriptionChannel>>>,
}

impl<H: ExchangeHandler, R: MessageRouter> ChannelBridgeHandler<H, R> {
    /// Create a new channel bridge handler
    pub fn new(inner: H, router: R) -> Self {
        Self {
            inner,
            router,
            subscriptions: Arc::new(RwLock::new(FxHashMap::default())),
        }
    }

    /// Register a subscription channel
    pub fn register_subscription(
        &self,
        key: String,
        sender: mpsc::UnboundedSender<Box<dyn Any + Send>>,
        subscription_type: SubscriptionType,
    ) {
        let channel = SubscriptionChannel {
            sender,
            subscription_type,
        };
        self.subscriptions.write().insert(key, channel);
    }

    /// Unregister a subscription
    pub fn unregister_subscription(&self, key: &str) -> Option<SubscriptionChannel> {
        self.subscriptions.write().remove(key)
    }

    /// Get the number of active subscriptions
    pub fn subscription_count(&self) -> usize {
        self.subscriptions.read().len()
    }

    /// Clear all subscriptions
    pub fn clear_subscriptions(&self) {
        self.subscriptions.write().clear();
    }

    /// Get inner handler reference
    pub const fn inner(&self) -> &H {
        &self.inner
    }

    /// Get inner handler mutable reference
    pub const fn inner_mut(&mut self) -> &mut H {
        &mut self.inner
    }
}

#[async_trait]
impl<H: ExchangeHandler, R: MessageRouter> MessageHandler for ChannelBridgeHandler<H, R> {
    async fn on_message(&mut self, message: Message) -> WebSocketResult<()> {
        // Let inner handler process first (for auth, heartbeat, etc.)
        self.inner.on_message(message.clone()).await?;

        // Route to appropriate channel
        if let Some(routed) = self.router.route_message(&message)? {
            let subs = self.subscriptions.read();
            if let Some(channel) = subs.get(&routed.subscription_key) {
                // Send to channel, ignore if receiver dropped
                let _ = channel.sender.send(routed.payload);
            }
        }

        Ok(())
    }

    async fn on_connected(&mut self) -> WebSocketResult<()> {
        self.inner.on_connected().await
    }

    async fn on_disconnected(&mut self) -> WebSocketResult<()> {
        self.inner.on_disconnected().await
    }

    async fn on_error(&mut self, error: WebSocketError) -> WebSocketResult<()> {
        self.inner.on_error(error).await
    }
}

#[async_trait]
impl<H: ExchangeHandler, R: MessageRouter> ExchangeHandler for ChannelBridgeHandler<H, R> {
    fn exchange(&self) -> crate::types::Exchange {
        self.inner.exchange()
    }

    async fn get_subscriptions(&self) -> WebSocketResult<Vec<Message>> {
        self.inner.get_subscriptions().await
    }

    async fn authenticate(&mut self) -> WebSocketResult<Option<Message>> {
        self.inner.authenticate().await
    }

    fn requires_auth(&self) -> bool {
        self.inner.requires_auth()
    }

    fn get_heartbeat(&self) -> Option<Message> {
        self.inner.get_heartbeat()
    }

    fn heartbeat_interval(&self) -> Option<std::time::Duration> {
        self.inner.heartbeat_interval()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::Exchange;

    struct TestRouter;

    impl MessageRouter for TestRouter {
        fn route_message(
            &self,
            message: &Message,
        ) -> WebSocketResult<Option<super::super::RoutedMessage>> {
            match message {
                Message::Text(text) if text.contains("trade") => {
                    Ok(Some(super::super::RoutedMessage {
                        subscription_key: "trades:TEST".to_string(),
                        subscription_type: SubscriptionType::Trades,
                        payload: Box::new(text.clone()),
                    }))
                }
                _ => Ok(None),
            }
        }
    }

    struct TestHandler;

    #[async_trait]
    impl MessageHandler for TestHandler {
        async fn on_message(&mut self, _message: Message) -> WebSocketResult<()> {
            Ok(())
        }
    }

    #[async_trait]
    impl ExchangeHandler for TestHandler {
        fn exchange(&self) -> Exchange {
            Exchange::Binance
        }

        async fn get_subscriptions(&self) -> WebSocketResult<Vec<Message>> {
            Ok(vec![])
        }
    }

    #[tokio::test]
    async fn test_channel_bridge() {
        let handler = TestHandler;
        let router = TestRouter;
        let mut bridge = ChannelBridgeHandler::new(handler, router);

        // Create channel
        let (tx, mut rx) = mpsc::unbounded_channel();
        bridge.register_subscription("trades:TEST".to_string(), tx, SubscriptionType::Trades);

        // Send trade message
        let trade_msg = Message::Text("trade data".into());
        bridge.on_message(trade_msg).await.unwrap();

        // Check received
        let received = rx.try_recv().unwrap();
        let text = received
            .downcast_ref::<smartstring::alias::String>()
            .unwrap();
        assert_eq!(text.as_str(), "trade data");

        // Non-trade message should not be routed
        let other_msg = Message::Text("other data".into());
        bridge.on_message(other_msg).await.unwrap();
        assert!(rx.try_recv().is_err());
    }
}