update metric

Author: Bhuvan506

kafka-streams-framework/src/main/java/org/hypertrace/core/kafkastreams/framework/punctuators/AbstractThrottledPunctuator.java

@@ -1,7 +1,7 @@
 package org.hypertrace.core.kafkastreams.framework.punctuators;
 
-import io.micrometer.core.instrument.Gauge;
 import io.micrometer.core.instrument.MeterRegistry;
+import io.micrometer.core.instrument.Tags;
 import java.time.Clock;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -10,6 +10,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicInteger;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.kafka.streams.KeyValue;
@@ -20,29 +21,24 @@
 
 @Slf4j
 public abstract class AbstractThrottledPunctuator<T> implements Punctuator {
+  private static final ConcurrentHashMap<String, AtomicInteger> totalEventCountGauge =
+      new ConcurrentHashMap<>();
+  private static final String TOTAL_EVENT_COUNT_GAUGE_NAME =
+      "abstract.throttled.punctuator.total.events.count";
   private final Clock clock;
   private final KeyValueStore<Long, List<T>> eventStore;
   private final ThrottledPunctuatorConfig config;
-  private final CustomMetrics customMetrics;
-
-  public AbstractThrottledPunctuator(
-      Clock clock, ThrottledPunctuatorConfig config, KeyValueStore<Long, List<T>> eventStore) {
-    this.clock = clock;
-    this.config = config;
-    this.eventStore = eventStore;
-    this.customMetrics = null;
-  }
+  private final MeterRegistry meterRegistry;
 
   public AbstractThrottledPunctuator(
       Clock clock,
       ThrottledPunctuatorConfig config,
       KeyValueStore<Long, List<T>> eventStore,
-      MeterRegistry meterRegistry,
-      String name) {
+      MeterRegistry meterRegistry) {
     this.clock = clock;
     this.config = config;
     this.eventStore = eventStore;
-    this.customMetrics = new CustomMetrics(meterRegistry, name);
+    this.meterRegistry = meterRegistry;
   }
 
   public void scheduleTask(long scheduleMs, T event) {
@@ -85,6 +81,7 @@ public final void punctuate(long timestamp) {
     long startTime = clock.millis();
     int totalProcessedWindows = 0;
     int totalProcessedTasks = 0;
+    int totalEventCount = 0;
 
     log.debug(
         "Processing tasks with throttling yield of {} until timestamp {}",
@@ -98,11 +95,11 @@ public final void punctuate(long timestamp) {
         totalProcessedWindows++;
         List<T> events = kv.value;
         long windowMs = kv.key;
+        totalEventCount += events.size();
         // collect all tasks to be rescheduled by key to perform bulk reschedules
         Map<Long, List<T>> rescheduledTasks = new HashMap<>();
         // loop through all events for this key until yield timeout is reached
         int i = 0;
-        int tasksBefore = totalProcessedTasks;
         for (; i < events.size() && !shouldYieldNow(startTime); i++) {
           T event = events.get(i);
           totalProcessedTasks++;
@@ -115,10 +112,6 @@ public final void punctuate(long timestamp) {
                           .computeIfAbsent(normalize(rescheduleTimestamp), (t) -> new ArrayList<>())
                           .add(event));
         }
-        if (customMetrics != null) {
-          customMetrics.setTaskProcessingStatus(
-              totalProcessedTasks > tasksBefore ? 0 : 1); // 0 if at least one task processed
-        }
         // process all reschedules
         rescheduledTasks.forEach(
             (newWindowMs, rescheduledEvents) -> {
@@ -146,14 +139,16 @@ public final void punctuate(long timestamp) {
         }
       }
     }
-    if (customMetrics != null) {
-      customMetrics.setEventStoreSize(currentTotalEventCount());
-    }
+    long timeTakenMs = clock.millis() - startTime;
+    boolean yielded = shouldYieldNow(startTime);
+    updateEventCountGauge(totalEventCount, yielded);
+
     log.debug(
-        "processed windows: {}, processed tasks: {}, time taken: {}",
+        "processed windows: {}, processed tasks: {}, total events: {}, time taken: {}ms",
         totalProcessedWindows,
         totalProcessedTasks,
-        clock.millis() - startTime);
+        totalEventCount,
+        timeTakenMs);
   }
 
   protected abstract TaskResult executeTask(long punctuateTimestamp, T object);
@@ -174,45 +169,22 @@ private long normalize(long timestamp) {
     return timestamp - (timestamp % config.getWindowMs());
   }
 
-  private int currentTotalEventCount() {
-    int count = 0;
-    try (KeyValueIterator<Long, List<T>> it = eventStore.all()) {
-      while (it.hasNext()) {
-        count += Optional.ofNullable(it.next().value).map(List::size).orElse(0);
-      }
+  private void updateEventCountGauge(int totalEventCount, boolean yielded) {
+    if (meterRegistry == null) {
+      return;
     }
-    return count;
-  }
 
-  protected final class CustomMetrics {
-    private final AtomicInteger eventStoreSize;
-    private final AtomicInteger noTasksProcessed;
-
-    public CustomMetrics(MeterRegistry meterRegistry, String name) {
-      this.noTasksProcessed = new AtomicInteger(1); // default to 1 = no tasks processed
-      this.eventStoreSize = new AtomicInteger(0);
-      if (meterRegistry != null) {
-        Gauge.builder("abstract.throttled.punctuator.task.size", eventStoreSize, AtomicInteger::get)
-            .description("Total number of scheduled tasks across all windows")
-            .tag("name", name)
-            .register(meterRegistry);
-        Gauge.builder(
-                "abstract.throttled.punctuator.no.tasks.processed",
-                noTasksProcessed,
-                AtomicInteger::get)
-            .description(
-                "1 if no tasks processed in last punctuate call, 0 if at least one processed")
-            .tag("name", name)
-            .register(meterRegistry);
-      }
-    }
+    String tagValue = String.valueOf(yielded);
 
-    public void setTaskProcessingStatus(int value) {
-      this.noTasksProcessed.set(value);
-    }
+    AtomicInteger gauge =
+        totalEventCountGauge.computeIfAbsent(
+            tagValue,
+            key -> {
+              AtomicInteger newGauge = new AtomicInteger(0);
+              meterRegistry.gauge(TOTAL_EVENT_COUNT_GAUGE_NAME, Tags.of("yielded", key), newGauge);
+              return newGauge;
+            });
 
-    public void setEventStoreSize(int value) {
-      this.eventStoreSize.set(value);
-    }
+    gauge.set(totalEventCount);
   }
 }

kafka-streams-framework/src/test/java/org/hypertrace/core/kafkastreams/framework/punctuators/AbstractThrottledPunctuatorTest.java

@@ -194,7 +194,7 @@ public TestPunctuator(
         Clock clock,
         ThrottledPunctuatorConfig config,
         KeyValueStore<Long, List<String>> objectStore) {
-      super(clock, config, objectStore);
+      super(clock, config, objectStore, null);
     }
 
     void setReturnResult(String object, TaskResult result) {