diff --git a/AGENTS.md b/AGENTS.md index 6e54d28..0b5b575 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -15,11 +15,10 @@ uv run cargo run -- test --reporter llm ## logging -A single knob (`-v` / `-q` flags or `TRYKE_LOG` env) drives both rust -and python verbosity. `RUST_LOG` is honored as a power-user override -for the rust side only. Default is `warn` on rust, off for workers -(no chatter unless asked). See `docs/guides/configuration.md` for the -full precedence chain. +A single knob (`-v` / `-q` flags or `TRYKE_LOG` env) drives both Rust +and Python verbosity. `RUST_LOG` is honored as a power-user override +for the Rust side only. The default is `warn` everywhere. See +`docs/guides/configuration.md` for the full precedence chain. ## running rust tests diff --git a/CHANGELOG.md b/CHANGELOG.md index 306ba24..955512c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -38,6 +38,7 @@ Released on 2026-07-04. - Add `--workers` to `tryke server` - Discover Python environments from `VIRTUAL_ENV`, Conda, and the project `.venv` +- Resolve one log level for Rust, reporter diagnostics, and Python workers ### Bug Fixes diff --git a/Cargo.lock b/Cargo.lock index eab729e..3b1a316 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -604,9 +604,45 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.32" +version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" +checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" + +[[package]] +name = "futures-macro" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6d3cde68c518367be28956066ddfef33813991b77a55005a69dae04bf3b10b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "futures-sink" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e34418ac499d6305c2fb5ad0ed2f6ac998c5f8ca209b4510f7f94242c647e307" + +[[package]] +name = "futures-task" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109" + +[[package]] +name = "futures-util" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa" +dependencies = [ + "futures-core", + "futures-macro", + "futures-task", + "pin-project-lite", + "slab", +] [[package]] name = "get-size-derive2" @@ -1759,6 +1795,12 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b2aa850e253778c88a04c3d7323b043aeda9d3e30d5971937c1855769763678e" +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + [[package]] name = "smallvec" version = "1.15.1" @@ -1952,6 +1994,20 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-util" +version = "0.7.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "494815d09bf52b5548659851081238f0ca39ff638363907596da739561c62c52" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "futures-util", + "pin-project-lite", + "tokio", +] + [[package]] name = "toml" version = "1.0.6+spec-1.1.0" @@ -2025,6 +2081,7 @@ dependencies = [ "serde_json", "tokio", "tokio-stream", + "tokio-util", "tryke_config", "tryke_discovery", "tryke_reporter", @@ -2040,7 +2097,6 @@ dependencies = [ name = "tryke_config" version = "0.0.30" dependencies = [ - "log", "serde", "tempfile", "toml", @@ -2106,6 +2162,7 @@ dependencies = [ "tempfile", "tokio", "tokio-stream", + "tokio-util", "tryke_config", "tryke_testing", "tryke_types", @@ -2124,6 +2181,7 @@ dependencies = [ "tempfile", "tokio", "tokio-stream", + "tokio-util", "tryke_discovery", "tryke_reporter", "tryke_runner", diff --git a/Cargo.toml b/Cargo.toml index 5be0446..ee6c9ac 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -40,6 +40,7 @@ tokio = { version = "1", features = [ "process", ] } tokio-stream = "0.1" +tokio-util = { version = "0.7", features = ["rt"] } notify = "6" notify-debouncer-mini = "0.4" toml = "1.0" diff --git a/crates/tryke/Cargo.toml b/crates/tryke/Cargo.toml index 082416b..4adce55 100644 --- a/crates/tryke/Cargo.toml +++ b/crates/tryke/Cargo.toml @@ -21,6 +21,7 @@ env_logger = "0.11" log = "0.4" tokio = { workspace = true, features = ["signal"] } tokio-stream = { workspace = true } +tokio-util = { workspace = true } tryke_config = { workspace = true } tryke_discovery = { workspace = true, features = ["filesystem"] } tryke_reporter = { workspace = true, features = ["terminal"] } diff --git a/crates/tryke/src/commands/clean.rs b/crates/tryke/src/commands/clean.rs index 42ca687..7e96f91 100644 --- a/crates/tryke/src/commands/clean.rs +++ b/crates/tryke/src/commands/clean.rs @@ -8,10 +8,12 @@ use crate::cli::{CleanArgs, GlobalArgs}; pub(crate) fn run_clean_command(args: CleanArgs, global: &GlobalArgs) -> Result { let cwd = env::current_dir()?; + let mut metadata = ProjectMetadata::new(args.root.as_deref().unwrap_or(&cwd)); metadata.apply_configuration_file(); metadata.apply_cli_args(args.project_options(global)); let project = Project::from_metadata(metadata); + let report = tryke_discovery::clean_project_cache(&project)?; if report.removed_entries == 0 { @@ -25,5 +27,6 @@ pub(crate) fn run_clean_command(args: CleanArgs, global: &GlobalArgs) -> Result< report.cache_dir.display() ); } + Ok(ExitStatus::Success) } diff --git a/crates/tryke/src/commands/server.rs b/crates/tryke/src/commands/server.rs index 5dfd7d4..4ac27c5 100644 --- a/crates/tryke/src/commands/server.rs +++ b/crates/tryke/src/commands/server.rs @@ -1,41 +1,44 @@ use std::env; use anyhow::Result; +use tokio_util::sync::CancellationToken; use tryke_config::{Project, ProjectMetadata}; use tryke_discovery::Discoverer; use tryke_runner::{WorkerPool, WorkerPoolOptions}; +use tryke_server::Server; use crate::ExitStatus; use crate::cli::{GlobalArgs, ServerArgs}; +use crate::logging::LogConfig; -pub(crate) fn run_server_command(args: ServerArgs, global: &GlobalArgs) -> Result { - let cli_filter = global.verbose.log_level_filter(); - let tryke_log = env::var("TRYKE_LOG").ok(); - let log_level = tryke_config::worker_log_level(tryke_log.as_deref(), cli_filter); +pub(crate) async fn run_server_command( + args: ServerArgs, + global: &GlobalArgs, + logging: LogConfig, + cancellation: CancellationToken, +) -> Result { let cwd = env::current_dir()?; + let mut metadata = ProjectMetadata::new(args.root.as_deref().unwrap_or(&cwd)); metadata.apply_configuration_file(); metadata.apply_cli_args(args.project_options(global)); let project = Project::from_metadata(metadata); - let runtime = tokio::runtime::Runtime::new()?; - - runtime.block_on(async move { - let worker_pool = WorkerPool::spawn( - &project, - WorkerPoolOptions { - size: args.workers, - python_path: None, - log_level, - warm: false, - }, - ) - .await; - - let discoverer = Discoverer::new(&project); - - tryke_server::Server::new(worker_pool, discoverer) - .serve() - .await - })?; + + let worker_pool = WorkerPool::spawn( + &project, + WorkerPoolOptions { + size: args.workers, + python_path: None, + log_level: logging.level(), + warm: false, + }, + ) + .await; + + let discoverer = Discoverer::new(&project); + let server = Server::new(worker_pool, discoverer); + + server.serve_with_cancellation(cancellation).await?; + Ok(ExitStatus::Success) } diff --git a/crates/tryke/src/commands/test.rs b/crates/tryke/src/commands/test.rs index 6c80ad6..b9fc474 100644 --- a/crates/tryke/src/commands/test.rs +++ b/crates/tryke/src/commands/test.rs @@ -6,11 +6,12 @@ use std::{ use anyhow::Result; use console::{Key, Term}; -use log::{LevelFilter, debug}; +use log::{LevelFilter, debug, warn}; use tokio_stream::StreamExt; +use tokio_util::sync::CancellationToken; use tryke_config::{Project, ProjectMetadata}; use tryke_discovery::{Discoverer, DiscoveryOptions}; -use tryke_reporter::{Reporter, Verbosity, build_reporter, reporter::WatchIdleInfo}; +use tryke_reporter::{Reporter, build_reporter, reporter::WatchIdleInfo}; use tryke_runner::{DistMode, WorkerPool, WorkerPoolOptions, partition_with_hooks}; use tryke_types::{ ChangedSelectionSummary, DiscoveryWarning, DiscoveryWarningKind, HookItem, RunSummary, @@ -21,6 +22,7 @@ use tryke_watcher::{FileChangeBatch, FileWatcher}; use super::CommandOrigin; use crate::ExitStatus; use crate::cli::{GlobalArgs, TestArgs}; +use crate::logging::LogConfig; #[derive(Debug, Eq, PartialEq)] enum Interruptible { @@ -30,14 +32,11 @@ enum Interruptible { async fn run_interruptibly( run: impl Future>, - interrupt: impl Future>, + cancellation: &CancellationToken, ) -> Result> { tokio::select! { biased; - signal = interrupt => { - signal?; - Ok(Interruptible::Interrupted) - }, + () = cancellation.cancelled() => Ok(Interruptible::Interrupted), result = run => result.map(Interruptible::Completed), } } @@ -59,16 +58,15 @@ fn resolve_interruptible( } } -pub(crate) fn run_test_command( +pub(crate) async fn run_test_command( args: TestArgs, global: &GlobalArgs, origin: CommandOrigin, + logging: LogConfig, + cancellation: CancellationToken, ) -> Result { - let cli_filter = global.verbose.log_level_filter(); - let tryke_log = env::var("TRYKE_LOG").ok(); - let rust_default = tryke_config::rust_log_default(tryke_log.as_deref(), cli_filter); - let log_level = tryke_config::worker_log_level(tryke_log.as_deref(), cli_filter); - let verbosity = Verbosity::from_level_filter(rust_default); + let log_level = logging.level(); + let verbosity = logging.reporter_verbosity(); let maxfail = if args.fail_fast { Some(1) @@ -77,7 +75,6 @@ pub(crate) fn run_test_command( }; let mut reporter = build_reporter(args.reporter.kind(), verbosity, global.no_progress); - let runtime = tokio::runtime::Runtime::new()?; let cwd = env::current_dir()?; let mut metadata = ProjectMetadata::new(args.root.as_deref().unwrap_or(&cwd)); @@ -97,20 +94,19 @@ pub(crate) fn run_test_command( TestFilter::from_args(&[], args.filter.as_deref(), args.markers.as_deref()) .map_err(|error| anyhow::anyhow!(error))?; - let result = runtime.block_on(run_interruptibly( - run_watch( - &mut *reporter, - &project, - log_level, - &test_filter, - maxfail, - args.workers, - args.dist.into(), - args.all, - args.now, - ), - tokio::signal::ctrl_c(), - )); + let result = run_watch( + &mut *reporter, + &project, + log_level, + &test_filter, + maxfail, + args.workers, + args.dist.into(), + args.all, + args.now, + &cancellation, + ) + .await; if resolve_interruptible(result, &mut *reporter)?.is_none() { return Ok(ExitStatus::Interrupted); @@ -158,21 +154,20 @@ pub(crate) fn run_test_command( // mode: args.snapshot_mode.to_wire(), // })?; - let result = runtime.block_on(run_interruptibly( - run_tests( - &mut *reporter, - &project, - log_level, - tests, - &discovered.hooks, - maxfail, - args.workers, - args.dist.into(), - Some(discovery_duration), - changed_selection, - ), - tokio::signal::ctrl_c(), - )); + let result = run_tests( + &mut *reporter, + &project, + log_level, + tests, + &discovered.hooks, + maxfail, + args.workers, + args.dist.into(), + Some(discovery_duration), + changed_selection, + &cancellation, + ) + .await; let Some(summary) = resolve_interruptible(result, &mut *reporter)? else { return Ok(ExitStatus::Interrupted); @@ -197,7 +192,8 @@ async fn run_tests( dist: DistMode, discovery_duration: Option, changed_selection: Option, -) -> Result { + cancellation: &CancellationToken, +) -> Result> { let pool = WorkerPool::spawn( project, WorkerPoolOptions { @@ -209,21 +205,34 @@ async fn run_tests( ) .await; - let summary = report_cycle( - reporter, - tests, - hooks, - &pool, - maxfail, - dist, - discovery_duration, - changed_selection, + let run_result = run_interruptibly( + report_cycle( + reporter, + tests, + hooks, + &pool, + maxfail, + dist, + discovery_duration, + changed_selection, + ), + cancellation, ) - .await?; + .await; + let shutdown_result = pool.shutdown().await; - pool.shutdown(); + combine_shutdown(run_result, shutdown_result) +} - Ok(summary) +fn combine_shutdown(result: Result, shutdown_result: Result<()>) -> Result { + match (result, shutdown_result) { + (Ok(value), Ok(())) => Ok(value), + (Ok(_), Err(shutdown_error)) => Err(shutdown_error), + (Err(error), Ok(())) => Err(error), + (Err(error), Err(shutdown_error)) => Err(error.context(format!( + "Worker pool shutdown also failed: {shutdown_error:#}" + ))), + } } fn flush_buffer( @@ -482,11 +491,6 @@ fn clear_watch_results(reporter: &mut dyn Reporter) { }); } -/// Run a single watch cycle. -/// -/// Test failures are non-fatal here: in watch -/// mode the whole point is to iterate on failing tests, so we discard -/// the run summary and any setup error and let the watcher keep running. async fn run_watch_cycle( reporter: &mut dyn Reporter, tests: Vec, @@ -496,7 +500,18 @@ async fn run_watch_cycle( dist: DistMode, discovery_duration: Option, ) { - pool.restart_workers().await; + // A worker that misses the restart deadline is still running the previous + // interpreter. Skipping the cycle costs the user a re-save; running it + // would report results from stale code as if they were current. + if let Err(error) = pool.restart_workers().await { + warn!("Watch: {error:#}"); + reporter.on_watch_idle(&WatchIdleInfo { + hint: "Worker restart failed — skipped this run. Waiting for file changes...", + start_time: None, + discovery_duration: None, + }); + return; + } if let Err(e) = report_cycle( reporter, @@ -510,17 +525,10 @@ async fn run_watch_cycle( ) .await { - debug!("watch: report_cycle errored: {e}"); + debug!("Watch: report_cycle errored: {e}"); } } -/// Perform startup discovery and, when `run_now` is set, the initial -/// test cycle. Discovery runs unconditionally so the import graph is -/// ready to answer "which tests are affected?" on the first file -/// change. When `run_now` is false we hand the reporter an idle frame -/// (header + Tests/Start/Discovery block + IDLE badge) so the -/// terminal communicates clearly that the watcher is alive and -/// waiting. async fn run_initial_cycle( reporter: &mut dyn Reporter, discoverer: &mut Discoverer, @@ -530,12 +538,6 @@ async fn run_initial_cycle( dist: DistMode, run_now: bool, ) { - // Arm before any reporter output so the deferred clear lands on - // the first warning, run-start, or idle frame — whichever fires - // first. The reporter's `flush_pending_clear` (called from each - // of those paths) consumes the flag, so warnings emitted just - // before `on_watch_idle` aren't wiped by a second clear inside - // the idle render. reporter.arm_clear(); let disc_start = Instant::now(); @@ -574,11 +576,8 @@ async fn run_watch( dist: DistMode, all_tests: bool, run_now: bool, -) -> Result<()> { - let root = project.root(); - let excludes = &project.discovery().exclude; - let mut discoverer = Discoverer::new(project); - + cancellation: &CancellationToken, +) -> Result> { let pool = WorkerPool::spawn( project, WorkerPoolOptions { @@ -590,11 +589,48 @@ async fn run_watch( ) .await; + let run_result = run_interruptibly( + run_watch_loop( + reporter, + project, + test_filter, + &pool, + maxfail, + dist, + all_tests, + run_now, + ), + cancellation, + ) + .await; + let shutdown_result = pool.shutdown().await; + + combine_shutdown(run_result, shutdown_result) +} + +#[expect( + clippy::too_many_arguments, + reason = "Watch options map directly to CLI flags; grouping into a struct would add indirection without clear benefit." +)] +async fn run_watch_loop( + reporter: &mut dyn Reporter, + project: &Project, + test_filter: &TestFilter, + pool: &WorkerPool, + maxfail: Option, + dist: DistMode, + all_tests: bool, + run_now: bool, +) -> Result<()> { + let root = project.root(); + let excludes = &project.discovery().exclude; + let mut discoverer = Discoverer::new(project); + run_initial_cycle( reporter, &mut discoverer, test_filter, - &pool, + pool, maxfail, dist, run_now, @@ -625,7 +661,7 @@ async fn run_watch( let hooks = discoverer.hooks(); let disc_dur = Some(disc_start.elapsed()); emit_discovery_warnings(reporter, &discoverer); - run_watch_cycle(reporter, tests, &hooks, &pool, maxfail, dist, disc_dur).await; + run_watch_cycle(reporter, tests, &hooks, pool, maxfail, dist, disc_dur).await; continue; } WatchLoopEvent::Command(WatchKeyAction::ClearResults) => { @@ -637,7 +673,7 @@ async fn run_watch( }; debug!( - "watch: file change batch — {} path(s) changed: {}", + "Watch: file change batch — {} path(s) changed: {}", paths.len(), paths .iter() @@ -646,48 +682,46 @@ async fn run_watch( .join(", ") ); - // Arm before the heavy rediscover so the previous cycle's - // output stays on screen while discovery + worker warmup - // happens. The reporter clears at the moment new content is - // about to land (warning, error, or run start), eliminating - // the blank-screen gap that's painful on large suites. reporter.arm_clear(); - // Time the full discovery work — `apply_changes` is the - // expensive part on large suites, so it has to be inside the - // measured window for `disc_dur` to mean anything. - let disc_start = Instant::now(); - let impact = discoverer.apply_changes(&paths); - let disc_dur = Some(disc_start.elapsed()); - if impact.paths.is_empty() { - debug!("watch: no eligible paths after discovery filtering"); + let discovery_start = Instant::now(); + let change_impact = discoverer.apply_changes(&paths); + let discovery_duration = Some(discovery_start.elapsed()); + + if change_impact.paths.is_empty() { + debug!("Watch: no eligible paths after discovery filtering"); continue; } - // When `--all` is set, rerun the full test set on every change instead - // of restricting to tests transitively affected by the changed files. - // Useful when the import graph misses dependencies (dynamic imports, - // string-referenced modules, external fixtures) or when debugging - // test ordering/flakiness. let raw_tests = if all_tests { discoverer.tests() } else { - impact.affected_tests + change_impact.affected_tests }; let tests = test_filter.apply(raw_tests); let hooks = discoverer.hooks(); + emit_discovery_warnings(reporter, &discoverer); - run_watch_cycle(reporter, tests, &hooks, &pool, maxfail, dist, disc_dur).await; + + run_watch_cycle( + reporter, + tests, + &hooks, + pool, + maxfail, + dist, + discovery_duration, + ) + .await; } - pool.shutdown(); Ok(()) } #[cfg(test)] mod interrupt_tests { - use std::future::{pending, ready}; + use std::future::ready; use super::*; @@ -710,25 +744,20 @@ mod interrupt_tests { #[tokio::test] async fn interruption_takes_priority_and_cleans_up_reporter() -> Result<()> { - let result = run_interruptibly( - ready(Ok::<_, anyhow::Error>(())), - ready(Ok::<(), std::io::Error>(())), - ) - .await; + let cancellation = CancellationToken::new(); + cancellation.cancel(); + let result = run_interruptibly(ready(Ok::<_, anyhow::Error>(())), &cancellation).await; let mut reporter = CleanupReporter::default(); - assert_eq!(resolve_interruptible(result, &mut reporter)?, None); + assert!(resolve_interruptible(result, &mut reporter)?.is_none()); assert_eq!(reporter.cleanup_calls, 1); Ok(()) } #[tokio::test] async fn completion_does_not_clean_up_reporter() -> Result<()> { - let result = run_interruptibly( - ready(Ok::<_, anyhow::Error>(42)), - pending::>(), - ) - .await; + let cancellation = CancellationToken::new(); + let result = run_interruptibly(ready(Ok::<_, anyhow::Error>(42)), &cancellation).await; let mut reporter = CleanupReporter::default(); assert_eq!(resolve_interruptible(result, &mut reporter)?, Some(42)); @@ -806,7 +835,7 @@ mod watch_tests { // Returns () — the important behavior is that it does NOT propagate the // underlying `report_cycle` Err that `tryke test` relies on for exit code. run_watch_cycle(&mut reporter, tests, &[], &pool, None, DistMode::Test, None).await; - pool.shutdown(); + pool.shutdown().await.expect("shut down worker pool"); Ok(()) } @@ -864,7 +893,7 @@ mod watch_tests { run_now, ) .await; - pool.shutdown(); + pool.shutdown().await.expect("shut down worker pool"); Ok(reporter) } @@ -906,8 +935,15 @@ mod watch_tests { #[cfg(test)] mod execution_tests { - use std::{io, path::PathBuf}; + use std::{ + io, + path::PathBuf, + sync::Arc, + time::{Duration, Instant}, + }; + use anyhow::Context as _; + use tokio::sync::Notify; use tryke_config::{Project, ProjectMetadata, TrykeOptions}; use tryke_discovery::{Discoverer, DiscoveryOptions}; use tryke_reporter::{ @@ -917,6 +953,28 @@ mod execution_tests { use super::*; + struct CancellationReporter { + started: Arc, + run_completes: usize, + cleanup_calls: usize, + } + + impl Reporter for CancellationReporter { + fn on_run_start(&mut self, _tests: &[tryke_types::TestItem]) { + self.started.notify_one(); + } + + fn on_test_complete(&mut self, _result: &tryke_types::TestResult) {} + + fn on_run_complete(&mut self, _summary: &RunSummary) { + self.run_completes += 1; + } + + fn cleanup(&mut self) { + self.cleanup_calls += 1; + } + } + fn configured_project(root: &std::path::Path) -> Project { let mut metadata = ProjectMetadata::new(root); metadata.apply_configuration_file(); @@ -960,6 +1018,7 @@ mod execution_tests { let fixture = TestProject::new()?; let project = configured_project(fixture.root()); let tests = discover_project(&project, DiscoveryOptions::default()); + let cancellation = CancellationToken::new(); let _ = run_tests( reporter, &project, @@ -971,6 +1030,7 @@ mod execution_tests { DistMode::Test, None, None, + &cancellation, ) .await; Ok(()) @@ -1040,7 +1100,7 @@ mod execution_tests { .await .is_ok() ); - pool.shutdown(); + pool.shutdown().await.expect("shut down worker pool"); Ok(()) } @@ -1064,7 +1124,7 @@ mod execution_tests { .await .is_ok() ); - pool.shutdown(); + pool.shutdown().await.expect("shut down worker pool"); Ok(()) } @@ -1081,6 +1141,7 @@ mod execution_tests { ..DiscoveryOptions::default() }, ); + let cancellation = CancellationToken::new(); assert!( run_tests( &mut reporter, @@ -1092,7 +1153,8 @@ mod execution_tests { 1, DistMode::Test, None, - None + None, + &cancellation, ) .await .is_ok() @@ -1100,6 +1162,58 @@ mod execution_tests { Ok(()) } + #[tokio::test] + async fn cancellation_interrupts_active_test_and_awaits_shutdown() -> anyhow::Result<()> { + let fixture = TestProject::with_files([( + "test_sleep.py", + "import time\nfrom tryke import test\n\n@test\ndef test_sleep():\n time.sleep(30)\n", + )])?; + let project = configured_project(fixture.root()); + let mut discoverer = Discoverer::new(&project); + let tests = discoverer.rediscover(); + let hooks = discoverer.hooks(); + + let started = Arc::new(Notify::new()); + let mut reporter = CancellationReporter { + started: Arc::clone(&started), + run_completes: 0, + cleanup_calls: 0, + }; + let cancellation = CancellationToken::new(); + let cancel = cancellation.clone(); + let cancel_task = tokio::spawn(async move { + started.notified().await; + tokio::time::sleep(Duration::from_millis(100)).await; + cancel.cancel(); + }); + + let start = Instant::now(); + let result = run_tests( + &mut reporter, + &project, + LevelFilter::Off, + tests, + &hooks, + None, + 1, + DistMode::Test, + None, + None, + &cancellation, + ) + .await; + cancel_task.await.context("Join cancellation task")?; + + assert!(resolve_interruptible(result, &mut reporter)?.is_none()); + assert_eq!(reporter.cleanup_calls, 1); + assert_eq!(reporter.run_completes, 0); + assert!( + start.elapsed() < Duration::from_secs(3), + "Cancellation should not wait for the sleeping Python test", + ); + Ok(()) + } + #[tokio::test] async fn integration_python_worker_runs_tests() -> io::Result<()> { let workspace_root = PathBuf::from(env!("CARGO_MANIFEST_DIR")) @@ -1160,7 +1274,7 @@ def test_failing(): ); } - pool.shutdown(); + pool.shutdown().await.expect("shut down worker pool"); Ok(()) } @@ -1202,7 +1316,7 @@ def test_failing(): result.is_ok(), "expected Ok when all tests pass, got {result:?}" ); - pool.shutdown(); + pool.shutdown().await.expect("shut down worker pool"); Ok(()) } @@ -1243,7 +1357,7 @@ def test_failing(): .expect("report_cycle should not error on test failures"); assert_eq!(summary.failed, 1, "expected one failed test"); assert_eq!(summary.passed, 0); - pool.shutdown(); + pool.shutdown().await.expect("shut down worker pool"); Ok(()) } } diff --git a/crates/tryke/src/lib.rs b/crates/tryke/src/lib.rs index c977e70..29f4cb1 100644 --- a/crates/tryke/src/lib.rs +++ b/crates/tryke/src/lib.rs @@ -1,8 +1,8 @@ mod cli; mod commands; +mod logging; use std::{ - env, io::{self, Write}, process::{ExitCode, Termination}, }; @@ -10,11 +10,13 @@ use std::{ use clap::{CommandFactory, Parser}; use console::style; use log::debug; +use tokio_util::sync::CancellationToken; use cli::{Cli, Commands}; use commands::{ CommandOrigin, run_clean_command, run_graph_command, run_server_command, run_test_command, }; +use logging::LogConfig; #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum ExitStatus { @@ -37,20 +39,14 @@ impl Termination for ExitStatus { } } -pub fn run() -> ExitStatus { - try_run().unwrap_or_else(report_error) +pub async fn run() -> ExitStatus { + try_run().await.unwrap_or_else(report_error) } -fn try_run() -> anyhow::Result { +async fn try_run() -> anyhow::Result { let cli = Cli::parse(); - let cli_filter = cli.global.verbose.log_level_filter(); - let tryke_log = env::var("TRYKE_LOG").ok(); - let rust_default = tryke_config::rust_log_default(tryke_log.as_deref(), cli_filter); - - env_logger::Builder::from_env( - env_logger::Env::default().default_filter_or(rust_default.as_str().to_ascii_lowercase()), - ) - .init(); + let logging = LogConfig::from_env(cli.global.verbose.log_level_filter())?; + logging.init_rust_logging(); debug!("{cli:?}"); @@ -64,8 +60,40 @@ fn try_run() -> anyhow::Result { let global = cli.global; match command { - Commands::Test(args) => run_test_command(args, &global, origin), - Commands::Server(args) => run_server_command(args, &global), + Commands::Test(args) => { + let cancellation = CancellationToken::new(); + let command = run_test_command(args, &global, origin, logging, cancellation.clone()); + tokio::pin!(command); + + tokio::select! { + biased; + signal = tokio::signal::ctrl_c() => { + signal?; + cancellation.cancel(); + command.await?; + + Ok(ExitStatus::Interrupted) + } + result = &mut command => result, + } + } + Commands::Server(args) => { + let cancellation = CancellationToken::new(); + let command = run_server_command(args, &global, logging, cancellation.clone()); + tokio::pin!(command); + + tokio::select! { + biased; + signal = tokio::signal::ctrl_c() => { + signal?; + cancellation.cancel(); + command.await?; + + Ok(ExitStatus::Interrupted) + } + result = &mut command => result, + } + } Commands::Clean(args) => run_clean_command(args, &global), Commands::Graph(args) => run_graph_command(args, &global), } diff --git a/crates/tryke/src/logging.rs b/crates/tryke/src/logging.rs new file mode 100644 index 0000000..0564f98 --- /dev/null +++ b/crates/tryke/src/logging.rs @@ -0,0 +1,137 @@ +use std::env; + +use anyhow::{Result, anyhow}; +use log::LevelFilter; +use tryke_reporter::Verbosity; + +const EXPECTED_LEVELS: &str = "off, error, warn, info, debug, or trace"; + +/// Process-wide logging configuration resolved from the CLI and environment. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct LogConfig { + level: LevelFilter, +} + +impl LogConfig { + pub(crate) fn from_env(cli_level: LevelFilter) -> Result { + match env::var("TRYKE_LOG") { + Ok(value) => Self::resolve(Some(&value), cli_level), + Err(env::VarError::NotPresent) => Self::resolve(None, cli_level), + Err(env::VarError::NotUnicode(_)) => Err(anyhow!( + "TRYKE_LOG must be valid Unicode; expected {EXPECTED_LEVELS}" + )), + } + } + + fn resolve(tryke_log: Option<&str>, cli_level: LevelFilter) -> Result { + let Some(value) = tryke_log else { + return Ok(Self { level: cli_level }); + }; + let level = value.trim().parse::().map_err(|_| { + anyhow!("invalid TRYKE_LOG value {value:?}; expected {EXPECTED_LEVELS}") + })?; + Ok(Self { level }) + } + + pub(crate) fn init_rust_logging(self) { + // `try_init` rather than `init`: installing a global logger twice is a + // benign no-op for us (only `try_run` calls this), and a panic here + // would take down the process before any command has run. + let _ = env_logger::Builder::from_env( + env_logger::Env::default().default_filter_or(self.level.as_str().to_ascii_lowercase()), + ) + .try_init(); + } + + pub(crate) fn level(self) -> LevelFilter { + self.level + } + + pub(crate) fn reporter_verbosity(self) -> Verbosity { + Verbosity::from_level_filter(self.level) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cli_level_is_used_without_tryke_log() { + for level in [ + LevelFilter::Off, + LevelFilter::Error, + LevelFilter::Warn, + LevelFilter::Info, + LevelFilter::Debug, + LevelFilter::Trace, + ] { + let config = LogConfig::resolve(None, level).expect("resolve log configuration"); + assert_eq!(config.level(), level); + } + } + + #[test] + fn tryke_log_overrides_cli_level() { + for (value, level) in [ + ("off", LevelFilter::Off), + ("error", LevelFilter::Error), + ("warn", LevelFilter::Warn), + ("info", LevelFilter::Info), + ("debug", LevelFilter::Debug), + ("trace", LevelFilter::Trace), + ] { + let config = + LogConfig::resolve(Some(value), LevelFilter::Off).expect("resolve TRYKE_LOG"); + assert_eq!(config.level(), level); + } + } + + #[test] + fn tryke_log_is_trimmed_and_case_insensitive() { + let config = + LogConfig::resolve(Some(" DeBuG "), LevelFilter::Warn).expect("resolve TRYKE_LOG"); + assert_eq!(config.level(), LevelFilter::Debug); + } + + #[test] + fn invalid_tryke_log_is_an_error() { + let error = LogConfig::resolve(Some("verbose"), LevelFilter::Warn) + .expect_err("invalid TRYKE_LOG should fail"); + assert_eq!( + error.to_string(), + "invalid TRYKE_LOG value \"verbose\"; expected off, error, warn, info, debug, or trace" + ); + } + + #[test] + fn reporter_verbosity_uses_resolved_level() { + assert!(matches!( + LogConfig { + level: LevelFilter::Off + } + .reporter_verbosity(), + Verbosity::Quiet + )); + assert!(matches!( + LogConfig { + level: LevelFilter::Error + } + .reporter_verbosity(), + Verbosity::Quiet + )); + assert!(matches!( + LogConfig { + level: LevelFilter::Warn + } + .reporter_verbosity(), + Verbosity::Normal + )); + for level in [LevelFilter::Info, LevelFilter::Debug, LevelFilter::Trace] { + assert!(matches!( + LogConfig { level }.reporter_verbosity(), + Verbosity::Verbose + )); + } + } +} diff --git a/crates/tryke/src/main.rs b/crates/tryke/src/main.rs index 35999c9..39b3a77 100644 --- a/crates/tryke/src/main.rs +++ b/crates/tryke/src/main.rs @@ -1,3 +1,4 @@ -fn main() -> tryke::ExitStatus { - tryke::run() +#[tokio::main] +async fn main() -> tryke::ExitStatus { + tryke::run().await } diff --git a/crates/tryke_config/Cargo.toml b/crates/tryke_config/Cargo.toml index b2c8c64..215c5d8 100644 --- a/crates/tryke_config/Cargo.toml +++ b/crates/tryke_config/Cargo.toml @@ -10,7 +10,6 @@ authors.workspace = true license.workspace = true [dependencies] -log = "0.4" serde = { workspace = true } toml = { workspace = true } diff --git a/crates/tryke_config/src/lib.rs b/crates/tryke_config/src/lib.rs index a18b1cc..c9a3d82 100644 --- a/crates/tryke_config/src/lib.rs +++ b/crates/tryke_config/src/lib.rs @@ -441,48 +441,6 @@ fn conda_environment_is_base(prefix: &Path) -> bool { prefix.file_name().is_none_or(|file_name| file_name != name) } -/// Default `RUST_LOG`-style filter directive when `RUST_LOG` is unset. -/// -/// `env_logger` honors `RUST_LOG` natively for fine-grained per-module -/// filtering; this only computes the fallback used when the user hasn't -/// set it. Precedence: `TRYKE_LOG` env > CLI flag. -/// -/// `TRYKE_LOG` is a bare level name (`off`/`error`/`warn`/`info`/`debug`/ -/// `trace`); `RUST_LOG`'s `tryke=info,hyper=warn` syntax is intentionally -/// not supported here — power users with that need just set `RUST_LOG`. -#[must_use] -pub fn rust_log_default(tryke_log_env: Option<&str>, cli: log::LevelFilter) -> log::LevelFilter { - parse_level(tryke_log_env).unwrap_or(cli) -} - -/// Level forwarded to spawned python workers via `TRYKE_LOG`. -/// -/// Precedence: `TRYKE_LOG` env (if set) > CLI flag (only when explicitly -/// more verbose than `Warn` — i.e., the user passed at least one `-v`). -/// Returns `Off` when the worker should stay silent; callers should not -/// set the env var on the child in that case so the worker preserves its -/// "no chatter unless asked" default. -/// -/// `RUST_LOG` is deliberately not consulted: it's a rust-specific -/// convention from `env_logger`, and silently translating its -/// per-module filter syntax into a python log level is a footgun. -/// `TRYKE_LOG` is the cross-language umbrella. -#[must_use] -pub fn worker_log_level(tryke_log_env: Option<&str>, cli: log::LevelFilter) -> log::LevelFilter { - if let Some(level) = parse_level(tryke_log_env) { - return level; - } - if cli > log::LevelFilter::Warn { - cli - } else { - log::LevelFilter::Off - } -} - -fn parse_level(s: Option<&str>) -> Option { - s.and_then(|v| v.trim().parse::().ok()) -} - fn parse_toml(contents: &str) -> Option { toml::from_str::(contents).ok()?.tool?.tryke } @@ -1092,60 +1050,4 @@ mod tests { let config = load_without_environment(dir.path(), TrykeOptions::default()); assert_eq!(config.python(), "C:foo\\python.exe"); } - - #[test] - fn rust_log_default_uses_tryke_log_env_when_set() { - let level = rust_log_default(Some("debug"), log::LevelFilter::Warn); - assert_eq!(level, log::LevelFilter::Debug); - } - - #[test] - fn rust_log_default_falls_back_to_cli_when_env_unset() { - let level = rust_log_default(None, log::LevelFilter::Info); - assert_eq!(level, log::LevelFilter::Info); - } - - #[test] - fn rust_log_default_falls_back_to_cli_when_env_unparseable() { - // Garbage values fall through rather than blowing up — RUST_LOG - // could carry per-module filters we don't try to interpret here. - let level = rust_log_default(Some("tryke=info,hyper=warn"), log::LevelFilter::Warn); - assert_eq!(level, log::LevelFilter::Warn); - } - - #[test] - fn worker_log_level_uses_tryke_log_env() { - let level = worker_log_level(Some("INFO"), log::LevelFilter::Warn); - assert_eq!(level, log::LevelFilter::Info); - } - - #[test] - fn worker_log_level_propagates_explicit_verbose_flag() { - let level = worker_log_level(None, log::LevelFilter::Debug); - assert_eq!(level, log::LevelFilter::Debug); - } - - #[test] - fn worker_log_level_stays_off_at_default_warn() { - // No env, no explicit `-v` → workers stay silent; preserves the - // pre-existing "no chatter unless asked" default for python. - let level = worker_log_level(None, log::LevelFilter::Warn); - assert_eq!(level, log::LevelFilter::Off); - } - - #[test] - fn worker_log_level_stays_off_when_quiet() { - // `-q` (Error) is even less verbose than the Warn default, so the - // worker definitely shouldn't be lit up. - let level = worker_log_level(None, log::LevelFilter::Error); - assert_eq!(level, log::LevelFilter::Off); - } - - #[test] - fn worker_log_level_env_wins_over_cli() { - // User explicitly set TRYKE_LOG, even though they also passed `-q` - // — env intent dominates the flag. - let level = worker_log_level(Some("info"), log::LevelFilter::Error); - assert_eq!(level, log::LevelFilter::Info); - } } diff --git a/crates/tryke_discovery/src/filesystem/cache.rs b/crates/tryke_discovery/src/filesystem/cache.rs index 18c9ca6..8cf554d 100644 --- a/crates/tryke_discovery/src/filesystem/cache.rs +++ b/crates/tryke_discovery/src/filesystem/cache.rs @@ -196,12 +196,12 @@ impl DiskCache { let entries = match Self::try_load(&path) { Ok(entries) => entries, Err(err) => { - trace!("discovery cache load failed ({err}): starting empty"); + trace!("Discovery cache load failed ({err}): starting empty"); HashMap::new() } }; debug!( - "discovery cache loaded {} entries from {}", + "Discovery cache loaded {} entries from {}", entries.len(), path.display() ); @@ -218,7 +218,7 @@ impl DiskCache { .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; if file.version != CACHE_VERSION { debug!( - "discovery cache version mismatch ({} vs {}): discarding", + "Discovery cache version mismatch ({} vs {}): discarding", file.version, CACHE_VERSION ); return Ok(HashMap::new()); @@ -283,7 +283,7 @@ impl DiskCache { fs::write(&tmp_path, &bytes)?; fs::rename(&tmp_path, path)?; debug!( - "discovery cache saved {} entries to {}", + "Discovery cache saved {} entries to {}", self.entries.len(), path.display() ); diff --git a/crates/tryke_discovery/src/filesystem/db.rs b/crates/tryke_discovery/src/filesystem/db.rs index 25ecf75..374942b 100644 --- a/crates/tryke_discovery/src/filesystem/db.rs +++ b/crates/tryke_discovery/src/filesystem/db.rs @@ -92,7 +92,7 @@ fn count_discover_file_execution(path: &std::path::Path) { pub(crate) fn parse_file(db: &dyn Db, file: SourceFile) -> ParsedAst { let path = file.path(db); trace!( - "parsing {}", + "Parsing {}", path.strip_prefix(file.root(db)).unwrap_or(path).display() ); ParsedAst::parse(file.text(db)) diff --git a/crates/tryke_discovery/src/filesystem/discoverer.rs b/crates/tryke_discovery/src/filesystem/discoverer.rs index 87287b8..c35a340 100644 --- a/crates/tryke_discovery/src/filesystem/discoverer.rs +++ b/crates/tryke_discovery/src/filesystem/discoverer.rs @@ -137,7 +137,7 @@ impl Discoverer { let mut paths = super::collect_python_files(&self.root, &self.excludes); paths.sort(); debug!( - "rediscover: found {} python files in {}", + "Rediscover: found {} Python files in {}", paths.len(), self.root.display() ); @@ -198,7 +198,7 @@ impl Discoverer { } } debug!( - "rediscover: cache hits {}/{} ({} parses pending)", + "Rediscover: cache hits {}/{} ({} parses pending)", hit_count, paths.len(), misses.len() @@ -256,7 +256,7 @@ impl Discoverer { } super::sort_tests(&mut tests); - debug!("rediscover: discovered {} tests total", tests.len()); + debug!("Rediscover: discovered {} tests total", tests.len()); tests } @@ -281,7 +281,7 @@ impl Discoverer { let paths = super::collect_python_files_restricted(&self.root, walk_roots, &self.excludes); debug!( - "rediscover_restricted: found {} python files across {} walk roots", + "Rediscover_restricted: found {} Python files across {} walk roots", paths.len(), walk_roots.len() ); @@ -326,7 +326,7 @@ impl Discoverer { } } debug!( - "rediscover_restricted: cache hits {}/{} ({} parses pending)", + "Rediscover_restricted: cache hits {}/{} ({} parses pending)", hit_count, paths.len(), misses.len() @@ -377,7 +377,7 @@ impl Discoverer { .collect(); super::sort_tests(&mut tests); debug!( - "rediscover_restricted: discovered {} tests total", + "Rediscover_restricted: discovered {} tests total", tests.len() ); tests @@ -441,11 +441,11 @@ impl Discoverer { fn upsert_source(&mut self, path: &Path, text: String) { if let Some(file) = self.inputs.get(path) { if file.text(&self.db) != &text { - trace!("rediscover: re-parsing changed file {}", path.display()); + trace!("Rediscover: re-parsing changed file {}", path.display()); file.set_text(&mut self.db).to(text); } } else { - trace!("rediscover: parsing new file {}", path.display()); + trace!("Rediscover: parsing new file {}", path.display()); let file = SourceFile::new( &self.db, text, @@ -536,7 +536,7 @@ impl Discoverer { pub fn rediscover_changed(&mut self, changed: &[PathBuf]) -> Vec { let changed = Self::canonicalize_paths(changed); debug!( - "rediscover_changed: processing {} changed paths", + "Rediscover_changed: processing {} changed paths", changed.len() ); let mut touched: Vec = Vec::new(); @@ -549,7 +549,7 @@ impl Discoverer { touched.push(path.clone()); } else { trace!( - "rediscover_changed: removing deleted file {}", + "Rediscover_changed: removing deleted file {}", path.display() ); self.import_graph.remove(path); @@ -590,7 +590,7 @@ impl Discoverer { .flat_map(|r| r.parsed.tests.clone()) .collect(); super::sort_tests(&mut tests); - debug!("rediscover_changed: {} tests after update", tests.len()); + debug!("Rediscover_changed: {} tests after update", tests.len()); tests } @@ -628,7 +628,7 @@ impl Discoverer { .collect(); modules.sort(); debug!( - "affected_modules: {:?} → {:?}", + "Affected_modules: {:?} → {:?}", changed .iter() .map(|p| p.display().to_string()) @@ -651,7 +651,7 @@ impl Discoverer { }) .collect(); debug!( - "tests_for_changed: {:?} → {} tests", + "Tests_for_changed: {:?} → {} tests", Self::canonicalize_paths(changed) .iter() .map(|p| p.display().to_string()) diff --git a/crates/tryke_discovery/src/filesystem/import_graph.rs b/crates/tryke_discovery/src/filesystem/import_graph.rs index 3ef725f..ea8acee 100644 --- a/crates/tryke_discovery/src/filesystem/import_graph.rs +++ b/crates/tryke_discovery/src/filesystem/import_graph.rs @@ -83,7 +83,7 @@ impl ImportGraph { for importer in importers { if visited.insert(importer.clone()) { trace!( - "import_graph: {} invalidated by change to {}", + "Import_graph: {} invalidated by change to {}", importer.display(), file.display() ); diff --git a/crates/tryke_discovery/src/filesystem/mod.rs b/crates/tryke_discovery/src/filesystem/mod.rs index 3123fe9..7327eda 100644 --- a/crates/tryke_discovery/src/filesystem/mod.rs +++ b/crates/tryke_discovery/src/filesystem/mod.rs @@ -113,7 +113,7 @@ pub(crate) fn discover_file_from_ast( parsed: &db::ParsedAst, ) -> tryke_types::DiscoveredFile { let Some(module) = parsed.syntax() else { - trace!("parse error in {}", file.display()); + trace!("Parse error in {}", file.display()); return tryke_types::DiscoveredFile::default(); }; let result = crate::source::discover_file_from_body( diff --git a/crates/tryke_discovery/src/selection.rs b/crates/tryke_discovery/src/selection.rs index c6c648e..95ebbb2 100644 --- a/crates/tryke_discovery/src/selection.rs +++ b/crates/tryke_discovery/src/selection.rs @@ -80,7 +80,7 @@ impl Discoverer { } if !options.paths.is_empty() && !options.changed && !options.changed_first { - debug!("path-restricted discovery: falling back to full discovery"); + debug!("Path-restricted discovery: falling back to full discovery"); } let all_tests = self.rediscover(); @@ -180,7 +180,7 @@ fn resolve_walk_roots(root: &Path, path_specs: &[PathSpec]) -> Option, - hook_cache: HashMap, - /// Most recent spawn or hook-replay failure, captured so - /// `run_single_test` can surface the real reason (and any worker - /// stderr) instead of the opaque "worker unavailable" placeholder. - /// Cleared once we have a live worker again so we don't replay a - /// stale error against an unrelated test. - last_failure: Option, -} - -impl WorkerState { - fn new() -> Self { - Self { - process: None, - hook_cache: HashMap::new(), - last_failure: None, - } - } -} - -/// Build the user-facing message for a worker error, appending captured -/// stderr (if any) so the python-side traceback is visible to the user -/// without needing to enable debug logging. The actual python crash — -/// e.g. `ModuleNotFoundError: No module named 'tryke'` when the worker -/// venv is missing the package — only ever appears on the worker's -/// stderr pipe, so suppressing it here is what makes spawn failures look -/// like an opaque "worker unavailable". -fn format_worker_failure(prefix: &str, err: &dyn std::fmt::Display, stderr: &str) -> String { - let mut msg = format!("{prefix}: {err}"); - let trimmed = stderr.trim(); - if !trimmed.is_empty() { - msg.push_str("\nworker stderr:\n"); - msg.push_str(trimmed); - } - msg -} - -enum WorkerMsg { - Unit(WorkUnit, mpsc::UnboundedSender), - Shutdown, -} - -/// Control messages delivered on a per-worker channel. -/// -/// `Ping` and `Restart` are fan-out operations: every worker must -/// receive exactly one. Routing them through the shared work-stealing -/// channel would let a single fast worker grab all N messages while -/// other workers remained on stale Python processes — defeating the -/// guarantee that watch/server-mode reloads run on a fresh interpreter. -/// A dedicated channel per worker eliminates that race. -enum WorkerCtrl { - Ping(oneshot::Sender<()>), - Restart(oneshot::Sender<()>), -} +const WORKER_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(1); pub struct WorkerPool { /// Sender channel for work units @@ -90,6 +30,39 @@ pub struct WorkerPool { /// /// One per-worker to distribute messages to all. ctrl_txs: Vec>, + + /// Pool-wide shutdown signal, observed ahead of both work and control messages. + shutdown: CancellationToken, + + /// Every worker task remains owned by the pool until shutdown completes. + /// + /// A `JoinSet` rather than `tokio_util`'s `TaskTracker`: the pool is a + /// fixed-size set of tasks whose `JoinError`s we want to report, and + /// `JoinSet` also gives us `abort_all` for the `Drop` path. `TaskTracker` + /// discards task outcomes, which would turn a panicking worker into a + /// silently short run instead of a shutdown error. + workers: JoinSet<()>, +} + +#[must_use = "dropping a worker run cancels its submitted work"] +pub struct WorkerRun { + results: UnboundedReceiverStream, + cancel: CancellationToken, +} + +impl Stream for WorkerRun { + type Item = TestResult; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + Pin::new(&mut this.results).poll_next(cx) + } +} + +impl Drop for WorkerRun { + fn drop(&mut self) { + self.cancel.cancel(); + } } #[derive(Clone, Copy, Debug)] @@ -101,7 +74,6 @@ pub struct WorkerPoolOptions<'a> { } impl WorkerPool { - /// Spawn a worker pool using the project's resolved root and Python interpreter. pub async fn spawn(project: &Project, options: WorkerPoolOptions<'_>) -> Self { Self::spawn_from_parts( options.size, @@ -114,17 +86,6 @@ impl WorkerPool { .await } - /// Spawns a pool whose workers receive `TRYKE_LOG=`. - /// - /// Pass `LevelFilter::Off` to leave workers silent (the env var is - /// then not set on the child, so the worker's - /// `_configure_logging_from_env` no-ops). Production callers use - /// `tryke_config::worker_log_level` to derive this from CLI flags - /// + the `TRYKE_LOG` env var. - /// - /// `python_path` overrides the default path of `root` plus its `python` - /// directory when present. If `warm` is true, this method also waits for - /// every Python subprocess to start before returning. pub async fn spawn_from_parts( size: usize, python_bin: &str, @@ -151,721 +112,314 @@ impl WorkerPool { ); let (work_tx, work_rx) = async_channel::unbounded(); let mut ctrl_txs = Vec::with_capacity(size); + let shutdown = CancellationToken::new(); + let mut workers = JoinSet::new(); for _ in 0..size { - let bin = python_bin.clone(); let work_rx = work_rx.clone(); + let shutdown = shutdown.clone(); let (ctrl_tx, ctrl_rx) = mpsc::unbounded_channel(); ctrl_txs.push(ctrl_tx); - - tokio::spawn(worker_task( - bin, + let worker = Worker::new( + python_bin.clone(), python_path.clone(), root.clone(), log_level, - work_rx, - ctrl_rx, - )); + ); + + workers.spawn(worker.run(work_rx, ctrl_rx, shutdown)); } - let pool = Self { work_tx, ctrl_txs }; + let pool = Self { + work_tx, + ctrl_txs, + shutdown, + workers, + }; - if warm { - pool.warm().await; + // Pre-warming is best-effort: a worker that misses the deadline just + // pays interpreter startup on its first unit. Unlike a missed restart + // it cannot leave stale code behind, so it must not fail construction. + if warm && let Err(error) = pool.warm().await { + warn!("{error:#}"); } pool } - /// Submit any number of `WorkUnit`s to the worker pool - /// - /// A `WorkUnit` is an atomic group of tests to be run sequentially on a single worker - /// Returns a stream - pub fn submit(&self, units: Vec) -> impl Stream + use<> { + pub fn submit(&self, units: Vec) -> WorkerRun { let (stream_tx, stream_rx) = mpsc::unbounded_channel(); + let cancel = CancellationToken::new(); for unit in units { - let _ = self - .work_tx - .send_blocking(WorkerMsg::Unit(unit, stream_tx.clone())); + // `try_send`, not `send_blocking`: every caller is async, and + // async-channel documents `send_blocking` as deadlock-prone in an + // async context. On an unbounded channel the only failure is a + // closed channel (the pool is shutting down); the unit's + // `result_tx` clone drops with the message, so the run's stream + // ends early instead of hanging. + let _ = self.work_tx.try_send(WorkerMsg::Unit { + unit, + result_tx: stream_tx.clone(), + cancel: cancel.clone(), + }); } - UnboundedReceiverStream::new(stream_rx) + WorkerRun { + results: UnboundedReceiverStream::new(stream_rx), + cancel, + } } - /// Send one ctrl message per worker and await every ack. - /// - /// `build` is the ctrl-variant constructor (e.g. `WorkerCtrl::Ping`) - /// — taking it as a function pointer lets `warm` and - /// `restart_workers` share this whole fan-out path. - /// - /// If a worker task has died (ctrl channel receiver dropped), its - /// `send` returns `Err`; we skip its ack rather than push a future - /// that will never resolve, which would hang the watcher/server. + /// Send `build`'s control message to every worker, returning the indices + /// of workers that did not acknowledge before the shared deadline. async fn fanout_ctrl_with_timeout( &self, build: fn(oneshot::Sender<()>) -> WorkerCtrl, timeout: Duration, - ) -> bool { - let mut ack_rxs = Vec::with_capacity(self.ctrl_txs.len()); - for ctrl_tx in &self.ctrl_txs { + ) -> Vec { + let mut pending = Vec::with_capacity(self.ctrl_txs.len()); + for (index, ctrl_tx) in self.ctrl_txs.iter().enumerate() { let (ack_tx, ack_rx) = oneshot::channel(); if ctrl_tx.send(build(ack_tx)).is_ok() { - ack_rxs.push(ack_rx); + pending.push((index, ack_rx)); } } - tokio::time::timeout(timeout, async { - for ack_rx in ack_rxs { - let _ = ack_rx.await; + + // One deadline for the whole fan-out: the messages are already in + // flight, so awaiting the acks in sequence costs no extra wall time. + let deadline = tokio::time::Instant::now() + timeout; + let mut unacked = Vec::new(); + for (index, ack_rx) in pending { + // A dropped ack sender means the worker task exited and killed its + // interpreter on the way out — nothing stale is left behind, so + // that is not a control failure. Only the deadline counts. + if tokio::time::timeout_at(deadline, ack_rx).await.is_err() { + unacked.push(index); } - }) - .await - .is_ok() + } + unacked } - async fn fanout_ctrl(&self, operation: &str, build: fn(oneshot::Sender<()>) -> WorkerCtrl) { - if !self - .fanout_ctrl_with_timeout(build, WORKER_CONTROL_TIMEOUT) - .await - { - // A control timeout means at least one worker did not ack a - // restart/warm, undermining the "fresh workers each run" - // guarantee — surface it at warn so it's visible by default. - warn!( - "worker control operation '{operation}' timed out after {WORKER_CONTROL_TIMEOUT:?}" - ); + /// `timeout` is a parameter rather than a direct read of + /// [`WORKER_CONTROL_TIMEOUT`] so the failure path stays testable without a + /// five-second wall-clock wait. + async fn fanout_ctrl( + &self, + operation: &str, + build: fn(oneshot::Sender<()>) -> WorkerCtrl, + timeout: Duration, + ) -> Result<()> { + let unacked = self.fanout_ctrl_with_timeout(build, timeout).await; + if unacked.is_empty() { + return Ok(()); } + Err(anyhow!( + "worker control operation '{operation}' timed out after {timeout:?}; \ + {}/{} workers did not acknowledge (workers {unacked:?})", + unacked.len(), + self.ctrl_txs.len(), + )) } - /// Replace every worker subprocess with a fresh, responsive process. + /// Replace every Python subprocess with a clean, pre-warmed process. /// - /// This is how watch and server mode pick up code changes: rather than - /// trying to mutate a live interpreter with `importlib.reload` (which is - /// brittle once classes/closures/decorator-bound state from the old - /// definitions are referenced from elsewhere), we drop the whole process - /// and let it re-import everything on the next `run_test`. The fresh - /// process replays cached `register_hooks` calls so fixtures keep - /// working — same path as crash recovery. - pub async fn restart_workers(&self) { - self.fanout_ctrl("restart", WorkerCtrl::Restart).await; - self.warm().await; - } - - /// Pre-spawn all worker processes in parallel so Python startup - /// latency is not on the critical path of the first tests. - async fn warm(&self) { - self.fanout_ctrl("warm", WorkerCtrl::Ping).await; - } - - pub fn shutdown(self) { - for _ in 0..self.ctrl_txs.len() { - let _ = self.work_tx.send_blocking(WorkerMsg::Shutdown); + /// Hook metadata belongs to work units, so the fresh processes stay + /// unconfigured until their next unit installs its current registrations. + /// + /// # Errors + /// + /// Returns an error if a worker fails to acknowledge the restart within + /// [`WORKER_CONTROL_TIMEOUT`]. A worker only misses that deadline while it + /// is still executing a unit, which means it is still holding the *old* + /// interpreter — so the caller must not present the next run's results as + /// reflecting current source. + pub async fn restart_workers(&self) -> Result<()> { + self.fanout_ctrl("restart", WorkerCtrl::Restart, WORKER_CONTROL_TIMEOUT) + .await?; + + // Warming is an optimization, not a correctness guarantee — see the + // note in `spawn_from_parts`. + if let Err(error) = self.warm().await { + warn!("{error:#}"); } - } -} -pub use tryke_types::path_to_module; - -async fn spawn_worker_process( - python_bin: &str, - path_refs: &[&Path], - root: &Path, - log_level: LevelFilter, -) -> Result { - let python_bin = python_bin.to_owned(); - let python_paths = path_refs - .iter() - .map(|path| (*path).to_path_buf()) - .collect::>(); - let root = root.to_path_buf(); - let spawn = tokio::task::spawn_blocking(move || { - let path_refs = python_paths - .iter() - .map(PathBuf::as_path) - .collect::>(); - WorkerProcess::spawn(&python_bin, &path_refs, &root, log_level) - }); - - match tokio::time::timeout(WORKER_SPAWN_TIMEOUT, spawn).await { - Ok(Ok(result)) => result, - Ok(Err(error)) => Err(anyhow!("worker spawn task failed: {error}")), - Err(_) => Err(anyhow!( - "worker process spawn timed out after {WORKER_SPAWN_TIMEOUT:?}" - )), + Ok(()) } -} -/// Ensure a worker process is live, spawning one if needed and replaying -/// every cached `register_hooks` call before returning it. Replay guarantees -/// that after a crash-and-respawn, subsequent tests still see their -/// fixtures — without replay, the fresh worker would have empty hook -/// metadata and silently skip `before_each` / `after_each`. -async fn ensure_worker<'a>( - state: &'a mut WorkerState, - python_bin: &str, - path_refs: &[&Path], - root: &Path, - log_level: LevelFilter, -) -> Option<&'a mut WorkerProcess> { - if state.process.is_some() { - return state.process.as_mut(); - } - trace!("worker_task: spawning process"); - let mut w = match spawn_worker_process(python_bin, path_refs, root, log_level).await { - Ok(w) => w, - Err(e) => { - let msg = format_worker_failure( - &format!("failed to spawn python worker ({python_bin} -m tryke.worker)"), - &e, - "", - ); - debug!("worker_task: {msg}"); - state.last_failure = Some(msg); - return None; - } - }; - for (module, params) in &state.hook_cache { - if let Err(e) = w.register_hooks(params.clone()).await { - // Drain stderr before dropping the dead worker — otherwise - // the python traceback that explains *why* replay failed - // (e.g. ModuleNotFoundError on the worker side) goes with - // it and the user sees only "Broken pipe". - let stderr_output = w.drain_stderr().await; - let msg = format_worker_failure( - &format!("hook replay failed for module {module}"), - &e, - &stderr_output, - ); - debug!("worker_task: {msg}"); - state.last_failure = Some(msg); - // Worker is in an inconsistent state (some modules registered, - // some not). Drop it so the next attempt starts from scratch - // rather than silently running tests without fixtures. - return None; - } + async fn warm(&self) -> Result<()> { + self.fanout_ctrl("warm", WorkerCtrl::Ping, WORKER_CONTROL_TIMEOUT) + .await } - state.last_failure = None; - state.process = Some(w); - state.process.as_mut() -} -/// Execute a single test on the worker. On any RPC error we respawn the -/// worker (replaying cached hooks) for the next test but do NOT retry the -/// failing test — a retry could double-execute side effects if the test -/// partially ran before the crash. The failing test is surfaced as -/// `TestOutcome::Error` with the worker's stderr attached for diagnosis. -async fn run_single_test( - state: &mut WorkerState, - python_bin: &str, - path_refs: &[&Path], - root: &Path, - log_level: LevelFilter, - test: tryke_types::TestItem, - result_tx: &mpsc::UnboundedSender, -) { - let Some(w) = ensure_worker(state, python_bin, path_refs, root, log_level).await else { - let message = state - .last_failure - .clone() - .unwrap_or_else(|| "worker unavailable (spawn or hook replay failed)".into()); - let _ = result_tx.send(TestResult { - test, - outcome: TestOutcome::Error { message }, - duration: Duration::ZERO, - stdout: String::new(), - stderr: String::new(), - }); - return; - }; - match w.run_test(&test).await { - Ok(result) => { - trace!("worker_task: test {} done", test.name); - let _ = result_tx.send(result); - } - Err(err) => { - debug!("worker_task: run_test error for {}: {err}", test.name); - let stderr_output = w.drain_stderr().await; - // Drop the dead worker; the next call to `ensure_worker` will - // spawn a fresh one and replay cached hooks so the remaining - // tests in this unit keep their fixtures. - state.process = None; - let message = format_worker_failure("worker error", &err, &stderr_output); - let _ = result_tx.send(TestResult { - test, - outcome: TestOutcome::Error { message }, - duration: Duration::ZERO, - stdout: String::new(), - stderr: stderr_output, - }); - } - } -} + /// Shut down every worker process and await every worker task. + /// + /// Workers get one second *in total* to terminate cleanly — they stop + /// concurrently, so a per-worker budget would multiply the wait by the + /// pool size. Tasks still running past that deadline are aborted. + /// + /// # Errors + /// + /// Returns an error if a worker task panicked, or if any task had to be + /// aborted because it outlived the shutdown deadline. + pub async fn shutdown(mut self) -> Result<()> { + self.shutdown.cancel(); -/// Send `register_hooks` to the worker for each unique module in the work -/// unit, caching the call so any respawn later in the unit can replay it. -async fn register_hooks_for_unit( - state: &mut WorkerState, - python_bin: &str, - path_refs: &[&Path], - root: &Path, - log_level: LevelFilter, - hooks: &[HookItem], - tests: &[tryke_types::TestItem], -) { - let mut seen = std::collections::HashSet::new(); - for test in tests { - if !seen.insert(test.module_path.clone()) { - continue; - } - let module_hooks: Vec = hooks - .iter() - .filter(|h| h.module_path == test.module_path) - .map(|h| crate::protocol::HookWire { - name: h.name.clone(), - per: serde_json::to_value(h.per) - .ok() - .and_then(|v| v.as_str().map(String::from)) - .unwrap_or_default(), - groups: h.groups.clone(), - depends_on: h.depends_on.clone(), - line_number: h.line_number, - }) - .collect(); - - if module_hooks.is_empty() { - continue; - } + let mut workers = std::mem::take(&mut self.workers); + let mut failures = Vec::new(); - let params = RegisterHooksParams { - module: test.module_path.clone(), - hooks: module_hooks, - }; - // Cache before sending so a respawn that races with this call can - // still replay the correct hooks. - state - .hook_cache - .insert(test.module_path.clone(), params.clone()); - - let Some(w) = ensure_worker(state, python_bin, path_refs, root, log_level).await else { - continue; - }; - if let Err(e) = w.register_hooks(params).await { - // Drain stderr before dropping so the python traceback that - // killed the worker reaches the user-facing test result via - // `last_failure`, instead of being lost with the process. - let stderr_output = w.drain_stderr().await; - let msg = format_worker_failure( - &format!("register_hooks failed for module {}", test.module_path), - &e, - &stderr_output, - ); - debug!("worker_task: {msg}"); - state.last_failure = Some(msg); - // Worker is potentially wedged; drop it so the next test - // forces a respawn-with-replay. - state.process = None; - } - } -} - -async fn handle_ctrl( - state: &mut WorkerState, - python_bin: &str, - path_refs: &[&Path], - root: &Path, - log_level: LevelFilter, - ctrl: WorkerCtrl, -) { - match ctrl { - WorkerCtrl::Ping(ack_tx) => { - trace!("worker_task: ping (pre-warm)"); - let _ = ensure_worker(state, python_bin, path_refs, root, log_level).await; - let _ = ack_tx.send(()); - } - WorkerCtrl::Restart(ack_tx) => { - trace!("worker_task: restart"); - if let Some(mut w) = state.process.take() { - w.shutdown().await; + let drained = tokio::time::timeout(WORKER_SHUTDOWN_TIMEOUT, async { + while let Some(result) = workers.join_next().await { + record_join_result(result, &mut failures); } - let _ = ack_tx.send(()); + }) + .await; + + if drained.is_err() { + // Anything still alive is wedged (most likely in Python teardown). + // `shutdown` aborts the remainder and awaits the aborts, so the + // child processes are killed by `WorkerProcess::drop` before we + // return rather than outliving the pool. + let remaining = workers.len(); + workers.shutdown().await; + failures.push(format!( + "{remaining} worker task(s) did not stop within \ + {WORKER_SHUTDOWN_TIMEOUT:?} and were aborted" + )); } - } -} -async fn handle_unit( - state: &mut WorkerState, - python_bin: &str, - path_refs: &[&Path], - root: &Path, - log_level: LevelFilter, - unit: WorkUnit, - result_tx: mpsc::UnboundedSender, -) { - if !unit.hooks.is_empty() { - register_hooks_for_unit( - state, - python_bin, - path_refs, - root, - log_level, - &unit.hooks, - &unit.tests, - ) - .await; - } - let finalize_modules: std::collections::HashSet = - unit.tests.iter().map(|t| t.module_path.clone()).collect(); - for test in unit.tests { - trace!("worker_task: running test {}", test.name); - run_single_test( - state, python_bin, path_refs, root, log_level, test, &result_tx, - ) - .await; - } - for module in finalize_modules { - if let Some(w) = state.process.as_mut() - && let Err(e) = w.finalize_hooks(module).await - { - debug!("worker_task: finalize_hooks failed: {e}"); + if failures.is_empty() { + Ok(()) + } else { + Err(anyhow!( + "Worker pool shutdown encountered task failures: {}", + failures.join("; ") + )) } } } -async fn worker_task( - python_bin: String, - python_path: Vec, - root: PathBuf, - log_level: LevelFilter, - work_rx: async_channel::Receiver, - mut ctrl_rx: mpsc::UnboundedReceiver, -) { - let path_refs: Vec<&Path> = python_path.iter().map(PathBuf::as_path).collect(); - let mut state = WorkerState::new(); - - loop { - // `biased` guarantees control messages take priority once the - // current Unit (if any) finishes. Without it, `select!` could - // keep picking Units off the shared queue while a Restart sat - // in this worker's ctrl channel — leaving the worker on a stale - // interpreter for arbitrarily long. - tokio::select! { - biased; - ctrl = ctrl_rx.recv() => { - let Some(ctrl) = ctrl else { break }; - handle_ctrl(&mut state, &python_bin, &path_refs, &root, log_level, ctrl).await; - } - msg = work_rx.recv() => { - match msg { - Ok(WorkerMsg::Unit(unit, result_tx)) => { - handle_unit( - &mut state, - &python_bin, - &path_refs, - &root, - log_level, - unit, - result_tx, - ) - .await; - } - Ok(WorkerMsg::Shutdown) | Err(_) => break, - } - } - } +impl Drop for WorkerPool { + fn drop(&mut self) { + self.shutdown.cancel(); + self.workers.abort_all(); } +} - if let Some(mut w) = state.process.take() { - w.shutdown().await; +fn record_join_result(result: std::result::Result<(), JoinError>, failures: &mut Vec) { + if let Err(error) = result { + let kind = if error.is_panic() { + "panicked" + } else { + "failed to join" + }; + failures.push(format!("worker task {kind}: {error}")); } } #[cfg(test)] mod tests { - use std::path::PathBuf; - + use log::LevelFilter; use tokio_stream::StreamExt; - use tryke_testing::{TestProject, python_bin as test_python_bin}; - use tryke_types::{FixturePer, HookItem, TestItem}; + use tryke_testing::{TestProject, python_bin as test_python_bin, workspace_root}; + use tryke_types::TestItem; use super::*; - use crate::schedule::WorkUnit; - - fn workspace_root() -> PathBuf { - PathBuf::from(env!("CARGO_MANIFEST_DIR")) - .join("../..") - .canonicalize() - .expect("workspace root") - } + use crate::schedule::{DistMode, partition_with_hooks}; fn python_package_dir() -> PathBuf { workspace_root().join("python") } - fn make_test_item(module: &str, name: &str, file: &std::path::Path) -> TestItem { - TestItem { - name: name.to_string(), - module_path: module.to_string(), - file_path: Some(file.to_path_buf()), - ..TestItem::default() + /// A pool wired up with real control channels but no worker tasks, so the + /// control plane can be exercised without paying for interpreter startup. + /// + /// The returned receivers must be kept alive by the caller: dropping a + /// `ctrl_rx` closes its channel, which `fanout_ctrl_with_timeout` treats as + /// "worker already gone" rather than as a missed acknowledgement. + fn control_only_pool( + size: usize, + ) -> ( + WorkerPool, + async_channel::Receiver, + Vec>, + ) { + let (work_tx, work_rx) = async_channel::unbounded(); + let mut senders = Vec::with_capacity(size); + let mut receivers = Vec::with_capacity(size); + for _ in 0..size { + let (ctrl_tx, ctrl_rx) = mpsc::unbounded_channel(); + senders.push(ctrl_tx); + receivers.push(ctrl_rx); } - } - - /// End-to-end crash-recovery test: a middle test crashes the worker; - /// the failure must surface as `TestOutcome::Error` for exactly that - /// test, subsequent tests in the unit must still run with their - /// fixtures (hooks replayed on respawn), and the crashing test must - /// NOT be retried (no double-execution of side effects). - #[tokio::test] - async fn worker_crash_replays_hooks_and_does_not_double_execute() { - let project = TestProject::new().expect("create test project"); - let crash_counter = project.root().join("CRASH_COUNT"); - let crash_counter_escaped = crash_counter.to_string_lossy().replace('\\', "\\\\"); - let source = format!( - r#"from tryke import test, fixture, Depends, expect - -@fixture -def counter() -> int: - return 42 - -@test -def test_first(n: int = Depends(counter)) -> None: - expect(n).to_equal(42) - -@test -def test_crasher() -> None: - import os - with open("{crash_counter_escaped}", "a") as f: - f.write("x") - f.flush() - os._exit(1) - -@test -def test_third(n: int = Depends(counter)) -> None: - expect(n).to_equal(42) -"# - ); - let test_file = project - .write("test_crash.py", source) - .expect("write test file"); - - let hook = HookItem { - name: "counter".into(), - module_path: "test_crash".into(), - per: FixturePer::Test, - groups: vec![], - depends_on: vec![], - line_number: None, - }; - let tests = vec![ - make_test_item("test_crash", "test_first", &test_file), - make_test_item("test_crash", "test_crasher", &test_file), - make_test_item("test_crash", "test_third", &test_file), - ]; - let unit = WorkUnit { - tests, - hooks: vec![hook], + let pool = WorkerPool { + work_tx, + ctrl_txs: senders, + shutdown: CancellationToken::new(), + workers: JoinSet::new(), }; - let python_path = [project.root().to_path_buf(), python_package_dir()]; - let pool = WorkerPool::spawn_from_parts( - 1, - &test_python_bin(), - project.root(), - Some(&python_path), - LevelFilter::Off, - true, - ) - .await; - - let mut results: Vec = pool.submit(vec![unit]).collect().await; - results.sort_by_key(|r| r.test.name.clone()); + (pool, work_rx, receivers) + } - assert_eq!(results.len(), 3, "expected 3 results, got {results:?}"); - // Sorted: test_crasher, test_first, test_third - let crasher = &results[0]; - let first = &results[1]; - let third = &results[2]; + #[tokio::test] + async fn fanout_ctrl_reports_every_worker_that_never_acks() { + let (pool, _work_rx, _ctrl_rxs) = control_only_pool(3); - assert_eq!(crasher.test.name, "test_crasher"); - assert!( - matches!(crasher.outcome, TestOutcome::Error { .. }), - "crasher should be Error, got {:?}", - crasher.outcome - ); - assert!( - matches!(first.outcome, TestOutcome::Passed), - "first should pass (fixture wired), got {:?}", - first.outcome - ); - assert!( - matches!(third.outcome, TestOutcome::Passed), - "third should pass after respawn+hook-replay, got {:?}", - third.outcome - ); + let unacked = pool + .fanout_ctrl_with_timeout(WorkerCtrl::Restart, Duration::from_millis(10)) + .await; - let count = std::fs::read_to_string(&crash_counter).unwrap_or_default(); assert_eq!( - count.len(), - 1, - "crashing test must run exactly once (no retry), got {count:?}" + unacked, + vec![0, 1, 2], + "a silent worker must be identified, not just counted" ); - - pool.shutdown(); } - /// Restarting the pool must yield a *fresh* Python interpreter — not - /// just an `importlib.reload`-mutated module. We prove this by - /// recording one tally mark per fresh import of the test module: the - /// module body increments a sidecar counter on every initial load. - /// Importlib.reload would re-run the body too, but in production it - /// leaves classes/closures bound to the old definitions in *other* - /// modules — the brittleness this rearchitecture exists to fix. - /// A second tally after `restart_workers` confirms a brand new - /// interpreter is in play. #[tokio::test] - async fn restart_workers_runs_module_body_on_fresh_interpreter() { - let project = TestProject::new().expect("create test project"); - - let counter_file = project.root().join("IMPORT_COUNT"); - let counter_escaped = counter_file.to_string_lossy().replace('\\', "\\\\"); - let source = format!( - r#"from tryke import test, expect - -with open("{counter_escaped}", "a") as f: - f.write("x") - f.flush() - -@test -def test_noop() -> None: - expect(1).to_equal(1) -"# - ); - let test_file = project - .write("test_restart_state.py", source) - .expect("write test file"); - - let make_unit = || WorkUnit { - tests: vec![make_test_item( - "test_restart_state", - "test_noop", - &test_file, - )], - hooks: vec![], - }; + async fn fanout_ctrl_treats_a_departed_worker_as_acknowledged() { + // A closed control channel means the worker task exited and killed its + // interpreter on the way out. Nothing stale survives it, so it must not + // be reported as a control failure. + let (pool, _work_rx, ctrl_rxs) = control_only_pool(2); + drop(ctrl_rxs); + + let unacked = pool + .fanout_ctrl_with_timeout(WorkerCtrl::Restart, Duration::from_millis(10)) + .await; - let python_path = [project.root().to_path_buf(), python_package_dir()]; - let pool = WorkerPool::spawn_from_parts( - 1, - &test_python_bin(), - project.root(), - Some(&python_path), - LevelFilter::Off, - true, - ) - .await; + assert!(unacked.is_empty(), "got {unacked:?}"); + } - let r1: Vec = pool.submit(vec![make_unit()]).collect().await; - assert_eq!(r1.len(), 1); - assert!( - matches!(r1[0].outcome, TestOutcome::Passed), - "first run should pass, got {:?}", - r1[0].outcome - ); + #[tokio::test] + async fn fanout_ctrl_error_names_the_operation_and_the_silent_workers() { + let (pool, _work_rx, _ctrl_rxs) = control_only_pool(2); - pool.restart_workers().await; + let error = pool + .fanout_ctrl("restart", WorkerCtrl::Restart, Duration::from_millis(10)) + .await + .expect_err("silent workers must fail the operation"); - let r2: Vec = pool.submit(vec![make_unit()]).collect().await; - assert_eq!(r2.len(), 1); + let message = format!("{error:#}"); + assert!(message.contains("'restart'"), "{message}"); assert!( - matches!(r2[0].outcome, TestOutcome::Passed), - "second run should pass on fresh interpreter, got {:?}", - r2[0].outcome - ); - - let count = std::fs::read_to_string(&counter_file).unwrap_or_default(); - assert_eq!( - count.len(), - 2, - "module body must run once per fresh interpreter \ - (1 initial + 1 after restart_workers); got {count:?}" + message.contains("2/2 workers did not acknowledge"), + "{message}" ); - - pool.shutdown(); } - /// When the worker python dies during startup (e.g. project venv - /// without `tryke` installed prints `ModuleNotFoundError` and - /// exits), the user-facing error must include the python stderr — - /// not just the opaque "worker unavailable" placeholder that used - /// to be all the user saw. Regression test for the diagnosability - /// fix. - /// - /// The unit carries a `HookItem` on purpose: with `hooks: vec![]`, - /// `handle_unit` skips `register_hooks_for_unit` entirely and the - /// failure would surface via `run_single_test`'s `run_test` error - /// path — which already existed before this PR. To exercise the - /// new code (`register_hooks_for_unit` stashing `last_failure` - /// after `drain_stderr`, then `ensure_worker`'s replay loop - /// stashing again on respawn, then `run_single_test` reading - /// `last_failure` instead of the opaque placeholder) the test - /// needs a hook so `register_hooks_for_unit` actually runs. - /// - /// Unix-only: simulates the failing python with a shell script. - /// The workspace venv's python has `tryke` installed in editable - /// mode and its `.pth` file is searched regardless of PYTHONPATH, - /// so we can't reproduce the missing-module case with a real - /// interpreter. A stub `python_bin` is sufficient — what we're - /// testing is the rust-side error propagation, not python's - /// resolution rules. - #[cfg(unix)] + /// `restart_workers` on a cold pool must start every process and + /// acknowledge within the control timeout. This matters because the file + /// watcher can fire before the user triggers any test run. #[tokio::test] - async fn worker_missing_tryke_surfaces_python_error_in_outcome() { - use std::os::unix::fs::PermissionsExt; - + async fn restart_workers_with_no_live_processes_acks() { let project = TestProject::new().expect("create test project"); - - let fake_python = project - .write( - "fake_python.sh", - "#!/bin/sh\n\ - echo \"$0: Error while finding module specification for \ - 'tryke.worker' (ModuleNotFoundError: No module named 'tryke')\" >&2\n\ - exit 1\n", - ) - .expect("write fake python"); - std::fs::set_permissions(&fake_python, std::fs::Permissions::from_mode(0o755)) - .expect("chmod fake python"); - - let test_file = project - .write("test_no_tryke.py", "def test_noop(): pass\n") - .expect("write test file"); - - // Hook on the same module forces `handle_unit` through - // `register_hooks_for_unit` → `ensure_worker` (spawn ok) → - // `register_hooks` (RPC error) → `drain_stderr` → stash - // `last_failure` → drop process. Then `run_single_test` → - // `ensure_worker` → respawn ok → replay `register_hooks` - // (RPC error) → stash `last_failure` → return None → - // `run_single_test` surfaces `last_failure` as the outcome - // message. Without this hook the new path isn't exercised. - let hook = HookItem { - name: "noop".into(), - module_path: "test_no_tryke".into(), - per: FixturePer::Test, - groups: vec![], - depends_on: vec![], - line_number: None, - }; - let unit = WorkUnit { - tests: vec![make_test_item("test_no_tryke", "test_noop", &test_file)], - hooks: vec![hook], - }; - - let python_path = [project.root().to_path_buf()]; + let python_path = [project.root().to_path_buf(), python_package_dir()]; let pool = WorkerPool::spawn_from_parts( - 1, - fake_python.to_str().expect("fake python path"), + 2, + &test_python_bin(), project.root(), Some(&python_path), LevelFilter::Off, @@ -873,39 +427,16 @@ def test_noop() -> None: ) .await; - let results: Vec = pool.submit(vec![unit]).collect().await; - assert_eq!(results.len(), 1); - let message = match &results[0].outcome { - TestOutcome::Error { message } => message.clone(), - other => panic!("expected Error outcome, got {other:?}"), - }; - // `run_single_test`'s `last_failure` path produces this prefix - // — proves we went through `ensure_worker`'s replay loop, not - // through `run_test`'s direct error path which would say - // "worker error: …". - assert!( - message.starts_with("hook replay failed for module test_no_tryke"), - "expected hook-replay prefix from ensure_worker.last_failure, got: {message}" - ); - assert!( - message.contains("No module named 'tryke'"), - "missing python traceback in error message: {message}" - ); - assert!( - message.contains("worker stderr:"), - "missing 'worker stderr:' header in error message: {message}" - ); + pool.restart_workers() + .await + .expect("restart_workers must ack within the control timeout"); - pool.shutdown(); + pool.shutdown().await.expect("clean shutdown"); } - /// `restart_workers` on a cold pool must start every process and - /// acknowledge within the control timeout. This matters because the file - /// watcher can fire before the user triggers any test run. #[tokio::test] - async fn restart_workers_with_no_live_processes_acks() { + async fn shutdown_joins_every_warmed_worker() { let project = TestProject::new().expect("create test project"); - let python_path = [project.root().to_path_buf(), python_package_dir()]; let pool = WorkerPool::spawn_from_parts( 2, @@ -913,29 +444,32 @@ def test_noop() -> None: project.root(), Some(&python_path), LevelFilter::Off, - false, + true, ) .await; - let restarted = - tokio::time::timeout(std::time::Duration::from_secs(10), pool.restart_workers()).await; - assert!(restarted.is_ok(), "restart_workers must ack within timeout"); - pool.shutdown(); + pool.shutdown() + .await + .expect("warmed workers must join cleanly"); } + /// Regression guard for the `send_blocking` → `try_send` change: submitting + /// to a pool whose work channel has closed must yield an empty stream + /// rather than parking the calling runtime thread. #[tokio::test] - async fn worker_control_fanout_times_out_when_a_worker_does_not_ack() { - let (work_tx, _work_rx) = async_channel::unbounded(); - let (ctrl_tx, _ctrl_rx) = mpsc::unbounded_channel(); - let pool = WorkerPool { - work_tx, - ctrl_txs: vec![ctrl_tx], - }; + async fn submit_to_a_closed_pool_ends_the_stream_instead_of_blocking() { + let (pool, work_rx, _ctrl_rxs) = control_only_pool(1); + drop(work_rx); + pool.work_tx.close(); - let acknowledged = pool - .fanout_ctrl_with_timeout(WorkerCtrl::Restart, Duration::from_millis(10)) - .await; + let units = partition_with_hooks(vec![TestItem::default()], &[], DistMode::Test).units; + assert!(!units.is_empty(), "test setup should produce a work unit"); + + let results: Vec = + tokio::time::timeout(Duration::from_secs(5), pool.submit(units).collect()) + .await + .expect("submit must not block on a closed channel"); - assert!(!acknowledged, "a missing worker ack must time out"); + assert!(results.is_empty(), "got {} results", results.len()); } } diff --git a/crates/tryke_runner/src/protocol.rs b/crates/tryke_runner/src/protocol.rs index c2f8452..dfe852c 100644 --- a/crates/tryke_runner/src/protocol.rs +++ b/crates/tryke_runner/src/protocol.rs @@ -30,6 +30,7 @@ //! worker never needs to re-walk the AST itself. use serde::{Deserialize, Serialize}; +use tryke_types::{HookItem, TestItem}; #[derive(Debug, Serialize)] #[serde(rename_all = "snake_case")] @@ -120,6 +121,32 @@ pub struct RegisterHooksParams { pub hooks: Vec, } +impl RegisterHooksParams { + /// Build the hook-registration payload for one test's module. + #[must_use] + pub fn for_test(test_item: &TestItem, hook_items: &[HookItem]) -> Self { + let hooks = hook_items + .iter() + .filter(|hook| hook.module_path == test_item.module_path) + .map(|hook| HookWire { + name: hook.name.clone(), + per: serde_json::to_value(hook.per) + .ok() + .and_then(|value| value.as_str().map(String::from)) + .unwrap_or_default(), + groups: hook.groups.clone(), + depends_on: hook.depends_on.clone(), + line_number: hook.line_number, + }) + .collect(); + + Self { + module: test_item.module_path.clone(), + hooks, + } + } +} + #[derive(Debug, Serialize)] pub struct FinalizeHooksParams { pub module: String, @@ -132,3 +159,45 @@ pub struct RunDoctestParams { } pub use tryke_types::{AssertionWire, RunTestResultWire}; + +#[cfg(test)] +mod tests { + use tryke_types::FixturePer; + + use super::*; + + #[test] + fn register_hooks_params_for_test_filters_and_maps_module_hooks() { + let test_item = TestItem { + module_path: "tests.test_math".into(), + ..TestItem::default() + }; + let matching = HookItem { + name: "database".into(), + module_path: test_item.module_path.clone(), + per: FixturePer::Scope, + groups: vec!["math".into()], + depends_on: vec!["connection".into()], + line_number: Some(12), + }; + let other = HookItem { + name: "other".into(), + module_path: "tests.test_other".into(), + per: FixturePer::Test, + groups: vec![], + depends_on: vec![], + line_number: None, + }; + + let params = RegisterHooksParams::for_test(&test_item, &[matching, other]); + + assert_eq!(params.module, test_item.module_path); + assert_eq!(params.hooks.len(), 1); + let hook = ¶ms.hooks[0]; + assert_eq!(hook.name, "database"); + assert_eq!(hook.per, "scope"); + assert_eq!(hook.groups, ["math"]); + assert_eq!(hook.depends_on, ["connection"]); + assert_eq!(hook.line_number, Some(12)); + } +} diff --git a/crates/tryke_runner/src/worker.rs b/crates/tryke_runner/src/worker.rs index 51fc598..4dfd248 100644 --- a/crates/tryke_runner/src/worker.rs +++ b/crates/tryke_runner/src/worker.rs @@ -1,871 +1,358 @@ -use std::collections::VecDeque; -use std::path::Path; -use std::sync::{Arc, Mutex}; +use std::path::PathBuf; use std::time::Duration; use anyhow::{Result, anyhow}; -use log::{debug, trace}; -use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader, BufWriter}; -use tokio::process::{Child, ChildStdin, ChildStdout, Command}; -use tryke_types::{TestItem, TestResult, convert_wire_result}; - -use crate::protocol::{ - FinalizeHooksParams, RPCRequest, RPCRequestMethod, RPCResponse, RegisterHooksParams, - RunDoctestParams, RunTestParams, RunTestResultWire, -}; - -/// Cap on retained worker-stderr bytes. Beyond this we keep the most recent -/// bytes and drop older ones — enough for diagnostic context on failures -/// without unbounded memory growth on workers that spew warnings. -const STDERR_RETAIN_BYTES: usize = 1 << 20; // 1 MiB - -pub struct WorkerProcess { - child: Child, - stdin: BufWriter, - stdout: BufReader, - /// Continuously-drained worker stderr. A background task reads the - /// child's stderr pipe into this buffer so the pipe never fills and - /// the worker can't block on a stderr write mid-RPC. - stderr_buf: Arc>>, - /// Handle to the stderr-drainer task. `drain_stderr` joins this - /// (with a short timeout) so any bytes still in the kernel pipe at - /// the moment of a worker failure end up in `stderr_buf` before we - /// snapshot it — without this, a worker that dies during startup - /// can lose its python traceback to a race with the RPC error path. - stderr_drainer: Option>, - next_id: u64, +use log::{LevelFilter, debug, trace}; +use tokio::sync::{mpsc, oneshot}; +use tokio_util::sync::CancellationToken; +use tryke_types::{TestOutcome, TestResult}; + +use crate::protocol::RegisterHooksParams; +use crate::schedule::WorkUnit; +use crate::worker_process::WorkerProcess; + +const WORKER_SPAWN_TIMEOUT: Duration = Duration::from_secs(5); + +/// One logical worker slot, which may replace its Python subprocess after a +/// crash, cancellation, or restart. +pub(crate) struct Worker { + python_bin: String, + python_path: Vec, + root: PathBuf, + log_level: LevelFilter, + process: Option, + /// Most recent spawn or hook-replay failure, captured so + /// `run_single_test` can surface the real reason (and any worker + /// stderr) instead of the opaque "worker unavailable" placeholder. + /// Cleared once we have a live worker again so we don't replay a + /// stale error against an unrelated test. + last_failure: Option, } -impl WorkerProcess { - /// Spawn a fresh worker process. - /// - /// `log_level` is forwarded as `TRYKE_LOG=` on the child env so - /// the python worker's `_configure_logging_from_env` lights up at the - /// same level as the rust process. Pass `LevelFilter::Off` to keep the - /// worker silent (no env var set), preserving the pre-existing - /// "no chatter unless asked" default. - /// - /// # Errors - /// Returns an error if the Python process cannot be spawned, if its stdio - /// pipes cannot be captured, or if the stderr drainer cannot be started. - pub fn spawn( - python_bin: &str, - python_path: &[&Path], - root: &Path, - log_level: log::LevelFilter, - ) -> Result { - debug!("spawning worker: {python_bin} -m tryke.worker (log={log_level})"); - let pythonpath = build_pythonpath(python_path); - let mut command = Command::new(python_bin); - command - .args(["-m", "tryke.worker"]) - .env("PYTHONPATH", &pythonpath) - .current_dir(root) - .stdin(std::process::Stdio::piped()) - .stdout(std::process::Stdio::piped()) - .stderr(std::process::Stdio::piped()); - if let Some(value) = worker_log_env_value(log_level) { - command.env("TRYKE_LOG", value); - } - let mut child = command.spawn()?; - let stdin = BufWriter::new(child.stdin.take().ok_or_else(|| anyhow!("no stdin"))?); - let stdout = BufReader::new(child.stdout.take().ok_or_else(|| anyhow!("no stdout"))?); - let stderr = child.stderr.take().ok_or_else(|| anyhow!("no stderr"))?; - debug!("worker spawned (pid {:?})", child.id()); - - // The worker can write to stderr at any time (asyncio default - // exception handler, library warnings, etc.). The kernel pipe - // buffer defaults to 64 KiB on Linux, so without an active - // reader the worker's next stderr write blocks once the pipe - // fills — and the worker stops responding to RPCs, surfacing - // as "tryke hangs at finalize_hooks". Spawn a drainer that - // keeps the pipe empty for the worker's lifetime. - let stderr_buf = Arc::new(Mutex::new(VecDeque::::new())); - let stderr_drainer = match spawn_stderr_drainer(stderr, Arc::clone(&stderr_buf)) { - Ok(handle) => handle, - Err(err) => { - if let Err(kill_err) = child.start_kill() { - debug!( - "failed to kill worker after stderr drainer setup error (pid {:?}): \ - {kill_err}", - child.id() - ); - } - return Err(err); - } - }; +pub(crate) enum WorkerMsg { + Unit { + unit: WorkUnit, + result_tx: mpsc::UnboundedSender, + cancel: CancellationToken, + }, +} - Ok(Self { - child, - stdin, - stdout, - stderr_buf, - stderr_drainer: Some(stderr_drainer), - next_id: 1, - }) +pub(crate) enum WorkerCtrl { + Ping(oneshot::Sender<()>), + Restart(oneshot::Sender<()>), +} + +fn format_worker_failure(prefix: &str, error: &dyn std::fmt::Display, stderr: &str) -> String { + let mut message = format!("{prefix}: {error}"); + let trimmed = stderr.trim(); + if !trimmed.is_empty() { + message.push_str("\nWorker stderr:\n"); + message.push_str(trimmed); } + message +} - async fn call serde::Deserialize<'de>>( - &mut self, - method: RPCRequestMethod, - params: Option, - ) -> Result { - let id = self.next_id; - self.next_id += 1; - let req = RPCRequest { - jsonrpc: "2.0", - id, - method, - params, - }; - let mut line = serde_json::to_string(&req)?; - line.push('\n'); - trace!("worker rpc -> {}", line.trim()); - self.stdin.write_all(line.as_bytes()).await?; - self.stdin.flush().await?; - // Read lines from stdout, skipping non-JSON garbage that a native - // library may have written to fd 1 during import (e.g. weasyprint via - // cffi). Collect leaked lines so we can surface them in errors. - let mut leaked_stdout: Vec = Vec::new(); - let resp: RPCResponse = loop { - let mut resp_line = String::new(); - let n = self.stdout.read_line(&mut resp_line).await?; - if n == 0 { - trace!("worker rpc: stdout EOF"); - return Err(if leaked_stdout.is_empty() { - anyhow!("worker process closed stdout") - } else { - anyhow!( - "worker process closed stdout after writing non-JSON output \ - (a library may have written to stdout during import):\n{}", - leaked_stdout.join("") - ) - }); - } - trace!("worker rpc <- {}", resp_line.trim()); - let trimmed = resp_line.trim(); - if !trimmed.is_empty() - && let Ok(resp) = serde_json::from_str::(trimmed) - { - if !leaked_stdout.is_empty() { - trace!( - "worker rpc: skipped {} non-JSON line(s) on stdout", - leaked_stdout.len() - ); - } - break resp; - } - leaked_stdout.push(resp_line); - if leaked_stdout.len() >= 50 { - return Err(anyhow!( - "expected JSON-RPC response from worker but got {} lines of \ - non-JSON output (a library may have written to stdout during \ - import):\n{}", - leaked_stdout.len(), - leaked_stdout.join("") - )); - } - }; - if let Some(err) = resp.error { - let detail = if let Some(tb) = &err.traceback { - format!("rpc error {}: {}\n{tb}", err.code, err.message) - } else { - format!("rpc error {}: {}", err.code, err.message) - }; - return Err(anyhow!(detail)); +impl Worker { + pub(crate) fn new( + python_bin: String, + python_path: Vec, + root: PathBuf, + log_level: LevelFilter, + ) -> Self { + Self { + python_bin, + python_path, + root, + log_level, + process: None, + last_failure: None, } - let val = resp.result.unwrap_or(serde_json::Value::Null); - Ok(serde_json::from_value(val)?) } - /// Run a discovered test or doctest in the worker process. - /// - /// # Errors - /// Returns an error if the request cannot be serialized, if worker I/O - /// fails, or if the worker returns a JSON-RPC error. - pub async fn run_test(&mut self, test: &TestItem) -> Result { - if let Some(object_path) = &test.doctest_object { - return self.run_doctest(test, object_path).await; + async fn spawn_process(&self) -> Result { + let python_bin = self.python_bin.clone(); + let python_paths = self.python_path.clone(); + let root = self.root.clone(); + let log_level = self.log_level; + let spawn = tokio::task::spawn_blocking(move || { + let path_refs = python_paths + .iter() + .map(PathBuf::as_path) + .collect::>(); + WorkerProcess::spawn(&python_bin, &path_refs, &root, log_level) + }); + + match tokio::time::timeout(WORKER_SPAWN_TIMEOUT, spawn).await { + Ok(Ok(result)) => result, + Ok(Err(error)) => Err(anyhow!("Worker spawn task failed: {error}")), + Err(_) => Err(anyhow!( + "Worker process spawn timed out after {WORKER_SPAWN_TIMEOUT:?}" + )), } - let params = serde_json::to_value(RunTestParams { - module: test.module_path.clone(), - function: test.name.clone(), - xfail: test.xfail.clone(), - groups: test.groups.clone(), - case_label: test.case_label.clone(), - })?; - let wire: RunTestResultWire = self.call(RPCRequestMethod::RunTest, Some(params)).await?; - Ok(convert_wire_result(test.clone(), wire)) } - /// Send hook metadata for a module to the Python worker. - /// Must be called before running any tests from that module. - /// - /// # Errors - /// Returns an error if hook metadata cannot be serialized or the worker - /// rejects the registration request. - pub async fn register_hooks(&mut self, params: RegisterHooksParams) -> Result<()> { - let value = serde_json::to_value(params)?; - self.call::(RPCRequestMethod::RegisterHooks, Some(value)) - .await?; - Ok(()) - } + /// Ensure a Python process is live, replaying only the active unit's hook + /// registrations when a replacement process is needed. + async fn ensure_process<'a>( + &'a mut self, + registrations: &[RegisterHooksParams], + ) -> Option<&'a mut WorkerProcess> { + if self.process.is_some() { + return self.process.as_mut(); + } - /// Tell the Python worker to run scope-level teardown for `per="scope"` - /// fixtures in a module. Must be called after all tests from that module - /// have run. - /// - /// # Errors - /// Returns an error if the finalize request cannot be serialized or the - /// worker reports a teardown failure. - pub async fn finalize_hooks(&mut self, module: String) -> Result<()> { - let value = serde_json::to_value(FinalizeHooksParams { module })?; - self.call::(RPCRequestMethod::FinalizeHooks, Some(value)) - .await?; - Ok(()) - } + trace!("Worker: spawning process"); - async fn run_doctest(&mut self, test: &TestItem, object_path: &str) -> Result { - let params = serde_json::to_value(RunDoctestParams { - module: test.module_path.clone(), - object_path: object_path.to_owned(), - })?; - let wire: RunTestResultWire = self - .call(RPCRequestMethod::RunDoctest, Some(params)) - .await?; - Ok(convert_wire_result(test.clone(), wire)) - } + let mut process = match self.spawn_process().await { + Ok(process) => process, + Err(error) => { + let message = format_worker_failure( + &format!( + "Failed to spawn Python worker ({} -m tryke.worker)", + self.python_bin + ), + &error, + "", + ); - /// Verify that the worker process is responsive. - /// - /// # Errors - /// Returns an error if the ping RPC fails or if the worker returns an - /// unexpected response. - pub async fn ping(&mut self) -> Result<()> { - let result: String = self.call(RPCRequestMethod::Ping, None).await?; - if result == "pong" { - Ok(()) - } else { - Err(anyhow!("unexpected ping response: {result}")) - } - } + debug!("Worker: {message}"); - /// Snapshot the buffered worker stderr and clear the buffer. - /// - /// Called on the error path when the worker is about to be - /// discarded. We kill the child first so the drainer task reaches - /// EOF, then await the drainer (with a short timeout) so any bytes - /// still in the kernel pipe at the time of the failure land in - /// `stderr_buf` before we snapshot. Without this, a python worker - /// that dies during startup (e.g. `ModuleNotFoundError: tryke`) - /// can lose its traceback to a race between the RPC's - /// `Broken pipe` and the drainer task being scheduled. - /// - /// # Panics - /// Panics only if the stderr-drainer task panicked while holding the - /// internal mutex (poisoning it). That task does no fallible work. - pub async fn drain_stderr(&mut self) -> String { - // Worker is about to be discarded; killing the child closes its - // stderr pipe which gives the drainer its EOF. Idempotent on a - // process that has already exited. - let _ = self.child.start_kill(); - if let Some(handle) = self.stderr_drainer.take() { - let _ = tokio::time::timeout(Duration::from_millis(500), handle).await; - } - let bytes: Vec = { - let mut g = self - .stderr_buf - .lock() - .expect("stderr buffer mutex poisoned"); - std::mem::take(&mut *g).into_iter().collect() + self.last_failure = Some(message); + + return None; + } }; - String::from_utf8_lossy(&bytes).into_owned() - } - pub async fn shutdown(&mut self) { - let _ = self.child.kill().await; - // Killing the child closes stderr; the drainer will see EOF and - // exit on its own — but `abort()` is the explicit, immediate - // signal that we're done with it, and prevents a leftover task - // from briefly holding the stderr FD + `stderr_buf` Arc past - // the worker's lifetime. - if let Some(handle) = self.stderr_drainer.take() { - handle.abort(); - } - } -} + for params in registrations { + if let Err(error) = process.register_hooks(params.clone()).await { + let stderr_output = process.drain_stderr().await; -impl Drop for WorkerProcess { - fn drop(&mut self) { - // Safety net: ensure the child process is killed when the worker is - // dropped (e.g. on the error-respawn path in pool.rs). start_kill() is - // the synchronous variant — safe to call on already-dead processes. - let _ = self.child.start_kill(); - // Dropping a Tokio JoinHandle detaches the task, so abort it - // explicitly to avoid an orphan drainer outliving the worker on - // respawn paths (the drainer holds stderr_buf + the stderr FD). - if let Some(handle) = self.stderr_drainer.take() { - handle.abort(); - } - } -} + let message = format_worker_failure( + &format!("Hook replay failed for module {}", params.module), + &error, + &stderr_output, + ); -/// Spawn a tokio task that continuously reads `stderr` into `buf` until -/// EOF or a read error, capping the buffer at `STDERR_RETAIN_BYTES`. -/// -/// Returns an error if no Tokio runtime is currently entered, rather than -/// panicking the way `tokio::spawn` would. This keeps the synchronous -/// `WorkerProcess::spawn` API safe to call from any context — callers in -/// non-async code receive a structured error instead of a panic. -fn spawn_stderr_drainer( - stderr: tokio::process::ChildStderr, - buf: Arc>>, -) -> Result> { - let handle = tokio::runtime::Handle::try_current() - .map_err(|e| anyhow!("WorkerProcess::spawn requires an active tokio runtime: {e}"))?; - Ok(handle.spawn(async move { - let mut reader = stderr; - let mut chunk = [0u8; 8192]; - loop { - match reader.read(&mut chunk).await { - Ok(0) => break, - Ok(n) => append_stderr(&buf, &chunk[..n]), - Err(e) => { - trace!("stderr drainer: read error: {e}"); - break; - } + debug!("Worker: {message}"); + + self.last_failure = Some(message); + + return None; } } - })) -} -/// Append `data` to `buf`, capping it at `STDERR_RETAIN_BYTES` by -/// dropping the oldest bytes. `VecDeque::drain` removes from the front -/// in `O(excess)` time, making steady-state appends `O(data.len())` -/// rather than `O(STDERR_RETAIN_BYTES)` as a contiguous `Vec` would require. -fn append_stderr(buf: &Mutex>, data: &[u8]) { - let mut g = buf.lock().expect("stderr buffer mutex poisoned"); - if data.len() >= STDERR_RETAIN_BYTES { - // A single chunk already fills the cap; keep only its tail. - g.clear(); - g.extend(data[data.len() - STDERR_RETAIN_BYTES..].iter().copied()); - return; - } - let total = g.len() + data.len(); - if total > STDERR_RETAIN_BYTES { - let excess = total - STDERR_RETAIN_BYTES; - g.drain(..excess); + self.last_failure = None; + self.process = Some(process); + self.process.as_mut() } - g.extend(data.iter().copied()); -} -/// Translate a resolved log level into the value placed on the spawned -/// worker's `TRYKE_LOG` env var, if any. -/// -/// `Off` returns `None` so the env var stays unset and the python -/// worker's `_configure_logging_from_env` no-ops. Anything else returns -/// the lowercase level name (`debug`, `info`, ...) which the worker's -/// `logging.getLevelName` understands once uppercased. -fn worker_log_env_value(log_level: log::LevelFilter) -> Option { - if log_level == log::LevelFilter::Off { - return None; - } - Some(log_level.as_str().to_ascii_lowercase()) -} + /// Install the active unit's complete hook metadata on an existing + /// process, or spawn a process and replay it there. + async fn prepare_unit(&mut self, registrations: &[RegisterHooksParams]) { + if self.process.is_none() { + let _ = self.ensure_process(registrations).await; + return; + } -fn build_pythonpath(extra: &[&Path]) -> String { - let existing = std::env::var("PYTHONPATH").unwrap_or_default(); - let mut parts: Vec = extra - .iter() - .map(|p| { - let s = p - .canonicalize() - .unwrap_or_else(|_| p.to_path_buf()) - .to_string_lossy() - .into_owned(); - // std::fs::canonicalize on windows produces \\?\ extended-length - // paths that python doesn't understand - #[cfg(windows)] - let s = s.strip_prefix(r"\\?\").unwrap_or(&s).to_string(); - s - }) - .collect(); - if !existing.is_empty() { - parts.push(existing); - } - let sep = if cfg!(windows) { ";" } else { ":" }; - parts.join(sep) -} + for params in registrations { + let registration_failure = { + let Some(process) = self.process.as_mut() else { + return; + }; + + match process.register_hooks(params.clone()).await { + Ok(()) => None, + Err(error) => { + let stderr_output = process.drain_stderr().await; + Some((error, stderr_output)) + } + } + }; -#[cfg(test)] -mod tests { - use std::path::PathBuf; + if let Some((error, stderr_output)) = registration_failure { + let message = format_worker_failure( + &format!("Hook registration failed for module {}", params.module), + &error, + &stderr_output, + ); - use tryke_types::{AssertionWire, ExpectedAssertion, TestOutcome}; + debug!("Worker: {message}"); - use super::*; + self.last_failure = Some(message); + self.process = None; - fn make_test_item() -> TestItem { - TestItem { - name: "test_add".into(), - module_path: "tests.test_math".into(), - file_path: Some(PathBuf::from("tests/test_math.py")), - line_number: Some(5), - display_name: None, - expected_assertions: vec![], - ..Default::default() + return; + } } - } - #[test] - fn worker_log_env_value_off_returns_none() { - // `Off` means: don't set TRYKE_LOG on the child env, preserving - // the python worker's "no chatter unless asked" default. - assert_eq!(worker_log_env_value(log::LevelFilter::Off), None); + self.last_failure = None; } - #[test] - fn worker_log_env_value_lowercases_level_name() { - // The python worker uppercases before passing to - // `logging.getLevelName`, so case actually doesn't matter, but - // shipping lowercase keeps env values consistent with how rust - // log levels render and avoids a 50/50 stylistic decision. - assert_eq!( - worker_log_env_value(log::LevelFilter::Info).as_deref(), - Some("info"), - ); - assert_eq!( - worker_log_env_value(log::LevelFilter::Debug).as_deref(), - Some("debug"), - ); - assert_eq!( - worker_log_env_value(log::LevelFilter::Warn).as_deref(), - Some("warn"), - ); - } + async fn run_single_test( + &mut self, + test: tryke_types::TestItem, + registrations: &[RegisterHooksParams], + result_tx: &mpsc::UnboundedSender, + ) { + let Some(process) = self.ensure_process(registrations).await else { + let message = self + .last_failure + .clone() + .unwrap_or_else(|| "Worker unavailable (spawn or hook replay failed)".into()); + + let _ = result_tx.send(TestResult { + test, + outcome: TestOutcome::Error { message }, + duration: Duration::ZERO, + stdout: String::new(), + stderr: String::new(), + }); - #[test] - fn convert_result_passed() { - let test = make_test_item(); - let wire = RunTestResultWire::Passed { - duration_ms: 10, - stdout: "out".into(), - stderr: "err".into(), + return; }; - let result = convert_wire_result(test, wire); - assert!(matches!(result.outcome, TestOutcome::Passed)); - assert_eq!(result.duration, Duration::from_millis(10)); - assert_eq!(result.stdout, "out"); - assert_eq!(result.stderr, "err"); - } - #[test] - fn convert_result_failed() { - let test = make_test_item(); - let wire = RunTestResultWire::Failed { - duration_ms: 5, - message: "expected 1 got 2".into(), - traceback: None, - assertions: vec![], - executed_lines: vec![], - stdout: String::new(), - stderr: String::new(), - }; - let result = convert_wire_result(test, wire); - assert!(matches!( - result.outcome, - TestOutcome::Failed { ref message, .. } if message == "expected 1 got 2" - )); + match process.run_test(&test).await { + Ok(result) => { + trace!("Worker: test {} done", test.name); + let _ = result_tx.send(result); + } + Err(error) => { + debug!("Worker: run_test error for {}: {error}", test.name); + let stderr_output = process.drain_stderr().await; + + // The next test reconstructs a replacement process from this + // active unit's registrations. + self.process = None; + + let message = format_worker_failure("Worker error", &error, &stderr_output); + + let _ = result_tx.send(TestResult { + test, + outcome: TestOutcome::Error { message }, + duration: Duration::ZERO, + stdout: String::new(), + stderr: stderr_output, + }); + } + } } - #[test] - fn convert_result_skipped() { - let test = make_test_item(); - let wire = RunTestResultWire::Skipped { - duration_ms: 0, - reason: Some("not ready".into()), - stdout: String::new(), - stderr: String::new(), - }; - let result = convert_wire_result(test, wire); - assert!(matches!( - result.outcome, - TestOutcome::Skipped { reason: Some(ref r) } if r == "not ready" - )); + fn registrations_for_unit(unit: &WorkUnit) -> Vec { + let mut seen = std::collections::HashSet::new(); + unit.tests + .iter() + .filter(|test| seen.insert(test.module_path.clone())) + .map(|test| RegisterHooksParams::for_test(test, &unit.hooks)) + .collect() } - #[test] - fn build_pythonpath_joins_paths() { - let dir_a = tempfile::tempdir().unwrap(); - let dir_b = tempfile::tempdir().unwrap(); - let a = build_pythonpath(&[dir_a.path()]); - let b = build_pythonpath(&[dir_b.path()]); - let result = build_pythonpath(&[dir_a.path(), dir_b.path()]); - let sep = if cfg!(windows) { ";" } else { ":" }; - assert_eq!(result, format!("{a}{sep}{b}")); - } + async fn handle_unit(&mut self, unit: WorkUnit, result_tx: mpsc::UnboundedSender) { + let registrations = Self::registrations_for_unit(&unit); - #[test] - fn convert_result_uses_discovered_multiline_assertion_source() { - let expression = - "t.expect(\n expr=actual,\n name=\"actual should be one\",\n).to_equal(other=1)"; - let test = TestItem { - expected_assertions: vec![ExpectedAssertion { - subject: "actual".into(), - matcher: "to_equal".into(), - negated: false, - args: vec!["other=1".into()], - line: 7, - label: Some("actual should be one".into()), - end_line: 10, - start_column: Some(4), - end_column: Some(18), - expression: expression.into(), - subject_span: expression - .find("actual") - .map(|offset| (offset, "actual".len())), - expected_arg_span: expression - .find("other=1") - .map(|offset| (offset, "other=1".len())), - expected_arg_value: Some("1".into()), - }], - ..Default::default() - }; - let result = convert_wire_result( - test, - RunTestResultWire::Failed { - duration_ms: 1, - message: "assertion failed".into(), - traceback: None, - assertions: vec![AssertionWire { - expression: ".to_equal(other=1)".into(), - expected: "1".into(), - received: "0".into(), - line: 10, - column: Some(6), - file: Some("tests/test_multiline.py".into()), - }], - executed_lines: vec![10], - stdout: String::new(), - stderr: String::new(), - }, - ); - let TestOutcome::Failed { - assertions, - executed_lines, - .. - } = result.outcome - else { - panic!("expected failed outcome"); - }; - let assertion = &assertions[0]; - assert_eq!(assertion.expression, expression); - assert_eq!(assertion.line, 7); - assert_eq!(assertion.span_offset, expression.find("actual").unwrap()); - assert_eq!( - assertion.expected_arg_span, - expression.find("other=1").map(|offset| (offset, 7)) - ); - assert_eq!(executed_lines, vec![7]); - } + self.prepare_unit(®istrations).await; - #[test] - fn convert_result_selects_inner_assertion_by_column() { - let test = TestItem { - expected_assertions: vec![ - ExpectedAssertion { - subject: "expect(0).to_equal(1)".into(), - matcher: "to_be_truthy".into(), - negated: false, - args: vec![], - line: 3, - end_line: 3, - start_column: Some(4), - end_column: Some(45), - expression: "expect(expect(0).to_equal(1)).to_be_truthy()".into(), - subject_span: Some((7, 21)), - expected_arg_span: None, - expected_arg_value: None, - label: None, - }, - ExpectedAssertion { - subject: "0".into(), - matcher: "to_equal".into(), - negated: false, - args: vec!["1".into()], - line: 3, - end_line: 3, - start_column: Some(11), - end_column: Some(32), - expression: "expect(0).to_equal(1)".into(), - subject_span: Some((7, 1)), - expected_arg_span: Some((19, 1)), - expected_arg_value: Some("1".into()), - label: None, - }, - ], - ..Default::default() - }; - let result = convert_wire_result( - test, - RunTestResultWire::Failed { - duration_ms: 1, - message: "assertion failed".into(), - traceback: None, - assertions: vec![AssertionWire { - expression: "expect(expect(0).to_equal(1)).to_be_truthy()".into(), - expected: "1".into(), - received: "0".into(), - line: 3, - column: Some(21), - file: None, - }], - executed_lines: vec![3], - stdout: String::new(), - stderr: String::new(), - }, - ); - let TestOutcome::Failed { assertions, .. } = result.outcome else { - panic!("expected failed outcome"); - }; - assert_eq!(assertions[0].expression, "expect(0).to_equal(1)"); - assert_eq!(assertions[0].span_offset, 7); - assert_eq!(assertions[0].expected_arg_span, Some((19, 1))); - } + for test in unit.tests { + trace!("Worker: running test {}", test.name); - #[test] - fn convert_result_matches_same_line_assertion_by_expected_value_without_column() { - let second_expression = "expect(b).to_equal(other=2)"; - let test = TestItem { - expected_assertions: vec![ - ExpectedAssertion { - subject: "a".into(), - matcher: "to_equal".into(), - negated: false, - args: vec!["1".into()], - line: 5, - end_line: 5, - expression: "expect(a).to_equal(1)".into(), - subject_span: Some((7, 1)), - expected_arg_span: Some((19, 1)), - label: None, - ..Default::default() - }, - ExpectedAssertion { - subject: "b".into(), - matcher: "to_equal".into(), - negated: false, - args: vec!["other=2".into()], - line: 5, - end_line: 5, - expression: second_expression.into(), - subject_span: Some((7, 1)), - expected_arg_span: second_expression - .find("other=2") - .map(|offset| (offset, "other=2".len())), - label: None, - ..Default::default() - }, - ], - ..Default::default() - }; - let result = convert_wire_result( - test, - RunTestResultWire::Failed { - duration_ms: 1, - message: "assertion failed".into(), - traceback: None, - assertions: vec![AssertionWire { - expression: "expect(a).to_equal(1); expect(b).to_equal(other=2)".into(), - expected: "2".into(), - received: "0".into(), - line: 5, - column: None, - file: None, - }], - executed_lines: vec![5], - stdout: String::new(), - stderr: String::new(), - }, - ); - let TestOutcome::Failed { assertions, .. } = result.outcome else { - panic!("expected failed outcome"); - }; - assert_eq!(assertions[0].expression, second_expression); - assert_eq!( - assertions[0].expected_arg_span, - second_expression - .find("other=2") - .map(|offset| (offset, "other=2".len())) - ); - } + self.run_single_test(test, ®istrations, &result_tx).await; + } - #[test] - fn convert_result_keeps_runtime_expression_for_ambiguous_same_line_without_column() { - let runtime_expression = "expect(a).to_equal(1); expect(b).to_equal(1)"; - let test = TestItem { - expected_assertions: vec![ - ExpectedAssertion { - subject: "a".into(), - matcher: "to_equal".into(), - negated: false, - args: vec!["1".into()], - line: 5, - end_line: 5, - expression: "expect(a).to_equal(1)".into(), - subject_span: Some((7, 1)), - label: None, - ..Default::default() - }, - ExpectedAssertion { - subject: "b".into(), - matcher: "to_equal".into(), - negated: false, - args: vec!["1".into()], - line: 5, - end_line: 5, - expression: "expect(b).to_equal(1)".into(), - subject_span: Some((7, 1)), - label: None, - ..Default::default() - }, - ], - ..Default::default() - }; - let result = convert_wire_result( - test, - RunTestResultWire::Failed { - duration_ms: 1, - message: "assertion failed".into(), - traceback: None, - assertions: vec![AssertionWire { - expression: runtime_expression.into(), - expected: "1".into(), - received: "0".into(), - line: 5, - column: None, - file: None, - }], - executed_lines: vec![5], - stdout: String::new(), - stderr: String::new(), - }, - ); - let TestOutcome::Failed { assertions, .. } = result.outcome else { - panic!("expected failed outcome"); - }; - assert_eq!(assertions[0].expression, runtime_expression); + for params in registrations { + if let Some(process) = self.process.as_mut() + && let Err(error) = process.finalize_hooks(params.module).await + { + debug!("Worker: finalize_hooks failed: {error}"); + } + } } - #[tokio::test] - async fn drop_kills_child_process() { - let mut child = tokio::process::Command::new("sleep") - .arg("60") - .spawn() - .expect("failed to spawn sleep"); - let pid = child.id().expect("missing pid"); - - // Wrap in a WorkerProcess-like drop: start_kill then drop - let _ = child.start_kill(); - drop(child); - - // Give the OS a moment to reap the process - tokio::time::sleep(Duration::from_millis(50)).await; - - // On Unix, sending signal 0 checks if the process exists - let status = std::process::Command::new("kill") - .args(["-0", &pid.to_string()]) - .status(); - assert!( - status.is_ok_and(|s| !s.success()), - "child process {pid} should be dead after start_kill + drop" - ); + async fn handle_control(&mut self, ctrl: WorkerCtrl) { + match ctrl { + WorkerCtrl::Ping(ack_tx) => { + trace!("Worker: ping (pre-warm)"); + let _ = self.ensure_process(&[]).await; + let _ = ack_tx.send(()); + } + WorkerCtrl::Restart(ack_tx) => { + trace!("Worker: restart"); + self.reset_process().await; + let _ = ack_tx.send(()); + } + } } - #[test] - fn append_stderr_caps_at_retain_bytes() { - let buf = Mutex::new(VecDeque::::new()); - // Fill with a recognisable head sequence, then push past the cap. - let head = vec![b'A'; STDERR_RETAIN_BYTES]; - append_stderr(&buf, &head); - let tail = vec![b'B'; 1024]; - append_stderr(&buf, &tail); - - let mut g = buf.lock().unwrap(); - assert_eq!(g.len(), STDERR_RETAIN_BYTES); - // Oldest A's are dropped; tail B's retained. - let bytes = g.make_contiguous(); - assert_eq!(&bytes[bytes.len() - tail.len()..], &tail[..]); - assert_eq!(bytes[0], b'A'); + async fn reset_process(&mut self) { + if let Some(mut process) = self.process.take() { + process.shutdown().await; + } } - #[test] - fn append_stderr_handles_single_oversized_write() { - let buf = Mutex::new(VecDeque::::new()); - let big = vec![b'X'; STDERR_RETAIN_BYTES + 4096]; - append_stderr(&buf, &big); - let mut g = buf.lock().unwrap(); - assert_eq!(g.len(), STDERR_RETAIN_BYTES); - assert!(g.make_contiguous().iter().all(|b| *b == b'X')); + async fn shutdown(mut self) { + self.reset_process().await; } - /// Reproducer for the "tryke hangs at `finalize_hooks`" bug: when the - /// worker writes more than the kernel pipe buffer (~64 KiB on Linux, - /// ~4 KiB on Windows) to stderr without anyone draining it, the - /// worker's next stderr write blocks and it stops responding to - /// RPCs. The drainer task spawned in `WorkerProcess::spawn` must - /// keep the pipe empty so RPC flow is unaffected. + /// Claim work units until the pool shuts down or its channels close. /// - /// Test contract is binary: either the drainer prevents the deadlock - /// and `done` arrives on stdout, or it doesn't. The cap-retention - /// property is covered by the `append_stderr_*` unit tests above and - /// is intentionally NOT asserted here — re-asserting it via a second - /// post-exit polling loop introduced wall-clock flakiness on slow - /// Windows runners without adding any signal a unit test wouldn't - /// catch. - /// - /// Uses `python3` (already a project dependency) for portability and - /// routes through the same `spawn_stderr_drainer` helper as - /// `WorkerProcess::spawn`, so a regression in the production drainer - /// path would deadlock this test too. - #[tokio::test] - async fn worker_continues_after_large_stderr_write() { - // 256 KiB is comfortably above both Linux's 64 KiB pipe buffer - // and Windows's 4 KiB default — enough to deadlock without a - // drainer. Larger values (we used to write 1 MiB) just make the - // test slower without proving anything additional. - const STDERR_BYTES: usize = 256 * 1024; - let mut child = tokio::process::Command::new("python3") - .args([ - "-c", - &format!( - "import sys; sys.stderr.write('X' * {STDERR_BYTES}); \ - sys.stderr.flush(); print('done', flush=True)" - ), - ]) - .stdin(std::process::Stdio::null()) - .stdout(std::process::Stdio::piped()) - .stderr(std::process::Stdio::piped()) - .spawn() - .expect("spawn python3"); - - let stderr = child.stderr.take().expect("no stderr"); - let stderr_buf: Arc>> = Arc::new(Mutex::new(VecDeque::new())); - spawn_stderr_drainer(stderr, Arc::clone(&stderr_buf)) - .expect("drainer requires tokio runtime"); - - let mut stdout = BufReader::new(child.stdout.take().expect("no stdout")); - let mut line = String::new(); - // 30s catches a real deadlock (which would never finish) while - // tolerating Windows python startup + scheduler jitter. Past 5s - // budgets occasionally tripped on Windows × 3.14 even though no - // deadlock occurred. - tokio::time::timeout(Duration::from_secs(30), stdout.read_line(&mut line)) - .await - .expect("stdout read timed out — drainer deadlocked") - .expect("stdout read errored"); - assert_eq!(line.trim(), "done"); - - let _ = child.wait().await; + /// `async_channel::Receiver::recv` is polled inside `select!` and is + /// therefore dropped whenever the shutdown or control branch wins. That is + /// safe with async-channel 2.x — a dropped `Recv` re-notifies another + /// listener rather than swallowing the unit — but unlike + /// `tokio::sync::mpsc` the crate does not document the guarantee, so a + /// channel swap here needs to re-check it. + pub(crate) async fn run( + mut self, + work_rx: async_channel::Receiver, + mut ctrl_rx: mpsc::UnboundedReceiver, + shutdown: CancellationToken, + ) { + 'worker: loop { + tokio::select! { + biased; + () = shutdown.cancelled() => break, + ctrl = ctrl_rx.recv() => { + let Some(ctrl) = ctrl else { break }; + self.handle_control(ctrl).await; + } + msg = work_rx.recv() => { + match msg { + Ok(WorkerMsg::Unit { + unit, + result_tx, + cancel, + }) => { + if cancel.is_cancelled() { + continue; + } + + let interrupted_by_shutdown = { + let unit_future = self.handle_unit(unit, result_tx); + + tokio::pin!(unit_future); + + tokio::select! { + biased; + () = shutdown.cancelled() => Some(true), + () = cancel.cancelled() => Some(false), + () = &mut unit_future => None, + } + }; + + if let Some(shutting_down) = interrupted_by_shutdown { + // The interrupted unit left an RPC half-written + // on the interpreter's stdin, so the process is + // no longer usable — drop it either way. + self.reset_process().await; + if shutting_down { + break 'worker; + } + } + } + Err(_) => break, + } + } + } + } + + self.shutdown().await; } } diff --git a/crates/tryke_runner/src/worker_process.rs b/crates/tryke_runner/src/worker_process.rs new file mode 100644 index 0000000..6376a8f --- /dev/null +++ b/crates/tryke_runner/src/worker_process.rs @@ -0,0 +1,867 @@ +use std::collections::VecDeque; +use std::path::Path; +use std::process::Stdio; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use anyhow::{Result, anyhow}; +use log::{debug, trace}; +use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader, BufWriter}; +use tokio::process::{Child, ChildStdin, ChildStdout, Command}; +use tryke_types::{TestItem, TestResult, convert_wire_result}; + +use crate::protocol::{ + FinalizeHooksParams, RPCRequest, RPCRequestMethod, RPCResponse, RegisterHooksParams, + RunDoctestParams, RunTestParams, RunTestResultWire, +}; + +/// Cap on retained worker-stderr bytes. Beyond this we keep the most recent +/// bytes and drop older ones — enough for diagnostic context on failures +/// without unbounded memory growth on workers that spew warnings. +const STDERR_RETAIN_BYTES: usize = 1 << 20; // 1 MiB + +pub struct WorkerProcess { + child: Child, + stdin: BufWriter, + stdout: BufReader, + /// Continuously-drained worker stderr. A background task reads the + /// child's stderr pipe into this buffer so the pipe never fills and + /// the worker can't block on a stderr write mid-RPC. + stderr_buf: Arc>>, + /// Handle to the stderr-drainer task. `drain_stderr` joins this + /// (with a short timeout) so any bytes still in the kernel pipe at + /// the moment of a worker failure end up in `stderr_buf` before we + /// snapshot it — without this, a worker that dies during startup + /// can lose its python traceback to a race with the RPC error path. + stderr_drainer: Option>, + next_id: u64, +} + +impl WorkerProcess { + /// Spawn a fresh worker process. + /// + /// `log_level` is forwarded as `TRYKE_LOG=` on the child env so + /// the Python worker's `_configure_logging_from_env` lights up at the + /// resolved Tryke level. `LevelFilter::Off` explicitly disables worker + /// logging. + /// + /// # Errors + /// Returns an error if the Python process cannot be spawned, if its stdio + /// pipes cannot be captured, or if the stderr drainer cannot be started. + pub fn spawn( + python_bin: &str, + python_path: &[&Path], + root: &Path, + log_level: log::LevelFilter, + ) -> Result { + debug!("Spawning worker: {python_bin} -m tryke.worker (log={log_level})"); + + let pythonpath = build_pythonpath(python_path); + let mut command = Command::new(python_bin); + + command + .args(["-m", "tryke.worker"]) + .env("PYTHONPATH", &pythonpath) + .current_dir(root) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .env("TRYKE_LOG", worker_log_env_value(log_level)); + + let mut child = command.spawn()?; + + let stdin = BufWriter::new(child.stdin.take().ok_or_else(|| anyhow!("No stdin"))?); + let stdout = BufReader::new(child.stdout.take().ok_or_else(|| anyhow!("No stdout"))?); + let stderr = child.stderr.take().ok_or_else(|| anyhow!("No stderr"))?; + + debug!("Worker spawned (pid {:?})", child.id()); + + // The worker can write to stderr at any time (asyncio default + // exception handler, library warnings, etc.). The kernel pipe + // buffer defaults to 64 KiB on Linux, so without an active + // reader the worker's next stderr write blocks once the pipe + // fills — and the worker stops responding to RPCs, surfacing + // as "tryke hangs at finalize_hooks". Spawn a drainer that + // keeps the pipe empty for the worker's lifetime. + let stderr_buf = Arc::new(Mutex::new(VecDeque::::new())); + let stderr_drainer = match spawn_stderr_drainer(stderr, Arc::clone(&stderr_buf)) { + Ok(handle) => handle, + Err(err) => { + if let Err(kill_err) = child.start_kill() { + debug!( + "Failed to kill worker after stderr drainer setup error (pid {:?}): \ + {kill_err}", + child.id() + ); + } + return Err(err); + } + }; + + Ok(Self { + child, + stdin, + stdout, + stderr_buf, + stderr_drainer: Some(stderr_drainer), + next_id: 1, + }) + } + + async fn call serde::Deserialize<'de>>( + &mut self, + method: RPCRequestMethod, + params: Option, + ) -> Result { + let id = self.next_id; + self.next_id += 1; + let req = RPCRequest { + jsonrpc: "2.0", + id, + method, + params, + }; + let mut line = serde_json::to_string(&req)?; + line.push('\n'); + trace!("Worker RPC -> {}", line.trim()); + self.stdin.write_all(line.as_bytes()).await?; + self.stdin.flush().await?; + + // Read lines from stdout, skipping non-JSON garbage that a native + // library may have written to fd 1 during import (e.g. weasyprint via + // cffi). Collect leaked lines so we can surface them in errors. + let mut leaked_stdout: Vec = Vec::new(); + + let resp: RPCResponse = loop { + let mut resp_line = String::new(); + let n = self.stdout.read_line(&mut resp_line).await?; + + if n == 0 { + trace!("Worker RPC: stdout EOF"); + return Err(if leaked_stdout.is_empty() { + anyhow!("Worker process closed stdout") + } else { + anyhow!( + "Worker process closed stdout after writing non-JSON output \ + (a library may have written to stdout during import):\n{}", + leaked_stdout.join("") + ) + }); + } + + trace!("Worker RPC <- {}", resp_line.trim()); + + let trimmed = resp_line.trim(); + + if !trimmed.is_empty() + && let Ok(resp) = serde_json::from_str::(trimmed) + { + if !leaked_stdout.is_empty() { + trace!( + "Worker RPC: skipped {} non-JSON line(s) on stdout", + leaked_stdout.len() + ); + } + break resp; + } + + leaked_stdout.push(resp_line); + + if leaked_stdout.len() >= 50 { + return Err(anyhow!( + "Expected JSON-RPC response from worker but got {} lines of \ + non-JSON output (a library may have written to stdout during \ + import):\n{}", + leaked_stdout.len(), + leaked_stdout.join("") + )); + } + }; + + if let Some(err) = resp.error { + let detail = if let Some(tb) = &err.traceback { + format!("RPC error {}: {}\n{tb}", err.code, err.message) + } else { + format!("RPC error {}: {}", err.code, err.message) + }; + return Err(anyhow!(detail)); + } + + let val = resp.result.unwrap_or(serde_json::Value::Null); + + Ok(serde_json::from_value(val)?) + } + + /// Run a discovered test or doctest in the worker process. + /// + /// # Errors + /// Returns an error if the request cannot be serialized, if worker I/O + /// fails, or if the worker returns a JSON-RPC error. + pub async fn run_test(&mut self, test: &TestItem) -> Result { + if let Some(object_path) = &test.doctest_object { + return self.run_doctest(test, object_path).await; + } + let params = serde_json::to_value(RunTestParams { + module: test.module_path.clone(), + function: test.name.clone(), + xfail: test.xfail.clone(), + groups: test.groups.clone(), + case_label: test.case_label.clone(), + })?; + let wire: RunTestResultWire = self.call(RPCRequestMethod::RunTest, Some(params)).await?; + Ok(convert_wire_result(test.clone(), wire)) + } + + /// Send hook metadata for a module to the Python worker. + /// Must be called before running any tests from that module. + /// + /// # Errors + /// Returns an error if hook metadata cannot be serialized or the worker + /// rejects the registration request. + pub async fn register_hooks(&mut self, params: RegisterHooksParams) -> Result<()> { + let value = serde_json::to_value(params)?; + self.call::(RPCRequestMethod::RegisterHooks, Some(value)) + .await?; + Ok(()) + } + + /// Tell the Python worker to run scope-level teardown for `per="scope"` + /// fixtures in a module. Must be called after all tests from that module + /// have run. + /// + /// # Errors + /// Returns an error if the finalize request cannot be serialized or the + /// worker reports a teardown failure. + pub async fn finalize_hooks(&mut self, module: String) -> Result<()> { + let value = serde_json::to_value(FinalizeHooksParams { module })?; + self.call::(RPCRequestMethod::FinalizeHooks, Some(value)) + .await?; + Ok(()) + } + + async fn run_doctest(&mut self, test: &TestItem, object_path: &str) -> Result { + let params = serde_json::to_value(RunDoctestParams { + module: test.module_path.clone(), + object_path: object_path.to_owned(), + })?; + let wire: RunTestResultWire = self + .call(RPCRequestMethod::RunDoctest, Some(params)) + .await?; + Ok(convert_wire_result(test.clone(), wire)) + } + + /// Verify that the worker process is responsive. + /// + /// # Errors + /// Returns an error if the ping RPC fails or if the worker returns an + /// unexpected response. + pub async fn ping(&mut self) -> Result<()> { + let result: String = self.call(RPCRequestMethod::Ping, None).await?; + if result == "pong" { + Ok(()) + } else { + Err(anyhow!("Unexpected ping response: {result}")) + } + } + + /// Snapshot the buffered worker stderr and clear the buffer. + /// + /// Called on the error path when the worker is about to be + /// discarded. We kill the child first so the drainer task reaches + /// EOF, then await the drainer (with a short timeout) so any bytes + /// still in the kernel pipe at the time of the failure land in + /// `stderr_buf` before we snapshot. Without this, a python worker + /// that dies during startup (e.g. `ModuleNotFoundError: tryke`) + /// can lose its traceback to a race between the RPC's + /// `Broken pipe` and the drainer task being scheduled. + /// + /// # Panics + /// Panics only if the stderr-drainer task panicked while holding the + /// internal mutex (poisoning it). That task does no fallible work. + pub async fn drain_stderr(&mut self) -> String { + // Worker is about to be discarded; killing the child closes its + // stderr pipe which gives the drainer its EOF. Idempotent on a + // process that has already exited. + let _ = self.child.start_kill(); + if let Some(handle) = self.stderr_drainer.take() { + let _ = tokio::time::timeout(Duration::from_millis(500), handle).await; + } + let bytes: Vec = { + let mut g = self + .stderr_buf + .lock() + .expect("stderr buffer mutex poisoned"); + std::mem::take(&mut *g).into_iter().collect() + }; + String::from_utf8_lossy(&bytes).into_owned() + } + + pub async fn shutdown(&mut self) { + let _ = self.child.kill().await; + // Killing the child closes stderr; the drainer will see EOF and + // exit on its own — but `abort()` is the explicit, immediate + // signal that we're done with it, and prevents a leftover task + // from briefly holding the stderr FD + `stderr_buf` Arc past + // the worker's lifetime. + if let Some(handle) = self.stderr_drainer.take() { + handle.abort(); + } + } +} + +impl Drop for WorkerProcess { + fn drop(&mut self) { + // Safety net: ensure the child process is killed when the worker is + // dropped (e.g. on the error-respawn path in worker.rs). start_kill() is + // the synchronous variant — safe to call on already-dead processes. + let _ = self.child.start_kill(); + // Dropping a Tokio JoinHandle detaches the task, so abort it + // explicitly to avoid an orphan drainer outliving the worker on + // respawn paths (the drainer holds stderr_buf + the stderr FD). + if let Some(handle) = self.stderr_drainer.take() { + handle.abort(); + } + } +} + +/// Spawn a tokio task that continuously reads `stderr` into `buf` until +/// EOF or a read error, capping the buffer at `STDERR_RETAIN_BYTES`. +/// +/// Returns an error if no Tokio runtime is currently entered, rather than +/// panicking the way `tokio::spawn` would. This keeps the synchronous +/// `WorkerProcess::spawn` API safe to call from any context — callers in +/// non-async code receive a structured error instead of a panic. +fn spawn_stderr_drainer( + stderr: tokio::process::ChildStderr, + buf: Arc>>, +) -> Result> { + let handle = tokio::runtime::Handle::try_current() + .map_err(|e| anyhow!("WorkerProcess::spawn requires an active tokio runtime: {e}"))?; + Ok(handle.spawn(async move { + let mut reader = stderr; + let mut chunk = [0u8; 8192]; + loop { + match reader.read(&mut chunk).await { + Ok(0) => break, + Ok(n) => append_stderr(&buf, &chunk[..n]), + Err(e) => { + trace!("Stderr drainer: read error: {e}"); + break; + } + } + } + })) +} + +/// Append `data` to `buf`, capping it at `STDERR_RETAIN_BYTES` by +/// dropping the oldest bytes. `VecDeque::drain` removes from the front +/// in `O(excess)` time, making steady-state appends `O(data.len())` +/// rather than `O(STDERR_RETAIN_BYTES)` as a contiguous `Vec` would require. +fn append_stderr(buf: &Mutex>, data: &[u8]) { + let mut g = buf.lock().expect("stderr buffer mutex poisoned"); + if data.len() >= STDERR_RETAIN_BYTES { + // A single chunk already fills the cap; keep only its tail. + g.clear(); + g.extend(data[data.len() - STDERR_RETAIN_BYTES..].iter().copied()); + return; + } + let total = g.len() + data.len(); + if total > STDERR_RETAIN_BYTES { + let excess = total - STDERR_RETAIN_BYTES; + g.drain(..excess); + } + g.extend(data.iter().copied()); +} + +/// Translate a resolved log level into the value placed on the spawned +/// worker's `TRYKE_LOG` environment variable. +/// +/// The canonical lowercase value keeps worker configuration deterministic, +/// including explicit propagation of `off`. +fn worker_log_env_value(log_level: log::LevelFilter) -> String { + log_level.as_str().to_ascii_lowercase() +} + +fn build_pythonpath(extra: &[&Path]) -> String { + let existing = std::env::var("PYTHONPATH").unwrap_or_default(); + let mut parts: Vec = extra + .iter() + .map(|p| { + let s = p + .canonicalize() + .unwrap_or_else(|_| p.to_path_buf()) + .to_string_lossy() + .into_owned(); + // std::fs::canonicalize on windows produces \\?\ extended-length + // paths that python doesn't understand + #[cfg(windows)] + let s = s.strip_prefix(r"\\?\").unwrap_or(&s).to_string(); + s + }) + .collect(); + if !existing.is_empty() { + parts.push(existing); + } + let sep = if cfg!(windows) { ";" } else { ":" }; + parts.join(sep) +} + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + + use tryke_types::{AssertionWire, ExpectedAssertion, TestOutcome}; + + use super::*; + + fn make_test_item() -> TestItem { + TestItem { + name: "test_add".into(), + module_path: "tests.test_math".into(), + file_path: Some(PathBuf::from("tests/test_math.py")), + line_number: Some(5), + display_name: None, + expected_assertions: vec![], + ..Default::default() + } + } + + #[test] + fn worker_log_env_value_canonicalizes_every_level() { + for (level, expected) in [ + (log::LevelFilter::Off, "off"), + (log::LevelFilter::Error, "error"), + (log::LevelFilter::Warn, "warn"), + (log::LevelFilter::Info, "info"), + (log::LevelFilter::Debug, "debug"), + (log::LevelFilter::Trace, "trace"), + ] { + assert_eq!(worker_log_env_value(level), expected); + } + } + + #[test] + fn convert_result_passed() { + let test = make_test_item(); + let wire = RunTestResultWire::Passed { + duration_ms: 10, + stdout: "out".into(), + stderr: "err".into(), + }; + let result = convert_wire_result(test, wire); + assert!(matches!(result.outcome, TestOutcome::Passed)); + assert_eq!(result.duration, Duration::from_millis(10)); + assert_eq!(result.stdout, "out"); + assert_eq!(result.stderr, "err"); + } + + #[test] + fn convert_result_failed() { + let test = make_test_item(); + let wire = RunTestResultWire::Failed { + duration_ms: 5, + message: "expected 1 got 2".into(), + traceback: None, + assertions: vec![], + executed_lines: vec![], + stdout: String::new(), + stderr: String::new(), + }; + let result = convert_wire_result(test, wire); + assert!(matches!( + result.outcome, + TestOutcome::Failed { ref message, .. } if message == "expected 1 got 2" + )); + } + + #[test] + fn convert_result_skipped() { + let test = make_test_item(); + let wire = RunTestResultWire::Skipped { + duration_ms: 0, + reason: Some("not ready".into()), + stdout: String::new(), + stderr: String::new(), + }; + let result = convert_wire_result(test, wire); + assert!(matches!( + result.outcome, + TestOutcome::Skipped { reason: Some(ref r) } if r == "not ready" + )); + } + + #[test] + fn build_pythonpath_joins_paths() { + let dir_a = tempfile::tempdir().unwrap(); + let dir_b = tempfile::tempdir().unwrap(); + let a = build_pythonpath(&[dir_a.path()]); + let b = build_pythonpath(&[dir_b.path()]); + let result = build_pythonpath(&[dir_a.path(), dir_b.path()]); + let sep = if cfg!(windows) { ";" } else { ":" }; + assert_eq!(result, format!("{a}{sep}{b}")); + } + + #[test] + fn convert_result_uses_discovered_multiline_assertion_source() { + let expression = + "t.expect(\n expr=actual,\n name=\"actual should be one\",\n).to_equal(other=1)"; + let test = TestItem { + expected_assertions: vec![ExpectedAssertion { + subject: "actual".into(), + matcher: "to_equal".into(), + negated: false, + args: vec!["other=1".into()], + line: 7, + label: Some("actual should be one".into()), + end_line: 10, + start_column: Some(4), + end_column: Some(18), + expression: expression.into(), + subject_span: expression + .find("actual") + .map(|offset| (offset, "actual".len())), + expected_arg_span: expression + .find("other=1") + .map(|offset| (offset, "other=1".len())), + expected_arg_value: Some("1".into()), + }], + ..Default::default() + }; + let result = convert_wire_result( + test, + RunTestResultWire::Failed { + duration_ms: 1, + message: "assertion failed".into(), + traceback: None, + assertions: vec![AssertionWire { + expression: ".to_equal(other=1)".into(), + expected: "1".into(), + received: "0".into(), + line: 10, + column: Some(6), + file: Some("tests/test_multiline.py".into()), + }], + executed_lines: vec![10], + stdout: String::new(), + stderr: String::new(), + }, + ); + let TestOutcome::Failed { + assertions, + executed_lines, + .. + } = result.outcome + else { + panic!("expected failed outcome"); + }; + let assertion = &assertions[0]; + assert_eq!(assertion.expression, expression); + assert_eq!(assertion.line, 7); + assert_eq!(assertion.span_offset, expression.find("actual").unwrap()); + assert_eq!( + assertion.expected_arg_span, + expression.find("other=1").map(|offset| (offset, 7)) + ); + assert_eq!(executed_lines, vec![7]); + } + + #[test] + fn convert_result_selects_inner_assertion_by_column() { + let test = TestItem { + expected_assertions: vec![ + ExpectedAssertion { + subject: "expect(0).to_equal(1)".into(), + matcher: "to_be_truthy".into(), + negated: false, + args: vec![], + line: 3, + end_line: 3, + start_column: Some(4), + end_column: Some(45), + expression: "expect(expect(0).to_equal(1)).to_be_truthy()".into(), + subject_span: Some((7, 21)), + expected_arg_span: None, + expected_arg_value: None, + label: None, + }, + ExpectedAssertion { + subject: "0".into(), + matcher: "to_equal".into(), + negated: false, + args: vec!["1".into()], + line: 3, + end_line: 3, + start_column: Some(11), + end_column: Some(32), + expression: "expect(0).to_equal(1)".into(), + subject_span: Some((7, 1)), + expected_arg_span: Some((19, 1)), + expected_arg_value: Some("1".into()), + label: None, + }, + ], + ..Default::default() + }; + let result = convert_wire_result( + test, + RunTestResultWire::Failed { + duration_ms: 1, + message: "assertion failed".into(), + traceback: None, + assertions: vec![AssertionWire { + expression: "expect(expect(0).to_equal(1)).to_be_truthy()".into(), + expected: "1".into(), + received: "0".into(), + line: 3, + column: Some(21), + file: None, + }], + executed_lines: vec![3], + stdout: String::new(), + stderr: String::new(), + }, + ); + let TestOutcome::Failed { assertions, .. } = result.outcome else { + panic!("expected failed outcome"); + }; + assert_eq!(assertions[0].expression, "expect(0).to_equal(1)"); + assert_eq!(assertions[0].span_offset, 7); + assert_eq!(assertions[0].expected_arg_span, Some((19, 1))); + } + + #[test] + fn convert_result_matches_same_line_assertion_by_expected_value_without_column() { + let second_expression = "expect(b).to_equal(other=2)"; + let test = TestItem { + expected_assertions: vec![ + ExpectedAssertion { + subject: "a".into(), + matcher: "to_equal".into(), + negated: false, + args: vec!["1".into()], + line: 5, + end_line: 5, + expression: "expect(a).to_equal(1)".into(), + subject_span: Some((7, 1)), + expected_arg_span: Some((19, 1)), + label: None, + ..Default::default() + }, + ExpectedAssertion { + subject: "b".into(), + matcher: "to_equal".into(), + negated: false, + args: vec!["other=2".into()], + line: 5, + end_line: 5, + expression: second_expression.into(), + subject_span: Some((7, 1)), + expected_arg_span: second_expression + .find("other=2") + .map(|offset| (offset, "other=2".len())), + label: None, + ..Default::default() + }, + ], + ..Default::default() + }; + let result = convert_wire_result( + test, + RunTestResultWire::Failed { + duration_ms: 1, + message: "assertion failed".into(), + traceback: None, + assertions: vec![AssertionWire { + expression: "expect(a).to_equal(1); expect(b).to_equal(other=2)".into(), + expected: "2".into(), + received: "0".into(), + line: 5, + column: None, + file: None, + }], + executed_lines: vec![5], + stdout: String::new(), + stderr: String::new(), + }, + ); + let TestOutcome::Failed { assertions, .. } = result.outcome else { + panic!("expected failed outcome"); + }; + assert_eq!(assertions[0].expression, second_expression); + assert_eq!( + assertions[0].expected_arg_span, + second_expression + .find("other=2") + .map(|offset| (offset, "other=2".len())) + ); + } + + #[test] + fn convert_result_keeps_runtime_expression_for_ambiguous_same_line_without_column() { + let runtime_expression = "expect(a).to_equal(1); expect(b).to_equal(1)"; + let test = TestItem { + expected_assertions: vec![ + ExpectedAssertion { + subject: "a".into(), + matcher: "to_equal".into(), + negated: false, + args: vec!["1".into()], + line: 5, + end_line: 5, + expression: "expect(a).to_equal(1)".into(), + subject_span: Some((7, 1)), + label: None, + ..Default::default() + }, + ExpectedAssertion { + subject: "b".into(), + matcher: "to_equal".into(), + negated: false, + args: vec!["1".into()], + line: 5, + end_line: 5, + expression: "expect(b).to_equal(1)".into(), + subject_span: Some((7, 1)), + label: None, + ..Default::default() + }, + ], + ..Default::default() + }; + let result = convert_wire_result( + test, + RunTestResultWire::Failed { + duration_ms: 1, + message: "assertion failed".into(), + traceback: None, + assertions: vec![AssertionWire { + expression: runtime_expression.into(), + expected: "1".into(), + received: "0".into(), + line: 5, + column: None, + file: None, + }], + executed_lines: vec![5], + stdout: String::new(), + stderr: String::new(), + }, + ); + let TestOutcome::Failed { assertions, .. } = result.outcome else { + panic!("expected failed outcome"); + }; + assert_eq!(assertions[0].expression, runtime_expression); + } + + #[tokio::test] + async fn drop_kills_child_process() { + let mut child = tokio::process::Command::new("sleep") + .arg("60") + .spawn() + .expect("failed to spawn sleep"); + let pid = child.id().expect("missing pid"); + + // Wrap in a WorkerProcess-like drop: start_kill then drop + let _ = child.start_kill(); + drop(child); + + // Give the OS a moment to reap the process + tokio::time::sleep(Duration::from_millis(50)).await; + + // On Unix, sending signal 0 checks if the process exists + let status = std::process::Command::new("kill") + .args(["-0", &pid.to_string()]) + .status(); + assert!( + status.is_ok_and(|s| !s.success()), + "child process {pid} should be dead after start_kill + drop" + ); + } + + #[test] + fn append_stderr_caps_at_retain_bytes() { + let buf = Mutex::new(VecDeque::::new()); + // Fill with a recognisable head sequence, then push past the cap. + let head = vec![b'A'; STDERR_RETAIN_BYTES]; + append_stderr(&buf, &head); + let tail = vec![b'B'; 1024]; + append_stderr(&buf, &tail); + + let mut g = buf.lock().unwrap(); + assert_eq!(g.len(), STDERR_RETAIN_BYTES); + // Oldest A's are dropped; tail B's retained. + let bytes = g.make_contiguous(); + assert_eq!(&bytes[bytes.len() - tail.len()..], &tail[..]); + assert_eq!(bytes[0], b'A'); + } + + #[test] + fn append_stderr_handles_single_oversized_write() { + let buf = Mutex::new(VecDeque::::new()); + let big = vec![b'X'; STDERR_RETAIN_BYTES + 4096]; + append_stderr(&buf, &big); + let mut g = buf.lock().unwrap(); + assert_eq!(g.len(), STDERR_RETAIN_BYTES); + assert!(g.make_contiguous().iter().all(|b| *b == b'X')); + } + + /// Reproducer for the "tryke hangs at `finalize_hooks`" bug: when the + /// worker writes more than the kernel pipe buffer (~64 KiB on Linux, + /// ~4 KiB on Windows) to stderr without anyone draining it, the + /// worker's next stderr write blocks and it stops responding to + /// RPCs. The drainer task spawned in `WorkerProcess::spawn` must + /// keep the pipe empty so RPC flow is unaffected. + /// + /// Test contract is binary: either the drainer prevents the deadlock + /// and `done` arrives on stdout, or it doesn't. The cap-retention + /// property is covered by the `append_stderr_*` unit tests above and + /// is intentionally NOT asserted here — re-asserting it via a second + /// post-exit polling loop introduced wall-clock flakiness on slow + /// Windows runners without adding any signal a unit test wouldn't + /// catch. + /// + /// Uses `python3` (already a project dependency) for portability and + /// routes through the same `spawn_stderr_drainer` helper as + /// `WorkerProcess::spawn`, so a regression in the production drainer + /// path would deadlock this test too. + #[tokio::test] + async fn worker_continues_after_large_stderr_write() { + // 256 KiB is comfortably above both Linux's 64 KiB pipe buffer + // and Windows's 4 KiB default — enough to deadlock without a + // drainer. Larger values (we used to write 1 MiB) just make the + // test slower without proving anything additional. + const STDERR_BYTES: usize = 256 * 1024; + let mut child = tokio::process::Command::new("python3") + .args([ + "-c", + &format!( + "import sys; sys.stderr.write('X' * {STDERR_BYTES}); \ + sys.stderr.flush(); print('done', flush=True)" + ), + ]) + .stdin(std::process::Stdio::null()) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .spawn() + .expect("spawn python3"); + + let stderr = child.stderr.take().expect("no stderr"); + let stderr_buf: Arc>> = Arc::new(Mutex::new(VecDeque::new())); + spawn_stderr_drainer(stderr, Arc::clone(&stderr_buf)) + .expect("drainer requires tokio runtime"); + + let mut stdout = BufReader::new(child.stdout.take().expect("no stdout")); + let mut line = String::new(); + // 30s catches a real deadlock (which would never finish) while + // tolerating Windows python startup + scheduler jitter. Past 5s + // budgets occasionally tripped on Windows × 3.14 even though no + // deadlock occurred. + tokio::time::timeout(Duration::from_secs(30), stdout.read_line(&mut line)) + .await + .expect("stdout read timed out — drainer deadlocked") + .expect("stdout read errored"); + assert_eq!(line.trim(), "done"); + + let _ = child.wait().await; + } +} diff --git a/crates/tryke_server/Cargo.toml b/crates/tryke_server/Cargo.toml index b09d3c4..40e8163 100644 --- a/crates/tryke_server/Cargo.toml +++ b/crates/tryke_server/Cargo.toml @@ -18,6 +18,7 @@ serde = { workspace = true } serde_json = { workspace = true } tokio = { workspace = true } tokio-stream = { workspace = true } +tokio-util = { workspace = true } tryke_discovery = { workspace = true, features = ["filesystem"] } tryke_reporter = { workspace = true, features = ["terminal"] } tryke_runner = { workspace = true } diff --git a/crates/tryke_server/src/handler.rs b/crates/tryke_server/src/handler.rs index 91ea964..adf10c0 100644 --- a/crates/tryke_server/src/handler.rs +++ b/crates/tryke_server/src/handler.rs @@ -10,6 +10,7 @@ use tokio::{ sync::{Mutex, mpsc}, }; use tokio_stream::StreamExt; +use tokio_util::sync::CancellationToken; use tryke_runner::{DistMode, WorkerPool, partition_with_hooks}; use tryke_types::filter::TestFilter; use tryke_types::{RunSummary, TestItem, TestOutcome}; @@ -20,8 +21,122 @@ use crate::protocol::{ RunParams, RunResponse, RunStartParams, TestCompleteParams, }; -/// Manages communication with a client over the given reader/writer -pub struct ConnectionHandler { +pub(crate) struct RunRequest { + id: Value, + params: RunParams, +} + +/// Owns the worker pool and executes queued runs in arrival order. +pub(crate) struct RunDispatcher { + discoverer: Arc>, + worker_pool: WorkerPool, + outbound_tx: mpsc::Sender, + run_rx: mpsc::Receiver, + cancellation: CancellationToken, +} + +impl RunDispatcher { + pub(crate) fn new( + discoverer: Arc>, + worker_pool: WorkerPool, + outbound_tx: mpsc::Sender, + run_rx: mpsc::Receiver, + cancellation: CancellationToken, + ) -> Self { + Self { + discoverer, + worker_pool, + outbound_tx, + run_rx, + cancellation, + } + } + + pub(crate) async fn run(mut self) -> anyhow::Result<()> { + let mut dispatch_result = Ok(()); + + loop { + let request = tokio::select! { + biased; + () = self.cancellation.cancelled() => break, + request = self.run_rx.recv() => { + let Some(request) = request else { + break; + }; + request + } + }; + + let execution = execute_run( + request.params, + &self.discoverer, + &self.outbound_tx, + &self.worker_pool, + ); + tokio::pin!(execution); + + let result = tokio::select! { + biased; + () = self.cancellation.cancelled() => break, + result = &mut execution => result, + }; + + let (run_id, summary) = match result { + Ok(result) => result, + Err(error) => { + dispatch_result = Err(error); + break; + } + }; + + let response = Response::new(request.id, RunResponse { run_id, summary }) + .into_json_line() + .context("Failed to serialize run response"); + + let response = match response { + Ok(response) => response, + Err(error) => { + dispatch_result = Err(error); + break; + } + }; + + let send_result = tokio::select! { + biased; + () = self.cancellation.cancelled() => break, + result = self.outbound_tx.send(response) => result, + }; + + if send_result.is_err() { + dispatch_result = Err(anyhow::anyhow!( + "Outbound channel closed while sending run response" + )); + break; + } + } + + let shutdown_result = self.worker_pool.shutdown().await; + + combine_shutdown(dispatch_result, shutdown_result) + } +} + +fn combine_shutdown( + result: anyhow::Result<()>, + shutdown_result: anyhow::Result<()>, +) -> anyhow::Result<()> { + match (result, shutdown_result) { + (Ok(()), Ok(())) => Ok(()), + (Ok(()), Err(shutdown_error)) => Err(shutdown_error), + (Err(error), Ok(())) => Err(error), + (Err(error), Err(shutdown_error)) => Err(error.context(format!( + "Worker pool shutdown also failed: {shutdown_error:#}" + ))), + } +} + +/// Manages communication with a client over the given reader/writer. +pub(crate) struct ConnectionHandler { /// Reader to read message off of. reader: R, @@ -31,8 +146,8 @@ pub struct ConnectionHandler { /// Test discoverer. discoverer: Arc>, - /// Worker pool. - worker_pool: Arc, + /// FIFO run queue. + run_tx: mpsc::Sender, /// Outbound message queue. /// @@ -40,8 +155,8 @@ pub struct ConnectionHandler { outbound_rx: mpsc::Receiver, outbound_tx: mpsc::Sender, - /// Serializes test runs that share the worker pool. - run_lock: Arc>, + /// Session-wide cancellation. + cancellation: CancellationToken, } impl ConnectionHandler @@ -49,23 +164,23 @@ where R: AsyncRead + Unpin, W: AsyncWrite + Unpin + Send + 'static, { - pub fn new( + pub(crate) fn new( reader: R, writer: W, discoverer: Arc>, outbound_rx: mpsc::Receiver, outbound_tx: mpsc::Sender, - worker_pool: Arc, - run_lock: Arc>, + run_tx: mpsc::Sender, + cancellation: CancellationToken, ) -> Self { Self { reader, writer, discoverer, - worker_pool, + run_tx, outbound_rx, outbound_tx, - run_lock, + cancellation, } } @@ -82,15 +197,12 @@ where reader, writer, discoverer, - worker_pool: pool, + run_tx, outbound_rx, outbound_tx, - run_lock, + cancellation, } = self; - // The writer task is the sole owner of the client output. Routing both - // responses and asynchronous notifications through it prevents - // concurrent writes from interleaving JSON-RPC messages. let mut writer_task = tokio::spawn(async move { let mut writer = BufWriter::new(writer); let mut outbound_rx = outbound_rx; @@ -112,9 +224,9 @@ where loop { line.clear(); - // A connection cannot make progress once its writer stops, so wait - // for either the next request or early writer termination. let read_result = tokio::select! { + biased; + () = cancellation.cancelled() => break, result = reader.read_line(&mut line) => result, result = &mut writer_task => { return result.context("outbound writer task failed")?; @@ -133,37 +245,44 @@ where } } - // Request handlers can enqueue notifications and optionally return - // a response. The response goes through the same queue so only the - // writer task ever touches the transport. - let response = - match handle_request(&line, &discoverer, &outbound_tx, &pool, &run_lock).await { - Ok(response) => response, - Err(_) if outbound_tx.is_closed() => { - // The writer owns the useful transport error. Await it - // instead of returning the secondary channel error. - return writer_task.await.context("outbound writer task failed")?; - } - Err(error) => { - // Stop the writer before returning a request-processing - // error so it cannot outlive the connection. - writer_task.abort(); - let _ = writer_task.await; - return Err(error); - } - }; + let request = handle_request(&line, &discoverer, &outbound_tx, &run_tx); + tokio::pin!(request); + + let request_result = tokio::select! { + biased; + () = cancellation.cancelled() => break, + result = &mut request => result, + }; - if let Some(bytes) = response - && outbound_tx.send(bytes).await.is_err() - { - // A closed receiver means the writer has already stopped. - return writer_task.await.context("outbound writer task failed")?; + let response = match request_result { + Ok(response) => response, + Err(_) if outbound_tx.is_closed() => { + // The writer owns the useful transport error. Await it + // instead of returning the secondary channel error. + return writer_task.await.context("outbound writer task failed")?; + } + Err(error) => { + // Stop the writer before returning a request-processing + // error so it cannot outlive the connection. + writer_task.abort(); + let _ = writer_task.await; + return Err(error); + } + }; + + if let Some(bytes) = response { + let send_result = tokio::select! { + biased; + () = cancellation.cancelled() => break, + result = outbound_tx.send(bytes) => result, + }; + if send_result.is_err() { + // A closed receiver means the writer has already stopped. + return writer_task.await.context("outbound writer task failed")?; + } } } - // Other server tasks retain outbound sender clones, so EOF alone does - // not close the queue. Cancel the writer explicitly; cancellation is - // therefore the expected successful shutdown result. writer_task.abort(); match writer_task.await { @@ -238,12 +357,12 @@ pub(crate) async fn apply_change( let impact = discoverer.lock().await.apply_changes(paths); if impact.paths.is_empty() { - debug!("apply_change: no eligible paths"); + debug!("Apply_change: no eligible paths"); return Ok(()); } debug!( - "apply_change: {} affected modules, {} affected tests", + "Apply_change: {} affected modules, {} affected tests", impact.affected_modules.len(), impact.affected_tests.len(), ); @@ -281,11 +400,14 @@ async fn execute_run( discoverer: &tokio::sync::Mutex, outbound_tx: &mpsc::Sender, pool: &WorkerPool, - run_lock: &Mutex<()>, ) -> anyhow::Result<(String, RunSummary)> { let run_id = run_params.run_id.clone(); - let _run_guard = run_lock.lock().await; - pool.restart_workers().await; + // Every server run must execute against current source. A worker that + // cannot be restarted is still holding the previous interpreter, so fail + // the run rather than reporting stale results as fresh ones. + pool.restart_workers() + .await + .context("Failed to restart workers before run")?; let discovery_start = Instant::now(); let (all_tests, hooks) = { let guard = discoverer.lock().await; @@ -378,12 +500,11 @@ async fn execute_run( /// /// # Errors /// Returns an error if an outbound message cannot be serialized or queued. -pub async fn handle_request( +pub(crate) async fn handle_request( line: &str, discoverer: &tokio::sync::Mutex, outbound_tx: &mpsc::Sender, - worker_pool: &WorkerPool, - run_lock: &Mutex<()>, + run_tx: &mpsc::Sender, ) -> anyhow::Result> { let Ok(req) = serde_json::from_str::(line.trim()) else { return Ok(None); @@ -410,45 +531,8 @@ pub async fn handle_request( Response::new(id, serde_json::json!({ "tests": tests })).into_json_line()? } RequestMethod::DidChange => { - let Some(params) = req.params else { - return Ok(Some( - ErrorResponse::new( - req.id, - INVALID_PARAMS, - "method 'did_change' requires params with paths".to_string(), - ) - .into_json_line()?, - )); - }; - let dc = match serde_json::from_value::(params) { - Ok(dc) => dc, - Err(e) => { - return Ok(Some( - ErrorResponse::new( - req.id, - INVALID_PARAMS, - format!("invalid params for 'did_change': {e}"), - ) - .into_json_line()?, - )); - } - }; - if dc.paths.is_empty() { - let tests = discoverer.lock().await.rediscover(); - debug!( - "did_change: empty paths — full rediscover, {} tests", - tests.len() - ); - send_notification( - outbound_tx, - NotificationMethod::DiscoverComplete, - DiscoverCompleteParams { tests }, - ) - .await?; - } else { - apply_change(discoverer, outbound_tx, &dc.paths).await?; - } - Response::new(id, "ok").into_json_line()? + return handle_did_change(req.id.clone(), req.params.clone(), discoverer, outbound_tx) + .await; } RequestMethod::Run => { let Some(params) = req.params else { @@ -474,9 +558,14 @@ pub async fn handle_request( )); } }; - let (run_id, summary) = - execute_run(run_params, discoverer, outbound_tx, worker_pool, run_lock).await?; - Response::new(id, RunResponse { run_id, summary }).into_json_line()? + run_tx + .send(RunRequest { + id, + params: run_params, + }) + .await + .map_err(|_| anyhow::anyhow!("Run dispatcher is closed"))?; + return Ok(None); } RequestMethod::Unknown(method) => ErrorResponse::new( req.id, @@ -488,6 +577,55 @@ pub async fn handle_request( Ok(Some(response)) } +async fn handle_did_change( + id: Option, + params: Option, + discoverer: &tokio::sync::Mutex, + outbound_tx: &mpsc::Sender, +) -> anyhow::Result> { + let Some(params) = params else { + return Ok(Some( + ErrorResponse::new( + id, + INVALID_PARAMS, + "method 'did_change' requires params with paths".to_string(), + ) + .into_json_line()?, + )); + }; + let dc = match serde_json::from_value::(params) { + Ok(dc) => dc, + Err(error) => { + return Ok(Some( + ErrorResponse::new( + id, + INVALID_PARAMS, + format!("invalid params for 'did_change': {error}"), + ) + .into_json_line()?, + )); + } + }; + if dc.paths.is_empty() { + let tests = discoverer.lock().await.rediscover(); + debug!( + "Did_change: empty paths — full rediscover, {} tests", + tests.len() + ); + send_notification( + outbound_tx, + NotificationMethod::DiscoverComplete, + DiscoverCompleteParams { tests }, + ) + .await?; + } else { + apply_change(discoverer, outbound_tx, &dc.paths).await?; + } + Ok(Some( + Response::new(id.unwrap_or(Value::Null), "ok").into_json_line()?, + )) +} + #[cfg(test)] mod tests { use std::{ @@ -549,24 +687,55 @@ mod tests { TestProject::new().expect("create test project") } - async fn make_pool() -> Arc { - Arc::new( - WorkerPool::spawn_from_parts( - 1, - &test_python_bin(), - std::path::Path::new("."), - None, - LevelFilter::Off, - false, - ) - .await, + async fn make_owned_pool() -> WorkerPool { + WorkerPool::spawn_from_parts( + 1, + &test_python_bin(), + std::path::Path::new("."), + None, + LevelFilter::Off, + false, ) + .await + } + + async fn make_pool() -> Arc { + Arc::new(make_owned_pool().await) } fn make_run_lock() -> Arc> { Arc::new(Mutex::new(())) } + /// Compatibility harness for request-level tests. Production requests are + /// queued by `super::handle_request`; the harness drains that queue and + /// executes the request so existing assertions can inspect a completed + /// response without constructing a full connection. + async fn handle_request( + line: &str, + discoverer: &Mutex, + outbound_tx: &mpsc::Sender, + worker_pool: &WorkerPool, + run_lock: &Mutex<()>, + ) -> anyhow::Result> { + let (run_tx, mut run_rx) = mpsc::channel(1); + let response = super::handle_request(line, discoverer, outbound_tx, &run_tx).await?; + if response.is_some() { + return Ok(response); + } + + let request = run_rx + .recv() + .await + .context("Queued request was not available to the test harness")?; + let _run_guard = run_lock.lock().await; + let (run_id, summary) = + execute_run(request.params, discoverer, outbound_tx, worker_pool).await?; + Ok(Some( + Response::new(request.id, RunResponse { run_id, summary }).into_json_line()?, + )) + } + #[tokio::test] async fn notification_reports_serialization_errors() { let (tx, _rx) = mpsc::channel(1); @@ -1306,21 +1475,31 @@ mod tests { &[], None, ))); - let pool = make_pool().await; + let pool = make_owned_pool().await; let (outbound_tx, outbound_rx) = mpsc::channel::(64); - let run_lock = make_run_lock(); + let (run_tx, run_rx) = mpsc::channel(8); + let cancellation = CancellationToken::new(); + let dispatcher = RunDispatcher::new( + Arc::clone(&disc), + pool, + outbound_tx.clone(), + run_rx, + cancellation.clone(), + ); + let dispatcher_task = tokio::spawn(dispatcher.run()); let (client, server_side) = tokio::io::duplex(1 << 16); let (server_r, server_w) = tokio::io::split(server_side); - tokio::spawn(async move { + let handler_cancellation = cancellation.clone(); + let handler_task = tokio::spawn(async move { ConnectionHandler::new( server_r, server_w, disc, outbound_rx, outbound_tx, - pool, - run_lock, + run_tx, + handler_cancellation, ) .run() .await @@ -1353,6 +1532,14 @@ mod tests { } let response = response.expect("loop exits only once the response is seen"); assert_eq!(response["result"]["run_id"], "r1"); + + cancellation.cancel(); + drop(client_w); + handler_task.await.expect("Join connection handler"); + dispatcher_task + .await + .expect("Join run dispatcher") + .expect("Shut down run dispatcher"); } #[tokio::test] @@ -1366,9 +1553,8 @@ mod tests { &[], None, ))); - let pool = make_pool().await; let (outbound_tx, outbound_rx) = mpsc::channel(1); - let run_lock = make_run_lock(); + let (run_tx, _run_rx) = mpsc::channel(1); let (mut client, server_reader) = tokio::io::duplex(64); let handler = ConnectionHandler::new( @@ -1377,8 +1563,8 @@ mod tests { discoverer, outbound_rx, outbound_tx, - pool, - run_lock, + run_tx, + CancellationToken::new(), ); client .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"ping\"}\n") diff --git a/crates/tryke_server/src/lib.rs b/crates/tryke_server/src/lib.rs index 0147caa..db8f559 100644 --- a/crates/tryke_server/src/lib.rs +++ b/crates/tryke_server/src/lib.rs @@ -1,4 +1,4 @@ -pub mod handler; +mod handler; pub mod protocol; pub mod server; diff --git a/crates/tryke_server/src/server.rs b/crates/tryke_server/src/server.rs index 65f5fb2..6cf59f9 100644 --- a/crates/tryke_server/src/server.rs +++ b/crates/tryke_server/src/server.rs @@ -1,499 +1,184 @@ use std::sync::Arc; +use anyhow::Context as _; use log::debug; -use tokio::{ - io::{AsyncRead, AsyncWrite, Stdin, Stdout}, - sync::{Mutex, mpsc}, -}; +use tokio::sync::{Mutex, mpsc}; +use tokio_util::sync::CancellationToken; use tryke_discovery::Discoverer; use tryke_runner::WorkerPool; -#[cfg(test)] -use tryke_watcher::FileChangeBatch; use tryke_watcher::FileWatcher; -use crate::handler::{ConnectionHandler, apply_change}; +use crate::handler::{ConnectionHandler, RunDispatcher, apply_change}; -enum WatchMode { - Filesystem, - #[cfg(test)] - Disabled, - #[cfg(test)] - Manual(mpsc::UnboundedReceiver), +enum SessionExit { + Cancelled, + Handler(anyhow::Result<()>), + Dispatcher(Result, tokio::task::JoinError>), } -pub struct Server { - reader: R, - writer: W, +pub struct Server { worker_pool: WorkerPool, discoverer: Discoverer, - watch_mode: WatchMode, } -impl Server { +impl Server { #[must_use] pub fn new(worker_pool: WorkerPool, discoverer: Discoverer) -> Self { - Self::with_transport( - worker_pool, - discoverer, - tokio::io::stdin(), - tokio::io::stdout(), - ) - } -} - -impl Server { - fn with_transport( - worker_pool: WorkerPool, - discoverer: Discoverer, - reader: R, - writer: W, - ) -> Self { Self { - reader, - writer, worker_pool, discoverer, - watch_mode: WatchMode::Filesystem, } } - #[cfg(test)] - fn without_file_watcher(mut self) -> Self { - self.watch_mode = WatchMode::Disabled; - self - } - - #[cfg(test)] - fn with_manual_file_watcher( - mut self, - changes: mpsc::UnboundedReceiver, - ) -> Self { - self.watch_mode = WatchMode::Manual(changes); - self - } -} - -impl Server -where - R: AsyncRead + Unpin, - W: AsyncWrite + Unpin + Send + 'static, -{ - /// Runs the server over its configured reader and writer. - /// - /// `Server::new` configures process stdin/stdout for the normal - /// editor-child-process protocol. Closing the reader shuts the server down. + /// Runs the server until the client closes its input or cancellation is + /// requested. /// /// # Errors - /// Returns an error if file watching cannot be initialized or the - /// client transport fails. - pub async fn serve(self) -> anyhow::Result<()> { + /// Returns an error if file watching cannot be initialized, the client + /// transport fails, or a server task cannot shut down cleanly. + pub async fn serve_with_cancellation( + self, + cancellation: CancellationToken, + ) -> anyhow::Result<()> { let Self { - reader, - writer, worker_pool, discoverer, - watch_mode, } = self; let root = discoverer.root().to_path_buf(); let excludes = discoverer.excludes().to_vec(); - let worker_pool = Arc::new(worker_pool); - let run_lock = Arc::new(Mutex::new(())); + let lifecycle = cancellation.child_token(); - // Everything sent to the client goes through this queue + // Everything sent to the client goes through this queue. let (outbound_tx, outbound_rx) = mpsc::channel(256); - // Initialize the discoverer and do its initial discovery/populate import graph/test cache + // Initialize the discoverer and populate its import graph and test cache. let discoverer = Arc::new(Mutex::new(discoverer)); discoverer.lock().await.rediscover(); - let disc_for_watcher = Arc::clone(&discoverer); - let outbound_for_watcher = outbound_tx.clone(); - let watcher_task = match watch_mode { - WatchMode::Filesystem => { - let mut watcher = FileWatcher::spawn(&root, &excludes)?; - Some(tokio::spawn(async move { - loop { - let batch = match watcher.next_batch().await { - Ok(Some(batch)) => batch, - Ok(None) => break, - Err(error) => { - debug!("server: stopping file watcher: {error}"); - break; - } - }; - if let Err(error) = - apply_change(&disc_for_watcher, &outbound_for_watcher, &batch.paths) - .await - { - debug!("server: stopping file-change notifications: {error}"); - break; - } - } - })) + let watcher = match FileWatcher::spawn(&root, &excludes) { + Ok(watcher) => watcher, + Err(error) => { + return match worker_pool.shutdown().await { + Ok(()) => Err(error), + Err(shutdown_error) => Err(error.context(format!( + "Worker pool shutdown also failed: {shutdown_error:#}" + ))), + }; } - #[cfg(test)] - WatchMode::Disabled => None, - #[cfg(test)] - WatchMode::Manual(mut changes) => Some(tokio::spawn(async move { - loop { - let Some(batch) = changes.recv().await else { - break; - }; - if let Err(error) = - apply_change(&disc_for_watcher, &outbound_for_watcher, &batch.paths).await - { - debug!("server: stopping file-change notifications: {error}"); - break; - } - } - })), }; + let watcher_task = tokio::spawn(watch_files( + watcher, + Arc::clone(&discoverer), + outbound_tx.clone(), + lifecycle.clone(), + )); - debug!("server: session started"); + let (run_tx, run_rx) = mpsc::channel(256); - // Initialize and run the connection handler - let handler = ConnectionHandler::new( - reader, - writer, + let dispatcher = RunDispatcher::new( Arc::clone(&discoverer), - outbound_rx, - outbound_tx, - Arc::clone(&worker_pool), - run_lock, + worker_pool, + outbound_tx.clone(), + run_rx, + lifecycle.clone(), ); - let handler_result = handler.run().await; + let mut dispatcher_task = tokio::spawn(dispatcher.run()); - debug!("server: session input closed — shutting down"); + debug!("Server: session started"); - if let Some(watcher_task) = watcher_task { - watcher_task.abort(); - let _ = watcher_task.await; - } - if let Ok(pool) = Arc::try_unwrap(worker_pool) { - pool.shutdown(); - } - handler_result - } -} - -#[cfg(test)] -mod tests { - use std::fs; - - use std::time::Duration; - - use log::LevelFilter; - use tokio::{ - io::{AsyncBufReadExt, AsyncWriteExt, BufReader, DuplexStream, ReadHalf, WriteHalf}, - time, - }; - use tryke_testing::{TestProject, python_bin as test_python_bin}; - - use super::*; - - type ClientWriter = WriteHalf; - type ClientReader = BufReader>; - - /// Spawn a server over an in-memory duplex pipe and return the client - /// halves of the session, mirroring how an editor owns the stdio of a - /// spawned `tryke server` child. - fn start_server() -> (ClientWriter, ClientReader, TestProject) { - start_server_inner(None) - } - - fn start_server_with_manual_file_watcher() -> ( - ClientWriter, - ClientReader, - TestProject, - mpsc::UnboundedSender, - ) { - let (changes_tx, changes_rx) = mpsc::unbounded_channel(); - let (writer, reader, directory) = start_server_inner(Some(changes_rx)); - (writer, reader, directory, changes_tx) - } + let handler = ConnectionHandler::new( + tokio::io::stdin(), + tokio::io::stdout(), + Arc::clone(&discoverer), + outbound_rx, + outbound_tx, + run_tx, + lifecycle.clone(), + ); - fn start_server_inner( - manual_changes: Option>, - ) -> (ClientWriter, ClientReader, TestProject) { - let dir = TestProject::new().expect("create test project"); - let root = dir.root().to_path_buf(); - let src_roots = vec![root.clone()]; - let python = test_python_bin(); - let (client, server_side) = tokio::io::duplex(1 << 16); - let (server_r, server_w) = tokio::io::split(server_side); - tokio::spawn(async move { - let worker_pool = - WorkerPool::spawn_from_parts(1, &python, &root, None, LevelFilter::Off, false) - .await; - let discoverer = Discoverer::from_parts(&root, src_roots, &[], None); - let server = Server::with_transport(worker_pool, discoverer, server_r, server_w); - let server = match manual_changes { - Some(changes) => server.with_manual_file_watcher(changes), - None => server.without_file_watcher(), - }; - server.serve().await.expect("server run"); - }); - let (client_r, client_w) = tokio::io::split(client); - (client_w, BufReader::new(client_r), dir) - } + let handler = handler.run(); + tokio::pin!(handler); - #[tokio::test] - async fn ping_pong() { - let (mut w, mut r, _dir) = start_server(); - w.write_all(b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"ping\"}\n") - .await - .unwrap(); - let mut line = String::new(); - // Generous timeout for loaded CI hosts. The worker pool starts cold, - // so ping itself does not wait for Python subprocess startup. - time::timeout(Duration::from_secs(30), r.read_line(&mut line)) - .await - .unwrap() - .unwrap(); - let val: serde_json::Value = serde_json::from_str(line.trim()).unwrap(); - assert_eq!(val["result"], "pong"); - } + let exit = tokio::select! { + biased; + () = lifecycle.cancelled() => SessionExit::Cancelled, + result = &mut handler => SessionExit::Handler(result), + result = &mut dispatcher_task => SessionExit::Dispatcher(result), + }; - #[tokio::test] - async fn stdin_eof_shuts_server_down() { - let dir = TestProject::new().expect("create test project"); - let root = dir.root().to_path_buf(); - let src_roots = vec![root.clone()]; - let python = test_python_bin(); - let (client, server_side) = tokio::io::duplex(1 << 16); - let (server_r, server_w) = tokio::io::split(server_side); - let handle = tokio::spawn(async move { - let worker_pool = - WorkerPool::spawn_from_parts(1, &python, &root, None, LevelFilter::Off, false) - .await; - let discoverer = Discoverer::from_parts(&root, src_roots, &[], None); - Server::with_transport(worker_pool, discoverer, server_r, server_w) - .without_file_watcher() - .serve() - .await - }); - // Closing the client end delivers EOF on the server's input — - // the LSP-style shutdown signal. - drop(client); - let result = time::timeout(Duration::from_secs(30), handle) - .await - .expect("server must shut down after EOF") - .expect("server task must not panic"); - assert!( - result.is_ok(), - "server must exit cleanly on EOF: {result:?}" - ); - } + lifecycle.cancel(); + debug!("Server: session stopping"); + + let (result, secondary_result, secondary_name) = match exit { + SessionExit::Cancelled => ( + handler.await, + dispatcher_task + .await + .context("Run dispatcher task failed") + .and_then(std::convert::identity), + "Run dispatcher", + ), + SessionExit::Handler(result) => ( + result, + dispatcher_task + .await + .context("Run dispatcher task failed") + .and_then(std::convert::identity), + "Run dispatcher", + ), + SessionExit::Dispatcher(result) => ( + result + .context("Run dispatcher task failed") + .and_then(std::convert::identity), + handler.await, + "Connection handler", + ), + }; + let result = match (result, secondary_result) { + (Ok(()), Ok(())) => Ok(()), + (Ok(()), Err(error)) | (Err(error), Ok(())) => Err(error), + (Err(error), Err(secondary_error)) => { + Err(error.context(format!("{secondary_name} also failed: {secondary_error:#}"))) + } + }; - fn match_body(value: &str) -> String { - format!( - "from tryke import describe, expect, test\n\ - \n\ - def match() -> str:\n\ - {INDENT}return \"{value}\"\n\ - \n\ - with describe(\"match\"):\n\ - {INDENT}@test(\"basic\")\n\ - {INDENT}def basic():\n\ - {INDENT}{INDENT}expect(match()).to_equal(\"set\")\n", - INDENT = " ", - ) - } + let watcher_result = watcher_task.await.context("File watcher task failed"); - /// Read JSON-RPC lines from `r` until one with an `id` field (the - /// response — notifications have no `id`). - async fn read_response(r: &mut ClientReader) -> serde_json::Value { - loop { - let mut line = String::new(); - time::timeout(Duration::from_secs(30), r.read_line(&mut line)) - .await - .unwrap() - .unwrap(); - let v: serde_json::Value = serde_json::from_str(line.trim()).unwrap(); - if v.get("id").is_some() { - return v; + match (result, watcher_result) { + (Ok(()), Ok(())) => Ok(()), + (Ok(()), Err(error)) | (Err(error), Ok(())) => Err(error), + (Err(error), Err(watcher_error)) => { + Err(error.context(format!("File watcher also failed: {watcher_error:#}"))) } } } +} - async fn read_notification(r: &mut ClientReader, method: &str) -> serde_json::Value { - loop { - let mut line = String::new(); - time::timeout(Duration::from_secs(30), r.read_line(&mut line)) - .await - .unwrap() - .unwrap(); - let value: serde_json::Value = serde_json::from_str(line.trim()).unwrap(); - if value["method"] == method { - return value; +async fn watch_files( + mut watcher: FileWatcher, + discoverer: Arc>, + outbound_tx: mpsc::Sender, + cancellation: CancellationToken, +) { + loop { + let batch = tokio::select! { + biased; + () = cancellation.cancelled() => break, + batch = watcher.next_batch() => batch, + }; + let batch = match batch { + Ok(Some(batch)) => batch, + Ok(None) => break, + Err(error) => { + debug!("Server: stopping file watcher: {error}"); + break; } - assert!( - value.get("id").is_none(), - "received unexpected response while waiting for {method}: {value}", - ); + }; + if let Err(error) = apply_change(&discoverer, &outbound_tx, &batch.paths).await { + debug!("Server: stopping file-change notifications: {error}"); + break; } } - - /// Send `did_change` then `run` on the SAME session — the invariant - /// that makes the in-band approach race-free. - async fn did_change_then_run( - w: &mut ClientWriter, - r: &mut ClientReader, - file: &std::path::Path, - rid: &str, - ) -> serde_json::Value { - // serde_json::to_string handles JSON escaping (Windows - // backslashes in the path would otherwise produce invalid JSON). - let mut dc = serde_json::to_string(&serde_json::json!({ - "jsonrpc": "2.0", - "id": 1, - "method": "did_change", - "params": { "paths": [file] }, - })) - .unwrap(); - dc.push('\n'); - w.write_all(dc.as_bytes()).await.unwrap(); - let _dc_resp = read_response(r).await; - - let run = format!( - "{{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"run\",\"params\":{{\"run_id\":\"{rid}\"}}}}\n" - ); - w.write_all(run.as_bytes()).await.unwrap(); - read_response(r).await - } - - /// Send `run` only (no `did_change`) — simulates a non-cooperating - /// client. Used to verify the FS-watcher fallback path. - async fn run_only(w: &mut ClientWriter, r: &mut ClientReader, rid: &str) -> serde_json::Value { - let run = format!( - "{{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"run\",\"params\":{{\"run_id\":\"{rid}\"}}}}\n" - ); - w.write_all(run.as_bytes()).await.unwrap(); - read_response(r).await - } - - /// Regression: a `run` issued *immediately* after a save (no sleep - /// to let the FS watcher catch up) must see fresh `sys.modules`, - /// not the previous cycle's cache. The client sends `did_change` - /// first so the server refreshes discovery synchronously before - /// the `run` line is read. - /// - /// Without the `did_change` step, phase 2 can use stale discovery - /// metadata even though workers are fresh for every run. - #[tokio::test] - async fn run_after_did_change_uses_fresh_module() { - let (mut w, mut r, dir) = start_server(); - let test_file = dir.root().join("test_match.py"); - - fs::write(&test_file, match_body("set")).unwrap(); - let resp = did_change_then_run(&mut w, &mut r, &test_file, "set").await; - let summary = &resp["result"]["summary"]; - assert_eq!( - summary["passed"].as_u64().unwrap_or(0), - 1, - "'set' baseline must pass — got summary={summary}", - ); - - // Flip to "st"; the assertion stays "set", so a fresh import must - // fail. A pass means the worker served its phase-1 cached module. - fs::write(&test_file, match_body("st")).unwrap(); - let resp = did_change_then_run(&mut w, &mut r, &test_file, "st").await; - let summary = &resp["result"]["summary"]; - let passed = summary["passed"].as_u64().unwrap_or(0); - let failed = summary["failed"].as_u64().unwrap_or(0); - let errors = summary["errors"].as_u64().unwrap_or(0); - assert!( - passed == 0 && (failed + errors) >= 1, - "file has match()->\"st\" but the run reported passed={passed}; \ - the worker served the stale phase-1 cache. summary={summary}", - ); - } - - /// A file-change event refreshes discovery for clients that do not send - /// `did_change`. The event is injected after the platform watcher boundary - /// so the test remains deterministic across operating systems. - #[tokio::test] - async fn manually_triggered_file_change_refreshes_discovery() { - let (mut w, mut r, dir, changes) = start_server_with_manual_file_watcher(); - let test_file = dir.root().join("test_match.py"); - - w.write_all(b"{\"jsonrpc\":\"2.0\",\"id\":0,\"method\":\"ping\"}\n") - .await - .unwrap(); - let _ = read_response(&mut r).await; - - fs::write(&test_file, match_body("set")).unwrap(); - changes - .send(FileChangeBatch { - paths: vec![test_file.clone()], - }) - .expect("send initial file change"); - let _ = read_notification(&mut r, "discover_complete").await; - let resp = run_only(&mut w, &mut r, "warm").await; - assert_eq!( - resp["result"]["summary"]["passed"].as_u64().unwrap_or(0), - 1, - "warm-up: 'set' body must pass — got {}", - resp["result"]["summary"] - ); - - fs::write(&test_file, match_body("st")).unwrap(); - changes - .send(FileChangeBatch { - paths: vec![test_file], - }) - .expect("send updated file change"); - let _ = read_notification(&mut r, "discover_complete").await; - let resp = run_only(&mut w, &mut r, "after_save").await; - let summary = &resp["result"]["summary"]; - let passed = summary["passed"].as_u64().unwrap_or(0); - let failed = summary["failed"].as_u64().unwrap_or(0); - assert!( - passed == 0 && failed >= 1, - "after the file-change event: 'st' body should fail the 'set' assertion — \ - got summary={summary}", - ); - } - - #[tokio::test] - async fn repeated_runs_reexecute_module_imports() { - let (mut w, mut r, dir) = start_server(); - let test_file = dir.root().join("test_import.py"); - let counter_file = dir.root().join("imports.txt"); - let counter_literal = serde_json::to_string(&counter_file).expect("serialize counter path"); - fs::write( - &test_file, - format!( - "from pathlib import Path\n\ - from tryke import test\n\ - \n\ - counter = Path({counter_literal})\n\ - previous = counter.read_text() if counter.exists() else \"\"\n\ - counter.write_text(previous + \"x\")\n\ - \n\ - @test\n\ - def test_import():\n\ - {INDENT}pass\n", - INDENT = " ", - ), - ) - .expect("write test file"); - - let first = did_change_then_run(&mut w, &mut r, &test_file, "first").await; - assert_eq!(first["result"]["summary"]["passed"], 1); - assert_eq!( - fs::read_to_string(&counter_file).expect("read import counter"), - "x", - ); - - let second = run_only(&mut w, &mut r, "second").await; - assert_eq!(second["result"]["summary"]["passed"], 1); - assert_eq!( - fs::read_to_string(&counter_file).expect("read import counter"), - "xx", - "each logical run must import test modules in a fresh Python process", - ); - } } diff --git a/crates/tryke_watcher/src/lib.rs b/crates/tryke_watcher/src/lib.rs index c17748e..4cccbea 100644 --- a/crates/tryke_watcher/src/lib.rs +++ b/crates/tryke_watcher/src/lib.rs @@ -62,7 +62,7 @@ impl ChangeQueue { paths.dedup(); let paths = self.change_filter.filter(&paths); if paths.is_empty() { - debug!("file watcher: change batch had no meaningful changes"); + debug!("File watcher: change batch had no meaningful changes"); continue; } return Ok(Some(FileChangeBatch { paths })); diff --git a/docs/guides/configuration.md b/docs/guides/configuration.md index 7e64e3d..9138e88 100644 --- a/docs/guides/configuration.md +++ b/docs/guides/configuration.md @@ -120,7 +120,7 @@ tryke server --cache-dir .cache/tryke ## Logging -Tryke has a single user-facing verbosity knob with a precedence chain spanning CLI flags, environment variables, and cross-language propagation to the python workers it spawns. +Tryke resolves one user-facing log level and applies it to Rust logging, reporter diagnostics, and every Python worker it spawns. ### CLI flags @@ -128,40 +128,37 @@ Tryke has a single user-facing verbosity knob with a precedence chain spanning C ### Environment variables -- **`TRYKE_LOG`** — the umbrella knob. Accepts a bare level name (`off`, `error`, `warn`, `info`, `debug`, `trace`) and propagates to **both** the rust process and every python worker it spawns. This is what you should set when you want one knob. -- **`RUST_LOG`** — power-user override for the rust side only. Honored natively by `env_logger`, so the standard per-module filter syntax (`tryke=debug,hyper=warn`) works. Does **not** propagate to python workers — its module-filter grammar doesn't map onto a python log level. +- **`TRYKE_LOG`** — the umbrella knob. Accepts a bare level name (`off`, `error`, `warn`, `info`, `debug`, `trace`) and overrides the CLI level everywhere. Values are case-insensitive and may have surrounding whitespace. An invalid value stops the command with an error instead of being silently ignored. +- **`RUST_LOG`** — power-user override for the Rust side only. Honored natively by `env_logger`, so the standard per-module filter syntax (`tryke=debug,hyper=warn`) works. Does **not** propagate to Python workers — its module-filter grammar doesn't map onto a Python log level. ### Precedence -**Rust log filter** (consumed by `env_logger`): +**Resolved Tryke level**: -1. `RUST_LOG` if set (wins natively). -2. `TRYKE_LOG` if set. -3. The CLI flag (`-v` / `-q`). -4. Default `warn`. +1. `TRYKE_LOG` if set. +2. The CLI flag (`-v` / `-q`). +3. Default `warn`. -**Python worker log** (spawned by tryke, configured by `TRYKE_LOG` on the worker env): +This level drives reporter diagnostics and is passed to Python workers as `TRYKE_LOG`. Python's standard library has no trace level, so workers map `trace` to `debug`. -1. `TRYKE_LOG` if set. -2. The CLI flag, **only** when explicitly more verbose than `warn` (i.e., the user passed at least one `-v`). Default `warn` does not light up workers — preserves the long-standing "no chatter unless asked" behavior. -3. Otherwise off. +The resolved level is also the default Rust filter. If `RUST_LOG` is set, it overrides that filter for Rust logging only; reporter and worker verbosity continue to use the resolved Tryke level. ### Examples ```bash -# Default: rust at warn, workers silent. +# Default: Rust and workers at warn. tryke test -# `-v` lights up both layers at info. +# `-v` raises Rust, reporter, and worker verbosity to info. tryke -v test -# Per-module rust filtering, workers stay silent. +# Per-module Rust filtering; reporter and workers remain at warn. RUST_LOG=tryke=debug,tryke_runner=trace tryke test # Single knob: both layers at debug, regardless of CLI flag. TRYKE_LOG=debug tryke test -# RUST_LOG wins for rust filtering; TRYKE_LOG still drives python. +# RUST_LOG wins for Rust filtering; TRYKE_LOG still drives Python. TRYKE_LOG=info RUST_LOG=tryke=warn tryke test ``` diff --git a/python/tryke/worker.py b/python/tryke/worker.py index 71ef12e..6e8c20d 100644 --- a/python/tryke/worker.py +++ b/python/tryke/worker.py @@ -41,10 +41,11 @@ 4. After every test in a module has run, the runner sends `finalize_hooks` and the executor runs `per="scope"` teardown. 5. In watch/server mode, file changes do not reach the worker over the - wire — the runner instead kills this subprocess and respawns it, - replaying `register_hooks` on the fresh process. `importlib.reload` - is not used; a clean interpreter is the only reliable way to drop - classes and closures captured under the old definitions. + wire — the runner instead kills this subprocess and respawns it. + Each work unit installs its current hook metadata before execution, + and only that active unit's metadata is replayed after a crash. + `importlib.reload` is not used; a clean interpreter is the only + reliable way to drop classes and closures captured under old definitions. """ from __future__ import annotations @@ -306,8 +307,8 @@ def _register_hooks( Any previously-cached :class:`HookExecutor` for this module is dropped so the next test rebuilds fixtures from the fresh - metadata — this matters when the runner re-registers the same - module (e.g. after a worker respawn during watch/server mode). + metadata — this matters when consecutive work units carry different + metadata for the same module. """ if not isinstance(hooks, list): return @@ -402,21 +403,19 @@ def _run_doctest( def _configure_logging_from_env() -> None: - """Opt-in worker logging via ``TRYKE_LOG``. + """Configure worker logging from the resolved ``TRYKE_LOG`` level. - Off by default so normal test runs don't emit anything on stderr. - The rust runner sets ``TRYKE_LOG=`` on the worker env when - ``-v`` (or ``TRYKE_LOG``) asks for cross-language verbosity, so - users typically don't set this directly. ``-q``/quiet does not - light up workers — workers stay silent unless the user explicitly - asked for more verbosity than the rust default ``warn``. + The Rust runner always sets ``TRYKE_LOG=`` on the worker + environment. A directly invoked worker remains off when the variable + is absent, and ``OFF`` explicitly disables logging. - Accepts ``DEBUG`` / ``INFO`` / ``WARN`` / ``ERROR`` / ``TRACE``. + Accepts ``OFF`` / ``ERROR`` / ``WARN`` / ``INFO`` / ``DEBUG`` / + ``TRACE``. Output goes to stderr so it never contaminates the JSON-RPC stream on stdout. """ level_name = os.environ.get("TRYKE_LOG", "").strip().upper() - if not level_name: + if not level_name or level_name == "OFF": return # Map TRACE to DEBUG since stdlib logging has no TRACE level. if level_name == "TRACE": @@ -435,7 +434,7 @@ def _configure_logging_from_env() -> None: def main() -> None: _configure_logging_from_env() - _log.debug("worker main: starting (pid=%d)", os.getpid()) + _log.debug("Worker main: starting (pid=%d)", os.getpid()) Worker(sys.stdin, sys.stdout).run()