Merge pull request #2297 from stonewhitener/terminate

Flush pending metric snapshots at activity completion

nb-apis/nb-api/src/main/java/io/nosqlbench/nb/api/engine/metrics/MetricsSnapshotScheduler.java

@@ -261,6 +261,12 @@ public final class MetricsSnapshotScheduler extends UnstartedPeriodicTaskCompone
 
     @Override
     protected void task() {
+        synchronized (scheduleLock) {
+            captureSnapshotLocked();
+        }
+    }
+
+    private void captureSnapshotLocked() {
         if (consumerIntervals.isEmpty()) {
             return;
         }
@@ -271,7 +277,10 @@ public final class MetricsSnapshotScheduler extends UnstartedPeriodicTaskCompone
         boolean includeHdrPayload = consumerIntervals.keySet().stream()
             .anyMatch(MetricsSnapshotConsumer::requiresHdrPayload);
         MetricsView snapshot = MetricsView.capture(metrics, baseIntervalMillis, includeHdrPayload);
-        processSnapshot(snapshot);
+        ScheduleNode base = scheduleNodes.get(baseIntervalMillis);
+        if (base != null) {
+            base.ingest(snapshot);
+        }
     }
 
     public void injectSnapshotForTesting(MetricsView view) {
@@ -285,29 +294,19 @@ public final class MetricsSnapshotScheduler extends UnstartedPeriodicTaskCompone
     /// Without this, short runs (where no periodic tick fires before completion) leave
     /// activity-level metrics out of the SQLite database entirely.
     ///
-    /// Unlike a periodic tick, this **bypasses the per-consumer interval bucketing** and
-    /// emits the captured view directly to every registered consumer. Otherwise a single
-    /// trigger would only contribute the base-interval amount toward each consumer's
-    /// accumulated time and consumers with longer intervals (e.g. SQLite at 30s vs CSV at
-    /// 5s) would still wait for their normal bucket boundary — defeating the point of the
-    /// final flush.
+    /// The final view is ingested through the normal cadence hierarchy, then every partial
+    /// bucket is flushed recursively. This preserves snapshots which were already pending
+    /// for slower consumers (e.g. SQLite at 30s vs CSV at 5s), while ensuring that the newly
+    /// captured delta is delivered exactly once.
     public void triggerSnapshot() {
-        if (consumerIntervals.isEmpty()) {
-            return;
-        }
-        List<NBMetric> metrics = new ArrayList<>(getParent().find().metrics());
-        if (metrics.isEmpty()) {
-            return;
-        }
-        boolean includeHdrPayload = consumerIntervals.keySet().stream()
-            .anyMatch(MetricsSnapshotConsumer::requiresHdrPayload);
-        MetricsView snapshot = MetricsView.capture(metrics, baseIntervalMillis, includeHdrPayload);
-        for (MetricsSnapshotConsumer consumer : consumerIntervals.keySet()) {
+        synchronized (scheduleLock) {
             try {
-                consumer.onMetricsSnapshot(snapshot);
-            } catch (Exception e) {
-                logger.warn("Final-snapshot consumer {} threw exception",
-                    consumer.getClass().getSimpleName(), e);
+                captureSnapshotLocked();
+            } finally {
+                ScheduleNode base = scheduleNodes.get(baseIntervalMillis);
+                if (base != null) {
+                    base.flushPending();
+                }
             }
         }
     }
@@ -382,6 +381,21 @@ public final class MetricsSnapshotScheduler extends UnstartedPeriodicTaskCompone
             }
         }
 
+        private void flushPending() {
+            if (pending != null) {
+                MetricsView ready = pending;
+                pending = null;
+                accumulatedMillis = 0L;
+                emit(ready);
+                for (ScheduleNode child : children) {
+                    child.ingest(ready);
+                }
+            }
+            for (ScheduleNode child : children) {
+                child.flushPending();
+            }
+        }
+
         private void emit(MetricsView view) {
             for (MetricsSnapshotConsumer consumer : consumers) {
                 try {

nb-apis/nb-api/src/test/java/io/nosqlbench/nb/api/engine/metrics/MetricsSnapshotSchedulerTest.java

@@ -21,6 +21,7 @@ package io.nosqlbench.nb.api.engine.metrics;
 import io.nosqlbench.nb.api.components.core.NBBaseComponent;
 import io.nosqlbench.nb.api.components.core.NBComponent;
 import io.nosqlbench.nb.api.engine.metrics.instruments.MetricCategory;
+import io.nosqlbench.nb.api.engine.metrics.instruments.NBFunctionGauge;
 import io.nosqlbench.nb.api.engine.metrics.instruments.NBMetricCounter;
 import io.nosqlbench.nb.api.engine.metrics.instruments.NBMetricHistogram;
 import io.nosqlbench.nb.api.engine.metrics.view.MetricsView;
@@ -223,6 +224,112 @@ public class MetricsSnapshotSchedulerTest {
         assertThat(payload).isInstanceOf(Histogram.class);
     }
 
+    @Test
+    public void testTriggerSnapshotFlushesPendingHierarchicalSnapshots() {
+        long baseInterval = 100_000L;
+        RecordingHdrConsumer baseConsumer = new RecordingHdrConsumer();
+        RecordingHdrConsumer mediumConsumer = new RecordingHdrConsumer();
+        RecordingHdrConsumer coarseConsumer = new RecordingHdrConsumer();
+
+        NBLabels labels = NBLabels.forKV("name", "hist_metric", "scenario", "scenario", "activity", "activity");
+        DeltaHdrHistogramReservoir reservoir = new DeltaHdrHistogramReservoir(labels, 3);
+        NBMetricHistogram liveHistogram = new NBMetricHistogram(
+            labels,
+            reservoir,
+            "hist",
+            "nanoseconds",
+            MetricCategory.Core
+        );
+        root.addComponentMetric(liveHistogram, MetricCategory.Core, "hist");
+
+        MetricsSnapshotScheduler.register(root, baseInterval, baseConsumer);
+        MetricsSnapshotScheduler.register(root, baseInterval * 2L, mediumConsumer);
+        MetricsSnapshotScheduler.register(root, baseInterval * 4L, coarseConsumer);
+        scheduler = MetricsSnapshotScheduler.lookup(root);
+
+        scheduler.injectSnapshotForTesting(histogramView(baseInterval, 10L, 20L));
+        scheduler.injectSnapshotForTesting(histogramView(baseInterval, 30L));
+        liveHistogram.update(40L);
+        liveHistogram.update(50L);
+
+        scheduler.triggerSnapshot();
+
+        assertThat(baseConsumer.snapshots).hasSize(3);
+        assertThat(histogramCount(baseConsumer.snapshots.get(0))).isEqualTo(2L);
+        assertThat(histogramCount(baseConsumer.snapshots.get(1))).isEqualTo(1L);
+        assertThat(histogramCount(baseConsumer.snapshots.get(2))).isEqualTo(2L);
+
+        assertThat(mediumConsumer.snapshots).hasSize(2);
+        assertThat(histogramCount(mediumConsumer.snapshots.get(0))).isEqualTo(3L);
+        assertThat(mediumConsumer.snapshots.get(0).intervalMillis()).isEqualTo(baseInterval * 2L);
+        assertThat(histogramCount(mediumConsumer.snapshots.get(1))).isEqualTo(2L);
+        assertThat(mediumConsumer.snapshots.get(1).intervalMillis()).isEqualTo(baseInterval);
+
+        assertThat(coarseConsumer.snapshots).hasSize(1);
+        assertThat(histogramCount(coarseConsumer.snapshots.getFirst())).isEqualTo(5L);
+        assertThat(coarseConsumer.snapshots.getFirst().intervalMillis()).isEqualTo(baseInterval * 3L);
+    }
+
+    @Test
+    public void testTriggerSnapshotFlushesPendingSnapshotsWhenCaptureFails() {
+        long baseInterval = 100_000L;
+        RecordingHdrConsumer baseConsumer = new RecordingHdrConsumer();
+        RecordingHdrConsumer coarseConsumer = new RecordingHdrConsumer();
+
+        MetricsSnapshotScheduler.register(root, baseInterval, baseConsumer);
+        MetricsSnapshotScheduler.register(root, baseInterval * 2L, coarseConsumer);
+        scheduler = MetricsSnapshotScheduler.lookup(root);
+
+        scheduler.injectSnapshotForTesting(histogramView(baseInterval, 10L));
+        NBFunctionGauge failingGauge = new NBFunctionGauge(
+            root,
+            () -> {
+                throw new IllegalStateException("capture failed");
+            },
+            "failing_gauge",
+            "failing gauge",
+            "",
+            MetricCategory.Core
+        );
+        root.addComponentMetric(failingGauge, MetricCategory.Core, "failing gauge");
+
+        assertThatThrownBy(scheduler::triggerSnapshot)
+            .isInstanceOf(IllegalStateException.class)
+            .hasMessage("capture failed");
+
+        assertThat(baseConsumer.snapshots).hasSize(1);
+        assertThat(coarseConsumer.snapshots).hasSize(1);
+        assertThat(histogramCount(coarseConsumer.snapshots.getFirst())).isEqualTo(1L);
+        assertThat(coarseConsumer.snapshots.getFirst().intervalMillis()).isEqualTo(baseInterval);
+    }
+
+    private MetricsView histogramView(long interval, long... values) {
+        NBLabels labels = NBLabels.forKV("name", "hist_metric", "scenario", "scenario", "activity", "activity");
+        DeltaHdrHistogramReservoir reservoir = new DeltaHdrHistogramReservoir(labels, 3);
+        NBMetricHistogram histogram = new NBMetricHistogram(
+            labels,
+            reservoir,
+            "hist",
+            "nanoseconds",
+            MetricCategory.Core
+        );
+        for (long value : values) {
+            histogram.update(value);
+        }
+        return MetricsView.capture(List.of(histogram), interval, true);
+    }
+
+    private long histogramCount(MetricsView view) {
+        return view.families().stream()
+            .filter(family -> family.familyName().equals("hist_metric"))
+            .flatMap(family -> family.samples().stream())
+            .filter(MetricsView.SummarySample.class::isInstance)
+            .map(MetricsView.SummarySample.class::cast)
+            .mapToLong(sample -> sample.statistics().count())
+            .findFirst()
+            .orElseThrow();
+    }
+
     private static final class RecordingReporter extends MetricsSnapshotReporterBase {
         private final List<MetricsView> snapshots = new ArrayList<>();