From bb558bad2b12ffce2e208ec4ddb3b511a15fac96 Mon Sep 17 00:00:00 2001 From: Hare Date: Wed, 19 Aug 2026 06:19:31 +0900 Subject: [PATCH] fix: share compiled prompt projection cache --- crates/worker-runtime/src/worker_backend.rs | 83 +++++++++++++---- crates/worker-runtime/src/worker_source.rs | 98 +++++++++++++++++++-- crates/worker/src/lib.rs | 6 +- crates/worker/src/worker.rs | 53 +++++++---- 4 files changed, 200 insertions(+), 40 deletions(-) diff --git a/crates/worker-runtime/src/worker_backend.rs b/crates/worker-runtime/src/worker_backend.rs index 963fe5c3..d8f0e2f4 100644 --- a/crates/worker-runtime/src/worker_backend.rs +++ b/crates/worker-runtime/src/worker_backend.rs @@ -219,14 +219,26 @@ impl WorkerObservationProvider for RuntimeGrantedWorkerObservationProvider { #[derive(Debug, Default)] pub(crate) struct WorkspacePromptProjectionCache { - active: Mutex>>, + active: Mutex>>, + fetch_gates: Mutex>>>, } impl WorkspacePromptProjectionCache { + pub(crate) fn fetch_gate(&self, workspace_id: &str) -> Result>, String> { + let mut gates = self + .fetch_gates + .lock() + .map_err(|_| "Workspace Prompt projection fetch gates lock was poisoned".to_string())?; + Ok(gates + .entry(workspace_id.to_string()) + .or_insert_with(|| Arc::new(Mutex::new(()))) + .clone()) + } + pub(crate) fn active( &self, workspace_id: &str, - ) -> Result>, String> { + ) -> Result>, String> { self.active .lock() .map(|active| active.get(workspace_id).cloned()) @@ -236,31 +248,40 @@ impl WorkspacePromptProjectionCache { pub(crate) fn observe( &self, projection: worker::WorkspacePromptProjection, - ) -> Result, String> { + ) -> Result, String> { projection.validate().map_err(|error| error.to_string())?; let workspace_id = projection.workspace_id.clone(); - let projection = Arc::new(projection); + let resolution = Arc::new( + worker::WorkspacePromptCatalogResolution::new(projection) + .map_err(|error| error.to_string())?, + ); let mut active = self .active .lock() .map_err(|_| "Workspace Prompt projection cache lock was poisoned".to_string())?; if let Some(current) = active.get(&workspace_id) { - if current.config_revision > projection.config_revision { + if current.projection.config_revision > resolution.projection.config_revision { return Ok(current.clone()); } - if current.config_revision == projection.config_revision - && (current.source_digest != projection.source_digest - || current.projection_digest != projection.projection_digest - || current.catalog.catalog_digest != projection.catalog.catalog_digest) + if current.projection.config_revision == resolution.projection.config_revision + && (current.projection.source_digest != resolution.projection.source_digest + || current.projection.projection_digest + != resolution.projection.projection_digest + || current.projection.catalog.catalog_digest + != resolution.projection.catalog.catalog_digest + || current.projection.catalog.schema_fingerprint + != resolution.projection.catalog.schema_fingerprint + || current.projection.catalog.toolchain_fingerprint + != resolution.projection.catalog.toolchain_fingerprint) { return Err(format!( "Workspace Prompt projection identity changed without a config revision transition: workspace={workspace_id} revision={}", - projection.config_revision + resolution.projection.config_revision )); } } - active.insert(workspace_id, projection.clone()); - Ok(projection) + active.insert(workspace_id, resolution.clone()); + Ok(resolution) } } @@ -677,7 +698,7 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory { ) .map_err(|error| error.to_string())?; let projection = self.prompt_projection_cache.observe(projection)?; - loader = loader.with_effective_catalog(projection.catalog.clone()); + loader = loader.with_effective_catalog(projection.projection.catalog.clone()); } } let flow_transition_enabled = manifest.feature.flow.enabled; @@ -1995,8 +2016,11 @@ mod tests { cache.observe(projection_v2).unwrap(); let active = cache.active("workspace-a").unwrap().unwrap(); - assert_eq!(active.config_revision, 9); - assert_eq!(active.catalog.catalog_digest, catalog_v2.catalog_digest); + assert_eq!(active.projection.config_revision, 9); + assert_eq!( + active.projection.catalog.catalog_digest, + catalog_v2.catalog_digest + ); } #[test] @@ -2033,6 +2057,35 @@ mod tests { ); } + #[test] + fn workspace_prompt_projection_cache_rejects_same_revision_schema_drift() { + let templates = BTreeMap::from([("default".to_string(), "prompt".to_string())]); + let first_catalog = + worker::EffectivePromptCatalog::new(templates.clone(), 8, "schema-a", "toolchain") + .unwrap(); + let drifted_catalog = + worker::EffectivePromptCatalog::new(templates, 8, "schema-b", "toolchain").unwrap(); + let first = worker::WorkspacePromptProjection::new( + "workspace-a", + "source-a", + first_catalog.catalog_digest.clone(), + first_catalog, + ) + .unwrap(); + let drifted = worker::WorkspacePromptProjection::new( + "workspace-a", + "source-a", + drifted_catalog.catalog_digest.clone(), + drifted_catalog, + ) + .unwrap(); + let cache = WorkspacePromptProjectionCache::default(); + + cache.observe(first).unwrap(); + let error = cache.observe(drifted).unwrap_err(); + assert!(error.contains("without a config revision transition")); + } + #[test] fn restart_restore_reconstructs_runtime_owned_worker_mutation_client() { let identity = RuntimeIdentityMaterial::generate("runtime-source").unwrap(); diff --git a/crates/worker-runtime/src/worker_source.rs b/crates/worker-runtime/src/worker_source.rs index 750dc67c..7454d53e 100644 --- a/crates/worker-runtime/src/worker_source.rs +++ b/crates/worker-runtime/src/worker_source.rs @@ -2,8 +2,8 @@ use std::sync::Arc; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use worker::{ - WorkspaceClient, WorkspaceClientError, WorkspacePromptProjection, WorkspaceRequest, - WorkspaceRequestMethod, WorkspaceResponse, + WorkspaceClient, WorkspaceClientError, WorkspacePromptCatalogResolution, + WorkspacePromptProjection, WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse, }; use crate::auth::{ @@ -287,6 +287,7 @@ fn execute_remote_worker_remove_http_blocking( Ok(WorkspaceResponse { status, body }) } +#[derive(Clone)] pub struct RuntimeOwnedWorkspaceClient { workspace_id: String, base_url: String, @@ -398,15 +399,29 @@ impl WorkspaceClient for RuntimeOwnedWorkspaceClient { fn current_prompt_projection( &self, - ) -> Result, WorkspaceClientError> { + ) -> Result, WorkspaceClientError> { let Some(cache) = self.prompt_projection_cache.as_ref() else { return Ok(None); }; - if let Some(projection) = cache + if let Some(resolution) = cache .active(&self.workspace_id) .map_err(WorkspaceClientError::Request)? { - return Ok(Some((*projection).clone())); + return Ok(Some((*resolution).clone())); + } + let fetch_gate = cache + .fetch_gate(&self.workspace_id) + .map_err(WorkspaceClientError::Request)?; + let _fetch_guard = fetch_gate.lock().map_err(|_| { + WorkspaceClientError::Request( + "Workspace Prompt projection fetch gate was poisoned".to_string(), + ) + })?; + if let Some(resolution) = cache + .active(&self.workspace_id) + .map_err(WorkspaceClientError::Request)? + { + return Ok(Some((*resolution).clone())); } let response = self.execute(WorkspaceRequest::get(format!( "/api/w/{}/config-sources/active/prompt-projection", @@ -602,9 +617,78 @@ mod tests { .with_prompt_projection_cache(cache); let projection = client.current_prompt_projection().unwrap().unwrap(); + let second = client.current_prompt_projection().unwrap().unwrap(); - assert_eq!(projection.config_revision, 3); - assert_eq!(projection.source_digest, "source-3"); + assert_eq!(projection.projection.config_revision, 3); + assert_eq!(projection.projection.source_digest, "source-3"); + assert!(Arc::ptr_eq(&projection.catalog, &second.catalog)); + } + + #[test] + fn concurrent_prompt_projection_miss_fetches_once_and_shares_catalog() { + use std::io::{Read, Write}; + use std::net::TcpListener; + use std::sync::Barrier; + + let catalog = worker::EffectivePromptCatalog::new( + std::collections::BTreeMap::from([( + "default".to_string(), + "shared prompt".to_string(), + )]), + 5, + "schema", + "toolchain", + ) + .unwrap(); + let projection = WorkspacePromptProjection::new( + "workspace-a", + "source-5", + catalog.catalog_digest.clone(), + catalog, + ) + .unwrap(); + let body = serde_json::to_string(&projection).unwrap(); + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + let server = std::thread::spawn(move || { + let (mut stream, _) = listener.accept().unwrap(); + let mut request = [0_u8; 4096]; + let _ = stream.read(&mut request).unwrap(); + write!( + stream, + "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + body.len(), + body + ) + .unwrap(); + }); + let cache = Arc::new(WorkspacePromptProjectionCache::default()); + let client = + RuntimeOwnedWorkspaceClient::new("workspace-a", base_url, "runtime-a", "worker-a") + .with_prompt_projection_cache(cache); + let barrier = Arc::new(Barrier::new(8)); + let threads = (0..8) + .map(|_| { + let client = client.clone(); + let barrier = barrier.clone(); + std::thread::spawn(move || { + barrier.wait(); + client.current_prompt_projection().unwrap().unwrap() + }) + }) + .collect::>(); + let resolutions = threads + .into_iter() + .map(|thread| thread.join().unwrap()) + .collect::>(); + server.join().unwrap(); + + let first = &resolutions[0].catalog; + assert!( + resolutions + .iter() + .all(|resolution| Arc::ptr_eq(first, &resolution.catalog)) + ); } #[test] diff --git a/crates/worker/src/lib.rs b/crates/worker/src/lib.rs index ad4fd207..76ced8ce 100644 --- a/crates/worker/src/lib.rs +++ b/crates/worker/src/lib.rs @@ -45,7 +45,7 @@ pub use shared_state::WorkerSharedState; pub use worker::{ LocalWorkingDirectory, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker, WorkerError, WorkerFilesystemAuthority, WorkerRunResult, WorkerWorkspaceContext, WorkspaceClient, - WorkspaceClientError, WorkspaceId, WorkspaceIdError, WorkspaceRequest, WorkspaceRequestMethod, - WorkspaceResponse, apply_worker_manifest, marker_workspace_client, - unavailable_workspace_client, + WorkspaceClientError, WorkspaceId, WorkspaceIdError, WorkspacePromptCatalogResolution, + WorkspaceRequest, WorkspaceRequestMethod, WorkspaceResponse, apply_worker_manifest, + marker_workspace_client, unavailable_workspace_client, }; diff --git a/crates/worker/src/worker.rs b/crates/worker/src/worker.rs index f7606b20..f618980a 100644 --- a/crates/worker/src/worker.rs +++ b/crates/worker/src/worker.rs @@ -212,6 +212,38 @@ pub enum WorkspaceClientError { Request(String), } +#[derive(Clone)] +pub struct WorkspacePromptCatalogResolution { + pub projection: Arc, + pub catalog: Arc, +} + +impl std::fmt::Debug for WorkspacePromptCatalogResolution { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("WorkspacePromptCatalogResolution") + .field("workspace_id", &self.projection.workspace_id) + .field("config_revision", &self.projection.config_revision) + .field("source_digest", &self.projection.source_digest) + .field("projection_digest", &self.projection.projection_digest) + .finish_non_exhaustive() + } +} + +impl WorkspacePromptCatalogResolution { + pub fn new(projection: WorkspacePromptProjection) -> Result { + projection.validate()?; + let catalog = PromptCatalog::load( + &PromptCatalogSource::builtins_only() + .with_effective_catalog(projection.catalog.clone()), + )?; + Ok(Self { + projection: Arc::new(projection), + catalog, + }) + } +} + /// Path-free Workspace operation authority injected by Runtime/host code. /// /// Workers receive this trait object rather than a Backend URL. The concrete @@ -229,7 +261,7 @@ pub trait WorkspaceClient: std::fmt::Debug + Send + Sync { /// launch/session state; this hook never reconstructs historical prompts. fn current_prompt_projection( &self, - ) -> Result, WorkspaceClientError> { + ) -> Result, WorkspaceClientError> { Ok(None) } @@ -296,7 +328,7 @@ impl WorkspaceClient for ReviewerChildWorkspaceClient { fn current_prompt_projection( &self, - ) -> Result, WorkspaceClientError> { + ) -> Result, WorkspaceClientError> { self.inner.current_prompt_projection() } @@ -1258,7 +1290,7 @@ impl Worker { if self.system_prompt_template.is_some() { return Ok(()); } - let Some(projection) = self + let Some(resolution) = self .workspace_context .client() .current_prompt_projection() @@ -1268,25 +1300,16 @@ impl Worker { else { return Ok(()); }; - projection - .validate() - .map_err(|source| WorkerError::WorkspacePromptProjection { - message: source.to_string(), - })?; + let projection = &resolution.projection; let current = self.prompts.load(); if current.projection().config_revision == projection.config_revision && current.projection().source_digest == projection.source_digest && current.projection().catalog_digest == projection.projection_digest + && Arc::ptr_eq(¤t, &resolution.catalog) { return Ok(()); } - let catalog = PromptCatalog::load( - &PromptCatalogSource::builtins_only().with_effective_catalog(projection.catalog), - ) - .map_err(|source| WorkerError::WorkspacePromptProjection { - message: source.to_string(), - })?; - self.prompts.store(catalog); + self.prompts.store(resolution.catalog); Ok(()) }