Replace Worker list revisions with ordered subscriptions

This commit is contained in:
2026-10-02 10:41:02 +09:00
parent 8f8c2676b4
commit 2c329072a2
24 changed files with 1320 additions and 963 deletions
+7 -14
View File
@@ -12,7 +12,7 @@ use serde::{Deserialize, Deserializer, Serialize};
use crate::Event as WorkerProtocolEvent; use crate::Event as WorkerProtocolEvent;
pub const SUBSCRIPTION_PROTOCOL_VERSION: u16 = 1; pub const SUBSCRIPTION_PROTOCOL_VERSION: u16 = 2;
/// Browser/runtime wire guardrails. These bounds are generated into the Workspace /// Browser/runtime wire guardrails. These bounds are generated into the Workspace
/// Browser validator so aggregate limits remain part of the Rust-owned contract. /// Browser validator so aggregate limits remain part of the Rust-owned contract.
pub const MAX_SUBSCRIPTION_FRAME_JSON_BYTES: usize = 16 * 1024 * 1024; pub const MAX_SUBSCRIPTION_FRAME_JSON_BYTES: usize = 16 * 1024 * 1024;
@@ -426,7 +426,6 @@ pub enum SubscriptionResponse {
request_id: SubscriptionRequestId, request_id: SubscriptionRequestId,
subscription_id: SubscriptionId, subscription_id: SubscriptionId,
selector: EventSubscriptionSelector, selector: EventSubscriptionSelector,
snapshot_revision: u64,
snapshot: SubscriptionSnapshot, snapshot: SubscriptionSnapshot,
}, },
Unsubscribed { Unsubscribed {
@@ -512,7 +511,6 @@ pub enum SubscriptionTerminationCode {
pub enum SubscriptionEvent { pub enum SubscriptionEvent {
Event { Event {
subscription_id: SubscriptionId, subscription_id: SubscriptionId,
subject_revision: u64,
payload: SubscriptionEventPayload, payload: SubscriptionEventPayload,
}, },
SubscriptionClosed { SubscriptionClosed {
@@ -666,9 +664,7 @@ pub struct SubscriptionWorker {
/// `Unavailable` preserves catalog membership without claiming that execution stopped. /// `Unavailable` preserves catalog membership without claiming that execution stopped.
#[serde(default)] #[serde(default)]
pub availability: SubscriptionWorkerAvailability, pub availability: SubscriptionWorkerAvailability,
/// Producer-owned monotonic revision for this Worker subject. /// Latest foreground state observed from the Worker. This remains
pub subject_revision: u64,
/// Latest revisioned foreground state observed from the Worker. This remains
/// absent until an authoritative Worker snapshot/event has been applied. /// absent until an authoritative Worker snapshot/event has been applied.
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub worker_state: Option<crate::WorkerStateSnapshot>, pub worker_state: Option<crate::WorkerStateSnapshot>,
@@ -970,7 +966,6 @@ mod tests {
runtime_id: None, runtime_id: None,
resource_key: None, resource_key: None,
availability: SubscriptionWorkerAvailability::Observed, availability: SubscriptionWorkerAvailability::Observed,
subject_revision: 0,
worker_state: None, worker_state: None,
state: SubscriptionWorkerState::Idle, state: SubscriptionWorkerState::Idle,
has_running_internal_workers: false, has_running_internal_workers: false,
@@ -1093,7 +1088,7 @@ mod tests {
assert_eq!( assert_eq!(
json, json,
serde_json::json!({ serde_json::json!({
"protocol_version": 1, "protocol_version": 2,
"frame": "request", "frame": "request",
"message": { "message": {
"method": "subscribe_events", "method": "subscribe_events",
@@ -1119,7 +1114,6 @@ mod tests {
request_id: request_id(), request_id: request_id(),
subscription_id: subscription_id(), subscription_id: subscription_id(),
selector: EventSubscriptionSelector::RuntimeWorkers, selector: EventSubscriptionSelector::RuntimeWorkers,
snapshot_revision: 7,
snapshot: SubscriptionSnapshot::Workers { snapshot: SubscriptionSnapshot::Workers {
workers: vec![worker("worker-1")], workers: vec![worker("worker-1")],
}, },
@@ -1158,10 +1152,10 @@ mod tests {
subscription_id: subscription_id(), subscription_id: subscription_id(),
}, },
)); ));
frame.protocol_version = 2; frame.protocol_version = 1;
assert!(matches!( assert!(matches!(
frame.validate(), frame.validate(),
Err(SubscriptionValidationError::UnsupportedProtocolVersion { actual: 2 }) Err(SubscriptionValidationError::UnsupportedProtocolVersion { actual: 1 })
)); ));
} }
@@ -1180,7 +1174,6 @@ mod tests {
let event = SubscriptionEvent::Event { let event = SubscriptionEvent::Event {
subscription_id: subscription_id(), subscription_id: subscription_id(),
subject_revision: 8,
payload: SubscriptionEventPayload::WorkerRemoved { payload: SubscriptionEventPayload::WorkerRemoved {
worker_id: worker_id("worker-2"), worker_id: worker_id("worker-2"),
runtime_id: None, runtime_id: None,
@@ -1211,7 +1204,7 @@ mod tests {
assert_eq!( assert_eq!(
serde_json::to_value(frame).unwrap(), serde_json::to_value(frame).unwrap(),
serde_json::json!({ serde_json::json!({
"protocol_version": 1, "protocol_version": 2,
"frame": "worker_protocol", "frame": "worker_protocol",
"message": { "message": {
"subscription_id": "subscription-1", "subscription_id": "subscription-1",
@@ -1253,7 +1246,7 @@ mod tests {
assert_eq!( assert_eq!(
serde_json::to_value(frame).unwrap(), serde_json::to_value(frame).unwrap(),
serde_json::json!({ serde_json::json!({
"protocol_version": 1, "protocol_version": 2,
"frame": "event", "frame": "event",
"message": { "message": {
"event": "subscription_closed", "event": "subscription_closed",
-3
View File
@@ -230,7 +230,6 @@ fn rust_serialized_subscription_frame_fixtures() -> Vec<SubscriptionFrame> {
runtime_id: Some("fixture-runtime".to_string()), runtime_id: Some("fixture-runtime".to_string()),
resource_key: Some("fixture-worker-resource".to_string()), resource_key: Some("fixture-worker-resource".to_string()),
availability: SubscriptionWorkerAvailability::default(), availability: SubscriptionWorkerAvailability::default(),
subject_revision: 1,
worker_state: None, worker_state: None,
state: SubscriptionWorkerState::Idle, state: SubscriptionWorkerState::Idle,
has_running_internal_workers: false, has_running_internal_workers: false,
@@ -250,7 +249,6 @@ fn rust_serialized_subscription_frame_fixtures() -> Vec<SubscriptionFrame> {
request_id, request_id,
subscription_id: subscription_id.clone(), subscription_id: subscription_id.clone(),
selector: EventSubscriptionSelector::WorkspaceWorkers, selector: EventSubscriptionSelector::WorkspaceWorkers,
snapshot_revision: 1,
snapshot: SubscriptionSnapshot::Workers { snapshot: SubscriptionSnapshot::Workers {
workers: vec![worker], workers: vec![worker],
}, },
@@ -259,7 +257,6 @@ fn rust_serialized_subscription_frame_fixtures() -> Vec<SubscriptionFrame> {
let pending = let pending =
SubscriptionFrame::new(SubscriptionFramePayload::Event(SubscriptionEvent::Event { SubscriptionFrame::new(SubscriptionFramePayload::Event(SubscriptionEvent::Event {
subscription_id, subscription_id,
subject_revision: 2,
payload: SubscriptionEventPayload::WorkerProtocol { payload: SubscriptionEventPayload::WorkerProtocol {
worker_id, worker_id,
event: Event::PendingSubmissionsChanged { event: Event::PendingSubmissionsChanged {
-2
View File
@@ -1334,7 +1334,6 @@ async fn runtime_protocol_ws_session(
request_id, request_id,
subscription_id: subscription_id.clone(), subscription_id: subscription_id.clone(),
selector: selector.clone(), selector: selector.clone(),
snapshot_revision: subscription.snapshot_revision(),
snapshot: subscription.snapshot().clone(), snapshot: subscription.snapshot().clone(),
}, },
)); ));
@@ -1354,7 +1353,6 @@ async fn runtime_protocol_ws_session(
SubscriptionFramePayload::Event( SubscriptionFramePayload::Event(
SubscriptionEvent::Event { SubscriptionEvent::Event {
subscription_id: event_subscription_id.clone(), subscription_id: event_subscription_id.clone(),
subject_revision: update.subject_revision,
payload: update.payload, payload: update.payload,
}, },
), ),
+4 -1
View File
@@ -36,5 +36,8 @@ pub mod workspace_request;
#[cfg(feature = "fs-store")] #[cfg(feature = "fs-store")]
pub use fs_store::{FsRuntimeStore, FsRuntimeStoreOptions}; pub use fs_store::{FsRuntimeStore, FsRuntimeStoreOptions};
pub use management::RuntimeOptions; pub use management::RuntimeOptions;
pub use runtime::{Runtime, RuntimeWorkspaceScope}; pub use runtime::{
Runtime, RuntimeEventSelectorSubscription, RuntimeSubscriptionRecvError,
RuntimeSubscriptionUpdate, RuntimeWorkspaceScope,
};
pub use session_store::UploadedFileUploadContext; pub use session_store::UploadedFileUploadContext;
+98 -60
View File
@@ -151,7 +151,6 @@ impl RuntimeWorkerRestoreResult {
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct RuntimeSubscriptionUpdate { pub struct RuntimeSubscriptionUpdate {
pub subject_revision: u64,
pub payload: SubscriptionEventPayload, pub payload: SubscriptionEventPayload,
} }
@@ -168,7 +167,6 @@ pub enum RuntimeSubscriptionRecvError {
pub struct RuntimeEventSelectorSubscription { pub struct RuntimeEventSelectorSubscription {
subscription_id: u64, subscription_id: u64,
selector: EventSubscriptionSelector, selector: EventSubscriptionSelector,
snapshot_revision: u64,
snapshot: SubscriptionSnapshot, snapshot: SubscriptionSnapshot,
receiver: mpsc::Receiver<RuntimeSubscriptionUpdate>, receiver: mpsc::Receiver<RuntimeSubscriptionUpdate>,
lagged: Arc<AtomicBool>, lagged: Arc<AtomicBool>,
@@ -180,10 +178,6 @@ impl RuntimeEventSelectorSubscription {
&self.selector &self.selector
} }
pub fn snapshot_revision(&self) -> u64 {
self.snapshot_revision
}
pub fn snapshot(&self) -> &SubscriptionSnapshot { pub fn snapshot(&self) -> &SubscriptionSnapshot {
&self.snapshot &self.snapshot
} }
@@ -1280,7 +1274,6 @@ impl Runtime {
state.persist_runtime_snapshot()?; state.persist_runtime_snapshot()?;
} }
let snapshot = state.subscription_snapshot(scope, &selector)?; let snapshot = state.subscription_snapshot(scope, &selector)?;
let snapshot_revision = state.subscription_revision;
let subscription_id = state.next_event_subscription_id; let subscription_id = state.next_event_subscription_id;
state.next_event_subscription_id = state.next_event_subscription_id.saturating_add(1); state.next_event_subscription_id = state.next_event_subscription_id.saturating_add(1);
let (sender, receiver) = mpsc::channel(SUBSCRIPTION_QUEUE_CAPACITY); let (sender, receiver) = mpsc::channel(SUBSCRIPTION_QUEUE_CAPACITY);
@@ -1298,7 +1291,6 @@ impl Runtime {
Ok(RuntimeEventSelectorSubscription { Ok(RuntimeEventSelectorSubscription {
subscription_id, subscription_id,
selector, selector,
snapshot_revision,
snapshot, snapshot,
receiver, receiver,
lagged, lagged,
@@ -2217,9 +2209,9 @@ impl Runtime {
if let Some(payload) = input_protocol_event(&input) { if let Some(payload) = input_protocol_event(&input) {
state.push_worker_observation_event(worker_ref.clone(), payload); state.push_worker_observation_event(worker_ref.clone(), payload);
} }
state.publish_worker_upsert(worker_ref.worker_id)?;
state.persist_runtime_snapshot()?; state.persist_runtime_snapshot()?;
state.persist_worker(&worker_ref.worker_id)?; state.persist_worker(&worker_ref.worker_id)?;
state.publish_worker_upsert(worker_ref.worker_id)?;
Ok(WorkerInteractionAck { Ok(WorkerInteractionAck {
worker_ref: worker_ref.clone(), worker_ref: worker_ref.clone(),
@@ -2500,9 +2492,9 @@ impl Runtime {
worker.workdir_attachments = workdir_attachments; worker.workdir_attachments = workdir_attachments;
worker.detail() worker.detail()
}; };
state.publish_worker_upsert(worker_ref.worker_id)?;
state.persist_runtime_snapshot()?; state.persist_runtime_snapshot()?;
state.persist_worker(&worker_ref.worker_id)?; state.persist_worker(&worker_ref.worker_id)?;
state.publish_worker_upsert(worker_ref.worker_id)?;
Ok(detail) Ok(detail)
} }
@@ -2541,8 +2533,8 @@ impl Runtime {
state state
.observation_events .observation_events
.retain(|event| event.worker_ref != *worker_ref); .retain(|event| event.worker_ref != *worker_ref);
state.publish_worker_removed(worker_ref.worker_id, workspace_id.as_deref())?;
state.persist_runtime_snapshot()?; state.persist_runtime_snapshot()?;
state.publish_worker_removed(worker_ref.worker_id, workspace_id.as_deref())?;
} }
Ok(()) Ok(())
} }
@@ -2564,9 +2556,9 @@ impl Runtime {
let mut state = self.lock()?; let mut state = self.lock()?;
let worker = state.worker_mut(worker_ref)?; let worker = state.worker_mut(worker_ref)?;
worker.apply_worker_state(&snapshot); worker.apply_worker_state(&snapshot);
state.publish_worker_upsert(worker_ref.worker_id)?;
state.persist_runtime_snapshot()?; state.persist_runtime_snapshot()?;
state.persist_worker(&worker_ref.worker_id)?; state.persist_worker(&worker_ref.worker_id)?;
state.publish_worker_upsert(worker_ref.worker_id)?;
Ok(()) Ok(())
} }
@@ -2870,8 +2862,8 @@ impl Runtime {
state state
.observation_events .observation_events
.retain(|event| event.worker_ref != *worker_ref); .retain(|event| event.worker_ref != *worker_ref);
state.publish_worker_removed(worker_ref.worker_id, removed_workspace_id.as_deref())?;
state.persist_runtime_snapshot()?; state.persist_runtime_snapshot()?;
state.publish_worker_removed(worker_ref.worker_id, removed_workspace_id.as_deref())?;
Ok(WorkerDeleteResult { Ok(WorkerDeleteResult {
worker_id: removed.worker_id, worker_id: removed.worker_id,
deleted: true, deleted: true,
@@ -2959,13 +2951,13 @@ impl Runtime {
let worker_state_changed = let worker_state_changed =
state.project_protocol_event_to_worker_state(worker_ref, &payload); state.project_protocol_event_to_worker_state(worker_ref, &payload);
let activity_changed = state.project_internal_worker_activity(worker_ref, &payload); let activity_changed = state.project_internal_worker_activity(worker_ref, &payload);
if worker_state_changed || activity_changed {
state.publish_worker_upsert(worker_ref.worker_id)?;
}
if worker_state_changed { if worker_state_changed {
state.persist_runtime_snapshot()?; state.persist_runtime_snapshot()?;
state.persist_worker(&worker_ref.worker_id)?; state.persist_worker(&worker_ref.worker_id)?;
} }
if worker_state_changed || activity_changed {
state.publish_worker_upsert(worker_ref.worker_id)?;
}
let event = state.push_worker_observation_event(worker_ref.clone(), payload); let event = state.push_worker_observation_event(worker_ref.clone(), payload);
Ok(event) Ok(event)
} }
@@ -3021,13 +3013,13 @@ impl Runtime {
state.project_protocol_event_to_worker_state(&worker_ref, &payload); state.project_protocol_event_to_worker_state(&worker_ref, &payload);
let activity_changed = let activity_changed =
state.project_internal_worker_activity(&worker_ref, &payload); state.project_internal_worker_activity(&worker_ref, &payload);
if worker_state_changed || activity_changed {
state.publish_worker_upsert(worker_ref.worker_id)?;
}
if worker_state_changed { if worker_state_changed {
state.persist_runtime_snapshot()?; state.persist_runtime_snapshot()?;
state.persist_worker(&worker_ref.worker_id)?; state.persist_worker(&worker_ref.worker_id)?;
} }
if worker_state_changed || activity_changed {
state.publish_worker_upsert(worker_ref.worker_id)?;
}
Ok(state.push_worker_observation_event(worker_ref, payload)) Ok(state.push_worker_observation_event(worker_ref, payload))
}), }),
) )
@@ -3054,13 +3046,13 @@ impl Runtime {
state.project_protocol_event_to_worker_state(&worker_ref, &payload); state.project_protocol_event_to_worker_state(&worker_ref, &payload);
let activity_changed = let activity_changed =
state.project_internal_worker_activity(&worker_ref, &payload); state.project_internal_worker_activity(&worker_ref, &payload);
if worker_state_changed || activity_changed {
state.publish_worker_upsert(worker_ref.worker_id)?;
}
if worker_state_changed { if worker_state_changed {
state.persist_runtime_snapshot()?; state.persist_runtime_snapshot()?;
state.persist_worker(&worker_ref.worker_id)?; state.persist_worker(&worker_ref.worker_id)?;
} }
if worker_state_changed || activity_changed {
state.publish_worker_upsert(worker_ref.worker_id)?;
}
Ok(state.push_worker_observation_event(worker_ref, payload)) Ok(state.push_worker_observation_event(worker_ref, payload))
}), }),
) )
@@ -3454,8 +3446,8 @@ impl Runtime {
logical_workdir_attachments_from_statuses(&workdir_attachments, claims); logical_workdir_attachments_from_statuses(&workdir_attachments, claims);
candidate.workdir_attachments = workdir_attachments; candidate.workdir_attachments = workdir_attachments;
state.workers.insert(worker_ref.worker_id, candidate); state.workers.insert(worker_ref.worker_id, candidate);
state.publish_worker_upsert(worker_ref.worker_id)?;
state.persist_worker(&worker_ref.worker_id)?; state.persist_worker(&worker_ref.worker_id)?;
state.publish_worker_upsert(worker_ref.worker_id)?;
Ok(()) Ok(())
} }
@@ -3702,8 +3694,6 @@ struct RuntimeState {
workspace_config_latest: BTreeMap<String, ConfigBundleRef>, workspace_config_latest: BTreeMap<String, ConfigBundleRef>,
workspace_config_fetch_gates: BTreeMap<String, Arc<Mutex<()>>>, workspace_config_fetch_gates: BTreeMap<String, Arc<Mutex<()>>>,
diagnostics: Vec<RuntimeDiagnostic>, diagnostics: Vec<RuntimeDiagnostic>,
subscription_revision: u64,
worker_subject_revisions: BTreeMap<WorkerId, u64>,
next_event_subscription_id: u64, next_event_subscription_id: u64,
subscriptions: BTreeMap<u64, SubscriptionSink>, subscriptions: BTreeMap<u64, SubscriptionSink>,
#[cfg(feature = "ws-server")] #[cfg(feature = "ws-server")]
@@ -3733,8 +3723,6 @@ impl RuntimeState {
workspace_config_latest: BTreeMap::new(), workspace_config_latest: BTreeMap::new(),
workspace_config_fetch_gates: BTreeMap::new(), workspace_config_fetch_gates: BTreeMap::new(),
diagnostics: Vec::new(), diagnostics: Vec::new(),
subscription_revision: 0,
worker_subject_revisions: BTreeMap::new(),
next_event_subscription_id: 1, next_event_subscription_id: 1,
subscriptions: BTreeMap::new(), subscriptions: BTreeMap::new(),
#[cfg(feature = "ws-server")] #[cfg(feature = "ws-server")]
@@ -3765,8 +3753,6 @@ impl RuntimeState {
workspace_config_latest: BTreeMap::new(), workspace_config_latest: BTreeMap::new(),
workspace_config_fetch_gates: BTreeMap::new(), workspace_config_fetch_gates: BTreeMap::new(),
diagnostics: Vec::new(), diagnostics: Vec::new(),
subscription_revision: 0,
worker_subject_revisions: BTreeMap::new(),
next_event_subscription_id: 1, next_event_subscription_id: 1,
subscriptions: BTreeMap::new(), subscriptions: BTreeMap::new(),
#[cfg(feature = "ws-server")] #[cfg(feature = "ws-server")]
@@ -3895,8 +3881,6 @@ impl RuntimeState {
workspace_config_fetch_gates: BTreeMap::new(), workspace_config_fetch_gates: BTreeMap::new(),
workspace_owners: persisted.workspace_owners, workspace_owners: persisted.workspace_owners,
diagnostics, diagnostics,
subscription_revision: 0,
worker_subject_revisions: BTreeMap::new(),
next_event_subscription_id: 1, next_event_subscription_id: 1,
subscriptions: BTreeMap::new(), subscriptions: BTreeMap::new(),
#[cfg(feature = "ws-server")] #[cfg(feature = "ws-server")]
@@ -4215,11 +4199,6 @@ impl RuntimeState {
runtime_id: None, runtime_id: None,
resource_key: None, resource_key: None,
availability: protocol::subscription::SubscriptionWorkerAvailability::Observed, availability: protocol::subscription::SubscriptionWorkerAvailability::Observed,
subject_revision: self
.worker_subject_revisions
.get(&worker.worker_id)
.copied()
.unwrap_or(0),
worker_state: worker.worker_state.clone(), worker_state: worker.worker_state.clone(),
state: subscription_worker_state(worker.status), state: subscription_worker_state(worker.status),
has_running_internal_workers: worker.internal_workers.has_running_worker(), has_running_internal_workers: worker.internal_workers.has_running_worker(),
@@ -4232,25 +4211,17 @@ impl RuntimeState {
} }
fn publish_worker_upsert(&mut self, worker_id: WorkerId) -> Result<(), RuntimeError> { fn publish_worker_upsert(&mut self, worker_id: WorkerId) -> Result<(), RuntimeError> {
self.subscription_revision = self.subscription_revision.saturating_add(1);
let subject_revision = {
let revision = self.worker_subject_revisions.entry(worker_id).or_insert(0);
*revision = revision.saturating_add(1);
*revision
};
let worker = self let worker = self
.workers .workers
.get(&worker_id) .get(&worker_id)
.ok_or(RuntimeError::WorkerNotFound { worker_id })?; .ok_or(RuntimeError::WorkerNotFound { worker_id })?;
let workspace_id = worker.workspace_id.clone(); let workspace_id = worker.workspace_id.clone();
let mut projected = self.subscription_worker(worker)?; let projected = self.subscription_worker(worker)?;
projected.subject_revision = subject_revision;
let projected_worker_id = projected.worker_id.clone(); let projected_worker_id = projected.worker_id.clone();
self.deliver_worker_subscription_update( self.deliver_worker_subscription_update(
&projected_worker_id, &projected_worker_id,
workspace_id.as_deref(), workspace_id.as_deref(),
RuntimeSubscriptionUpdate { RuntimeSubscriptionUpdate {
subject_revision,
payload: SubscriptionEventPayload::WorkerUpserted { worker: projected }, payload: SubscriptionEventPayload::WorkerUpserted { worker: projected },
}, },
); );
@@ -4262,19 +4233,12 @@ impl RuntimeState {
worker_id: WorkerId, worker_id: WorkerId,
workspace_id: Option<&str>, workspace_id: Option<&str>,
) -> Result<(), RuntimeError> { ) -> Result<(), RuntimeError> {
self.subscription_revision = self.subscription_revision.saturating_add(1);
let subject_revision = {
let revision = self.worker_subject_revisions.entry(worker_id).or_insert(0);
*revision = revision.saturating_add(1);
*revision
};
let worker_id = SubscriptionWorkerId::new(worker_id.to_string()) let worker_id = SubscriptionWorkerId::new(worker_id.to_string())
.map_err(subscription_validation_error)?; .map_err(subscription_validation_error)?;
self.deliver_worker_subscription_update( self.deliver_worker_subscription_update(
&worker_id, &worker_id,
workspace_id, workspace_id,
RuntimeSubscriptionUpdate { RuntimeSubscriptionUpdate {
subject_revision,
payload: SubscriptionEventPayload::WorkerRemoved { payload: SubscriptionEventPayload::WorkerRemoved {
worker_id: worker_id.clone(), worker_id: worker_id.clone(),
runtime_id: None, runtime_id: None,
@@ -4417,8 +4381,8 @@ impl RuntimeState {
worker_ref: Some(worker_ref.clone()), worker_ref: Some(worker_ref.clone()),
}); });
self.workers.insert(worker_ref.worker_id, candidate); self.workers.insert(worker_ref.worker_id, candidate);
self.publish_worker_upsert(worker_ref.worker_id)?;
self.persist_runtime_snapshot()?; self.persist_runtime_snapshot()?;
self.publish_worker_upsert(worker_ref.worker_id)?;
Ok(()) Ok(())
} }
@@ -6637,7 +6601,6 @@ mod tests {
let mut subscription = runtime let mut subscription = runtime
.subscribe_event_selector(EventSubscriptionSelector::RuntimeWorkers) .subscribe_event_selector(EventSubscriptionSelector::RuntimeWorkers)
.unwrap(); .unwrap();
assert_eq!(subscription.snapshot_revision(), 0);
let SubscriptionSnapshot::Workers { workers } = subscription.snapshot() else { let SubscriptionSnapshot::Workers { workers } = subscription.snapshot() else {
panic!("runtime_workers must return a Worker snapshot"); panic!("runtime_workers must return a Worker snapshot");
}; };
@@ -6645,11 +6608,9 @@ mod tests {
let created = runtime.create_worker(task_request("live")).unwrap(); let created = runtime.create_worker(task_request("live")).unwrap();
let update = receive_subscription_update(&mut subscription).unwrap(); let update = receive_subscription_update(&mut subscription).unwrap();
assert_eq!(update.subject_revision, 1);
match update.payload { match update.payload {
SubscriptionEventPayload::WorkerUpserted { worker } => { SubscriptionEventPayload::WorkerUpserted { worker } => {
assert_eq!(worker.worker_id.as_str(), created.worker_id.to_string()); assert_eq!(worker.worker_id.as_str(), created.worker_id.to_string());
assert_eq!(worker.subject_revision, 1);
assert_eq!(worker.state, SubscriptionWorkerState::Idle); assert_eq!(worker.state, SubscriptionWorkerState::Idle);
} }
payload => panic!("unexpected subscription payload: {payload:?}"), payload => panic!("unexpected subscription payload: {payload:?}"),
@@ -6666,7 +6627,6 @@ mod tests {
) )
.unwrap(); .unwrap();
let update = receive_subscription_update(&mut subscription).unwrap(); let update = receive_subscription_update(&mut subscription).unwrap();
assert_eq!(update.subject_revision, 2);
match update.payload { match update.payload {
SubscriptionEventPayload::WorkerUpserted { worker } => { SubscriptionEventPayload::WorkerUpserted { worker } => {
assert_eq!(worker.state, SubscriptionWorkerState::Idle); assert_eq!(worker.state, SubscriptionWorkerState::Idle);
@@ -6685,7 +6645,6 @@ mod tests {
) )
.unwrap(); .unwrap();
let update = receive_subscription_update(&mut subscription).unwrap(); let update = receive_subscription_update(&mut subscription).unwrap();
assert_eq!(update.subject_revision, 3);
match update.payload { match update.payload {
SubscriptionEventPayload::WorkerUpserted { worker } => { SubscriptionEventPayload::WorkerUpserted { worker } => {
assert_eq!(worker.state, SubscriptionWorkerState::Idle); assert_eq!(worker.state, SubscriptionWorkerState::Idle);
@@ -6727,7 +6686,6 @@ mod tests {
) )
.unwrap(); .unwrap();
let update = receive_subscription_update(&mut subscription).unwrap(); let update = receive_subscription_update(&mut subscription).unwrap();
assert_eq!(update.subject_revision, 4);
match update.payload { match update.payload {
SubscriptionEventPayload::WorkerUpserted { worker } => { SubscriptionEventPayload::WorkerUpserted { worker } => {
assert_eq!(worker.state, SubscriptionWorkerState::Idle); assert_eq!(worker.state, SubscriptionWorkerState::Idle);
@@ -6738,7 +6696,6 @@ mod tests {
runtime.stop_worker(&created.worker_ref, None).unwrap(); runtime.stop_worker(&created.worker_ref, None).unwrap();
let update = receive_subscription_update(&mut subscription).unwrap(); let update = receive_subscription_update(&mut subscription).unwrap();
assert_eq!(update.subject_revision, 5);
match update.payload { match update.payload {
SubscriptionEventPayload::WorkerUpserted { worker } => { SubscriptionEventPayload::WorkerUpserted { worker } => {
assert_eq!(worker.worker_id.as_str(), created.worker_id.to_string()); assert_eq!(worker.worker_id.as_str(), created.worker_id.to_string());
@@ -6750,6 +6707,87 @@ mod tests {
runtime.stop_runtime().unwrap(); runtime.stop_runtime().unwrap();
} }
#[test]
fn concurrent_snapshot_and_worker_creation_cross_the_boundary_exactly_once() {
let runtime = runtime_with_backend();
let barrier = Arc::new(Barrier::new(2));
let subscribing_runtime = runtime.clone();
let subscribing_barrier = barrier.clone();
let subscriber = std::thread::spawn(move || {
subscribing_barrier.wait();
subscribing_runtime
.subscribe_event_selector(EventSubscriptionSelector::RuntimeWorkers)
.unwrap()
});
let creating_runtime = runtime.clone();
let creating_barrier = barrier.clone();
let creator = std::thread::spawn(move || {
creating_barrier.wait();
creating_runtime
.create_worker(task_request("concurrent-boundary"))
.unwrap()
});
let mut subscription = subscriber.join().unwrap();
let created = creator.join().unwrap();
let SubscriptionSnapshot::Workers { workers } = subscription.snapshot() else {
panic!("runtime_workers must return a Worker snapshot");
};
let in_snapshot = workers
.iter()
.any(|worker| worker.worker_id.as_str() == created.worker_id.to_string());
let in_event = match subscription.receiver.try_recv() {
Ok(RuntimeSubscriptionUpdate {
payload: SubscriptionEventPayload::WorkerUpserted { worker },
}) => worker.worker_id.as_str() == created.worker_id.to_string(),
Ok(update) => panic!("unexpected subscription update: {update:?}"),
Err(mpsc::error::TryRecvError::Empty) => false,
Err(error) => panic!("unexpected subscription receive error: {error}"),
};
assert_ne!(
in_snapshot, in_event,
"the concurrent creation must appear in the snapshot or its ordered suffix, never both"
);
}
#[cfg(feature = "fs-store")]
#[test]
fn restarted_runtime_subscription_starts_from_persisted_authoritative_snapshot() {
let root = tempfile::tempdir().unwrap();
let options = crate::fs_store::FsRuntimeStoreOptions {
root: root.path().join("runtime"),
runtime_id: "subscription-restart".to_string(),
display_name: None,
};
let runtime = Runtime::with_fs_store_and_execution_backend(
options.clone(),
Arc::new(TestExecutionBackend::default()),
)
.unwrap();
runtime.store_config_bundle(test_bundle()).unwrap();
let worker = runtime
.create_worker(task_request("persisted-subscription-worker"))
.unwrap();
drop(runtime);
let restarted = Runtime::with_fs_store_and_execution_backend(
options,
Arc::new(TestExecutionBackend::default()),
)
.unwrap();
let subscription = restarted
.subscribe_event_selector(EventSubscriptionSelector::RuntimeWorkers)
.unwrap();
let SubscriptionSnapshot::Workers { workers } = subscription.snapshot() else {
panic!("runtime_workers must return a Worker snapshot");
};
assert!(workers.iter().any(|candidate| {
candidate.worker_id.as_str() == worker.worker_id.to_string()
&& candidate.availability
== protocol::subscription::SubscriptionWorkerAvailability::Observed
}));
}
#[test] #[test]
fn worker_lifecycle_subscription_delivers_only_selected_workers() { fn worker_lifecycle_subscription_delivers_only_selected_workers() {
let runtime = runtime_with_backend(); let runtime = runtime_with_backend();
@@ -914,31 +914,15 @@ CREATE TABLE worker_registry_observations (
worker_id TEXT NOT NULL, worker_id TEXT NOT NULL,
availability TEXT NOT NULL CHECK (availability IN ('observed', 'unavailable')), availability TEXT NOT NULL CHECK (availability IN ('observed', 'unavailable')),
worker_json TEXT, worker_json TEXT,
connection_generation INTEGER NOT NULL,
subject_revision INTEGER NOT NULL,
snapshot_revision INTEGER NOT NULL,
projection_revision INTEGER NOT NULL,
observed_at TEXT NOT NULL, observed_at TEXT NOT NULL,
PRIMARY KEY (workspace_id, runtime_id, worker_id), PRIMARY KEY (workspace_id, runtime_id, worker_id),
FOREIGN KEY (workspace_id, worker_id) FOREIGN KEY (workspace_id, worker_id)
REFERENCES worker_registry(workspace_id, worker_id) ON DELETE CASCADE REFERENCES worker_registry(workspace_id, worker_id) ON DELETE CASCADE
); );
CREATE TABLE worker_registry_projection_cursors (
workspace_id TEXT NOT NULL,
runtime_id TEXT NOT NULL,
connection_generation INTEGER NOT NULL,
snapshot_revision INTEGER NOT NULL,
PRIMARY KEY (workspace_id, runtime_id)
);
CREATE TABLE worker_registry_projection_revisions (
workspace_id TEXT PRIMARY KEY,
revision INTEGER NOT NULL
);
CREATE TABLE worker_registry_projection_removals ( CREATE TABLE worker_registry_projection_removals (
workspace_id TEXT NOT NULL, workspace_id TEXT NOT NULL,
runtime_id TEXT NOT NULL, runtime_id TEXT NOT NULL,
worker_id TEXT NOT NULL, worker_id TEXT NOT NULL,
projection_revision INTEGER NOT NULL,
PRIMARY KEY (workspace_id, runtime_id, worker_id) PRIMARY KEY (workspace_id, runtime_id, worker_id)
); );
CREATE TABLE worker_registry_projection_diagnostics ( CREATE TABLE worker_registry_projection_diagnostics (
@@ -33,26 +33,19 @@ pub enum RuntimeSubscriptionBrokerError {
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub enum BrokerSubscriptionEvent { pub enum BrokerSubscriptionEvent {
Snapshot { Snapshot {
connection_generation: u64,
snapshot_revision: u64,
snapshot: SubscriptionSnapshot, snapshot: SubscriptionSnapshot,
}, },
Event { Event {
connection_generation: u64,
subject_revision: u64,
payload: SubscriptionEventPayload, payload: SubscriptionEventPayload,
}, },
Disconnected { Disconnected {
connection_generation: u64,
message: String, message: String,
}, },
Rejected { Rejected {
connection_generation: u64,
code: SubscriptionRejectionCode, code: SubscriptionRejectionCode,
message: String, message: String,
}, },
Closed { Closed {
connection_generation: u64,
code: SubscriptionTerminationCode, code: SubscriptionTerminationCode,
message: String, message: String,
}, },
@@ -60,7 +53,6 @@ pub enum BrokerSubscriptionEvent {
#[derive(Clone, Debug, Default, PartialEq, Eq)] #[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct RuntimeSubscriptionBrokerStatus { pub struct RuntimeSubscriptionBrokerStatus {
pub connection_generation: u64,
pub connected: bool, pub connected: bool,
pub desired_selectors: usize, pub desired_selectors: usize,
pub upstream_subscriptions: usize, pub upstream_subscriptions: usize,
@@ -82,7 +74,30 @@ impl BrokerSubscription {
&self.selector &self.selector
} }
pub async fn recv(&mut self) -> Option<BrokerSubscriptionEvent> { pub async fn recv(&mut self) -> Option<BrokerSubscriptionEvent> {
self.receiver.recv().await loop {
if let Some(event) = self.receiver.recv().await {
return Some(event);
}
if self.restart().is_err() {
return None;
}
}
}
pub(crate) fn restart(&mut self) -> Result<(), RuntimeSubscriptionBrokerError> {
self.commands
.send(Command::Unsubscribe(self.downstream_id))
.map_err(|_| RuntimeSubscriptionBrokerError::Closed)?;
let (events, receiver) = mpsc::channel(DOWNSTREAM_QUEUE_CAPACITY);
self.commands
.send(Command::Subscribe {
downstream_id: self.downstream_id,
selector: self.selector.clone(),
events,
})
.map_err(|_| RuntimeSubscriptionBrokerError::Closed)?;
self.receiver = receiver;
Ok(())
} }
} }
@@ -120,10 +135,7 @@ impl RuntimeSubscriptionBroker {
pub fn register_remote_runtime(&self, config: RemoteRuntimeConfig) -> u64 { pub fn register_remote_runtime(&self, config: RemoteRuntimeConfig) -> u64 {
let generation = self.next_generation.fetch_add(1, Ordering::Relaxed); let generation = self.next_generation.fetch_add(1, Ordering::Relaxed);
let (commands, receiver) = mpsc::unbounded_channel(); let (commands, receiver) = mpsc::unbounded_channel();
let status = Arc::new(RwLock::new(RuntimeSubscriptionBrokerStatus { let status = Arc::new(RwLock::new(RuntimeSubscriptionBrokerStatus::default()));
connection_generation: generation,
..Default::default()
}));
let registration = Registration { let registration = Registration {
generation, generation,
commands: commands.clone(), commands: commands.clone(),
@@ -156,7 +168,6 @@ impl RuntimeSubscriptionBroker {
let generation = self.next_generation.fetch_add(1, Ordering::Relaxed); let generation = self.next_generation.fetch_add(1, Ordering::Relaxed);
let (commands, receiver) = mpsc::unbounded_channel(); let (commands, receiver) = mpsc::unbounded_channel();
let status = Arc::new(RwLock::new(RuntimeSubscriptionBrokerStatus { let status = Arc::new(RwLock::new(RuntimeSubscriptionBrokerStatus {
connection_generation: generation,
connected: true, connected: true,
..Default::default() ..Default::default()
})); }));
@@ -257,7 +268,6 @@ impl RuntimeSubscriptionBroker {
.clone(); .clone();
if !initial_status.connected { if !initial_status.connected {
let _ = initial_events.try_send(BrokerSubscriptionEvent::Disconnected { let _ = initial_events.try_send(BrokerSubscriptionEvent::Disconnected {
connection_generation: initial_status.connection_generation,
message: "Runtime subscription connection is not currently available".to_string(), message: "Runtime subscription connection is not currently available".to_string(),
}); });
} }
@@ -286,8 +296,7 @@ struct SelectorState {
downstreams: HashMap<u64, mpsc::Sender<BrokerSubscriptionEvent>>, downstreams: HashMap<u64, mpsc::Sender<BrokerSubscriptionEvent>>,
upstream_id: Option<SubscriptionId>, upstream_id: Option<SubscriptionId>,
pending: bool, pending: bool,
snapshot: Option<(u64, SubscriptionSnapshot)>, snapshot: Option<SubscriptionSnapshot>,
revisions: HashMap<String, u64>,
} }
impl SelectorState { impl SelectorState {
fn new() -> Self { fn new() -> Self {
@@ -296,7 +305,6 @@ impl SelectorState {
upstream_id: None, upstream_id: None,
pending: false, pending: false,
snapshot: None, snapshot: None,
revisions: HashMap::new(),
} }
} }
} }
@@ -334,11 +342,9 @@ impl State {
selector.upstream_id = None; selector.upstream_id = None;
selector.pending = false; selector.pending = false;
selector.snapshot = None; selector.snapshot = None;
selector.revisions.clear();
broadcast( broadcast(
&mut selector.downstreams, &mut selector.downstreams,
BrokerSubscriptionEvent::Disconnected { BrokerSubscriptionEvent::Disconnected {
connection_generation: self.generation,
message: message.clone(), message: message.clone(),
}, },
); );
@@ -347,12 +353,77 @@ impl State {
} }
struct EmbeddedEntry { struct EmbeddedEntry {
lifetime: u64,
downstreams: HashMap<u64, mpsc::Sender<BrokerSubscriptionEvent>>, downstreams: HashMap<u64, mpsc::Sender<BrokerSubscriptionEvent>>,
snapshot_revision: u64,
snapshot: SubscriptionSnapshot, snapshot: SubscriptionSnapshot,
task: tokio::task::JoinHandle<()>, task: tokio::task::JoinHandle<()>,
} }
enum EmbeddedUpdate {
Event {
selector: EventSubscriptionSelector,
lifetime: u64,
payload: SubscriptionEventPayload,
},
Closed {
selector: EventSubscriptionSelector,
lifetime: u64,
code: SubscriptionTerminationCode,
message: String,
},
}
fn start_embedded_subscription(
runtime: &worker_runtime::Runtime,
runtime_id: &str,
selector: EventSubscriptionSelector,
lifetime: u64,
updates: mpsc::UnboundedSender<EmbeddedUpdate>,
) -> Result<(SubscriptionSnapshot, tokio::task::JoinHandle<()>), worker_runtime::error::RuntimeError>
{
let mut subscription = runtime.subscribe_event_selector(selector.clone())?;
let snapshot = project_snapshot_runtime(subscription.snapshot().clone(), runtime_id);
let task_runtime_id = runtime_id.to_string();
let task = tokio::spawn(async move {
loop {
match subscription.recv().await {
Ok(update) => {
let payload = project_payload_runtime(update.payload, &task_runtime_id);
if updates
.send(EmbeddedUpdate::Event {
selector: selector.clone(),
lifetime,
payload,
})
.is_err()
{
return;
}
}
Err(worker_runtime::RuntimeSubscriptionRecvError::Lagged) => {
let _ = updates.send(EmbeddedUpdate::Closed {
selector,
lifetime,
code: SubscriptionTerminationCode::Lagged,
message: "embedded Runtime subscription lagged; resubscribing from a fresh snapshot".to_string(),
});
return;
}
Err(worker_runtime::RuntimeSubscriptionRecvError::Closed) => {
let _ = updates.send(EmbeddedUpdate::Closed {
selector,
lifetime,
code: SubscriptionTerminationCode::ServerShutdown,
message: "embedded Runtime subscription closed".to_string(),
});
return;
}
}
}
});
Ok((snapshot, task))
}
async fn run_embedded_connection( async fn run_embedded_connection(
runtime_id: String, runtime_id: String,
runtime: worker_runtime::Runtime, runtime: worker_runtime::Runtime,
@@ -364,6 +435,7 @@ async fn run_embedded_connection(
let (updates, mut update_receiver) = mpsc::unbounded_channel(); let (updates, mut update_receiver) = mpsc::unbounded_channel();
let mut entries = HashMap::<EventSubscriptionSelector, EmbeddedEntry>::new(); let mut entries = HashMap::<EventSubscriptionSelector, EmbeddedEntry>::new();
let mut downstream_index = HashMap::<u64, EventSubscriptionSelector>::new(); let mut downstream_index = HashMap::<u64, EventSubscriptionSelector>::new();
let mut next_lifetime = 1_u64;
loop { loop {
tokio::select! { tokio::select! {
command = commands.recv() => match command { command = commands.recv() => match command {
@@ -371,30 +443,38 @@ async fn run_embedded_connection(
downstream_index.insert(downstream_id, selector.clone()); downstream_index.insert(downstream_id, selector.clone());
if let Some(entry) = entries.get_mut(&selector) { if let Some(entry) = entries.get_mut(&selector) {
let _ = events.try_send(BrokerSubscriptionEvent::Snapshot { let _ = events.try_send(BrokerSubscriptionEvent::Snapshot {
connection_generation: generation,
snapshot_revision: entry.snapshot_revision,
snapshot: entry.snapshot.clone(), snapshot: entry.snapshot.clone(),
}); });
entry.downstreams.insert(downstream_id, events); entry.downstreams.insert(downstream_id, events);
} else { } else {
match runtime.subscribe_event_selector(selector.clone()) { let lifetime = next_lifetime;
Ok(mut subscription) => { next_lifetime = next_lifetime.saturating_add(1);
let snapshot_revision = subscription.snapshot_revision(); match start_embedded_subscription(
let snapshot = project_snapshot_runtime(subscription.snapshot().clone(), &runtime_id); &runtime,
let _ = events.try_send(BrokerSubscriptionEvent::Snapshot { connection_generation: generation, snapshot_revision, snapshot: snapshot.clone() }); &runtime_id,
let sender = updates.clone(); selector.clone(),
let task_selector = selector.clone(); lifetime,
let task_runtime_id = runtime_id.clone(); updates.clone(),
let task = tokio::spawn(async move { ) {
while let Ok(update) = subscription.recv().await { Ok((snapshot, task)) => {
let payload = project_payload_runtime(update.payload, &task_runtime_id); let _ = events.try_send(BrokerSubscriptionEvent::Snapshot {
if sender.send((task_selector.clone(), update.subject_revision, payload)).is_err() { break; } snapshot: snapshot.clone(),
}
}); });
entries.insert(selector, EmbeddedEntry { downstreams: HashMap::from([(downstream_id, events)]), snapshot_revision, snapshot, task }); entries.insert(
selector,
EmbeddedEntry {
lifetime,
downstreams: HashMap::from([(downstream_id, events)]),
snapshot,
task,
},
);
} }
Err(error) => { Err(error) => {
let _ = events.try_send(BrokerSubscriptionEvent::Rejected { connection_generation: generation, code: SubscriptionRejectionCode::UnsupportedSelector, message: error.to_string() }); let _ = events.try_send(BrokerSubscriptionEvent::Rejected {
code: SubscriptionRejectionCode::UnsupportedSelector,
message: error.to_string(),
});
} }
} }
} }
@@ -407,7 +487,7 @@ async fn run_embedded_connection(
} }
Some(Command::Shutdown(replacement)) => { Some(Command::Shutdown(replacement)) => {
for entry in entries.values_mut() { 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}") }); broadcast(&mut entry.downstreams, BrokerSubscriptionEvent::Closed { code: SubscriptionTerminationCode::ServerShutdown, message: format!("embedded Runtime registration {generation} was replaced by {replacement}") });
entry.task.abort(); entry.task.abort();
} }
return; return;
@@ -415,16 +495,68 @@ async fn run_embedded_connection(
None => return, None => return,
}, },
update = update_receiver.recv() => { update = update_receiver.recv() => {
let Some((selector, subject_revision, payload)) = update else { return; }; let Some(update) = update else { return; };
if let Some(entry) = entries.get_mut(&selector) { match update {
entry.snapshot_revision = entry.snapshot_revision.saturating_add(1); EmbeddedUpdate::Event { selector, lifetime, payload } => {
if let Some(entry) = entries
.get_mut(&selector)
.filter(|entry| entry.lifetime == lifetime)
{
apply_event_to_cached_snapshot(&mut entry.snapshot, &payload); apply_event_to_cached_snapshot(&mut entry.snapshot, &payload);
broadcast(&mut entry.downstreams, BrokerSubscriptionEvent::Event { connection_generation: generation, subject_revision, payload }); broadcast(&mut entry.downstreams, BrokerSubscriptionEvent::Event { payload });
}
}
EmbeddedUpdate::Closed { selector, lifetime, code, message } => {
if !entries
.get(&selector)
.is_some_and(|entry| entry.lifetime == lifetime)
{
continue;
}
let Some(mut entry) = entries.remove(&selector) else { continue; };
broadcast(
&mut entry.downstreams,
BrokerSubscriptionEvent::Closed { code, message },
);
if entry.downstreams.is_empty() {
continue;
}
let new_lifetime = next_lifetime;
next_lifetime = next_lifetime.saturating_add(1);
match start_embedded_subscription(
&runtime,
&runtime_id,
selector.clone(),
new_lifetime,
updates.clone(),
) {
Ok((snapshot, task)) => {
broadcast(
&mut entry.downstreams,
BrokerSubscriptionEvent::Snapshot {
snapshot: snapshot.clone(),
},
);
entry.snapshot = snapshot;
entry.task = task;
entry.lifetime = new_lifetime;
entries.insert(selector, entry);
}
Err(error) => {
broadcast(
&mut entry.downstreams,
BrokerSubscriptionEvent::Rejected {
code: SubscriptionRejectionCode::Internal,
message: error.to_string(),
},
);
}
}
}
} }
} }
} }
*status.write().expect("broker status poisoned") = RuntimeSubscriptionBrokerStatus { *status.write().expect("broker status poisoned") = RuntimeSubscriptionBrokerStatus {
connection_generation: generation,
connected: true, connected: true,
desired_selectors: entries.len(), desired_selectors: entries.len(),
upstream_subscriptions: entries.len(), upstream_subscriptions: entries.len(),
@@ -597,10 +729,8 @@ async fn apply_online(
.selectors .selectors
.entry(selector.clone()) .entry(selector.clone())
.or_insert_with(SelectorState::new); .or_insert_with(SelectorState::new);
if let Some((revision, snapshot)) = &entry.snapshot { if let Some(snapshot) = &entry.snapshot {
let _ = events.try_send(BrokerSubscriptionEvent::Snapshot { let _ = events.try_send(BrokerSubscriptionEvent::Snapshot {
connection_generation: state.generation,
snapshot_revision: *revision,
snapshot: snapshot.clone(), snapshot: snapshot.clone(),
}); });
} }
@@ -696,7 +826,6 @@ async fn handle_frame(
request_id, request_id,
subscription_id, subscription_id,
selector, selector,
snapshot_revision,
snapshot, snapshot,
}) => { }) => {
if state.pending.remove(&request_id) != Some(selector.clone()) { if state.pending.remove(&request_id) != Some(selector.clone()) {
@@ -706,18 +835,13 @@ async fn handle_frame(
let entry = state.selectors.get_mut(&selector).ok_or(())?; let entry = state.selectors.get_mut(&selector).ok_or(())?;
entry.pending = false; entry.pending = false;
entry.upstream_id = Some(subscription_id.clone()); entry.upstream_id = Some(subscription_id.clone());
entry.snapshot = Some((snapshot_revision, snapshot.clone())); entry.snapshot = Some(snapshot.clone());
entry.revisions = snapshot_revisions(&snapshot);
state state
.upstream_index .upstream_index
.insert(subscription_id, selector.clone()); .insert(subscription_id, selector.clone());
broadcast( broadcast(
&mut entry.downstreams, &mut entry.downstreams,
BrokerSubscriptionEvent::Snapshot { BrokerSubscriptionEvent::Snapshot { snapshot },
connection_generation: state.generation,
snapshot_revision,
snapshot,
},
); );
if entry.downstreams.is_empty() { if entry.downstreams.is_empty() {
maybe_unsubscribe(socket, state, selector).await?; maybe_unsubscribe(socket, state, selector).await?;
@@ -734,18 +858,13 @@ async fn handle_frame(
if let Some(mut entry) = state.selectors.remove(&selector) { if let Some(mut entry) = state.selectors.remove(&selector) {
broadcast( broadcast(
&mut entry.downstreams, &mut entry.downstreams,
BrokerSubscriptionEvent::Rejected { BrokerSubscriptionEvent::Rejected { code, message },
connection_generation: state.generation,
code,
message,
},
); );
} }
} }
} }
SubscriptionFramePayload::Event(SubscriptionEvent::Event { SubscriptionFramePayload::Event(SubscriptionEvent::Event {
subscription_id, subscription_id,
subject_revision,
payload, payload,
}) => { }) => {
let selector = state let selector = state
@@ -756,24 +875,12 @@ async fn handle_frame(
payload.validate_for_selector(&selector).map_err(|_| ())?; payload.validate_for_selector(&selector).map_err(|_| ())?;
let payload = project_payload_runtime(payload, &state.runtime_id); let payload = project_payload_runtime(payload, &state.runtime_id);
let entry = state.selectors.get_mut(&selector).ok_or(())?; let entry = state.selectors.get_mut(&selector).ok_or(())?;
if let Some(subject) = event_subject(&payload) { if let Some(snapshot) = entry.snapshot.as_mut() {
let revision = entry.revisions.entry(subject).or_insert(0);
if subject_revision <= *revision {
return Ok(());
}
*revision = subject_revision;
}
if let Some((snapshot_revision, snapshot)) = entry.snapshot.as_mut() {
*snapshot_revision = snapshot_revision.saturating_add(1);
apply_event_to_cached_snapshot(snapshot, &payload); apply_event_to_cached_snapshot(snapshot, &payload);
} }
broadcast( broadcast(
&mut entry.downstreams, &mut entry.downstreams,
BrokerSubscriptionEvent::Event { BrokerSubscriptionEvent::Event { payload },
connection_generation: state.generation,
subject_revision,
payload,
},
); );
} }
SubscriptionFramePayload::Event(SubscriptionEvent::SubscriptionClosed { SubscriptionFramePayload::Event(SubscriptionEvent::SubscriptionClosed {
@@ -785,14 +892,9 @@ async fn handle_frame(
let should_resubscribe = if let Some(entry) = state.selectors.get_mut(&selector) { let should_resubscribe = if let Some(entry) = state.selectors.get_mut(&selector) {
entry.upstream_id = None; entry.upstream_id = None;
entry.snapshot = None; entry.snapshot = None;
entry.revisions.clear();
broadcast( broadcast(
&mut entry.downstreams, &mut entry.downstreams,
BrokerSubscriptionEvent::Closed { BrokerSubscriptionEvent::Closed { code, message },
connection_generation: state.generation,
code,
message,
},
); );
!entry.downstreams.is_empty() !entry.downstreams.is_empty()
} else { } else {
@@ -822,10 +924,9 @@ fn close_all(state: &mut State, replacement: u64) {
broadcast( broadcast(
&mut entry.downstreams, &mut entry.downstreams,
BrokerSubscriptionEvent::Closed { BrokerSubscriptionEvent::Closed {
connection_generation: state.generation,
code: SubscriptionTerminationCode::ServerShutdown, code: SubscriptionTerminationCode::ServerShutdown,
message: format!( message: format!(
"Runtime connection generation {} was fenced by generation {replacement}", "Runtime registration {} was replaced by {replacement}",
state.generation state.generation
), ),
}, },
@@ -846,51 +947,6 @@ fn broadcast(
downstreams.retain(|id, _| !closed.contains(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<(), ()> { async fn send_frame(socket: &mut RuntimeSocket, frame: SubscriptionFrame) -> Result<(), ()> {
frame.validate().map_err(|_| ())?; frame.validate().map_err(|_| ())?;
socket socket
@@ -979,7 +1035,6 @@ fn runtime_token(
} }
fn update_status(status: &RwLock<RuntimeSubscriptionBrokerStatus>, state: &State, connected: bool) { fn update_status(status: &RwLock<RuntimeSubscriptionBrokerStatus>, state: &State, connected: bool) {
*status.write().expect("broker status poisoned") = RuntimeSubscriptionBrokerStatus { *status.write().expect("broker status poisoned") = RuntimeSubscriptionBrokerStatus {
connection_generation: state.generation,
connected, connected,
desired_selectors: state desired_selectors: state
.selectors .selectors
@@ -264,16 +264,9 @@ async fn replacing_runtime_registration_fences_the_old_generation() {
assert!(second_generation > first_generation); assert!(second_generation > first_generation);
assert!(matches!( assert!(matches!(
next_event(&mut old).await, next_event(&mut old).await,
BrokerSubscriptionEvent::Closed { BrokerSubscriptionEvent::Closed { .. }
connection_generation,
..
} if connection_generation == first_generation
)); ));
let status = wait_for_status(&broker, "runtime-test", |status| { wait_for_status(&broker, "runtime-test", |status| status.connected).await;
status.connection_generation == second_generation && status.connected
})
.await;
assert_eq!(status.connection_generation, second_generation);
server.abort(); server.abort();
} }
@@ -430,6 +423,75 @@ async fn embedded_runtime_uses_in_process_subscription_source() {
server.abort(); server.abort();
} }
#[tokio::test]
async fn embedded_downstream_overflow_recovers_through_a_fresh_snapshot() {
let (runtime, _config, server) = fixture().await;
let worker = runtime.list_workers().unwrap().remove(0);
let broker = RuntimeSubscriptionBroker::new("local");
broker.register_embedded_runtime("embedded-worker-runtime", runtime.clone());
let mut subscription = broker
.subscribe(
"embedded-worker-runtime",
EventSubscriptionSelector::RuntimeWorkers,
)
.unwrap();
assert!(matches!(
next_snapshot(&mut subscription).await,
BrokerSubscriptionEvent::Snapshot { .. }
));
let final_command_id = (DOWNSTREAM_QUEUE_CAPACITY as u64) * 3;
for command_id in 1..=final_command_id {
let mut snapshot = worker.worker_state.clone().unwrap();
snapshot.last_command_id = command_id;
runtime
.observe_worker_event(
&worker.worker_ref,
protocol::Event::WorkerState { snapshot },
)
.unwrap();
}
let mut saw_recovery_snapshot = false;
let mut converged = false;
tokio::time::timeout(Duration::from_secs(5), async {
while !converged {
match subscription
.recv()
.await
.expect("broker subscription closed")
{
BrokerSubscriptionEvent::Snapshot {
snapshot: SubscriptionSnapshot::Workers { workers },
} => {
saw_recovery_snapshot = true;
converged = workers.iter().any(|candidate| {
candidate.worker_id.as_str() == worker.worker_id.to_string()
&& candidate
.worker_state
.as_ref()
.is_some_and(|state| state.last_command_id == final_command_id)
});
}
BrokerSubscriptionEvent::Event {
payload: SubscriptionEventPayload::WorkerUpserted { worker: candidate },
} if saw_recovery_snapshot => {
converged = candidate.worker_id.as_str() == worker.worker_id.to_string()
&& candidate
.worker_state
.as_ref()
.is_some_and(|state| state.last_command_id == final_command_id);
}
_ => {}
}
}
})
.await
.expect("overflow recovery did not converge");
assert!(saw_recovery_snapshot);
server.abort();
}
#[tokio::test] #[tokio::test]
async fn strict_runtime_subscription_rejects_private_endpoint_before_websocket_connect() { async fn strict_runtime_subscription_rejects_private_endpoint_before_websocket_connect() {
let config = RemoteRuntimeConfig::new( let config = RemoteRuntimeConfig::new(
+67 -73
View File
@@ -1501,10 +1501,11 @@ impl WorkerRemovalService {
/// and WorkerRemove use the retention commit below; compensation preserves /// and WorkerRemove use the retention commit below; compensation preserves
/// its pre-create semantics while sharing the atomic catalog contract. /// its pre-create semantics while sharing the atomic catalog contract.
fn commit_catalog_removal(&self, target: &RuntimeWorkerRef) -> crate::Result<()> { fn commit_catalog_removal(&self, target: &RuntimeWorkerRef) -> crate::Result<()> {
let commit = self self.worker_projection.publish_ordered(|| {
.store self.store
.delete_worker_registry(&self.workspace_id, target)?; .delete_worker_registry(&self.workspace_id, target)
self.worker_projection.publish_commit(commit) .map(|commit| ((), commit))
})
} }
async fn execute_cleanup_removal(&self, candidate: &CleanupWorkerCandidate) -> ApiResult<()> { async fn execute_cleanup_removal(&self, candidate: &CleanupWorkerCandidate) -> ApiResult<()> {
@@ -1631,18 +1632,17 @@ impl WorkerRemovalService {
)); ));
} }
}; };
match self.store.commit_worker_removal( match self.worker_projection.publish_ordered(|| {
self.store
.commit_worker_removal(
&self.workspace_id, &self.workspace_id,
&prepared.plan.operation_id, &prepared.plan.operation_id,
&prepared.plan.input_fingerprint, &prepared.plan.input_fingerprint,
&result, &result,
) { )
Ok(commit) => { .map(|commit| ((), commit.catalog))
self.worker_projection }) {
.publish_commit(commit.catalog) Ok(()) => Ok(worker_remove_success_response(target)),
.map_err(|error| error.to_string())?;
Ok(worker_remove_success_response(target))
}
Err(error) => Ok(worker_retention_error_response(error)), Err(error) => Ok(worker_retention_error_response(error)),
} }
} }
@@ -1764,12 +1764,12 @@ impl WorkerRemovalService {
.map_err(|_| "Worker removal recovery authority is unavailable".to_string())?; .map_err(|_| "Worker removal recovery authority is unavailable".to_string())?;
if let Some(prepared) = prepared { if let Some(prepared) = prepared {
if prepared.plan.state == crate::retention::WorkerRemovalPlanState::Succeeded { if prepared.plan.state == crate::retention::WorkerRemovalPlanState::Succeeded {
let commit = self
.store
.recover_succeeded_worker_removal_catalog(&self.workspace_id, target)
.map_err(|error| error.to_string())?;
self.worker_projection self.worker_projection
.publish_commit(commit) .publish_ordered(|| {
self.store
.recover_succeeded_worker_removal_catalog(&self.workspace_id, target)
.map(|commit| ((), commit))
})
.map_err(|error| error.to_string())?; .map_err(|error| error.to_string())?;
return Ok(worker_remove_success_response(target)); return Ok(worker_remove_success_response(target));
} }
@@ -1947,17 +1947,17 @@ impl WorkerRemovalService {
)); ));
} }
}; };
match self.store.commit_worker_removal( match self.worker_projection.publish_ordered(|| {
self.store
.commit_worker_removal(
&self.workspace_id, &self.workspace_id,
&plan.operation_id, &plan.operation_id,
&plan.input_fingerprint, &plan.input_fingerprint,
&retention_result, &retention_result,
) { )
Ok(commit) => { .map(|commit| ((), commit.catalog))
self.worker_projection }) {
.publish_commit(commit.catalog) Ok(()) => {}
.map_err(|error| error.to_string())?;
}
Err(error) => { Err(error) => {
let _ = self.store.fail_worker_removal( let _ = self.store.fail_worker_removal(
&self.workspace_id, &self.workspace_id,
@@ -3454,17 +3454,19 @@ impl WorkspaceApi {
}; };
let bound_at = Utc::now().to_rfc3339_opts(SecondsFormat::Nanos, true); let bound_at = Utc::now().to_rfc3339_opts(SecondsFormat::Nanos, true);
let runtime_run_id = crate::backend_job::runtime_run_id(&reservation.attempt.attempt_id); let runtime_run_id = crate::backend_job::runtime_run_id(&reservation.attempt.attempt_id);
let (commit, newly_bound) = self.store.bind_backend_job_attempt_worker( let newly_bound = self.worker_projection.publish_ordered(|| {
self.store
.bind_backend_job_attempt_worker(
&self.config.workspace_id, &self.config.workspace_id,
&request.job_id, &request.job_id,
&reservation.attempt.attempt_id, &reservation.attempt.attempt_id,
&worker.worker, &worker.worker,
Some(&runtime_run_id), Some(&runtime_run_id),
&bound_at, &bound_at,
)?; )
if newly_bound { .map(|(commit, newly_bound)| (newly_bound, commit))
self.worker_projection.publish_commit(commit)?; })?;
} else { if !newly_bound {
let job = self let job = self
.store .store
.get_backend_job(&self.config.workspace_id, &request.job_id)? .get_backend_job(&self.config.workspace_id, &request.job_id)?
@@ -19435,10 +19437,11 @@ async fn scoped_start_workspace_orchestrator(
deleted.diagnostics, deleted.diagnostics,
)); ));
} }
let removal = api api.worker_projection.publish_ordered(|| {
.store api.store
.delete_worker_registry(&api.config.workspace_id, &existing.worker)?; .delete_worker_registry(&api.config.workspace_id, &existing.worker)
api.worker_projection.publish_commit(removal)?; .map(|commit| ((), commit))
})?;
disposition = "recreated"; disposition = "recreated";
} }
@@ -21768,22 +21771,23 @@ async fn set_worker_retention(
} }
} }
let retention_state = if pinned { "pinned" } else { "normal" }; let retention_state = if pinned { "pinned" } else { "normal" };
let commit = api.store.update_worker_retention( let changed = api.worker_projection.publish_ordered(|| {
api.store
.update_worker_retention(
&api.config.workspace_id, &api.config.workspace_id,
&worker_ref, &worker_ref,
retention_state, retention_state,
now_registry_timestamp().as_str(), now_registry_timestamp().as_str(),
)?; )
if commit.changes.is_empty() { .map(|commit| (!commit.changes.is_empty(), commit))
})?;
if !changed {
return Err(cleanup_api_error( return Err(cleanup_api_error(
runtime_id.as_str(), runtime_id.as_str(),
"workspace_worker_retention_unknown_worker", "workspace_worker_retention_unknown_worker",
"Worker is not known to the Backend registry", "Worker is not known to the Backend registry",
)); ));
} }
api.worker_projection
.publish_commit(commit)
.map_err(ApiError::from)?;
Ok(Json(WorkerRetentionResponse { Ok(Json(WorkerRetentionResponse {
workspace_id: api.config.workspace_id, workspace_id: api.config.workspace_id,
runtime_id: worker_ref.runtime_id, runtime_id: worker_ref.runtime_id,
@@ -26680,7 +26684,7 @@ fn project_observed_workspace_workers(
fn workers_response(api: WorkspaceApi) -> ApiResult<server_api::WorkerListResponse> { fn workers_response(api: WorkspaceApi) -> ApiResult<server_api::WorkerListResponse> {
let limit = api.config.max_records.min(200); let limit = api.config.max_records.min(200);
let (_, projections) = api let projections = api
.store .store
.worker_registry_projection_snapshot(&api.config.workspace_id, limit)?; .worker_registry_projection_snapshot(&api.config.workspace_id, limit)?;
let workdir_records = api let workdir_records = api
@@ -27478,9 +27482,12 @@ fn record_worker_summary(
created_at: timestamp.clone(), created_at: timestamp.clone(),
updated_at: timestamp, updated_at: timestamp,
}; };
let catalog_commit = api.store.upsert_worker_registry(&record)?;
api.worker_projection api.worker_projection
.publish_commit(catalog_commit) .publish_ordered(|| {
api.store
.upsert_worker_registry(&record)
.map(|commit| ((), commit))
})
.map_err(ApiError::from)?; .map_err(ApiError::from)?;
if let Ok(worker_id) = if let Ok(worker_id) =
protocol::subscription::SubscriptionWorkerId::new(worker.worker.worker_id.clone()) protocol::subscription::SubscriptionWorkerId::new(worker.worker.worker_id.clone())
@@ -27496,7 +27503,6 @@ fn record_worker_summary(
runtime_id: Some(worker.worker.runtime_id.clone()), runtime_id: Some(worker.worker.runtime_id.clone()),
resource_key: None, resource_key: None,
availability: protocol::subscription::SubscriptionWorkerAvailability::Observed, availability: protocol::subscription::SubscriptionWorkerAvailability::Observed,
subject_revision: 0,
worker_state: worker.worker_state.clone(), worker_state: worker.worker_state.clone(),
state, state,
has_running_internal_workers: false, has_running_internal_workers: false,
@@ -29427,7 +29433,6 @@ mod tests {
runtime_id: Some("runtime-a".to_string()), runtime_id: Some("runtime-a".to_string()),
resource_key: Some("W-1".to_string()), resource_key: Some("W-1".to_string()),
availability, availability,
subject_revision: 4,
worker_state: Some(protocol::WorkerStateSnapshot::initial()), worker_state: Some(protocol::WorkerStateSnapshot::initial()),
state: protocol::subscription::SubscriptionWorkerState::Idle, state: protocol::subscription::SubscriptionWorkerState::Idle,
has_running_internal_workers: false, has_running_internal_workers: false,
@@ -29438,10 +29443,6 @@ mod tests {
workdir_attachments: Vec::new(), workdir_attachments: Vec::new(),
}, },
availability, availability,
connection_generation: 2,
subject_revision: 4,
snapshot_revision: 7,
projection_revision: 9,
observed_at: "2026-09-24T00:00:01Z".to_string(), observed_at: "2026-09-24T00:00:01Z".to_string(),
}), }),
job: None, job: None,
@@ -33018,7 +33019,6 @@ mod tests {
runtime_id: Some(EMBEDDED_WORKER_RUNTIME_ID.to_string()), runtime_id: Some(EMBEDDED_WORKER_RUNTIME_ID.to_string()),
resource_key: None, resource_key: None,
availability: protocol::subscription::SubscriptionWorkerAvailability::Observed, availability: protocol::subscription::SubscriptionWorkerAvailability::Observed,
subject_revision: 1,
worker_state: None, worker_state: None,
state: protocol::subscription::SubscriptionWorkerState::Paused, state: protocol::subscription::SubscriptionWorkerState::Paused,
has_running_internal_workers: false, has_running_internal_workers: false,
@@ -44053,7 +44053,8 @@ mod tests {
let summary = api.runtime.worker(&target).unwrap(); let summary = api.runtime.worker(&target).unwrap();
sync_worker_observation(&api, &summary).unwrap(); sync_worker_observation(&api, &summary).unwrap();
seed_worker_control_grant(&api, &source, &target, "embedded-valid-proof"); seed_worker_control_grant(&api, &source, &target, "embedded-valid-proof");
let mut subscriber = api.worker_projection.subscribe(); let mut subscriber = api.worker_projection.subscribe_ordered(10).unwrap();
let _ = subscriber.take_snapshot();
let response = WorkerRemovalService::new(&api) let response = WorkerRemovalService::new(&api)
.execute_async( .execute_async(
@@ -44086,7 +44087,7 @@ mod tests {
.unwrap() .unwrap()
.is_none() .is_none()
); );
let removal = tokio::time::timeout(std::time::Duration::from_secs(1), async { tokio::time::timeout(std::time::Duration::from_secs(1), async {
loop { loop {
let event = subscriber.recv().await.expect("projection event"); let event = subscriber.recv().await.expect("projection event");
if event.changes.iter().any(|change| { if event.changes.iter().any(|change| {
@@ -44098,7 +44099,6 @@ mod tests {
}) })
.await .await
.expect("WorkerRemove catalog removal broadcast"); .expect("WorkerRemove catalog removal broadcast");
assert!(removal.revision > 0);
} }
#[tokio::test] #[tokio::test]
@@ -44135,7 +44135,8 @@ mod tests {
let worker = spawned.worker.unwrap().worker; let worker = spawned.worker.unwrap().worker;
let summary = api.runtime.worker(&worker).unwrap(); let summary = api.runtime.worker(&worker).unwrap();
sync_worker_observation(&api, &summary).unwrap(); sync_worker_observation(&api, &summary).unwrap();
let mut subscriber = api.worker_projection.subscribe(); let mut subscriber = api.worker_projection.subscribe_ordered(10).unwrap();
let _ = subscriber.take_snapshot();
let diagnostics = compensate_failed_worker_spawn( let diagnostics = compensate_failed_worker_spawn(
&api, &api,
&summary, &summary,
@@ -44151,7 +44152,7 @@ mod tests {
.unwrap() .unwrap()
.is_none() .is_none()
); );
let removal = tokio::time::timeout(std::time::Duration::from_secs(1), async { tokio::time::timeout(std::time::Duration::from_secs(1), async {
loop { loop {
let event = subscriber.recv().await.expect("projection event"); let event = subscriber.recv().await.expect("projection event");
if event.changes.iter().any(|change| { if event.changes.iter().any(|change| {
@@ -44163,7 +44164,6 @@ mod tests {
}) })
.await .await
.expect("spawn compensation catalog removal broadcast"); .expect("spawn compensation catalog removal broadcast");
assert!(removal.revision > 0);
sync_worker_observation(&api, &summary).unwrap(); sync_worker_observation(&api, &summary).unwrap();
assert!( assert!(
api.store api.store
@@ -46005,7 +46005,8 @@ mod tests {
.expect("singleton Worker cleanup candidate") .expect("singleton Worker cleanup candidate")
.clone(); .clone();
assert_eq!(candidate.blocking_reason, None); assert_eq!(candidate.blocking_reason, None);
let mut subscriber = api.worker_projection.subscribe(); let mut subscriber = api.worker_projection.subscribe_ordered(10).unwrap();
let _ = subscriber.take_snapshot();
WorkerRemovalService::new(&api) WorkerRemovalService::new(&api)
.execute_cleanup_removal(&candidate) .execute_cleanup_removal(&candidate)
@@ -46129,7 +46130,8 @@ mod tests {
.iter() .iter()
.any(|item| item.worker_id.as_str() == worker_id) .any(|item| item.worker_id.as_str() == worker_id)
); );
let mut subscriber = api.worker_projection.subscribe(); let mut subscriber = api.worker_projection.subscribe_ordered(10).unwrap();
let _ = subscriber.take_snapshot();
let plan = build_runtime_cleanup_plan(&api, EMBEDDED_WORKER_RUNTIME_ID) let plan = build_runtime_cleanup_plan(&api, EMBEDDED_WORKER_RUNTIME_ID)
.unwrap_or_else(|error| panic!("cleanup plan: {}", error.error)); .unwrap_or_else(|error| panic!("cleanup plan: {}", error.error));
let candidate = plan let candidate = plan
@@ -46173,20 +46175,18 @@ mod tests {
event.changes.as_slice(), event.changes.as_slice(),
[crate::store::WorkerCatalogChange::Removed(removed)] if removed == &worker [crate::store::WorkerCatalogChange::Removed(removed)] if removed == &worker
)); ));
let acting_revision = tokio::time::timeout( tokio::time::timeout(
std::time::Duration::from_secs(1), std::time::Duration::from_secs(1),
next_workspace_worker_removal(&mut acting_socket, worker_id.as_str()), next_workspace_worker_removal(&mut acting_socket, worker_id.as_str()),
) )
.await .await
.expect("acting Workspace Worker subscription removal"); .expect("acting Workspace Worker subscription removal");
let observing_revision = tokio::time::timeout( tokio::time::timeout(
std::time::Duration::from_secs(1), std::time::Duration::from_secs(1),
next_workspace_worker_removal(&mut observing_socket, worker_id.as_str()), next_workspace_worker_removal(&mut observing_socket, worker_id.as_str()),
) )
.await .await
.expect("independent Workspace Worker subscription removal"); .expect("independent Workspace Worker subscription removal");
assert_eq!(acting_revision, event.revision);
assert_eq!(observing_revision, event.revision);
assert!( assert!(
api.store api.store
.get_worker_registry(&api.config.workspace_id, &worker) .get_worker_registry(&api.config.workspace_id, &worker)
@@ -46202,11 +46202,10 @@ mod tests {
recovered.plan.state, recovered.plan.state,
crate::retention::WorkerRemovalPlanState::Succeeded crate::retention::WorkerRemovalPlanState::Succeeded
); );
let (snapshot_revision, snapshot) = api let snapshot = api
.store .store
.worker_registry_projection_snapshot(&api.config.workspace_id, 100) .worker_registry_projection_snapshot(&api.config.workspace_id, 100)
.unwrap(); .unwrap();
assert_eq!(snapshot_revision, event.revision);
assert!( assert!(
snapshot snapshot
.iter() .iter()
@@ -46228,8 +46227,8 @@ mod tests {
.all(|item| item.worker_id.as_str() != worker_id) .all(|item| item.worker_id.as_str() != worker_id)
); );
// A retry replays the durable fence at the same revision, healing an // A retry replays the durable fence, healing an in-process broadcast
// in-process broadcast omitted after the first DB commit. // omitted after the first DB commit.
WorkerRemovalService::new(&api) WorkerRemovalService::new(&api)
.commit_catalog_removal(&worker) .commit_catalog_removal(&worker)
.unwrap(); .unwrap();
@@ -46237,7 +46236,6 @@ mod tests {
.await .await
.expect("catalog removal replay") .expect("catalog removal replay")
.expect("projection replay"); .expect("projection replay");
assert_eq!(replay.revision, event.revision);
assert!(matches!( assert!(matches!(
replay.changes.as_slice(), replay.changes.as_slice(),
[crate::store::WorkerCatalogChange::Removed(removed)] if removed == &worker [crate::store::WorkerCatalogChange::Removed(removed)] if removed == &worker
@@ -49669,10 +49667,7 @@ mod tests {
} }
} }
async fn next_workspace_worker_removal( async fn next_workspace_worker_removal(socket: &mut TestWebSocket, expected_worker_id: &str) {
socket: &mut TestWebSocket,
expected_worker_id: &str,
) -> u64 {
loop { loop {
let Message::Text(text) = socket.next().await.unwrap().unwrap() else { let Message::Text(text) = socket.next().await.unwrap().unwrap() else {
continue; continue;
@@ -49681,7 +49676,6 @@ mod tests {
serde_json::from_str(text.as_str()).unwrap(); serde_json::from_str(text.as_str()).unwrap();
if let protocol::subscription::SubscriptionFramePayload::Event( if let protocol::subscription::SubscriptionFramePayload::Event(
protocol::subscription::SubscriptionEvent::Event { protocol::subscription::SubscriptionEvent::Event {
subject_revision,
payload: payload:
protocol::subscription::SubscriptionEventPayload::WorkerRemoved { protocol::subscription::SubscriptionEventPayload::WorkerRemoved {
worker_id, worker_id,
@@ -49692,7 +49686,7 @@ mod tests {
) = frame.payload ) = frame.payload
&& worker_id.as_str() == expected_worker_id && worker_id.as_str() == expected_worker_id
{ {
return subject_revision; return;
} }
} }
} }
@@ -49812,7 +49806,7 @@ mod tests {
.unwrap(); .unwrap();
assert_eq!(spawned.state, InternalWorkerOperationState::Accepted); assert_eq!(spawned.state, InternalWorkerOperationState::Accepted);
let worker_id = spawned.worker.unwrap().worker.worker_id; let worker_id = spawned.worker.unwrap().worker.worker_id;
let (_, projections) = api let projections = api
.store .store
.worker_registry_projection_snapshot(TEST_WORKSPACE_ID, 10) .worker_registry_projection_snapshot(TEST_WORKSPACE_ID, 10)
.unwrap(); .unwrap();
File diff suppressed because it is too large Load Diff
+283 -75
View File
@@ -11,7 +11,10 @@ use tokio::task::JoinHandle;
use crate::runtime_subscription::{ use crate::runtime_subscription::{
BrokerSubscriptionEvent, RuntimeSubscriptionBroker, RuntimeSubscriptionBrokerError, BrokerSubscriptionEvent, RuntimeSubscriptionBroker, RuntimeSubscriptionBrokerError,
}; };
use crate::store::{ControlPlaneStore, WorkerCatalogChange, WorkerRegistryProjectionCommit}; use crate::store::{
ControlPlaneStore, WorkerCatalogChange, WorkerRegistryProjectionCommit,
WorkerRegistryProjectionRecord,
};
use worker_runtime::identity::RuntimeWorkerRef; use worker_runtime::identity::RuntimeWorkerRef;
const PROJECTION_EVENT_CAPACITY: usize = 256; const PROJECTION_EVENT_CAPACITY: usize = 256;
@@ -20,21 +23,42 @@ pub type WorkerProjectionChange = WorkerCatalogChange;
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct WorkerProjectionEvent { pub struct WorkerProjectionEvent {
pub revision: u64,
pub changes: Vec<WorkerProjectionChange>, pub changes: Vec<WorkerProjectionChange>,
} }
/// An authoritative durable snapshot followed by changes committed after that snapshot.
pub struct WorkerProjectionSubscription {
snapshot: Vec<WorkerRegistryProjectionRecord>,
receiver: broadcast::Receiver<WorkerProjectionEvent>,
}
impl WorkerProjectionSubscription {
pub fn take_snapshot(&mut self) -> Vec<WorkerRegistryProjectionRecord> {
std::mem::take(&mut self.snapshot)
}
pub async fn recv(&mut self) -> Result<WorkerProjectionEvent, broadcast::error::RecvError> {
self.receiver.recv().await
}
#[cfg(test)]
pub fn try_recv(&mut self) -> Result<WorkerProjectionEvent, broadcast::error::TryRecvError> {
self.receiver.try_recv()
}
}
/// Owns the durable Runtime-to-Workspace Worker projection and its independent /// Owns the durable Runtime-to-Workspace Worker projection and its independent
/// low-frequency publication stream. /// low-frequency publication stream.
/// ///
/// One Runtime subscription is kept alive per registered Runtime even when no /// The ordering gate covers both durable commits and subscription snapshot creation.
/// browser is connected. Browser subscriptions therefore never become the /// Every subscription therefore starts with one authoritative snapshot and receives
/// authority for populating `worker_registry` observations. /// only commits that completed after it, in the same order in which they were stored.
pub struct WorkerProjectionService { pub struct WorkerProjectionService {
workspace_id: String, workspace_id: String,
store: Arc<dyn ControlPlaneStore>, store: Arc<dyn ControlPlaneStore>,
events: broadcast::Sender<WorkerProjectionEvent>, events: broadcast::Sender<WorkerProjectionEvent>,
connection_generation_epoch: u64, ordering: Mutex<()>,
attach_gate: tokio::sync::Mutex<()>,
tasks: Mutex<HashMap<String, JoinHandle<()>>>, tasks: Mutex<HashMap<String, JoinHandle<()>>>,
} }
@@ -45,8 +69,8 @@ impl WorkerProjectionService {
workspace_id: workspace_id.into(), workspace_id: workspace_id.into(),
store, store,
events, events,
connection_generation_epoch: Utc::now().timestamp_nanos_opt().unwrap_or_default().max(0) ordering: Mutex::new(()),
as u64, attach_gate: tokio::sync::Mutex::new(()),
tasks: Mutex::new(HashMap::new()), tasks: Mutex::new(HashMap::new()),
}) })
} }
@@ -62,8 +86,38 @@ impl WorkerProjectionService {
} }
} }
pub fn subscribe(&self) -> broadcast::Receiver<WorkerProjectionEvent> { pub fn subscribe_ordered(&self, limit: usize) -> crate::Result<WorkerProjectionSubscription> {
self.events.subscribe() let _ordering = self
.ordering
.lock()
.expect("Worker projection ordering lock poisoned");
let receiver = self.events.subscribe();
let snapshot_limit = limit.saturating_add(1);
let mut snapshot = self
.store
.worker_registry_projection_snapshot(self.workspace_id.as_str(), snapshot_limit)?;
if snapshot.len() > limit {
return Err(crate::Error::Store(format!(
"authoritative Worker projection exceeds the configured snapshot limit of {limit} records"
)));
}
snapshot.shrink_to_fit();
Ok(WorkerProjectionSubscription { snapshot, receiver })
}
/// Serialize one durable projection mutation with its publication. The closure
/// must return the exact projection payload captured by the same transaction.
pub(crate) fn publish_ordered<T, E>(
&self,
mutation: impl FnOnce() -> std::result::Result<(T, WorkerRegistryProjectionCommit), E>,
) -> std::result::Result<T, E> {
let _ordering = self
.ordering
.lock()
.expect("Worker projection ordering lock poisoned");
let (result, commit) = mutation()?;
self.publish_commit(commit);
Ok(result)
} }
pub async fn attach_runtime( pub async fn attach_runtime(
@@ -71,6 +125,20 @@ impl WorkerProjectionService {
broker: &RuntimeSubscriptionBroker, broker: &RuntimeSubscriptionBroker,
runtime_id: &str, runtime_id: &str,
) -> Result<(), RuntimeSubscriptionBrokerError> { ) -> Result<(), RuntimeSubscriptionBrokerError> {
// Serialize replacement and wait for the old consumer to stop before the new
// lifetime can apply its snapshot. Late notifications can therefore never
// cross a Runtime ownership boundary.
let _attach = self.attach_gate.lock().await;
let previous = self
.tasks
.lock()
.expect("Worker projection task lock poisoned")
.remove(runtime_id);
if let Some(previous) = previous {
previous.abort();
let _ = previous.await;
}
let mut subscription = let mut subscription =
broker.subscribe(runtime_id, EventSubscriptionSelector::RuntimeWorkers)?; broker.subscribe(runtime_id, EventSubscriptionSelector::RuntimeWorkers)?;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2); let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
@@ -80,12 +148,14 @@ impl WorkerProjectionService {
break; break;
}; };
let initialized = matches!(event, BrokerSubscriptionEvent::Snapshot { .. }); let initialized = matches!(event, BrokerSubscriptionEvent::Snapshot { .. });
if let Err(error) = self.apply_broker_event(runtime_id, event).await { if let Err(error) = self.apply_broker_event(runtime_id, event) {
tracing::warn!( tracing::warn!(
runtime_id, runtime_id,
error = %error, error = %error,
"failed to initialize durable Worker projection" "failed to initialize durable Worker projection; restarting from a fresh snapshot"
); );
subscription.restart()?;
continue;
} }
if initialized { if initialized {
break; break;
@@ -99,26 +169,22 @@ impl WorkerProjectionService {
let Some(service) = service.upgrade() else { let Some(service) = service.upgrade() else {
return; return;
}; };
if let Err(error) = service if let Err(error) = service.apply_broker_event(task_runtime_id.as_str(), event) {
.apply_broker_event(task_runtime_id.as_str(), event)
.await
{
tracing::warn!( tracing::warn!(
runtime_id = %task_runtime_id, runtime_id = %task_runtime_id,
error = %error, error = %error,
"failed to update durable Worker projection" "failed to update durable Worker projection; restarting from a fresh snapshot"
); );
if subscription.restart().is_err() {
return;
}
} }
} }
}); });
if let Some(previous) = self self.tasks
.tasks
.lock() .lock()
.expect("Worker projection task lock poisoned") .expect("Worker projection task lock poisoned")
.insert(runtime_id, task) .insert(runtime_id, task);
{
previous.abort();
}
Ok(()) Ok(())
} }
@@ -128,110 +194,96 @@ impl WorkerProjectionService {
worker: &protocol::subscription::SubscriptionWorker, worker: &protocol::subscription::SubscriptionWorker,
observed_at: &str, observed_at: &str,
) -> crate::Result<()> { ) -> crate::Result<()> {
let commit = self.store.apply_worker_registry_observation( self.publish_ordered(|| {
self.store
.apply_worker_registry_observation(
self.workspace_id.as_str(), self.workspace_id.as_str(),
runtime_id, runtime_id,
0,
worker, worker,
observed_at, observed_at,
)?; )
self.publish_commit(commit) .map(|commit| ((), commit))
})
} }
pub fn refresh(&self, worker: &RuntimeWorkerRef) -> crate::Result<()> { pub fn refresh(&self, worker: &RuntimeWorkerRef) -> crate::Result<()> {
let commit = self self.publish_ordered(|| {
.store self.store
.commit_worker_registry_projection_refresh(self.workspace_id.as_str(), worker)?; .commit_worker_registry_projection_refresh(self.workspace_id.as_str(), worker)
self.publish_commit(commit) .map(|commit| ((), commit))
})
} }
async fn apply_broker_event( fn apply_broker_event(
&self, &self,
runtime_id: &str, runtime_id: &str,
event: BrokerSubscriptionEvent, event: BrokerSubscriptionEvent,
) -> crate::Result<()> { ) -> crate::Result<()> {
let observed_at = Utc::now().to_rfc3339(); let observed_at = Utc::now().to_rfc3339();
let projection_generation = |connection_generation: u64| { self.publish_ordered(|| {
self.connection_generation_epoch
.saturating_add(connection_generation)
};
let commit = match event { let commit = match event {
BrokerSubscriptionEvent::Snapshot { BrokerSubscriptionEvent::Snapshot {
connection_generation,
snapshot_revision,
snapshot: SubscriptionSnapshot::Workers { workers }, snapshot: SubscriptionSnapshot::Workers { workers },
} => self.store.reconcile_worker_registry_snapshot( } => self.store.reconcile_worker_registry_snapshot(
self.workspace_id.as_str(), self.workspace_id.as_str(),
runtime_id, runtime_id,
projection_generation(connection_generation),
snapshot_revision,
workers.as_slice(), workers.as_slice(),
observed_at.as_str(), observed_at.as_str(),
)?, )?,
BrokerSubscriptionEvent::Snapshot { .. } => return Ok(()), BrokerSubscriptionEvent::Snapshot { .. } => {
return Ok((
(),
WorkerRegistryProjectionCommit {
changes: Vec::new(),
},
));
}
BrokerSubscriptionEvent::Event { BrokerSubscriptionEvent::Event {
connection_generation,
payload: SubscriptionEventPayload::WorkerUpserted { worker }, payload: SubscriptionEventPayload::WorkerUpserted { worker },
..
} => self.store.apply_worker_registry_observation( } => self.store.apply_worker_registry_observation(
self.workspace_id.as_str(), self.workspace_id.as_str(),
runtime_id, runtime_id,
projection_generation(connection_generation),
&worker, &worker,
observed_at.as_str(), observed_at.as_str(),
)?, )?,
BrokerSubscriptionEvent::Event { BrokerSubscriptionEvent::Event {
connection_generation,
subject_revision,
payload: SubscriptionEventPayload::WorkerRemoved { worker_id, .. }, payload: SubscriptionEventPayload::WorkerRemoved { worker_id, .. },
} => self.store.mark_worker_registry_observation_unavailable( } => self.store.mark_worker_registry_observation_unavailable(
self.workspace_id.as_str(), self.workspace_id.as_str(),
runtime_id, runtime_id,
projection_generation(connection_generation),
Some(worker_id.as_str()), Some(worker_id.as_str()),
Some(subject_revision),
observed_at.as_str(), observed_at.as_str(),
)?, )?,
BrokerSubscriptionEvent::Event { .. } => return Ok(()), BrokerSubscriptionEvent::Event { .. } => {
BrokerSubscriptionEvent::Disconnected { return Ok((
connection_generation, (),
.. WorkerRegistryProjectionCommit {
changes: Vec::new(),
},
));
} }
| BrokerSubscriptionEvent::Rejected { BrokerSubscriptionEvent::Disconnected { .. }
connection_generation, | BrokerSubscriptionEvent::Rejected { .. }
.. | BrokerSubscriptionEvent::Closed { .. } => {
} self.store.mark_worker_registry_observation_unavailable(
| BrokerSubscriptionEvent::Closed {
connection_generation,
..
} => self.store.mark_worker_registry_observation_unavailable(
self.workspace_id.as_str(), self.workspace_id.as_str(),
runtime_id, runtime_id,
projection_generation(connection_generation),
None,
None, None,
observed_at.as_str(), observed_at.as_str(),
)?, )?
}
}; };
self.publish_commit(commit) Ok(((), commit))
})
} }
/// Broadcast an already committed catalog payload. Persistence always happens fn publish_commit(&self, commit: WorkerRegistryProjectionCommit) {
/// before this call; a process crash therefore closes subscribers, whose next
/// subscription starts from the durable snapshot. Retrying a removal returns
/// the same durable change and can heal an in-process missed broadcast.
pub(crate) fn publish_commit(
&self,
commit: WorkerRegistryProjectionCommit,
) -> crate::Result<()> {
if commit.changes.is_empty() { if commit.changes.is_empty() {
return Ok(()); return;
} }
let _ = self.events.send(WorkerProjectionEvent { let _ = self.events.send(WorkerProjectionEvent {
revision: commit.revision,
changes: commit.changes, changes: commit.changes,
}); });
Ok(())
} }
} }
@@ -240,3 +292,159 @@ impl Drop for WorkerProjectionService {
self.shutdown(); self.shutdown();
} }
} }
#[cfg(test)]
mod tests {
use std::sync::{Arc, Barrier, Mutex};
use crate::store::{
AccountRecord, ControlPlaneStore, SqliteWorkspaceStore, WorkerCatalogChange,
WorkerRegistryRecord, WorkspaceRecord,
};
use worker_runtime::identity::RuntimeWorkerRef;
use super::WorkerProjectionService;
const WORKSPACE_ID: &str = "projection-ordering-test";
async fn fixture() -> (
Arc<SqliteWorkspaceStore>,
Arc<WorkerProjectionService>,
tempfile::TempDir,
) {
let temp = tempfile::tempdir().unwrap();
let store = Arc::new(SqliteWorkspaceStore::open(temp.path().join("workspace.db")).unwrap());
store
.upsert_account(&AccountRecord {
account_id: "owner".to_string(),
kind: "user".to_string(),
handle: "owner".to_string(),
display_name: "Owner".to_string(),
created_at: "1".to_string(),
updated_at: "1".to_string(),
})
.unwrap();
store
.upsert_workspace(&WorkspaceRecord {
workspace_id: WORKSPACE_ID.to_string(),
owner_account_id: "owner".to_string(),
display_name: "Projection ordering".to_string(),
state: "active".to_string(),
created_at: "1".to_string(),
updated_at: "1".to_string(),
})
.await
.unwrap();
let service = WorkerProjectionService::new(WORKSPACE_ID, store.clone());
(store, service, temp)
}
fn worker(runtime_id: &str, worker_id: &str) -> WorkerRegistryRecord {
WorkerRegistryRecord {
workspace_id: WORKSPACE_ID.to_string(),
worker: RuntimeWorkerRef::new(runtime_id, worker_id),
display_name: format!("{runtime_id}-{worker_id}"),
profile: Some("builtin:coder".to_string()),
retention_state: "normal".to_string(),
transcript_ref: None,
session_ref: None,
summary_ref: None,
diagnostics_ref: None,
created_at: "1".to_string(),
updated_at: "1".to_string(),
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn snapshot_registration_is_atomic_with_concurrent_publication() {
let (store, service, _temp) = fixture().await;
let barrier = Arc::new(Barrier::new(2));
let subscriber_service = service.clone();
let subscriber_barrier = barrier.clone();
let subscriber = std::thread::spawn(move || {
subscriber_barrier.wait();
subscriber_service.subscribe_ordered(10).unwrap()
});
let publisher_service = service.clone();
let publisher_barrier = barrier.clone();
let publisher_store = store.clone();
let record = worker("runtime-a", "worker-a");
let publisher = std::thread::spawn(move || {
publisher_barrier.wait();
publisher_service
.publish_ordered(|| {
publisher_store
.upsert_worker_registry(&record)
.map(|commit| ((), commit))
})
.unwrap();
});
let mut subscription = subscriber.join().unwrap();
publisher.join().unwrap();
let in_snapshot = subscription
.take_snapshot()
.iter()
.any(|record| record.registry.worker.worker_id == "worker-a");
let in_event = match subscription.try_recv() {
Ok(event) => matches!(
event.changes.as_slice(),
[WorkerCatalogChange::Upsert(record)]
if record.registry.worker.worker_id == "worker-a"
),
Err(tokio::sync::broadcast::error::TryRecvError::Empty) => false,
Err(error) => panic!("unexpected projection receive error: {error}"),
};
assert_ne!(
in_snapshot, in_event,
"the commit must cross the boundary exactly once"
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn concurrent_commits_are_published_in_durable_execution_order() {
let (store, service, _temp) = fixture().await;
let mut subscription = service.subscribe_ordered(10).unwrap();
assert!(subscription.take_snapshot().is_empty());
let barrier = Arc::new(Barrier::new(3));
let durable_order = Arc::new(Mutex::new(Vec::new()));
let mut publishers = Vec::new();
for worker_id in ["worker-a", "worker-b"] {
let service = service.clone();
let store = store.clone();
let barrier = barrier.clone();
let durable_order = durable_order.clone();
let record = worker("runtime-a", worker_id);
publishers.push(std::thread::spawn(move || {
barrier.wait();
service
.publish_ordered(|| {
durable_order
.lock()
.unwrap()
.push(record.worker.worker_id.clone());
store
.upsert_worker_registry(&record)
.map(|commit| ((), commit))
})
.unwrap();
}));
}
barrier.wait();
for publisher in publishers {
publisher.join().unwrap();
}
let published_order = [
subscription.try_recv().unwrap(),
subscription.try_recv().unwrap(),
]
.into_iter()
.map(|event| match event.changes.as_slice() {
[WorkerCatalogChange::Upsert(record)] => record.registry.worker.worker_id.clone(),
changes => panic!("unexpected projection changes: {changes:?}"),
})
.collect::<Vec<_>>();
assert_eq!(published_order, *durable_order.lock().unwrap());
}
}
@@ -73,10 +73,8 @@ const WORKSPACE_DELETION_PURGE_TABLES: &[&str] = &[
"worker_orphan_diagnostics", "worker_orphan_diagnostics",
"worker_registry", "worker_registry",
"worker_registry_observations", "worker_registry_observations",
"worker_registry_projection_cursors",
"worker_registry_projection_diagnostics", "worker_registry_projection_diagnostics",
"worker_registry_projection_removals", "worker_registry_projection_removals",
"worker_registry_projection_revisions",
"worker_removal_operations", "worker_removal_operations",
"worker_retention_audit_events", "worker_retention_audit_events",
"worker_session_archives", "worker_session_archives",
@@ -271,7 +271,6 @@ async fn run_worker_protocol(
worker_id: worker_id.clone(), worker_id: worker_id.clone(),
runtime_id: Some(runtime_id), runtime_id: Some(runtime_id),
}, },
snapshot_revision: 0,
snapshot: SubscriptionSnapshot::WorkerProtocol { snapshot: SubscriptionSnapshot::WorkerProtocol {
worker_id: worker_id.clone(), worker_id: worker_id.clone(),
events: vec![initial_event], events: vec![initial_event],
@@ -284,13 +283,10 @@ async fn run_worker_protocol(
{ {
return; return;
} }
let mut subject_revision = 0_u64;
while let Some(event) = events.recv().await { while let Some(event) = events.recv().await {
subject_revision = subject_revision.saturating_add(1);
let frame = let frame =
SubscriptionFrame::new(SubscriptionFramePayload::Event(SubscriptionEvent::Event { SubscriptionFrame::new(SubscriptionFramePayload::Event(SubscriptionEvent::Event {
subscription_id: subscription_id.clone(), subscription_id: subscription_id.clone(),
subject_revision,
payload: SubscriptionEventPayload::WorkerProtocol { payload: SubscriptionEventPayload::WorkerProtocol {
worker_id: worker_id.clone(), worker_id: worker_id.clone(),
event, event,
@@ -333,12 +329,10 @@ async fn run_workspace_workers(
subscription_id: SubscriptionId, subscription_id: SubscriptionId,
outbound: mpsc::Sender<WsMessage>, outbound: mpsc::Sender<WsMessage>,
) { ) {
// Subscribe before reading so commits racing with the snapshot are replayed. let Ok(mut projection) = api
let mut events = api.worker_projection.subscribe(); .worker_projection
let Ok((snapshot_revision, records)) = api.store.worker_registry_projection_snapshot( .subscribe_ordered(api.config.max_records.max(1))
&api.config.workspace_id, else {
api.config.max_records.max(1),
) else {
let _ = send_rejected( let _ = send_rejected(
&outbound, &outbound,
request_id, request_id,
@@ -348,7 +342,8 @@ async fn run_workspace_workers(
.await; .await;
return; return;
}; };
let mut workers = records let mut workers = projection
.take_snapshot()
.into_iter() .into_iter()
.filter_map(|record| project_registry_worker(&api, record)) .filter_map(|record| project_registry_worker(&api, record))
.collect::<Vec<_>>(); .collect::<Vec<_>>();
@@ -360,7 +355,6 @@ async fn run_workspace_workers(
request_id, request_id,
subscription_id: subscription_id.clone(), subscription_id: subscription_id.clone(),
selector: EventSubscriptionSelector::WorkspaceWorkers, selector: EventSubscriptionSelector::WorkspaceWorkers,
snapshot_revision,
snapshot: SubscriptionSnapshot::Workers { workers }, snapshot: SubscriptionSnapshot::Workers { workers },
}, },
)), )),
@@ -372,16 +366,14 @@ async fn run_workspace_workers(
} }
loop { loop {
match events.recv().await { match projection.recv().await {
Ok(event) if event.revision <= snapshot_revision => continue,
Ok(event) => { Ok(event) => {
for change in event.changes { for change in event.changes {
let payload = match change { let payload = match change {
crate::worker_projection::WorkerProjectionChange::Upsert(record) => { crate::worker_projection::WorkerProjectionChange::Upsert(record) => {
let Some(mut worker) = project_registry_worker(&api, record) else { let Some(worker) = project_registry_worker(&api, record) else {
continue; continue;
}; };
worker.subject_revision = event.revision;
SubscriptionEventPayload::WorkerUpserted { worker } SubscriptionEventPayload::WorkerUpserted { worker }
} }
crate::worker_projection::WorkerProjectionChange::Removed(worker) => { crate::worker_projection::WorkerProjectionChange::Removed(worker) => {
@@ -396,7 +388,7 @@ async fn run_workspace_workers(
} }
} }
}; };
if send_event(&outbound, &subscription_id, event.revision, payload) if send_event(&outbound, &subscription_id, payload)
.await .await
.is_err() .is_err()
{ {
@@ -431,7 +423,6 @@ fn project_registry_worker(
let mut worker = if let Some(observation) = record.observation { let mut worker = if let Some(observation) = record.observation {
let mut worker = observation.worker; let mut worker = observation.worker;
worker.availability = observation.availability; worker.availability = observation.availability;
worker.subject_revision = observation.projection_revision;
worker worker
} else { } else {
SubscriptionWorker { SubscriptionWorker {
@@ -442,7 +433,6 @@ fn project_registry_worker(
runtime_id: Some(runtime_id.clone()), runtime_id: Some(runtime_id.clone()),
resource_key: record.resource_key.clone(), resource_key: record.resource_key.clone(),
availability: protocol::subscription::SubscriptionWorkerAvailability::Unavailable, availability: protocol::subscription::SubscriptionWorkerAvailability::Unavailable,
subject_revision: 0,
worker_state: None, worker_state: None,
state: protocol::subscription::SubscriptionWorkerState::Stopped, state: protocol::subscription::SubscriptionWorkerState::Stopped,
has_running_internal_workers: false, has_running_internal_workers: false,
@@ -559,14 +549,12 @@ fn sort_workers(workers: &mut [SubscriptionWorker]) {
async fn send_event( async fn send_event(
outbound: &mpsc::Sender<WsMessage>, outbound: &mpsc::Sender<WsMessage>,
subscription_id: &SubscriptionId, subscription_id: &SubscriptionId,
subject_revision: u64,
payload: SubscriptionEventPayload, payload: SubscriptionEventPayload,
) -> Result<(), ()> { ) -> Result<(), ()> {
send_frame( send_frame(
outbound, outbound,
SubscriptionFrame::new(SubscriptionFramePayload::Event(SubscriptionEvent::Event { SubscriptionFrame::new(SubscriptionFramePayload::Event(SubscriptionEvent::Event {
subscription_id: subscription_id.clone(), subscription_id: subscription_id.clone(),
subject_revision,
payload, payload,
})), })),
) )
@@ -617,7 +605,6 @@ mod tests {
runtime_id: Some("runtime-1".to_string()), runtime_id: Some("runtime-1".to_string()),
resource_key: Some("W-1".to_string()), resource_key: Some("W-1".to_string()),
availability: protocol::subscription::SubscriptionWorkerAvailability::Observed, availability: protocol::subscription::SubscriptionWorkerAvailability::Observed,
subject_revision: 1,
worker_state: None, worker_state: None,
state: protocol::subscription::SubscriptionWorkerState::Idle, state: protocol::subscription::SubscriptionWorkerState::Idle,
has_running_internal_workers: false, has_running_internal_workers: false,
@@ -227,7 +227,7 @@ Deno.serve({ hostname: "127.0.0.1", port }, async (request) => {
if (frame?.message?.method !== "subscribe_events") return; if (frame?.message?.method !== "subscribe_events") return;
const requestId = frame.message.params.request_id; const requestId = frame.message.params.request_id;
socket.send(JSON.stringify({ socket.send(JSON.stringify({
protocol_version: 1, protocol_version: 2,
frame: "response", frame: "response",
message: { message: {
result: "subscribed", result: "subscribed",
@@ -235,7 +235,6 @@ Deno.serve({ hostname: "127.0.0.1", port }, async (request) => {
request_id: requestId, request_id: requestId,
subscription_id: "fixture-workers", subscription_id: "fixture-workers",
selector: frame.message.params.selector, selector: frame.message.params.selector,
snapshot_revision: 1,
snapshot: { topic: "workers", data: { workers: [] } }, snapshot: { topic: "workers", data: { workers: [] } },
}, },
}, },
@@ -129,7 +129,7 @@ Deno.serve({ hostname: "127.0.0.1", port }, async (request) => {
if (frame?.message?.method !== "subscribe_events") return; if (frame?.message?.method !== "subscribe_events") return;
const requestId = frame.message.params.request_id; const requestId = frame.message.params.request_id;
socket.send(JSON.stringify({ socket.send(JSON.stringify({
protocol_version: 1, protocol_version: 2,
frame: "response", frame: "response",
message: { message: {
result: "subscribed", result: "subscribed",
@@ -137,7 +137,6 @@ Deno.serve({ hostname: "127.0.0.1", port }, async (request) => {
request_id: requestId, request_id: requestId,
subscription_id: "fixture-workers", subscription_id: "fixture-workers",
selector: { topic: "workspace_workers" }, selector: { topic: "workspace_workers" },
snapshot_revision: 1,
snapshot: { topic: "workers", data: { workers: [] } }, snapshot: { topic: "workers", data: { workers: [] } },
}, },
}, },
@@ -138,7 +138,7 @@ Deno.serve({ hostname: "127.0.0.1", port }, async (request) => {
const frame = JSON.parse(String(event.data)); const frame = JSON.parse(String(event.data));
if (frame?.message?.method !== "subscribe_events") return; if (frame?.message?.method !== "subscribe_events") return;
socket.send(JSON.stringify({ socket.send(JSON.stringify({
protocol_version: 1, protocol_version: 2,
frame: "response", frame: "response",
message: { message: {
result: "subscribed", result: "subscribed",
@@ -146,7 +146,6 @@ Deno.serve({ hostname: "127.0.0.1", port }, async (request) => {
request_id: frame.message.params.request_id, request_id: frame.message.params.request_id,
subscription_id: "fixture-workers", subscription_id: "fixture-workers",
selector: { topic: "workspace_workers" }, selector: { topic: "workspace_workers" },
snapshot_revision: 1,
snapshot: { topic: "workers", data: { workers: [] } }, snapshot: { topic: "workers", data: { workers: [] } },
}, },
}, },
File diff suppressed because one or more lines are too long
+3 -7
View File
@@ -304,11 +304,7 @@ resource_key?: string | null,
*/ */
availability: SubscriptionWorkerAvailability, availability: SubscriptionWorkerAvailability,
/** /**
* Producer-owned monotonic revision for this Worker subject. * Latest foreground state observed from the Worker. This remains
*/
subject_revision: number,
/**
* Latest revisioned foreground state observed from the Worker. This remains
* absent until an authoritative Worker snapshot/event has been applied. * absent until an authoritative Worker snapshot/event has been applied.
*/ */
worker_state?: WorkerStateSnapshot | null, worker_state?: WorkerStateSnapshot | null,
@@ -336,9 +332,9 @@ export type SubscriptionRequest = { "method": "subscribe_events", "params": { re
export type SubscriptionWorkerProtocolMethod = { subscription_id: SubscriptionId, method: Method, }; export type SubscriptionWorkerProtocolMethod = { subscription_id: SubscriptionId, method: Method, };
export type SubscriptionResponse = { "result": "subscribed", "payload": { request_id: SubscriptionRequestId, subscription_id: SubscriptionId, selector: EventSubscriptionSelector, snapshot_revision: number, snapshot: SubscriptionSnapshot, } } | { "result": "unsubscribed", "payload": { request_id: SubscriptionRequestId, subscription_id: SubscriptionId, } } | { "result": "subscription_rejected", "payload": { request_id: SubscriptionRequestId, subscription_id?: SubscriptionId | null, code: SubscriptionRejectionCode, message: string, } }; export type SubscriptionResponse = { "result": "subscribed", "payload": { request_id: SubscriptionRequestId, subscription_id: SubscriptionId, selector: EventSubscriptionSelector, snapshot: SubscriptionSnapshot, } } | { "result": "unsubscribed", "payload": { request_id: SubscriptionRequestId, subscription_id: SubscriptionId, } } | { "result": "subscription_rejected", "payload": { request_id: SubscriptionRequestId, subscription_id?: SubscriptionId | null, code: SubscriptionRejectionCode, message: string, } };
export type SubscriptionEvent = { "event": "event", "data": { subscription_id: SubscriptionId, subject_revision: number, payload: SubscriptionEventPayload, } } | { "event": "subscription_closed", "data": { subscription_id: SubscriptionId, code: SubscriptionTerminationCode, message: string, } }; export type SubscriptionEvent = { "event": "event", "data": { subscription_id: SubscriptionId, payload: SubscriptionEventPayload, } } | { "event": "subscription_closed", "data": { subscription_id: SubscriptionId, code: SubscriptionTerminationCode, message: string, } };
export type SubscriptionFramePayload = { "frame": "request", "message": SubscriptionRequest } | { "frame": "response", "message": SubscriptionResponse } | { "frame": "event", "message": SubscriptionEvent } | { "frame": "worker_protocol", "message": SubscriptionWorkerProtocolMethod }; export type SubscriptionFramePayload = { "frame": "request", "message": SubscriptionRequest } | { "frame": "response", "message": SubscriptionResponse } | { "frame": "event", "message": SubscriptionEvent } | { "frame": "worker_protocol", "message": SubscriptionWorkerProtocolMethod };
@@ -83,7 +83,7 @@ test('invalid inbound frame closes before dispatch and reconnects for a fresh sn
first.message( first.message(
JSON.stringify({ JSON.stringify({
protocol_version: 1, protocol_version: 2,
frame: 'event', frame: 'event',
message: { message: {
event: 'subscription_closed', event: 'subscription_closed',
@@ -141,7 +141,7 @@ test('valid subscription closure remains isolated and resubscribes only its sele
: { topic: 'workspace_workdirs', data: { workdirs: [] } }; : { topic: 'workspace_workdirs', data: { workdirs: [] } };
socket.message( socket.message(
JSON.stringify({ JSON.stringify({
protocol_version: 1, protocol_version: 2,
frame: 'response', frame: 'response',
message: { message: {
result: 'subscribed', result: 'subscribed',
@@ -149,7 +149,6 @@ test('valid subscription closure remains isolated and resubscribes only its sele
request_id: request.message.params.request_id, request_id: request.message.params.request_id,
subscription_id: `subscription-${index + 1}`, subscription_id: `subscription-${index + 1}`,
selector, selector,
snapshot_revision: 1,
snapshot, snapshot,
}, },
}, },
@@ -160,7 +159,7 @@ test('valid subscription closure remains isolated and resubscribes only its sele
const sentBeforeClosure = socket.sent.length; const sentBeforeClosure = socket.sent.length;
socket.message( socket.message(
JSON.stringify({ JSON.stringify({
protocol_version: 1, protocol_version: 2,
frame: 'event', frame: 'event',
message: { message: {
event: 'subscription_closed', event: 'subscription_closed',
@@ -138,7 +138,7 @@ export class WorkspaceMultiplexer {
subscription.requestId = requestId; subscription.requestId = requestId;
this.#requests.set(requestId, subscription.clientId); this.#requests.set(requestId, subscription.clientId);
this.#send({ this.#send({
protocol_version: 1, protocol_version: 2,
frame: 'request', frame: 'request',
message: { message: {
method: 'subscribe_events', method: 'subscribe_events',
@@ -213,7 +213,7 @@ export class WorkspaceMultiplexer {
const subscription = this.#subscriptions.get(clientId); const subscription = this.#subscriptions.get(clientId);
if (!subscription?.subscriptionId) throw new Error('Worker protocol subscription is not open'); if (!subscription?.subscriptionId) throw new Error('Worker protocol subscription is not open');
this.#send({ this.#send({
protocol_version: 1, protocol_version: 2,
frame: 'worker_protocol', frame: 'worker_protocol',
message: { subscription_id: subscription.subscriptionId, method }, message: { subscription_id: subscription.subscriptionId, method },
}); });
@@ -225,7 +225,7 @@ export class WorkspaceMultiplexer {
this.#subscriptions.delete(clientId); this.#subscriptions.delete(clientId);
if (subscription.subscriptionId && this.#socket?.readyState === WebSocket.OPEN) { if (subscription.subscriptionId && this.#socket?.readyState === WebSocket.OPEN) {
this.#send({ this.#send({
protocol_version: 1, protocol_version: 2,
frame: 'request', frame: 'request',
message: { message: {
method: 'unsubscribe_events', method: 'unsubscribe_events',
@@ -16,7 +16,7 @@ function assert(condition: unknown, message: string): asserts condition {
function subscribedFrame(overrides: Record<string, unknown> = {}): Record<string, unknown> { function subscribedFrame(overrides: Record<string, unknown> = {}): Record<string, unknown> {
return { return {
protocol_version: 1, protocol_version: 2,
frame: 'response', frame: 'response',
message: { message: {
result: 'subscribed', result: 'subscribed',
@@ -24,7 +24,6 @@ function subscribedFrame(overrides: Record<string, unknown> = {}): Record<string
request_id: 'request-1', request_id: 'request-1',
subscription_id: 'subscription-1', subscription_id: 'subscription-1',
selector: { topic: 'workspace_workers' }, selector: { topic: 'workspace_workers' },
snapshot_revision: 1,
snapshot: { topic: 'workers', data: { workers: [] } }, snapshot: { topic: 'workers', data: { workers: [] } },
...overrides, ...overrides,
}, },
@@ -86,7 +85,6 @@ Deno.test('generated subscription validator accepts an external Workdir Worker',
runtime_id: 'runtime-1', runtime_id: 'runtime-1',
resource_key: 'worker-resource-1', resource_key: 'worker-resource-1',
availability: 'observed', availability: 'observed',
subject_revision: 1,
state: 'idle', state: 'idle',
has_running_internal_workers: false, has_running_internal_workers: false,
workspace_id: 'workspace-1', workspace_id: 'workspace-1',
@@ -123,8 +121,8 @@ Deno.test('generated subscription validator rejects malformed and newer shapes',
'unknown nested field', 'unknown nested field',
subscribedFrame({ snapshot: { topic: 'workers', data: { workers: [], future: true } } }), subscribedFrame({ snapshot: { topic: 'workers', data: { workers: [], future: true } } }),
], ],
['wrong field type', subscribedFrame({ snapshot_revision: '1' })], ['wrong field type', { ...subscribedFrame(), protocol_version: '2' }],
['unsafe integer', subscribedFrame({ snapshot_revision: Number.MAX_SAFE_INTEGER + 1 })], ['unsafe integer', { ...subscribedFrame(), protocol_version: Number.MAX_SAFE_INTEGER + 1 }],
[ [
'missing serialized Worker defaults', 'missing serialized Worker defaults',
subscribedFrame({ subscribedFrame({
@@ -132,23 +130,22 @@ Deno.test('generated subscription validator rejects malformed and newer shapes',
topic: 'workers', topic: 'workers',
data: { data: {
workers: [ workers: [
{ worker_id: 'worker-1', subject_revision: 1, state: 'idle' }, { worker_id: 'worker-1', state: 'idle' },
], ],
}, },
}, },
}), }),
], ],
['unsupported version', { ...subscribedFrame(), protocol_version: 2 }], ['unsupported version', { ...subscribedFrame(), protocol_version: 1 }],
[ [
'newer Worker event', 'newer Worker event',
{ {
protocol_version: 1, protocol_version: 2,
frame: 'event', frame: 'event',
message: { message: {
event: 'event', event: 'event',
data: { data: {
subscription_id: 'subscription-1', subscription_id: 'subscription-1',
subject_revision: 2,
payload: { payload: {
event: 'worker_protocol', event: 'worker_protocol',
data: { worker_id: 'worker-1', event: { event: 'future_event' } }, data: { worker_id: 'worker-1', event: { event: 'future_event' } },
@@ -160,13 +157,12 @@ Deno.test('generated subscription validator rejects malformed and newer shapes',
[ [
'out-of-range Rust u32', 'out-of-range Rust u32',
{ {
protocol_version: 1, protocol_version: 2,
frame: 'event', frame: 'event',
message: { message: {
event: 'event', event: 'event',
data: { data: {
subscription_id: 'subscription-1', subscription_id: 'subscription-1',
subject_revision: 2,
payload: { payload: {
event: 'worker_protocol', event: 'worker_protocol',
data: { data: {
@@ -207,7 +203,7 @@ Deno.test('generated subscription validator enforces Rust semantic bounds and se
[ [
'oversized rejection message', 'oversized rejection message',
{ {
protocol_version: 1, protocol_version: 2,
frame: 'response', frame: 'response',
message: { message: {
result: 'subscription_rejected', result: 'subscription_rejected',
@@ -242,7 +238,6 @@ Deno.test('generated subscription validator enforces Rust semantic bounds and se
worker_id: 'worker-1', worker_id: 'worker-1',
runtime_id: 'runtime-1', runtime_id: 'runtime-1',
availability: 'observed', availability: 'observed',
subject_revision: 1,
state: 'idle', state: 'idle',
has_running_internal_workers: false, has_running_internal_workers: false,
}, },
@@ -250,7 +245,6 @@ Deno.test('generated subscription validator enforces Rust semantic bounds and se
worker_id: 'worker-1', worker_id: 'worker-1',
runtime_id: 'runtime-1', runtime_id: 'runtime-1',
availability: 'observed', availability: 'observed',
subject_revision: 2,
state: 'running', state: 'running',
has_running_internal_workers: false, has_running_internal_workers: false,
}, },
@@ -262,7 +256,7 @@ Deno.test('generated subscription validator enforces Rust semantic bounds and se
[ [
'server-to-Browser direction mismatch', 'server-to-Browser direction mismatch',
{ {
protocol_version: 1, protocol_version: 2,
frame: 'request', frame: 'request',
message: { message: {
method: 'unsubscribe_events', method: 'unsubscribe_events',
@@ -281,13 +275,12 @@ Deno.test('generated subscription validator enforces Rust semantic bounds and se
Deno.test('aggregate limits include keys in unconstrained Worker JSON values', () => { Deno.test('aggregate limits include keys in unconstrained Worker JSON values', () => {
const oversizedKey = 'k'.repeat(MAX_SUBSCRIPTION_STRING_BYTES + 1); const oversizedKey = 'k'.repeat(MAX_SUBSCRIPTION_STRING_BYTES + 1);
const result = decode({ const result = decode({
protocol_version: 1, protocol_version: 2,
frame: 'event', frame: 'event',
message: { message: {
event: 'event', event: 'event',
data: { data: {
subscription_id: 'subscription-1', subscription_id: 'subscription-1',
subject_revision: 2,
payload: { payload: {
event: 'worker_protocol', event: 'worker_protocol',
data: { data: {
@@ -306,13 +299,12 @@ Deno.test('aggregate limits include keys in unconstrained Worker JSON values', (
Deno.test('stateful selector validation fences a routed event before projection', () => { Deno.test('stateful selector validation fences a routed event before projection', () => {
const frame: SubscriptionFrame = { const frame: SubscriptionFrame = {
protocol_version: 1, protocol_version: 2,
frame: 'event', frame: 'event',
message: { message: {
event: 'event', event: 'event',
data: { data: {
subscription_id: 'subscription-1', subscription_id: 'subscription-1',
subject_revision: 2,
payload: { payload: {
event: 'worker_protocol', event: 'worker_protocol',
data: { worker_id: 'worker-2', event: { event: 'thinking_start' } }, data: { worker_id: 'worker-2', event: { event: 'thinking_start' } },
@@ -43,7 +43,7 @@ function snapshot(
availability: SubscriptionWorker["availability"], availability: SubscriptionWorker["availability"],
): SubscriptionFrame { ): SubscriptionFrame {
return { return {
protocol_version: 1, protocol_version: 2,
frame: "response", frame: "response",
message: { message: {
result: "subscribed", result: "subscribed",
@@ -51,7 +51,6 @@ function snapshot(
request_id: "request-1", request_id: "request-1",
subscription_id: "subscription-1", subscription_id: "subscription-1",
selector: { topic: "workspace_workers" }, selector: { topic: "workspace_workers" },
snapshot_revision: 1,
snapshot: { snapshot: {
topic: "workers", topic: "workers",
data: { data: {
@@ -60,7 +59,6 @@ function snapshot(
runtime_id: "runtime-a", runtime_id: "runtime-a",
resource_key: "W-1", resource_key: "W-1",
availability, availability,
subject_revision: 1,
state, state,
// An absent foreground snapshot must not prevent stopping an active catalog Worker. // An absent foreground snapshot must not prevent stopping an active catalog Worker.
worker_state: null, worker_state: null,
@@ -6,55 +6,55 @@ import type {
export type WorkspaceWorkersProjection = { export type WorkspaceWorkersProjection = {
workers: Map<string, SubscriptionWorker>; workers: Map<string, SubscriptionWorker>;
revisions: Map<string, number>; subscriptionId: string | null;
}; };
export function createWorkspaceWorkersProjection(): WorkspaceWorkersProjection { export function createWorkspaceWorkersProjection(): WorkspaceWorkersProjection {
return { workers: new Map(), revisions: new Map() }; return { workers: new Map(), subscriptionId: null };
} }
export function applyWorkspaceWorkersFrame( export function applyWorkspaceWorkersFrame(
projection: WorkspaceWorkersProjection, projection: WorkspaceWorkersProjection,
frame: SubscriptionFrame, frame: SubscriptionFrame,
): void { ): void {
if (frame.protocol_version !== 1) throw new Error('unsupported Worker subscription protocol'); if (frame.protocol_version !== 2) throw new Error('unsupported Worker subscription protocol');
if (frame.frame === 'response' && frame.message.result === 'subscribed') { if (frame.frame === 'response' && frame.message.result === 'subscribed') {
if (frame.message.payload.selector.topic !== 'workspace_workers') return; if (frame.message.payload.selector.topic !== 'workspace_workers') return;
const snapshot = frame.message.payload.snapshot; const snapshot = frame.message.payload.snapshot;
if (snapshot.topic !== 'workers') throw new Error('workspace_workers returned a non-Worker snapshot'); if (snapshot.topic !== 'workers') throw new Error('workspace_workers returned a non-Worker snapshot');
projection.workers.clear();
projection.revisions.clear(); // Build the replacement before mutating the live projection so one malformed
// Worker cannot partially clear the previous subscription lifetime.
const workers = new Map<string, SubscriptionWorker>();
for (const worker of snapshot.data.workers) { for (const worker of snapshot.data.workers) {
const key = workerKey(worker.runtime_id, worker.worker_id); workers.set(workerKey(worker.runtime_id, worker.worker_id), worker);
projection.workers.set(key, worker);
projection.revisions.set(key, worker.subject_revision);
} }
projection.workers = workers;
projection.subscriptionId = frame.message.payload.subscription_id;
return; return;
} }
if (frame.frame !== 'event' || frame.message.event !== 'event') return; if (frame.frame !== 'event' || frame.message.event !== 'event') return;
applyPayload(projection, frame.message.data.subject_revision, frame.message.data.payload); const subscriptionId = frame.message.data.subscription_id;
if (!projection.subscriptionId) {
throw new Error('workspace_workers event arrived before its snapshot');
}
if (subscriptionId !== projection.subscriptionId) return;
applyPayload(projection, frame.message.data.payload);
} }
function applyPayload( function applyPayload(
projection: WorkspaceWorkersProjection, projection: WorkspaceWorkersProjection,
subjectRevision: number,
payload: SubscriptionEventPayload, payload: SubscriptionEventPayload,
): void { ): void {
if (payload.event === 'worker_upserted') { if (payload.event === 'worker_upserted') {
const worker = payload.data.worker; const worker = payload.data.worker;
const key = workerKey(worker.runtime_id, worker.worker_id); projection.workers.set(workerKey(worker.runtime_id, worker.worker_id), worker);
if (subjectRevision <= (projection.revisions.get(key) ?? 0)) return;
projection.revisions.set(key, subjectRevision);
projection.workers.set(key, worker);
} else if (payload.event === 'worker_removed') { } else if (payload.event === 'worker_removed') {
const key = workerKey(payload.data.runtime_id, payload.data.worker_id); projection.workers.delete(workerKey(payload.data.runtime_id, payload.data.worker_id));
if (subjectRevision <= (projection.revisions.get(key) ?? 0)) return;
projection.revisions.set(key, subjectRevision);
projection.workers.delete(key);
} }
} }
function workerKey(runtimeId: string | null | undefined, workerId: string): string { function workerKey(runtimeId: string | null | undefined, workerId: string): string {
if (!runtimeId) throw new Error('Workspace Worker projection is missing runtime_id'); if (!runtimeId) throw new Error('Workspace Worker projection is missing runtime_id');
return `${runtimeId}:${workerId}`; return JSON.stringify([runtimeId, workerId]);
} }
@@ -1,15 +1,29 @@
import type { SubscriptionFrame, SubscriptionWorker } from '$lib/generated/protocol'; import type {
SubscriptionEventPayload,
SubscriptionFrame,
SubscriptionWorker,
} from '$lib/generated/protocol';
import { liveWorkerState } from './worker-state';
import {
applyWorkspaceWorkersFrame,
createWorkspaceWorkersProjection,
} from './worker-subscription-model';
function assertEquals(actual: unknown, expected: unknown): void { function assertEquals(actual: unknown, expected: unknown): void {
if (JSON.stringify(actual) !== JSON.stringify(expected)) { if (JSON.stringify(actual) !== JSON.stringify(expected)) {
throw new Error(`expected ${JSON.stringify(expected)}, got ${JSON.stringify(actual)}`); throw new Error(`expected ${JSON.stringify(expected)}, got ${JSON.stringify(actual)}`);
} }
} }
import { liveWorkerState } from './worker-state';
import { function assertThrows(fn: () => void, message: string): void {
applyWorkspaceWorkersFrame, try {
createWorkspaceWorkersProjection, fn();
} from './worker-subscription-model'; } catch {
return;
}
throw new Error(message);
}
declare const Deno: { declare const Deno: {
test(name: string, fn: () => void | Promise<void>): void; test(name: string, fn: () => void | Promise<void>): void;
@@ -18,7 +32,6 @@ declare const Deno: {
function worker( function worker(
runtimeId: string, runtimeId: string,
workerId: string, workerId: string,
revision: number,
hasRunningInternalWorkers = false, hasRunningInternalWorkers = false,
availability: 'observed' | 'unavailable' = 'observed', availability: 'observed' | 'unavailable' = 'observed',
): SubscriptionWorker { ): SubscriptionWorker {
@@ -26,7 +39,6 @@ function worker(
worker_id: workerId, worker_id: workerId,
runtime_id: runtimeId, runtime_id: runtimeId,
availability, availability,
subject_revision: revision,
state: 'idle', state: 'idle',
has_running_internal_workers: hasRunningInternalWorkers, has_running_internal_workers: hasRunningInternalWorkers,
workspace_id: 'workspace-test', workspace_id: 'workspace-test',
@@ -36,15 +48,52 @@ function worker(
}; };
} }
function key(runtimeId: string, workerId: string): string {
return JSON.stringify([runtimeId, workerId]);
}
function snapshot(
subscriptionId: string,
workers: SubscriptionWorker[],
): SubscriptionFrame {
return {
protocol_version: 2,
frame: 'response',
message: {
result: 'subscribed',
payload: {
request_id: `request-${subscriptionId}`,
subscription_id: subscriptionId,
selector: { topic: 'workspace_workers' },
snapshot: { topic: 'workers', data: { workers } },
},
},
};
}
function event(
subscriptionId: string,
payload: SubscriptionEventPayload,
): SubscriptionFrame {
return {
protocol_version: 2,
frame: 'event',
message: {
event: 'event',
data: { subscription_id: subscriptionId, payload },
},
};
}
Deno.test('Worker list state uses the authoritative live snapshot separately from lifecycle', () => { Deno.test('Worker list state uses the authoritative live snapshot separately from lifecycle', () => {
const active = worker('runtime-a', 'worker-1', 1); const active = worker('runtime-a', 'worker-1');
active.worker_state = { active.worker_state = {
last_command_id: 1, last_command_id: 1,
state: { kind: 'busy', state: { kind: 'run', state: 'paused' } }, state: { kind: 'busy', state: { kind: 'run', state: 'paused' } },
}; };
assertEquals(liveWorkerState(active), 'paused'); assertEquals(liveWorkerState(active), 'paused');
const unavailable = worker('runtime-a', 'worker-2', 1, false, 'unavailable'); const unavailable = worker('runtime-a', 'worker-2', false, 'unavailable');
unavailable.worker_state = { unavailable.worker_state = {
last_command_id: 2, last_command_id: 2,
state: { kind: 'busy', state: { kind: 'run', state: 'running' } }, state: { kind: 'busy', state: { kind: 'run', state: 'running' } },
@@ -67,127 +116,118 @@ Deno.test('Worker list state uses the authoritative live snapshot separately fro
Deno.test('workspace Worker snapshot keeps equal local ids from different Runtimes', () => { Deno.test('workspace Worker snapshot keeps equal local ids from different Runtimes', () => {
const projection = createWorkspaceWorkersProjection(); const projection = createWorkspaceWorkersProjection();
const frame: SubscriptionFrame = { applyWorkspaceWorkersFrame(
protocol_version: 1, projection,
frame: 'response', snapshot('subscription-1', [worker('runtime-a', '1'), worker('runtime-b', '1')]),
message: { );
result: 'subscribed', assertEquals([...projection.workers.keys()].sort(), [key('runtime-a', '1'), key('runtime-b', '1')]);
payload: {
request_id: 'request-1',
subscription_id: 'subscription-1',
selector: { topic: 'workspace_workers' },
snapshot_revision: 1,
snapshot: {
topic: 'workers',
data: { workers: [worker('runtime-a', '1', 1), worker('runtime-b', '1', 1)] },
},
},
},
};
applyWorkspaceWorkersFrame(projection, frame);
assertEquals([...projection.workers.keys()].sort(), ['runtime-a:1', 'runtime-b:1']);
}); });
Deno.test('workspace Worker snapshot and live update preserve canonical Job metadata', () => { Deno.test('workspace Worker snapshot and ordered live update preserve canonical Job metadata', () => {
const projection = createWorkspaceWorkersProjection(); const projection = createWorkspaceWorkersProjection();
const snapshotWorker = worker('runtime-a', 'job-worker', 1); const snapshotWorker = worker('runtime-a', 'job-worker');
snapshotWorker.job = { snapshotWorker.job = {
job_id: 'check:T-1:r1', job_id: 'check:T-1:r1',
attempt_id: 'check:T-1:r1:attempt:1', attempt_id: 'check:T-1:r1:attempt:1',
purpose: 'ticket_item_check', purpose: 'ticket_item_check',
}; };
applyWorkspaceWorkersFrame(projection, { applyWorkspaceWorkersFrame(projection, snapshot('subscription-1', [snapshotWorker]));
protocol_version: 1, assertEquals(projection.workers.get(key('runtime-a', 'job-worker'))?.job, snapshotWorker.job);
frame: 'response',
message: {
result: 'subscribed',
payload: {
request_id: 'request-1',
subscription_id: 'subscription-1',
selector: { topic: 'workspace_workers' },
snapshot_revision: 1,
snapshot: { topic: 'workers', data: { workers: [snapshotWorker] } },
},
},
});
assertEquals(projection.workers.get('runtime-a:job-worker')?.job, snapshotWorker.job);
const updated = { ...snapshotWorker, subject_revision: 2 }; const updated = { ...snapshotWorker };
updated.job = { ...snapshotWorker.job, attempt_id: 'check:T-1:r1:attempt:2' }; updated.job = { ...snapshotWorker.job, attempt_id: 'check:T-1:r1:attempt:2' };
applyWorkspaceWorkersFrame(projection, { applyWorkspaceWorkersFrame(
protocol_version: 1, projection,
frame: 'event', event('subscription-1', { event: 'worker_upserted', data: { worker: updated } }),
message: { );
event: 'event', assertEquals(projection.workers.get(key('runtime-a', 'job-worker'))?.job, updated.job);
data: {
subscription_id: 'subscription-1',
subject_revision: 2,
payload: { event: 'worker_upserted', data: { worker: updated } },
},
},
});
assertEquals(projection.workers.get('runtime-a:job-worker')?.job, updated.job);
}); });
Deno.test('workspace Worker reducer ignores stale events and removes composite subject', () => { Deno.test('fresh subscription snapshot replaces old state and fences delayed old events', () => {
const projection = createWorkspaceWorkersProjection(); const projection = createWorkspaceWorkersProjection();
projection.workers.set('runtime-a:1', worker('runtime-a', '1', 3)); applyWorkspaceWorkersFrame(
projection.revisions.set('runtime-a:1', 3); projection,
snapshot('subscription-old', [worker('runtime-a', 'old'), worker('runtime-b', 'same')]),
);
applyWorkspaceWorkersFrame(
projection,
event('subscription-old', {
event: 'worker_upserted',
data: { worker: worker('runtime-a', 'before-reconnect') },
}),
);
applyWorkspaceWorkersFrame(projection, { applyWorkspaceWorkersFrame(
protocol_version: 1, projection,
frame: 'event', snapshot('subscription-new', [worker('runtime-a', 'current')]),
message: { );
event: 'event', applyWorkspaceWorkersFrame(
data: { projection,
subscription_id: 'subscription-1', event('subscription-old', {
subject_revision: 2, event: 'worker_upserted',
payload: { event: 'worker_upserted', data: { worker: worker('runtime-a', '1', 2) } }, data: { worker: worker('runtime-a', 'late-old-event') },
}, }),
}, );
});
assertEquals(projection.revisions.get('runtime-a:1'), 3);
applyWorkspaceWorkersFrame(projection, { assertEquals(projection.subscriptionId, 'subscription-new');
protocol_version: 1, assertEquals([...projection.workers.keys()], [key('runtime-a', 'current')]);
frame: 'event', });
message: {
event: 'event', Deno.test('workspace Worker reducer applies ordered updates and composite removal', () => {
data: { const projection = createWorkspaceWorkersProjection();
subscription_id: 'subscription-1', applyWorkspaceWorkersFrame(
subject_revision: 4, projection,
payload: { snapshot('subscription-1', [worker('runtime-a', '1'), worker('runtime-b', '1')]),
);
const updated = worker('runtime-a', '1');
updated.state = 'running';
applyWorkspaceWorkersFrame(
projection,
event('subscription-1', { event: 'worker_upserted', data: { worker: updated } }),
);
applyWorkspaceWorkersFrame(
projection,
event('subscription-1', {
event: 'worker_removed', event: 'worker_removed',
data: { worker_id: '1', runtime_id: 'runtime-a' }, data: { worker_id: '1', runtime_id: 'runtime-a' },
}, }),
}, );
},
}); assertEquals([...projection.workers.keys()], [key('runtime-b', '1')]);
assertEquals(projection.workers.size, 0); });
assertEquals(projection.revisions.get('runtime-a:1'), 4);
Deno.test('workspace Worker event before a snapshot is rejected', () => {
const projection = createWorkspaceWorkersProjection();
assertThrows(
() =>
applyWorkspaceWorkersFrame(
projection,
event('subscription-1', {
event: 'worker_upserted',
data: { worker: worker('runtime-a', '1') },
}),
),
'event-before-snapshot must be rejected',
);
}); });
Deno.test('fatal child stop replaces the running-child sidebar projection', () => { Deno.test('fatal child stop replaces the running-child sidebar projection', () => {
const projection = createWorkspaceWorkersProjection(); const projection = createWorkspaceWorkersProjection();
projection.workers.set('runtime-a:1', worker('runtime-a', '1', 1, true)); applyWorkspaceWorkersFrame(
projection.revisions.set('runtime-a:1', 1); projection,
snapshot('subscription-1', [worker('runtime-a', '1', true)]),
applyWorkspaceWorkersFrame(projection, { );
protocol_version: 1, applyWorkspaceWorkersFrame(
frame: 'event', projection,
message: { event('subscription-1', {
event: 'event',
data: {
subscription_id: 'subscription-1',
subject_revision: 2,
payload: {
event: 'worker_upserted', event: 'worker_upserted',
data: { worker: worker('runtime-a', '1', 2, false) }, data: { worker: worker('runtime-a', '1', false) },
}, }),
}, );
},
});
assertEquals(projection.workers.get('runtime-a:1')?.has_running_internal_workers, false); assertEquals(
assertEquals(projection.revisions.get('runtime-a:1'), 2); projection.workers.get(key('runtime-a', '1'))?.has_running_internal_workers,
false,
);
}); });