552 lines
22 KiB
Rust
552 lines
22 KiB
Rust
use std::collections::{BTreeMap, HashMap, HashSet};
|
|
|
|
use axum::extract::ws::{Message as WsMessage, WebSocket};
|
|
use futures::{SinkExt, StreamExt};
|
|
use protocol::subscription::{
|
|
EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame,
|
|
SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest,
|
|
SubscriptionResponse, SubscriptionSnapshot, SubscriptionTerminationCode, SubscriptionWorker,
|
|
};
|
|
use tokio::sync::mpsc;
|
|
use worker_runtime::identity::RuntimeWorkerRef;
|
|
|
|
use crate::runtime_subscription::{BrokerSubscriptionEvent, RuntimeSubscriptionBroker};
|
|
use crate::server::{WorkspaceApi, connect_workspace_worker_protocol};
|
|
|
|
const OUTBOUND_CAPACITY: usize = 256;
|
|
|
|
struct ActiveSubscription {
|
|
task: tokio::task::JoinHandle<()>,
|
|
methods: Option<mpsc::Sender<protocol::Method>>,
|
|
}
|
|
|
|
pub(crate) async fn serve_workspace_subscription(api: WorkspaceApi, socket: WebSocket) {
|
|
let broker = api.runtime_subscription_broker().clone();
|
|
let (mut socket_sender, mut socket_receiver) = socket.split();
|
|
let (control_outbound, mut control_receiver) = mpsc::channel::<WsMessage>(OUTBOUND_CAPACITY);
|
|
let (protocol_outbound, mut protocol_receiver) = mpsc::channel::<WsMessage>(OUTBOUND_CAPACITY);
|
|
let writer = tokio::spawn(async move {
|
|
loop {
|
|
let message = tokio::select! {
|
|
biased;
|
|
message = control_receiver.recv() => message,
|
|
message = protocol_receiver.recv() => message,
|
|
};
|
|
let Some(message) = message else { break };
|
|
if socket_sender.send(message).await.is_err() {
|
|
break;
|
|
}
|
|
}
|
|
});
|
|
let mut next_subscription_id = 1_u64;
|
|
let mut subscriptions = HashMap::<SubscriptionId, ActiveSubscription>::new();
|
|
|
|
while let Some(message) = socket_receiver.next().await {
|
|
let Ok(message) = message else { break };
|
|
match message {
|
|
WsMessage::Text(text) => {
|
|
let Ok(frame) = serde_json::from_str::<SubscriptionFrame>(text.as_str()) else {
|
|
break;
|
|
};
|
|
if frame.validate().is_err() {
|
|
break;
|
|
}
|
|
subscriptions.retain(|_, subscription| !subscription.task.is_finished());
|
|
match frame.payload {
|
|
SubscriptionFramePayload::Request(SubscriptionRequest::SubscribeEvents {
|
|
request_id,
|
|
selector,
|
|
}) => {
|
|
let subscription_id = SubscriptionId::new(format!(
|
|
"workspace-subscription-{next_subscription_id}"
|
|
))
|
|
.expect("generated Workspace subscription id is valid");
|
|
next_subscription_id = next_subscription_id.saturating_add(1);
|
|
match selector {
|
|
EventSubscriptionSelector::WorkspaceWorkers => {
|
|
let task = tokio::spawn(run_workspace_workers(
|
|
broker.clone(),
|
|
request_id,
|
|
subscription_id.clone(),
|
|
control_outbound.clone(),
|
|
));
|
|
subscriptions.insert(
|
|
subscription_id,
|
|
ActiveSubscription {
|
|
task,
|
|
methods: None,
|
|
},
|
|
);
|
|
}
|
|
EventSubscriptionSelector::WorkerProtocol {
|
|
worker_id,
|
|
runtime_id: Some(runtime_id),
|
|
} => {
|
|
let worker = RuntimeWorkerRef::new(&runtime_id, worker_id.as_str());
|
|
match connect_workspace_worker_protocol(&api, &worker).await {
|
|
Ok(connection) => {
|
|
let methods = connection.methods.clone();
|
|
let task = tokio::spawn(run_worker_protocol(
|
|
request_id,
|
|
subscription_id.clone(),
|
|
runtime_id,
|
|
worker_id,
|
|
connection.events,
|
|
control_outbound.clone(),
|
|
protocol_outbound.clone(),
|
|
));
|
|
subscriptions.insert(
|
|
subscription_id,
|
|
ActiveSubscription {
|
|
task,
|
|
methods: Some(methods),
|
|
},
|
|
);
|
|
}
|
|
Err(error) => {
|
|
let _ = send_rejected(
|
|
&control_outbound,
|
|
request_id,
|
|
SubscriptionRejectionCode::ResourceNotFound,
|
|
error.to_string(),
|
|
)
|
|
.await;
|
|
}
|
|
}
|
|
}
|
|
_ => {
|
|
let _ = send_rejected(
|
|
&control_outbound, request_id, SubscriptionRejectionCode::UnsupportedSelector,
|
|
"Workspace clients may subscribe only to workspace_workers or a runtime-scoped worker_protocol selector".to_string(),
|
|
).await;
|
|
}
|
|
}
|
|
}
|
|
SubscriptionFramePayload::Request(SubscriptionRequest::UnsubscribeEvents {
|
|
request_id,
|
|
subscription_id,
|
|
}) => {
|
|
if let Some(subscription) = subscriptions.remove(&subscription_id) {
|
|
subscription.task.abort();
|
|
}
|
|
if send_frame(
|
|
&control_outbound,
|
|
SubscriptionFrame::new(SubscriptionFramePayload::Response(
|
|
SubscriptionResponse::Unsubscribed {
|
|
request_id,
|
|
subscription_id,
|
|
},
|
|
)),
|
|
)
|
|
.await
|
|
.is_err()
|
|
{
|
|
break;
|
|
}
|
|
}
|
|
SubscriptionFramePayload::WorkerProtocol(message) => {
|
|
let Some(methods) = subscriptions
|
|
.get(&message.subscription_id)
|
|
.and_then(|value| value.methods.clone())
|
|
else {
|
|
break;
|
|
};
|
|
if methods.send(message.method).await.is_err() {
|
|
break;
|
|
}
|
|
}
|
|
SubscriptionFramePayload::Response(_) | SubscriptionFramePayload::Event(_) => {
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
WsMessage::Ping(value) => {
|
|
if control_outbound.send(WsMessage::Pong(value)).await.is_err() {
|
|
break;
|
|
}
|
|
}
|
|
WsMessage::Pong(_) => {}
|
|
WsMessage::Close(_) | WsMessage::Binary(_) => break,
|
|
}
|
|
}
|
|
|
|
for (_, subscription) in subscriptions {
|
|
subscription.task.abort();
|
|
}
|
|
drop(control_outbound);
|
|
drop(protocol_outbound);
|
|
let _ = writer.await;
|
|
}
|
|
|
|
async fn send_rejected(
|
|
outbound: &mpsc::Sender<WsMessage>,
|
|
request_id: protocol::subscription::SubscriptionRequestId,
|
|
code: SubscriptionRejectionCode,
|
|
message: String,
|
|
) -> Result<(), ()> {
|
|
send_frame(
|
|
outbound,
|
|
SubscriptionFrame::new(SubscriptionFramePayload::Response(
|
|
SubscriptionResponse::SubscriptionRejected {
|
|
request_id,
|
|
subscription_id: None,
|
|
code,
|
|
message,
|
|
},
|
|
)),
|
|
)
|
|
.await
|
|
}
|
|
|
|
async fn run_worker_protocol(
|
|
request_id: protocol::subscription::SubscriptionRequestId,
|
|
subscription_id: SubscriptionId,
|
|
runtime_id: String,
|
|
worker_id: protocol::subscription::SubscriptionWorkerId,
|
|
mut events: mpsc::Receiver<protocol::Event>,
|
|
control_outbound: mpsc::Sender<WsMessage>,
|
|
protocol_outbound: mpsc::Sender<WsMessage>,
|
|
) {
|
|
if send_frame(
|
|
&control_outbound,
|
|
SubscriptionFrame::new(SubscriptionFramePayload::Response(
|
|
SubscriptionResponse::Subscribed {
|
|
request_id,
|
|
subscription_id: subscription_id.clone(),
|
|
selector: EventSubscriptionSelector::WorkerProtocol {
|
|
worker_id: worker_id.clone(),
|
|
runtime_id: Some(runtime_id),
|
|
},
|
|
snapshot_revision: 0,
|
|
snapshot: SubscriptionSnapshot::WorkerProtocol {
|
|
worker_id: worker_id.clone(),
|
|
events: Vec::new(),
|
|
},
|
|
},
|
|
)),
|
|
)
|
|
.await
|
|
.is_err()
|
|
{
|
|
return;
|
|
}
|
|
let mut subject_revision = 0_u64;
|
|
while let Some(event) = events.recv().await {
|
|
subject_revision = subject_revision.saturating_add(1);
|
|
let frame =
|
|
SubscriptionFrame::new(SubscriptionFramePayload::Event(SubscriptionEvent::Event {
|
|
subscription_id: subscription_id.clone(),
|
|
subject_revision,
|
|
payload: SubscriptionEventPayload::WorkerProtocol {
|
|
worker_id: worker_id.clone(),
|
|
event,
|
|
},
|
|
}));
|
|
if try_send_frame(&protocol_outbound, frame).is_err() {
|
|
let _ = send_frame(
|
|
&control_outbound,
|
|
SubscriptionFrame::new(SubscriptionFramePayload::Event(
|
|
SubscriptionEvent::SubscriptionClosed {
|
|
subscription_id: subscription_id.clone(),
|
|
code: SubscriptionTerminationCode::Lagged,
|
|
message:
|
|
"Worker protocol subscriber lagged; resubscribe for a fresh snapshot"
|
|
.to_string(),
|
|
},
|
|
)),
|
|
)
|
|
.await;
|
|
return;
|
|
}
|
|
}
|
|
let _ = send_frame(
|
|
&control_outbound,
|
|
SubscriptionFrame::new(SubscriptionFramePayload::Event(
|
|
SubscriptionEvent::SubscriptionClosed {
|
|
subscription_id,
|
|
code: SubscriptionTerminationCode::ResourceGone,
|
|
message: "Worker protocol stream closed".to_string(),
|
|
},
|
|
)),
|
|
)
|
|
.await;
|
|
}
|
|
|
|
async fn run_workspace_workers(
|
|
broker: RuntimeSubscriptionBroker,
|
|
request_id: protocol::subscription::SubscriptionRequestId,
|
|
subscription_id: SubscriptionId,
|
|
outbound: mpsc::Sender<WsMessage>,
|
|
) {
|
|
let runtime_ids = broker.runtime_ids();
|
|
let mut pending = runtime_ids.iter().cloned().collect::<HashSet<_>>();
|
|
let (events, mut event_receiver) = mpsc::channel(OUTBOUND_CAPACITY);
|
|
let mut upstreams = tokio::task::JoinSet::new();
|
|
for runtime_id in runtime_ids {
|
|
let Ok(mut subscription) =
|
|
broker.subscribe(&runtime_id, EventSubscriptionSelector::RuntimeWorkers)
|
|
else {
|
|
pending.remove(&runtime_id);
|
|
continue;
|
|
};
|
|
let sender = events.clone();
|
|
upstreams.spawn(async move {
|
|
while let Some(event) = subscription.recv().await {
|
|
if sender.send((runtime_id.clone(), event)).await.is_err() {
|
|
break;
|
|
}
|
|
}
|
|
});
|
|
}
|
|
drop(events);
|
|
|
|
let mut workers = HashMap::<String, BTreeMap<String, SubscriptionWorker>>::new();
|
|
while !pending.is_empty() {
|
|
let Some((runtime_id, event)) = event_receiver.recv().await else {
|
|
return;
|
|
};
|
|
match event {
|
|
BrokerSubscriptionEvent::Snapshot { snapshot, .. } => {
|
|
install_snapshot(&mut workers, &runtime_id, snapshot);
|
|
pending.remove(&runtime_id);
|
|
}
|
|
BrokerSubscriptionEvent::Disconnected { .. }
|
|
| BrokerSubscriptionEvent::Rejected { .. }
|
|
| BrokerSubscriptionEvent::Closed { .. } => {
|
|
pending.remove(&runtime_id);
|
|
}
|
|
BrokerSubscriptionEvent::Event { .. } => {}
|
|
}
|
|
}
|
|
|
|
let mut revisions = HashMap::<RuntimeWorkerRef, u64>::new();
|
|
let mut initial_workers = Vec::new();
|
|
for (runtime_id, runtime) in &mut workers {
|
|
for worker in runtime.values_mut() {
|
|
let worker_ref = RuntimeWorkerRef::new(runtime_id, worker.worker_id.as_str());
|
|
worker.subject_revision = next_revision(&mut revisions, &worker_ref);
|
|
initial_workers.push(worker.clone());
|
|
}
|
|
}
|
|
sort_workers(&mut initial_workers);
|
|
if send_frame(
|
|
&outbound,
|
|
SubscriptionFrame::new(SubscriptionFramePayload::Response(
|
|
SubscriptionResponse::Subscribed {
|
|
request_id,
|
|
subscription_id: subscription_id.clone(),
|
|
selector: EventSubscriptionSelector::WorkspaceWorkers,
|
|
snapshot_revision: 1,
|
|
snapshot: SubscriptionSnapshot::Workers {
|
|
workers: initial_workers,
|
|
},
|
|
},
|
|
)),
|
|
)
|
|
.await
|
|
.is_err()
|
|
{
|
|
return;
|
|
}
|
|
|
|
while let Some((runtime_id, event)) = event_receiver.recv().await {
|
|
match event {
|
|
BrokerSubscriptionEvent::Snapshot { snapshot, .. } => {
|
|
let removed = workers.remove(&runtime_id).unwrap_or_default();
|
|
for worker in removed.values() {
|
|
let worker_ref = RuntimeWorkerRef::new(&runtime_id, worker.worker_id.as_str());
|
|
let revision = next_revision(&mut revisions, &worker_ref);
|
|
if send_event(
|
|
&outbound,
|
|
&subscription_id,
|
|
revision,
|
|
SubscriptionEventPayload::WorkerRemoved {
|
|
worker_id: worker.worker_id.clone(),
|
|
runtime_id: Some(runtime_id.clone()),
|
|
},
|
|
)
|
|
.await
|
|
.is_err()
|
|
{
|
|
return;
|
|
}
|
|
}
|
|
install_snapshot(&mut workers, &runtime_id, snapshot);
|
|
if let Some(current) = workers.get_mut(&runtime_id) {
|
|
for worker in current.values_mut() {
|
|
let worker_ref =
|
|
RuntimeWorkerRef::new(&runtime_id, worker.worker_id.as_str());
|
|
let revision = next_revision(&mut revisions, &worker_ref);
|
|
worker.subject_revision = revision;
|
|
if send_event(
|
|
&outbound,
|
|
&subscription_id,
|
|
revision,
|
|
SubscriptionEventPayload::WorkerUpserted {
|
|
worker: worker.clone(),
|
|
},
|
|
)
|
|
.await
|
|
.is_err()
|
|
{
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
BrokerSubscriptionEvent::Event { payload, .. } => match payload {
|
|
SubscriptionEventPayload::WorkerUpserted { mut worker } => {
|
|
worker.runtime_id = Some(runtime_id.clone());
|
|
let worker_ref = RuntimeWorkerRef::new(&runtime_id, worker.worker_id.as_str());
|
|
let revision = next_revision(&mut revisions, &worker_ref);
|
|
worker.subject_revision = revision;
|
|
workers
|
|
.entry(runtime_id)
|
|
.or_default()
|
|
.insert(worker.worker_id.to_string(), worker.clone());
|
|
if send_event(
|
|
&outbound,
|
|
&subscription_id,
|
|
revision,
|
|
SubscriptionEventPayload::WorkerUpserted { worker },
|
|
)
|
|
.await
|
|
.is_err()
|
|
{
|
|
return;
|
|
}
|
|
}
|
|
SubscriptionEventPayload::WorkerRemoved { worker_id, .. } => {
|
|
workers
|
|
.entry(runtime_id.clone())
|
|
.or_default()
|
|
.remove(worker_id.as_str());
|
|
let worker_ref = RuntimeWorkerRef::new(&runtime_id, worker_id.as_str());
|
|
let revision = next_revision(&mut revisions, &worker_ref);
|
|
if send_event(
|
|
&outbound,
|
|
&subscription_id,
|
|
revision,
|
|
SubscriptionEventPayload::WorkerRemoved {
|
|
worker_id,
|
|
runtime_id: Some(runtime_id),
|
|
},
|
|
)
|
|
.await
|
|
.is_err()
|
|
{
|
|
return;
|
|
}
|
|
}
|
|
_ => {}
|
|
},
|
|
BrokerSubscriptionEvent::Disconnected { .. } => {}
|
|
BrokerSubscriptionEvent::Rejected { code, message, .. } => {
|
|
let _ = send_frame(
|
|
&outbound,
|
|
SubscriptionFrame::new(SubscriptionFramePayload::Event(
|
|
SubscriptionEvent::SubscriptionClosed {
|
|
subscription_id: subscription_id.clone(),
|
|
code: rejection_termination(code),
|
|
message,
|
|
},
|
|
)),
|
|
)
|
|
.await;
|
|
return;
|
|
}
|
|
BrokerSubscriptionEvent::Closed { code, message, .. } => {
|
|
let _ = send_frame(
|
|
&outbound,
|
|
SubscriptionFrame::new(SubscriptionFramePayload::Event(
|
|
SubscriptionEvent::SubscriptionClosed {
|
|
subscription_id: subscription_id.clone(),
|
|
code,
|
|
message,
|
|
},
|
|
)),
|
|
)
|
|
.await;
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
fn install_snapshot(
|
|
workers: &mut HashMap<String, BTreeMap<String, SubscriptionWorker>>,
|
|
runtime_id: &str,
|
|
snapshot: SubscriptionSnapshot,
|
|
) {
|
|
let SubscriptionSnapshot::Workers {
|
|
workers: snapshot_workers,
|
|
} = snapshot
|
|
else {
|
|
return;
|
|
};
|
|
let mut projected = BTreeMap::new();
|
|
for mut worker in snapshot_workers {
|
|
worker.runtime_id = Some(runtime_id.to_string());
|
|
projected.insert(worker.worker_id.to_string(), worker);
|
|
}
|
|
workers.insert(runtime_id.to_string(), projected);
|
|
}
|
|
|
|
async fn send_event(
|
|
outbound: &mpsc::Sender<WsMessage>,
|
|
subscription_id: &SubscriptionId,
|
|
subject_revision: u64,
|
|
payload: SubscriptionEventPayload,
|
|
) -> Result<(), ()> {
|
|
send_frame(
|
|
outbound,
|
|
SubscriptionFrame::new(SubscriptionFramePayload::Event(SubscriptionEvent::Event {
|
|
subscription_id: subscription_id.clone(),
|
|
subject_revision,
|
|
payload,
|
|
})),
|
|
)
|
|
.await
|
|
}
|
|
|
|
fn try_send_frame(outbound: &mpsc::Sender<WsMessage>, frame: SubscriptionFrame) -> Result<(), ()> {
|
|
frame.validate().map_err(|_| ())?;
|
|
let text = serde_json::to_string(&frame).map_err(|_| ())?;
|
|
outbound
|
|
.try_send(WsMessage::Text(text.into()))
|
|
.map_err(|_| ())
|
|
}
|
|
|
|
async fn send_frame(
|
|
outbound: &mpsc::Sender<WsMessage>,
|
|
frame: SubscriptionFrame,
|
|
) -> Result<(), ()> {
|
|
frame.validate().map_err(|_| ())?;
|
|
outbound
|
|
.send(WsMessage::Text(
|
|
serde_json::to_string(&frame).map_err(|_| ())?.into(),
|
|
))
|
|
.await
|
|
.map_err(|_| ())
|
|
}
|
|
|
|
fn next_revision(revisions: &mut HashMap<RuntimeWorkerRef, u64>, worker: &RuntimeWorkerRef) -> u64 {
|
|
let revision = revisions.entry(worker.clone()).or_insert(0);
|
|
*revision = revision.saturating_add(1);
|
|
*revision
|
|
}
|
|
fn sort_workers(workers: &mut [SubscriptionWorker]) {
|
|
workers.sort_by(|left, right| {
|
|
left.runtime_id
|
|
.cmp(&right.runtime_id)
|
|
.then_with(|| left.worker_id.cmp(&right.worker_id))
|
|
});
|
|
}
|
|
fn rejection_termination(code: SubscriptionRejectionCode) -> SubscriptionTerminationCode {
|
|
match code {
|
|
SubscriptionRejectionCode::Unauthorized => SubscriptionTerminationCode::Unauthorized,
|
|
SubscriptionRejectionCode::ResourceNotFound => SubscriptionTerminationCode::ResourceGone,
|
|
_ => SubscriptionTerminationCode::ServerShutdown,
|
|
}
|
|
}
|