[ISSUE #9816] Fix calculate consumer lag with opentelemetry (#9873)

Author: lizhimins

broker/src/main/java/org/apache/rocketmq/broker/longpolling/PopCommandCallback.java

@@ -17,33 +17,31 @@
 
 package org.apache.rocketmq.broker.longpolling;
 
+import java.util.concurrent.CompletableFuture;
 import java.util.function.BiConsumer;
-import java.util.function.Consumer;
 import org.apache.rocketmq.broker.metrics.ConsumerLagCalculator;
 import org.apache.rocketmq.remoting.CommandCallback;
 
 public class PopCommandCallback implements CommandCallback {
 
     private final BiConsumer<ConsumerLagCalculator.ProcessGroupInfo,
-        Consumer<ConsumerLagCalculator.CalculateLagResult>> biConsumer;
-
+        CompletableFuture<ConsumerLagCalculator.CalculateLagResult>> biConsumer;
     private final ConsumerLagCalculator.ProcessGroupInfo info;
-    private final Consumer<ConsumerLagCalculator.CalculateLagResult> lagRecorder;
-
+    private final CompletableFuture<ConsumerLagCalculator.CalculateLagResult> future;
 
     public PopCommandCallback(
         BiConsumer<ConsumerLagCalculator.ProcessGroupInfo,
-                    Consumer<ConsumerLagCalculator.CalculateLagResult>> biConsumer,
+            CompletableFuture<ConsumerLagCalculator.CalculateLagResult>> biConsumer,
         ConsumerLagCalculator.ProcessGroupInfo info,
-        Consumer<ConsumerLagCalculator.CalculateLagResult> lagRecorder) {
+        CompletableFuture<ConsumerLagCalculator.CalculateLagResult> future) {
 
         this.biConsumer = biConsumer;
         this.info = info;
-        this.lagRecorder = lagRecorder;
+        this.future = future;
     }
 
     @Override
     public void accept() {
-        biConsumer.accept(info, lagRecorder);
+        biConsumer.accept(info, future);
     }
 }

broker/src/main/java/org/apache/rocketmq/broker/metrics/BrokerMetricsManager.java

@@ -673,14 +673,8 @@ private void initLagAndDlqMetrics() {
         consumerLagMessages = brokerMeter.gaugeBuilder(GAUGE_CONSUMER_LAG_MESSAGES)
             .setDescription("Consumer lag messages")
             .ofLongs()
-            .buildWithCallback(measurement ->
-                consumerLagCalculator.calculateLag(result -> {
-                    // Note: 'record' method uses HashMap which may cause
-                    // concurrent access issues when Pull thread executes Pop callbacks.
-                    synchronized (this) {
-                        measurement.record(result.lag, buildLagAttributes(result));
-                    }
-                }));
+            .buildWithCallback(measurement -> consumerLagCalculator.calculateLag(result ->
+                measurement.record(result.lag, buildLagAttributes(result))));
 
         consumerLagLatency = brokerMeter.gaugeBuilder(GAUGE_CONSUMER_LAG_LATENCY)
             .setDescription("Consumer lag time")

broker/src/main/java/org/apache/rocketmq/broker/metrics/ConsumerLagCalculator.java

@@ -16,8 +16,13 @@
  */
 package org.apache.rocketmq.broker.metrics;
 
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.TimeUnit;
+import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 import org.apache.rocketmq.broker.BrokerController;
@@ -211,20 +216,43 @@ private void processAllGroup(Consumer<ProcessGroupInfo> consumer) {
     }
 
     public void calculateLag(Consumer<CalculateLagResult> lagRecorder) {
+
+        List<CompletableFuture<CalculateLagResult>> futures = new ArrayList<>();
+
+        BiConsumer<ConsumerLagCalculator.ProcessGroupInfo,
+            CompletableFuture<ConsumerLagCalculator.CalculateLagResult>> biConsumer =
+                (info, future) -> calculate(info, future::complete);
+
         processAllGroup(info -> {
             if (info.group == null || info.topic == null) {
                 return;
             }
-
+            CompletableFuture<CalculateLagResult> future = new CompletableFuture<>();
             if (info.isPop && brokerConfig.isEnableNotifyBeforePopCalculateLag()) {
                 if (popLongPollingService.notifyMessageArriving(info.topic, -1, info.group,
-                    true, null, 0, null, null, new PopCommandCallback(this::calculate, info, lagRecorder))) {
+                    true, null, 0, null, null,
+                    new PopCommandCallback(biConsumer, info, future))) {
+                    futures.add(future);
                     return;
                 }
             }
-
             calculate(info, lagRecorder);
         });
+
+        // Set the maximum wait time to 10 seconds to avoid indefinite blocking
+        // in case of a fast fail that causes the future to not complete its execution.
+        try {
+            CompletableFuture.allOf(futures.toArray(
+                new CompletableFuture[0])).get(10, TimeUnit.SECONDS);
+
+            futures.forEach(future -> {
+                if (future.isDone() && !future.isCompletedExceptionally()) {
+                    lagRecorder.accept(future.join());
+                }
+            });
+        } catch (Exception e) {
+            LOGGER.error("Calculate lag timeout after 10 seconds", e);
+        }
     }
 
     public void calculate(ProcessGroupInfo info, Consumer<CalculateLagResult> lagRecorder) {