Wire stats aggregation, fix batch doc counting, extract publish_to_sinks (#118)

Author: makarski

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "mstream"
-version = "0.44.0"
+version = "0.45.0"
 edition = "2024"
 authors = ["Simon Makarski"]
 

src/api/handler/system.rs

@@ -46,13 +46,13 @@ pub async fn stats(State(state): State<AppState>) -> Result<impl IntoResponse, A
 
     let uptime = state.start_time.elapsed().as_secs();
 
-    // TODO: wire total_docs_processed / total_bytes_transferred once
-    // in-engine metrics collection is implemented
+    let (total_events, total_bytes, _) = jm.aggregate_metrics();
+
     Ok((
         StatusCode::OK,
         Json(SystemStats {
-            total_docs_processed: 0,
-            total_bytes_transferred: 0,
+            total_events_processed: total_events,
+            total_bytes_processed: total_bytes,
             uptime_seconds: uptime,
             running_jobs: counts.running,
             stopped_jobs: counts.stopped,

src/api/types.rs

@@ -310,8 +310,8 @@ pub struct HealthStatus {
 
 #[derive(Serialize)]
 pub struct SystemStats {
-    pub total_docs_processed: u64,
-    pub total_bytes_transferred: u64,
+    pub total_events_processed: u64,
+    pub total_bytes_processed: u64,
     pub uptime_seconds: u64,
     pub running_jobs: usize,
     pub stopped_jobs: usize,
@@ -688,15 +688,15 @@ mod tests {
         fn system_stats_serializes() {
             assert_serializes_to(
                 &SystemStats {
-                    total_docs_processed: 1000,
-                    total_bytes_transferred: 50000,
+                    total_events_processed: 1000,
+                    total_bytes_processed: 50000,
                     uptime_seconds: 3600,
                     running_jobs: 3,
                     stopped_jobs: 1,
                     error_jobs: 2,
                 },
                 json!({
-                    "total_docs_processed": 1000, "total_bytes_transferred": 50000,
+                    "total_events_processed": 1000, "total_bytes_processed": 50000,
                     "uptime_seconds": 3600, "running_jobs": 3, "stopped_jobs": 1, "error_jobs": 2,
                 }),
             );

src/job_manager/mod.rs

@@ -66,8 +66,9 @@ impl JobMetricsCounter {
         }
     }
 
-    pub fn record_success(&self, bytes: u64) {
-        self.events_processed.fetch_add(1, Ordering::Relaxed);
+    pub fn record_success(&self, doc_count: u64, bytes: u64) {
+        self.events_processed
+            .fetch_add(doc_count, Ordering::Relaxed);
         self.bytes_processed.fetch_add(bytes, Ordering::Relaxed);
         self.last_processed_at
             .store(chrono::Utc::now().timestamp_millis(), Ordering::Relaxed);
@@ -259,6 +260,18 @@ impl JobManager {
         self.running_jobs.get(name).map(|jc| &jc.metrics)
     }
 
+    /// Sums metrics across all running jobs.
+    pub fn aggregate_metrics(&self) -> (u64, u64, u64) {
+        self.running_jobs.values().fold((0, 0, 0), |acc, jc| {
+            let m = &jc.metrics;
+            (
+                acc.0 + m.events_processed.load(Ordering::Relaxed),
+                acc.1 + m.bytes_processed.load(Ordering::Relaxed),
+                acc.2 + m.errors.load(Ordering::Relaxed),
+            )
+        })
+    }
+
     /// Returns job counts grouped by state.
     pub async fn job_state_counts(&self) -> anyhow::Result<JobStateCounts> {
         let jobs = self.job_store.list_all().await?;
@@ -686,8 +699,8 @@ mod tests {
     #[test]
     fn metrics_counter_record_success() {
         let m = JobMetricsCounter::new();
-        m.record_success(256);
-        m.record_success(512);
+        m.record_success(1, 256);
+        m.record_success(1, 512);
 
         assert_eq!(m.events_processed.load(Ordering::Relaxed), 2);
         assert_eq!(m.bytes_processed.load(Ordering::Relaxed), 768);
@@ -709,13 +722,22 @@ mod tests {
     #[test]
     fn metrics_counter_mixed_success_and_errors() {
         let m = JobMetricsCounter::new();
-        m.record_success(100);
+        m.record_success(1, 100);
         m.record_error();
-        m.record_success(200);
+        m.record_success(1, 200);
         m.record_error();
 
         assert_eq!(m.events_processed.load(Ordering::Relaxed), 2);
         assert_eq!(m.bytes_processed.load(Ordering::Relaxed), 300);
         assert_eq!(m.errors.load(Ordering::Relaxed), 2);
     }
+
+    #[test]
+    fn metrics_counter_batch_counts_all_docs() {
+        let m = JobMetricsCounter::new();
+        m.record_success(100, 5000);
+
+        assert_eq!(m.events_processed.load(Ordering::Relaxed), 100);
+        assert_eq!(m.bytes_processed.load(Ordering::Relaxed), 5000);
+    }
 }

src/provision/pipeline/processor.rs

@@ -66,7 +66,7 @@ impl EventHandler {
                             .map_err(|err| anyhow!("failed to apply source schema: {}", err))
                     })?;
 
-                    if let Err(err) = self.process_event(event).await {
+                    if let Err(err) = self.process_event(event, 1).await {
                         error!(job_name = %self.pipeline.name, "failed to process event: {}", err)
                     }
                 }
@@ -127,6 +127,7 @@ impl EventHandler {
         let source_encoding = batch[0].encoding.clone();
         let attributes = batch[0].attributes.clone();
         let last_cursor = batch.last().and_then(|event| event.cursor.clone());
+        let doc_count = batch.len() as u64;
         let payloads: Vec<Vec<u8>> = batch.drain(..).map(|event| event.raw_bytes).collect();
 
         let batch_event = block_in_place(|| {
@@ -150,37 +151,55 @@ impl EventHandler {
             batch_event.raw_bytes.len()
         );
 
-        if let Err(err) = self.process_event(batch_event).await {
+        if let Err(err) = self.process_event(batch_event, doc_count).await {
             error!(job_name = %self.pipeline.name, "failed to process event: {}", err)
         }
 
         Ok(())
     }
 
-    async fn process_event(&mut self, source_event: SourceEvent) -> anyhow::Result<()> {
+    async fn process_event(
+        &mut self,
+        source_event: SourceEvent,
+        doc_count: u64,
+    ) -> anyhow::Result<()> {
         let event_bytes = source_event.raw_bytes.len() as u64;
-        let event_holder = self.apply_middlewares(source_event).await?;
-        let mut all_sinks_succeeded = true;
 
-        // Extract cursor before the loop for checkpointing
+        let event_holder = match self.apply_middlewares(source_event).await {
+            Ok(holder) => holder,
+            Err(err) => {
+                self.record_event_outcome(false, doc_count, event_bytes);
+                return Err(err);
+            }
+        };
+
         let cursor = event_holder.cursor.clone();
+        let all_sinks_succeeded = self.publish_to_sinks(event_holder).await;
+
+        self.record_event_outcome(all_sinks_succeeded, doc_count, event_bytes);
+
+        if self.pipeline.with_checkpoints && all_sinks_succeeded {
+            self.save_checkpoint(cursor).await?;
+        }
 
+        Ok(())
+    }
+
+    async fn publish_to_sinks(&mut self, event_holder: SourceEvent) -> bool {
+        let mut all_succeeded = true;
         let mut event_option = Some(event_holder);
         let sinks_len = self.pipeline.sinks.len();
 
         for (i, sink_def) in self.pipeline.sinks.iter_mut().enumerate() {
             let sink_event_result = block_in_place(|| {
                 let event = if i == sinks_len - 1 {
-                    // Move the event on the last iteration to avoid clone
                     event_option.take().expect("event should be present")
                 } else {
-                    // Clone for intermediate sinks
                     event_option
                         .as_ref()
                         .expect("event should be present")
                         .clone()
                 };
-                // apply sink schema
                 event.apply_schema(Some(&sink_def.config.output_encoding), &sink_def.schema)
             });
 
@@ -193,16 +212,13 @@ impl EventHandler {
                     error!(
                         job_name = %self.pipeline.name,
                         "failed to encode for sink {}/{}: {}",
-                        sink_def.config.service_name,
-                        sink_def.config.resource,
-                        err
+                        sink_def.config.service_name, sink_def.config.resource, err
                     );
-                    all_sinks_succeeded = false;
+                    all_succeeded = false;
                     continue;
                 }
             };
 
-            // maybe we need a config in publisher, ie explode batch
             match sink_def
                 .sink
                 .publish(sink_event, sink_def.config.resource.clone(), None)
@@ -212,38 +228,28 @@ impl EventHandler {
                     info!(
                         job_name = %self.pipeline.name,
                         "published to {}/{}: {}",
-                        sink_def.config.service_name,
-                        sink_def.config.resource,
-                        message
+                        sink_def.config.service_name, sink_def.config.resource, message
                     );
                 }
                 Err(err) => {
                     error!(
                         job_name = %self.pipeline.name,
                         "failed to publish to {}/{}: {:#}",
-                        sink_def.config.service_name,
-                        sink_def.config.resource,
-                        err
+                        sink_def.config.service_name, sink_def.config.resource, err
                     );
-                    all_sinks_succeeded = false;
+                    all_succeeded = false;
                     continue;
                 }
             };
         }
 
-        self.record_event_outcome(all_sinks_succeeded, event_bytes);
-
-        if self.pipeline.with_checkpoints && all_sinks_succeeded {
-            self.save_checkpoint(cursor).await?;
-        }
-
-        Ok(())
+        all_succeeded
     }
 
-    fn record_event_outcome(&self, success: bool, bytes: u64) {
+    fn record_event_outcome(&self, success: bool, doc_count: u64, bytes: u64) {
         if let Some(ref m) = self.metrics {
             if success {
-                m.record_success(bytes);
+                m.record_success(doc_count, bytes);
             } else {
                 m.record_error();
             }