116 lines
4.5 KiB
Rust
116 lines
4.5 KiB
Rust
#[derive(Debug, Deserialize)]
|
|
struct WsQuery {
|
|
token: Option<String>,
|
|
}
|
|
async fn websocket(
|
|
State(state): State<AppState>,
|
|
Query(query): Query<WsQuery>,
|
|
ws: WebSocketUpgrade,
|
|
) -> Result<Response, AppError> {
|
|
let expected = state.config.app_token.trim();
|
|
if !expected.is_empty() && query.token.as_deref() != Some(expected) {
|
|
return Err(AppError::Unauthorized);
|
|
}
|
|
Ok(ws.on_upgrade(move |socket| websocket_loop(state, socket)))
|
|
}
|
|
|
|
fn control_plan_ws_message(snapshot: &crate::state::ControlPlanSnapshot) -> Value {
|
|
json!({
|
|
"event": "control_plan.updated",
|
|
"timestamp": Utc::now(),
|
|
"data": {
|
|
"revision": snapshot.revision,
|
|
"plan": snapshot.plan.as_ref(),
|
|
}
|
|
})
|
|
}
|
|
|
|
async fn send_ws_bootstrap(state: &AppState, socket: &mut WebSocket) -> Result<Option<u64>, ()> {
|
|
let (message, revision) = match build_bootstrap(state).await {
|
|
Ok(data) => {
|
|
let revision = Some(data.control_plan_revision);
|
|
(
|
|
json!({"event":"bootstrap","timestamp":Utc::now(),"data":data}),
|
|
revision,
|
|
)
|
|
}
|
|
Err(err) => (
|
|
json!({"event":"error","timestamp":Utc::now(),"data":{"message":err.to_string()}}),
|
|
None,
|
|
),
|
|
};
|
|
socket
|
|
.send(Message::Text(message.to_string()))
|
|
.await
|
|
.map_err(|_| ())?;
|
|
Ok(revision)
|
|
}
|
|
|
|
async fn websocket_loop(state: AppState, mut socket: WebSocket) {
|
|
// Subscribe before building the bootstrap so state changes during bootstrap generation are
|
|
// queued and can be applied immediately after the first frame.
|
|
let mut receiver = state.events.subscribe();
|
|
let mut control_plan = state.subscribe_control_plan();
|
|
|
|
let bootstrap_revision = match send_ws_bootstrap(&state, &mut socket).await {
|
|
Ok(revision) => revision,
|
|
Err(()) => return,
|
|
};
|
|
let mut last_control_plan_revision = bootstrap_revision;
|
|
|
|
// If the watch value is exactly the plan embedded in bootstrap, mark it seen to avoid a
|
|
// duplicate control_plan.updated frame. A newer revision remains pending and is sent below.
|
|
let current_revision = control_plan
|
|
.borrow()
|
|
.as_ref()
|
|
.map(|snapshot| snapshot.revision);
|
|
if current_revision == bootstrap_revision {
|
|
control_plan.borrow_and_update();
|
|
}
|
|
|
|
loop {
|
|
tokio::select! {
|
|
event = receiver.recv() => {
|
|
match event {
|
|
Ok(event) => {
|
|
if let Ok(text) = serde_json::to_string(&event) {
|
|
if socket.send(Message::Text(text)).await.is_err() { break; }
|
|
}
|
|
}
|
|
Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => {
|
|
tracing::warn!(skipped, "websocket client lagged; sending full bootstrap resync");
|
|
// Drop the retained stale backlog before taking the replacement snapshot.
|
|
// Events created while bootstrap is built are queued on this fresh receiver.
|
|
receiver = state.events.subscribe();
|
|
match send_ws_bootstrap(&state, &mut socket).await {
|
|
Ok(revision) => last_control_plan_revision = revision,
|
|
Err(()) => break,
|
|
}
|
|
}
|
|
Err(_) => break,
|
|
}
|
|
}
|
|
changed = control_plan.changed() => {
|
|
if changed.is_err() { break; }
|
|
let snapshot = control_plan.borrow_and_update().clone();
|
|
if let Some(snapshot) = snapshot {
|
|
if last_control_plan_revision == Some(snapshot.revision) {
|
|
continue;
|
|
}
|
|
let text = control_plan_ws_message(snapshot.as_ref()).to_string();
|
|
if socket.send(Message::Text(text)).await.is_err() { break; }
|
|
last_control_plan_revision = Some(snapshot.revision);
|
|
}
|
|
}
|
|
message = socket.next() => {
|
|
match message {
|
|
Some(Ok(Message::Ping(value))) => { if socket.send(Message::Pong(value)).await.is_err() { break; } }
|
|
Some(Ok(Message::Text(text))) if text == "ping" => { if socket.send(Message::Text("pong".into())).await.is_err() { break; } }
|
|
Some(Ok(Message::Close(_))) | None | Some(Err(_)) => break,
|
|
_ => {}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|