diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index f085be3793..c310077e6a 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -72,6 +72,7 @@ use std::num::NonZeroUsize; use std::ops::Deref; use std::sync::Arc; use std::sync::LazyLock; +use std::sync::OnceLock; use crate::context::GuardianContextMode; @@ -949,12 +950,13 @@ const ORIGINAL_IMAGE_PATCH_SIZE: u32 = 32; const ORIGINAL_IMAGE_MAX_PATCHES: usize = 10_000; const ORIGINAL_IMAGE_ESTIMATE_CACHE_SIZE: usize = 32; -static ORIGINAL_IMAGE_ESTIMATE_CACHE: LazyLock>> = - LazyLock::new(|| { - BlockingLruCache::new( - NonZeroUsize::new(ORIGINAL_IMAGE_ESTIMATE_CACHE_SIZE).unwrap_or(NonZeroUsize::MIN), - ) - }); +type OriginalImageEstimateCache = BlockingLruCache<[u8; 20], Arc>>>; + +static ORIGINAL_IMAGE_ESTIMATE_CACHE: LazyLock = LazyLock::new(|| { + BlockingLruCache::new( + NonZeroUsize::new(ORIGINAL_IMAGE_ESTIMATE_CACHE_SIZE).unwrap_or(NonZeroUsize::MIN), + ) +}); fn estimate_response_item_model_visible_bytes(item: &ResponseItem) -> i64 { match item { @@ -1111,7 +1113,7 @@ fn parse_base64_image_data_url(url: &str) -> Option<&str> { fn estimate_original_image_bytes(image_url: &str) -> Option { let key = sha1_digest(image_url.as_bytes()); - ORIGINAL_IMAGE_ESTIMATE_CACHE.get_or_insert_with(key, || { + ORIGINAL_IMAGE_ESTIMATE_CACHE.get_or_init(key, || { let payload = match parse_base64_image_data_url(image_url) { Some(payload) => payload, None => { diff --git a/codex-rs/utils/audio/src/lib.rs b/codex-rs/utils/audio/src/lib.rs index 53567952e0..0f1f53b8bb 100644 --- a/codex-rs/utils/audio/src/lib.rs +++ b/codex-rs/utils/audio/src/lib.rs @@ -11,7 +11,9 @@ use codex_utils_cache::sha1_digest; use codex_utils_string::approx_token_count; use std::io::Cursor; use std::num::NonZeroUsize; +use std::sync::Arc; use std::sync::LazyLock; +use std::sync::OnceLock; use symphonia::core::formats::FormatOptions; use symphonia::core::formats::TrackType; use symphonia::core::formats::probe::Hint; @@ -30,7 +32,7 @@ const MAX_PROMPT_AUDIO_BASE64_BYTES: usize = MAX_PROMPT_AUDIO_INPUT_BYTES.div_ce const AUDIO_TOKEN_ESTIMATE_CACHE_SIZE: usize = 32; const AUDIO_TOKENS_PER_SECOND: f64 = 10.0; -static AUDIO_TOKEN_ESTIMATE_CACHE: LazyLock> = +static AUDIO_TOKEN_ESTIMATE_CACHE: LazyLock>>> = LazyLock::new(|| { BlockingLruCache::new( NonZeroUsize::new(AUDIO_TOKEN_ESTIMATE_CACHE_SIZE).unwrap_or(NonZeroUsize::MIN), @@ -145,7 +147,7 @@ fn canonical_audio_mime(mime: &str) -> Option<&'static str> { /// Estimates audio tokens from decoded duration, falling back to the data URL size. pub fn estimate_audio_token_count(audio_url: &str) -> usize { let key = sha1_digest(audio_url.as_bytes()); - AUDIO_TOKEN_ESTIMATE_CACHE.get_or_insert_with(key, || { + AUDIO_TOKEN_ESTIMATE_CACHE.get_or_init(key, || { let Some(duration_seconds) = audio_duration_seconds(audio_url) else { return approx_token_count(audio_url); }; diff --git a/codex-rs/utils/cache/src/initialization_tests.rs b/codex-rs/utils/cache/src/initialization_tests.rs new file mode 100644 index 0000000000..23040f704e --- /dev/null +++ b/codex-rs/utils/cache/src/initialization_tests.rs @@ -0,0 +1,51 @@ +//! Checks that initialization releases the global lock and survives LRU eviction. + +use super::BlockingLruCache; +use std::num::NonZeroUsize; +use std::sync::mpsc; +use std::time::Duration; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn initialization_keeps_other_keys_available_through_eviction() { + let cache = BlockingLruCache::new(NonZeroUsize::new(/*n*/ 2).expect("capacity")); + assert_eq!(cache.get_or_init("hot", || 3), 3); + let handle = tokio::runtime::Handle::current(); + let (started_tx, started_rx) = mpsc::channel(); + let (release_tx, release_rx) = mpsc::channel(); + + std::thread::scope(|scope| { + let cache = &cache; + let initializing = scope.spawn(move || { + let _entered = handle.enter(); + cache.get_or_init("busy", || { + started_tx.send(()).expect("signal initialization"); + release_rx + .recv_timeout(Duration::from_secs(/*secs*/ 5)) + .expect("other keys should finish before initialization is released"); + 7 + }) + }); + started_rx + .recv_timeout(Duration::from_secs(/*secs*/ 5)) + .expect("initialization should start"); + + assert_eq!(cache.get_or_init("hot", || panic!("cached value")), 3); + assert_eq!(cache.get_or_init("other", || 5), 5); + assert!(cache.get(&"busy").is_none()); + + // The evicted initializer still owns its cell. A new resident cell can + // independently initialize the same deterministic value for this key. + assert_eq!(cache.get_or_init("busy", || 7), 7); + release_tx.send(()).expect("release initialization"); + assert_eq!(initializing.join().expect("initialization thread"), 7); + assert_eq!( + cache.with_mut(|entries| { + entries + .iter() + .map(|(key, value)| (*key, *value.get().expect("initialized value"))) + .collect::>() + }), + vec![("busy", 7), ("other", 5)] + ); + }); +} diff --git a/codex-rs/utils/cache/src/lib.rs b/codex-rs/utils/cache/src/lib.rs index ed4b926687..5375a0f4d4 100644 --- a/codex-rs/utils/cache/src/lib.rs +++ b/codex-rs/utils/cache/src/lib.rs @@ -1,6 +1,8 @@ use std::borrow::Borrow; use std::hash::Hash; use std::num::NonZeroUsize; +use std::sync::Arc; +use std::sync::OnceLock; use lru::LruCache; use sha1::Digest; @@ -119,6 +121,25 @@ where } } +impl BlockingLruCache>> +where + K: Eq + Hash, + V: Clone, +{ + /// Returns the cached value, initializing it outside the global cache lock. + /// + /// Concurrent callers share initialization while the entry remains cached. + /// An in-flight entry can be evicted and initialized again, so factories must + /// be deterministic. Initialization and same-key waits use a blocking region. + pub fn get_or_init(&self, key: K, value: impl FnOnce() -> V) -> V { + let entry = self.get_or_insert_with(key, || Arc::new(OnceLock::new())); + if let Some(value) = entry.get() { + return value.clone(); + } + tokio::task::block_in_place(|| entry.get_or_init(value).clone()) + } +} + fn lock_if_runtime(m: &Mutex>) -> Option>> where K: Eq + Hash, @@ -141,6 +162,10 @@ pub fn sha1_digest(bytes: &[u8]) -> [u8; 20] { out } +#[cfg(test)] +#[path = "initialization_tests.rs"] +mod initialization_tests; + #[cfg(test)] mod tests { use super::BlockingLruCache;