Initial boilerplate scaffold for continuum-proxy
This commit is contained in:
@@ -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(())
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Messages sent from this gateway up to `continuum-backend`'s `/ws/edge/v1` route.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ClientMessage {
|
||||
Hello {
|
||||
gateway_id: String,
|
||||
farm_id: String,
|
||||
token: String,
|
||||
version: &'static str,
|
||||
},
|
||||
Heartbeat {
|
||||
seq: u64,
|
||||
},
|
||||
Telemetry {
|
||||
printer_id: String,
|
||||
payload: serde_json::Value,
|
||||
},
|
||||
TaskResult {
|
||||
task_id: String,
|
||||
ok: bool,
|
||||
error: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
/// Messages received from `continuum-backend` over the same connection.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(tag = "type", rename_all = "snake_case")]
|
||||
pub enum ServerMessage {
|
||||
HelloAck {
|
||||
session_id: String,
|
||||
},
|
||||
HeartbeatAck {
|
||||
seq: u64,
|
||||
},
|
||||
Task {
|
||||
task_id: String,
|
||||
kind: String,
|
||||
payload: serde_json::Value,
|
||||
},
|
||||
}
|
||||
Reference in New Issue
Block a user