diff --git a/crates/trapfalld/src/digest.rs b/crates/trapfalld/src/digest.rs index cb25a50..a5d2919 100644 --- a/crates/trapfalld/src/digest.rs +++ b/crates/trapfalld/src/digest.rs @@ -45,11 +45,24 @@ impl DigestTask { loop { tokio::select! { - Some(event) = self.rx.recv() => { - tracing::debug!("Digest received event: project={}", event.project_id); - buffer.push(event); - if buffer.len() >= FLUSH_THRESHOLD { - Self::flush(&store, &mut buffer, &self.ws_tx, &self.alert_tx).await; + maybe_event = self.rx.recv() => { + match maybe_event { + Some(event) => { + tracing::debug!("Digest received event: project={}", event.project_id); + buffer.push(event); + if buffer.len() >= FLUSH_THRESHOLD { + Self::flush(&store, &mut buffer, &self.ws_tx, &self.alert_tx).await; + } + } + None => { + // All senders dropped and the queue is fully drained — + // flush the partial batch and stop. + tracing::info!("Digest channel closed, flushing {} buffered events", buffer.len()); + if !buffer.is_empty() { + Self::flush(&store, &mut buffer, &self.ws_tx, &self.alert_tx).await; + } + break; + } } } _ = interval.tick() => { @@ -57,13 +70,6 @@ impl DigestTask { Self::flush(&store, &mut buffer, &self.ws_tx, &self.alert_tx).await; } } - else => { - // Channel closed — drain remaining - if !buffer.is_empty() { - Self::flush(&store, &mut buffer, &self.ws_tx, &self.alert_tx).await; - } - break; - } } } @@ -267,4 +273,83 @@ mod tests { }; assert_eq!(derive_title(&ev), "Unknown Error"); } + + fn make_event(project_id: &str, n: usize) -> IngestEvent { + use trapfall_proto::*; + IngestEvent { + project_id: project_id.into(), + fingerprint: Fingerprint::new("drain-hash".into(), vec![]), + event: Event { + event_id: format!("e{n}"), + level: Level::Error, + platform: None, + release: None, + environment: None, + server_name: None, + breadcrumbs: Breadcrumbs::default(), + exception: None, + message: Some("drain test".into()), + tags: serde_json::Value::Null, + extra: serde_json::Value::Null, + contexts: serde_json::Value::Null, + timestamp: None, + }, + received_at: "2026-01-01T00:00:00Z".into(), + } + } + + async fn test_store() -> Store { + let backend = trapfall_db::open_database("sqlite::memory:").await.unwrap(); + trapfall_db::run_sqlite_migrations(backend.sqlite_pool().unwrap()).await.unwrap(); + Store::new(backend) + } + + /// Events queued (fewer than the flush threshold) when the channel closes + /// must all be persisted and the task must exit on its own. + #[tokio::test] + async fn drains_partial_batch_and_exits_on_channel_close() { + let store = test_store().await; + let project = store.create_project("drain", "Drain").await.unwrap(); + let (tx, rx) = mpsc::channel(256); + let handle = tokio::spawn(DigestTask::new(store.clone(), rx).run()); + + let n = 7; + for i in 0..n { + tx.send(make_event(&project.id, i)).await.unwrap(); + } + drop(tx); + + tokio::time::timeout(std::time::Duration::from_secs(1), handle) + .await + .expect("digest task must exit after channel close") + .unwrap() + .unwrap(); + + let issues = store.list_issues(&project.id, 10, 0).await.unwrap(); + assert_eq!(issues.len(), 1); + assert_eq!(store.count_events(&issues[0].id).await.unwrap(), n as i64); + } + + /// More events than the flush threshold, plus a remainder, all persisted. + #[tokio::test] + async fn drains_multiple_batches_on_channel_close() { + let store = test_store().await; + let project = store.create_project("drain2", "Drain2").await.unwrap(); + let (tx, rx) = mpsc::channel(256); + let n = FLUSH_THRESHOLD * 2 + 3; + for i in 0..n { + tx.send(make_event(&project.id, i)).await.unwrap(); + } + drop(tx); + + let handle = tokio::spawn(DigestTask::new(store.clone(), rx).run()); + tokio::time::timeout(std::time::Duration::from_secs(5), handle) + .await + .expect("digest task must exit after channel close") + .unwrap() + .unwrap(); + + let issues = store.list_issues(&project.id, 10, 0).await.unwrap(); + assert_eq!(store.count_events(&issues[0].id).await.unwrap(), n as i64); + } } diff --git a/crates/trapfalld/src/main.rs b/crates/trapfalld/src/main.rs index 62319ec..f25c13b 100644 --- a/crates/trapfalld/src/main.rs +++ b/crates/trapfalld/src/main.rs @@ -248,10 +248,79 @@ async fn run_server(store: Store, listen: String, db_url: String) -> Result<()> // HTTP server let listener = tokio::net::TcpListener::bind(&listen).await?; info!("Listening on {listen}"); - axum::serve(listener, trapfalld::server::router(state)).await?; + // Serving the router consumes `state` (and with it the only `ingest_tx`), + // so once graceful shutdown completes the digest channel is closed and the + // digest task drains its queue and flushes the final batch before exiting. + // Open connections (e.g. long-lived WebSockets) can stall graceful shutdown, + // so give them a bounded grace period and then drop the server regardless. + let stop = std::sync::Arc::new(tokio::sync::Notify::new()); + let stop_rx = stop.clone(); + let serve = axum::serve(listener, trapfalld::server::router(state)).with_graceful_shutdown(async move { + stop_rx.notified().await; + }); + let mut server = Box::pin(std::future::IntoFuture::into_future(serve)); + tokio::select! { + res = &mut server => { + if let Err(e) = res { + tracing::error!("HTTP server error: {e}"); + } + } + _ = async { + shutdown_signal().await; + stop.notify_one(); + tokio::time::sleep(HTTP_GRACE_TIMEOUT).await; + } => tracing::warn!("HTTP graceful shutdown timed out after {HTTP_GRACE_TIMEOUT:?}; dropping open connections"), + } + // Drop the server (and the router state holding `ingest_tx`) so the digest channel closes. + drop(server); - digest_handle.abort(); + info!("HTTP server stopped, draining digest pipeline"); retention_handle.abort(); + match tokio::time::timeout(DIGEST_DRAIN_TIMEOUT, digest_handle).await { + Ok(Ok(())) => info!("Digest pipeline drained"), + Ok(Err(e)) => tracing::error!("Digest task join error: {e}"), + Err(_) => { + tracing::error!("Digest drain timed out after {DIGEST_DRAIN_TIMEOUT:?}; remaining events may be lost") + } + } + // The digest task owned the ws sender; once it exits the bridge ends on its own. bridge_handle.abort(); Ok(()) } + +/// Upper bound on how long shutdown waits for open HTTP connections to finish. +const HTTP_GRACE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); + +/// Upper bound on how long shutdown waits for the digest task to flush. +const DIGEST_DRAIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10); + +/// Resolves when SIGINT (Ctrl-C) or SIGTERM is received. +async fn shutdown_signal() { + let ctrl_c = async { + if let Err(e) = tokio::signal::ctrl_c().await { + tracing::error!("Failed to install Ctrl-C handler: {e}"); + std::future::pending::<()>().await; + } + }; + + #[cfg(unix)] + let terminate = async { + match tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) { + Ok(mut sig) => { + sig.recv().await; + } + Err(e) => { + tracing::error!("Failed to install SIGTERM handler: {e}"); + std::future::pending::<()>().await; + } + } + }; + #[cfg(not(unix))] + let terminate = std::future::pending::<()>(); + + tokio::select! { + _ = ctrl_c => {}, + _ = terminate => {}, + } + info!("Shutdown signal received"); +}