fix(digest): drain queue and flush on channel close; graceful shutdown

crates/trapfalld/src/digest.rs

@@ -45,25 +45,31 @@ 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() => {
                     if !buffer.is_empty() {
                         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);
+    }
 }

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");
+}