From 0308efbf7a6019797a161e5fd2dfbdd79d016325 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 21 Oct 2026 01:52:39 +0800 Subject: [PATCH] Run synchronous evaluation callbacks on the managed blocking pool --- src/eval/blocking.rs | 68 +++++++++++++++++ src/eval/case_runner.rs | 107 +++++++++++++++++--------- src/eval/definition.rs | 42 ++++++---- src/eval/executor.rs | 33 +++++--- src/eval/mod.rs | 1 + src/eval/run.rs | 19 ++++- src/eval/scorer.rs | 47 +++++++++++- src/eval/task.rs | 33 +++++++- src/eval/types.rs | 14 ++++ tests/blocking_eval.rs | 165 ++++++++++++++++++++++++++++++++++++++++ tests/eager_eval.rs | 125 ++++++++++++++++++++++++++++++ 11 files changed, 586 insertions(+), 68 deletions(-) create mode 100644 src/eval/blocking.rs create mode 100644 tests/blocking_eval.rs diff --git a/src/eval/blocking.rs b/src/eval/blocking.rs new file mode 100644 index 0000000..e216996 --- /dev/null +++ b/src/eval/blocking.rs @@ -0,0 +1,68 @@ +//! Managed blocking invocation on the run's fixed execution binding. +use crate::internal::execution::ExecutionBinding; + +pub(crate) async fn invoke( + binding: Option<&ExecutionBinding>, + work: impl FnOnce() -> anyhow::Result + Send + 'static, +) -> anyhow::Result { + let binding = binding.ok_or_else(|| { + anyhow::anyhow!("blocking callbacks require a managed evaluation executor") + })?; + let task = binding.spawn_blocking(work); + // Abort queued work on run cancellation. Tokio cannot interrupt a closure + // already running; its owned data remains valid until that closure returns. + struct AbortOnDrop(tokio::task::AbortHandle); + impl Drop for AbortOnDrop { + fn drop(&mut self) { + self.0.abort(); + } + } + let _abort = AbortOnDrop(task.abort_handle()); + task.await? +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }; + + #[test] + fn dropping_an_invocation_cancels_queued_work_before_it_can_start() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .max_blocking_threads(1) + .build() + .unwrap(); + let binding = ExecutionBinding::external(runtime.handle().clone()); + let (release, waiting) = std::sync::mpsc::channel(); + let (entered, started) = std::sync::mpsc::channel(); + let occupied = binding.spawn_blocking(move || { + entered.send(()).unwrap(); + waiting + .recv_timeout(std::time::Duration::from_secs(3)) + .unwrap() + }); + started + .recv_timeout(std::time::Duration::from_secs(3)) + .unwrap(); + let called = Arc::new(AtomicBool::new(false)); + runtime.block_on(async { + { + let called = called.clone(); + let callback = invoke(Some(&binding), move || { + called.store(true, Ordering::SeqCst); + Ok(()) + }); + futures::pin_mut!(callback); + assert!(futures::poll!(callback).is_pending()); + } + release.send(()).unwrap(); + occupied.await.unwrap(); + binding.spawn_blocking(|| ()).await.unwrap(); + }); + assert!(!called.load(Ordering::SeqCst)); + } +} diff --git a/src/eval/case_runner.rs b/src/eval/case_runner.rs index ec39207..718368d 100644 --- a/src/eval/case_runner.rs +++ b/src/eval/case_runner.rs @@ -1,29 +1,40 @@ -//! One case's computation, independent of scheduling, recording and reporting. -use super::{Case, CaseError, EvalResult, Scorer, ScorerArgs, Task, TaskHooks}; -use crate::SpanHandle; +//! One case's computation, independent of recording and reporting. +use super::{Case, CaseError, EvalDefinition, EvalResult, ScorerArgs, TaskHooks}; +use crate::{internal::execution::ExecutionBinding, SpanHandle}; use futures::{stream::FuturesOrdered, StreamExt}; +use std::sync::Arc; pub(crate) struct CaseRunner; impl CaseRunner { - pub async fn run( + pub async fn run( case: Case, - task: &dyn Task, - scorers: &[Box>], + definition: &EvalDefinition, span: Option<&SpanHandle>, - ) -> EvalResult - where - I: Send + Sync, - O: Send + Sync, - { + binding: Option<&ExecutionBinding>, + ) -> EvalResult { + let case = Arc::new(case); let hooks = TaskHooks::for_case( span, case.expected.as_ref(), case.metadata.clone(), case.tags.clone(), ); + let task = definition.task(); + let computed = if task.is_blocking() { + let case = case.clone(); + let state = hooks.rebind(None, None); + let span = span.cloned(); + super::blocking::invoke(binding, move || { + let hooks = state.rebind(span.as_ref(), case.expected.as_ref()); + task.run_blocking(&case.input, &hooks) + }) + .await + } else { + task.run(&case.input, &hooks).await + }; let mut errors = Vec::new(); - let output = match task.run(&case.input, &hooks).await { - Ok(output) => Some(output), + let output = match computed { + Ok(output) => Some(Arc::new(output)), Err(source) => { errors.push(CaseError::Task { source }); None @@ -31,24 +42,34 @@ impl CaseRunner { }; let mut scores = Vec::new(); if let Some(output) = &output { - // Drive scorers concurrently while keeping definition order in the - // outcome, regardless of their relative completion times. let mut pending = FuturesOrdered::new(); - for scorer in scorers { - let args = ScorerArgs { - input: &case.input, - output, - expected: case.expected.as_ref(), - metadata: case.metadata.as_ref(), - }; + for scorer in definition.scorers() { + let case = case.clone(); + let output = output.clone(); + let scorer = scorer.clone(); pending.push_back(async move { - scorer - .score(args) - .await - .map_err(|source| CaseError::Scorer { - name: scorer.name().to_owned(), - source, + let name = scorer.name().to_owned(); + let result = if scorer.is_blocking() { + super::blocking::invoke(binding, move || { + scorer.score_blocking(ScorerArgs { + input: &case.input, + output: &output, + expected: case.expected.as_ref(), + metadata: case.metadata.as_ref(), + }) }) + .await + } else { + scorer + .score(ScorerArgs { + input: &case.input, + output: &output, + expected: case.expected.as_ref(), + metadata: case.metadata.as_ref(), + }) + .await + }; + result.map_err(|source| CaseError::Scorer { name, source }) }); } while let Some(result) = pending.next().await { @@ -61,6 +82,13 @@ impl CaseRunner { let metadata = hooks.get_metadata(); let tags = hooks.get_tags(); drop(hooks); + // All owned invocations have joined. Borrowed callback arguments cannot + // escape, so no input/output clones are required to return ownership. + let case = + Arc::try_unwrap(case).unwrap_or_else(|_| unreachable!("case invocations joined")); + let output = output.map(|output| { + Arc::try_unwrap(output).unwrap_or_else(|_| unreachable!("scorers joined")) + }); EvalResult { input: case.input, output, @@ -80,7 +108,7 @@ impl CaseRunner { #[cfg(test)] mod tests { use super::*; - use crate::eval::{Score, Scores, SimpleFnTask}; + use crate::eval::{Score, Scorer, Scores, Task}; use async_trait::async_trait; use serde_json::json; use std::sync::{ @@ -123,7 +151,12 @@ mod tests { xact_id: Some("version".into()), created: Some("timestamp".into()), }; - let result = futures::executor::block_on(CaseRunner::run(case, &LocalTask, &[], None)); + let result = futures::executor::block_on(CaseRunner::run( + case, + &EvalDefinition::new(Box::new(LocalTask), vec![]).unwrap(), + None, + None, + )); assert_eq!(result.output.unwrap().0, 6); assert_eq!(result.case_id.as_deref(), Some("case")); assert_eq!(result.xact_id.as_deref(), Some("version")); @@ -174,8 +207,8 @@ mod tests { input: 1, ..Default::default() }, - &FailTask, - &scorers, + &EvalDefinition::new(Box::new(FailTask), scorers).unwrap(), + None, None, ) .await; @@ -233,8 +266,14 @@ mod tests { input: 1, ..Default::default() }, - &SimpleFnTask::new(|input: &usize| Ok(*input)), - &scorers, + &EvalDefinition::new( + Box::new(crate::eval::task(|input: &usize, _hooks| { + Box::pin(async move { Ok(*input) }) + })), + scorers, + ) + .unwrap(), + None, None, ), ) diff --git a/src/eval/definition.rs b/src/eval/definition.rs index 11adddd..f4c48c8 100644 --- a/src/eval/definition.rs +++ b/src/eval/definition.rs @@ -15,8 +15,8 @@ pub struct EvalDefinition { } struct DefinitionInner { - task: Box>, - scorers: Vec>>, + task: Arc>, + scorers: Vec>>, } impl Clone for EvalDefinition { @@ -37,7 +37,7 @@ impl EvalDefinition { scorers: Vec>>, ) -> Result { let mut names = std::collections::HashSet::new(); - let mut validated: Vec>> = Vec::with_capacity(scorers.len()); + let mut validated: Vec>> = Vec::with_capacity(scorers.len()); for (index, scorer) in scorers.into_iter().enumerate() { let name = scorer.name().to_owned(); if name.trim().is_empty() { @@ -46,11 +46,11 @@ impl EvalDefinition { if !names.insert(name.clone()) { return Err(EvalConfigError::DuplicateScorerName { name }); } - validated.push(Box::new(NamedScorer { name, scorer })); + validated.push(Arc::new(NamedScorer { name, scorer })); } Ok(Self { inner: Arc::new(DefinitionInner { - task, + task: Arc::from(task), scorers: validated, }), }) @@ -63,11 +63,11 @@ impl EvalDefinition { self.inner.scorers.iter().map(|scorer| scorer.name()) } - pub(crate) fn task(&self) -> &dyn Task { - self.inner.task.as_ref() + pub(crate) fn task(&self) -> Arc> { + self.inner.task.clone() } - pub(crate) fn scorers(&self) -> &[Box>] { + pub(crate) fn scorers(&self) -> &[Arc>] { &self.inner.scorers } } @@ -79,10 +79,18 @@ struct NamedScorer { #[async_trait] impl Scorer for NamedScorer { + fn is_blocking(&self) -> bool { + self.scorer.is_blocking() + } + fn name(&self) -> &str { &self.name } + fn score_blocking(&self, args: ScorerArgs<'_, I, O>) -> anyhow::Result { + self.scorer.score_blocking(args) + } + async fn score(&self, args: ScorerArgs<'_, I, O>) -> anyhow::Result { self.scorer.score(args).await } @@ -91,7 +99,7 @@ impl Scorer for NamedScorer { #[cfg(test)] mod tests { use super::*; - use crate::eval::{case_runner::CaseRunner, task, Case, CaseError, SimpleFnTask}; + use crate::eval::{case_runner::CaseRunner, task, Case, CaseError}; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; // No Clone, Serialize, equality or runtime is needed for computation. @@ -133,8 +141,8 @@ mod tests { }; let (one, two) = futures::executor::block_on(async { futures::join!( - CaseRunner::run(case(1, 10), first.task(), first.scorers(), None), - CaseRunner::run(case(2, 20), other.task(), other.scorers(), None), + CaseRunner::run(case(1, 10), &first, None, None), + CaseRunner::run(case(2, 20), &other, None, None), ) }); assert_eq!(calls.load(Ordering::SeqCst), 2); @@ -145,8 +153,8 @@ mod tests { // Reuse the same callbacks after both prior invocations have completed. let (one, two) = futures::executor::block_on(async { futures::join!( - CaseRunner::run(case(3, 30), first.task(), first.scorers(), None), - CaseRunner::run(case(4, 40), other.task(), other.scorers(), None), + CaseRunner::run(case(3, 30), &first, None, None), + CaseRunner::run(case(4, 40), &other, None, None), ) }); assert_eq!(one.output.unwrap().0, 33); @@ -173,7 +181,9 @@ mod tests { fn scorer_identity_is_captured_once_for_every_run() { let changed = Arc::new(AtomicBool::new(false)); let definition = EvalDefinition::new( - Box::new(SimpleFnTask::new(|input: &usize| Ok(*input))), + Box::new(task(|input: &usize, _hooks| { + Box::pin(async move { Ok(*input) }) + })), vec![Box::new(ChangingName(changed.clone()))], ) .unwrap(); @@ -185,8 +195,8 @@ mod tests { input: 1, ..Default::default() }, - definition.task(), - definition.scorers(), + &definition, + None, None, )); assert!( diff --git a/src/eval/executor.rs b/src/eval/executor.rs index 372939c..98e4b3a 100644 --- a/src/eval/executor.rs +++ b/src/eval/executor.rs @@ -55,10 +55,11 @@ struct CaseOutcome { pub(crate) struct EvalExecutor; impl EvalExecutor { - pub async fn run( + pub async fn run( mut plan: EvalPlan, recorder: Option<&dyn EvalRecorder>, reporter: &mut impl EvalReporter, + binding: Option<&crate::internal::execution::ExecutionBinding>, ) -> std::result::Result, EvalRunError> { reporter.started(&plan.name); let dataset = plan.source.provenance(); @@ -102,7 +103,7 @@ impl EvalExecutor { pending.push(async move { ( index, - Self::run_case(case, definition, recorder, index).await, + Self::run_case(case, definition, recorder, index, binding).await, ) }); if pending.len() >= plan.concurrency { @@ -153,11 +154,12 @@ impl EvalExecutor { } } - async fn run_case( + async fn run_case( case: Case, definition: &EvalDefinition, recorder: Option<&dyn EvalRecorder>, index: usize, + binding: Option<&crate::internal::execution::ExecutionBinding>, ) -> CaseOutcome { let mut failures = Vec::new(); let recording = if let Some(recorder) = recorder { @@ -172,7 +174,7 @@ impl EvalExecutor { None }; let span = recording.as_ref().and_then(|recording| recording.span()); - let mut result = CaseRunner::run(case, definition.task(), definition.scorers(), span).await; + let mut result = CaseRunner::run(case, definition, span, binding).await; if let Some(recording) = recording { failures.extend(recording.finish(&mut result).await); } @@ -203,7 +205,7 @@ impl EvalExecutor { #[cfg(test)] mod tests { use super::*; - use crate::eval::{task, CaseSource, SimpleFnTask, SourceError}; + use crate::eval::{task, CaseSource, SourceError}; use std::{ cell::Cell, sync::{ @@ -290,7 +292,7 @@ mod tests { ) .unwrap(); let mut reporter = Observed::default(); - let error = futures::executor::block_on(EvalExecutor::run(plan, None, &mut reporter)) + let error = futures::executor::block_on(EvalExecutor::run(plan, None, &mut reporter, None)) .err() .expect("source should fail"); assert!( @@ -356,8 +358,9 @@ mod tests { ) .unwrap(); let mut reporter = Observed::default(); - let summary = futures::executor::block_on(EvalExecutor::run(plan, None, &mut reporter)) - .expect("local computation should succeed"); + let summary = + futures::executor::block_on(EvalExecutor::run(plan, None, &mut reporter, None)) + .expect("local computation should succeed"); assert_eq!(summary.total_cases, count); assert_eq!(summary.successful_cases, count); assert_eq!( @@ -422,7 +425,9 @@ mod tests { .source(cases.into()) .definition( EvalDefinition::new( - Box::new(SimpleFnTask::new(|input: &usize| Ok(input + 10))), + Box::new(task(|input: &usize, _hooks| { + Box::pin(async move { Ok(input + 10) }) + })), vec![], ) .unwrap(), @@ -432,9 +437,13 @@ mod tests { ) .unwrap(); let mut reporter = Observed::default(); - let error = - futures::executor::block_on(EvalExecutor::run(plan, Some(&recorder), &mut reporter)) - .unwrap_err(); + let error = futures::executor::block_on(EvalExecutor::run( + plan, + Some(&recorder), + &mut reporter, + None, + )) + .unwrap_err(); assert_eq!(error.failures.len(), 3); assert!(matches!( &error.failures[0], diff --git a/src/eval/mod.rs b/src/eval/mod.rs index 1da6849..2b293c5 100644 --- a/src/eval/mod.rs +++ b/src/eval/mod.rs @@ -49,6 +49,7 @@ //! # } //! ``` +mod blocking; mod bt_runner; mod case_runner; mod definition; diff --git a/src/eval/run.rs b/src/eval/run.rs index 4625872..cb784ed 100644 --- a/src/eval/run.rs +++ b/src/eval/run.rs @@ -82,11 +82,18 @@ pub fn run( return Completion::rejected(EvalRunError::new(EvalRunFailure::Runtime(source))) } }; + let callbacks = binding.clone(); binding.response(async move { let mut reporter = TerminalReporter { quiet: options.quiet, }; - EvalExecutor::run(EvalPlan::new(options)?, None, &mut reporter).await + EvalExecutor::run( + EvalPlan::new(options)?, + None, + &mut reporter, + Some(&callbacks), + ) + .await }) } Target::Recorded { client, submit } => { @@ -103,6 +110,7 @@ fn recorded, ) -> RunCompletion { let binding = client.execution_binding(); + let callbacks = binding.clone(); binding.response_try(|| { let operation = client.span_operation()?; Ok(async move { @@ -117,8 +125,13 @@ fn recorded: Send + Sync { /// Name of this scorer (used in results) fn name(&self) -> &str; + /// Whether evaluation runs should invoke this callback on their blocking pool. + /// Override together with `score_blocking` for synchronous work; async scorers use the defaults. + fn is_blocking(&self) -> bool { + false + } + + /// Synchronous callback body used when `is_blocking()` returns true. + /// The evaluation runner calls this on its selected blocking pool. + fn score_blocking(&self, _args: ScorerArgs<'_, I, O>) -> Result { + anyhow::bail!("scorer has no synchronous callback") + } + /// Score the output async fn score(&self, args: ScorerArgs<'_, I, O>) -> Result; } @@ -48,7 +60,8 @@ where /// Wrapper for sync scoring functions /// -/// This allows sync functions to be used as scorers in both async and sync evaluators. +/// Evaluation runs dispatch this callback on their selected runtime's blocking pool. +/// Direct calls to `Scorer::score` execute the callback on the calling thread. pub struct SyncFnScorer where F: Fn(ScorerArgs<'_, I, O>) -> Result + Send + Sync, @@ -78,11 +91,19 @@ where I: Send + Sync, O: Send + Sync, { + fn is_blocking(&self) -> bool { + true + } + fn name(&self) -> &str { &self.name } async fn score(&self, args: ScorerArgs<'_, I, O>) -> Result { + self.score_blocking(args) + } + + fn score_blocking(&self, args: ScorerArgs<'_, I, O>) -> Result { (self.func)(args) } } @@ -116,11 +137,19 @@ where O: PartialEq + Send + Sync, I: Send + Sync, { + fn is_blocking(&self) -> bool { + true + } + fn name(&self) -> &str { &self.name } async fn score(&self, args: ScorerArgs<'_, I, O>) -> Result { + self.score_blocking(args) + } + + fn score_blocking(&self, args: ScorerArgs<'_, I, O>) -> Result { let score = match args.expected { Some(expected) if expected == args.output => 1.0, Some(_) => 0.0, @@ -205,11 +234,19 @@ impl Scorer for LevenshteinScorer where I: Send + Sync, { + fn is_blocking(&self) -> bool { + true + } + fn name(&self) -> &str { &self.name } async fn score(&self, args: ScorerArgs<'_, I, String>) -> Result { + self.score_blocking(args) + } + + fn score_blocking(&self, args: ScorerArgs<'_, I, String>) -> Result { let expected = match args.expected { Some(e) => e, None => return Ok(vec![Score::new(self.name.clone(), 0.0)]), @@ -274,11 +311,19 @@ impl Scorer for NumericDiffScorer where I: Send + Sync, { + fn is_blocking(&self) -> bool { + true + } + fn name(&self) -> &str { &self.name } async fn score(&self, args: ScorerArgs<'_, I, f64>) -> Result { + self.score_blocking(args) + } + + fn score_blocking(&self, args: ScorerArgs<'_, I, f64>) -> Result { let expected = match args.expected { Some(e) => e, None => return Ok(vec![Score::new(self.name.clone(), 0.0)]), diff --git a/src/eval/task.rs b/src/eval/task.rs index 18ebb90..c3ca4ee 100644 --- a/src/eval/task.rs +++ b/src/eval/task.rs @@ -12,6 +12,18 @@ use super::types::TaskHooks; /// Callback errors use `anyhow::Result` to preserve application error causes. #[async_trait] pub trait Task: Send + Sync { + /// Whether evaluation runs should invoke this callback on their blocking pool. + /// Override together with `run_blocking` for synchronous work; async tasks use the defaults. + fn is_blocking(&self) -> bool { + false + } + + /// Synchronous callback body used when `is_blocking()` returns true. + /// The evaluation runner calls this on its selected blocking pool. + fn run_blocking(&self, _input: &I, _hooks: &TaskHooks<'_, I, O>) -> Result { + anyhow::bail!("task has no synchronous callback") + } + /// Execute the task on the given input async fn run(&self, input: &I, hooks: &TaskHooks<'_, I, O>) -> Result; } @@ -19,7 +31,8 @@ pub trait Task: Send + Sync { /// Wrapper for sync task functions /// /// This allows sync functions to be used as tasks in both async and sync evaluators. -/// The function runs synchronously but is wrapped in an async trait implementation. +/// Evaluation runs dispatch this callback on their selected runtime's blocking pool. +/// Direct calls to `Task::run` execute the callback on the calling thread. pub struct SyncFnTask where F: Fn(&I, &TaskHooks<'_, I, O>) -> Result + Send + Sync, @@ -47,7 +60,15 @@ where I: Send + Sync, O: Send + Sync, { + fn is_blocking(&self) -> bool { + true + } + async fn run(&self, input: &I, hooks: &TaskHooks<'_, I, O>) -> Result { + self.run_blocking(input, hooks) + } + + fn run_blocking(&self, input: &I, hooks: &TaskHooks<'_, I, O>) -> Result { (self.func)(input, hooks) } } @@ -74,7 +95,7 @@ where } } -/// Simple wrapper for tasks that don't need hooks +/// Synchronous task without hooks, dispatched on the run's managed blocking pool. pub struct SimpleFnTask where F: Fn(&I) -> Result + Send + Sync, @@ -102,7 +123,15 @@ where I: Send + Sync, O: Send + Sync, { + fn is_blocking(&self) -> bool { + true + } + async fn run(&self, input: &I, _hooks: &TaskHooks<'_, I, O>) -> Result { + self.run_blocking(input, _hooks) + } + + fn run_blocking(&self, input: &I, _hooks: &TaskHooks<'_, I, O>) -> Result { (self.func)(input) } } diff --git a/src/eval/types.rs b/src/eval/types.rs index 8b34da0..49b427c 100644 --- a/src/eval/types.rs +++ b/src/eval/types.rs @@ -241,6 +241,20 @@ impl<'a, I, O> TaskHooks<'a, I, O> { } } + pub(crate) fn rebind<'b>( + &self, + span: Option<&'b SpanHandle>, + expected: Option<&'b O>, + ) -> TaskHooks<'b, I, O> { + TaskHooks { + span, + expected, + metadata: self.metadata.clone(), + tags: self.tags.clone(), + _phantom: std::marker::PhantomData, + } + } + /// Add metadata key-value pair pub fn add_metadata(&self, key: impl Into, value: Value) { if let Ok(mut metadata) = self.metadata.lock() { diff --git a/tests/blocking_eval.rs b/tests/blocking_eval.rs new file mode 100644 index 0000000..b745a81 --- /dev/null +++ b/tests/blocking_eval.rs @@ -0,0 +1,165 @@ +use braintrust_sdk_rust::eval::*; +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +struct Input(usize); +struct Output(usize); +fn options( + task: Box>, + scorers: Vec>>, +) -> EvalOpts { + EvalOpts::builder() + .name("blocking".into()) + .source( + vec![Case { + input: Input(3), + expected: Some(Output(6)), + metadata: None, + tags: None, + id: None, + xact_id: None, + created: None, + }] + .into(), + ) + .definition(EvalDefinition::new(task, scorers).unwrap()) + .quiet(true) + .runtime(tokio::runtime::Handle::current()) + .build() +} + +#[tokio::test(flavor = "current_thread")] +async fn synchronous_task_uses_selected_blocking_pool_and_preserves_shared_hooks_without_clones() { + let worker = std::thread::current().id(); + let selected = tokio::runtime::Handle::current().id(); + let entered = Arc::new(tokio::sync::Notify::new()); + let started = entered.clone(); + let (release, waiting) = std::sync::mpsc::channel(); + let waiting = Mutex::new(waiting); + let task = SyncFnTask::new(move |input: &Input, hooks: &TaskHooks<'_, Input, Output>| { + assert_ne!(std::thread::current().id(), worker); + assert_eq!(tokio::runtime::Handle::current().id(), selected); + hooks.add_metadata("blocking", serde_json::json!(true)); + hooks.add_tag("from-blocking"); + started.notify_one(); + waiting + .lock() + .unwrap() + .recv_timeout(Duration::from_secs(2))?; + Ok(Output(input.0 + hooks.expected.unwrap().0)) + }); + // An async case on the same single-thread executor must make progress while + // the synchronous task is blocked; it releases the callback through a channel. + let release = tokio::spawn(async move { + entered.notified().await; + release.send(()).unwrap(); + }); + let summary = run(EvalRequest::local(options(Box::new(task), vec![]))) + .await + .unwrap(); + release.await.unwrap(); + assert!(summary.results[0].errors.is_empty()); + assert_eq!(summary.results[0].output.as_ref().unwrap().0, 9); + assert_eq!( + summary.results[0].metadata.as_ref().unwrap()["blocking"], + true + ); + assert_eq!( + summary.results[0].tags.as_ref().unwrap(), + &["from-blocking"] + ); +} + +#[tokio::test(flavor = "current_thread")] +async fn synchronous_scorers_run_concurrently_and_keep_definition_order_without_clones() { + let worker = std::thread::current().id(); + let selected = tokio::runtime::Handle::current().id(); + let (entered, mut arrivals) = tokio::sync::mpsc::unbounded_channel(); + let mut releases = Vec::new(); + let mut scorers: Vec>> = Vec::new(); + for name in ["first", "second"] { + let entered = entered.clone(); + let (release, waiting) = std::sync::mpsc::channel(); + releases.push(release); + let waiting = Mutex::new(waiting); + scorers.push(Box::new(SyncFnScorer::new( + name, + move |args: ScorerArgs<'_, Input, Output>| { + assert_ne!(std::thread::current().id(), worker); + assert_eq!(tokio::runtime::Handle::current().id(), selected); + assert_eq!(args.input.0, 3); + assert_eq!(args.output.0, 6); + assert_eq!(args.expected.unwrap().0, 6); + entered.send(name).unwrap(); + waiting + .lock() + .unwrap() + .recv_timeout(Duration::from_secs(2))?; + Ok(vec![Score::new(name, 1.0)]) + }, + ))); + } + let task = SimpleFnTask::new(|input: &Input| Ok(Output(input.0 * 2))); + let token = run(EvalRequest::local(options(Box::new(task), scorers))); + let mut names = vec![ + arrivals.recv().await.unwrap(), + arrivals.recv().await.unwrap(), + ]; + names.sort(); + assert_eq!(names, ["first", "second"]); + for release in releases.into_iter().rev() { + release.send(()).unwrap(); + } + let summary = token.await.unwrap(); + assert!(summary.results[0].errors.is_empty()); + assert_eq!( + summary.results[0] + .scores + .iter() + .map(|score| score.name.as_str()) + .collect::>(), + ["first", "second"] + ); +} + +#[tokio::test] +async fn synchronous_panics_become_case_errors_and_later_cases_continue() { + let opts = EvalOpts::builder() + .name("panic".into()) + .source( + (0..2) + .map(|input| Case { + input, + ..Default::default() + }) + .collect::>>() + .into(), + ) + .definition( + EvalDefinition::new( + Box::new(SimpleFnTask::new(|input: &usize| { + if *input == 0 { + panic!("synchronous application panic"); + } + Ok(*input) + })), + vec![], + ) + .unwrap(), + ) + .quiet(true) + .build(); + let summary = run(EvalRequest::local(opts)).await.unwrap(); + assert_eq!(summary.failed_cases, 1); + assert_eq!(summary.successful_cases, 1); + let [CaseError::Task { source }] = summary.results[0].errors.as_slice() else { + panic!("expected task error") + }; + assert!(source + .downcast_ref::() + .unwrap() + .is_panic()); + assert_eq!(summary.results[1].output, Some(1)); +} diff --git a/tests/eager_eval.rs b/tests/eager_eval.rs index 64853fd..f7a956b 100644 --- a/tests/eager_eval.rs +++ b/tests/eager_eval.rs @@ -391,3 +391,128 @@ async fn shutdown_deadline_cancels_pending_run_and_revokes_escaped_permission() Err(WriteError::Admission(EnqueueError::Closed)) )); } + +#[tokio::test] +async fn running_blocking_callback_keeps_owned_data_after_timeout_but_loses_write_permission() { + #[derive(serde::Serialize)] + struct OwnedInput { + value: usize, + #[serde(skip)] + drops: Arc, + #[serde(skip)] + dropped: Arc, + } + impl Drop for OwnedInput { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::SeqCst); + self.dropped.notify_one(); + } + } + let server = MockServer::start().await; + let client = client(&server).await; + let entered = Arc::new(tokio::sync::Notify::new()); + let drops = Arc::new(AtomicUsize::new(0)); + let dropped = Arc::new(tokio::sync::Notify::new()); + let (release, waiting) = std::sync::mpsc::channel(); + let waiting = Mutex::new(waiting); + let (report, observed) = tokio::sync::oneshot::channel(); + let report = Mutex::new(Some(report)); + let callback = SyncFnTask::new({ + let entered = entered.clone(); + move |input: &OwnedInput, hooks: &TaskHooks<'_, OwnedInput, usize>| { + entered.notify_one(); + waiting + .lock() + .unwrap() + .recv_timeout(Duration::from_secs(3))?; + let rejected = hooks.span.unwrap().log(SpanLog::default()).is_err(); + report + .lock() + .unwrap() + .take() + .unwrap() + .send((input.value, rejected)) + .unwrap(); + Ok(input.value) + } + }); + let opts = options( + EvalDefinition::new(Box::new(callback), vec![]).unwrap(), + vec![Case { + input: OwnedInput { + value: 17, + drops: drops.clone(), + dropped: dropped.clone(), + }, + expected: None, + metadata: None, + tags: None, + id: None, + xact_id: None, + created: None, + }], + ); + let token = run(EvalRequest::recorded(opts, client.clone())); + entered.notified().await; + assert!(client.shutdown(Duration::from_millis(20)).await.is_err()); + assert!(matches!( + token + .await + .err() + .expect("run must time out") + .failures + .as_slice(), + [EvalRunFailure::Execution(ExecutionError::TimedOut)] + )); + assert_eq!( + drops.load(Ordering::SeqCst), + 0, + "running callback still owns the input" + ); + release.send(()).unwrap(); + assert_eq!(observed.await.unwrap(), (17, true)); + dropped.notified().await; + assert_eq!(drops.load(Ordering::SeqCst), 1); +} + +#[tokio::test(flavor = "current_thread")] +async fn synchronous_callback_can_wait_for_recording_without_blocking_delivery() { + let server = MockServer::start().await; + let client = client(&server).await; + let worker = std::thread::current().id(); + let callback = SyncFnTask::new(move |input: &usize, hooks: &TaskHooks<'_, usize, usize>| { + assert_ne!(std::thread::current().id(), worker); + futures::executor::block_on( + hooks + .span + .unwrap() + .log( + SpanLog::builder() + .metadata( + serde_json::json!({"sync_write": true}) + .as_object() + .unwrap() + .clone(), + ) + .build() + .unwrap(), + ) + .into_future(), + )?; + Ok(*input) + }); + let token = run(EvalRequest::recorded( + options( + EvalDefinition::new(Box::new(callback), vec![]).unwrap(), + vec![Case { + input: 9, + ..Default::default() + }], + ), + client.clone(), + )); + let shutdown = client.shutdown(Duration::from_secs(5)); + let (result, shutdown) = tokio::join!(token, shutdown); + assert_eq!(result.unwrap().results[0].output, Some(9)); + shutdown.unwrap(); +}