Support Multiple Upstream WebSocket Connections for HA Setups

Author: angel-ding-cb

src/main.rs

@@ -36,7 +36,7 @@ struct Args {
     listen_addr: SocketAddr,
 
     #[arg(long, env, help = "WebSocket URI of the upstream server to connect to")]
-    upstream_ws: Uri,
+    upstream_ws: Vec<Uri>,
 
     #[arg(
         long,
@@ -144,6 +144,14 @@ async fn main() {
             .expect("failed to setup Prometheus endpoint")
     }
 
+    // Validate that we have at least one upstream URI
+    if args.upstream_ws.is_empty() {
+        error!(message = "no upstream URIs provided");
+        panic!("No upstream URIs provided");
+    }
+
+    info!(message = "using upstream URIs", uris = ?args.upstream_ws);
+
     let metrics = Arc::new(Metrics::default());
     let metrics_clone = metrics.clone();
 
@@ -165,14 +173,29 @@ async fn main() {
     };
 
     let token = CancellationToken::new();
+    let mut subscriber_tasks = Vec::new();
+
+    // Start a subscriber for each upstream URI
+    for (index, uri) in args.upstream_ws.iter().enumerate() {
+        let uri_clone = uri.clone();
+        let listener_clone = listener.clone();
+        let token_clone = token.clone();
+        let metrics_clone = metrics.clone();
+
+        let mut subscriber = WebsocketSubscriber::new(
+            uri_clone,
+            listener_clone,
+            args.subscriber_max_interval,
+            metrics_clone,
+        );
+        
+        let task = tokio::spawn(async move {
+            info!(message = "starting subscriber", index = index, uri = uri_clone.to_string());
+            subscriber.run(token_clone).await;
+        });
 
-    let mut subscriber = WebsocketSubscriber::new(
-        args.upstream_ws,
-        listener,
-        args.subscriber_max_interval,
-        metrics.clone(),
-    );
-    let subscriber_task = subscriber.run(token.clone());
+        subscriber_tasks.push(task);
+    }
 
     let registry = Registry::new(sender, metrics.clone());
 
@@ -194,8 +217,8 @@ async fn main() {
     let mut terminate = signal(SignalKind::terminate()).unwrap();
 
     tokio::select! {
-        _ = subscriber_task => {
-            info!("subscriber task terminated");
+        _ = futures::future::join_all(subscriber_tasks) => {
+            info!("all subscriber tasks terminated");
             token.cancel();
         },
         _ = server_task => {

src/metrics.rs

@@ -29,4 +29,17 @@ pub struct Metrics {
 
     #[metric(describe = "Count of messages received from the upstream source")]
     pub upstream_messages: Gauge,
+
+    // New metrics for multiple upstream connections
+    #[metric(describe = "Number of active upstream connections")]
+    pub upstream_connections: Gauge,
+
+    #[metric(describe = "Number of upstream connection attempts")]
+    pub upstream_connection_attempts: Counter,
+
+    #[metric(describe = "Number of successful upstream connections")]
+    pub upstream_connection_successes: Counter,
+
+    #[metric(describe = "Number of failed upstream connection attempts")]
+    pub upstream_connection_failures: Counter,
 }

src/subscriber.rs

@@ -40,27 +40,52 @@ where
     }
 
     pub async fn run(&mut self, token: CancellationToken) {
-        info!("starting upstream subscription");
+        // Added the URI to the log message for better identification
+        info!(
+            message = "starting upstream subscription",
+            uri = self.uri.to_string()
+        );
         loop {
             select! {
                 _ = token.cancelled() => {
-                    info!("cancelled upstream subscription");
+                    info!(
+                        message = "cancelled upstream subscription",
+                        uri = self.uri.to_string()
+                    );
                     return;
                 }
                 result = self.connect_and_listen() => {
                     match result {
                         Ok(()) => {
-                            info!(message="upstream connection closed");
+                            info!(
+                                message = "upstream connection closed",
+                                uri = self.uri.to_string()
+                            );
                         }
                         Err(e) => {
-                            error!(message="upstream websocket error", error=e.to_string());
+                            // Added URI to the error log for better debugging
+                            error!(
+                                message = "upstream websocket error", 
+                                uri = self.uri.to_string(), 
+                                error = e.to_string()
+                            );
                             self.metrics.upstream_errors.increment(1);
+                            // Decrement the active connections count when connection fails
+                            self.metrics.upstream_connections.decrement(1);
 
                             if let Some(duration) = self.backoff.next_backoff() {
-                                warn!(message="recconecting", seconds=duration.as_secs());
+                                // Added URI to the warning message
+                                warn!(
+                                    message = "reconnecting", 
+                                    uri = self.uri.to_string(), 
+                                    seconds = duration.as_secs()
+                                );
                                 select! {
                                     _ = token.cancelled() => {
-                                        info!(message="cancelled subscriber during backoff");
+                                        info!(
+                                            message = "cancelled subscriber during backoff", 
+                                            uri = self.uri.to_string()
+                                        );
                                         return
                                     }
                                     _ = tokio::time::sleep(duration) => {}
@@ -79,8 +104,31 @@ where
             uri = self.uri.to_string()
         );
 
-        let (ws_stream, _) = connect_async(&self.uri).await?;
-        info!(message = "websocket connection established");
+        // Increment connection attempts counter for metrics
+        self.metrics.upstream_connection_attempts.increment(1);
+
+        // Modified connection with success/failure metrics tracking
+        let (ws_stream, _) = match connect_async(&self.uri).await {
+            Ok(connection) => {
+                // Track successful connections
+                self.metrics.upstream_connection_successes.increment(1);
+                connection
+            },
+            Err(e) => {
+                // Track failed connections
+                self.metrics.upstream_connection_failures.increment(1);
+                return Err(e);
+            }
+        };
+
+        info!(
+            message = "websocket connection established", 
+            uri = self.uri.to_string()
+        );
+
+        // Increment active connections counter
+        self.metrics.upstream_connections.increment(1);
+        // Reset backoff timer on successful connection
         self.backoff.reset();
 
         let (_, mut read) = ws_stream.split();
@@ -89,12 +137,20 @@ where
             match message {
                 Ok(msg) => {
                     let text = msg.to_text()?;
-                    trace!(message = "received message", payload = text);
+                    trace!(
+                        message = "received message", 
+                        uri = self.uri.to_string(), 
+                        payload = text
+                    );
                     self.metrics.upstream_messages.increment(1);
                     (self.handler)(text.into());
                 }
                 Err(e) => {
-                    error!(message = "error receiving message", error = e.to_string());
+                    error!(
+                        message = "error receiving message", 
+                        uri = self.uri.to_string(), 
+                        error = e.to_string()
+                    );
                     return Err(e);
                 }
             }