diff --git a/crates/workspace-server/src/hosts.rs b/crates/workspace-server/src/hosts.rs index 386b92b1..dbbe3ba3 100644 --- a/crates/workspace-server/src/hosts.rs +++ b/crates/workspace-server/src/hosts.rs @@ -1427,6 +1427,10 @@ impl EmbeddedWorkerRuntime { self } + pub(crate) fn subscription_runtime(&self) -> worker_runtime::Runtime { + self.runtime.clone() + } + pub fn from_runtime(workspace_id: impl AsRef, runtime: worker_runtime::Runtime) -> Self { let workspace_id = workspace_id.as_ref().to_string(); Self { diff --git a/crates/workspace-server/src/runtime_subscription.rs b/crates/workspace-server/src/runtime_subscription.rs index 7ea983ba..94002e5f 100644 --- a/crates/workspace-server/src/runtime_subscription.rs +++ b/crates/workspace-server/src/runtime_subscription.rs @@ -148,6 +148,45 @@ impl RuntimeSubscriptionBroker { generation } + pub fn register_embedded_runtime( + &self, + runtime_id: impl Into, + runtime: worker_runtime::Runtime, + ) -> u64 { + let runtime_id = runtime_id.into(); + let generation = self.next_generation.fetch_add(1, Ordering::Relaxed); + let (commands, receiver) = mpsc::unbounded_channel(); + let status = Arc::new(RwLock::new(RuntimeSubscriptionBrokerStatus { + connection_generation: generation, + connected: true, + ..Default::default() + })); + let previous = self + .registrations + .write() + .expect("broker registry poisoned") + .insert( + runtime_id.clone(), + Registration { + generation, + commands: commands.clone(), + status: status.clone(), + }, + ); + if let Some(previous) = previous { + let _ = previous.commands.send(Command::Shutdown(generation)); + } + tokio::spawn(run_embedded_connection( + runtime_id, + runtime, + self.workspace_id.to_string(), + generation, + receiver, + status, + )); + generation + } + pub fn unregister_runtime(&self, runtime_id: &str) { if let Some(registration) = self .registrations @@ -282,6 +321,112 @@ impl State { } } +struct EmbeddedEntry { + downstreams: HashMap>, + snapshot_revision: u64, + snapshot: SubscriptionSnapshot, + task: tokio::task::JoinHandle<()>, +} + +async fn run_embedded_connection( + runtime_id: String, + runtime: worker_runtime::Runtime, + _workspace_id: String, + generation: u64, + mut commands: mpsc::UnboundedReceiver, + status: Arc>, +) { + let (updates, mut update_receiver) = mpsc::unbounded_channel(); + let mut entries = HashMap::::new(); + let mut downstream_index = HashMap::::new(); + loop { + tokio::select! { + command = commands.recv() => match command { + Some(Command::Subscribe { downstream_id, selector, events }) => { + downstream_index.insert(downstream_id, selector.clone()); + if let Some(entry) = entries.get_mut(&selector) { + let _ = events.try_send(BrokerSubscriptionEvent::Snapshot { + connection_generation: generation, + snapshot_revision: entry.snapshot_revision, + snapshot: entry.snapshot.clone(), + }); + entry.downstreams.insert(downstream_id, events); + } else { + match runtime.subscribe_event_selector(selector.clone()) { + Ok(mut subscription) => { + let snapshot_revision = subscription.snapshot_revision(); + let snapshot = project_snapshot_runtime(subscription.snapshot().clone(), &runtime_id); + let _ = events.try_send(BrokerSubscriptionEvent::Snapshot { connection_generation: generation, snapshot_revision, snapshot: snapshot.clone() }); + let sender = updates.clone(); + let task_selector = selector.clone(); + let task_runtime_id = runtime_id.clone(); + let task = tokio::spawn(async move { + while let Ok(update) = subscription.recv().await { + let payload = project_payload_runtime(update.payload, &task_runtime_id); + if sender.send((task_selector.clone(), update.subject_revision, payload)).is_err() { break; } + } + }); + entries.insert(selector, EmbeddedEntry { downstreams: HashMap::from([(downstream_id, events)]), snapshot_revision, snapshot, task }); + } + Err(error) => { + let _ = events.try_send(BrokerSubscriptionEvent::Rejected { connection_generation: generation, code: SubscriptionRejectionCode::UnsupportedSelector, message: error.to_string() }); + } + } + } + } + Some(Command::Unsubscribe(id)) => { + if let Some(selector) = downstream_index.remove(&id) { + let empty = entries.get_mut(&selector).is_some_and(|entry| { entry.downstreams.remove(&id); entry.downstreams.is_empty() }); + if empty { if let Some(entry) = entries.remove(&selector) { entry.task.abort(); } } + } + } + Some(Command::Shutdown(replacement)) => { + for entry in entries.values_mut() { + broadcast(&mut entry.downstreams, BrokerSubscriptionEvent::Closed { connection_generation: generation, code: SubscriptionTerminationCode::ServerShutdown, message: format!("embedded Runtime generation {generation} was fenced by {replacement}") }); + entry.task.abort(); + } + return; + } + None => return, + }, + update = update_receiver.recv() => { + let Some((selector, subject_revision, payload)) = update else { return; }; + if let Some(entry) = entries.get_mut(&selector) { + broadcast(&mut entry.downstreams, BrokerSubscriptionEvent::Event { connection_generation: generation, subject_revision, payload }); + } + } + } + *status.write().expect("broker status poisoned") = RuntimeSubscriptionBrokerStatus { + connection_generation: generation, + connected: true, + desired_selectors: entries.len(), + upstream_subscriptions: entries.len(), + }; + } +} + +fn project_snapshot_runtime( + mut snapshot: SubscriptionSnapshot, + runtime_id: &str, +) -> SubscriptionSnapshot { + if let SubscriptionSnapshot::Workers { workers } = &mut snapshot { + for worker in workers { + worker.runtime_id = Some(runtime_id.to_string()); + } + } + snapshot +} + +fn project_payload_runtime( + mut payload: SubscriptionEventPayload, + runtime_id: &str, +) -> SubscriptionEventPayload { + if let SubscriptionEventPayload::WorkerUpserted { worker } = &mut payload { + worker.runtime_id = Some(runtime_id.to_string()); + } + payload +} + async fn run_connection( config: RemoteRuntimeConfig, workspace_id: String, diff --git a/crates/workspace-server/src/runtime_subscription_tests.rs b/crates/workspace-server/src/runtime_subscription_tests.rs index 57bd9e18..78a53ab0 100644 --- a/crates/workspace-server/src/runtime_subscription_tests.rs +++ b/crates/workspace-server/src/runtime_subscription_tests.rs @@ -276,3 +276,40 @@ async fn reconnect_resubscribes_and_replaces_state_from_fresh_snapshot() { assert_eq!(workers.len(), 2); restarted.abort(); } + +#[tokio::test] +async fn embedded_runtime_uses_in_process_subscription_source() { + 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(); + let BrokerSubscriptionEvent::Snapshot { snapshot, .. } = next_snapshot(&mut subscription).await + else { + panic!("expected embedded snapshot"); + }; + let SubscriptionSnapshot::Workers { workers } = snapshot else { + panic!("expected Worker snapshot"); + }; + assert_eq!( + workers[0].runtime_id.as_deref(), + Some("embedded-worker-runtime") + ); + runtime + .observe_worker_event( + &worker.worker_ref, + protocol::Event::Status { + status: protocol::WorkerStatus::Running, + }, + ) + .unwrap(); + assert!(matches!(next_event(&mut subscription).await, + BrokerSubscriptionEvent::Event { payload: SubscriptionEventPayload::WorkerUpserted { worker }, .. } + if worker.runtime_id.as_deref() == Some("embedded-worker-runtime") && worker.state == SubscriptionWorkerState::Running)); + server.abort(); +} diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 51100750..4941c47e 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -308,19 +308,20 @@ impl WorkspaceApi { .await?; import_configured_repositories(store.as_ref(), &config)?; config.repositories = load_configured_repositories_from_store(store.as_ref(), &config)?; - let runtime = RuntimeRegistry::for_workspace( - EmbeddedWorkerRuntime::new_fs_store_with_execution_backend( - config.workspace_id.clone(), - config.embedded_runtime_store_root.clone(), - execution_backend, - ) - .map(|runtime| runtime.with_resource_broker(resource_broker.clone())) - .map_err(|err| { - crate::Error::Store(format!("invalid embedded Worker backend: {err}")) - })?, - ); + let embedded_runtime = EmbeddedWorkerRuntime::new_fs_store_with_execution_backend( + config.workspace_id.clone(), + config.embedded_runtime_store_root.clone(), + execution_backend, + ) + .map(|runtime| runtime.with_resource_broker(resource_broker.clone())) + .map_err(|err| crate::Error::Store(format!("invalid embedded Worker backend: {err}")))?; + let embedded_subscription_runtime = embedded_runtime.subscription_runtime(); + let embedded_runtime_id = EMBEDDED_WORKER_RUNTIME_ID.to_string(); + let runtime = RuntimeRegistry::for_workspace(embedded_runtime); let runtime_subscription_broker = RuntimeSubscriptionBroker::new(config.workspace_id.clone()); + runtime_subscription_broker + .register_embedded_runtime(embedded_runtime_id, embedded_subscription_runtime); for remote_config in config.remote_runtime_sources.iter().cloned() { let remote_runtime = RemoteWorkerRuntime::new( remote_config.clone(),