From 16b7af6bfe5e04933aa34c9779c064c2fcbeea37 Mon Sep 17 00:00:00 2001 From: Adam Perry Date: Thu, 9 Jul 2026 07:32:29 +0000 Subject: [PATCH] perf: load skill roots concurrently (cherry picked from commit 7a3b81e5d8510cc3a91af5328417d5aeed20c7af) --- codex-rs/core-skills/src/root_loader.rs | 135 ++++-- codex-rs/core-skills/src/root_loader_tests.rs | 453 ++++++++++++++++++ 2 files changed, 561 insertions(+), 27 deletions(-) create mode 100644 codex-rs/core-skills/src/root_loader_tests.rs diff --git a/codex-rs/core-skills/src/root_loader.rs b/codex-rs/core-skills/src/root_loader.rs index a31c8a2753..df3c3980dd 100644 --- a/codex-rs/core-skills/src/root_loader.rs +++ b/codex-rs/core-skills/src/root_loader.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use std::sync::Mutex; use codex_utils_plugins::PluginSkillRoot; +use futures::StreamExt; use crate::SkillLoadOutcome; use crate::loader::SkillRoot; @@ -12,6 +13,21 @@ use crate::loader::SkillRootSnapshot; use crate::loader::load_skill_root; use crate::model::SkillFileSystemsByPath; +const MAX_CONCURRENT_ROOT_SCANS: usize = 4; + +struct PendingRootLoad { + root: SkillRoot, + root_index: usize, + duplicate_root_indices: Vec, + cache_key: Option, +} + +struct CompletedRootLoad { + root_index: usize, + duplicate_root_indices: Vec, + snapshot: SkillRootSnapshot, +} + /// Parsed plugin skill-root snapshots produced by one plugin load. /// /// Clones share the same snapshots. The plugins manager stores them with the corresponding loaded @@ -44,8 +60,11 @@ pub(crate) async fn load_and_merge_skill_roots( where I: IntoIterator, { - let mut root_snapshots = Vec::new(); - for root in roots { + let roots = roots.into_iter().collect::>(); + let root_count = roots.len(); + let mut pending_loads = Vec::::with_capacity(root_count); + let mut pending_load_by_cache_key = HashMap::::new(); + for (root_index, root) in roots.into_iter().enumerate() { let cache_key = match ( root.plugin_id.clone(), root.plugin_namespace.clone(), @@ -59,34 +78,92 @@ where }), _ => None, }; - let cached_snapshot = cache_key.as_ref().and_then(|cache_key| { - let plugin_skill_snapshots = plugin_skill_snapshots?; - plugin_skill_snapshots - .snapshots_by_root - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .get(cache_key) - .cloned() + // Preserve sequential cache semantics: with a cache, the first occurrence scans and later + // occurrences reuse its snapshot. Without a cache, duplicate roots remain independent. + if plugin_skill_snapshots.is_some() + && let Some(cache_key) = cache_key.as_ref() + && let Some(pending_index) = pending_load_by_cache_key.get(cache_key) + { + pending_loads[*pending_index] + .duplicate_root_indices + .push(root_index); + continue; + } + + if plugin_skill_snapshots.is_some() + && let Some(cache_key) = cache_key.as_ref() + { + pending_load_by_cache_key.insert(cache_key.clone(), pending_loads.len()); + } + pending_loads.push(PendingRootLoad { + root, + root_index, + duplicate_root_indices: Vec::new(), + cache_key, }); - let snapshot = match cached_snapshot { - Some(snapshot) => snapshot, - None => { - let snapshot = load_skill_root(root).await; - if let Some(plugin_skill_snapshots) = plugin_skill_snapshots - && let Some(cache_key) = cache_key - { - plugin_skill_snapshots - .snapshots_by_root - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .insert(cache_key, snapshot.clone()); - } - snapshot - } - }; - root_snapshots.push(snapshot); } + // Unordered buffering lets fast host and cached roots free slots while an earlier executor + // root is still pending. The indexed snapshots below restore the original merge order. + let completed_loads = futures::stream::iter(pending_loads) + .map(|pending| async move { + let PendingRootLoad { + root, + root_index, + duplicate_root_indices, + cache_key, + } = pending; + let cached_snapshot = cache_key.as_ref().and_then(|cache_key| { + let plugin_skill_snapshots = plugin_skill_snapshots?; + plugin_skill_snapshots + .snapshots_by_root + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get(cache_key) + .cloned() + }); + match cached_snapshot { + Some(snapshot) => CompletedRootLoad { + root_index, + duplicate_root_indices, + snapshot, + }, + None => { + let snapshot = load_skill_root(root).await; + if let Some(plugin_skill_snapshots) = plugin_skill_snapshots + && let Some(cache_key) = cache_key + { + plugin_skill_snapshots + .snapshots_by_root + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(cache_key, snapshot.clone()); + } + CompletedRootLoad { + root_index, + duplicate_root_indices, + snapshot, + } + } + } + }) + .buffer_unordered(MAX_CONCURRENT_ROOT_SCANS) + .collect::>() + .await; + + let mut indexed_root_snapshots = Vec::with_capacity(root_count); + for completed in completed_loads { + for root_index in completed.duplicate_root_indices { + indexed_root_snapshots.push((root_index, completed.snapshot.clone())); + } + indexed_root_snapshots.push((completed.root_index, completed.snapshot)); + } + indexed_root_snapshots.sort_unstable_by_key(|(root_index, _)| *root_index); + let root_snapshots = indexed_root_snapshots + .into_iter() + .map(|(_, snapshot)| snapshot) + .collect(); + merge_skill_root_snapshots(root_snapshots) } @@ -156,3 +233,7 @@ fn merge_skill_root_snapshots(snapshots: Vec) -> SkillLoadOut outcome } + +#[cfg(test)] +#[path = "root_loader_tests.rs"] +mod tests; diff --git a/codex-rs/core-skills/src/root_loader_tests.rs b/codex-rs/core-skills/src/root_loader_tests.rs new file mode 100644 index 0000000000..fef8dfd2de --- /dev/null +++ b/codex-rs/core-skills/src/root_loader_tests.rs @@ -0,0 +1,453 @@ +use std::fs; +use std::path::Path; +use std::path::PathBuf; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use codex_exec_server::CopyOptions; +use codex_exec_server::CreateDirectoryOptions; +use codex_exec_server::ExecutorFileSystem; +use codex_exec_server::ExecutorFileSystemFuture; +use codex_exec_server::FileMetadata; +use codex_exec_server::FileSystemReadStream; +use codex_exec_server::FileSystemSandboxContext; +use codex_exec_server::LOCAL_FS; +use codex_exec_server::ReadDirectoryEntry; +use codex_exec_server::RemoveOptions; +use codex_exec_server::WalkOptions; +use codex_exec_server::WalkOutcome; +use codex_protocol::protocol::SkillScope; +use codex_utils_absolute_path::test_support::PathBufExt; +use codex_utils_path_uri::PathUri; +use pretty_assertions::assert_eq; +use tokio::sync::Notify; +use tokio::sync::Semaphore; + +use super::MAX_CONCURRENT_ROOT_SCANS; +use super::PluginSkillSnapshots; +use super::load_and_merge_skill_roots; +use crate::SkillError; +use crate::loader::SkillRoot; + +const TEST_WAIT: Duration = Duration::from_secs(/*secs*/ 5); + +struct ControlledFileSystem { + inner: Arc, + walk_gate: Option>, + walks_started: AtomicUsize, + walk_started: Notify, + blocked_read_root: Option, + blocked_read_gate: Arc, + observed_read_root: Option, + observed_read: AtomicBool, + observed_read_notify: Notify, +} + +impl ControlledFileSystem { + fn new(inner: Arc) -> Self { + Self { + inner, + walk_gate: None, + walks_started: AtomicUsize::new(/*v*/ 0), + walk_started: Notify::new(), + blocked_read_root: None, + blocked_read_gate: Arc::new(Semaphore::new(/*permits*/ 0)), + observed_read_root: None, + observed_read: AtomicBool::new(/*v*/ false), + observed_read_notify: Notify::new(), + } + } + + fn with_walk_gate(mut self, walk_gate: Arc) -> Self { + self.walk_gate = Some(walk_gate); + self + } + + fn with_blocked_read_root(mut self, blocked_read_root: PathUri) -> Self { + self.blocked_read_root = Some(blocked_read_root); + self + } + + fn with_observed_read_root(mut self, observed_read_root: PathUri) -> Self { + self.observed_read_root = Some(observed_read_root); + self + } + + fn release_blocked_read(&self) { + self.blocked_read_gate.add_permits(/*n*/ 1); + } + + fn walks_started(&self) -> usize { + self.walks_started.load(Ordering::Acquire) + } + + async fn wait_for_walks(&self, expected: usize) { + tokio::time::timeout(TEST_WAIT, async { + loop { + let notified = self.walk_started.notified(); + if self.walks_started() >= expected { + break; + } + notified.await; + } + }) + .await + .expect("expected skill-root walks to start"); + } + + async fn wait_for_observed_read(&self) { + tokio::time::timeout(TEST_WAIT, async { + loop { + let notified = self.observed_read_notify.notified(); + if self.observed_read.load(Ordering::Acquire) { + break; + } + notified.await; + } + }) + .await + .expect("expected the unblocked skill root to read its skill"); + } +} + +impl ExecutorFileSystem for ControlledFileSystem { + fn canonicalize<'a>( + &'a self, + path: &'a PathUri, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, PathUri> { + self.inner.canonicalize(path, sandbox) + } + + fn read_file<'a>( + &'a self, + path: &'a PathUri, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, Vec> { + let should_block = path.basename().as_deref() == Some("SKILL.md") + && self + .blocked_read_root + .as_ref() + .is_some_and(|root| path.starts_with(root)); + let should_observe = path.basename().as_deref() == Some("SKILL.md") + && self + .observed_read_root + .as_ref() + .is_some_and(|root| path.starts_with(root)); + Box::pin(async move { + if should_block { + self.blocked_read_gate + .acquire() + .await + .expect("blocked read gate should remain open") + // Consume the permit so one release advances exactly one gated operation. + .forget(); + } + let result = self.inner.read_file(path, sandbox).await; + if should_observe { + self.observed_read.store(/*val*/ true, Ordering::Release); + self.observed_read_notify.notify_waiters(); + } + result + }) + } + + fn read_file_stream<'a>( + &'a self, + path: &'a PathUri, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, FileSystemReadStream> { + self.inner.read_file_stream(path, sandbox) + } + + fn write_file<'a>( + &'a self, + path: &'a PathUri, + contents: Vec, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, ()> { + self.inner.write_file(path, contents, sandbox) + } + + fn create_directory<'a>( + &'a self, + path: &'a PathUri, + options: CreateDirectoryOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, ()> { + self.inner.create_directory(path, options, sandbox) + } + + fn get_metadata<'a>( + &'a self, + path: &'a PathUri, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, FileMetadata> { + self.inner.get_metadata(path, sandbox) + } + + fn read_directory<'a>( + &'a self, + path: &'a PathUri, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, Vec> { + self.inner.read_directory(path, sandbox) + } + + fn walk<'a>( + &'a self, + path: &'a PathUri, + options: WalkOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, WalkOutcome> { + self.walks_started.fetch_add(/*val*/ 1, Ordering::AcqRel); + self.walk_started.notify_waiters(); + Box::pin(async move { + if let Some(walk_gate) = &self.walk_gate { + walk_gate + .acquire() + .await + .expect("walk gate should remain open") + // Consume the permit so one release advances exactly one gated operation. + .forget(); + } + self.inner.walk(path, options, sandbox).await + }) + } + + fn remove<'a>( + &'a self, + path: &'a PathUri, + options: RemoveOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, ()> { + self.inner.remove(path, options, sandbox) + } + + fn copy<'a>( + &'a self, + source_path: &'a PathUri, + destination_path: &'a PathUri, + options: CopyOptions, + sandbox: Option<&'a FileSystemSandboxContext>, + ) -> ExecutorFileSystemFuture<'a, ()> { + self.inner + .copy(source_path, destination_path, options, sandbox) + } +} + +fn write_skill(root: &Path, directory: &str, contents: &str) -> PathBuf { + let skill_path = root.join(directory).join("SKILL.md"); + fs::create_dir_all( + skill_path + .parent() + .expect("skill path should have a parent"), + ) + .expect("create skill directory"); + fs::write(&skill_path, contents).expect("write skill"); + skill_path +} + +fn plain_root( + path: &Path, + scope: SkillScope, + file_system: Arc, +) -> SkillRoot { + SkillRoot { + path: path.to_path_buf().abs(), + scope, + file_system, + plugin_id: None, + plugin_namespace: Some("test".to_string()), + plugin_root: None, + } +} + +fn plugin_root( + skills_root: &Path, + plugin_root: &Path, + file_system: Arc, +) -> SkillRoot { + SkillRoot { + path: skills_root.to_path_buf().abs(), + scope: SkillScope::User, + file_system, + plugin_id: Some("demo@test".to_string()), + plugin_namespace: Some("demo".to_string()), + plugin_root: Some(plugin_root.to_path_buf().abs()), + } +} + +#[tokio::test] +async fn loads_roots_unordered_but_merges_errors_in_input_order() { + const ROOT_COUNT: usize = MAX_CONCURRENT_ROOT_SCANS + 1; + + let temp = tempfile::tempdir().expect("tempdir"); + let roots = (0..ROOT_COUNT) + .map(|index| temp.path().join(format!("root-{index}"))) + .collect::>(); + for root in &roots { + fs::create_dir_all(root).expect("create root"); + } + let first_skill = write_skill(&roots[0], "broken", "missing frontmatter"); + let second_skill = write_skill(&roots[1], "broken", "also missing frontmatter"); + + let file_system = Arc::new( + ControlledFileSystem::new(Arc::clone(&LOCAL_FS)) + .with_blocked_read_root(PathUri::from_host_native_path(&roots[0]).expect("first root")) + .with_observed_read_root( + PathUri::from_host_native_path(&roots[1]).expect("second root"), + ), + ); + let root_file_system: Arc = file_system.clone(); + let skill_roots = roots + .iter() + .enumerate() + .map(|(index, root)| { + plain_root( + root, + if index == 0 { + SkillScope::Repo + } else { + SkillScope::User + }, + Arc::clone(&root_file_system), + ) + }) + .collect::>(); + + let load = tokio::spawn(async move { + load_and_merge_skill_roots(skill_roots, /*plugin_skill_snapshots*/ None).await + }); + file_system.wait_for_observed_read().await; + file_system.wait_for_walks(ROOT_COUNT).await; + file_system.release_blocked_read(); + let outcome = load.await.expect("skill-root load should finish"); + + assert_eq!(outcome.skills, Vec::new()); + assert_eq!( + outcome.errors, + vec![ + SkillError { + path: dunce::canonicalize(first_skill) + .expect("canonical first skill") + .abs(), + message: "missing YAML frontmatter delimited by ---".to_string(), + }, + SkillError { + path: dunce::canonicalize(second_skill) + .expect("canonical second skill") + .abs(), + message: "missing YAML frontmatter delimited by ---".to_string(), + }, + ] + ); +} + +#[tokio::test] +async fn limits_concurrent_root_scans() { + const ROOT_COUNT: usize = MAX_CONCURRENT_ROOT_SCANS + 1; + + let temp = tempfile::tempdir().expect("tempdir"); + let roots = (0..ROOT_COUNT) + .map(|index| temp.path().join(format!("root-{index}"))) + .collect::>(); + for root in &roots { + fs::create_dir_all(root).expect("create root"); + } + + let walk_gate = Arc::new(Semaphore::new(/*permits*/ 0)); + let file_system = Arc::new( + ControlledFileSystem::new(Arc::clone(&LOCAL_FS)).with_walk_gate(Arc::clone(&walk_gate)), + ); + let root_file_system: Arc = file_system.clone(); + let skill_roots = roots + .iter() + .map(|root| plain_root(root, SkillScope::User, Arc::clone(&root_file_system))) + .collect::>(); + + let load = tokio::spawn(async move { + load_and_merge_skill_roots(skill_roots, /*plugin_skill_snapshots*/ None).await + }); + file_system.wait_for_walks(MAX_CONCURRENT_ROOT_SCANS).await; + assert_eq!(file_system.walks_started(), MAX_CONCURRENT_ROOT_SCANS); + + walk_gate.add_permits(/*n*/ 1); + file_system.wait_for_walks(ROOT_COUNT).await; + assert_eq!(file_system.walks_started(), ROOT_COUNT); + + walk_gate.add_permits(MAX_CONCURRENT_ROOT_SCANS); + let outcome = load.await.expect("skill-root load should finish"); + assert_eq!(outcome.skills, Vec::new()); + assert_eq!(outcome.errors, Vec::new()); +} + +#[tokio::test] +async fn duplicate_plugin_roots_share_work_only_when_snapshots_are_available() { + let temp = tempfile::tempdir().expect("tempdir"); + let skills_root = temp.path().join("skills"); + write_skill( + &skills_root, + "demo", + "---\nname: demo\ndescription: demo skill\n---\n", + ); + let broken_skill = write_skill(&skills_root, "broken", "missing frontmatter"); + + let cached_file_system = Arc::new(ControlledFileSystem::new(Arc::clone(&LOCAL_FS))); + let cached_root_file_system: Arc = cached_file_system.clone(); + let snapshots = PluginSkillSnapshots::for_plugin_load(); + let cached_outcome = load_and_merge_skill_roots( + [ + plugin_root( + &skills_root, + temp.path(), + Arc::clone(&cached_root_file_system), + ), + plugin_root( + &skills_root, + temp.path(), + Arc::clone(&cached_root_file_system), + ), + ], + Some(&snapshots), + ) + .await; + assert_eq!(cached_file_system.walks_started(), 1); + + let uncached_file_system = Arc::new(ControlledFileSystem::new(Arc::clone(&LOCAL_FS))); + let uncached_root_file_system: Arc = uncached_file_system.clone(); + let uncached_outcome = load_and_merge_skill_roots( + [ + plugin_root( + &skills_root, + temp.path(), + Arc::clone(&uncached_root_file_system), + ), + plugin_root( + &skills_root, + temp.path(), + Arc::clone(&uncached_root_file_system), + ), + ], + /*plugin_skill_snapshots*/ None, + ) + .await; + assert_eq!(uncached_file_system.walks_started(), 2); + + assert_eq!(cached_outcome.skills, uncached_outcome.skills); + assert_eq!(cached_outcome.errors, uncached_outcome.errors); + assert_eq!(cached_outcome.skills.len(), 1); + let expected_error = SkillError { + path: dunce::canonicalize(broken_skill) + .expect("canonical broken skill") + .abs(), + message: "missing YAML frontmatter delimited by ---".to_string(), + }; + assert_eq!( + cached_outcome.errors, + vec![expected_error.clone(), expected_error] + ); +}