use std::cell::RefCell;
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use crate::value::{DictMap, VmError, VmValue};
mod metadata_snapshot;
static TURN_EPOCH: AtomicU64 = AtomicU64::new(0);
thread_local! {
static TURN_STABLE_HOST_CACHE: RefCell<HashMap<String, (u64, VmValue)>> =
RefCell::new(HashMap::new());
}
fn current_epoch() -> u64 {
TURN_EPOCH.load(Ordering::Acquire)
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum TurnCacheDisposition {
StableRead,
Invalidates,
Live,
}
fn disposition(capability: &str, operation: &str) -> TurnCacheDisposition {
match (capability, operation) {
("runtime", "pipeline_input") | ("project", "metadata_get") => {
TurnCacheDisposition::StableRead
}
("project", "metadata_set" | "metadata_save" | "metadata_refresh_hashes") => {
TurnCacheDisposition::Invalidates
}
_ => TurnCacheDisposition::Live,
}
}
fn is_turn_stable(capability: &str, operation: &str) -> bool {
disposition(capability, operation) == TurnCacheDisposition::StableRead
}
fn invalidates_turn_stable_reads(capability: &str, operation: &str) -> bool {
disposition(capability, operation) == TurnCacheDisposition::Invalidates
}
pub(crate) struct InvalidationScope {
invalidates: bool,
}
impl Drop for InvalidationScope {
fn drop(&mut self) {
if self.invalidates {
reset();
}
}
}
pub(crate) fn invalidation_scope(capability: &str, operation: &str) -> InvalidationScope {
let invalidates = invalidates_turn_stable_reads(capability, operation);
if invalidates {
reset();
}
InvalidationScope { invalidates }
}
fn cache_key(capability: &str, operation: &str, params: &DictMap) -> String {
if params.is_empty() {
return format!("{capability}.{operation}");
}
let json = crate::llm::helpers::vm_value_to_json(&VmValue::dict(params.clone()));
format!(
"{capability}.{operation}#{}",
serde_json::to_string(&json).unwrap_or_default()
)
}
pub(crate) async fn cached_or<F, Fut>(
capability: &str,
operation: &str,
params: &DictMap,
dispatch: F,
) -> Result<Option<VmValue>, VmError>
where
F: FnOnce() -> Fut,
Fut: std::future::Future<Output = Result<Option<VmValue>, VmError>>,
{
if !is_turn_stable(capability, operation) {
return dispatch().await;
}
if let Some(cached) = lookup(capability, operation, params) {
return Ok(Some(cached));
}
let dispatch_epoch = current_epoch();
let result = dispatch().await?;
if let Some(value) = &result {
store_at_epoch(capability, operation, params, value, dispatch_epoch);
}
Ok(result)
}
pub(crate) async fn cached_metadata_or<F, Fut>(
params: &DictMap,
dispatch: F,
) -> Result<Option<VmValue>, VmError>
where
F: FnOnce(DictMap) -> Fut,
Fut: std::future::Future<Output = Result<Option<VmValue>, VmError>>,
{
metadata_snapshot::cached_or(params, current_epoch(), dispatch).await
}
pub fn lookup(capability: &str, operation: &str, params: &DictMap) -> Option<VmValue> {
if !is_turn_stable(capability, operation) {
return None;
}
let epoch = current_epoch();
if capability == "project" && operation == "metadata_get" {
return metadata_snapshot::lookup(params, epoch);
}
let key = cache_key(capability, operation, params);
TURN_STABLE_HOST_CACHE.with(|cache| {
cache
.borrow()
.get(&key)
.filter(|(written, _)| *written == epoch)
.map(|(_, value)| value.clone())
})
}
pub fn store(capability: &str, operation: &str, params: &DictMap, value: &VmValue) {
store_at_epoch(capability, operation, params, value, current_epoch());
}
fn store_at_epoch(
capability: &str,
operation: &str,
params: &DictMap,
value: &VmValue,
dispatch_epoch: u64,
) {
if !is_turn_stable(capability, operation) {
return;
}
if current_epoch() != dispatch_epoch {
return;
}
if capability == "project" && operation == "metadata_get" {
metadata_snapshot::store(params, value, dispatch_epoch);
return;
}
let key = cache_key(capability, operation, params);
TURN_STABLE_HOST_CACHE.with(|cache| {
cache
.borrow_mut()
.insert(key, (dispatch_epoch, value.clone()));
});
}
pub fn lookup_by_name(name: &str, params: &DictMap) -> Option<VmValue> {
let (capability, operation) = name.split_once('.')?;
lookup(capability, operation, params)
}
pub fn store_by_name(name: &str, params: &DictMap, value: &VmValue) {
if let Some((capability, operation)) = name.split_once('.') {
store(capability, operation, params, value);
}
}
pub(crate) fn reset() {
TURN_EPOCH.fetch_add(1, Ordering::AcqRel);
reset_local();
}
pub(crate) fn reset_local() {
TURN_STABLE_HOST_CACHE.with(|cache| cache.borrow_mut().clear());
metadata_snapshot::reset_local();
}
#[cfg(test)]
pub(crate) fn epoch_test_lock() -> &'static std::sync::Mutex<()> {
static LOCK: std::sync::OnceLock<std::sync::Mutex<()>> = std::sync::OnceLock::new();
LOCK.get_or_init(|| std::sync::Mutex::new(()))
}
#[cfg(test)]
mod tests {
use std::sync::{Arc, Mutex};
use super::super::{
clear_host_call_bridge, dispatch_host_operation, reset_host_state, set_host_call_bridge,
HostCallBridge,
};
use super::reset;
use crate::value::{DictMap, VmValue};
struct CountingRuntimeBridge {
counts: Arc<Mutex<std::collections::HashMap<(String, String), usize>>>,
}
struct VersionedMetadataBridge {
generation: Arc<std::sync::atomic::AtomicUsize>,
counts: Arc<Mutex<std::collections::HashMap<(String, String), usize>>>,
}
impl HostCallBridge for VersionedMetadataBridge {
fn dispatch<'a>(
&'a self,
capability: &'a str,
operation: &'a str,
params: &'a DictMap,
) -> super::super::HostCallDispatchFuture<'a> {
*self
.counts
.lock()
.unwrap()
.entry((capability.to_string(), operation.to_string()))
.or_insert(0) += 1;
if capability == "project" && operation == "metadata_set" {
self.generation
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
return super::super::host_call_ready(Ok(Some(VmValue::Nil)));
}
let generation = self.generation.load(std::sync::atomic::Ordering::SeqCst);
assert!(
params.get("namespace").is_none(),
"metadata snapshots must be fetched without a namespace projection"
);
let namespace = |name: &str| {
(
crate::value::intern_key(name),
VmValue::dict(DictMap::from_iter([(
crate::value::intern_key("generation"),
VmValue::Int(generation as i64),
)])),
)
};
super::super::host_call_ready(Ok(Some(VmValue::dict(DictMap::from_iter([
namespace("facts"),
namespace("test"),
])))))
}
}
impl HostCallBridge for CountingRuntimeBridge {
fn dispatch<'a>(
&'a self,
capability: &'a str,
operation: &'a str,
_params: &'a DictMap,
) -> super::super::HostCallDispatchFuture<'a> {
*self
.counts
.lock()
.unwrap()
.entry((capability.to_string(), operation.to_string()))
.or_insert(0) += 1;
super::super::host_call_ready(Ok(Some(VmValue::String(arcstr::ArcStr::from(format!(
"{capability}.{operation}"
))))))
}
}
fn run_async<F, Fut>(test: F)
where
F: FnOnce() -> Fut,
Fut: std::future::Future<Output = ()>,
{
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("runtime");
rt.block_on(async {
let local = tokio::task::LocalSet::new();
local.run_until(test()).await;
});
}
#[test]
fn turn_stable_host_capability_is_fetched_once_per_turn() {
let _guard = super::epoch_test_lock()
.lock()
.unwrap_or_else(|e| e.into_inner());
run_async(|| async {
reset_host_state();
let counts = Arc::new(Mutex::new(std::collections::HashMap::new()));
set_host_call_bridge(Arc::new(CountingRuntimeBridge {
counts: counts.clone(),
}));
let count = |cap: &str, op: &str| -> usize {
counts
.lock()
.unwrap()
.get(&(cap.to_string(), op.to_string()))
.copied()
.unwrap_or(0)
};
for _ in 0..20 {
dispatch_host_operation("runtime", "pipeline_input", &DictMap::new())
.await
.expect("pipeline_input");
}
assert_eq!(
count("runtime", "pipeline_input"),
1,
"20 same-turn reads must hit the host exactly once"
);
for _ in 0..3 {
dispatch_host_operation("runtime", "record_run", &DictMap::new())
.await
.expect("record_run");
}
assert_eq!(
count("runtime", "record_run"),
3,
"writes/non-stable ops must never be served from the turn memo"
);
reset();
for _ in 0..20 {
dispatch_host_operation("runtime", "pipeline_input", &DictMap::new())
.await
.expect("pipeline_input");
}
assert_eq!(
count("runtime", "pipeline_input"),
2,
"a new turn must re-fetch once, not serve the prior turn's value"
);
clear_host_call_bridge();
});
}
#[test]
fn metadata_namespaces_share_a_snapshot_and_writes_invalidate_inherited_values() {
let _guard = super::epoch_test_lock()
.lock()
.unwrap_or_else(|e| e.into_inner());
run_async(|| async {
reset_host_state();
let counts = Arc::new(Mutex::new(std::collections::HashMap::new()));
let generation = Arc::new(std::sync::atomic::AtomicUsize::new(0));
set_host_call_bridge(Arc::new(VersionedMetadataBridge {
generation,
counts: counts.clone(),
}));
let count = |op: &str| -> usize {
counts
.lock()
.unwrap()
.get(&("project".to_string(), op.to_string()))
.copied()
.unwrap_or(0)
};
let descendant_facts = DictMap::from_iter([
(
crate::value::intern_key("dir"),
VmValue::String(arcstr::ArcStr::from("src/nested")),
),
(
crate::value::intern_key("namespace"),
VmValue::String(arcstr::ArcStr::from("facts")),
),
]);
let descendant_test = DictMap::from_iter([
(
crate::value::intern_key("dir"),
VmValue::String(arcstr::ArcStr::from("src/nested")),
),
(
crate::value::intern_key("namespace"),
VmValue::String(arcstr::ArcStr::from("test")),
),
]);
for _ in 0..100 {
let value = dispatch_host_operation("project", "metadata_get", &descendant_facts)
.await
.expect("metadata_get");
assert!(matches!(
value.as_dict().and_then(|fields| fields.get("generation")),
Some(VmValue::Int(0))
));
}
assert_eq!(
count("metadata_get"),
1,
"100 exact reads must dispatch once"
);
dispatch_host_operation("project", "metadata_get", &descendant_test)
.await
.expect("parameter-distinct metadata_get");
assert_eq!(
count("metadata_get"),
1,
"sibling namespaces must project from one directory snapshot"
);
let ancestor_write = DictMap::from_iter([
(
crate::value::intern_key("dir"),
VmValue::String(arcstr::ArcStr::from("src")),
),
(
crate::value::intern_key("namespace"),
VmValue::String(arcstr::ArcStr::from("facts")),
),
(
crate::value::intern_key("value"),
VmValue::dict(DictMap::new()),
),
]);
dispatch_host_operation("project", "metadata_set", &ancestor_write)
.await
.expect("metadata_set");
assert_eq!(count("metadata_set"), 1);
let refreshed = dispatch_host_operation("project", "metadata_get", &descendant_facts)
.await
.expect("read after ancestor write");
assert!(
matches!(
refreshed
.as_dict()
.and_then(|fields| fields.get("generation")),
Some(VmValue::Int(1))
),
"an ancestor write must invalidate a cached descendant read"
);
assert_eq!(count("metadata_get"), 2);
reset();
dispatch_host_operation("project", "metadata_get", &descendant_facts)
.await
.expect("next-turn metadata_get");
assert_eq!(count("metadata_get"), 3, "the next turn must re-read once");
clear_host_call_bridge();
});
}
#[test]
fn every_canonical_metadata_mutator_invalidates_turn_stable_reads() {
for operation in ["metadata_set", "metadata_save", "metadata_refresh_hashes"] {
assert!(
super::invalidates_turn_stable_reads("project", operation),
"project.{operation} must invalidate the metadata read memo"
);
}
for operation in ["metadata_get", "metadata_inspect", "metadata_stale"] {
assert!(
!super::invalidates_turn_stable_reads("project", operation),
"read-only project.{operation} must not open a new epoch"
);
}
}
#[test]
fn mutation_scope_invalidates_before_and_after_every_return_path() {
let _guard = super::epoch_test_lock()
.lock()
.unwrap_or_else(|e| e.into_inner());
let before = super::current_epoch();
{
let _scope = super::invalidation_scope("project", "metadata_set");
assert!(
super::current_epoch() > before,
"mutation must invalidate before dispatch"
);
}
let after_mutation = super::current_epoch();
assert!(
after_mutation > before + 1,
"scope drop must invalidate after dispatch"
);
{
let _scope = super::invalidation_scope("project", "metadata_get");
}
assert_eq!(
super::current_epoch(),
after_mutation,
"read-only dispatch must not invalidate the memo"
);
}
#[test]
fn turn_boundary_on_another_thread_invalidates_this_thread() {
let _guard = super::epoch_test_lock()
.lock()
.unwrap_or_else(|e| e.into_inner());
let params = DictMap::new();
let cached = VmValue::String(arcstr::ArcStr::from("turn-1"));
super::store("runtime", "pipeline_input", ¶ms, &cached);
assert!(
super::lookup("runtime", "pipeline_input", ¶ms).is_some(),
"same-turn read must hit"
);
std::thread::spawn(reset).join().expect("reset thread");
assert!(
super::lookup("runtime", "pipeline_input", ¶ms).is_none(),
"a turn boundary observed on another thread must invalidate this thread's entry"
);
}
#[test]
fn pre_mutation_read_cannot_poison_the_post_mutation_epoch() {
let _guard = super::epoch_test_lock()
.lock()
.unwrap_or_else(|e| e.into_inner());
reset();
let params = DictMap::from_iter([(
crate::value::intern_key("dir"),
VmValue::String(arcstr::ArcStr::from("src")),
)]);
let dispatch_epoch = super::current_epoch();
reset();
super::store_at_epoch(
"project",
"metadata_get",
¶ms,
&VmValue::dict(DictMap::from_iter([(
crate::value::intern_key("facts"),
VmValue::dict(DictMap::new()),
)])),
dispatch_epoch,
);
assert!(
super::lookup("project", "metadata_get", ¶ms).is_none(),
"an old dispatch result must not become the new epoch's cached value"
);
}
#[test]
fn store_ignores_non_turn_stable_operations() {
let _guard = super::epoch_test_lock()
.lock()
.unwrap_or_else(|e| e.into_inner());
let params = DictMap::new();
let value = VmValue::String(arcstr::ArcStr::from("live"));
super::store("session", "active_roots", ¶ms, &value);
assert!(
super::lookup("session", "active_roots", ¶ms).is_none(),
"non-allowlisted reads must never be served from the memo"
);
}
#[test]
fn dotted_name_helpers_share_the_split_pair_entry() {
let _guard = super::epoch_test_lock()
.lock()
.unwrap_or_else(|e| e.into_inner());
reset();
let params = DictMap::new();
let value = VmValue::String(arcstr::ArcStr::from("shared"));
super::store_by_name("runtime.pipeline_input", ¶ms, &value);
assert_eq!(
super::lookup("runtime", "pipeline_input", ¶ms).map(|v| v.display()),
Some("shared".to_string()),
"store_by_name must populate the entry lookup() reads"
);
assert_eq!(
super::lookup_by_name("runtime.pipeline_input", ¶ms).map(|v| v.display()),
Some("shared".to_string()),
"lookup_by_name must read it back"
);
assert!(
super::lookup_by_name("no-separator", ¶ms).is_none(),
"a name without a capability separator must not panic or match"
);
}
}