942 lines
33 KiB
Rust
942 lines
33 KiB
Rust
use std::collections::{HashMap, HashSet};
|
|
use std::sync::atomic::{AtomicU64, Ordering};
|
|
use std::sync::{Arc, RwLock};
|
|
use std::time::Duration;
|
|
|
|
use futures::{SinkExt, StreamExt};
|
|
use protocol::subscription::{
|
|
EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame,
|
|
SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest,
|
|
SubscriptionRequestId, SubscriptionResponse, SubscriptionSnapshot, SubscriptionTerminationCode,
|
|
};
|
|
use tokio::sync::mpsc;
|
|
use tokio_tungstenite::connect_async;
|
|
use tokio_tungstenite::tungstenite::Message;
|
|
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
|
|
use worker_runtime::auth::{CapabilityTokenSigner, capability_claims};
|
|
|
|
use crate::hosts::RemoteRuntimeConfig;
|
|
|
|
const DOWNSTREAM_QUEUE_CAPACITY: usize = 256;
|
|
const RECONNECT_DELAY: Duration = Duration::from_millis(100);
|
|
|
|
type RuntimeSocket =
|
|
tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
|
|
|
|
#[derive(Debug, thiserror::Error)]
|
|
pub enum RuntimeSubscriptionBrokerError {
|
|
#[error("unknown Runtime {0:?}")]
|
|
UnknownRuntime(String),
|
|
#[error("Runtime subscription broker command channel closed")]
|
|
Closed,
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub enum BrokerSubscriptionEvent {
|
|
Snapshot {
|
|
connection_generation: u64,
|
|
snapshot_revision: u64,
|
|
snapshot: SubscriptionSnapshot,
|
|
},
|
|
Event {
|
|
connection_generation: u64,
|
|
subject_revision: u64,
|
|
payload: SubscriptionEventPayload,
|
|
},
|
|
Disconnected {
|
|
connection_generation: u64,
|
|
message: String,
|
|
},
|
|
Rejected {
|
|
connection_generation: u64,
|
|
code: SubscriptionRejectionCode,
|
|
message: String,
|
|
},
|
|
Closed {
|
|
connection_generation: u64,
|
|
code: SubscriptionTerminationCode,
|
|
message: String,
|
|
},
|
|
}
|
|
|
|
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
|
pub struct RuntimeSubscriptionBrokerStatus {
|
|
pub connection_generation: u64,
|
|
pub connected: bool,
|
|
pub desired_selectors: usize,
|
|
pub upstream_subscriptions: usize,
|
|
}
|
|
|
|
pub struct BrokerSubscription {
|
|
downstream_id: u64,
|
|
runtime_id: String,
|
|
selector: EventSubscriptionSelector,
|
|
receiver: mpsc::Receiver<BrokerSubscriptionEvent>,
|
|
commands: mpsc::UnboundedSender<Command>,
|
|
}
|
|
|
|
impl BrokerSubscription {
|
|
pub fn runtime_id(&self) -> &str {
|
|
&self.runtime_id
|
|
}
|
|
pub fn selector(&self) -> &EventSubscriptionSelector {
|
|
&self.selector
|
|
}
|
|
pub async fn recv(&mut self) -> Option<BrokerSubscriptionEvent> {
|
|
self.receiver.recv().await
|
|
}
|
|
}
|
|
|
|
impl Drop for BrokerSubscription {
|
|
fn drop(&mut self) {
|
|
let _ = self.commands.send(Command::Unsubscribe(self.downstream_id));
|
|
}
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
struct Registration {
|
|
generation: u64,
|
|
commands: mpsc::UnboundedSender<Command>,
|
|
status: Arc<RwLock<RuntimeSubscriptionBrokerStatus>>,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct RuntimeSubscriptionBroker {
|
|
workspace_id: Arc<str>,
|
|
next_generation: Arc<AtomicU64>,
|
|
next_downstream: Arc<AtomicU64>,
|
|
registrations: Arc<RwLock<HashMap<String, Registration>>>,
|
|
}
|
|
|
|
impl RuntimeSubscriptionBroker {
|
|
pub fn new(workspace_id: impl Into<String>) -> Self {
|
|
Self {
|
|
workspace_id: Arc::from(workspace_id.into()),
|
|
next_generation: Arc::new(AtomicU64::new(1)),
|
|
next_downstream: Arc::new(AtomicU64::new(1)),
|
|
registrations: Arc::new(RwLock::new(HashMap::new())),
|
|
}
|
|
}
|
|
|
|
pub fn register_remote_runtime(&self, config: RemoteRuntimeConfig) -> u64 {
|
|
let generation = self.next_generation.fetch_add(1, Ordering::Relaxed);
|
|
let (commands, receiver) = mpsc::unbounded_channel();
|
|
let status = Arc::new(RwLock::new(RuntimeSubscriptionBrokerStatus {
|
|
connection_generation: generation,
|
|
..Default::default()
|
|
}));
|
|
let registration = Registration {
|
|
generation,
|
|
commands: commands.clone(),
|
|
status: status.clone(),
|
|
};
|
|
let previous = self
|
|
.registrations
|
|
.write()
|
|
.expect("broker registry poisoned")
|
|
.insert(config.runtime_id.clone(), registration);
|
|
if let Some(previous) = previous {
|
|
let _ = previous.commands.send(Command::Shutdown(generation));
|
|
}
|
|
tokio::spawn(run_connection(
|
|
config,
|
|
self.workspace_id.to_string(),
|
|
generation,
|
|
receiver,
|
|
status,
|
|
));
|
|
generation
|
|
}
|
|
|
|
pub fn register_embedded_runtime(
|
|
&self,
|
|
runtime_id: impl Into<String>,
|
|
runtime: worker_runtime::Runtime,
|
|
) -> u64 {
|
|
let runtime_id = runtime_id.into();
|
|
let generation = self.next_generation.fetch_add(1, Ordering::Relaxed);
|
|
let (commands, receiver) = mpsc::unbounded_channel();
|
|
let status = Arc::new(RwLock::new(RuntimeSubscriptionBrokerStatus {
|
|
connection_generation: generation,
|
|
connected: true,
|
|
..Default::default()
|
|
}));
|
|
let previous = self
|
|
.registrations
|
|
.write()
|
|
.expect("broker registry poisoned")
|
|
.insert(
|
|
runtime_id.clone(),
|
|
Registration {
|
|
generation,
|
|
commands: commands.clone(),
|
|
status: status.clone(),
|
|
},
|
|
);
|
|
if let Some(previous) = previous {
|
|
let _ = previous.commands.send(Command::Shutdown(generation));
|
|
}
|
|
tokio::spawn(run_embedded_connection(
|
|
runtime_id,
|
|
runtime,
|
|
self.workspace_id.to_string(),
|
|
generation,
|
|
receiver,
|
|
status,
|
|
));
|
|
generation
|
|
}
|
|
|
|
pub fn unregister_runtime(&self, runtime_id: &str) {
|
|
if let Some(registration) = self
|
|
.registrations
|
|
.write()
|
|
.expect("broker registry poisoned")
|
|
.remove(runtime_id)
|
|
{
|
|
let _ = registration
|
|
.commands
|
|
.send(Command::Shutdown(registration.generation.saturating_add(1)));
|
|
}
|
|
}
|
|
|
|
pub fn runtime_ids(&self) -> Vec<String> {
|
|
let mut runtime_ids = self
|
|
.registrations
|
|
.read()
|
|
.expect("broker registry poisoned")
|
|
.keys()
|
|
.cloned()
|
|
.collect::<Vec<_>>();
|
|
runtime_ids.sort();
|
|
runtime_ids
|
|
}
|
|
|
|
pub fn status(&self, runtime_id: &str) -> Option<RuntimeSubscriptionBrokerStatus> {
|
|
let status = self
|
|
.registrations
|
|
.read()
|
|
.expect("broker registry poisoned")
|
|
.get(runtime_id)?
|
|
.status
|
|
.clone();
|
|
Some(status.read().expect("broker status poisoned").clone())
|
|
}
|
|
|
|
pub fn subscribe(
|
|
&self,
|
|
runtime_id: &str,
|
|
selector: EventSubscriptionSelector,
|
|
) -> Result<BrokerSubscription, RuntimeSubscriptionBrokerError> {
|
|
selector
|
|
.validate()
|
|
.map_err(|_| RuntimeSubscriptionBrokerError::Closed)?;
|
|
let registration = self
|
|
.registrations
|
|
.read()
|
|
.expect("broker registry poisoned")
|
|
.get(runtime_id)
|
|
.cloned()
|
|
.ok_or_else(|| {
|
|
RuntimeSubscriptionBrokerError::UnknownRuntime(runtime_id.to_string())
|
|
})?;
|
|
let downstream_id = self.next_downstream.fetch_add(1, Ordering::Relaxed);
|
|
let (events, receiver) = mpsc::channel(DOWNSTREAM_QUEUE_CAPACITY);
|
|
let initial_events = events.clone();
|
|
registration
|
|
.commands
|
|
.send(Command::Subscribe {
|
|
downstream_id,
|
|
selector: selector.clone(),
|
|
events,
|
|
})
|
|
.map_err(|_| RuntimeSubscriptionBrokerError::Closed)?;
|
|
let initial_status = registration
|
|
.status
|
|
.read()
|
|
.expect("broker status poisoned")
|
|
.clone();
|
|
if !initial_status.connected {
|
|
let _ = initial_events.try_send(BrokerSubscriptionEvent::Disconnected {
|
|
connection_generation: initial_status.connection_generation,
|
|
message: "Runtime subscription connection is not currently available".to_string(),
|
|
});
|
|
}
|
|
Ok(BrokerSubscription {
|
|
downstream_id,
|
|
runtime_id: runtime_id.to_string(),
|
|
selector,
|
|
receiver,
|
|
commands: registration.commands,
|
|
})
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
enum Command {
|
|
Subscribe {
|
|
downstream_id: u64,
|
|
selector: EventSubscriptionSelector,
|
|
events: mpsc::Sender<BrokerSubscriptionEvent>,
|
|
},
|
|
Unsubscribe(u64),
|
|
Shutdown(u64),
|
|
}
|
|
|
|
struct SelectorState {
|
|
downstreams: HashMap<u64, mpsc::Sender<BrokerSubscriptionEvent>>,
|
|
upstream_id: Option<SubscriptionId>,
|
|
pending: bool,
|
|
snapshot: Option<(u64, SubscriptionSnapshot)>,
|
|
revisions: HashMap<String, u64>,
|
|
}
|
|
impl SelectorState {
|
|
fn new() -> Self {
|
|
Self {
|
|
downstreams: HashMap::new(),
|
|
upstream_id: None,
|
|
pending: false,
|
|
snapshot: None,
|
|
revisions: HashMap::new(),
|
|
}
|
|
}
|
|
}
|
|
|
|
struct State {
|
|
runtime_id: String,
|
|
generation: u64,
|
|
next_request: u64,
|
|
selectors: HashMap<EventSubscriptionSelector, SelectorState>,
|
|
downstream_index: HashMap<u64, EventSubscriptionSelector>,
|
|
pending: HashMap<SubscriptionRequestId, EventSubscriptionSelector>,
|
|
upstream_index: HashMap<SubscriptionId, EventSubscriptionSelector>,
|
|
}
|
|
impl State {
|
|
fn new(runtime_id: String, generation: u64) -> Self {
|
|
Self {
|
|
runtime_id,
|
|
generation,
|
|
next_request: 1,
|
|
selectors: HashMap::new(),
|
|
downstream_index: HashMap::new(),
|
|
pending: HashMap::new(),
|
|
upstream_index: HashMap::new(),
|
|
}
|
|
}
|
|
fn request_id(&mut self) -> SubscriptionRequestId {
|
|
let id = self.next_request;
|
|
self.next_request = self.next_request.saturating_add(1);
|
|
SubscriptionRequestId::new(format!("server-{}-{id}", self.generation)).unwrap()
|
|
}
|
|
fn disconnected(&mut self, message: String) {
|
|
self.pending.clear();
|
|
self.upstream_index.clear();
|
|
for selector in self.selectors.values_mut() {
|
|
selector.upstream_id = None;
|
|
selector.pending = false;
|
|
selector.snapshot = None;
|
|
selector.revisions.clear();
|
|
broadcast(
|
|
&mut selector.downstreams,
|
|
BrokerSubscriptionEvent::Disconnected {
|
|
connection_generation: self.generation,
|
|
message: message.clone(),
|
|
},
|
|
);
|
|
}
|
|
}
|
|
}
|
|
|
|
struct EmbeddedEntry {
|
|
downstreams: HashMap<u64, mpsc::Sender<BrokerSubscriptionEvent>>,
|
|
snapshot_revision: u64,
|
|
snapshot: SubscriptionSnapshot,
|
|
task: tokio::task::JoinHandle<()>,
|
|
}
|
|
|
|
async fn run_embedded_connection(
|
|
runtime_id: String,
|
|
runtime: worker_runtime::Runtime,
|
|
_workspace_id: String,
|
|
generation: u64,
|
|
mut commands: mpsc::UnboundedReceiver<Command>,
|
|
status: Arc<RwLock<RuntimeSubscriptionBrokerStatus>>,
|
|
) {
|
|
let (updates, mut update_receiver) = mpsc::unbounded_channel();
|
|
let mut entries = HashMap::<EventSubscriptionSelector, EmbeddedEntry>::new();
|
|
let mut downstream_index = HashMap::<u64, EventSubscriptionSelector>::new();
|
|
loop {
|
|
tokio::select! {
|
|
command = commands.recv() => match command {
|
|
Some(Command::Subscribe { downstream_id, selector, events }) => {
|
|
downstream_index.insert(downstream_id, selector.clone());
|
|
if let Some(entry) = entries.get_mut(&selector) {
|
|
let _ = events.try_send(BrokerSubscriptionEvent::Snapshot {
|
|
connection_generation: generation,
|
|
snapshot_revision: entry.snapshot_revision,
|
|
snapshot: entry.snapshot.clone(),
|
|
});
|
|
entry.downstreams.insert(downstream_id, events);
|
|
} else {
|
|
match runtime.subscribe_event_selector(selector.clone()) {
|
|
Ok(mut subscription) => {
|
|
let snapshot_revision = subscription.snapshot_revision();
|
|
let snapshot = project_snapshot_runtime(subscription.snapshot().clone(), &runtime_id);
|
|
let _ = events.try_send(BrokerSubscriptionEvent::Snapshot { connection_generation: generation, snapshot_revision, snapshot: snapshot.clone() });
|
|
let sender = updates.clone();
|
|
let task_selector = selector.clone();
|
|
let task_runtime_id = runtime_id.clone();
|
|
let task = tokio::spawn(async move {
|
|
while let Ok(update) = subscription.recv().await {
|
|
let payload = project_payload_runtime(update.payload, &task_runtime_id);
|
|
if sender.send((task_selector.clone(), update.subject_revision, payload)).is_err() { break; }
|
|
}
|
|
});
|
|
entries.insert(selector, EmbeddedEntry { downstreams: HashMap::from([(downstream_id, events)]), snapshot_revision, snapshot, task });
|
|
}
|
|
Err(error) => {
|
|
let _ = events.try_send(BrokerSubscriptionEvent::Rejected { connection_generation: generation, code: SubscriptionRejectionCode::UnsupportedSelector, message: error.to_string() });
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Some(Command::Unsubscribe(id)) => {
|
|
if let Some(selector) = downstream_index.remove(&id) {
|
|
let empty = entries.get_mut(&selector).is_some_and(|entry| { entry.downstreams.remove(&id); entry.downstreams.is_empty() });
|
|
if empty { if let Some(entry) = entries.remove(&selector) { entry.task.abort(); } }
|
|
}
|
|
}
|
|
Some(Command::Shutdown(replacement)) => {
|
|
for entry in entries.values_mut() {
|
|
broadcast(&mut entry.downstreams, BrokerSubscriptionEvent::Closed { connection_generation: generation, code: SubscriptionTerminationCode::ServerShutdown, message: format!("embedded Runtime generation {generation} was fenced by {replacement}") });
|
|
entry.task.abort();
|
|
}
|
|
return;
|
|
}
|
|
None => return,
|
|
},
|
|
update = update_receiver.recv() => {
|
|
let Some((selector, subject_revision, payload)) = update else { return; };
|
|
if let Some(entry) = entries.get_mut(&selector) {
|
|
broadcast(&mut entry.downstreams, BrokerSubscriptionEvent::Event { connection_generation: generation, subject_revision, payload });
|
|
}
|
|
}
|
|
}
|
|
*status.write().expect("broker status poisoned") = RuntimeSubscriptionBrokerStatus {
|
|
connection_generation: generation,
|
|
connected: true,
|
|
desired_selectors: entries.len(),
|
|
upstream_subscriptions: entries.len(),
|
|
};
|
|
}
|
|
}
|
|
|
|
fn project_snapshot_runtime(
|
|
mut snapshot: SubscriptionSnapshot,
|
|
runtime_id: &str,
|
|
) -> SubscriptionSnapshot {
|
|
if let SubscriptionSnapshot::Workers { workers } = &mut snapshot {
|
|
for worker in workers {
|
|
worker.runtime_id = Some(runtime_id.to_string());
|
|
}
|
|
}
|
|
snapshot
|
|
}
|
|
|
|
fn project_payload_runtime(
|
|
mut payload: SubscriptionEventPayload,
|
|
runtime_id: &str,
|
|
) -> SubscriptionEventPayload {
|
|
match &mut payload {
|
|
SubscriptionEventPayload::WorkerUpserted { worker } => {
|
|
worker.runtime_id = Some(runtime_id.to_string());
|
|
}
|
|
SubscriptionEventPayload::WorkerRemoved {
|
|
runtime_id: projected_runtime_id,
|
|
..
|
|
} => {
|
|
*projected_runtime_id = Some(runtime_id.to_string());
|
|
}
|
|
_ => {}
|
|
}
|
|
payload
|
|
}
|
|
|
|
async fn run_connection(
|
|
config: RemoteRuntimeConfig,
|
|
workspace_id: String,
|
|
generation: u64,
|
|
mut commands: mpsc::UnboundedReceiver<Command>,
|
|
status: Arc<RwLock<RuntimeSubscriptionBrokerStatus>>,
|
|
) {
|
|
let mut state = State::new(config.runtime_id.clone(), generation);
|
|
let mut disconnect_notified = false;
|
|
loop {
|
|
update_status(&status, &state, false);
|
|
let connecting = connect_runtime(&config, &workspace_id);
|
|
tokio::pin!(connecting);
|
|
let connection = loop {
|
|
tokio::select! {
|
|
command = commands.recv() => match command {
|
|
Some(Command::Shutdown(replacement)) => { close_all(&mut state, replacement); return; }
|
|
Some(command) => { apply_offline(&mut state, command); update_status(&status, &state, false); }
|
|
None => return,
|
|
},
|
|
connected = &mut connecting => break connected,
|
|
}
|
|
};
|
|
let mut socket = match connection {
|
|
Ok(socket) => socket,
|
|
Err(error) => {
|
|
if !disconnect_notified {
|
|
state.disconnected(error);
|
|
disconnect_notified = true;
|
|
}
|
|
tokio::time::sleep(RECONNECT_DELAY).await;
|
|
continue;
|
|
}
|
|
};
|
|
if resubscribe_all(&mut socket, &mut state).await.is_err() {
|
|
state.disconnected("failed to restore Runtime subscriptions".into());
|
|
disconnect_notified = true;
|
|
tokio::time::sleep(RECONNECT_DELAY).await;
|
|
continue;
|
|
}
|
|
update_status(&status, &state, true);
|
|
let reason = loop {
|
|
tokio::select! {
|
|
command = commands.recv() => match command {
|
|
Some(Command::Shutdown(replacement)) => { let _ = socket.close(None).await; close_all(&mut state, replacement); return; }
|
|
Some(command) => if apply_online(&mut socket, &mut state, command).await.is_err() { break "failed to apply Runtime subscription command".into(); },
|
|
None => return,
|
|
},
|
|
message = socket.next() => match message {
|
|
Some(Ok(Message::Text(text))) => match serde_json::from_str::<SubscriptionFrame>(text.as_str()) {
|
|
Ok(frame) if frame.validate().is_ok() => if handle_frame(&mut socket, &mut state, frame).await.is_err() { break "invalid Runtime subscription transition".into(); },
|
|
_ => break "Runtime returned an invalid subscription frame".into(),
|
|
},
|
|
Some(Ok(Message::Ping(value))) => if socket.send(Message::Pong(value)).await.is_err() { break "Runtime pong failed".into(); },
|
|
Some(Ok(Message::Pong(_))) => {},
|
|
Some(Ok(Message::Close(_))) | None => break "Runtime subscription connection closed".into(),
|
|
Some(Ok(Message::Binary(_) | Message::Frame(_))) => break "Runtime returned a non-text subscription frame".into(),
|
|
Some(Err(error)) => break format!("Runtime subscription connection failed: {error}"),
|
|
}
|
|
}
|
|
update_status(&status, &state, true);
|
|
};
|
|
state.disconnected(reason);
|
|
disconnect_notified = true;
|
|
update_status(&status, &state, false);
|
|
tokio::time::sleep(RECONNECT_DELAY).await;
|
|
}
|
|
}
|
|
|
|
fn apply_offline(state: &mut State, command: Command) {
|
|
match command {
|
|
Command::Subscribe {
|
|
downstream_id,
|
|
selector,
|
|
events,
|
|
} => {
|
|
state
|
|
.downstream_index
|
|
.insert(downstream_id, selector.clone());
|
|
state
|
|
.selectors
|
|
.entry(selector)
|
|
.or_insert_with(SelectorState::new)
|
|
.downstreams
|
|
.insert(downstream_id, events);
|
|
}
|
|
Command::Unsubscribe(id) => remove_downstream(state, id),
|
|
Command::Shutdown(_) => unreachable!(),
|
|
}
|
|
}
|
|
|
|
async fn apply_online(
|
|
socket: &mut RuntimeSocket,
|
|
state: &mut State,
|
|
command: Command,
|
|
) -> Result<(), ()> {
|
|
match command {
|
|
Command::Subscribe {
|
|
downstream_id,
|
|
selector,
|
|
events,
|
|
} => {
|
|
state
|
|
.downstream_index
|
|
.insert(downstream_id, selector.clone());
|
|
let entry = state
|
|
.selectors
|
|
.entry(selector.clone())
|
|
.or_insert_with(SelectorState::new);
|
|
if let Some((revision, snapshot)) = &entry.snapshot {
|
|
let _ = events.try_send(BrokerSubscriptionEvent::Snapshot {
|
|
connection_generation: state.generation,
|
|
snapshot_revision: *revision,
|
|
snapshot: snapshot.clone(),
|
|
});
|
|
}
|
|
entry.downstreams.insert(downstream_id, events);
|
|
if entry.upstream_id.is_none() && !entry.pending {
|
|
send_subscribe(socket, state, selector).await?;
|
|
}
|
|
}
|
|
Command::Unsubscribe(id) => {
|
|
let selector = state.downstream_index.get(&id).cloned();
|
|
remove_downstream(state, id);
|
|
if let Some(selector) = selector {
|
|
maybe_unsubscribe(socket, state, selector).await?;
|
|
}
|
|
}
|
|
Command::Shutdown(_) => unreachable!(),
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn resubscribe_all(socket: &mut RuntimeSocket, state: &mut State) -> Result<(), ()> {
|
|
let selectors = state
|
|
.selectors
|
|
.iter()
|
|
.filter(|(_, value)| !value.downstreams.is_empty())
|
|
.map(|(key, _)| key.clone())
|
|
.collect::<Vec<_>>();
|
|
for selector in selectors {
|
|
send_subscribe(socket, state, selector).await?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn send_subscribe(
|
|
socket: &mut RuntimeSocket,
|
|
state: &mut State,
|
|
selector: EventSubscriptionSelector,
|
|
) -> Result<(), ()> {
|
|
let request_id = state.request_id();
|
|
send_frame(
|
|
socket,
|
|
SubscriptionFrame::new(SubscriptionFramePayload::Request(
|
|
SubscriptionRequest::SubscribeEvents {
|
|
request_id: request_id.clone(),
|
|
selector: selector.clone(),
|
|
},
|
|
)),
|
|
)
|
|
.await?;
|
|
state.pending.insert(request_id, selector.clone());
|
|
state.selectors.get_mut(&selector).unwrap().pending = true;
|
|
Ok(())
|
|
}
|
|
|
|
async fn maybe_unsubscribe(
|
|
socket: &mut RuntimeSocket,
|
|
state: &mut State,
|
|
selector: EventSubscriptionSelector,
|
|
) -> Result<(), ()> {
|
|
let Some(entry) = state.selectors.get(&selector) else {
|
|
return Ok(());
|
|
};
|
|
if !entry.downstreams.is_empty() {
|
|
return Ok(());
|
|
}
|
|
if let Some(subscription_id) = entry.upstream_id.clone() {
|
|
let request_id = state.request_id();
|
|
send_frame(
|
|
socket,
|
|
SubscriptionFrame::new(SubscriptionFramePayload::Request(
|
|
SubscriptionRequest::UnsubscribeEvents {
|
|
request_id,
|
|
subscription_id: subscription_id.clone(),
|
|
},
|
|
)),
|
|
)
|
|
.await?;
|
|
state.upstream_index.remove(&subscription_id);
|
|
state.selectors.remove(&selector);
|
|
} else if !entry.pending {
|
|
state.selectors.remove(&selector);
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
async fn handle_frame(
|
|
socket: &mut RuntimeSocket,
|
|
state: &mut State,
|
|
frame: SubscriptionFrame,
|
|
) -> Result<(), ()> {
|
|
match frame.payload {
|
|
SubscriptionFramePayload::Response(SubscriptionResponse::Subscribed {
|
|
request_id,
|
|
subscription_id,
|
|
selector,
|
|
snapshot_revision,
|
|
snapshot,
|
|
}) => {
|
|
if state.pending.remove(&request_id) != Some(selector.clone()) {
|
|
return Err(());
|
|
}
|
|
let snapshot = project_snapshot_runtime(snapshot, &state.runtime_id);
|
|
let entry = state.selectors.get_mut(&selector).ok_or(())?;
|
|
entry.pending = false;
|
|
entry.upstream_id = Some(subscription_id.clone());
|
|
entry.snapshot = Some((snapshot_revision, snapshot.clone()));
|
|
entry.revisions = snapshot_revisions(&snapshot);
|
|
state
|
|
.upstream_index
|
|
.insert(subscription_id, selector.clone());
|
|
broadcast(
|
|
&mut entry.downstreams,
|
|
BrokerSubscriptionEvent::Snapshot {
|
|
connection_generation: state.generation,
|
|
snapshot_revision,
|
|
snapshot,
|
|
},
|
|
);
|
|
if entry.downstreams.is_empty() {
|
|
maybe_unsubscribe(socket, state, selector).await?;
|
|
}
|
|
}
|
|
SubscriptionFramePayload::Response(SubscriptionResponse::Unsubscribed { .. }) => {}
|
|
SubscriptionFramePayload::Response(SubscriptionResponse::SubscriptionRejected {
|
|
request_id,
|
|
code,
|
|
message,
|
|
..
|
|
}) => {
|
|
if let Some(selector) = state.pending.remove(&request_id) {
|
|
if let Some(mut entry) = state.selectors.remove(&selector) {
|
|
broadcast(
|
|
&mut entry.downstreams,
|
|
BrokerSubscriptionEvent::Rejected {
|
|
connection_generation: state.generation,
|
|
code,
|
|
message,
|
|
},
|
|
);
|
|
}
|
|
}
|
|
}
|
|
SubscriptionFramePayload::Event(SubscriptionEvent::Event {
|
|
subscription_id,
|
|
subject_revision,
|
|
payload,
|
|
}) => {
|
|
let selector = state
|
|
.upstream_index
|
|
.get(&subscription_id)
|
|
.cloned()
|
|
.ok_or(())?;
|
|
payload.validate_for_selector(&selector).map_err(|_| ())?;
|
|
let payload = project_payload_runtime(payload, &state.runtime_id);
|
|
let entry = state.selectors.get_mut(&selector).ok_or(())?;
|
|
if let Some(subject) = event_subject(&payload) {
|
|
let revision = entry.revisions.entry(subject).or_insert(0);
|
|
if subject_revision <= *revision {
|
|
return Ok(());
|
|
}
|
|
*revision = subject_revision;
|
|
}
|
|
broadcast(
|
|
&mut entry.downstreams,
|
|
BrokerSubscriptionEvent::Event {
|
|
connection_generation: state.generation,
|
|
subject_revision,
|
|
payload,
|
|
},
|
|
);
|
|
}
|
|
SubscriptionFramePayload::Event(SubscriptionEvent::SubscriptionClosed {
|
|
subscription_id,
|
|
code,
|
|
message,
|
|
}) => {
|
|
let selector = state.upstream_index.remove(&subscription_id).ok_or(())?;
|
|
let should_resubscribe = if let Some(entry) = state.selectors.get_mut(&selector) {
|
|
entry.upstream_id = None;
|
|
entry.snapshot = None;
|
|
entry.revisions.clear();
|
|
broadcast(
|
|
&mut entry.downstreams,
|
|
BrokerSubscriptionEvent::Closed {
|
|
connection_generation: state.generation,
|
|
code,
|
|
message,
|
|
},
|
|
);
|
|
!entry.downstreams.is_empty()
|
|
} else {
|
|
false
|
|
};
|
|
if should_resubscribe {
|
|
send_subscribe(socket, state, selector).await?;
|
|
}
|
|
}
|
|
SubscriptionFramePayload::Request(_) | SubscriptionFramePayload::WorkerProtocol(_) => {
|
|
return Err(());
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn remove_downstream(state: &mut State, id: u64) {
|
|
if let Some(selector) = state.downstream_index.remove(&id) {
|
|
if let Some(entry) = state.selectors.get_mut(&selector) {
|
|
entry.downstreams.remove(&id);
|
|
}
|
|
}
|
|
}
|
|
|
|
fn close_all(state: &mut State, replacement: u64) {
|
|
for entry in state.selectors.values_mut() {
|
|
broadcast(
|
|
&mut entry.downstreams,
|
|
BrokerSubscriptionEvent::Closed {
|
|
connection_generation: state.generation,
|
|
code: SubscriptionTerminationCode::ServerShutdown,
|
|
message: format!(
|
|
"Runtime connection generation {} was fenced by generation {replacement}",
|
|
state.generation
|
|
),
|
|
},
|
|
);
|
|
}
|
|
}
|
|
|
|
fn broadcast(
|
|
downstreams: &mut HashMap<u64, mpsc::Sender<BrokerSubscriptionEvent>>,
|
|
event: BrokerSubscriptionEvent,
|
|
) {
|
|
let mut closed = HashSet::new();
|
|
for (id, sender) in downstreams.iter() {
|
|
if sender.try_send(event.clone()).is_err() {
|
|
closed.insert(*id);
|
|
}
|
|
}
|
|
downstreams.retain(|id, _| !closed.contains(id));
|
|
}
|
|
|
|
fn snapshot_revisions(snapshot: &SubscriptionSnapshot) -> HashMap<String, u64> {
|
|
match snapshot {
|
|
SubscriptionSnapshot::Workers { workers } => workers
|
|
.iter()
|
|
.map(|worker| {
|
|
(
|
|
format!(
|
|
"{}:{}",
|
|
worker.runtime_id.as_deref().unwrap_or_default(),
|
|
worker.worker_id
|
|
),
|
|
worker.subject_revision,
|
|
)
|
|
})
|
|
.collect(),
|
|
SubscriptionSnapshot::WorkerProtocol { worker_id, .. } => {
|
|
HashMap::from([(worker_id.to_string(), 0)])
|
|
}
|
|
SubscriptionSnapshot::WorkspaceWorkdirs { .. } => HashMap::new(),
|
|
}
|
|
}
|
|
fn event_subject(payload: &SubscriptionEventPayload) -> Option<String> {
|
|
Some(match payload {
|
|
SubscriptionEventPayload::WorkerUpserted { worker } => format!(
|
|
"{}:{}",
|
|
worker.runtime_id.as_deref().unwrap_or_default(),
|
|
worker.worker_id
|
|
),
|
|
SubscriptionEventPayload::WorkerRemoved {
|
|
worker_id,
|
|
runtime_id,
|
|
} => format!(
|
|
"{}:{}",
|
|
runtime_id.as_deref().unwrap_or_default(),
|
|
worker_id
|
|
),
|
|
SubscriptionEventPayload::WorkerProtocol { worker_id, .. } => worker_id.to_string(),
|
|
SubscriptionEventPayload::WorkdirUpserted { workdir } => {
|
|
workdir.working_directory_id.to_string()
|
|
}
|
|
SubscriptionEventPayload::WorkdirRemoved {
|
|
working_directory_id,
|
|
} => working_directory_id.to_string(),
|
|
})
|
|
}
|
|
async fn send_frame(socket: &mut RuntimeSocket, frame: SubscriptionFrame) -> Result<(), ()> {
|
|
frame.validate().map_err(|_| ())?;
|
|
socket
|
|
.send(Message::Text(
|
|
serde_json::to_string(&frame).map_err(|_| ())?.into(),
|
|
))
|
|
.await
|
|
.map_err(|_| ())
|
|
}
|
|
|
|
async fn connect_runtime(
|
|
config: &RemoteRuntimeConfig,
|
|
workspace_id: &str,
|
|
) -> Result<RuntimeSocket, String> {
|
|
let endpoint = runtime_endpoint(&config.base_url);
|
|
let mut request = endpoint
|
|
.into_client_request()
|
|
.map_err(|error| format!("invalid Runtime subscription endpoint: {error}"))?;
|
|
if let Some(token) = runtime_token(config, workspace_id)? {
|
|
request.headers_mut().insert(
|
|
"authorization",
|
|
format!("Bearer {token}")
|
|
.parse()
|
|
.map_err(|error| format!("invalid Runtime authorization header: {error}"))?,
|
|
);
|
|
}
|
|
connect_async(request)
|
|
.await
|
|
.map(|(socket, _)| socket)
|
|
.map_err(|error| format!("failed to connect Runtime subscription endpoint: {error}"))
|
|
}
|
|
fn runtime_endpoint(base_url: &str) -> String {
|
|
let base = base_url.trim_end_matches('/');
|
|
if let Some(rest) = base.strip_prefix("https://") {
|
|
format!("wss://{rest}/v1/protocol/ws")
|
|
} else if let Some(rest) = base.strip_prefix("http://") {
|
|
format!("ws://{rest}/v1/protocol/ws")
|
|
} else {
|
|
format!("{base}/v1/protocol/ws")
|
|
}
|
|
}
|
|
fn runtime_token(
|
|
config: &RemoteRuntimeConfig,
|
|
workspace_id: &str,
|
|
) -> Result<Option<String>, String> {
|
|
let Some(auth) = config.auth.as_ref() else {
|
|
return Ok(config.bearer_token.clone());
|
|
};
|
|
let signer = CapabilityTokenSigner::new(&auth.server_id, &auth.server_private_key);
|
|
let claims = capability_claims(
|
|
&auth.server_id,
|
|
&config.runtime_id,
|
|
workspace_id,
|
|
vec!["workers:list".into()],
|
|
300,
|
|
)
|
|
.map_err(|error| error.to_string())?;
|
|
signer
|
|
.sign(&claims)
|
|
.map(Some)
|
|
.map_err(|error| error.to_string())
|
|
}
|
|
fn update_status(status: &RwLock<RuntimeSubscriptionBrokerStatus>, state: &State, connected: bool) {
|
|
*status.write().expect("broker status poisoned") = RuntimeSubscriptionBrokerStatus {
|
|
connection_generation: state.generation,
|
|
connected,
|
|
desired_selectors: state
|
|
.selectors
|
|
.values()
|
|
.filter(|value| !value.downstreams.is_empty())
|
|
.count(),
|
|
upstream_subscriptions: state.upstream_index.len(),
|
|
};
|
|
}
|
|
|
|
#[cfg(test)]
|
|
#[path = "runtime_subscription_tests.rs"]
|
|
mod tests;
|