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]