From 48eb5b4fca4b972e20f4f5204bc7237afe5700a9 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 14 Oct 2026 03:05:02 +0800 Subject: [PATCH 1/5] Consolidate durable subscription lifecycle and retention --- bt-daemon/src/delivery_ledger.rs | 55 +-- bt-daemon/src/dispatch.rs | 23 +- bt-daemon/src/journal.rs | 49 ++- bt-daemon/src/pipeline/consumer.rs | 96 ++++- bt-daemon/src/pipeline/journal.rs | 155 +++++++- bt-daemon/src/pipeline/routing.rs | 11 - bt-daemon/src/pipeline/translation.rs | 47 ++- bt-daemon/src/server/auth.rs | 42 ++ bt-daemon/src/server/ingress.rs | 7 +- bt-daemon/src/server/maintenance.rs | 414 ++++++++++++++++++-- bt-daemon/src/server/mod.rs | 6 +- bt-daemon/src/server/recovery.rs | 44 +++ bt-daemon/src/server/sessions.rs | 3 +- bt-daemon/src/transcript_import/pipeline.rs | 2 +- bt-daemon/src/transcript_mirror.rs | 8 +- bt-daemon/tests/pipeline.rs | 78 ++++ 16 files changed, 931 insertions(+), 109 deletions(-) diff --git a/bt-daemon/src/delivery_ledger.rs b/bt-daemon/src/delivery_ledger.rs index 6b7a7ce8..84853439 100644 --- a/bt-daemon/src/delivery_ledger.rs +++ b/bt-daemon/src/delivery_ledger.rs @@ -263,8 +263,7 @@ impl std::error::Error for CommitFailure { } /// Wrap a sink with destination-scoped terminal-span suppression. If durable -/// ledger state cannot be loaded, callers retain normal delivery rather than -/// risking lost trace data. +/// ledger state cannot be loaded, retain input until its recovery condition is met. pub(crate) struct LedgerSink { inner: Box, ledger: Option, @@ -284,15 +283,10 @@ impl LedgerSink { session_id: &str, config: Option<&SessionConfig>, bt_version: &str, - ) -> Self { + ) -> anyhow::Result { let loaded = match (data_dir, config) { (Some(data_dir), Some(config)) if config.destination.is_some() => { - DeliveryLedger::load(data_dir, source, session_id, config) - .await - .map_err(|error| { - tracing::warn!(source, session_id, "delivery ledger unavailable: {error}") - }) - .ok() + Some(DeliveryLedger::load(data_dir, source, session_id, config).await?) } _ => None, }; @@ -300,7 +294,7 @@ impl LedgerSink { Some((ledger, origins)) => (Some(ledger), origins), None => (None, HashMap::new()), }; - Self { + Ok(Self { inner, ledger, span_origins, @@ -309,7 +303,7 @@ impl LedgerSink { source_version: None, bt_version: bt_version.to_owned(), attached_parent: config.and_then(|config| config.attached_span_ids().0), - } + }) } } @@ -604,7 +598,8 @@ mod tests { Some(&config), "bt1", ) - .await; + .await + .unwrap(); first.set_capture_versions(None, Some("native1")); let turn = SpanRow { span_id: "turn-1".into(), @@ -652,7 +647,8 @@ mod tests { Some(&config), "bt2", ) - .await; + .await + .unwrap(); recovered.set_capture_versions(Some("v2"), Some("native2")); recovered .emit(&[ @@ -936,7 +932,8 @@ mod tests { Some(&config("project-a")), "test", ) - .await; + .await + .unwrap(); first.emit(&[terminal("span-1")]).await.unwrap(); first.flush().await.unwrap(); assert_eq!(first_output.lock().len(), 1); @@ -955,7 +952,8 @@ mod tests { Some(&same_destination), "test", ) - .await; + .await + .unwrap(); assert_eq!(repeated.emit(&[terminal("span-1")]).await.unwrap(), 0); assert_eq!(repeated.emit(&[partial_merge("span-1")]).await.unwrap(), 0); repeated.flush().await.unwrap(); @@ -974,7 +972,8 @@ mod tests { Some(&config("project-a")), "test", ) - .await; + .await + .unwrap(); assert_eq!(first.emit(&[terminal("span-1")]).await.unwrap(), 1); first.flush().await.unwrap(); @@ -987,7 +986,8 @@ mod tests { Some(&config("project-a")), "test", ) - .await; + .await + .unwrap(); assert_eq!( second .emit(&[late_merge("span-1", "output")]) @@ -1017,7 +1017,8 @@ mod tests { Some(&config("project-a")), "test", ) - .await; + .await + .unwrap(); assert_eq!( third.emit(&[late_merge("span-1", "output")]).await.unwrap(), 0 @@ -1040,7 +1041,8 @@ mod tests { Some(&config("project-a")), "test", ) - .await; + .await + .unwrap(); assert_eq!( first .emit(&[late_terminal_merge("root", "session:stop:3", 3)]) @@ -1059,7 +1061,8 @@ mod tests { Some(&config("project-a")), "test", ) - .await; + .await + .unwrap(); assert_eq!( resumed .emit(&[late_terminal_merge("root", "session:stop:3", 3)]) @@ -1092,7 +1095,8 @@ mod tests { Some(&config("project-a")), "test", ) - .await; + .await + .unwrap(); initial.emit(&[terminal("span-1")]).await.unwrap(); initial.flush().await.unwrap(); @@ -1108,7 +1112,8 @@ mod tests { Some(&config("project-b")), "test", ) - .await; + .await + .unwrap(); assert_eq!(replay.emit(&[terminal("span-1")]).await.unwrap(), 1); replay.flush().await.unwrap(); assert_eq!(replay_output.lock().len(), 1); @@ -1125,7 +1130,8 @@ mod tests { Some(&config("project-a")), "test", ) - .await; + .await + .unwrap(); let mut authoritative = terminal("span-1"); let SpanOp::Insert(row) = &mut authoritative else { unreachable!(); @@ -1143,7 +1149,8 @@ mod tests { Some(&config("project-a")), "test", ) - .await; + .await + .unwrap(); assert_eq!(repeated.emit(&[terminal("span-1")]).await.unwrap(), 0); assert_eq!( repeated diff --git a/bt-daemon/src/dispatch.rs b/bt-daemon/src/dispatch.rs index 132f3efe..04ffb939 100644 --- a/bt-daemon/src/dispatch.rs +++ b/bt-daemon/src/dispatch.rs @@ -1,7 +1,4 @@ -//! Per-session dispatch. Each session owns an ordered queue and a single actor -//! task that runs its translator + sink serially, so events for one session -//! are processed strictly in arrival order. Different sessions run -//! concurrently. +//! Session facade over ordered source publication and independent route delivery. //! //! Hook acknowledgement happens before this layer, immediately after the raw //! event is durably journaled. Translation, correlation, and delivery here are @@ -12,7 +9,6 @@ use crate::pipeline::{ translation::Translation, }; use crate::sink::SinkFactory; -use crate::translate::Registry; use crate::wire::Envelope; use parking_lot::Mutex; use std::path::PathBuf; @@ -70,17 +66,17 @@ pub struct Session { impl Session { /// Spawn a session's actor task and return its handle. - pub fn spawn( - options: SessionOptions, - _translators: Arc, - sink_factory: Arc, - ) -> Arc { + pub fn spawn(options: SessionOptions, sink_factory: Arc) -> Arc { let (tx, mut rx) = mpsc::channel(QUEUE_CAPACITY); let counters = Arc::new(Counters::default()); let recovery = Arc::new(Mutex::new(None)); let last_error = Arc::new(Mutex::new(None)); let permalink = Arc::new(Mutex::new(None)); let source = options.source.clone(); + let observer_key = format!( + "{}\u{1f}{}\u{1f}{}", + options.source, options.session_id, options.cursor.route + ); let route = Consumer::spawn( options.translation.journal.clone(), options.cursor, @@ -103,6 +99,7 @@ impl Session { let source_recovery = producer.recovery.clone(); let source_progress = producer.progress.clone(); let errors = last_error.clone(); + let source_errors = source_error.clone(); let queued = counters.clone(); tokio::spawn(async move { while let Some(message) = rx.recv().await { @@ -113,7 +110,7 @@ impl Session { route.configure(config); } if let Err(error) = producer.advance(through).await { - *errors.lock() = Some(error.to_string()); + *source_errors.lock() = Some(error.to_string()); } } SessionMsg::Configure(config, reply) => { @@ -134,6 +131,7 @@ impl Session { let route = route.clone(); let queued = queued.clone(); let errors = errors.clone(); + let source_errors = source_errors.clone(); tokio::spawn(async move { let mut reply = reply; let waiting = async { @@ -146,7 +144,7 @@ impl Session { } }, Err(error) => { - *errors.lock() = Some(error); + *source_errors.lock() = Some(error); queued.queued.load(Ordering::Relaxed).max(1) } } @@ -171,6 +169,7 @@ impl Session { ) .await; } + producer.unobserve(observer_key.clone()).await; route.shutdown().await; if producer.routes.load(Ordering::Relaxed) == 0 { producer.shutdown().await; diff --git a/bt-daemon/src/journal.rs b/bt-daemon/src/journal.rs index 7af6c6a9..c301022b 100644 --- a/bt-daemon/src/journal.rs +++ b/bt-daemon/src/journal.rs @@ -599,7 +599,7 @@ pub async fn gc_old_journals(data_dir: &Path, max_age: std::time::Duration) { } None => false, }; - if old && !protected { + if old && !protected && legacy_collectible(data_dir, &path).await { if let Err(e) = tokio::fs::remove_file(&path).await { tracing::warn!(path = %path.display(), "failed to remove stale journal: {e}"); } @@ -607,6 +607,53 @@ pub async fn gc_old_journals(data_dir: &Path, max_age: std::time::Duration) { } } +// Pre-SpanOp journals can be collected only with explicit completion evidence. +// Unknown or uncheckpointed history stays available for migration/reconstruction. +async fn legacy_collectible(data_dir: &Path, path: &Path) -> bool { + let Ok(Some(mut reader)) = + JournalReader::open(path, JournalReader::recorded_len(path).await).await + else { + return false; + }; + let mut observed = std::collections::HashMap::::new(); + let mut completed = std::collections::HashMap::::new(); + loop { + match reader.next_record().await { + Ok(Some(entry)) => match entry.record { + JournalRecord::Event(event) => { + let partition = data_dir.join("pipeline").join(format!( + "{}.sqlite", + crate::ids::session_storage_id(&event.source, &event.session_id) + )); + if tokio::fs::try_exists(partition).await.unwrap_or(true) { + return false; + } + let Some(route) = event.route else { + return false; + }; + observed.insert( + serde_json::to_string(&route).unwrap_or_default(), + entry.through, + ); + } + JournalRecord::DeliveryCheckpoint { route, through } => { + completed + .entry(serde_json::to_string(&route).unwrap_or_default()) + .and_modify(|previous| *previous = (*previous).max(through)) + .or_insert(through); + } + JournalRecord::Control(_) => return false, + }, + Ok(None) => break, + Err(_) => return false, + } + } + !observed.is_empty() + && observed + .iter() + .all(|(route, through)| completed.get(route).copied().unwrap_or(0) >= *through) +} + /// Reconstruct a translator-usable [`Envelope`] from a redacted journal entry. /// The live token is gone (redacted), so `auth.token` is empty — fine for /// rebuilding translator state; the sink must be re-supplied live credentials diff --git a/bt-daemon/src/pipeline/consumer.rs b/bt-daemon/src/pipeline/consumer.rs index 81f74333..475a045f 100644 --- a/bt-daemon/src/pipeline/consumer.rs +++ b/bt-daemon/src/pipeline/consumer.rs @@ -59,7 +59,7 @@ impl ConsumerOptions { Some(config), &self.version, ) - .await; + .await?; let mut sink = RoutingSink::spawn( Box::new(ledger), self.source.clone(), @@ -354,8 +354,37 @@ impl Consumer { }; let current = settings.borrow_and_update().clone(); if sink.is_none() { + journal + .bind(&cursor.route, ¤t) + .await + .map_err(|error| { + if error + .downcast_ref::() + .is_some() + { + Box::new(( + Failure::legacy( + error.to_string(), + cursor.completed_position() + 1, + ¤t, + ), + Work::Resubmit, + )) + } else { + options.storage(&error, &cursor, attempts, Work::Resubmit) + } + })?; sink = Some(options.sink(¤t).await.map_err(|error| { - options.routing(&error, &cursor, ¤t, attempts) + if error.downcast_ref::().is_some() + || error.downcast_ref::().is_some() + { + let mut held = + options.storage(&error, &cursor, attempts, Work::Resubmit); + held.0.boundary = Boundary::Ledger; + held + } else { + options.routing(&error, &cursor, ¤t, attempts) + } })?); } let routing = sink.as_mut().ok_or_else(|| { @@ -707,6 +736,69 @@ mod tests { assert_eq!(record.submitted.load(Ordering::SeqCst), 2); consumer.shutdown().await; } + #[tokio::test] + async fn real_sdk_rejection_holds_until_relevant_configuration_changes() { + use wiremock::{ + matchers::{method, path}, + Mock, MockServer, ResponseTemplate, + }; + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/version")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) + .mount(&server) + .await; + let rejection = Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(400).set_body_string("invalid row")) + .expect(1) + .mount_as_scoped(&server) + .await; + let tmp = tempfile::tempdir().unwrap(); + let journal = SpanJournal::open(&tmp.path().join("spans.sqlite")) + .await + .unwrap(); + stored(&journal).await; + let route = SessionRoute { + destination: Some(crate::wire::TraceDestination::ProjectLogs { + project_id: Some("project".into()), + project_name: None, + }), + ..Default::default() + }; + let mut opts = options(tmp.path(), &route, Arc::new(Recording::default())); + opts.config.auth.api_url = Some(server.uri()); + opts.config.auth.app_url = Some(server.uri()); + opts.factory = Arc::new(crate::BraintrustSinkFactory::new( + crate::BraintrustSinkConfig::default(), + )); + let mut config = opts.config.clone(); + let consumer = Consumer::spawn( + journal.clone(), + journal.subscribe(&route, 0).await.unwrap(), + opts, + ); + wait_for(|| consumer.error.lock().unwrap().is_some()).await; + assert_eq!(*consumer.progress.borrow(), 0); + drop(rejection); + Mock::given(method("POST")) + .and(path("/logs3")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) + .expect(1) + .mount(&server) + .await; + // A new credential cannot cure a permanently rejected span body. + config.auth.token = "replacement-token".into(); + consumer.configure(config.clone()); + tokio::time::sleep(std::time::Duration::from_millis(500)).await; + assert_eq!(*consumer.progress.borrow(), 0); + config.tags.push("corrected".into()); + consumer.configure(config); + wait_for(|| *consumer.progress.borrow() == 1).await; + assert!(consumer.error.lock().unwrap().is_none()); + consumer.shutdown().await; + } + #[tokio::test] async fn editing_a_predecessor_releases_a_later_plugins_hold() { let tmp = tempfile::tempdir().unwrap(); diff --git a/bt-daemon/src/pipeline/journal.rs b/bt-daemon/src/pipeline/journal.rs index a9a12940..e4e5c8f8 100644 --- a/bt-daemon/src/pipeline/journal.rs +++ b/bt-daemon/src/pipeline/journal.rs @@ -1,6 +1,9 @@ //! Concrete durable channel: one ordered SpanOp log and a head per route. use super::{Completion, Delivery}; -use crate::{translate::SpanOp, wire::SessionRoute}; +use crate::{ + translate::SpanOp, + wire::{SessionConfig, SessionRoute}, +}; use rusqlite::{params, Connection, OptionalExtension}; use serde::{Deserialize, Serialize}; use std::{ @@ -9,6 +12,17 @@ use std::{ }; use tokio::sync::watch; +#[derive(Debug)] +pub(crate) struct BindingChanged; +impl std::fmt::Display for BindingChanged { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str( + "credential resolution changed the durable head's pinned backend or organization", + ) + } +} +impl std::error::Error for BindingChanged {} + #[derive(Clone)] pub(crate) struct SpanJournal { connection: Arc>, @@ -62,28 +76,32 @@ impl SpanJournal { let db = Connection::open(path)?; db.busy_timeout(std::time::Duration::from_secs(5))?; let version: u32 = db.query_row("PRAGMA user_version", [], |row| row.get(0))?; - anyhow::ensure!(version <= 2, "unsupported SpanOp journal schema {version}"); - db.execute_batch( - "PRAGMA journal_mode=WAL; PRAGMA synchronous=FULL; BEGIN IMMEDIATE; - CREATE TABLE IF NOT EXISTS batches ( - input INTEGER PRIMARY KEY, committed INTEGER NOT NULL DEFAULT 0, - observation INTEGER NOT NULL DEFAULT 1); - CREATE TABLE IF NOT EXISTS ops ( - position INTEGER PRIMARY KEY AUTOINCREMENT, - input INTEGER NOT NULL, value BLOB NOT NULL); - CREATE INDEX IF NOT EXISTS ops_input ON ops(input); - CREATE TABLE IF NOT EXISTS heads ( - route TEXT PRIMARY KEY, position INTEGER NOT NULL DEFAULT 0, - failure TEXT); - CREATE TABLE IF NOT EXISTS recovery (id INTEGER PRIMARY KEY CHECK(id=1),failure TEXT); - ", - )?; - if version == 1 { + anyhow::ensure!(version <= 3, "unsupported SpanOp journal schema {version}"); + db.execute_batch("PRAGMA journal_mode=WAL; PRAGMA synchronous=FULL;")?; + if version < 3 { db.execute_batch( - "ALTER TABLE batches ADD COLUMN observation INTEGER NOT NULL DEFAULT 1", + "BEGIN IMMEDIATE; + CREATE TABLE IF NOT EXISTS batches ( + input INTEGER PRIMARY KEY, committed INTEGER NOT NULL DEFAULT 0, + observation INTEGER NOT NULL DEFAULT 1); + CREATE TABLE IF NOT EXISTS ops ( + position INTEGER PRIMARY KEY AUTOINCREMENT, + input INTEGER NOT NULL, value BLOB NOT NULL); + CREATE INDEX IF NOT EXISTS ops_input ON ops(input); + CREATE TABLE IF NOT EXISTS heads ( + route TEXT PRIMARY KEY, position INTEGER NOT NULL DEFAULT 0, + failure TEXT); + CREATE TABLE IF NOT EXISTS recovery (id INTEGER PRIMARY KEY CHECK(id=1),failure TEXT); + ", )?; + if version == 1 { + db.execute_batch( + "ALTER TABLE batches ADD COLUMN observation INTEGER NOT NULL DEFAULT 1", + )?; + } + db.execute_batch("ALTER TABLE heads ADD COLUMN binding TEXT")?; + db.execute_batch("PRAGMA user_version=3; COMMIT")?; } - db.execute_batch("PRAGMA user_version=2; COMMIT")?; Ok(db) }) .await??; @@ -280,6 +298,61 @@ impl SpanJournal { }).await } + /// Tokens can change; a durable head's backend and organization cannot. + pub async fn bind(&self, route: &str, config: &SessionConfig) -> anyhow::Result<()> { + let route = route.to_owned(); + let binding = serde_json::to_string(&serde_json::json!({ + "api_url": config.auth.api_url.as_deref().unwrap_or("https://api.braintrust.dev"), + "org_id":config.auth.org_id,"org_name":config.auth.org_name, + }))?; + self.with_db(move |db| { + let previous: Option = + db.query_row("SELECT binding FROM heads WHERE route=?", [&route], |r| { + r.get(0) + })?; + anyhow::ensure!( + previous + .as_ref() + .is_none_or(|previous| previous == &binding), + BindingChanged + ); + if previous.is_none() { + db.execute( + "UPDATE heads SET binding=? WHERE route=?", + params![binding, route], + )?; + } + Ok(()) + }) + .await + } + pub async fn routes(&self) -> anyhow::Result> { + self.with_db(|db| { + let mut query = db.prepare("SELECT route FROM heads ORDER BY route")?; + let rows = query + .query_map([], |r| r.get::<_, String>(0))? + .map(|route| Ok(serde_json::from_str(&route?)?)) + .collect(); + rows + }) + .await + } + /// Whole-partition retention only; an empty subscriber set is not completion. + pub async fn collectible(&self, input: u64) -> anyhow::Result { + self.with_db(move |db| Ok(db.query_row("SELECT EXISTS(SELECT 1 FROM heads) + AND NOT EXISTS(SELECT 1 FROM heads WHERE failure IS NOT NULL OR position < (SELECT coalesce(max(position),0) FROM ops JOIN batches USING(input) WHERE committed=1)) + AND NOT EXISTS(SELECT 1 FROM recovery WHERE failure IS NOT NULL) + AND NOT EXISTS(SELECT 1 FROM batches WHERE committed=0) + AND (SELECT coalesce(max(input),0) FROM batches WHERE committed=1)>=?",[input],|r|r.get(0))?)).await + } + pub async fn close_storage(&self) -> anyhow::Result<()> { + self.with_db(|db| { + db.execute_batch("PRAGMA wal_checkpoint(TRUNCATE)")?; + Ok(()) + }) + .await + } + pub async fn source_failure(&self) -> anyhow::Result> { self.with_db(|db| { Ok(db @@ -648,4 +721,46 @@ mod tests { legacy.tags.push("different route".into()); assert!(!journal.route_complete(&legacy, 10).await.unwrap()); } + + #[tokio::test] + async fn schema_two_heads_upgrade_and_keep_a_nonsecret_backend_binding() { + let tmp = tempfile::tempdir().unwrap(); + let path = tmp.path().join("old.sqlite"); + let db = Connection::open(&path).unwrap(); + db.execute_batch("CREATE TABLE heads(route TEXT PRIMARY KEY,position INTEGER NOT NULL DEFAULT 0,failure TEXT); PRAGMA user_version=2;").unwrap(); + let route = route("canonical"); + let key = serde_json::to_string(&route).unwrap(); + db.execute("INSERT INTO heads(route,position) VALUES(?,7)", [&key]) + .unwrap(); + drop(db); + let journal = SpanJournal::open(&path).await.unwrap(); + let mut config = route.with_auth(crate::wire::BackendAuth { + token: "secret-a".into(), + api_url: Some("https://one.test".into()), + app_url: None, + org_id: Some("org".into()), + org_name: None, + }); + journal.bind(&key, &config).await.unwrap(); + config.auth.token = "secret-b".into(); + journal.bind(&key, &config).await.unwrap(); + config.auth.api_url = Some("https://other.test".into()); + assert!(journal + .bind(&key, &config) + .await + .unwrap_err() + .downcast_ref::() + .is_some()); + let restored = journal.subscribe(&route, 0).await.unwrap(); + assert_eq!(restored.completed_position(), 7); + journal + .with_db(|db| { + let binding: String = + db.query_row("SELECT binding FROM heads", [], |r| r.get(0))?; + assert!(!binding.contains("secret-")); + Ok(()) + }) + .await + .unwrap(); + } } diff --git a/bt-daemon/src/pipeline/routing.rs b/bt-daemon/src/pipeline/routing.rs index 6ddb1ee0..86efe4cb 100644 --- a/bt-daemon/src/pipeline/routing.rs +++ b/bt-daemon/src/pipeline/routing.rs @@ -9,7 +9,6 @@ use tokio::sync::{mpsc, oneshot}; struct Request { config: SessionConfig, - versions: (Option, Option), operation: Operation, reply: oneshot::Sender, } @@ -53,7 +52,6 @@ impl Completion for Forward { pub(crate) struct RoutingSink { tx: mpsc::Sender>, config: Option, - versions: (Option, Option), pending: bool, permalink: Option, } @@ -73,10 +71,6 @@ impl RoutingSink { while let Ok(Some(delivery)) = InputPort::receive(&mut rx).await { let request = delivery.value(); sink.configure(&request.config); - sink.set_capture_versions( - request.versions.0.as_deref(), - request.versions.1.as_deref(), - ); let result = match &request.operation { Operation::Emit(ops) => { let mut ops = ops.clone(); @@ -150,7 +144,6 @@ impl RoutingSink { Self { tx, config: None, - versions: (None, None), pending: false, permalink: None, } @@ -164,7 +157,6 @@ impl RoutingSink { let (reply, receive) = oneshot::channel(); let delivery = Delivery::new(operation, Forward).map(|operation| Request { config, - versions: self.versions.clone(), operation, reply, }); @@ -187,9 +179,6 @@ impl Sink for RoutingSink { fn configure(&mut self, config: &SessionConfig) { self.config = Some(config.clone()); } - fn set_capture_versions(&mut self, plugin: Option<&str>, source: Option<&str>) { - self.versions = (plugin.map(str::to_owned), source.map(str::to_owned)); - } async fn emit(&mut self, ops: &[SpanOp]) -> anyhow::Result { self.request(Operation::Emit(ops.to_vec())).await } diff --git a/bt-daemon/src/pipeline/translation.rs b/bt-daemon/src/pipeline/translation.rs index ea98f134..2cedbf4d 100644 --- a/bt-daemon/src/pipeline/translation.rs +++ b/bt-daemon/src/pipeline/translation.rs @@ -12,7 +12,7 @@ use crate::{ }; use parking_lot::Mutex; use std::{ - collections::HashMap, + collections::{HashMap, HashSet}, path::PathBuf, sync::{ atomic::{AtomicU64, AtomicUsize, Ordering}, @@ -21,6 +21,7 @@ use std::{ }; use tokio::sync::{mpsc, oneshot, Notify}; +#[derive(Clone)] pub(crate) struct Observation { pub key: String, pub route: SessionRoute, @@ -30,6 +31,7 @@ pub(crate) struct Observation { enum Command { Observe(Box), Advance(u64), + Unobserve(String), Stop, Barrier(oneshot::Sender>), Control(ControlKind, oneshot::Sender>), @@ -101,6 +103,7 @@ impl Translation { observers: Vec::new(), origins: HashMap::new(), capture: None, + pending_snapshots: HashSet::new(), bounds: HashMap::new(), position: 0, committed, @@ -133,6 +136,9 @@ impl Translation { .await .map_err(|_| anyhow::anyhow!("translation actor is gone")) } + pub async fn unobserve(&self, key: String) { + let _ = self.tx.send(Command::Unobserve(key)).await; + } pub async fn advance(&self, through: u64) -> anyhow::Result<()> { self.progress .requested @@ -186,6 +192,7 @@ struct Producer { observers: Vec, origins: HashMap, capture: Option, + pending_snapshots: HashSet, bounds: HashMap, position: u64, committed: u64, @@ -235,6 +242,7 @@ impl Producer { Some(Command::Barrier(reply))=> { let _=reply.send(self.boundary().await); }, Some(Command::Observe(observer))=> { let _=self.observe_history(&observer).await; self.observers.push(*observer); }, Some(Command::Advance(_))=>{}, + Some(Command::Unobserve(key))=>self.observers.retain(|observer| observer.key != key), Some(Command::Control(kind,reply))=> { let result=self.append_control(kind).await; if let Err(error)=result { tracing::warn!(%error,"ordered control could not be persisted while source is recovering"); } @@ -260,12 +268,29 @@ impl Producer { attempts = failure.attempt.saturating_add(1); self.state = State::Processing; let result = async { + // Retry a failed correlation write before releasing the + // canonical candidate, even if registry observation is idempotent. + for key in self.pending_snapshots.clone() { + self.phase = Boundary::Journal; + crate::server::persist_active_parent_snapshot( + &self.options.data_dir, + &key, + &self.options.correlation, + ) + .await + .map_err(anyhow::Error::msg)?; + self.pending_snapshots.remove(&key); + } if self.publication.is_some() { self.publish().await?; } else if self.progress.requested.load(Ordering::Relaxed) > self.committed { self.reader = None; self.translator = None; self.position = 0; + } else { + for observer in self.observers.clone() { + self.observe_history(&observer).await?; + } } self.advance(self.progress.requested.load(Ordering::Relaxed)) .await @@ -299,6 +324,10 @@ impl Producer { result } Command::Advance(through) => self.advance(through).await, + Command::Unobserve(key) => { + self.observers.retain(|observer| observer.key != key); + Ok(()) + } Command::Barrier(reply) => { let _ = reply.send(self.boundary().await); Ok(()) @@ -563,7 +592,13 @@ impl Producer { self.phase = Boundary::Journal; if pending.observation { for observer in &self.observers { - Self::observe(&self.options, ops, observer).await?; + Self::observe( + &self.options, + &mut self.pending_snapshots, + ops, + observer, + ) + .await?; } } if publish { @@ -606,12 +641,13 @@ impl Producer { }; after = through; let ops = rows.into_iter().map(|(_, op)| op).collect::>(); - Self::observe(&self.options, &ops, observer).await?; + Self::observe(&self.options, &mut self.pending_snapshots, &ops, observer).await?; } Ok(()) } async fn observe( options: &TranslationOptions, + pending_snapshots: &mut HashSet, ops: &[SpanOp], observer: &Observation, ) -> anyhow::Result<()> { @@ -621,6 +657,9 @@ impl Producer { .correlation .observe_ops(&observer.key, &observer.route, &observer.config, &ops) { + pending_snapshots.insert(observer.key.clone()); + } + if pending_snapshots.contains(&observer.key) { crate::server::persist_active_parent_snapshot( &options.data_dir, &observer.key, @@ -628,6 +667,7 @@ impl Producer { ) .await .map_err(anyhow::Error::msg)?; + pending_snapshots.remove(&observer.key); options.correlation_changed.notify_one(); } Ok(()) @@ -747,6 +787,7 @@ mod tests { observers: Vec::new(), origins: HashMap::new(), capture: None, + pending_snapshots: HashSet::new(), bounds: HashMap::new(), position: 0, committed: 0, diff --git a/bt-daemon/src/server/auth.rs b/bt-daemon/src/server/auth.rs index a9783d1c..f87d5453 100644 --- a/bt-daemon/src/server/auth.rs +++ b/bt-daemon/src/server/auth.rs @@ -33,7 +33,49 @@ pub(crate) async fn resolve_route_auth( Ok(lease) } +/// A route refresh updates the same in-memory lease cache used by later hooks. +/// Persisted head identity remains independent of the credential itself. +struct RouteAuthProvider { + provider: Arc, + key: DeliveryKey, + states: Arc>>, +} +#[async_trait::async_trait] +impl AuthProvider for RouteAuthProvider { + async fn resolve( + &self, + selection: &AuthSelection, + reason: AuthResolveReason, + ) -> anyhow::Result { + let lease = self.provider.resolve(selection, reason).await?; + let mut states = self.states.lock().await; + if let Some(state) = states.get_mut(&self.key) { + anyhow::ensure!( + lease.selection == state.lease.selection, + "credential refresh changed the pinned auth selection" + ); + anyhow::ensure!( + lease.auth.api_url == state.lease.auth.api_url + && lease.auth.org_id == state.lease.auth.org_id + && lease.auth.org_name == state.lease.auth.org_name, + "credential refresh changed the pinned backend or organization" + ); + state.lease = lease.clone(); + } + Ok(lease) + } +} + impl Daemon { + pub(super) fn route_auth_provider(&self, key: &DeliveryKey) -> Option> { + self.auth_provider.as_ref().map(|provider| { + Arc::new(RouteAuthProvider { + provider: provider.clone(), + key: key.clone(), + states: self.session_auth.clone(), + }) as Arc + }) + } pub(super) async fn configure_event(&self, env: &mut Envelope) -> anyhow::Result { let Some(provider) = &self.auth_provider else { anyhow::bail!("daemon host has no Braintrust auth provider"); diff --git a/bt-daemon/src/server/ingress.rs b/bt-daemon/src/server/ingress.rs index cda3162d..a20056d6 100644 --- a/bt-daemon/src/server/ingress.rs +++ b/bt-daemon/src/server/ingress.rs @@ -217,8 +217,11 @@ pub(super) async fn accept_resolved_event( payload.insert("_bt_transcript_replay".to_string(), serde_json::json!(true)); } } - let schedule_flush = crate::hook::should_flush_ingress_event(&env) - || (env.source == "claude-code" && env.event == "TranscriptUpdate"); + // Startup resumes an existing boundary rather than inventing a new + // native/control observation that would require translator reconstruction. + let schedule_flush = journal_through > replay_through + && (crate::hook::should_flush_ingress_event(&env) + || (env.source == "claude-code" && env.event == "TranscriptUpdate")); let source = env.source.clone(); let event = env.event.clone(); let session_id = env.session_id.clone(); diff --git a/bt-daemon/src/server/maintenance.rs b/bt-daemon/src/server/maintenance.rs index ea5ded27..401d89ec 100644 --- a/bt-daemon/src/server/maintenance.rs +++ b/bt-daemon/src/server/maintenance.rs @@ -9,42 +9,244 @@ pub(super) const RETENTION: Duration = Duration::from_secs(7 * 24 * 60 * 60); /// How often that retention is enforced while the daemon keeps running. pub(super) const GC_INTERVAL: Duration = Duration::from_secs(60 * 60); -pub(super) async fn collect_garbage(data_dir: &std::path::Path) { - journal::gc_old_journals(data_dir, RETENTION).await; - journal::gc_old_managed_runs(data_dir, RETENTION).await; - // Durable partitions may still need any recorded transcript prefix. - let mut partitions = tokio::fs::read_dir(data_dir.join("pipeline")).await.ok(); - let mut retained = false; - if let Some(entries) = partitions.as_mut() { - while let Ok(Some(entry)) = entries.next_entry().await { - retained |= entry.path().extension().and_then(|ext| ext.to_str()) == Some("sqlite"); +pub(super) async fn collect_garbage(daemon: &Daemon) { + collect_partitions(daemon, RETENTION).await; + journal::gc_old_journals(&daemon.data_dir, RETENTION).await; + journal::gc_old_managed_runs(&daemon.data_dir, RETENTION).await; + // Protect references from every retained input log, including completed + // partitions still inside retention and sessions with no registered heads. + let Ok(protected) = retained_references(&daemon.data_dir).await else { + return; + }; + crate::transcript_mirror::gc_old_mirrors(&daemon.data_dir, RETENTION, &protected).await; + gc_old_correlation_dir(&daemon.data_dir.join("correlation"), RETENTION, &protected).await; + gc_old_correlation_dir( + &daemon.data_dir.join("correlation/parents"), + RETENTION, + &protected, + ) + .await; +} + +async fn old(path: &std::path::Path, age: Duration) -> bool { + tokio::fs::metadata(path) + .await + .ok() + .and_then(|m| m.modified().ok()) + .and_then(|modified| SystemTime::now().duration_since(modified).ok()) + .is_some_and(|elapsed| elapsed > age) +} + +async fn collect_partitions(daemon: &Daemon, age: Duration) { + let Ok(mut entries) = tokio::fs::read_dir(journal::journal_dir(&daemon.data_dir)).await else { + return; + }; + while let Ok(Some(entry)) = entries.next_entry().await { + let path = entry.path(); + if path.extension().and_then(|v| v.to_str()) != Some("ndjson") || !old(&path, age).await { + continue; + } + let recorded = journal::JournalReader::recorded_len(&path).await; + let Ok(Some(mut reader)) = journal::JournalReader::open(&path, recorded).await else { + continue; + }; + let Ok(Some(first)) = reader.next_entry().await else { + continue; + }; + let id = crate::ids::session_storage_id(&first.source, &first.session_id); + // Legacy mixed-source logs are handled by their own conservative reader. + if path != journal::source_journal_path(&daemon.data_dir, &first.source, &first.session_id) + { + continue; + } + let storage = daemon + .data_dir + .join("pipeline") + .join(format!("{id}.sqlite")); + if !old(&storage, age).await { + continue; + } + let wal = storage.with_extension("sqlite-wal"); + if tokio::fs::try_exists(&wal).await.unwrap_or(true) && !old(&wal, age).await { + continue; + } + let lock = daemon.session_lock(&first.source, &first.session_id); + let _guard = lock.lock().await; + let namespace = crate::ids::session_namespace(&first.source, &first.session_id); + if daemon.journals.lock().unwrap().contains_key(&namespace) + || daemon.translations.lock().unwrap().contains_key(&namespace) + { + continue; + } + if journal::JournalReader::recorded_len(&path).await != recorded { + continue; + } + let Ok(partition) = crate::pipeline::journal::SpanJournal::open(&storage).await else { + continue; + }; + let Ok(routes) = partition.routes().await else { + continue; + }; + let pending = daemon + .pending_sessions + .lock() + .unwrap() + .values() + .any(|state| { + state.events.iter().any(|event| { + event.env.source == first.source && event.env.session_id == first.session_id + }) + }); + if pending || !captured_routes_registered(&daemon.data_dir, &path, recorded, &routes).await + { + continue; + } + + if routes.iter().any(|route| { + DeliveryKey::new(&first.source, &first.session_id, route) + .is_ok_and(|key| daemon.correlation.has_active_tools(&key.correlation_key())) + }) { + continue; + } + if !partition.collectible(recorded).await.unwrap_or(false) + || partition.close_storage().await.is_err() + { + continue; + } + drop(partition); + drop(reader); + // Delete the output partition first. A failed input removal can only + // cause idempotent replay; it cannot hide new captures behind an old receipt. + if tokio::fs::remove_file(&storage).await.is_err() { + continue; + } + for sidecar in [wal, storage.with_extension("sqlite-shm"), path] { + if let Err(error) = tokio::fs::remove_file(&sidecar).await { + if error.kind() != std::io::ErrorKind::NotFound { + tracing::warn!(%error,path=%sidecar.display(),"expired partition cleanup failed"); + } + } } } - if !retained { - crate::transcript_mirror::gc_old_mirrors(data_dir, RETENTION).await; +} + +// Authentication/discovery may have prevented subscription registration. Such +// captured routes are obligations too, even if all existing heads are complete. +async fn captured_routes_registered( + data_dir: &std::path::Path, + path: &std::path::Path, + recorded: u64, + heads: &[SessionRoute], +) -> bool { + let Ok(Some(mut reader)) = journal::JournalReader::open(path, recorded).await else { + return false; + }; + loop { + match reader.next_record().await { + Ok(Some(record)) => { + if let journal::JournalRecord::Event(event) = record.record { + let env = journal::envelope_from_redacted(event); + let route = read_correlation_state(data_dir, &automatic_link_key(&env)) + .await + .and_then(|state| state.linked_route) + .or(env.route); + if route + .as_ref() + .is_none_or(|route| !heads.iter().any(|head| head.same_route(route))) + { + return false; + } + } + } + Ok(None) => return true, + Err(_) => return false, + } } - gc_old_correlation_states(data_dir, RETENTION).await; } -pub(super) async fn gc_old_correlation_states(data_dir: &std::path::Path, max_age: Duration) { - gc_old_correlation_dir(&data_dir.join("correlation"), max_age).await; - gc_old_correlation_dir(&data_dir.join("correlation").join("parents"), max_age).await; +async fn retained_references(data_dir: &std::path::Path) -> anyhow::Result> { + fn mirrors(value: &Value, protected: &mut HashSet) { + match value { + Value::Object(fields) => { + for (key, value) in fields { + if key == "mirror" { + if let Some(path) = value.as_str() { + protected.insert(path.into()); + } + } + mirrors(value, protected); + } + } + Value::Array(values) => { + for value in values { + mirrors(value, protected); + } + } + _ => {} + } + } + let mut protected = HashSet::new(); + let mut sources = HashSet::new(); + let mut entries = match tokio::fs::read_dir(journal::journal_dir(data_dir)).await { + Ok(entries) => entries, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(protected), + Err(error) => return Err(error.into()), + }; + while let Some(entry) = entries.next_entry().await? { + if entry.path().extension().and_then(|v| v.to_str()) != Some("ndjson") { + continue; + } + let mut reader = journal::JournalReader::open( + &entry.path(), + journal::JournalReader::recorded_len(&entry.path()).await, + ) + .await?; + while let Some(record) = match reader.as_mut() { + Some(reader) => reader.next_record().await?, + None => None, + } { + if let journal::JournalRecord::Event(event) = record.record { + mirrors(&event.payload, &mut protected); + if let Some(route) = &event.route { + let key = DeliveryKey::new(&event.source, &event.session_id, route)? + .correlation_key(); + protected.insert(correlation_state_path(data_dir, &key)); + protected.insert(active_parent_snapshot_path(data_dir, &key)); + } + let storage = data_dir.join("pipeline").join(format!( + "{}.sqlite", + crate::ids::session_storage_id(&event.source, &event.session_id) + )); + if sources.insert((event.source.clone(), event.session_id.clone())) + && tokio::fs::try_exists(&storage).await? + { + let partition = crate::pipeline::journal::SpanJournal::open(&storage).await?; + for route in partition.routes().await? { + let key = DeliveryKey::new(&event.source, &event.session_id, &route)? + .correlation_key(); + protected.insert(correlation_state_path(data_dir, &key)); + protected.insert(active_parent_snapshot_path(data_dir, &key)); + } + } + } + } + } + Ok(protected) } -pub(super) async fn gc_old_correlation_dir(dir: &std::path::Path, max_age: Duration) { +async fn gc_old_correlation_dir( + dir: &std::path::Path, + age: Duration, + protected: &HashSet, +) { let Ok(mut entries) = tokio::fs::read_dir(dir).await else { return; }; - let now = SystemTime::now(); while let Ok(Some(entry)) = entries.next_entry().await { - let old = entry - .metadata() - .await - .ok() - .and_then(|metadata| metadata.modified().ok()) - .and_then(|modified| now.duration_since(modified).ok()) - .is_some_and(|age| age > max_age); - if old && entry.file_type().await.is_ok_and(|kind| kind.is_file()) { + if !protected.contains(&entry.path()) + && entry.file_type().await.is_ok_and(|kind| kind.is_file()) + && old(&entry.path(), age).await + { let _ = tokio::fs::remove_file(entry.path()).await; } } @@ -79,7 +281,7 @@ pub(super) fn spawn_gc(daemon: Arc) { _ = daemon.shutdown.notified() => return, _ = tokio::time::sleep(GC_INTERVAL) => {} } - collect_garbage(&daemon.data_dir).await; + collect_garbage(&daemon).await; daemon .correlation .prune_retired(crate::process::process_is_alive); @@ -196,3 +398,163 @@ pub(super) async fn probe_alive(endpoint: &std::path::Path) -> bool { Ok(Ok(Some(_))) ) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + translate::{SpanOp, SpanRow}, + wire::Envelope, + }; + + fn event(session: &str, route: SessionRoute) -> Envelope { + Envelope { + source: "debug".into(), + session_id: session.into(), + event: "event".into(), + ts_ms: 1, + payload: serde_json::json!({}), + route: Some(route), + config: None, + capture: None, + managed_run_id: None, + plugin_version: None, + source_version: None, + } + } + #[tokio::test] + async fn retention_requires_complete_inactive_partitions_and_nonempty_heads() { + for scenario in [ + "no_heads", + "unresolved", + "unregistered", + "held", + "source_held", + "untranslated", + "invisible", + "active", + "within_retention", + "complete", + ] { + let tmp = tempfile::tempdir().unwrap(); + let env = event(scenario, SessionRoute::default()); + let raw = journal::source_journal_path(tmp.path(), &env.source, &env.session_id); + let mut writer = JournalWriter::open_path(&raw).await.unwrap(); + let mut through = writer.append(&env).await.unwrap(); + if scenario == "unregistered" { + let mut other = env.clone(); + other + .route + .as_mut() + .unwrap() + .tags + .push("not registered".into()); + through = writer.append(&other).await.unwrap(); + } + drop(writer); + let storage = tmp.path().join("pipeline").join(format!( + "{}.sqlite", + crate::ids::session_storage_id(&env.source, &env.session_id) + )); + let partition = crate::pipeline::journal::SpanJournal::open(&storage) + .await + .unwrap(); + let input = if scenario == "untranslated" { + through - 1 + } else { + through + }; + partition.begin(input, true).await.unwrap(); + partition + .append( + input, + &[SpanOp::Insert(SpanRow { + span_id: "stored".into(), + ..Default::default() + })], + ) + .await + .unwrap(); + partition.commit(input).await.unwrap(); + if scenario != "no_heads" { + let mut head = partition + .subscribe(env.route.as_ref().unwrap(), 0) + .await + .unwrap(); + if scenario != "unresolved" { + head.next_batch(8) + .await + .unwrap() + .unwrap() + .into_completed() + .await + .ok() + .unwrap(); + } + if scenario == "held" { + partition + .set_failure(&head.route, Some("held".into())) + .await + .unwrap(); + } + } + if scenario == "source_held" { + partition + .set_source_failure(Some("held".into())) + .await + .unwrap(); + } + if scenario == "invisible" { + partition.begin(through + 1, true).await.unwrap(); + } + drop(partition); + let daemon = Daemon::new( + super::super::debug_serve_options("test", tmp.path()), + tmp.path().into(), + ); + if scenario == "active" { + daemon + .journal_writer_for(&env.source, &env.session_id) + .await + .unwrap(); + } + let age = if scenario == "within_retention" { + RETENTION + } else { + Duration::ZERO + }; + collect_partitions(&daemon, age).await; + assert_eq!(storage.exists(), scenario != "complete", "{scenario}"); + assert_eq!(raw.exists(), scenario != "complete", "{scenario}"); + drain_all(&daemon).await; + } + } + #[tokio::test] + async fn retained_inputs_protect_only_their_reconstruction_references() { + let tmp = tempfile::tempdir().unwrap(); + let dir = crate::transcript_mirror::mirror_dir(tmp.path()); + std::fs::create_dir_all(&dir).unwrap(); + let retained = dir.join("retained.jsonl"); + let obsolete = dir.join("obsolete.jsonl"); + std::fs::write(&retained, "retained").unwrap(); + std::fs::write(&obsolete, "obsolete").unwrap(); + let mut env = event("test", SessionRoute::default()); + env.payload = serde_json::json!({"_bt_transcript_mirror":{"mirror": retained,"through":8}}); + let raw = journal::source_journal_path(tmp.path(), "debug", "test"); + let mut writer = JournalWriter::open_path(&raw).await.unwrap(); + writer.append(&env).await.unwrap(); + drop(writer); + let key = DeliveryKey::new("debug", "test", env.route.as_ref().unwrap()) + .unwrap() + .correlation_key(); + let parent = active_parent_snapshot_path(tmp.path(), &key); + std::fs::create_dir_all(parent.parent().unwrap()).unwrap(); + std::fs::write(&parent, "retained parent").unwrap(); + let protected = retained_references(tmp.path()).await.unwrap(); + crate::transcript_mirror::gc_old_mirrors(tmp.path(), Duration::ZERO, &protected).await; + gc_old_correlation_dir(parent.parent().unwrap(), Duration::ZERO, &protected).await; + assert!(retained.exists()); + assert!(parent.exists()); + assert!(!obsolete.exists()); + } +} diff --git a/bt-daemon/src/server/mod.rs b/bt-daemon/src/server/mod.rs index 64284699..fba1ef20 100644 --- a/bt-daemon/src/server/mod.rs +++ b/bt-daemon/src/server/mod.rs @@ -108,7 +108,7 @@ pub struct Daemon { translators: Arc, sink_factory: Arc, auth_provider: Option>, - session_auth: tokio::sync::Mutex>, + session_auth: Arc>>, route_aliases: Mutex>, /// Serializes only capture-side journal appends for a source session. session_locks: Mutex>>>, @@ -151,7 +151,7 @@ impl Daemon { translators: opts.translators, sink_factory: opts.sink_factory, auth_provider: opts.auth_provider, - session_auth: tokio::sync::Mutex::new(HashMap::new()), + session_auth: Arc::new(tokio::sync::Mutex::new(HashMap::new())), route_aliases: Mutex::new(HashMap::new()), session_locks: Mutex::new(HashMap::new()), dispatch_locks: Mutex::new(HashMap::new()), @@ -241,9 +241,9 @@ pub async fn run(args: ServeArgs, opts: ServeOptions) -> anyhow::Result<()> { tracing::info!(socket = %socket.display(), "bt-daemon listening"); let daemon = Daemon::new(opts, data_dir); - collect_garbage(&daemon.data_dir).await; restore_active_parent_snapshots(&daemon.data_dir, &daemon.correlation).await; restore_pending_sessions(&daemon).await; + collect_garbage(&daemon).await; recover_unprocessed_journals(&daemon).await; spawn_pending_reconciler(daemon.clone()); let idle_timeout = Duration::from_secs(args.idle_timeout_secs); diff --git a/bt-daemon/src/server/recovery.rs b/bt-daemon/src/server/recovery.rs index 2eb6d4c2..ede9305c 100644 --- a/bt-daemon/src/server/recovery.rs +++ b/bt-daemon/src/server/recovery.rs @@ -73,6 +73,50 @@ pub(super) async fn recover_unprocessed_journals(daemon: &Arc) { } before = through; } + // Heads are the authoritative subscriptions, including dormant routes + // and canonical profile IDs absent from older captured envelopes. + let templates: Vec<_> = latest_by_route.values().cloned().collect(); + let mut sources = HashSet::new(); + for template in templates { + if !sources.insert((template.env.source.clone(), template.env.session_id.clone())) { + continue; + } + let span_path = daemon.data_dir.join("pipeline").join(format!( + "{}.sqlite", + crate::ids::session_storage_id(&template.env.source, &template.env.session_id) + )); + if !tokio::fs::try_exists(&span_path).await.unwrap_or(false) { + continue; + } + let Ok(journal) = crate::pipeline::journal::SpanJournal::open(&span_path).await else { + continue; + }; + let Ok(routes) = journal.routes().await else { + continue; + }; + for route in routes { + let matching = latest_by_route + .iter() + .find(|(_, event)| { + event.env.source == template.env.source + && event.env.session_id == template.env.session_id + && event + .env + .route + .as_ref() + .is_some_and(|recorded| recorded.same_route(&route)) + }) + .map(|(key, _)| key.clone()); + let mut event = matching + .and_then(|key| latest_by_route.remove(&key)) + .unwrap_or_else(|| template.clone()); + event.env.route = Some(route.clone()); + latest_by_route.insert( + serde_json::to_string(&route).expect("persisted route"), + event, + ); + } + } for mut event in latest_by_route.into_values() { if pending.contains(&( event.env.source.clone(), diff --git a/bt-daemon/src/server/sessions.rs b/bt-daemon/src/server/sessions.rs index dc2ff2f3..93da136f 100644 --- a/bt-daemon/src/server/sessions.rs +++ b/bt-daemon/src/server/sessions.rs @@ -128,14 +128,13 @@ impl Daemon { SessionOptions { translation, cursor, - auth_provider: self.auth_provider.clone(), + auth_provider: self.route_auth_provider(key), session_id: env.session_id.clone(), source: env.source.clone(), bt_version: self.version.clone(), config, data_dir: self.data_dir.clone(), }, - self.translators.clone(), self.sink_factory.clone(), ); map.insert(key.clone(), session.clone()); diff --git a/bt-daemon/src/transcript_import/pipeline.rs b/bt-daemon/src/transcript_import/pipeline.rs index d158a8fb..7ccfdf12 100644 --- a/bt-daemon/src/transcript_import/pipeline.rs +++ b/bt-daemon/src/transcript_import/pipeline.rs @@ -319,7 +319,7 @@ impl ImportProcessor { env.config.as_ref(), &self.opts.version, ) - .await, + .await?, ); self.sessions.insert( sid.clone(), diff --git a/bt-daemon/src/transcript_mirror.rs b/bt-daemon/src/transcript_mirror.rs index ec2ca783..2a098485 100644 --- a/bt-daemon/src/transcript_mirror.rs +++ b/bt-daemon/src/transcript_mirror.rs @@ -165,7 +165,11 @@ pub async fn capture_generation( /// Best-effort age-based collection, mirroring journal GC. Mirrors are only /// useful for as long as their journal survives. -pub async fn gc_old_mirrors(data_dir: &Path, max_age: std::time::Duration) { +pub(crate) async fn gc_old_mirrors( + data_dir: &Path, + max_age: std::time::Duration, + protected: &std::collections::HashSet, +) { let dir = mirror_dir(data_dir); let Ok(mut entries) = tokio::fs::read_dir(&dir).await else { return; @@ -186,7 +190,7 @@ pub async fn gc_old_mirrors(data_dir: &Path, max_age: std::time::Duration) { .and_then(|meta| meta.modified().ok()) .and_then(|modified| now.duration_since(modified).ok()) .is_some_and(|age| age > max_age); - if old { + if old && !protected.contains(&path) { let _ = tokio::fs::remove_file(&path).await; } } diff --git a/bt-daemon/tests/pipeline.rs b/bt-daemon/tests/pipeline.rs index 8c143e5c..856d10f2 100644 --- a/bt-daemon/tests/pipeline.rs +++ b/bt-daemon/tests/pipeline.rs @@ -3524,6 +3524,84 @@ async fn missing_credentials_recover_without_another_native_event() { handle.await.unwrap(); } +#[tokio::test] +async fn startup_restores_a_durable_head_absent_from_captured_routes() { + let provider = Arc::new(TestAuthProvider { + calls: Mutex::new(Vec::new()), + fail: false, + first_lease_expired: false, + }); + let (data_dir, socket, handle, _tmp) = start_routed_daemon(provider.clone()).await; + let env = envelope("durable-subscription", "SessionStart", 1); + forward_envelope(&env, &socket, &dummy_host(), false) + .await + .unwrap(); + assert!( + flush_session("durable-subscription", &socket, 5000) + .await + .unwrap() + .flushed + ); + shutdown(&socket).await; + handle.await.unwrap(); + let path = std::fs::read_dir(data_dir.join("pipeline")) + .unwrap() + .filter_map(Result::ok) + .map(|entry| entry.path()) + .find(|path| path.extension().is_some_and(|ext| ext == "sqlite")) + .unwrap(); + let mut dormant = env.route.unwrap(); + dormant.tags = vec!["dormant".into()]; + let db = rusqlite::Connection::open(&path).unwrap(); + db.execute( + "INSERT INTO heads(route,position) VALUES(?,0)", + [serde_json::to_string(&dormant).unwrap()], + ) + .unwrap(); + drop(db); + let recording = Arc::new(RouteRecordingSinkFactory::default()); + let handle = start_daemon_at_with( + data_dir.clone(), + socket.clone(), + provider, + recording.clone(), + ) + .await; + tokio::time::timeout(Duration::from_secs(5), async { + loop { + let sinks = recording.sinks.lock().unwrap().clone(); + if sinks + .iter() + .any(|sink| sink.emitted.load(std::sync::atomic::Ordering::SeqCst) > 0) + { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + assert_eq!( + recording.sinks.lock().unwrap().len(), + 1, + "only the dormant head has unresolved work" + ); + let status = run_status(StatusArgs { + socket: Some(socket.clone()), + session_id: Some("durable-subscription".into()), + }) + .await + .unwrap() + .unwrap(); + assert_eq!(status.sessions.len(), 1); + assert_eq!( + status.sessions[0].route.as_ref().unwrap().tags, + dormant.tags + ); + shutdown(&socket).await; + handle.await.unwrap(); +} + #[tokio::test] async fn clean_shutdown_delivers_the_shared_retirement_tail_to_every_route() { use std::sync::atomic::{AtomicUsize, Ordering}; From 3ae80f5ef278aa58c92cb3b5bb1953e7639c2b30 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 14 Oct 2026 11:24:19 +0800 Subject: [PATCH 2/5] Close shared journal storage after cancelled reads on shutdown --- bt-daemon/src/pipeline/journal.rs | 55 +++++++++++++++++++++++++-- bt-daemon/src/pipeline/translation.rs | 11 ++++-- 2 files changed, 59 insertions(+), 7 deletions(-) diff --git a/bt-daemon/src/pipeline/journal.rs b/bt-daemon/src/pipeline/journal.rs index e4e5c8f8..8ab057d9 100644 --- a/bt-daemon/src/pipeline/journal.rs +++ b/bt-daemon/src/pipeline/journal.rs @@ -25,7 +25,7 @@ impl std::error::Error for BindingChanged {} #[derive(Clone)] pub(crate) struct SpanJournal { - connection: Arc>, + connection: Arc>>, changed: watch::Sender, } @@ -107,7 +107,7 @@ impl SpanJournal { .await??; let (changed, _) = watch::channel(0); Ok(Self { - connection: Arc::new(Mutex::new(connection)), + connection: Arc::new(Mutex::new(Some(connection))), changed, }) } @@ -121,7 +121,27 @@ impl SpanJournal { let mut db = db .lock() .map_err(|_| anyhow::anyhow!("SpanOp storage lock poisoned"))?; - operation(&mut db) + operation( + db.as_mut() + .ok_or_else(|| anyhow::anyhow!("SpanOp journal is closed"))?, + ) + }) + .await? + } + + /// Closing the shared cell also covers cancelled reads whose blocking + /// jobs can outlive their actor. Queued jobs cannot reopen the connection. + pub async fn close(&self) -> anyhow::Result<()> { + let connection = self.connection.clone(); + tokio::task::spawn_blocking(move || -> anyhow::Result<()> { + let database = connection + .lock() + .map_err(|_| anyhow::anyhow!("SpanOp storage lock poisoned"))? + .take(); + if let Some(database) = database { + database.close().map_err(|(_, error)| error)?; + } + Ok(()) }) .await? } @@ -763,4 +783,33 @@ mod tests { .await .unwrap(); } + #[tokio::test] + async fn close_releases_storage_even_when_a_blocking_read_was_cancelled() { + let tmp = tempfile::tempdir().unwrap(); + let path = tmp.path().join("spans.sqlite"); + let journal = SpanJournal::open(&path).await.unwrap(); + let reader = journal.clone(); + let (entered, started) = tokio::sync::oneshot::channel(); + let (release, blocked) = std::sync::mpsc::channel(); + let read = tokio::spawn(async move { + reader + .with_db(move |_| { + let _ = entered.send(()); + blocked.recv().unwrap(); + Ok(()) + }) + .await + }); + started.await.unwrap(); + read.abort(); + assert!(read.await.unwrap_err().is_cancelled()); + let closing = journal.clone(); + let close = tokio::spawn(async move { closing.close().await }); + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + assert!(!close.is_finished()); + release.send(()).unwrap(); + close.await.unwrap().unwrap(); + assert!(journal.tail().await.is_err()); + std::fs::remove_file(path).unwrap(); + } } diff --git a/bt-daemon/src/pipeline/translation.rs b/bt-daemon/src/pipeline/translation.rs index 2cedbf4d..0966d709 100644 --- a/bt-daemon/src/pipeline/translation.rs +++ b/bt-daemon/src/pipeline/translation.rs @@ -50,7 +50,7 @@ impl Progress { pub(crate) struct Translation { tx: mpsc::Sender, - task: Mutex>>, + task: tokio::sync::Mutex>>, pub journal: SpanJournal, pub error: Arc>>, pub recovery: Arc>>, @@ -114,7 +114,7 @@ impl Translation { let task = tokio::spawn(actor.run(rx)); Ok(Arc::new(Self { tx, - task: Mutex::new(Some(task)), + task: tokio::sync::Mutex::new(Some(task)), journal, error, recovery, @@ -124,10 +124,13 @@ impl Translation { } pub async fn shutdown(&self) { let _ = self.tx.send(Command::Stop).await; - let task = self.task.lock().take(); - if let Some(task) = task { + let mut task = self.task.lock().await; + if let Some(task) = task.take() { let _ = task.await; } + if let Err(error) = self.journal.close().await { + tracing::warn!(%error, "SpanOp journal could not be closed during source shutdown"); + } } pub async fn observe(&self, observer: Observation) -> anyhow::Result<()> { self.routes.fetch_add(1, Ordering::Relaxed); From a876cf9472dd3ca980226d849bf6ec9466557fcd Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 14 Oct 2026 21:31:13 +0800 Subject: [PATCH 3/5] Preserve route isolation during refresh, recovery, and retention --- bt-daemon/src/server/auth.rs | 105 ++++++++++++++++++--- bt-daemon/src/server/maintenance.rs | 84 ++++++++++------- bt-daemon/src/server/recovery.rs | 141 +++++++++++++++++++++++++--- 3 files changed, 267 insertions(+), 63 deletions(-) diff --git a/bt-daemon/src/server/auth.rs b/bt-daemon/src/server/auth.rs index f87d5453..5b047bf4 100644 --- a/bt-daemon/src/server/auth.rs +++ b/bt-daemon/src/server/auth.rs @@ -77,9 +77,9 @@ impl Daemon { }) } pub(super) async fn configure_event(&self, env: &mut Envelope) -> anyhow::Result { - let Some(provider) = &self.auth_provider else { + if self.auth_provider.is_none() { anyhow::bail!("daemon host has no Braintrust auth provider"); - }; + } let requested_route = env .route .clone() @@ -118,6 +118,9 @@ impl Daemon { } }; + let provider = self + .route_auth_provider(&key) + .expect("auth provider checked"); let lease = resolve_route_auth( provider.as_ref(), &selection, @@ -217,7 +220,7 @@ impl Daemon { &self, key: &DeliveryKey, ) -> anyhow::Result<()> { - let Some(provider) = &self.auth_provider else { + let Some(provider) = self.route_auth_provider(key) else { return Ok(()); }; let Some(state) = self.session_auth.lock().await.get(key).cloned() else { @@ -230,17 +233,7 @@ impl Daemon { let lease = provider .resolve(&selection, AuthResolveReason::Expiring) .await?; - if lease.selection != state.lease.selection { - anyhow::bail!("credential refresh changed the session auth selection"); - } - let config = state.route.with_auth(lease.auth.clone()); - self.session_auth.lock().await.insert( - key.clone(), - SessionAuthState { - route: state.route, - lease, - }, - ); + let config = state.route.with_auth(lease.auth); let session = { self.sessions.lock().get(key).cloned() }; if let Some(session) = session { session.configure(config).await?; @@ -248,3 +241,87 @@ impl Daemon { Ok(()) } } + +#[cfg(test)] +mod tests { + use super::*; + + struct Refresh(AuthLease); + #[async_trait::async_trait] + impl AuthProvider for Refresh { + async fn resolve( + &self, + _: &AuthSelection, + reason: AuthResolveReason, + ) -> anyhow::Result { + assert_eq!(reason, AuthResolveReason::Expiring); + Ok(self.0.clone()) + } + } + + #[tokio::test] + async fn expiring_refreshes_preserve_backend_binding_and_accept_new_tokens() { + for trigger in ["flush", "event"] { + for change in ["backend", "org_id", "org_name", "token"] { + let tmp = tempfile::tempdir().unwrap(); + let route = SessionRoute { + destination: Some(crate::wire::TraceDestination::ProjectLogs { + project_id: Some("project".into()), + project_name: None, + }), + ..Default::default() + }; + let lease = AuthLease { + selection: route.auth.clone(), + auth: BackendAuth { + token: "old".into(), + api_url: Some("https://original.test".into()), + app_url: None, + org_id: Some("original-id".into()), + org_name: Some("original-name".into()), + }, + expires_at_ms: Some(0), + }; + let mut refreshed = lease.clone(); + refreshed.auth.token = "new".into(); + match change { + "backend" => refreshed.auth.api_url = Some("https://other.test".into()), + "org_id" => refreshed.auth.org_id = Some("other-id".into()), + "org_name" => refreshed.auth.org_name = Some("other-name".into()), + _ => {} + } + let mut options = debug_serve_options("test", tmp.path()); + options.auth_provider = Some(Arc::new(Refresh(refreshed))); + let daemon = Daemon::new(options, tmp.path().into()); + let key = DeliveryKey::new("debug", "test", &route).unwrap(); + daemon.session_auth.lock().await.insert( + key.clone(), + SessionAuthState { + route: route.clone(), + lease: lease.clone(), + }, + ); + let mut env: Envelope = serde_json::from_value(serde_json::json!({ + "source":"debug", "session_id":"test", "event":"event", "ts_ms":1, + "payload":{}, "route":route, + })) + .unwrap(); + let result = if trigger == "flush" { + daemon.refresh_session_before_flush(&key).await + } else { + daemon.configure_event(&mut env).await.map(|_| ()) + }; + assert_eq!(result.is_ok(), change == "token", "{trigger}: {change}"); + let cached = daemon.session_auth.lock().await[&key].lease.clone(); + assert_eq!(cached.auth.api_url, lease.auth.api_url); + assert_eq!(cached.auth.org_id, lease.auth.org_id); + assert_eq!(cached.auth.org_name, lease.auth.org_name); + assert_eq!( + cached.auth.token, + if change == "token" { "new" } else { "old" } + ); + super::super::maintenance::drain_all(&daemon).await; + } + } + } +} diff --git a/bt-daemon/src/server/maintenance.rs b/bt-daemon/src/server/maintenance.rs index 401d89ec..1b97d79f 100644 --- a/bt-daemon/src/server/maintenance.rs +++ b/bt-daemon/src/server/maintenance.rs @@ -207,22 +207,20 @@ async fn retained_references(data_dir: &std::path::Path) -> anyhow::Result) { - let pending: HashSet<(String, String, u64)> = daemon + let pending: Vec = daemon .pending_sessions .lock() .values() - .flat_map(|state| { - state.events.iter().map(|event| { - ( - event.env.source.clone(), - event.env.session_id.clone(), - event.journal_through, - ) - }) - }) + .flat_map(|state| state.events.iter().cloned()) .collect(); let Ok(mut entries) = tokio::fs::read_dir(journal::journal_dir(&daemon.data_dir)).await else { return; @@ -118,11 +110,16 @@ pub(super) async fn recover_unprocessed_journals(daemon: &Arc) { } } for mut event in latest_by_route.into_values() { - if pending.contains(&( - event.env.source.clone(), - event.env.session_id.clone(), - event.journal_through, - )) { + if pending.iter().any(|pending| { + pending.env.source == event.env.source + && pending.env.session_id == event.env.session_id + && pending.journal_through == event.journal_through + && match (&pending.env.route, &event.env.route) { + (Some(a), Some(b)) => a.same_route(b), + (None, None) => true, + _ => false, + } + }) { continue; } let recovered_route = recovered_delivery_route(daemon, &event.env).await; @@ -202,3 +199,117 @@ pub(super) async fn recover_unprocessed_journals(daemon: &Arc) { ); } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::translate::{SpanOp, SpanRow}; + + struct TestAuth; + #[async_trait::async_trait] + impl AuthProvider for TestAuth { + async fn resolve( + &self, + selection: &AuthSelection, + _: AuthResolveReason, + ) -> anyhow::Result { + Ok(AuthLease { + selection: selection.clone(), + auth: BackendAuth { + token: "test".into(), + api_url: None, + app_url: None, + org_id: None, + org_name: None, + }, + expires_at_ms: None, + }) + } + } + + #[tokio::test] + async fn pending_capture_does_not_suppress_an_unrelated_durable_head() { + let tmp = tempfile::tempdir().unwrap(); + let route = SessionRoute { + destination: Some(crate::wire::TraceDestination::ProjectLogs { + project_id: Some("project".into()), + project_name: None, + }), + ..Default::default() + }; + let env: Envelope = serde_json::from_value(serde_json::json!({ + "source":"debug", "session_id":"test", "event":"event", "ts_ms":1, + "payload":{}, "route":route, + })) + .unwrap(); + let raw = journal::source_journal_path(tmp.path(), "debug", "test"); + let mut writer = JournalWriter::open_path(&raw).await.unwrap(); + let through = writer.append(&env).await.unwrap(); + drop(writer); + let path = tmp.path().join("pipeline").join(format!( + "{}.sqlite", + crate::ids::session_storage_id("debug", "test"), + )); + let spans = crate::pipeline::journal::SpanJournal::open(&path) + .await + .unwrap(); + spans.begin(through, true).await.unwrap(); + spans + .append( + through, + &[SpanOp::Insert(SpanRow { + span_id: "stored".into(), + name: "stored".into(), + root_span_id: "stored".into(), + end_ms: Some(1), + ..Default::default() + })], + ) + .await + .unwrap(); + spans.commit(through).await.unwrap(); + spans.subscribe(&route, 0).await.unwrap(); + let mut dormant = route.clone(); + dormant.tags.push("dormant".into()); + spans.subscribe(&dormant, 0).await.unwrap(); + + let mut options = debug_serve_options("test", tmp.path()); + options.auth_provider = Some(Arc::new(TestAuth)); + let daemon = Daemon::new(options, tmp.path().into()); + let pending_key = automatic_link_key(&env); + daemon.pending_sessions.lock().unwrap().insert( + pending_key.clone(), + PendingSession { + deadline: Some(Instant::now() + Duration::from_secs(60)), + events: vec![PendingEvent { + env, + replay_through: 0, + journal_through: through, + }], + ..Default::default() + }, + ); + recover_unprocessed_journals(&daemon).await; + daemon.ingress_barrier().await; + tokio::time::timeout(Duration::from_secs(5), async { + while !spans.route_complete(&dormant, through).await.unwrap() { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + assert!(!spans.route_complete(&route, through).await.unwrap()); + assert_eq!( + daemon.pending_sessions.lock().unwrap()[&pending_key] + .events + .len(), + 1 + ); + assert!(daemon + .sessions + .lock() + .unwrap() + .contains_key(&DeliveryKey::new("debug", "test", &dormant).unwrap())); + super::super::maintenance::drain_all(&daemon).await; + } +} From 32a1f3ccaede532d989669519893e8ee4f85157b Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Wed, 14 Oct 2026 21:40:48 +0800 Subject: [PATCH 4/5] Wait for shutdown receipt before disconnecting test client --- bt-daemon/tests/pipeline.rs | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/bt-daemon/tests/pipeline.rs b/bt-daemon/tests/pipeline.rs index 856d10f2..f75b8ff3 100644 --- a/bt-daemon/tests/pipeline.rs +++ b/bt-daemon/tests/pipeline.rs @@ -1750,7 +1750,24 @@ async fn shutdown_drains_then_retries_late_capture_on_replacement() { let shutdown_socket = socket.clone(); let shutdown_task = tokio::spawn(async move { shutdown_daemon(&shutdown_socket).await }); - tokio::time::sleep(Duration::from_millis(50)).await; + // Wait for proof that the server received shutdown and began draining. + // A fixed sleep can expire before the client sends its request on busy CI. + let journal_path = source_journal_path(&data_dir, "debug", "drained-shutdown"); + tokio::time::timeout(Duration::from_secs(5), async { + loop { + let journal = tokio::fs::read_to_string(&journal_path).await.unwrap(); + if journal.lines().any(|line| { + serde_json::from_str::(line).is_ok_and(|record| { + record["_bt_record_type"] == "pipeline_control" && record["kind"] == "retire" + }) + }) { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("daemon did not begin shutdown"); assert!( !shutdown_task.is_finished(), "shutdown returned before the sink drained" @@ -1770,7 +1787,7 @@ async fn shutdown_drains_then_retries_late_capture_on_replacement() { shutdown_task.abort(); let _ = shutdown_task.await; gate.notify_one(); - tokio::time::timeout(Duration::from_secs(2), handle) + tokio::time::timeout(Duration::from_secs(5), handle) .await .expect("daemon stayed alive after the shutdown client disconnected") .unwrap(); From d6894551815d033c0fe5eab99b1205eca012b566 Mon Sep 17 00:00:00 2001 From: Stephen Belanger Date: Sat, 17 Oct 2026 00:24:07 +0800 Subject: [PATCH 5/5] Align lifecycle recovery with fallible ledgers and panic bans --- bt-daemon/src/delivery_ledger.rs | 12 ++++++++---- bt-daemon/src/pipeline/consumer.rs | 4 ++-- bt-daemon/src/server/auth.rs | 2 +- bt-daemon/src/server/maintenance.rs | 19 +++++++------------ bt-daemon/src/server/recovery.rs | 22 ++++++++++------------ 5 files changed, 28 insertions(+), 31 deletions(-) diff --git a/bt-daemon/src/delivery_ledger.rs b/bt-daemon/src/delivery_ledger.rs index 84853439..8464e008 100644 --- a/bt-daemon/src/delivery_ledger.rs +++ b/bt-daemon/src/delivery_ledger.rs @@ -540,7 +540,8 @@ mod tests { Some(&config), "test", ) - .await; + .await + .unwrap(); sink.emit(&[terminal("one")]).await.unwrap(); let path = sink.ledger.as_ref().unwrap().path.clone(); let original = std::fs::read(&path).unwrap(); @@ -715,7 +716,8 @@ mod tests { Some(&config), "bt1", ) - .await; + .await + .unwrap(); sink.set_capture_versions(Some("v1"), Some("native1")); let path = sink.ledger.as_ref().unwrap().path.clone(); // Block persistence: nothing may reach the sink until origins are saved. @@ -750,7 +752,8 @@ mod tests { Some(&config), "bt2", ) - .await; + .await + .unwrap(); assert_eq!(recovered.span_origins.len(), 256); assert_eq!( recovered.span_origins["span-0"].plugin_version.as_deref(), @@ -815,7 +818,8 @@ mod tests { Some(&config), if is_child { "bt2" } else { "bt1" }, ) - .await; + .await + .unwrap(); sink.set_capture_versions( Some(if is_child { "v2" } else { "v1" }), Some(if is_child { "native2" } else { "native1" }), diff --git a/bt-daemon/src/pipeline/consumer.rs b/bt-daemon/src/pipeline/consumer.rs index 475a045f..850ff6eb 100644 --- a/bt-daemon/src/pipeline/consumer.rs +++ b/bt-daemon/src/pipeline/consumer.rs @@ -778,7 +778,7 @@ mod tests { journal.subscribe(&route, 0).await.unwrap(), opts, ); - wait_for(|| consumer.error.lock().unwrap().is_some()).await; + wait_for(|| consumer.error.lock().is_some()).await; assert_eq!(*consumer.progress.borrow(), 0); drop(rejection); Mock::given(method("POST")) @@ -795,7 +795,7 @@ mod tests { config.tags.push("corrected".into()); consumer.configure(config); wait_for(|| *consumer.progress.borrow() == 1).await; - assert!(consumer.error.lock().unwrap().is_none()); + assert!(consumer.error.lock().is_none()); consumer.shutdown().await; } diff --git a/bt-daemon/src/server/auth.rs b/bt-daemon/src/server/auth.rs index 5b047bf4..e1be30fe 100644 --- a/bt-daemon/src/server/auth.rs +++ b/bt-daemon/src/server/auth.rs @@ -120,7 +120,7 @@ impl Daemon { let provider = self .route_auth_provider(&key) - .expect("auth provider checked"); + .ok_or_else(|| anyhow::anyhow!("daemon host has no Braintrust auth provider"))?; let lease = resolve_route_auth( provider.as_ref(), &selection, diff --git a/bt-daemon/src/server/maintenance.rs b/bt-daemon/src/server/maintenance.rs index 1b97d79f..d9f1a5e0 100644 --- a/bt-daemon/src/server/maintenance.rs +++ b/bt-daemon/src/server/maintenance.rs @@ -73,8 +73,8 @@ async fn collect_partitions(daemon: &Daemon, age: Duration) { let lock = daemon.session_lock(&first.source, &first.session_id); let _guard = lock.lock().await; let namespace = crate::ids::session_namespace(&first.source, &first.session_id); - if daemon.journals.lock().unwrap().contains_key(&namespace) - || daemon.translations.lock().unwrap().contains_key(&namespace) + if daemon.journals.lock().contains_key(&namespace) + || daemon.translations.lock().contains_key(&namespace) { continue; } @@ -87,16 +87,11 @@ async fn collect_partitions(daemon: &Daemon, age: Duration) { let Ok(routes) = partition.routes().await else { continue; }; - let pending = daemon - .pending_sessions - .lock() - .unwrap() - .values() - .any(|state| { - state.events.iter().any(|event| { - event.env.source == first.source && event.env.session_id == first.session_id - }) - }); + let pending = daemon.pending_sessions.lock().values().any(|state| { + state.events.iter().any(|event| { + event.env.source == first.source && event.env.session_id == first.session_id + }) + }); if pending || !captured_routes_registered(&daemon.data_dir, &path, recorded, &routes).await { continue; diff --git a/bt-daemon/src/server/recovery.rs b/bt-daemon/src/server/recovery.rs index 92e3cd2f..5ecc63ef 100644 --- a/bt-daemon/src/server/recovery.rs +++ b/bt-daemon/src/server/recovery.rs @@ -87,6 +87,13 @@ pub(super) async fn recover_unprocessed_journals(daemon: &Arc) { continue; }; for route in routes { + let route_key = match serde_json::to_string(&route) { + Ok(key) => key, + Err(error) => { + tracing::warn!(%error, "durable route cannot be encoded for recovery"); + continue; + } + }; let matching = latest_by_route .iter() .find(|(_, event)| { @@ -103,10 +110,7 @@ pub(super) async fn recover_unprocessed_journals(daemon: &Arc) { .and_then(|key| latest_by_route.remove(&key)) .unwrap_or_else(|| template.clone()); event.env.route = Some(route.clone()); - latest_by_route.insert( - serde_json::to_string(&route).expect("persisted route"), - event, - ); + latest_by_route.insert(route_key, event); } } for mut event in latest_by_route.into_values() { @@ -277,7 +281,7 @@ mod tests { options.auth_provider = Some(Arc::new(TestAuth)); let daemon = Daemon::new(options, tmp.path().into()); let pending_key = automatic_link_key(&env); - daemon.pending_sessions.lock().unwrap().insert( + daemon.pending_sessions.lock().insert( pending_key.clone(), PendingSession { deadline: Some(Instant::now() + Duration::from_secs(60)), @@ -299,16 +303,10 @@ mod tests { .await .unwrap(); assert!(!spans.route_complete(&route, through).await.unwrap()); - assert_eq!( - daemon.pending_sessions.lock().unwrap()[&pending_key] - .events - .len(), - 1 - ); + assert_eq!(daemon.pending_sessions.lock()[&pending_key].events.len(), 1); assert!(daemon .sessions .lock() - .unwrap() .contains_key(&DeliveryKey::new("debug", "test", &dormant).unwrap())); super::super::maintenance::drain_all(&daemon).await; }