Close the gaps three probes left open

The owner check and the unlink were separate steps, so a start landing
between them still lost its metadata to the backend it replaced. Both
now run under the same lock as the write, which is the only writer.

Probe output is drained into a bounded tail rather than kept whole. The
binary being probed is a possibly damaged install by construction, so an
import stuck printing is exactly the input this path can meet, and the
payload is the last JSON line either way.

RECORD can outlive the files it lists: an interrupted replace deletes the
package and leaves the metadata behind. The top-level names it records
are stat'd now. 298 distributions in a working venv, none flagged, 44ms.
This commit is contained in:
Daniel Han 2026-07-27 18:02:06 +00:00
commit 984cafc1ea
5 changed files with 94 additions and 18 deletions

View file

@ -268,6 +268,9 @@ impl BackendOwnerState {
/// Only while the file is still ours. Cleanup after an exit can take seconds,
/// and a start in that window already wrote the next metadata to this path.
pub(crate) fn remove(self) {
// Held across the check and the unlink: a start that lands between the
// two would otherwise have its file deleted by the backend it replaced.
let _guard = lock_metadata();
if let Ok(Some(on_disk)) = read_metadata(&self.path) {
if on_disk.token_sha256 != self.metadata.token_sha256 {
warn!(
@ -316,7 +319,16 @@ impl BackendOwnerState {
}
}
/// Serializes activation against cleanup. Both run in this process, and the
/// check in remove() only holds if no write can land inside it.
static METADATA_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
fn lock_metadata() -> std::sync::MutexGuard<'static, ()> {
METADATA_LOCK.lock().unwrap_or_else(|e| e.into_inner())
}
fn write_metadata(path: &Path, metadata: &DesktopBackendMetadata) -> Result<(), String> {
let _guard = lock_metadata();
let parent = path
.parent()
.ok_or_else(|| "desktop owner metadata path has no parent".to_string())?;

View file

@ -712,6 +712,8 @@ sleep 120
async fn managed_runtime_probe_survives_more_stdout_than_the_pipe_holds() {
// The probe imports the whole backend, so output can exceed the 64 KiB
// pipe buffer, and waiting before draining wedges the child mid-write.
// This CLI writes ~1 MiB before the payload, so the drain both keeps up
// and keeps only the tail, which still has to contain the payload.
let _cache_guard = MANAGED_CAPABILITY_CACHE_TEST_LOCK.lock().await;
let _cache_home = ManagedCapabilityCacheHome::new("runtime-noisy");
remove_managed_capability_cache();
@ -720,7 +722,7 @@ sleep 120
"runtime-noisy",
r#"#!/bin/sh
if [ "$1" = "studio" ] && [ "$2" = "desktop-runtime-check" ] && [ "$3" = "--json" ]; then
awk 'BEGIN { while (i++ < 4000) print "import chatter on stdout" }'
awk 'BEGIN { while (i++ < 40000) print "import chatter on stdout" }'
printf '{"runtime_ready":true}\n'
exit 0
fi

View file

@ -303,6 +303,30 @@ async fn run_cli_probe(bin: &Path, args: &[&str]) -> bool {
}
}
/// Room for the payload and what preceded it. Bounded because the binary being
/// probed is a possibly damaged install, and keeping every byte lets an import
/// stuck printing grow this buffer for the whole timeout.
const PROBE_TAIL_BYTES: usize = 64 * 1024;
/// Drained while the wait runs: the backend import can print more than the pipe
/// holds, and waiting first deadlocks, timing out a healthy install into repair.
fn drain_probe_tail(mut stdout: tokio::process::ChildStdout) -> tokio::task::JoinHandle<Vec<u8>> {
tokio::spawn(async move {
let mut tail = Vec::new();
let mut chunk = [0u8; 8192];
while let Ok(read) = stdout.read(&mut chunk).await {
if read == 0 {
break;
}
tail.extend_from_slice(&chunk[..read]);
if tail.len() > PROBE_TAIL_BYTES {
tail.drain(..tail.len() - PROBE_TAIL_BYTES);
}
}
tail
})
}
/// The probe could not answer: it would not run, or exited without a payload.
/// Distinct from a CLI that is simply too old to know the command.
const RUNTIME_PROBE_FAILED: &str = "studio_runtime_probe_failed";
@ -353,17 +377,11 @@ async fn probe_cli_runtime(bin: &Path) -> Result<(), String> {
);
return Err(RUNTIME_PROBE_FAILED.to_string());
};
let Some(mut stdout) = child.stdout.take() else {
let Some(stdout) = child.stdout.take() else {
return Err(RUNTIME_PROBE_FAILED.to_string());
};
// Drain while waiting: the backend import can print more than the pipe holds,
// and waiting first deadlocks, timing out a healthy install into repair.
let reader = tokio::spawn(async move {
let mut buffer = Vec::new();
let _ = stdout.read_to_end(&mut buffer).await;
buffer
});
let reader = drain_probe_tail(stdout);
let status = match tokio::time::timeout(Duration::from_secs(10), child.wait()).await {
Ok(Ok(status)) => status,
@ -445,16 +463,11 @@ async fn probe_cli_capability(bin: &Path) -> Option<DesktopCapability> {
);
return None;
};
let Some(mut stdout) = child.stdout.take() else {
let Some(stdout) = child.stdout.take() else {
return None;
};
// Same reason as the runtime probe above.
let reader = tokio::spawn(async move {
let mut buffer = Vec::new();
let _ = stdout.read_to_end(&mut buffer).await;
buffer
});
let reader = drain_probe_tail(stdout);
match tokio::time::timeout(Duration::from_secs(10), child.wait()).await {
Ok(Ok(status)) if status.success() => {}

View file

@ -18,7 +18,7 @@ import types
import urllib.error
import urllib.request
from datetime import datetime, timezone
from pathlib import Path
from pathlib import Path, PurePosixPath
from typing import List, Optional
import typer
@ -297,6 +297,23 @@ def _studio_requirement_roots(run_mod):
return roots
def _recorded_packages_missing(installed) -> bool:
"""Whether RECORD lists top-level packages the venv no longer has.
An interrupted replace removes the package and can leave its RECORD, so
reading one is not proof the modules are there. Stats the top-level names
only: a handful per distribution, and the failure is never partial.
"""
packages = set()
for recorded in installed.files or ():
top = PurePosixPath(str(recorded)).parts[0]
# .dist-info is the metadata itself; .data and ../ land outside the tree.
if top == ".." or top.endswith((".dist-info", ".data")):
continue
packages.add(top)
return any(not installed.locate_file(name).exists() for name in packages)
def _missing_studio_requirement(run_mod):
"""First requirement studio.txt needs that this venv cannot supply.
@ -318,7 +335,7 @@ def _missing_studio_requirement(run_mod):
return requirement.name
# Metadata outlives the package: wheels store METADATA first, so a killed
# unpack leaves a version behind. RECORD is last, so its absence marks it.
if installed.files is None:
if installed.files is None or _recorded_packages_missing(installed):
return requirement.name
# Only studio.txt pins are enforced: install_python_stack uses --no-deps,
# so transitive bounds read unsatisfied in venvs that work. A prerelease

View file

@ -252,3 +252,35 @@ def test_a_root_pin_is_checked_even_when_a_dependency_names_it_first(monkeypatch
studio.desktop_runtime_check(_json_output = True)
assert json.loads(capsys.readouterr().out)["module"] == "huggingface-hub"
def test_a_record_without_its_package_is_reported_missing(monkeypatch, capsys, tmp_path):
"""An interrupted replace can delete the package and leave its RECORD, so a
readable file list is not proof the modules are there."""
studio = importlib.import_module("unsloth_cli.commands.studio")
backend = tmp_path / "backend"
requirements = backend / "requirements"
requirements.mkdir(parents = True)
(requirements / "studio.txt").write_text("structlog\n", encoding = "utf-8")
run_mod = SimpleNamespace(__file__ = str(backend / "run.py"))
monkeypatch.setattr(studio, "_load_run_module", lambda: run_mod)
site_packages = tmp_path / "site-packages"
(site_packages / "structlog-25.1.0.dist-info").mkdir(parents = True)
installed = SimpleNamespace(
version = "25.1.0",
requires = None,
files = ["structlog/__init__.py", "structlog-25.1.0.dist-info/RECORD"],
locate_file = lambda name: site_packages / name,
)
monkeypatch.setattr(
importlib.import_module("importlib.metadata"), "distribution", lambda _name: installed,
)
with pytest.raises(typer.Exit):
studio.desktop_runtime_check(_json_output = True)
assert json.loads(capsys.readouterr().out)["module"] == "structlog"
(site_packages / "structlog").mkdir()
studio.desktop_runtime_check(_json_output = True)
assert json.loads(capsys.readouterr().out) == {"runtime_ready": True}