Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

Replace Worker list revisions with ordered subscriptions

Browse Source
This commit is contained in:
Hare
2026-10-02 10:41:02 +09:00
parent 8f8c2676b4
commit 2c329072a2
24 changed files with 1320 additions and 963 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/workspace-server/src/latest_schema.sql
-16
View File
@@ -914,31 +914,15 @@ CREATE TABLE worker_registry_observations (
worker_id TEXT NOT NULL,
availability TEXT NOT NULL CHECK (availability IN ('observed', 'unavailable')),
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,
PRIMARY KEY (workspace_id, runtime_id, worker_id),
FOREIGN KEY (workspace_id, worker_id)
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 (
workspace_id TEXT NOT NULL,
runtime_id TEXT NOT NULL,
worker_id TEXT NOT NULL,
projection_revision INTEGER NOT NULL,
PRIMARY KEY (workspace_id, runtime_id, worker_id)
);
CREATE TABLE worker_registry_projection_diagnostics (
crates/workspace-server/src/runtime_subscription.rs
+185 -130
View File
@@ -33,26 +33,19 @@ pub enum RuntimeSubscriptionBrokerError {
#[derive(Clone, Debug)]
pub enum BrokerSubscriptionEvent {
Snapshot {
connection_generation: u64,
snapshot_revision: u64,
snapshot: SubscriptionSnapshot,
},
Event {
connection_generation: u64,
subject_revision: u64,
payload: SubscriptionEventPayload,
},
Disconnected {
connection_generation: u64,
message: String,
},
Rejected {
connection_generation: u64,
code: SubscriptionRejectionCode,
message: String,
},
Closed {
connection_generation: u64,
code: SubscriptionTerminationCode,
message: String,
},
@@ -60,7 +53,6 @@ pub enum BrokerSubscriptionEvent {
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct RuntimeSubscriptionBrokerStatus {
pub connection_generation: u64,
pub connected: bool,
pub desired_selectors: usize,
pub upstream_subscriptions: usize,
@@ -82,7 +74,30 @@ impl BrokerSubscription {
&self.selector
}
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 {
let generation = self.next_generation.fetch_add(1, Ordering::Relaxed);
let (commands, receiver) = mpsc::unbounded_channel();
let status = Arc::new(RwLock::new(RuntimeSubscriptionBrokerStatus {
connection_generation: generation,
..Default::default()
}));
let status = Arc::new(RwLock::new(RuntimeSubscriptionBrokerStatus::default()));
let registration = Registration {
generation,
commands: commands.clone(),
@@ -156,7 +168,6 @@ impl RuntimeSubscriptionBroker {
let generation = self.next_generation.fetch_add(1, Ordering::Relaxed);
let (commands, receiver) = mpsc::unbounded_channel();
let status = Arc::new(RwLock::new(RuntimeSubscriptionBrokerStatus {
connection_generation: generation,
connected: true,
..Default::default()
}));
@@ -257,7 +268,6 @@ impl RuntimeSubscriptionBroker {
.clone();
if !initial_status.connected {
let _ = initial_events.try_send(BrokerSubscriptionEvent::Disconnected {
connection_generation: initial_status.connection_generation,
message: "Runtime subscription connection is not currently available".to_string(),
});
}
@@ -286,8 +296,7 @@ struct SelectorState {
downstreams: HashMap<u64, mpsc::Sender<BrokerSubscriptionEvent>>,
upstream_id: Option<SubscriptionId>,
pending: bool,
snapshot: Option<(u64, SubscriptionSnapshot)>,
revisions: HashMap<String, u64>,
snapshot: Option<SubscriptionSnapshot>,
}
impl SelectorState {
fn new() -> Self {
@@ -296,7 +305,6 @@ impl SelectorState {
upstream_id: None,
pending: false,
snapshot: None,
revisions: HashMap::new(),
}
}
}
@@ -334,11 +342,9 @@ impl State {
selector.upstream_id = None;
selector.pending = false;
selector.snapshot = None;
selector.revisions.clear();
broadcast(
&mut selector.downstreams,
BrokerSubscriptionEvent::Disconnected {
connection_generation: self.generation,
message: message.clone(),
},
);
@@ -347,12 +353,77 @@ impl State {
}
struct EmbeddedEntry {
lifetime: u64,
downstreams: HashMap<u64, mpsc::Sender<BrokerSubscriptionEvent>>,
snapshot_revision: u64,
snapshot: SubscriptionSnapshot,
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(
runtime_id: String,
runtime: worker_runtime::Runtime,
@@ -364,6 +435,7 @@ async fn run_embedded_connection(
let (updates, mut update_receiver) = mpsc::unbounded_channel();
let mut entries = HashMap::<EventSubscriptionSelector, EmbeddedEntry>::new();
let mut downstream_index = HashMap::<u64, EventSubscriptionSelector>::new();
let mut next_lifetime = 1_u64;
loop {
tokio::select! {
command = commands.recv() => match command {
@@ -371,30 +443,38 @@ async fn run_embedded_connection(
downstream_index.insert(downstream_id, selector.clone());
if let Some(entry) = entries.get_mut(&selector) {
let _ = events.try_send(BrokerSubscriptionEvent::Snapshot {
connection_generation: generation,
snapshot_revision: entry.snapshot_revision,
snapshot: entry.snapshot.clone(),
});
entry.downstreams.insert(downstream_id, events);
} else {
match runtime.subscribe_event_selector(selector.clone()) {
Ok(mut subscription) => {
let snapshot_revision = subscription.snapshot_revision();
let snapshot = project_snapshot_runtime(subscription.snapshot().clone(), &runtime_id);
let _ = events.try_send(BrokerSubscriptionEvent::Snapshot { connection_generation: generation, snapshot_revision, snapshot: snapshot.clone() });
let sender = updates.clone();
let task_selector = selector.clone();
let task_runtime_id = runtime_id.clone();
let task = tokio::spawn(async move {
while let Ok(update) = subscription.recv().await {
let payload = project_payload_runtime(update.payload, &task_runtime_id);
if sender.send((task_selector.clone(), update.subject_revision, payload)).is_err() { break; }
}
let lifetime = next_lifetime;
next_lifetime = next_lifetime.saturating_add(1);
match start_embedded_subscription(
&runtime,
&runtime_id,
selector.clone(),
lifetime,
updates.clone(),
) {
Ok((snapshot, task)) => {
let _ = events.try_send(BrokerSubscriptionEvent::Snapshot {
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) => {
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)) => {
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();
}
return;
@@ -415,16 +495,68 @@ async fn run_embedded_connection(
None => return,
},
update = update_receiver.recv() => {
let Some((selector, subject_revision, payload)) = update else { return; };
if let Some(entry) = entries.get_mut(&selector) {
entry.snapshot_revision = entry.snapshot_revision.saturating_add(1);
apply_event_to_cached_snapshot(&mut entry.snapshot, &payload);
broadcast(&mut entry.downstreams, BrokerSubscriptionEvent::Event { connection_generation: generation, subject_revision, payload });
let Some(update) = update else { return; };
match update {
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);
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 {
connection_generation: generation,
connected: true,
desired_selectors: entries.len(),
upstream_subscriptions: entries.len(),
@@ -597,10 +729,8 @@ async fn apply_online(
.selectors
.entry(selector.clone())
.or_insert_with(SelectorState::new);
if let Some((revision, snapshot)) = &entry.snapshot {
if let Some(snapshot) = &entry.snapshot {
let _ = events.try_send(BrokerSubscriptionEvent::Snapshot {
connection_generation: state.generation,
snapshot_revision: *revision,
snapshot: snapshot.clone(),
});
}
@@ -696,7 +826,6 @@ async fn handle_frame(
request_id,
subscription_id,
selector,
snapshot_revision,
snapshot,
}) => {
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(())?;
entry.pending = false;
entry.upstream_id = Some(subscription_id.clone());
entry.snapshot = Some((snapshot_revision, snapshot.clone()));
entry.revisions = snapshot_revisions(&snapshot);
entry.snapshot = Some(snapshot.clone());
state
.upstream_index
.insert(subscription_id, selector.clone());
broadcast(
&mut entry.downstreams,
BrokerSubscriptionEvent::Snapshot {
connection_generation: state.generation,
snapshot_revision,
snapshot,
},
BrokerSubscriptionEvent::Snapshot { snapshot },
);
if entry.downstreams.is_empty() {
maybe_unsubscribe(socket, state, selector).await?;
@@ -734,18 +858,13 @@ async fn handle_frame(
if let Some(mut entry) = state.selectors.remove(&selector) {
broadcast(
&mut entry.downstreams,
BrokerSubscriptionEvent::Rejected {
connection_generation: state.generation,
code,
message,
},
BrokerSubscriptionEvent::Rejected { code, message },
);
}
}
}
SubscriptionFramePayload::Event(SubscriptionEvent::Event {
subscription_id,
subject_revision,
payload,
}) => {
let selector = state
@@ -756,24 +875,12 @@ async fn handle_frame(
payload.validate_for_selector(&selector).map_err(|_| ())?;
let payload = project_payload_runtime(payload, &state.runtime_id);
let entry = state.selectors.get_mut(&selector).ok_or(())?;
if let Some(subject) = event_subject(&payload) {
let revision = entry.revisions.entry(subject).or_insert(0);
if subject_revision <= *revision {
return Ok(());
}
*revision = subject_revision;
}
if let Some((snapshot_revision, snapshot)) = entry.snapshot.as_mut() {
*snapshot_revision = snapshot_revision.saturating_add(1);
if let Some(snapshot) = entry.snapshot.as_mut() {
apply_event_to_cached_snapshot(snapshot, &payload);
}
broadcast(
&mut entry.downstreams,
BrokerSubscriptionEvent::Event {
connection_generation: state.generation,
subject_revision,
payload,
},
BrokerSubscriptionEvent::Event { payload },
);
}
SubscriptionFramePayload::Event(SubscriptionEvent::SubscriptionClosed {
@@ -785,14 +892,9 @@ async fn handle_frame(
let should_resubscribe = if let Some(entry) = state.selectors.get_mut(&selector) {
entry.upstream_id = None;
entry.snapshot = None;
entry.revisions.clear();
broadcast(
&mut entry.downstreams,
BrokerSubscriptionEvent::Closed {
connection_generation: state.generation,
code,
message,
},
BrokerSubscriptionEvent::Closed { code, message },
);
!entry.downstreams.is_empty()
} else {
@@ -822,10 +924,9 @@ fn close_all(state: &mut State, replacement: u64) {
broadcast(
&mut entry.downstreams,
BrokerSubscriptionEvent::Closed {
connection_generation: state.generation,
code: SubscriptionTerminationCode::ServerShutdown,
message: format!(
"Runtime connection generation {} was fenced by generation {replacement}",
"Runtime registration {} was replaced by {replacement}",
state.generation
),
},
@@ -846,51 +947,6 @@ fn broadcast(
downstreams.retain(|id, _| !closed.contains(id));
}
fn snapshot_revisions(snapshot: &SubscriptionSnapshot) -> HashMap<String, u64> {
match snapshot {
SubscriptionSnapshot::Workers { workers } => workers
.iter()
.map(|worker| {
(
format!(
"{}:{}",
worker.runtime_id.as_deref().unwrap_or_default(),
worker.worker_id
),
worker.subject_revision,
)
})
.collect(),
SubscriptionSnapshot::WorkerProtocol { worker_id, .. } => {
HashMap::from([(worker_id.to_string(), 0)])
}
SubscriptionSnapshot::WorkspaceWorkdirs { .. } => HashMap::new(),
}
}
fn event_subject(payload: &SubscriptionEventPayload) -> Option<String> {
Some(match payload {
SubscriptionEventPayload::WorkerUpserted { worker } => format!(
"{}:{}",
worker.runtime_id.as_deref().unwrap_or_default(),
worker.worker_id
),
SubscriptionEventPayload::WorkerRemoved {
worker_id,
runtime_id,
} => format!(
"{}:{}",
runtime_id.as_deref().unwrap_or_default(),
worker_id
),
SubscriptionEventPayload::WorkerProtocol { worker_id, .. } => worker_id.to_string(),
SubscriptionEventPayload::WorkdirUpserted { workdir } => {
workdir.working_directory_id.to_string()
}
SubscriptionEventPayload::WorkdirRemoved {
working_directory_id,
} => working_directory_id.to_string(),
})
}
async fn send_frame(socket: &mut RuntimeSocket, frame: SubscriptionFrame) -> Result<(), ()> {
frame.validate().map_err(|_| ())?;
socket
@@ -979,7 +1035,6 @@ fn runtime_token(
}
fn update_status(status: &RwLock<RuntimeSubscriptionBrokerStatus>, state: &State, connected: bool) {
*status.write().expect("broker status poisoned") = RuntimeSubscriptionBrokerStatus {
connection_generation: state.generation,
connected,
desired_selectors: state
.selectors
crates/workspace-server/src/runtime_subscription_tests.rs
+71 -9
View File
@@ -264,16 +264,9 @@ async fn replacing_runtime_registration_fences_the_old_generation() {
assert!(second_generation > first_generation);
assert!(matches!(
next_event(&mut old).await,
BrokerSubscriptionEvent::Closed {
connection_generation,
..
} if connection_generation == first_generation
BrokerSubscriptionEvent::Closed { .. }
));
let status = wait_for_status(&broker, "runtime-test", |status| {
status.connection_generation == second_generation && status.connected
})
.await;
assert_eq!(status.connection_generation, second_generation);
wait_for_status(&broker, "runtime-test", |status| status.connected).await;
server.abort();
}
@@ -430,6 +423,75 @@ async fn embedded_runtime_uses_in_process_subscription_source() {
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]
async fn strict_runtime_subscription_rejects_private_endpoint_before_websocket_connect() {
let config = RemoteRuntimeConfig::new(
crates/workspace-server/src/server.rs
+85 -91
View File
@@ -1501,10 +1501,11 @@ impl WorkerRemovalService {
/// and WorkerRemove use the retention commit below; compensation preserves
/// its pre-create semantics while sharing the atomic catalog contract.
fn commit_catalog_removal(&self, target: &RuntimeWorkerRef) -> crate::Result<()> {
let commit = self
.store
.delete_worker_registry(&self.workspace_id, target)?;
self.worker_projection.publish_commit(commit)
self.worker_projection.publish_ordered(|| {
self.store
.delete_worker_registry(&self.workspace_id, target)
.map(|commit| ((), commit))
})
}
async fn execute_cleanup_removal(&self, candidate: &CleanupWorkerCandidate) -> ApiResult<()> {
@@ -1631,18 +1632,17 @@ impl WorkerRemovalService {
));
}
};
match self.store.commit_worker_removal(
&self.workspace_id,
&prepared.plan.operation_id,
&prepared.plan.input_fingerprint,
&result,
) {
Ok(commit) => {
self.worker_projection
.publish_commit(commit.catalog)
.map_err(|error| error.to_string())?;
Ok(worker_remove_success_response(target))
}
match self.worker_projection.publish_ordered(|| {
self.store
.commit_worker_removal(
&self.workspace_id,
&prepared.plan.operation_id,
&prepared.plan.input_fingerprint,
&result,
)
.map(|commit| ((), commit.catalog))
}) {
Ok(()) => Ok(worker_remove_success_response(target)),
Err(error) => Ok(worker_retention_error_response(error)),
}
}
@@ -1764,12 +1764,12 @@ impl WorkerRemovalService {
.map_err(|_| "Worker removal recovery authority is unavailable".to_string())?;
if let Some(prepared) = prepared {
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
.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())?;
return Ok(worker_remove_success_response(target));
}
@@ -1947,17 +1947,17 @@ impl WorkerRemovalService {
));
}
};
match self.store.commit_worker_removal(
&self.workspace_id,
&plan.operation_id,
&plan.input_fingerprint,
&retention_result,
) {
Ok(commit) => {
self.worker_projection
.publish_commit(commit.catalog)
.map_err(|error| error.to_string())?;
}
match self.worker_projection.publish_ordered(|| {
self.store
.commit_worker_removal(
&self.workspace_id,
&plan.operation_id,
&plan.input_fingerprint,
&retention_result,
)
.map(|commit| ((), commit.catalog))
}) {
Ok(()) => {}
Err(error) => {
let _ = self.store.fail_worker_removal(
&self.workspace_id,
@@ -3454,17 +3454,19 @@ impl WorkspaceApi {
};
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 (commit, newly_bound) = self.store.bind_backend_job_attempt_worker(
&self.config.workspace_id,
&request.job_id,
&reservation.attempt.attempt_id,
&worker.worker,
Some(&runtime_run_id),
&bound_at,
)?;
if newly_bound {
self.worker_projection.publish_commit(commit)?;
} else {
let newly_bound = self.worker_projection.publish_ordered(|| {
self.store
.bind_backend_job_attempt_worker(
&self.config.workspace_id,
&request.job_id,
&reservation.attempt.attempt_id,
&worker.worker,
Some(&runtime_run_id),
&bound_at,
)
.map(|(commit, newly_bound)| (newly_bound, commit))
})?;
if !newly_bound {
let job = self
.store
.get_backend_job(&self.config.workspace_id, &request.job_id)?
@@ -19435,10 +19437,11 @@ async fn scoped_start_workspace_orchestrator(
deleted.diagnostics,
));
}
let removal = api
.store
.delete_worker_registry(&api.config.workspace_id, &existing.worker)?;
api.worker_projection.publish_commit(removal)?;
api.worker_projection.publish_ordered(|| {
api.store
.delete_worker_registry(&api.config.workspace_id, &existing.worker)
.map(|commit| ((), commit))
})?;
disposition = "recreated";
}
@@ -21768,22 +21771,23 @@ async fn set_worker_retention(
}
}
let retention_state = if pinned { "pinned" } else { "normal" };
let commit = api.store.update_worker_retention(
&api.config.workspace_id,
&worker_ref,
retention_state,
now_registry_timestamp().as_str(),
)?;
if commit.changes.is_empty() {
let changed = api.worker_projection.publish_ordered(|| {
api.store
.update_worker_retention(
&api.config.workspace_id,
&worker_ref,
retention_state,
now_registry_timestamp().as_str(),
)
.map(|commit| (!commit.changes.is_empty(), commit))
})?;
if !changed {
return Err(cleanup_api_error(
runtime_id.as_str(),
"workspace_worker_retention_unknown_worker",
"Worker is not known to the Backend registry",
));
}
api.worker_projection
.publish_commit(commit)
.map_err(ApiError::from)?;
Ok(Json(WorkerRetentionResponse {
workspace_id: api.config.workspace_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> {
let limit = api.config.max_records.min(200);
let (_, projections) = api
let projections = api
.store
.worker_registry_projection_snapshot(&api.config.workspace_id, limit)?;
let workdir_records = api
@@ -27478,9 +27482,12 @@ fn record_worker_summary(
created_at: timestamp.clone(),
updated_at: timestamp,
};
let catalog_commit = api.store.upsert_worker_registry(&record)?;
api.worker_projection
.publish_commit(catalog_commit)
.publish_ordered(|| {
api.store
.upsert_worker_registry(&record)
.map(|commit| ((), commit))
})
.map_err(ApiError::from)?;
if let Ok(worker_id) =
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()),
resource_key: None,
availability: protocol::subscription::SubscriptionWorkerAvailability::Observed,
subject_revision: 0,
worker_state: worker.worker_state.clone(),
state,
has_running_internal_workers: false,
@@ -29427,7 +29433,6 @@ mod tests {
runtime_id: Some("runtime-a".to_string()),
resource_key: Some("W-1".to_string()),
availability,
subject_revision: 4,
worker_state: Some(protocol::WorkerStateSnapshot::initial()),
state: protocol::subscription::SubscriptionWorkerState::Idle,
has_running_internal_workers: false,
@@ -29438,10 +29443,6 @@ mod tests {
workdir_attachments: Vec::new(),
},
availability,
connection_generation: 2,
subject_revision: 4,
snapshot_revision: 7,
projection_revision: 9,
observed_at: "2026-09-24T00:00:01Z".to_string(),
}),
job: None,
@@ -33018,7 +33019,6 @@ mod tests {
runtime_id: Some(EMBEDDED_WORKER_RUNTIME_ID.to_string()),
resource_key: None,
availability: protocol::subscription::SubscriptionWorkerAvailability::Observed,
subject_revision: 1,
worker_state: None,
state: protocol::subscription::SubscriptionWorkerState::Paused,
has_running_internal_workers: false,
@@ -44053,7 +44053,8 @@ mod tests {
let summary = api.runtime.worker(&target).unwrap();
sync_worker_observation(&api, &summary).unwrap();
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)
.execute_async(
@@ -44086,7 +44087,7 @@ mod tests {
.unwrap()
.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 {
let event = subscriber.recv().await.expect("projection event");
if event.changes.iter().any(|change| {
@@ -44098,7 +44099,6 @@ mod tests {
})
.await
.expect("WorkerRemove catalog removal broadcast");
assert!(removal.revision > 0);
}
#[tokio::test]
@@ -44135,7 +44135,8 @@ mod tests {
let worker = spawned.worker.unwrap().worker;
let summary = api.runtime.worker(&worker).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(
&api,
&summary,
@@ -44151,7 +44152,7 @@ mod tests {
.unwrap()
.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 {
let event = subscriber.recv().await.expect("projection event");
if event.changes.iter().any(|change| {
@@ -44163,7 +44164,6 @@ mod tests {
})
.await
.expect("spawn compensation catalog removal broadcast");
assert!(removal.revision > 0);
sync_worker_observation(&api, &summary).unwrap();
assert!(
api.store
@@ -46005,7 +46005,8 @@ mod tests {
.expect("singleton Worker cleanup candidate")
.clone();
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)
.execute_cleanup_removal(&candidate)
@@ -46129,7 +46130,8 @@ mod tests {
.iter()
.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)
.unwrap_or_else(|error| panic!("cleanup plan: {}", error.error));
let candidate = plan
@@ -46173,20 +46175,18 @@ mod tests {
event.changes.as_slice(),
[crate::store::WorkerCatalogChange::Removed(removed)] if removed == &worker
));
let acting_revision = tokio::time::timeout(
tokio::time::timeout(
std::time::Duration::from_secs(1),
next_workspace_worker_removal(&mut acting_socket, worker_id.as_str()),
)
.await
.expect("acting Workspace Worker subscription removal");
let observing_revision = tokio::time::timeout(
tokio::time::timeout(
std::time::Duration::from_secs(1),
next_workspace_worker_removal(&mut observing_socket, worker_id.as_str()),
)
.await
.expect("independent Workspace Worker subscription removal");
assert_eq!(acting_revision, event.revision);
assert_eq!(observing_revision, event.revision);
assert!(
api.store
.get_worker_registry(&api.config.workspace_id, &worker)
@@ -46202,11 +46202,10 @@ mod tests {
recovered.plan.state,
crate::retention::WorkerRemovalPlanState::Succeeded
);
let (snapshot_revision, snapshot) = api
let snapshot = api
.store
.worker_registry_projection_snapshot(&api.config.workspace_id, 100)
.unwrap();
assert_eq!(snapshot_revision, event.revision);
assert!(
snapshot
.iter()
@@ -46228,8 +46227,8 @@ mod tests {
.all(|item| item.worker_id.as_str() != worker_id)
);
// A retry replays the durable fence at the same revision, healing an
// in-process broadcast omitted after the first DB commit.
// A retry replays the durable fence, healing an in-process broadcast
// omitted after the first DB commit.
WorkerRemovalService::new(&api)
.commit_catalog_removal(&worker)
.unwrap();
@@ -46237,7 +46236,6 @@ mod tests {
.await
.expect("catalog removal replay")
.expect("projection replay");
assert_eq!(replay.revision, event.revision);
assert!(matches!(
replay.changes.as_slice(),
[crate::store::WorkerCatalogChange::Removed(removed)] if removed == &worker
@@ -49669,10 +49667,7 @@ mod tests {
}
}
async fn next_workspace_worker_removal(
socket: &mut TestWebSocket,
expected_worker_id: &str,
) -> u64 {
async fn next_workspace_worker_removal(socket: &mut TestWebSocket, expected_worker_id: &str) {
loop {
let Message::Text(text) = socket.next().await.unwrap().unwrap() else {
continue;
@@ -49681,7 +49676,6 @@ mod tests {
serde_json::from_str(text.as_str()).unwrap();
if let protocol::subscription::SubscriptionFramePayload::Event(
protocol::subscription::SubscriptionEvent::Event {
subject_revision,
payload:
protocol::subscription::SubscriptionEventPayload::WorkerRemoved {
worker_id,
@@ -49692,7 +49686,7 @@ mod tests {
) = frame.payload
&& worker_id.as_str() == expected_worker_id
{
return subject_revision;
return;
}
}
}
@@ -49812,7 +49806,7 @@ mod tests {
.unwrap();
assert_eq!(spawned.state, InternalWorkerOperationState::Accepted);
let worker_id = spawned.worker.unwrap().worker.worker_id;
let (_, projections) = api
let projections = api
.store
.worker_registry_projection_snapshot(TEST_WORKSPACE_ID, 10)
.unwrap();
crates/workspace-server/src/store.rs
+344 -326
View File
File diff suppressed because it is too large Load Diff
crates/workspace-server/src/worker_projection.rs
+318 -110
View File
@@ -11,7 +11,10 @@ use tokio::task::JoinHandle;
use crate::runtime_subscription::{
BrokerSubscriptionEvent, RuntimeSubscriptionBroker, RuntimeSubscriptionBrokerError,
};
use crate::store::{ControlPlaneStore, WorkerCatalogChange, WorkerRegistryProjectionCommit};
use crate::store::{
ControlPlaneStore, WorkerCatalogChange, WorkerRegistryProjectionCommit,
WorkerRegistryProjectionRecord,
};
use worker_runtime::identity::RuntimeWorkerRef;
const PROJECTION_EVENT_CAPACITY: usize = 256;
@@ -20,21 +23,42 @@ pub type WorkerProjectionChange = WorkerCatalogChange;
#[derive(Debug, Clone)]
pub struct WorkerProjectionEvent {
pub revision: u64,
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
/// low-frequency publication stream.
///
/// One Runtime subscription is kept alive per registered Runtime even when no
/// browser is connected. Browser subscriptions therefore never become the
/// authority for populating `worker_registry` observations.
/// The ordering gate covers both durable commits and subscription snapshot creation.
/// Every subscription therefore starts with one authoritative snapshot and receives
/// only commits that completed after it, in the same order in which they were stored.
pub struct WorkerProjectionService {
workspace_id: String,
store: Arc<dyn ControlPlaneStore>,
events: broadcast::Sender<WorkerProjectionEvent>,
connection_generation_epoch: u64,
ordering: Mutex<()>,
attach_gate: tokio::sync::Mutex<()>,
tasks: Mutex<HashMap<String, JoinHandle<()>>>,
}
@@ -45,8 +69,8 @@ impl WorkerProjectionService {
workspace_id: workspace_id.into(),
store,
events,
connection_generation_epoch: Utc::now().timestamp_nanos_opt().unwrap_or_default().max(0)
as u64,
ordering: Mutex::new(()),
attach_gate: tokio::sync::Mutex::new(()),
tasks: Mutex::new(HashMap::new()),
})
}
@@ -62,8 +86,38 @@ impl WorkerProjectionService {
}
}
pub fn subscribe(&self) -> broadcast::Receiver<WorkerProjectionEvent> {
self.events.subscribe()
pub fn subscribe_ordered(&self, limit: usize) -> crate::Result<WorkerProjectionSubscription> {
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(
@@ -71,6 +125,20 @@ impl WorkerProjectionService {
broker: &RuntimeSubscriptionBroker,
runtime_id: &str,
) -> 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 =
broker.subscribe(runtime_id, EventSubscriptionSelector::RuntimeWorkers)?;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
@@ -80,12 +148,14 @@ impl WorkerProjectionService {
break;
};
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!(
runtime_id,
error = %error,
"failed to initialize durable Worker projection"
"failed to initialize durable Worker projection; restarting from a fresh snapshot"
);
subscription.restart()?;
continue;
}
if initialized {
break;
@@ -99,26 +169,22 @@ impl WorkerProjectionService {
let Some(service) = service.upgrade() else {
return;
};
if let Err(error) = service
.apply_broker_event(task_runtime_id.as_str(), event)
.await
{
if let Err(error) = service.apply_broker_event(task_runtime_id.as_str(), event) {
tracing::warn!(
runtime_id = %task_runtime_id,
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
.tasks
self.tasks
.lock()
.expect("Worker projection task lock poisoned")
.insert(runtime_id, task)
{
previous.abort();
}
.insert(runtime_id, task);
Ok(())
}
@@ -128,110 +194,96 @@ impl WorkerProjectionService {
worker: &protocol::subscription::SubscriptionWorker,
observed_at: &str,
) -> crate::Result<()> {
let commit = self.store.apply_worker_registry_observation(
self.workspace_id.as_str(),
runtime_id,
0,
worker,
observed_at,
)?;
self.publish_commit(commit)
self.publish_ordered(|| {
self.store
.apply_worker_registry_observation(
self.workspace_id.as_str(),
runtime_id,
worker,
observed_at,
)
.map(|commit| ((), commit))
})
}
pub fn refresh(&self, worker: &RuntimeWorkerRef) -> crate::Result<()> {
let commit = self
.store
.commit_worker_registry_projection_refresh(self.workspace_id.as_str(), worker)?;
self.publish_commit(commit)
self.publish_ordered(|| {
self.store
.commit_worker_registry_projection_refresh(self.workspace_id.as_str(), worker)
.map(|commit| ((), commit))
})
}
async fn apply_broker_event(
fn apply_broker_event(
&self,
runtime_id: &str,
event: BrokerSubscriptionEvent,
) -> crate::Result<()> {
let observed_at = Utc::now().to_rfc3339();
let projection_generation = |connection_generation: u64| {
self.connection_generation_epoch
.saturating_add(connection_generation)
};
let commit = match event {
BrokerSubscriptionEvent::Snapshot {
connection_generation,
snapshot_revision,
snapshot: SubscriptionSnapshot::Workers { workers },
} => self.store.reconcile_worker_registry_snapshot(
self.workspace_id.as_str(),
runtime_id,
projection_generation(connection_generation),
snapshot_revision,
workers.as_slice(),
observed_at.as_str(),
)?,
BrokerSubscriptionEvent::Snapshot { .. } => return Ok(()),
BrokerSubscriptionEvent::Event {
connection_generation,
payload: SubscriptionEventPayload::WorkerUpserted { worker },
..
} => self.store.apply_worker_registry_observation(
self.workspace_id.as_str(),
runtime_id,
projection_generation(connection_generation),
&worker,
observed_at.as_str(),
)?,
BrokerSubscriptionEvent::Event {
connection_generation,
subject_revision,
payload: SubscriptionEventPayload::WorkerRemoved { worker_id, .. },
} => self.store.mark_worker_registry_observation_unavailable(
self.workspace_id.as_str(),
runtime_id,
projection_generation(connection_generation),
Some(worker_id.as_str()),
Some(subject_revision),
observed_at.as_str(),
)?,
BrokerSubscriptionEvent::Event { .. } => return Ok(()),
BrokerSubscriptionEvent::Disconnected {
connection_generation,
..
}
| BrokerSubscriptionEvent::Rejected {
connection_generation,
..
}
| BrokerSubscriptionEvent::Closed {
connection_generation,
..
} => self.store.mark_worker_registry_observation_unavailable(
self.workspace_id.as_str(),
runtime_id,
projection_generation(connection_generation),
None,
None,
observed_at.as_str(),
)?,
};
self.publish_commit(commit)
self.publish_ordered(|| {
let commit = match event {
BrokerSubscriptionEvent::Snapshot {
snapshot: SubscriptionSnapshot::Workers { workers },
} => self.store.reconcile_worker_registry_snapshot(
self.workspace_id.as_str(),
runtime_id,
workers.as_slice(),
observed_at.as_str(),
)?,
BrokerSubscriptionEvent::Snapshot { .. } => {
return Ok((
(),
WorkerRegistryProjectionCommit {
changes: Vec::new(),
},
));
}
BrokerSubscriptionEvent::Event {
payload: SubscriptionEventPayload::WorkerUpserted { worker },
} => self.store.apply_worker_registry_observation(
self.workspace_id.as_str(),
runtime_id,
&worker,
observed_at.as_str(),
)?,
BrokerSubscriptionEvent::Event {
payload: SubscriptionEventPayload::WorkerRemoved { worker_id, .. },
} => self.store.mark_worker_registry_observation_unavailable(
self.workspace_id.as_str(),
runtime_id,
Some(worker_id.as_str()),
observed_at.as_str(),
)?,
BrokerSubscriptionEvent::Event { .. } => {
return Ok((
(),
WorkerRegistryProjectionCommit {
changes: Vec::new(),
},
));
}
BrokerSubscriptionEvent::Disconnected { .. }
| BrokerSubscriptionEvent::Rejected { .. }
| BrokerSubscriptionEvent::Closed { .. } => {
self.store.mark_worker_registry_observation_unavailable(
self.workspace_id.as_str(),
runtime_id,
None,
observed_at.as_str(),
)?
}
};
Ok(((), commit))
})
}
/// Broadcast an already committed catalog payload. Persistence always happens
/// 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<()> {
fn publish_commit(&self, commit: WorkerRegistryProjectionCommit) {
if commit.changes.is_empty() {
return Ok(());
return;
}
let _ = self.events.send(WorkerProjectionEvent {
revision: commit.revision,
changes: commit.changes,
});
Ok(())
}
}
@@ -240,3 +292,159 @@ impl Drop for WorkerProjectionService {
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());
}
}
crates/workspace-server/src/workspace_deletion.rs
-2
View File
@@ -73,10 +73,8 @@ const WORKSPACE_DELETION_PURGE_TABLES: &[&str] = &[
"worker_orphan_diagnostics",
"worker_registry",
"worker_registry_observations",
"worker_registry_projection_cursors",
"worker_registry_projection_diagnostics",
"worker_registry_projection_removals",
"worker_registry_projection_revisions",
"worker_removal_operations",
"worker_retention_audit_events",
"worker_session_archives",
crates/workspace-server/src/workspace_subscription.rs
+9 -22
View File
@@ -271,7 +271,6 @@ async fn run_worker_protocol(
worker_id: worker_id.clone(),
runtime_id: Some(runtime_id),
},
snapshot_revision: 0,
snapshot: SubscriptionSnapshot::WorkerProtocol {
worker_id: worker_id.clone(),
events: vec![initial_event],
@@ -284,13 +283,10 @@ async fn run_worker_protocol(
{
return;
}
let mut subject_revision = 0_u64;
while let Some(event) = events.recv().await {
subject_revision = subject_revision.saturating_add(1);
let frame =
SubscriptionFrame::new(SubscriptionFramePayload::Event(SubscriptionEvent::Event {
subscription_id: subscription_id.clone(),
subject_revision,
payload: SubscriptionEventPayload::WorkerProtocol {
worker_id: worker_id.clone(),
event,
@@ -333,12 +329,10 @@ async fn run_workspace_workers(
subscription_id: SubscriptionId,
outbound: mpsc::Sender<WsMessage>,
) {
// Subscribe before reading so commits racing with the snapshot are replayed.
let mut events = api.worker_projection.subscribe();
let Ok((snapshot_revision, records)) = api.store.worker_registry_projection_snapshot(
&api.config.workspace_id,
api.config.max_records.max(1),
) else {
let Ok(mut projection) = api
.worker_projection
.subscribe_ordered(api.config.max_records.max(1))
else {
let _ = send_rejected(
&outbound,
request_id,
@@ -348,7 +342,8 @@ async fn run_workspace_workers(
.await;
return;
};
let mut workers = records
let mut workers = projection
.take_snapshot()
.into_iter()
.filter_map(|record| project_registry_worker(&api, record))
.collect::<Vec<_>>();
@@ -360,7 +355,6 @@ async fn run_workspace_workers(
request_id,
subscription_id: subscription_id.clone(),
selector: EventSubscriptionSelector::WorkspaceWorkers,
snapshot_revision,
snapshot: SubscriptionSnapshot::Workers { workers },
},
)),
@@ -372,16 +366,14 @@ async fn run_workspace_workers(
}
loop {
match events.recv().await {
Ok(event) if event.revision <= snapshot_revision => continue,
match projection.recv().await {
Ok(event) => {
for change in event.changes {
let payload = match change {
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;
};
worker.subject_revision = event.revision;
SubscriptionEventPayload::WorkerUpserted { 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
.is_err()
{
@@ -431,7 +423,6 @@ fn project_registry_worker(
let mut worker = if let Some(observation) = record.observation {
let mut worker = observation.worker;
worker.availability = observation.availability;
worker.subject_revision = observation.projection_revision;
worker
} else {
SubscriptionWorker {
@@ -442,7 +433,6 @@ fn project_registry_worker(
runtime_id: Some(runtime_id.clone()),
resource_key: record.resource_key.clone(),
availability: protocol::subscription::SubscriptionWorkerAvailability::Unavailable,
subject_revision: 0,
worker_state: None,
state: protocol::subscription::SubscriptionWorkerState::Stopped,
has_running_internal_workers: false,
@@ -559,14 +549,12 @@ fn sort_workers(workers: &mut [SubscriptionWorker]) {
async fn send_event(
outbound: &mpsc::Sender<WsMessage>,
subscription_id: &SubscriptionId,
subject_revision: u64,
payload: SubscriptionEventPayload,
) -> Result<(), ()> {
send_frame(
outbound,
SubscriptionFrame::new(SubscriptionFramePayload::Event(SubscriptionEvent::Event {
subscription_id: subscription_id.clone(),
subject_revision,
payload,
})),
)
@@ -617,7 +605,6 @@ mod tests {
runtime_id: Some("runtime-1".to_string()),
resource_key: Some("W-1".to_string()),
availability: protocol::subscription::SubscriptionWorkerAvailability::Observed,
subject_revision: 1,
worker_state: None,
state: protocol::subscription::SubscriptionWorkerState::Idle,
has_running_internal_workers: false,