diff --git a/crates/workspace-server/src/latest_schema.sql b/crates/workspace-server/src/latest_schema.sql index 5b3c7f27a..7cd9de68c 100644 --- a/crates/workspace-server/src/latest_schema.sql +++ b/crates/workspace-server/src/latest_schema.sql @@ -798,6 +798,13 @@ 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 ( diagnostic_id INTEGER PRIMARY KEY AUTOINCREMENT, workspace_id TEXT NOT NULL, diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index 1c572944c..cdbe81cf2 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -68,6 +68,13 @@ CREATE TABLE IF NOT EXISTS worker_registry_projection_revisions ( workspace_id TEXT PRIMARY KEY, revision INTEGER NOT NULL ); +CREATE TABLE IF NOT EXISTS 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 IF NOT EXISTS worker_registry_projection_diagnostics ( diagnostic_id INTEGER PRIMARY KEY AUTOINCREMENT, workspace_id TEXT NOT NULL, @@ -5589,6 +5596,10 @@ impl ControlPlaneStore for SqliteWorkspaceStore { &record.worker.worker_id, &record.created_at, )?; + tx.execute( + "DELETE FROM worker_registry_projection_removals WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3", + params![record.workspace_id, record.worker.runtime_id, record.worker.worker_id], + )?; tx.commit()?; Ok(()) }) @@ -5698,7 +5709,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore { worker_projection_cursor(&tx, workspace_id, runtime_id)?; if connection_generation < current_generation || (connection_generation == current_generation - && snapshot_revision <= current_snapshot) + && snapshot_revision < current_snapshot) { let revision = current_worker_projection_revision(&tx, workspace_id)?; tx.commit()?; @@ -5903,8 +5914,26 @@ impl ControlPlaneStore for SqliteWorkspaceStore { ) -> Result { self.with_conn_mut(|conn| { let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?; + let published_revision = tx + .query_row( + "SELECT projection_revision FROM worker_registry_projection_removals WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3", + params![workspace_id, worker.runtime_id, worker.worker_id], + |row| row.get::<_, i64>(0), + ) + .optional()?; + if let Some(revision) = published_revision { + tx.commit()?; + return Ok(WorkerRegistryProjectionCommit { + revision: revision.max(0) as u64, + changed_workers: Vec::new(), + }); + } let changed_workers = vec![worker.clone()]; let revision = commit_worker_projection_changes(&tx, workspace_id, &changed_workers)?; + tx.execute( + "INSERT INTO worker_registry_projection_removals (workspace_id, runtime_id, worker_id, projection_revision) VALUES (?1, ?2, ?3, ?4)", + params![workspace_id, worker.runtime_id, worker.worker_id, revision as i64], + )?; tx.commit()?; Ok(WorkerRegistryProjectionCommit { revision, @@ -8714,11 +8743,15 @@ fn upsert_worker_observation( params![workspace_id, runtime_id, worker.worker_id.as_str()], |row| Ok((row.get::<_, String>(0)?, row.get::<_, Option>(1)?, row.get::<_, i64>(2)?, row.get::<_, i64>(3)?)), ).optional()?; - if let Some((_, _, generation, revision)) = existing.as_ref() { + if let Some((old_availability, _, generation, revision)) = existing.as_ref() { let generation = (*generation).max(0) as u64; let revision = (*revision).max(0) as u64; + let restores_unavailable = old_availability == "unavailable" + && availability == SubscriptionWorkerAvailability::Observed; if connection_generation < generation - || (connection_generation == generation && worker.subject_revision <= revision) + || (connection_generation == generation + && (worker.subject_revision < revision + || (worker.subject_revision == revision && !restores_unavailable))) { return Ok(false); } @@ -14946,6 +14979,31 @@ INSERT INTO worker_registry ( .reconcile_worker_registry_snapshot("local-dev", "embedded", 1, 11, &[], "4") .unwrap(); assert!(duplicate.changed_workers.is_empty()); + + let reconnected = store + .reconcile_worker_registry_snapshot( + "local-dev", + "embedded", + 1, + 11, + &[observed.clone()], + "5", + ) + .unwrap(); + assert_eq!(reconnected.changed_workers, vec![catalog.worker.clone()]); + let projection = store + .worker_registry_projection("local-dev", &catalog.worker) + .unwrap() + .unwrap(); + assert_eq!( + projection.observation.unwrap().availability, + SubscriptionWorkerAvailability::Observed, + "an equal-revision reconnect snapshot must restore availability" + ); + let duplicate_reconnect = store + .reconcile_worker_registry_snapshot("local-dev", "embedded", 1, 11, &[observed], "6") + .unwrap(); + assert!(duplicate_reconnect.changed_workers.is_empty()); } #[tokio::test] @@ -14984,6 +15042,15 @@ INSERT INTO worker_registry ( .delete_worker_registry("local-dev", &worker_ref) .unwrap() ); + let first_removal = store + .publish_worker_registry_removal("local-dev", &worker_ref) + .unwrap(); + assert_eq!(first_removal.changed_workers, vec![worker_ref.clone()]); + let duplicate_removal = store + .publish_worker_registry_removal("local-dev", &worker_ref) + .unwrap(); + assert!(duplicate_removal.changed_workers.is_empty()); + assert_eq!(duplicate_removal.revision, first_removal.revision); let event = SubscriptionWorker { worker_id: SubscriptionWorkerId::new("removed").unwrap(), runtime_id: Some("embedded".to_string()), diff --git a/crates/workspace-server/src/worker_projection.rs b/crates/workspace-server/src/worker_projection.rs index 928e3f071..36b01a001 100644 --- a/crates/workspace-server/src/worker_projection.rs +++ b/crates/workspace-server/src/worker_projection.rs @@ -155,6 +155,9 @@ impl WorkerProjectionService { let commit = self .store .publish_worker_registry_removal(self.workspace_id.as_str(), &worker)?; + if commit.changed_workers.is_empty() { + return Ok(()); + } let _ = self.events.send(WorkerProjectionEvent { revision: commit.revision, changes: vec![WorkerProjectionChange::Removed(worker)], diff --git a/crates/workspace-server/src/workspace_deletion.rs b/crates/workspace-server/src/workspace_deletion.rs index e0d9b91b9..7cf98789f 100644 --- a/crates/workspace-server/src/workspace_deletion.rs +++ b/crates/workspace-server/src/workspace_deletion.rs @@ -70,6 +70,7 @@ const WORKSPACE_DELETION_PURGE_TABLES: &[&str] = &[ "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", diff --git a/crates/workspace-server/src/workspace_subscription.rs b/crates/workspace-server/src/workspace_subscription.rs index 0ec7541fe..f84c79df6 100644 --- a/crates/workspace-server/src/workspace_subscription.rs +++ b/crates/workspace-server/src/workspace_subscription.rs @@ -314,7 +314,7 @@ async fn run_workspace_workers( }; let mut workers = records .into_iter() - .filter_map(project_registry_worker) + .filter_map(|record| project_registry_worker(&api, record)) .collect::>(); sort_workers(&mut workers); if send_frame( @@ -342,7 +342,7 @@ async fn run_workspace_workers( for change in event.changes { let payload = match change { crate::worker_projection::WorkerProjectionChange::Upsert(record) => { - let Some(mut worker) = project_registry_worker(record) else { + let Some(mut worker) = project_registry_worker(&api, record) else { continue; }; worker.subject_revision = event.revision; @@ -388,6 +388,7 @@ async fn run_workspace_workers( } fn project_registry_worker( + api: &WorkspaceApi, record: crate::store::WorkerRegistryProjectionRecord, ) -> Option { let runtime_id = record.registry.worker.runtime_id.clone(); @@ -422,9 +423,27 @@ fn project_registry_worker( worker.workspace_id = Some(record.registry.workspace_id); worker.display_name = Some(record.registry.display_name); worker.profile = record.registry.profile; + if !project_repository_key(api, &mut worker) { + return None; + } Some(worker) } +fn project_repository_key(api: &WorkspaceApi, worker: &mut SubscriptionWorker) -> bool { + let Some(repository_id) = worker.repository_id.take() else { + worker.repository_key = None; + return true; + }; + let Ok(Some(repository)) = api + .store + .get_repository(&api.config.workspace_id, &repository_id) + else { + return false; + }; + worker.repository_key = Some(repository.repository_key); + true +} + fn sort_workers(workers: &mut [SubscriptionWorker]) { workers.sort_by(|left, right| { left.runtime_id