diff --git a/bt-daemon/Cargo.lock b/bt-daemon/Cargo.lock index 05fbde0f..fe6c4914 100644 --- a/bt-daemon/Cargo.lock +++ b/bt-daemon/Cargo.lock @@ -239,7 +239,7 @@ dependencies = [ [[package]] name = "braintrust-sdk-rust" version = "0.1.0-alpha.2" -source = "git+https://github.com/braintrustdata/braintrust-sdk-rust?rev=e85ab580f52174efa9f90b9069a5fb2f5460a20d#e85ab580f52174efa9f90b9069a5fb2f5460a20d" +source = "git+https://github.com/braintrustdata/braintrust-sdk-rust?rev=f2216fa0ceee132e02f1392c0e5af89ae762e9c8#f2216fa0ceee132e02f1392c0e5af89ae762e9c8" dependencies = [ "anyhow", "arc-swap", diff --git a/bt-daemon/Cargo.toml b/bt-daemon/Cargo.toml index 6d8ed9a5..b058b2bf 100644 --- a/bt-daemon/Cargo.toml +++ b/bt-daemon/Cargo.toml @@ -22,7 +22,7 @@ expect_used = "deny" [dependencies] # Keep this pinned exactly, matching the dependency policy used by `bt`. -braintrust-sdk-rust = { git = "https://github.com/braintrustdata/braintrust-sdk-rust", rev = "e85ab580f52174efa9f90b9069a5fb2f5460a20d" } +braintrust-sdk-rust = { git = "https://github.com/braintrustdata/braintrust-sdk-rust", rev = "f2216fa0ceee132e02f1392c0e5af89ae762e9c8" } base64 = "0.22" anyhow = "1" parking_lot = "0.12" diff --git a/bt-daemon/src/dispatch.rs b/bt-daemon/src/dispatch.rs index 61fc8a5a..55896cd9 100644 --- a/bt-daemon/src/dispatch.rs +++ b/bt-daemon/src/dispatch.rs @@ -9,6 +9,8 @@ use crate::delivery_ledger::LedgerSink; use crate::journal::JournalWriter; +use crate::pipeline::plugins::{PluginOutcome, Plugins}; +use crate::pipeline::{Completion, Delivery, Stage}; use crate::sink::SinkFactory; use crate::translate::{Registry, SessionCtx}; use crate::wire::{Envelope, SessionRoute}; @@ -629,12 +631,13 @@ impl SessionActor { .checkpoint_and_flush(&mut translator, &mut sink, &ctx) .await; if checkpointable { - self.record_delivery_checkpoint( - &mut sink, - pending_through, - &mut acknowledged_through, - ) - .await; + checkpointable &= self + .record_delivery_checkpoint( + &mut sink, + pending_through, + &mut acknowledged_through, + ) + .await; } let _ = reply.send(self.counters.queued.load(Ordering::Relaxed)); } @@ -648,12 +651,13 @@ impl SessionActor { ) .await; if checkpointable { - self.record_delivery_checkpoint( - &mut sink, - pending_through, - &mut acknowledged_through, - ) - .await; + checkpointable &= self + .record_delivery_checkpoint( + &mut sink, + pending_through, + &mut acknowledged_through, + ) + .await; } let _ = reply.send(self.counters.queued.load(Ordering::Relaxed)); } @@ -725,58 +729,38 @@ impl SessionActor { .as_ref() .map(|config| config.span_plugins.as_slice()) .unwrap_or_default(); + let mut stage = Plugins { + paths: plugin_paths.to_vec(), + source: self.source.clone(), + session_id: self.session_id.clone(), + }; let mut processed = Vec::with_capacity(ops.len()); for op in &ops { - match crate::span_processor::process( - plugin_paths, - op, - &self.source, - &self.session_id, - ) { - Ok(result) => { - if let Some(failure) = result.failure { - delivered = false; - if failure.newly_seen { - if let Err(error) = crate::plugin_diagnostics::record( - &self.data_dir, - &self.source, - &failure.path, - &failure.message, - ) { - tracing::warn!( - session_id = %self.session_id, - "failed to persist span plugin diagnostic: {error}" - ); - } - self.set_error(format!( - "span plugin {} failed; span operations are being discarded: {}", - failure.path.display(), - failure.message - )); - } - } - if let Some(op) = result.op { - processed.push(op); - } - } - Err(error) => { + match stage.process(op).await { + PluginOutcome::Prepared(op) => processed.push(*op), + PluginOutcome::Filtered => {} + PluginOutcome::Failed(failure) => { delivered = false; - if let Some(plugin) = plugin_paths.first() { - if let Err(diagnostic_error) = crate::plugin_diagnostics::record( + if failure.newly_seen { + if let Err(error) = crate::plugin_diagnostics::record( &self.data_dir, &self.source, - plugin, - &error.to_string(), + &failure.path, + &failure.message, ) { - tracing::warn!( - session_id = %self.session_id, - "failed to persist span plugin diagnostic: {diagnostic_error}" - ); + tracing::warn!(session_id = %self.session_id, + "failed to persist span plugin diagnostic: {error}"); } + self.set_error(format!( + "span plugin {} failed; span operations are being discarded: {}", + failure.path.display(), + failure.message + )); } - self.set_error(format!( - "span plugin processor failed; span operation discarded: {error}" - )); + } + PluginOutcome::Unavailable(error) => { + delivered = false; + self.set_error(format!("span plugin processor failed: {error}")); } } } @@ -953,38 +937,36 @@ impl SessionActor { } } - async fn flush_sink(&self, sink: &mut Box) -> bool { - if let Err(e) = sink.flush().await { - self.set_error(format!("sink flush failed: {e}")); - return false; - } - self.refresh_permalink(sink.as_ref()); - !sink.has_pending_delivery() - } - + /// False means delivery was rejected and this actor must never advance its + /// contiguous event boundary. Deferred delivery and failed checkpoint writes + /// leave accepted work eligible for a later checkpoint. async fn record_delivery_checkpoint( &self, sink: &mut Box, through: u64, acknowledged_through: &mut u64, - ) { - if !self.flush_sink(sink).await { - return; + ) -> bool { + if let Err(error) = sink.flush().await { + self.set_error(format!("sink flush failed: {error}")); + return false; } - if through <= *acknowledged_through { - return; + self.refresh_permalink(sink.as_ref()); + if sink.has_pending_delivery() || through <= *acknowledged_through { + return true; } - if let Err(error) = self - .journal - .lock() - .await - .append_delivery_checkpoint(&self.route, through) - .await - { - self.set_error(format!("delivery checkpoint failed: {error}")); - } else { - *acknowledged_through = through; + let pending = Delivery::new( + through, + EventCompletion { + journal: self.journal.clone(), + route: self.route.clone(), + through, + }, + ); + match pending.into_completed().await { + Ok(through) => *acknowledged_through = through, + Err(failed) => self.set_error(format!("delivery checkpoint failed: {}", failed.error)), } + true } fn refresh_permalink(&self, sink: &dyn crate::sink::Sink) { @@ -1015,6 +997,23 @@ pub(crate) fn is_tool_lifecycle_event(event: &str) -> bool { ) } +struct EventCompletion { + journal: Arc>, + route: SessionRoute, + through: u64, +} + +#[async_trait::async_trait] +impl Completion for EventCompletion { + async fn complete(&mut self) -> anyhow::Result<()> { + self.journal + .lock() + .await + .append_delivery_checkpoint(&self.route, self.through) + .await + } +} + #[cfg(test)] mod tests { use super::*; @@ -1151,16 +1150,8 @@ mod tests { .await .unwrap(); let updates = session.join("updates.jsonl"); - tokio::fs::write(&updates, b"{\"update\":1}\n") - .await - .unwrap(); - let writer_path = updates.clone(); - let writer = tokio::spawn(async move { - tokio::time::sleep(std::time::Duration::from_millis(25)).await; - tokio::fs::write(&writer_path, b"{\"update\":1}\n{\"update\":2}\n") - .await - .unwrap(); - }); + let initial = b"{\"update\":1}\n"; + tokio::fs::write(&updates, initial).await.unwrap(); let mut env = Envelope { source: "grok".into(), source_version: None, @@ -1177,8 +1168,31 @@ mod tests { config: None, }; - hydrate_transcript_reference(tmp.path(), &mut env).await; - writer.await.unwrap(); + let mirror = crate::transcript_mirror::mirror_path( + tmp.path(), + &crate::ids::session_namespace("grok", "session-1"), + updates.to_str().unwrap(), + ); + { + let hydration = hydrate_transcript_reference(tmp.path(), &mut env); + tokio::pin!(hydration); + tokio::select! { + biased; + () = async { + // Append after the first copy, before polling hydration again. + // A delayed task can miss the entire capture window on busy CI. + while !std::fs::metadata(&mirror).is_ok_and(|meta| meta.len() == initial.len() as u64) { + tokio::task::yield_now().await; + } + std::fs::write(&updates, b"{\"update\":1}\n{\"update\":2}\n").unwrap(); + } => {}, + () = &mut hydration => panic!("terminal hydration did not wait for another capture"), + _ = tokio::time::sleep(std::time::Duration::from_secs(5)) => { + panic!("initial transcript capture did not complete"); + }, + } + hydration.await; + } let update_mirror = &env.payload["_bt_grok_transcript_mirrors"]["updates"]; assert_eq!( diff --git a/bt-daemon/src/lib.rs b/bt-daemon/src/lib.rs index 2df4e679..61a3f58f 100644 --- a/bt-daemon/src/lib.rs +++ b/bt-daemon/src/lib.rs @@ -24,6 +24,7 @@ mod hook; mod ids; mod journal; mod managed_run; +mod pipeline; mod plugin_diagnostics; pub(crate) mod process; mod route; diff --git a/bt-daemon/src/pipeline/delivery.rs b/bt-daemon/src/pipeline/delivery.rs new file mode 100644 index 00000000..d6cd35cb --- /dev/null +++ b/bt-daemon/src/pipeline/delivery.rs @@ -0,0 +1,122 @@ +/// A message owns its completion obligation until terminal handling succeeds. +pub(crate) struct Delivery { + value: T, + completion: Box, +} + +// async_trait adds a redundant must_use annotation to its boxed futures. +#[allow(clippy::double_must_use)] +#[async_trait::async_trait] +pub(crate) trait Completion: Send { + async fn complete(&mut self) -> anyhow::Result<()>; +} + +pub(crate) struct Unfinished { + #[cfg_attr(not(test), allow(dead_code))] + pub delivery: Delivery, + pub error: anyhow::Error, +} + +impl Delivery { + pub fn new(value: T, completion: impl Completion + 'static) -> Self { + Self { + value, + completion: Box::new(completion), + } + } + + #[cfg_attr(not(test), allow(dead_code))] + pub fn value(&self) -> &T { + &self.value + } + + #[cfg_attr(not(test), allow(dead_code))] + pub fn map(self, map: impl FnOnce(T) -> U) -> Delivery { + Delivery { + value: map(self.value), + completion: self.completion, + } + } + + pub async fn into_completed(mut self) -> Result> { + match self.completion.complete().await { + Ok(()) => Ok(self.value), + Err(error) => Err(Unfinished { + delivery: self, + error, + }), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }; + use tokio::sync::mpsc; + + struct Count(Arc); + #[async_trait::async_trait] + impl Completion for Count { + async fn complete(&mut self) -> anyhow::Result<()> { + self.0.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + } + + #[tokio::test] + async fn mapping_and_channel_forwarding_preserve_completion() { + let count = Arc::new(AtomicUsize::new(0)); + let (tx, mut rx) = mpsc::channel(1); + tx.send(Delivery::new(7, Count(count.clone())).map(|n| n.to_string())) + .await + .unwrap_or_else(|_| panic!("channel closed")); + let message = rx.recv().await.unwrap(); + assert_eq!(message.value(), "7"); + assert_eq!(count.load(Ordering::SeqCst), 0); + assert_eq!(message.into_completed().await.ok().unwrap(), "7"); + assert_eq!(count.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn failed_send_returns_the_obligation_and_drop_does_not_complete() { + let count = Arc::new(AtomicUsize::new(0)); + let (tx, rx) = mpsc::channel(1); + drop(rx); + let pending = tx + .send(Delivery::new(7, Count(count.clone()))) + .await + .err() + .unwrap(); + drop(pending); + assert_eq!(count.load(Ordering::SeqCst), 0); + } + + struct Retry { + first: bool, + } + #[async_trait::async_trait] + impl Completion for Retry { + async fn complete(&mut self) -> anyhow::Result<()> { + if std::mem::take(&mut self.first) { + anyhow::bail!("head write failed"); + } + Ok(()) + } + } + + #[tokio::test] + async fn failed_completion_returns_the_value_and_can_retry() { + let failed = Delivery::new(7, Retry { first: true }) + .into_completed() + .await + .err() + .unwrap(); + assert_eq!(failed.error.to_string(), "head write failed"); + assert_eq!(*failed.delivery.value(), 7); + assert_eq!(failed.delivery.into_completed().await.ok().unwrap(), 7); + } +} diff --git a/bt-daemon/src/pipeline/mod.rs b/bt-daemon/src/pipeline/mod.rs new file mode 100644 index 00000000..6180279e --- /dev/null +++ b/bt-daemon/src/pipeline/mod.rs @@ -0,0 +1,16 @@ +//! Typed stage contracts and completion-preserving channel messages. + +pub(crate) mod delivery; +pub(crate) mod plugins; + +pub(crate) use delivery::{Completion, Delivery}; + +// async_trait adds a redundant must_use annotation to its boxed futures. +#[allow(clippy::double_must_use)] +#[async_trait::async_trait] +pub(crate) trait Stage: Send { + type Input: Sync; + type Outcome: Send; + + async fn process(&mut self, input: &Self::Input) -> Self::Outcome; +} diff --git a/bt-daemon/src/pipeline/plugins.rs b/bt-daemon/src/pipeline/plugins.rs new file mode 100644 index 00000000..ddce96b9 --- /dev/null +++ b/bt-daemon/src/pipeline/plugins.rs @@ -0,0 +1,34 @@ +use super::Stage; +use crate::{span_processor, translate::SpanOp}; +use std::path::PathBuf; + +/// The outcome says what happened; the caller chooses the next stage. +pub(crate) enum PluginOutcome { + Prepared(Box), + Filtered, + Failed(span_processor::PluginFailure), + Unavailable(anyhow::Error), +} + +pub(crate) struct Plugins { + pub paths: Vec, + pub source: String, + pub session_id: String, +} + +#[async_trait::async_trait] +impl Stage for Plugins { + type Input = SpanOp; + type Outcome = PluginOutcome; + + async fn process(&mut self, input: &SpanOp) -> PluginOutcome { + match span_processor::process(&self.paths, input, &self.source, &self.session_id) { + Ok(result) => match (result.op, result.failure) { + (_, Some(failure)) => PluginOutcome::Failed(failure), + (Some(op), None) => PluginOutcome::Prepared(Box::new(op)), + (None, None) => PluginOutcome::Filtered, + }, + Err(error) => PluginOutcome::Unavailable(error), + } + } +} diff --git a/bt-daemon/src/sink/braintrust.rs b/bt-daemon/src/sink/braintrust.rs index 7bca1248..100a1ecd 100644 --- a/bt-daemon/src/sink/braintrust.rs +++ b/bt-daemon/src/sink/braintrust.rs @@ -1,9 +1,8 @@ //! The Braintrust sink: maps sink-neutral [`SpanOp`]s onto `braintrust-sdk-rust`. //! -//! Multi-profile: a session's backend URLs come from its own config (bt -//! resolves them per profile), so clients are built lazily and cached by -//! `(api_url, app_url)` — sessions on the same instance share a client, -//! sessions on different instances get their own. Within a client, each +//! Each routing branch owns its checked SDK queue. Clients are built lazily +//! using its configured backend URLs; flush results cannot include another +//! branch's rejected operations. Within a client, each //! session's token/org travel per span (`span_builder_with_credentials`) and //! never leak across sessions. Span ids are the translator's deterministic //! UUIDv5 strings, reused as the SDK `row_id` merge key so journal replay @@ -57,20 +56,14 @@ impl CachedClient { if *queued_bytes < QUEUE_FLUSH_BYTE_BUDGET { return Ok(()); } - self.client - .flush() - .await - .map_err(|error| anyhow::anyhow!("braintrust flush failed: {error}"))?; + self.client.flush().await.map_err(anyhow::Error::new)?; *queued_bytes = 0; Ok(()) } async fn flush(&self) -> anyhow::Result<()> { let mut queued_bytes = self.queued_bytes.lock().await; - self.client - .flush() - .await - .map_err(|error| anyhow::anyhow!("braintrust flush failed: {error}"))?; + self.client.flush().await.map_err(anyhow::Error::new)?; *queued_bytes = 0; Ok(()) } @@ -122,6 +115,7 @@ impl ClientCache { } let client = BraintrustClient::builder() .skip_login(true) + .checked_delivery(true) .api_url(api_url.to_string()) .app_url(app_url.to_string()) .build() @@ -140,7 +134,6 @@ impl ClientCache { /// Hands out a per-session sink over the shared client pool. pub struct BraintrustSinkFactory { - cache: Arc, default_api_url: Option, default_app_url: Option, } @@ -148,7 +141,6 @@ pub struct BraintrustSinkFactory { impl BraintrustSinkFactory { pub fn new(cfg: BraintrustSinkConfig) -> Self { Self { - cache: Arc::new(ClientCache::new()), default_api_url: cfg.api_url, default_app_url: cfg.app_url, } @@ -158,7 +150,9 @@ impl BraintrustSinkFactory { impl SinkFactory for BraintrustSinkFactory { fn create(&self, _session_id: &str, source: &str) -> anyhow::Result> { Ok(Box::new(BraintrustSink { - cache: self.cache.clone(), + // Acceptance belongs to this routing branch. A shared SDK queue can + // attribute another branch's rejected batch to this branch's flush. + cache: Arc::new(ClientCache::new()), default_api_url: self.default_api_url.clone(), default_app_url: self.default_app_url.clone(), source: source.to_string(), @@ -732,6 +726,7 @@ mod tests { .await; let client = BraintrustClient::builder() .skip_login(true) + .checked_delivery(true) .api_url(server.uri()) .app_url(server.uri()) .environment("test", Some("sink-test")) diff --git a/bt-daemon/src/trace_runtime.rs b/bt-daemon/src/trace_runtime.rs index 2e1ad6f0..114f1614 100644 --- a/bt-daemon/src/trace_runtime.rs +++ b/bt-daemon/src/trace_runtime.rs @@ -17,7 +17,7 @@ use crate::{ }; use async_trait::async_trait; use std::ffi::OsString; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::Duration; @@ -489,18 +489,26 @@ fn daemon_warnings( } async fn doctor_output(host: &TraceHostContext, args: DoctorArgs) -> DoctorCommandOutput { - doctor_output_at(host, args, &paths::socket_path(None)).await + let source = args.agent.identity().id; + doctor_output_at( + host, + args, + &paths::socket_path(None), + paths::agent_settings_path(source, None), + crate::settings::AgentSettings::load(source), + ) + .await } async fn doctor_output_at( host: &TraceHostContext, args: DoctorArgs, socket: &Path, + settings_path: PathBuf, + settings: crate::settings::AgentSettings, ) -> DoctorCommandOutput { let source = args.agent.identity().id; - let settings_path = paths::agent_settings_path(source, None); let settings_present = settings_path.exists(); - let settings = crate::settings::AgentSettings::load(source); let enabled = settings.tracing_enabled(); let mut warnings = Vec::new(); @@ -1304,7 +1312,16 @@ mod tests { agent: DoctorAgent, ) -> DoctorCommandOutput { let host = test_host(Arc::new(shell)); - doctor_output_at(&host, DoctorArgs { agent }, &self.socket).await + // Use fixture settings rather than the developer's installed route + // or a managed run's invocation settings. + doctor_output_at( + &host, + DoctorArgs { agent }, + &self.socket, + self._dir.path().join("braintrust.json"), + crate::settings::AgentSettings::default(), + ) + .await } async fn forward_event(&self, source: &str, session_id: &str, route: SessionRoute) { @@ -1512,6 +1529,8 @@ mod tests { agent: DoctorAgent::Codex, }, &socket, + temp.path().join("braintrust.json"), + crate::settings::AgentSettings::default(), ) .await; @@ -1535,6 +1554,8 @@ mod tests { agent: DoctorAgent::Codex, }, &test_endpoint(temp.path()), + temp.path().join("braintrust.json"), + crate::settings::AgentSettings::default(), ) .await; assert_eq!(output.daemon.status, DaemonStatus::NotRunning); diff --git a/bt-daemon/tests/braintrust_sink.rs b/bt-daemon/tests/braintrust_sink.rs index 46b843a5..98c24903 100644 --- a/bt-daemon/tests/braintrust_sink.rs +++ b/bt-daemon/tests/braintrust_sink.rs @@ -815,3 +815,91 @@ async fn experiment_sessions_use_experiment_object_type_and_id() { "experiment spans were routed as project logs: {bodies}" ); } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn rejected_backend_batch_is_not_a_successful_flush() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/version")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(400).set_body_string("invalid row")) + .mount(&server) + .await; + let factory = BraintrustSinkFactory::new(BraintrustSinkConfig::default()); + let mut sink = factory.create("rejected", "debug").unwrap(); + let mut config = session_config(&server.uri()); + config.destination = Some(TraceDestination::ProjectLogs { + project_id: Some("project".into()), + project_name: None, + }); + sink.configure(&config); + sink.emit(&[SpanOp::Insert(row( + "rejected", + "rejected", + &[], + "rejected", + SpanType::Task, + 1, + Some(2), + ))]) + .await + .unwrap(); + let flushed = sink.flush().await; + assert!( + flushed.is_err(), + "a rejected SDK batch must hold completion" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn rejected_route_does_not_poison_a_sibling_flush_on_the_same_backend() { + use wiremock::matchers::header; + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/version")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({}))) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .and(header("authorization", "Bearer sk-bad")) + .respond_with(ResponseTemplate::new(400)) + .mount(&server) + .await; + Mock::given(method("POST")) + .and(path("/logs3")) + .and(header("authorization", "Bearer sk-good")) + .respond_with(ResponseTemplate::new(200)) + .mount(&server) + .await; + let factory = BraintrustSinkFactory::new(BraintrustSinkConfig::default()); + let mut bad = factory.create("bad", "debug").unwrap(); + let mut good = factory.create("good", "debug").unwrap(); + for (sink, token) in [(&mut bad, "sk-bad"), (&mut good, "sk-good")] { + let mut config = session_config(&server.uri()); + config.auth.token = token.into(); + config.destination = Some(TraceDestination::ProjectLogs { + project_id: Some("project".into()), + project_name: None, + }); + sink.configure(&config); + sink.emit(&[SpanOp::Insert(row( + token, + token, + &[], + token, + SpanType::Task, + 1, + Some(2), + ))]) + .await + .unwrap(); + } + let (bad_result, good_result) = tokio::join!(bad.flush(), good.flush()); + assert!(bad_result.is_err()); + good_result.unwrap(); +} diff --git a/bt-daemon/tests/pipeline.rs b/bt-daemon/tests/pipeline.rs index 2a285244..06f81e45 100644 --- a/bt-daemon/tests/pipeline.rs +++ b/bt-daemon/tests/pipeline.rs @@ -569,6 +569,98 @@ async fn shutdown(socket: &Path) { shutdown_daemon(socket).await.unwrap(); } +/// A failed flush consumes its batch, just like checked SDK delivery. An empty +/// subsequent flush must not acknowledge that lost batch's source events. +#[tokio::test] +async fn rejected_flush_blocks_checkpoints_until_journal_replay() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + #[derive(Default)] + struct RejectOnce { + flushes: AtomicUsize, + emitted: AtomicUsize, + } + struct RejectOnceSink(Arc); + struct RejectOnceFactory(Arc); + impl SinkFactory for RejectOnceFactory { + fn create(&self, _: &str, _: &str) -> anyhow::Result> { + Ok(Box::new(RejectOnceSink(self.0.clone()))) + } + } + #[async_trait] + impl Sink for RejectOnceSink { + async fn emit(&mut self, ops: &[SpanOp]) -> anyhow::Result { + self.0.emitted.fetch_add(ops.len(), Ordering::SeqCst); + Ok(ops.len() as u64) + } + async fn flush(&mut self) -> anyhow::Result<()> { + if self.0.flushes.fetch_add(1, Ordering::SeqCst) == 0 { + anyhow::bail!("backend rejected and consumed the batch"); + } + Ok(()) + } + } + + let tmp = tempfile::tempdir().unwrap(); + let data_dir = tmp.path().join("data"); + let socket = test_endpoint(tmp.path()); + let sink = Arc::new(RejectOnce::default()); + let start = || { + let args = ServeArgs { + socket: Some(socket.clone()), + data_dir: Some(data_dir.clone()), + idle_timeout_secs: 0, + session_idle_timeout_secs: 0, + }; + let opts = ServeOptions { + version: "test".into(), + translators: Arc::new(Registry::default_agents()), + sink_factory: Arc::new(RejectOnceFactory(sink.clone())), + auth_provider: Some(Arc::new(TestAuthProvider { + calls: Mutex::new(Vec::new()), + fail: false, + first_lease_expired: false, + })), + }; + tokio::spawn(async move { run_serve(args, opts).await.unwrap() }) + }; + let handle = start(); + wait_for(&socket).await; + let session = "rejected-flush"; + forward_envelope( + &envelope(session, "SessionStart", 1), + &socket, + &dummy_host(), + false, + ) + .await + .unwrap(); + flush_session(session, &socket, 5_000).await.unwrap(); + assert_eq!(sink.flushes.load(Ordering::SeqCst), 1); + forward_envelope(&envelope(session, "Stop", 2), &socket, &dummy_host(), false) + .await + .unwrap(); + flush_session(session, &socket, 5_000).await.unwrap(); + shutdown(&socket).await; + handle.await.unwrap(); + + let journal_path = source_journal_path(&data_dir, "debug", session); + let journal = std::fs::read_to_string(&journal_path).unwrap(); + assert!(!journal.contains("delivery_checkpoint"), "{journal}"); + assert_eq!(sink.flushes.load(Ordering::SeqCst), 1); + let emitted_before_restart = sink.emitted.load(Ordering::SeqCst); + + let handle = start(); + wait_for(&socket).await; + flush_session(session, &socket, 5_000).await.unwrap(); + assert!(sink.emitted.load(Ordering::SeqCst) > emitted_before_restart); + assert!(std::fs::read_to_string(journal_path) + .unwrap() + .contains("delivery_checkpoint")); + shutdown(&socket).await; + handle.await.unwrap(); +} + #[tokio::test] async fn routed_sessions_resolve_multiple_profiles_without_journaling_credentials() { let provider = Arc::new(TestAuthProvider {