Initial boilerplate scaffold for continuum-proxy

src/uplink/mod.rs

@@ -0,0 +1,113 @@
+mod protocol;
+
+pub use protocol::{ClientMessage, ServerMessage};
+
+use std::time::Duration;
+
+use futures_util::{SinkExt, StreamExt};
+use tokio::sync::mpsc;
+use tokio_tungstenite::tungstenite::Message;
+use tracing::{info, warn};
+
+use crate::config::Config;
+
+const MIN_BACKOFF: Duration = Duration::from_secs(1);
+const MAX_BACKOFF: Duration = Duration::from_secs(60);
+const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(15);
+
+/// Runs the resilient uplink loop forever: connect, authenticate, exchange
+/// heartbeats and telemetry/task messages, and reconnect with exponential
+/// backoff (plus jitter) whenever the connection drops.
+pub async fn run(
+    config: Config,
+    mut telemetry_rx: mpsc::Receiver<ClientMessage>,
+    task_tx: mpsc::Sender<ServerMessage>,
+) {
+    let mut backoff = MIN_BACKOFF;
+
+    loop {
+        info!(url = %config.uplink_url, "connecting to cloud uplink");
+
+        match connect_and_serve(&config, &mut telemetry_rx, &task_tx).await {
+            Ok(()) => {
+                info!("uplink connection closed cleanly");
+                backoff = MIN_BACKOFF;
+            }
+            Err(err) => {
+                warn!(?err, backoff_secs = backoff.as_secs(), "uplink connection failed, retrying");
+            }
+        }
+
+        let jitter = Duration::from_millis(rand::random::<u64>() % 500);
+        tokio::time::sleep(backoff + jitter).await;
+        backoff = (backoff * 2).min(MAX_BACKOFF);
+    }
+}
+
+async fn connect_and_serve(
+    config: &Config,
+    telemetry_rx: &mut mpsc::Receiver<ClientMessage>,
+    task_tx: &mpsc::Sender<ServerMessage>,
+) -> anyhow::Result<()> {
+    let (ws_stream, _resp) = tokio_tungstenite::connect_async(&config.uplink_url).await?;
+    let (mut write, mut read) = ws_stream.split();
+
+    send(&mut write, &ClientMessage::Hello {
+        gateway_id: config.gateway_id.clone(),
+        farm_id: config.farm_id.clone(),
+        token: config.gateway_token.clone(),
+        version: env!("CARGO_PKG_VERSION"),
+    })
+    .await?;
+
+    let mut heartbeat = tokio::time::interval(HEARTBEAT_INTERVAL);
+    let mut seq: u64 = 0;
+    heartbeat.tick().await; // consume the immediate first tick
+
+    loop {
+        tokio::select! {
+            _ = heartbeat.tick() => {
+                seq += 1;
+                send(&mut write, &ClientMessage::Heartbeat { seq }).await?;
+            }
+
+            Some(msg) = telemetry_rx.recv() => {
+                send(&mut write, &msg).await?;
+            }
+
+            frame = read.next() => {
+                match frame {
+                    Some(Ok(Message::Text(text))) => {
+                        match serde_json::from_str::<ServerMessage>(&text) {
+                            Ok(server_msg) => {
+                                if task_tx.send(server_msg).await.is_err() {
+                                    return Ok(());
+                                }
+                            }
+                            Err(err) => warn!(?err, "failed to decode server message"),
+                        }
+                    }
+                    Some(Ok(Message::Ping(payload))) => {
+                        write.send(Message::Pong(payload)).await?;
+                    }
+                    Some(Ok(Message::Close(frame))) => {
+                        info!(?frame, "server closed uplink");
+                        return Ok(());
+                    }
+                    Some(Ok(_)) => {}
+                    Some(Err(err)) => return Err(err.into()),
+                    None => return Ok(()),
+                }
+            }
+        }
+    }
+}
+
+async fn send(
+    write: &mut (impl SinkExt<Message, Error = tokio_tungstenite::tungstenite::Error> + Unpin),
+    msg: &ClientMessage,
+) -> anyhow::Result<()> {
+    let payload = serde_json::to_string(msg)?;
+    write.send(Message::Text(payload)).await?;
+    Ok(())
+}