Skip to content

Fix env-keyed probe memo written three ways (#1888) #1895

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Mikola Lysenko (mikolalysenko) merged 4 commits into from
Nov 2, 2026
Merged

Fix env-keyed probe memo written three ways (#1888) #1895

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
0dc82d8
Start fix for #1888
claude Oct 30, 2026
6fbb64b
Test python site memo keeps env-moved answer
claude Oct 30, 2026
9ff63ee
Share one env-keyed probe memo across crawlers
claude Oct 30, 2026
baa250a
Address review on the shared probe memo
claude Oct 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
170 changes: 2 additions & 168 deletions crates/socket-patch-core/src/crawlers/composer_crawler.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use super::types::{CrawledPackage, CrawlerOptions};
use crate::patch::path_safety;
use crate::utils::composer_version::composer_versions_equivalent;
use crate::utils::fs::{is_dir, is_dir_sync, is_file, run_blocking};
use crate::utils::process::{CommandRunner, GlobalProbeRunner};
use crate::utils::process::{CommandRunner, EnvKeyedMemo, GlobalProbeRunner};
use crate::utils::relpath::normalize_lexically;

#[cfg(test)]
Expand Down Expand Up @@ -474,85 +474,10 @@ pub fn parse_composer_home_output(stdout: &str) -> Option<PathBuf> {
}
}

/// A process-wide memo of one subprocess answer, reused only while the
/// whole environment (every variable, and the working directory) is exactly
/// what it was when the answer was produced — the subprocess's only inputs
/// the CLI controls. A different environment re-runs it and replaces the
/// entry; concurrent first callers may each run it.
///
/// Only an answer is kept. A FAILED probe (spawn error under fd pressure,
/// a non-zero exit, empty stdout) is asked again by the next caller —
/// global discovery runs at least twice per command (the crawl, then the
/// package-lookup pass), so a transient first failure must not stick. So
/// is an answer the environment changed under, which the key would
/// misfile. Same rule as the `gem env`
/// memo in `ruby_crawler`.
struct EnvKeyedMemo {
slot: std::sync::Mutex<Option<(EnvKey, String)>>,
}

type EnvKey = (
Vec<(std::ffi::OsString, std::ffi::OsString)>,
Option<PathBuf>,
);

/// What the subprocess inherits: the (sorted) environment and the working
/// directory. Sorted so a `remove_var`/`set_var` round-trip that restores
/// an identical environment still matches (`vars_os` yields `environ`'s
/// order, which such a round-trip permutes) — the same key `gem_env_key`
/// builds.
fn env_key() -> EnvKey {
let mut vars: Vec<_> = std::env::vars_os().collect();
vars.sort();
(vars, std::env::current_dir().ok())
}

impl EnvKeyedMemo {
const fn new() -> Self {
Self {
slot: std::sync::Mutex::new(None),
}
}

fn get_or_run(&self, run: impl FnOnce() -> Option<String>) -> Option<String> {
let key = env_key();
let ran_under = key.clone();
self.get_or_run_keyed(key, move || {
let answer = run();
// Keepable only if the environment the subprocess actually ran
// under is still the one we keyed on.
let env_unchanged = env_key() == ran_under;
(answer, env_unchanged)
})
}

/// The memoized ask. `run` hands back its answer and whether that
/// answer is keepable; a `None` answer is never kept. Split from
/// [`Self::get_or_run`] so tests can drive both against a key of their
/// own.
fn get_or_run_keyed(
&self,
key: EnvKey,
run: impl FnOnce() -> (Option<String>, bool),
) -> Option<String> {
let lock = || self.slot.lock().unwrap_or_else(|e| e.into_inner());
if let Some((cached_key, answer)) = lock().as_ref() {
if *cached_key == key {
return Some(answer.clone());
}
}
let (answer, keepable) = run();
if let (Some(answer), true) = (&answer, keepable) {
*lock() = Some((key, answer.clone()));
}
answer
}
}

/// `composer global config home` boots PHP and Composer (0.3-1 s), and
/// global discovery asks for it on every pass (the crawl, then each
/// package-lookup pass of the same command).
static COMPOSER_GLOBAL_HOME: EnvKeyedMemo = EnvKeyedMemo::new();
static COMPOSER_GLOBAL_HOME: EnvKeyedMemo<String> = EnvKeyedMemo::new();

/// Get the Composer home directory.
///
Expand Down Expand Up @@ -2478,96 +2403,5 @@ mod tests {
"vacuous fixtures: {crawled}/{found}"
);
}

/// The `composer global config home` memo re-runs only when the
/// environment changed.
#[test]
fn env_keyed_memo_reuses_only_an_identical_environment() {
use std::ffi::OsString;
let env = |path: &str| -> EnvKey {
(
vec![(OsString::from("PATH"), OsString::from(path))],
Some(PathBuf::from("/cwd")),
)
};
let memo = EnvKeyedMemo::new();
let mut runs = 0;
let mut run = |answer: Option<&str>| {
runs += 1;
(answer.map(str::to_string), true)
};
assert_eq!(
memo.get_or_run_keyed(env("/a"), || run(Some("home"))),
Some("home".into())
);
assert_eq!(
memo.get_or_run_keyed(env("/a"), || run(Some("x"))),
Some("home".into())
);
// A different environment misses and replaces the entry.
assert_eq!(
memo.get_or_run_keyed(env("/b"), || run(Some("other"))),
Some("other".into())
);
assert_eq!(
memo.get_or_run_keyed(env("/b"), || run(Some("x"))),
Some("other".into())
);
assert_eq!(
memo.get_or_run_keyed(env("/a"), || run(Some("again"))),
Some("again".into())
);
assert_eq!(runs, 3);
}

/// A failed probe is not kept: the asker gets its own `None` and
/// the next caller asks again, as every caller did before the memo
/// (global discovery runs the probe once for the crawl and once for
/// the package-lookup pass). Neither is an answer the environment
/// changed under.
#[test]
fn env_keyed_memo_asks_again_after_a_failed_probe() {
use std::ffi::OsString;
let key: EnvKey = (
vec![(OsString::from("PATH"), OsString::from("/a"))],
Some(PathBuf::from("/cwd")),
);
let memo = EnvKeyedMemo::new();
let runs = std::cell::Cell::new(0);
let run = |answer: Option<&str>, keepable: bool| {
runs.set(runs.get() + 1);
(answer.map(str::to_string), keepable)
};

// A failure, then an answer whose environment moved: both are
// handed back, neither is kept.
assert_eq!(memo.get_or_run_keyed(key.clone(), || run(None, true)), None);
assert_eq!(
memo.get_or_run_keyed(key.clone(), || run(Some("moved"), false)),
Some("moved".into())
);
assert_eq!(runs.get(), 2);

// The next caller still asks, and a complete answer is kept.
assert_eq!(
memo.get_or_run_keyed(key.clone(), || run(Some("home"), true)),
Some("home".into())
);
assert_eq!(
memo.get_or_run_keyed(key.clone(), || run(None, true)),
Some("home".into())
);
assert_eq!(runs.get(), 3);
}

/// The memo key is order-insensitive: an environment restored to
/// the same variables still hits, whatever order `vars_os` yields.
#[test]
fn env_key_is_sorted() {
let key = super::super::env_key();
let mut sorted = key.0.clone();
sorted.sort();
assert_eq!(key.0, sorted);
}
}
}
108 changes: 8 additions & 100 deletions crates/socket-patch-core/src/crawlers/python_crawler.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ use super::types::{CrawledPackage, CrawlerOptions};
use crate::utils::fs::{
is_dir, is_dir_sync, read_regular_to_string, read_regular_to_string_sync, run_blocking,
};
use crate::utils::process::{CommandRunner, SystemCommandRunner};
use crate::utils::process::{CommandRunner, EnvKeyedMemo, SystemCommandRunner};

#[cfg(test)]
mod oracle;
Expand Down Expand Up @@ -2762,59 +2762,14 @@ fn pipenv_workon_home_venvs(
found
}

/// What the interpreter query below depends on besides the interpreter
/// itself: the whole process environment (`PATH` picks the interpreter;
/// `HOME`, `PYTHONUSERBASE`, `PYTHONHOME`, … shape its answer) and the
/// working directory.
type SiteQueryKey = (
Vec<(std::ffi::OsString, std::ffi::OsString)>,
Option<PathBuf>,
);

fn site_query_key() -> SiteQueryKey {
let mut env: Vec<_> = std::env::vars_os().collect();
env.sort();
(env, std::env::current_dir().ok())
}

/// The last successful `(key, answer)` of an expensive query, re-run
/// whenever the key differs. A failed query (`None`) is never stored, so the
/// next ask runs it again exactly as an unmemoized caller would: a transient
/// spawn failure during the crawl cannot blank the vendor phase's answer.
struct KeyedMemo<K, V> {
slot: std::sync::Mutex<Option<(K, V)>>,
}

impl<K: PartialEq, V: Clone> KeyedMemo<K, V> {
const fn new() -> Self {
Self {
slot: std::sync::Mutex::new(None),
}
}

fn get_or_run(&self, key: K, run: impl FnOnce() -> Option<V>) -> Option<V> {
let lock = || {
self.slot
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
};
if let Some((cached, value)) = lock().as_ref() {
if *cached == key {
return Some(value.clone());
}
}
let value = run()?;
*lock() = Some((key, value.clone()));
Some(value)
}
}

/// The interpreter's site-packages answer (the `find_python_command` probe
/// plus the `site` query — two process spawns), per [`SiteQueryKey`]: a
/// plus the `site` query — two process spawns), per environment (`PATH`
/// picks the interpreter; `HOME`, `PYTHONUSERBASE`, `PYTHONHOME`, … shape
/// its answer; see [`EnvKeyedMemo`]): a
/// scan asks it once for its crawl and again when it vendors, in the same
/// environment, so the second ask re-spawned both processes for the same
/// output. Only a successful answer is kept (see [`KeyedMemo`]).
static SITE_QUERY_MEMO: KeyedMemo<SiteQueryKey, String> = KeyedMemo::new();
/// output.
static SITE_QUERY_MEMO: EnvKeyedMemo<String> = EnvKeyedMemo::new();

/// The unmemoized interpreter query: probe for a python and ask it for its
/// site-packages. [`SITE_QUERY_MEMO`] wraps this; tests call it directly as
Expand Down Expand Up @@ -2853,8 +2808,7 @@ pub async fn get_global_python_site_packages() -> Vec<PathBuf> {
}

// 1. Ask Python for site-packages (subprocesses: on the blocking pool)
let site_output =
run_blocking(|| SITE_QUERY_MEMO.get_or_run(site_query_key(), run_site_query)).await;
let site_output = run_blocking(|| SITE_QUERY_MEMO.get_or_run(run_site_query)).await;
if let Some(stdout) = site_output {
for p in parse_python_site_packages_output(&stdout) {
add_path(p, &mut seen, &mut results);
Expand Down Expand Up @@ -8048,52 +8002,6 @@ G=
}
}

/// DC-6: the interpreter query memo answers from its slot only for the
/// exact key it was filled under, re-runs on any other key (a changed
/// environment or working directory), and never keeps a failed query's
/// `None`: the next ask with the same key runs the query again, as an
/// unmemoized caller would.
#[test]
fn keyed_memo_reruns_only_on_a_different_key() {
use std::sync::atomic::{AtomicUsize, Ordering};
let memo: KeyedMemo<(u32, &str), String> = KeyedMemo::new();
let runs = AtomicUsize::new(0);
let run = |out: Option<&str>| {
runs.fetch_add(1, Ordering::SeqCst);
out.map(str::to_string)
};
assert_eq!(
memo.get_or_run((1, "a"), || run(Some("x"))),
Some("x".into())
);
assert_eq!(
memo.get_or_run((1, "a"), || run(Some("stale"))),
Some("x".into())
);
assert_eq!(runs.load(Ordering::SeqCst), 1);
// A failure is returned but not stored: the same key re-runs.
assert_eq!(memo.get_or_run((2, "a"), || run(None)), None);
assert_eq!(runs.load(Ordering::SeqCst), 2);
assert_eq!(
memo.get_or_run((2, "a"), || run(Some("y"))),
Some("y".into())
);
assert_eq!(runs.load(Ordering::SeqCst), 3);
assert_eq!(
memo.get_or_run((2, "a"), || run(Some("stale"))),
Some("y".into())
);
assert_eq!(runs.load(Ordering::SeqCst), 3);
// A failure under a new key leaves nothing cached for the old one
// to be served from, and the old key re-runs.
assert_eq!(memo.get_or_run((3, "a"), || run(None)), None);
assert_eq!(
memo.get_or_run((1, "a"), || run(Some("z"))),
Some("z".into())
);
assert_eq!(runs.load(Ordering::SeqCst), 5);
}

/// DC-6: the memoized site query is wired to the whole environment.
/// Under two different `PYTHONUSERBASE` values (the interpreter's
/// `getusersitepackages` follows it) the memoized answer equals the
Expand Down Expand Up @@ -8123,7 +8031,7 @@ G=
}
}
}
let memoized = || SITE_QUERY_MEMO.get_or_run(site_query_key(), run_site_query);
let memoized = || SITE_QUERY_MEMO.get_or_run(run_site_query);

let base_a = tempfile::tempdir().unwrap();
let base_b = tempfile::tempdir().unwrap();
Expand Down
Loading
Loading