Skip to content

Commit e6ef4c1

Browse files
Fix slow hosted Poetry lock rewrite (#1473) (#1475)
* Start fix for #1473 Assisted-by: Claude Code:claude-opus-5-5 * Speed up the batched Poetry and PDM lock rewrite A hosted scan that pins N packages in one poetry.lock or pdm.lock searched the whole lock text two or three times per changed fragment to locate it and check it is unique, and the Poetry redirect then parsed the whole lock once more just to read [metadata] lock-version for the stale-install warning. On the 400-package bench those passes were about a fifth of the scan's instructions. The batch now finds and counts every changed fragment in one Aho-Corasick pass over the original and one over the result, and hands the lock version it read from its own parse back to the Poetry redirect. Output is unchanged: the batch-vs-step equivalence tests still pass, and a self-overlapping fragment only ever sends the batch to the step-by-step path, never accepts what that path would refuse. Refs #1473 Assisted-by: Claude Code:claude-opus-5-5 * Find lock fragments with a small automaton Building one Aho-Corasick automaton over whole package units cost more than the per-fragment scans it replaced. Search for each fragment's first 64 bytes instead and compare the full fragment at each hit: the same single pass, with an automaton a few KB in size. On the 400-package poetry/hosted bench the scan now runs 403M instructions instead of 506M on main, with a byte-identical lock and report. Refs #1473 Assisted-by: Claude Code:claude-opus-5-5 * Tighten the lock rewrite counters' docs Say exactly which sites the test-only parse and scan counters instrument, count renders through the same helper, note that only [metadata] lock-version is untouched by package rewrites, and stop the Poetry redirect's local shadowing the lock_version function. No behavior change. Refs #1473 Assisted-by: Claude Code:claude-opus-5-5 --------- Co-authored-by: Claude <noreply@anthropic.com>
1 parent d141050 commit e6ef4c1

4 files changed

Lines changed: 203 additions & 64 deletions

File tree

‎crates/socket-patch-core/src/patch/redirect/pdm.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -480,7 +480,7 @@ mod parse_reuse_equivalence_tests {
480480
/// re-parse, not one per package (#762).
481481
#[test]
482482
fn many_patches_render_the_lock_once() {
483-
use crate::utils::lock_fragments::RENDERS;
483+
use crate::utils::lock_fragments::{RENDERS, SCANS};
484484
let url = |name: &str| {
485485
format!("https://patch.socket.dev/patch/pypi/{name}/a/{name}-1.26.18-py3-none-any.whl")
486486
};
@@ -512,6 +512,7 @@ mod parse_reuse_equivalence_tests {
512512
.collect();
513513
let files = BTreeMap::from([("pdm.lock".to_string(), lock)]);
514514
RENDERS.with(|renders| renders.set(0));
515+
SCANS.with(|scans| scans.set(0));
515516
let mut got = RewriteResult::default();
516517
rewrite(&files, &deps, &mut got);
517518
if got.confirmed_pdm_uuids.len() != 12 {
@@ -523,6 +524,9 @@ mod parse_reuse_equivalence_tests {
523524
1,
524525
"{fixture} crlf={crlf}"
525526
);
527+
// One pass over the original and one over the result find
528+
// and count every dep's fragments (#1473).
529+
assert_eq!(SCANS.with(|scans| scans.get()), 2, "{fixture} crlf={crlf}");
526530
}
527531
}
528532
assert!(checked >= 16, "only {checked} locks landed every dep");

‎crates/socket-patch-core/src/patch/redirect/poetry.rs‎

Lines changed: 45 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,8 @@ use super::{DepOverride, FileEdit, RewriteResult, RewriteWarning};
1111
use crate::crawlers::python_crawler::canonicalize_pypi_name;
1212
use crate::utils::poetry_lock::{
1313
generated_by_version, lock_version, lock_wiring, pyproject_source_constraints,
14-
rewrite_poetry_lock_all, source_constraint_advisory, source_constraint_breaks_writer,
15-
source_constraint_refusal, LockBatch, LockStep, PoetryLockDep,
14+
rewrite_poetry_lock_all_with_version, source_constraint_advisory,
15+
source_constraint_breaks_writer, source_constraint_refusal, LockBatch, LockStep, PoetryLockDep,
1616
};
1717

1818
/// Whether the lock was written by a Poetry release older than 1.4. Those
@@ -21,15 +21,25 @@ use crate::utils::poetry_lock::{
2121
/// serving the upstream bytes after the redirect. Formats `0`/`1.0`/`1.1` are
2222
/// only written by such releases; lock `2.0` is written by 1.3 through 1.8, so
2323
/// the `@generated by Poetry X.Y.Z` header (present from 1.4) decides there.
24-
fn pre_1_4_writer(lock_text: &str) -> Option<&'static str> {
25-
let lock = lock_text.parse::<DocumentMut>().ok()?;
26-
match lock_version(&lock) {
27-
Ok("0") => Some("0"),
28-
Ok("1.0") => Some("1.0"),
29-
Ok("1.1") => Some("1.1"),
30-
Ok("2.0") if !matches!(generated_by_version(lock_text), Some(v) if v >= (1, 4)) => {
31-
Some("2.0")
24+
///
25+
/// `version` is the lock's `lock_version` when the caller already has it;
26+
/// otherwise the lock is parsed for it.
27+
fn pre_1_4_writer(lock_text: &str, version: Option<&str>) -> Option<&'static str> {
28+
let parsed;
29+
let version = match version {
30+
Some(version) => version,
31+
None => {
32+
#[cfg(test)]
33+
crate::utils::lock_fragments::tick(&crate::utils::lock_fragments::PARSES);
34+
parsed = lock_text.parse::<DocumentMut>().ok()?;
35+
lock_version(&parsed).ok()?
3236
}
37+
};
38+
match version {
39+
"0" => Some("0"),
40+
"1.0" => Some("1.0"),
41+
"1.1" => Some("1.1"),
42+
"2.0" if !matches!(generated_by_version(lock_text), Some(v) if v >= (1, 4)) => Some("2.0"),
3343
_ => None,
3444
}
3545
}
@@ -116,13 +126,17 @@ pub(super) fn rewrite_poetry(
116126
})
117127
.collect();
118128
// Every dep is rewritten over one parse and one render of the lock.
119-
let LockBatch {
120-
text: content,
121-
steps,
122-
} = rewrite_poetry_lock_all(original, &deps);
129+
let (
130+
LockBatch {
131+
text: content,
132+
steps,
133+
},
134+
batch_lock_version,
135+
) = rewrite_poetry_lock_all_with_version(original, &deps);
123136
let mut stale_warned = false;
124137
// `pre_1_4_writer` reads only the first line and `[metadata]
125-
// lock-version`, which no package rewrite touches: judged once.
138+
// lock-version`, which no package rewrite touches: judged once, from
139+
// the batch's parse when it has one.
126140
let mut writer_format: Option<Option<&'static str>> = None;
127141
for (&(dep, _), step) in usable.iter().zip(steps) {
128142
match step {
@@ -152,9 +166,9 @@ pub(super) fn rewrite_poetry(
152166
});
153167
}
154168
if !stale_warned {
155-
if let Some(format) =
156-
*writer_format.get_or_insert_with(|| pre_1_4_writer(&content))
157-
{
169+
if let Some(format) = *writer_format.get_or_insert_with(|| {
170+
pre_1_4_writer(&content, batch_lock_version.as_deref())
171+
}) {
158172
stale_warned = true;
159173
// Poetry 1.0 installs url sources through pip, which reads
160174
// the hash from the URL fragment; pip 22.3–23.0 take the
@@ -495,7 +509,7 @@ mod equivalence_tests {
495509
/// re-parse, not one per package (#760).
496510
#[test]
497511
fn many_patches_render_the_lock_once() {
498-
use crate::utils::lock_fragments::RENDERS;
512+
use crate::utils::lock_fragments::{PARSES, RENDERS, SCANS};
499513
for version in VERSIONS.iter().filter(|version| !version.starts_with("0.")) {
500514
for crlf in [false, true] {
501515
let mut lock = grown(version, 11);
@@ -516,13 +530,25 @@ mod equivalence_tests {
516530
let mut deps = vec![dep("urllib3", "1.26.18", Some(SHA), 0)];
517531
deps.extend((0..11).map(|i| dep(&format!("pkg{i}"), "1.26.18", Some(SHA), i + 1)));
518532
RENDERS.with(|renders| renders.set(0));
533+
PARSES.with(|parses| parses.set(0));
534+
SCANS.with(|scans| scans.set(0));
519535
let got = run(rewrite_poetry, &files, &deps);
520536
assert_eq!(got.confirmed_python_lock_uuids.len(), 12, "{version}");
521537
assert_eq!(
522538
RENDERS.with(|renders| renders.get()),
523539
1,
524540
"{version} crlf={crlf}"
525541
);
542+
// The batch parses the lock and its rendering, and the
543+
// stale-install check reads the batch's parse instead of its
544+
// own; the fragments are found and counted in one pass over
545+
// the original and one over the result (#1473).
546+
assert_eq!(
547+
PARSES.with(|parses| parses.get()),
548+
2,
549+
"{version} crlf={crlf}"
550+
);
551+
assert_eq!(SCANS.with(|scans| scans.get()), 2, "{version} crlf={crlf}");
526552
}
527553
}
528554
}

‎crates/socket-patch-core/src/utils/lock_fragments.rs‎

Lines changed: 116 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,19 @@ thread_local! {
1818
/// Whole-lock renders this thread's rewrites took, each followed by a
1919
/// full re-parse: what a hosted rewrite of N deps must not pay N times.
2020
pub(crate) static RENDERS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
21+
/// Whole-lock parses at the instrumented sites: [`rewrite_batch`]'s
22+
/// parse of the original and of its rendering, and the Poetry
23+
/// redirect's stale-install check when it has to parse for itself.
24+
pub(crate) static PARSES: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
25+
/// Whole-text fragment searches [`rewrite_batch`] took to find or count
26+
/// its changed fragments: what a batch of N deps must not pay N times.
27+
pub(crate) static SCANS: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
28+
}
29+
30+
/// Count one tick of a test counter ([`RENDERS`], [`PARSES`], [`SCANS`]).
31+
#[cfg(test)]
32+
pub(crate) fn tick(counter: &'static std::thread::LocalKey<std::cell::Cell<usize>>) {
33+
counter.with(|count| count.set(count.get() + 1));
2134
}
2235

2336
/// Takes `name`'s fragments of `text` from its (spanned) parse.
@@ -220,7 +233,7 @@ pub(crate) fn finish<'a>(
220233
fragments_in: FragmentsIn,
221234
) -> Result<FragmentRewrite<'a>, String> {
222235
#[cfg(test)]
223-
RENDERS.with(|renders| renders.set(renders.get() + 1));
236+
tick(&RENDERS);
224237
let before = before?;
225238
let rendered = crate::utils::python_lock::preserve_line_endings(text, rendered);
226239
let after_doc = toml_edit::Document::parse(rendered).map_err(|e| e.to_string())?;
@@ -271,6 +284,45 @@ pub(crate) fn finish<'a>(
271284
})
272285
}
273286

287+
/// The start of each of `needles` in `haystack`, found in one pass; `None`
288+
/// unless every needle occurs exactly once.
289+
///
290+
/// The pass looks for each needle's first [`PREFIX`] bytes (a small
291+
/// automaton, where the whole fragments would make a large one) and checks
292+
/// the full needle at every place its prefix occurs. Every occurrence starts
293+
/// with its prefix, so none is missed. They are counted overlapping, so a
294+
/// needle that overlaps itself may count more than its non-overlapping
295+
/// `matches` would; that only ever sends the batch to the step-by-step
296+
/// rewrite, never accepts one that path would refuse.
297+
fn sole_occurrences(haystack: &str, needles: &[&str]) -> Option<Vec<usize>> {
298+
/// Enough of a lock fragment to name its package.
299+
const PREFIX: usize = 64;
300+
if needles.is_empty() {
301+
return Some(Vec::new());
302+
}
303+
if needles.iter().any(|needle| needle.is_empty()) {
304+
return None;
305+
}
306+
#[cfg(test)]
307+
tick(&SCANS);
308+
let prefixes = needles
309+
.iter()
310+
.map(|needle| &needle.as_bytes()[..needle.len().min(PREFIX)]);
311+
let finder = aho_corasick::AhoCorasick::new(prefixes).ok()?;
312+
let haystack = haystack.as_bytes();
313+
let mut starts: Vec<Option<usize>> = vec![None; needles.len()];
314+
for found in finder.find_overlapping_iter(haystack) {
315+
let index = found.pattern().as_usize();
316+
if !haystack[found.start()..].starts_with(needles[index].as_bytes()) {
317+
continue;
318+
}
319+
if starts[index].replace(found.start()).is_some() {
320+
return None;
321+
}
322+
}
323+
starts.into_iter().collect()
324+
}
325+
274326
/// What one dep's rewrite did to a lock, whether it ran in a
275327
/// [`rewrite_batch`] or step by step.
276328
#[derive(Debug, Clone, PartialEq, Eq)]
@@ -318,6 +370,8 @@ pub(crate) fn rewrite_batch<P>(
318370
if text.contains("\r\n") && text.replace("\r\n", "").contains('\n') {
319371
return None;
320372
}
373+
#[cfg(test)]
374+
tick(&PARSES);
321375
let original = toml_edit::Document::parse(text.to_owned()).ok()?;
322376
let mut lock = original.clone().into_mut();
323377
let mut steps = Vec::with_capacity(names.len());
@@ -347,14 +401,13 @@ pub(crate) fn rewrite_batch<P>(
347401
});
348402
}
349403
#[cfg(test)]
350-
RENDERS.with(|renders| renders.set(renders.get() + 1));
404+
tick(&RENDERS);
351405
let rendered = crate::utils::python_lock::preserve_line_endings(text, lock.to_string());
406+
#[cfg(test)]
407+
tick(&PARSES);
352408
let after_doc = toml_edit::Document::parse(rendered).ok()?;
353409
let rendered = after_doc.raw();
354410
let file_terminator = majority_terminator(text);
355-
// Each changed fragment as the byte range of `text` it replaces (trimmed
356-
// to what differs) and the replacement.
357-
let mut splices: Vec<(std::ops::Range<usize>, &str)> = Vec::new();
358411
let mut edits_of: Vec<(usize, Vec<(String, String)>)> = Vec::new();
359412
for (index, before) in &rewritten {
360413
let after = fragments_in(&after_doc, rendered, names[*index]).ok()?;
@@ -364,40 +417,42 @@ pub(crate) fn rewrite_batch<P>(
364417
let mut edits = Vec::new();
365418
for (old, new) in before.iter().zip(after) {
366419
let new = respell(old, &new, file_terminator);
367-
if *old == new {
368-
continue;
369-
}
370-
if text.matches(old.as_str()).count() != 1 {
371-
return None;
420+
if *old != new {
421+
edits.push((old.clone(), new));
372422
}
373-
edits.push((old.clone(), new));
374423
}
375424
edits_of.push((*index, edits));
376425
}
377-
for (_, edits) in &edits_of {
378-
for (old, new) in edits {
379-
let at = text.find(old.as_str())?;
380-
// Only the bytes that differ are replaced, so the units' shared
381-
// boundaries (a unit's fragment ends at the next unit's header)
382-
// never overlap.
383-
let common =
384-
|a: &mut dyn Iterator<Item = (u8, u8)>| a.take_while(|(x, y)| x == y).count();
385-
let mut prefix = common(&mut old.bytes().zip(new.bytes()));
386-
while !old.is_char_boundary(prefix) || !new.is_char_boundary(prefix) {
387-
prefix -= 1;
388-
}
389-
let mut suffix =
390-
common(&mut old[prefix..].bytes().rev().zip(new[prefix..].bytes().rev()));
391-
while !old.is_char_boundary(old.len() - suffix)
392-
|| !new.is_char_boundary(new.len() - suffix)
393-
{
394-
suffix -= 1;
395-
}
396-
splices.push((
397-
at + prefix..at + old.len() - suffix,
398-
&new[prefix..new.len() - suffix],
399-
));
426+
// Where each changed fragment sits in `text`, all found in one pass; a
427+
// fragment that is not there exactly once ends the batch, as it refuses
428+
// the step-by-step rewrite.
429+
let olds: Vec<&str> = edits_of
430+
.iter()
431+
.flat_map(|(_, edits)| edits.iter().map(|(old, _)| old.as_str()))
432+
.collect();
433+
let starts = sole_occurrences(text, &olds)?;
434+
// Each changed fragment as the byte range of `text` it replaces (trimmed
435+
// to what differs) and the replacement.
436+
let mut splices: Vec<(std::ops::Range<usize>, &str)> = Vec::new();
437+
let edits = edits_of.iter().flat_map(|(_, edits)| edits);
438+
for ((old, new), at) in edits.zip(starts) {
439+
// Only the bytes that differ are replaced, so the units' shared
440+
// boundaries (a unit's fragment ends at the next unit's header)
441+
// never overlap.
442+
let common = |a: &mut dyn Iterator<Item = (u8, u8)>| a.take_while(|(x, y)| x == y).count();
443+
let mut prefix = common(&mut old.bytes().zip(new.bytes()));
444+
while !old.is_char_boundary(prefix) || !new.is_char_boundary(prefix) {
445+
prefix -= 1;
400446
}
447+
let mut suffix = common(&mut old[prefix..].bytes().rev().zip(new[prefix..].bytes().rev()));
448+
while !old.is_char_boundary(old.len() - suffix) || !new.is_char_boundary(new.len() - suffix)
449+
{
450+
suffix -= 1;
451+
}
452+
splices.push((
453+
at + prefix..at + old.len() - suffix,
454+
&new[prefix..new.len() - suffix],
455+
));
401456
}
402457
splices.sort_by_key(|(range, _)| (range.start, range.end));
403458
let mut result = String::with_capacity(rendered.len());
@@ -414,13 +469,14 @@ pub(crate) fn rewrite_batch<P>(
414469
if result != rendered {
415470
return None;
416471
}
472+
// Every replacement must be unique in the result too, so its rollback
473+
// can't hit the wrong place.
474+
let news: Vec<&str> = edits_of
475+
.iter()
476+
.flat_map(|(_, edits)| edits.iter().map(|(_, new)| new.as_str()))
477+
.collect();
478+
sole_occurrences(&result, &news)?;
417479
for (index, edits) in edits_of {
418-
if edits
419-
.iter()
420-
.any(|(_, new)| result.matches(new.as_str()).count() != 1)
421-
{
422-
return None;
423-
}
424480
if !edits.is_empty() {
425481
steps[index] = LockStep::Rewritten(edits);
426482
}
@@ -450,6 +506,26 @@ mod tests {
450506
Ok(vec![text[span].to_string()])
451507
}
452508

509+
#[test]
510+
fn sole_occurrences_finds_each_needle_once_or_refuses() {
511+
// Long needles sharing their first 64+ bytes are told apart by the
512+
// full comparison.
513+
let shared = "x".repeat(80);
514+
let (a, b) = (format!("{shared}A"), format!("{shared}B"));
515+
let haystack = format!("[{a}]\n[{b}]\n");
516+
assert_eq!(
517+
sole_occurrences(&haystack, &[&b, &a, "]\n["]),
518+
Some(vec![haystack.find(&b).unwrap(), 1, a.len() + 1])
519+
);
520+
assert_eq!(sole_occurrences(&haystack, &[]), Some(vec![]));
521+
// Missing, repeated, empty and self-overlapping needles all refuse.
522+
assert_eq!(sole_occurrences(&haystack, &[&a, "missing"]), None);
523+
assert_eq!(sole_occurrences(&haystack, &[&shared]), None);
524+
assert_eq!(sole_occurrences(&haystack, &["]\n"]), None);
525+
assert_eq!(sole_occurrences(&haystack, &[""]), None);
526+
assert_eq!(sole_occurrences("aaa", &["aa"]), None);
527+
}
528+
453529
#[test]
454530
fn pair_fragments_refuses_a_shape_change() {
455531
for what in ["Poetry", "PDM"] {

0 commit comments

Comments
 (0)