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

fix: share compiled prompt projection cache

Browse Source
This commit is contained in:
Hare
2026-08-19 06:19:31 +09:00
parent fcc7c49ff1
commit bb558bad2b
4 changed files with 200 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker-runtime/src/worker_backend.rs
+68 -15
View File
@@ -219,14 +219,26 @@ impl WorkerObservationProvider for RuntimeGrantedWorkerObservationProvider {
#[derive(Debug, Default)]
pub(crate) struct WorkspacePromptProjectionCache {
active: Mutex<HashMap<String, Arc<worker::WorkspacePromptProjection>>>,
active: Mutex<HashMap<String, Arc<worker::WorkspacePromptCatalogResolution>>>,
fetch_gates: Mutex<HashMap<String, Arc<Mutex<()>>>>,
}
impl WorkspacePromptProjectionCache {
pub(crate) fn fetch_gate(&self, workspace_id: &str) -> Result<Arc<Mutex<()>>, 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<Option<Arc<worker::WorkspacePromptProjection>>, String> {
) -> Result<Option<Arc<worker::WorkspacePromptCatalogResolution>>, 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<Arc<worker::WorkspacePromptProjection>, String> {
) -> Result<Arc<worker::WorkspacePromptCatalogResolution>, 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();
crates/worker-runtime/src/worker_source.rs
+91 -7
View File
@@ -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<Option<WorkspacePromptProjection>, WorkspaceClientError> {
) -> Result<Option<WorkspacePromptCatalogResolution>, 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::<Vec<_>>();
let resolutions = threads
.into_iter()
.map(|thread| thread.join().unwrap())
.collect::<Vec<_>>();
server.join().unwrap();
let first = &resolutions[0].catalog;
assert!(
resolutions
.iter()
.all(|resolution| Arc::ptr_eq(first, &resolution.catalog))
);
}
#[test]
crates/worker/src/lib.rs
+3 -3
View File
@@ -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,
};
crates/worker/src/worker.rs
+38 -15
View File
@@ -212,6 +212,38 @@ pub enum WorkspaceClientError {
Request(String),
}
#[derive(Clone)]
pub struct WorkspacePromptCatalogResolution {
pub projection: Arc<WorkspacePromptProjection>,
pub catalog: Arc<PromptCatalog>,
}
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<Self, CatalogError> {
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<Option<WorkspacePromptProjection>, WorkspaceClientError> {
) -> Result<Option<WorkspacePromptCatalogResolution>, WorkspaceClientError> {
Ok(None)
}
@@ -296,7 +328,7 @@ impl WorkspaceClient for ReviewerChildWorkspaceClient {
fn current_prompt_projection(
&self,
) -> Result<Option<WorkspacePromptProjection>, WorkspaceClientError> {
) -> Result<Option<WorkspacePromptCatalogResolution>, WorkspaceClientError> {
self.inner.current_prompt_projection()
}
@@ -1258,7 +1290,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
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<C: LlmClient, St: Store> Worker<C, St> {
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(&current, &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(())
}