feat(s3stream): record accumulated value in CounterMetric (#2047)

Signed-off-by: Shichao Nie <[email protected]>
AutoMQ · Oct 9, 2024 · f7ffd88 · f7ffd88
1 parent de557c8
commit f7ffd88
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 5 deletions.
diff --git a/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/CounterMetric.java b/s3stream/src/main/java/com/automq/stream/s3/metrics/wrapper/CounterMetric.java
@@ -14,13 +14,15 @@
 import com.automq.stream.s3.metrics.MetricsConfig;
 import com.automq.stream.s3.metrics.MetricsLevel;
 
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Supplier;
 
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.metrics.LongCounter;
 
 public class CounterMetric extends ConfigurableMetric {
     private final Supplier<LongCounter> longCounterSupplier;
+    private final AtomicLong total = new AtomicLong(0);
 
     public CounterMetric(MetricsConfig metricsConfig, Supplier<LongCounter> longCounterSupplier) {
         super(metricsConfig, Attributes.empty());
@@ -33,10 +35,15 @@ public CounterMetric(MetricsConfig metricsConfig, Attributes extraAttributes, Su
     }
 
     public boolean add(MetricsLevel metricsLevel, long value) {
+        total.addAndGet(value);
         if (metricsLevel.isWithin(this.metricsLevel)) {
             longCounterSupplier.get().add(value, attributes);
             return true;
         }
         return false;
     }
+
+    public long getValue() {
+        return total.get();
+    }
 }
diff --git a/s3stream/src/main/java/com/automq/stream/s3/network/AsyncNetworkBandwidthLimiter.java b/s3stream/src/main/java/com/automq/stream/s3/network/AsyncNetworkBandwidthLimiter.java
@@ -146,7 +146,7 @@ private void forceConsume(long size) {
 
     public CompletableFuture<Void> consume(ThrottleStrategy throttleStrategy, long size) {
         CompletableFuture<Void> cf = new CompletableFuture<>();
-        cf.whenComplete((v, e) -> logMetrics(size, throttleStrategy));
+        cf.whenComplete((v, e) -> NetworkStats.getInstance().networkUsageTotalStats(type, throttleStrategy).add(MetricsLevel.INFO, size));
         if (Objects.requireNonNull(throttleStrategy) == ThrottleStrategy.BYPASS) {
             forceConsume(size);
             cf.complete(null);
@@ -171,10 +171,6 @@ private void reduceToken(long size) {
         this.availableTokens = Math.max(-maxTokens, availableTokens - size);
     }
 
-    private void logMetrics(long size, ThrottleStrategy strategy) {
-        NetworkStats.getInstance().networkUsageTotalStats(type, strategy).add(MetricsLevel.INFO, size);
-    }
-
     public enum Type {
         INBOUND("Inbound"),
         OUTBOUND("Outbound");