server: broker runtime event subscriptions
This commit is contained in:
@@ -31,7 +31,7 @@ sha2.workspace = true
|
||||
thiserror.workspace = true
|
||||
ticket.workspace = true
|
||||
memory.workspace = true
|
||||
tokio = { workspace = true, features = ["fs", "macros", "net", "rt-multi-thread", "sync"] }
|
||||
tokio = { workspace = true, features = ["fs", "macros", "net", "rt-multi-thread", "sync", "time"] }
|
||||
tokio-tungstenite.workspace = true
|
||||
worker.workspace = true
|
||||
worker-runtime.workspace = true
|
||||
|
||||
@@ -19,6 +19,7 @@ pub mod records;
|
||||
pub use records::ticket_api_typescript;
|
||||
pub mod repositories;
|
||||
pub mod resource_broker;
|
||||
pub mod runtime_subscription;
|
||||
pub mod server;
|
||||
pub mod skills;
|
||||
pub mod store;
|
||||
|
||||
@@ -0,0 +1,737 @@
|
||||
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 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 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);
|
||||
registration
|
||||
.commands
|
||||
.send(Command::Subscribe {
|
||||
downstream_id,
|
||||
selector: selector.clone(),
|
||||
events,
|
||||
})
|
||||
.map_err(|_| RuntimeSubscriptionBrokerError::Closed)?;
|
||||
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 {
|
||||
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(generation: u64) -> Self {
|
||||
Self {
|
||||
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(),
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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(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 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 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(_) => 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| (worker.worker_id.to_string(), 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 } => worker.worker_id.to_string(),
|
||||
SubscriptionEventPayload::WorkerRemoved { 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;
|
||||
@@ -0,0 +1,278 @@
|
||||
use super::*;
|
||||
use protocol::subscription::{SubscriptionWorkerIds, SubscriptionWorkerState};
|
||||
use worker_runtime::Runtime;
|
||||
use worker_runtime::catalog::{
|
||||
CreateWorkerRequest, ProfileSelector, ProfileSourceArchiveHttpRef, ProfileSourceArchiveSource,
|
||||
};
|
||||
use worker_runtime::execution::{
|
||||
WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation, WorkerExecutionResult,
|
||||
WorkerExecutionRunState, WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult,
|
||||
};
|
||||
use worker_runtime::profile_archive::{ProfileSourceArchiveRef, ProfileSourceGraphSummary};
|
||||
|
||||
#[derive(Debug)]
|
||||
struct TestExecutionBackend;
|
||||
|
||||
impl WorkerExecutionBackend for TestExecutionBackend {
|
||||
fn backend_id(&self) -> &str {
|
||||
"runtime-subscription-test"
|
||||
}
|
||||
|
||||
fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult {
|
||||
WorkerExecutionSpawnResult::connected(
|
||||
WorkerExecutionHandle::new(request.worker_ref, self.backend_id()),
|
||||
WorkerExecutionRunState::Idle,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
fn dispatch_input(
|
||||
&self,
|
||||
_handle: &WorkerExecutionHandle,
|
||||
_input: worker_runtime::interaction::WorkerInput,
|
||||
) -> WorkerExecutionResult {
|
||||
WorkerExecutionResult::accepted(
|
||||
WorkerExecutionOperation::Input,
|
||||
WorkerExecutionRunState::Busy,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
const TOKEN: &str = "runtime-subscription-test-token";
|
||||
|
||||
fn create_request(name: &str) -> CreateWorkerRequest {
|
||||
CreateWorkerRequest {
|
||||
idempotency_key: None,
|
||||
idempotency_fingerprint: None,
|
||||
profile: ProfileSelector::Builtin("builtin:companion".to_string()),
|
||||
display_name: Some(name.to_string()),
|
||||
config_bundle: None,
|
||||
profile_source: ProfileSourceArchiveSource::Http {
|
||||
location: ProfileSourceArchiveHttpRef {
|
||||
url: "http://127.0.0.1/profiles/test".to_string(),
|
||||
etag: None,
|
||||
archive: ProfileSourceArchiveRef {
|
||||
id: "test-profile-source".to_string(),
|
||||
digest: "test-digest".to_string(),
|
||||
size_bytes: 0,
|
||||
source_graph: ProfileSourceGraphSummary {
|
||||
source_count: 0,
|
||||
total_source_bytes: 0,
|
||||
entrypoints: std::collections::BTreeMap::new(),
|
||||
import_count: 0,
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
initial_input: None,
|
||||
working_directory_request: None,
|
||||
working_directory: None,
|
||||
workspace_api: None,
|
||||
}
|
||||
}
|
||||
|
||||
async fn start_runtime_server(
|
||||
listener: tokio::net::TcpListener,
|
||||
runtime: Runtime,
|
||||
) -> tokio::task::JoinHandle<()> {
|
||||
let router = worker_runtime::http_server::runtime_http_router(runtime, TOKEN.to_string());
|
||||
tokio::spawn(async move {
|
||||
let _ = axum::serve(listener, router).await;
|
||||
})
|
||||
}
|
||||
|
||||
async fn fixture() -> (Runtime, RemoteRuntimeConfig, tokio::task::JoinHandle<()>) {
|
||||
let runtime = Runtime::with_execution_backend(
|
||||
worker_runtime::RuntimeOptions::default(),
|
||||
std::sync::Arc::new(TestExecutionBackend),
|
||||
)
|
||||
.unwrap();
|
||||
runtime
|
||||
.create_worker_scoped(
|
||||
&worker_runtime::RuntimeWorkspaceScope::new("local", "local-token"),
|
||||
create_request("fixture"),
|
||||
)
|
||||
.unwrap();
|
||||
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let task = start_runtime_server(listener, runtime.clone()).await;
|
||||
let config = RemoteRuntimeConfig::new(
|
||||
"runtime-test",
|
||||
"Runtime test",
|
||||
format!("http://{address}"),
|
||||
Some(TOKEN.to_string()),
|
||||
);
|
||||
(runtime, config, task)
|
||||
}
|
||||
|
||||
async fn next_event(subscription: &mut BrokerSubscription) -> BrokerSubscriptionEvent {
|
||||
tokio::time::timeout(Duration::from_secs(5), subscription.recv())
|
||||
.await
|
||||
.expect("subscription event timed out")
|
||||
.expect("subscription closed")
|
||||
}
|
||||
|
||||
async fn next_snapshot(subscription: &mut BrokerSubscription) -> BrokerSubscriptionEvent {
|
||||
loop {
|
||||
let event = next_event(subscription).await;
|
||||
if matches!(event, BrokerSubscriptionEvent::Snapshot { .. }) {
|
||||
return event;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait_for_status(
|
||||
broker: &RuntimeSubscriptionBroker,
|
||||
runtime_id: &str,
|
||||
predicate: impl Fn(&RuntimeSubscriptionBrokerStatus) -> bool,
|
||||
) -> RuntimeSubscriptionBrokerStatus {
|
||||
tokio::time::timeout(Duration::from_secs(5), async {
|
||||
loop {
|
||||
if let Some(status) = broker.status(runtime_id) {
|
||||
if predicate(&status) {
|
||||
return status;
|
||||
}
|
||||
}
|
||||
tokio::time::sleep(Duration::from_millis(10)).await;
|
||||
}
|
||||
})
|
||||
.await
|
||||
.expect("broker status timed out")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn equal_downstream_selectors_share_one_upstream_subscription() {
|
||||
let (runtime, config, server) = fixture().await;
|
||||
let worker = runtime.list_workers().unwrap().remove(0);
|
||||
let broker = RuntimeSubscriptionBroker::new("local");
|
||||
broker.register_remote_runtime(config);
|
||||
let selector = EventSubscriptionSelector::WorkerLifecycle {
|
||||
worker_ids: SubscriptionWorkerIds::new([
|
||||
protocol::subscription::SubscriptionWorkerId::new(
|
||||
worker.worker_ref.worker_id.to_string(),
|
||||
)
|
||||
.unwrap(),
|
||||
])
|
||||
.unwrap(),
|
||||
};
|
||||
let mut first = broker.subscribe("runtime-test", selector.clone()).unwrap();
|
||||
let mut second = broker.subscribe("runtime-test", selector).unwrap();
|
||||
assert!(matches!(
|
||||
next_snapshot(&mut first).await,
|
||||
BrokerSubscriptionEvent::Snapshot { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
next_snapshot(&mut second).await,
|
||||
BrokerSubscriptionEvent::Snapshot { .. }
|
||||
));
|
||||
let status = wait_for_status(&broker, "runtime-test", |status| {
|
||||
status.upstream_subscriptions == 1
|
||||
})
|
||||
.await;
|
||||
assert_eq!(status.desired_selectors, 1);
|
||||
|
||||
runtime
|
||||
.observe_worker_event(
|
||||
&worker.worker_ref,
|
||||
protocol::Event::Status {
|
||||
status: protocol::WorkerStatus::Running,
|
||||
},
|
||||
)
|
||||
.unwrap();
|
||||
for subscription in [&mut first, &mut second] {
|
||||
assert!(matches!(
|
||||
next_event(subscription).await,
|
||||
BrokerSubscriptionEvent::Event {
|
||||
payload: SubscriptionEventPayload::WorkerUpserted { ref worker },
|
||||
..
|
||||
} if worker.state == SubscriptionWorkerState::Running
|
||||
));
|
||||
}
|
||||
|
||||
drop(first);
|
||||
tokio::task::yield_now().await;
|
||||
assert_eq!(
|
||||
broker
|
||||
.status("runtime-test")
|
||||
.unwrap()
|
||||
.upstream_subscriptions,
|
||||
1
|
||||
);
|
||||
drop(second);
|
||||
wait_for_status(&broker, "runtime-test", |status| {
|
||||
status.upstream_subscriptions == 0 && status.desired_selectors == 0
|
||||
})
|
||||
.await;
|
||||
server.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn replacing_runtime_registration_fences_the_old_generation() {
|
||||
let (_runtime, config, server) = fixture().await;
|
||||
let broker = RuntimeSubscriptionBroker::new("local");
|
||||
let first_generation = broker.register_remote_runtime(config.clone());
|
||||
let mut old = broker
|
||||
.subscribe("runtime-test", EventSubscriptionSelector::RuntimeWorkers)
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
next_snapshot(&mut old).await,
|
||||
BrokerSubscriptionEvent::Snapshot { .. }
|
||||
));
|
||||
|
||||
let second_generation = broker.register_remote_runtime(config);
|
||||
assert!(second_generation > first_generation);
|
||||
assert!(matches!(
|
||||
next_event(&mut old).await,
|
||||
BrokerSubscriptionEvent::Closed {
|
||||
connection_generation,
|
||||
..
|
||||
} if connection_generation == first_generation
|
||||
));
|
||||
let status = wait_for_status(&broker, "runtime-test", |status| {
|
||||
status.connection_generation == second_generation && status.connected
|
||||
})
|
||||
.await;
|
||||
assert_eq!(status.connection_generation, second_generation);
|
||||
server.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reconnect_resubscribes_and_replaces_state_from_fresh_snapshot() {
|
||||
let (runtime, config, server) = fixture().await;
|
||||
let address = config
|
||||
.base_url
|
||||
.strip_prefix("http://")
|
||||
.unwrap()
|
||||
.parse::<std::net::SocketAddr>()
|
||||
.unwrap();
|
||||
server.abort();
|
||||
tokio::task::yield_now().await;
|
||||
|
||||
let broker = RuntimeSubscriptionBroker::new("local");
|
||||
broker.register_remote_runtime(config);
|
||||
let mut subscription = broker
|
||||
.subscribe("runtime-test", EventSubscriptionSelector::RuntimeWorkers)
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
next_event(&mut subscription).await,
|
||||
BrokerSubscriptionEvent::Disconnected { .. }
|
||||
));
|
||||
|
||||
runtime
|
||||
.create_worker_scoped(
|
||||
&worker_runtime::RuntimeWorkspaceScope::new("local", "local-token"),
|
||||
create_request("after-reconnect"),
|
||||
)
|
||||
.unwrap();
|
||||
let listener = tokio::net::TcpListener::bind(address).await.unwrap();
|
||||
let restarted = start_runtime_server(listener, runtime).await;
|
||||
let event = next_snapshot(&mut subscription).await;
|
||||
let BrokerSubscriptionEvent::Snapshot { snapshot, .. } = event else {
|
||||
panic!("expected fresh snapshot after reconnect");
|
||||
};
|
||||
let SubscriptionSnapshot::Workers { workers } = snapshot else {
|
||||
panic!("expected Worker snapshot");
|
||||
};
|
||||
assert_eq!(workers.len(), 2);
|
||||
restarted.abort();
|
||||
}
|
||||
@@ -90,6 +90,7 @@ use crate::repositories::{
|
||||
RepositoryRegistryReader, RepositorySummary,
|
||||
};
|
||||
use crate::resource_broker::BackendResourceBroker;
|
||||
use crate::runtime_subscription::RuntimeSubscriptionBroker;
|
||||
use crate::skills;
|
||||
use crate::store::{
|
||||
AccountRecord, ApiTokenRecord, AuthChallengeRecord, BrowserSessionRecord, ControlPlaneStore,
|
||||
@@ -248,6 +249,7 @@ pub struct WorkspaceApi {
|
||||
runtime: Arc<RuntimeRegistry>,
|
||||
companion: Arc<CompanionConsole>,
|
||||
observation_proxy: BackendObservationProxy,
|
||||
runtime_subscription_broker: RuntimeSubscriptionBroker,
|
||||
resource_broker: BackendResourceBroker,
|
||||
credential_operation_lock: Arc<std::sync::Mutex<()>>,
|
||||
}
|
||||
@@ -317,19 +319,21 @@ impl WorkspaceApi {
|
||||
crate::Error::Store(format!("invalid embedded Worker backend: {err}"))
|
||||
})?,
|
||||
);
|
||||
let runtime_subscription_broker =
|
||||
RuntimeSubscriptionBroker::new(config.workspace_id.clone());
|
||||
for remote_config in config.remote_runtime_sources.iter().cloned() {
|
||||
runtime.register(
|
||||
RemoteWorkerRuntime::new(
|
||||
remote_config,
|
||||
config.workspace_id.clone(),
|
||||
config
|
||||
.backend_base_url
|
||||
.clone()
|
||||
.unwrap_or_else(|| "http://127.0.0.1:8787".to_string()),
|
||||
)
|
||||
.map(|host| host.with_resource_broker(resource_broker.clone()))
|
||||
.map_err(|err| err.into_error())?,
|
||||
);
|
||||
let remote_runtime = RemoteWorkerRuntime::new(
|
||||
remote_config.clone(),
|
||||
config.workspace_id.clone(),
|
||||
config
|
||||
.backend_base_url
|
||||
.clone()
|
||||
.unwrap_or_else(|| "http://127.0.0.1:8787".to_string()),
|
||||
)
|
||||
.map(|host| host.with_resource_broker(resource_broker.clone()))
|
||||
.map_err(|err| err.into_error())?;
|
||||
runtime.register(remote_runtime);
|
||||
runtime_subscription_broker.register_remote_runtime(remote_config);
|
||||
}
|
||||
let runtime = Arc::new(runtime);
|
||||
let companion = Arc::new(CompanionConsole::disabled());
|
||||
@@ -344,6 +348,7 @@ impl WorkspaceApi {
|
||||
runtime,
|
||||
companion,
|
||||
observation_proxy,
|
||||
runtime_subscription_broker,
|
||||
resource_broker,
|
||||
credential_operation_lock: Arc::new(std::sync::Mutex::new(())),
|
||||
})
|
||||
@@ -353,6 +358,10 @@ impl WorkspaceApi {
|
||||
self.config.workspace_id.as_str()
|
||||
}
|
||||
|
||||
pub fn runtime_subscription_broker(&self) -> &RuntimeSubscriptionBroker {
|
||||
&self.runtime_subscription_broker
|
||||
}
|
||||
|
||||
fn mint_worker_workspace_credential(
|
||||
&self,
|
||||
runtime_id: &str,
|
||||
|
||||
Reference in New Issue
Block a user