From 745ba45ea4cc260c1019c58eeb17dfcccf06b09b Mon Sep 17 00:00:00 2001
From: Albumen Kevin <jhq0812@gmail.com>
Date: Wed, 27 Dec 2023 23:20:10 +0800
Subject: [PATCH] Fix MethodKeyRt override & Error enhance (#13580)

Co-authored-by: xiaosheng <songxiaosheng@apache.org>
---
 .../dubbo/metrics/data/RtStatComposite.java   |  4 +++-
 .../model/sample/GaugeMetricSample.java       | 24 ++++++++++++++++++-
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/dubbo-metrics/dubbo-metrics-api/src/main/java/org/apache/dubbo/metrics/data/RtStatComposite.java b/dubbo-metrics/dubbo-metrics-api/src/main/java/org/apache/dubbo/metrics/data/RtStatComposite.java
index 0a3001f2ecc..729afdacde6 100644
--- a/dubbo-metrics/dubbo-metrics-api/src/main/java/org/apache/dubbo/metrics/data/RtStatComposite.java
+++ b/dubbo-metrics/dubbo-metrics-api/src/main/java/org/apache/dubbo/metrics/data/RtStatComposite.java
@@ -160,7 +160,9 @@ private List<Action> calServiceRtActions(Invocation invocation, String registryO
     public void calcMethodKeyRt(Invocation invocation, String registryOpType, Long responseTime) {
         List<Action> actions;
 
-        if (invocation.getServiceModel() != null && invocation.getServiceModel().getServiceMetadata() != null) {
+        if (getServiceLevel()
+                && invocation.getServiceModel() != null
+                && invocation.getServiceModel().getServiceMetadata() != null) {
             Map<String, Object> attributeMap =
                     invocation.getServiceModel().getServiceMetadata().getAttributeMap();
             Map<String, List<Action>> cache = (Map<String, List<Action>>) attributeMap.get("MethodKeyRt");
diff --git a/dubbo-metrics/dubbo-metrics-api/src/main/java/org/apache/dubbo/metrics/model/sample/GaugeMetricSample.java b/dubbo-metrics/dubbo-metrics-api/src/main/java/org/apache/dubbo/metrics/model/sample/GaugeMetricSample.java
index bc17fd511af..d998e0be37c 100644
--- a/dubbo-metrics/dubbo-metrics-api/src/main/java/org/apache/dubbo/metrics/model/sample/GaugeMetricSample.java
+++ b/dubbo-metrics/dubbo-metrics-api/src/main/java/org/apache/dubbo/metrics/model/sample/GaugeMetricSample.java
@@ -16,18 +16,25 @@
  */
 package org.apache.dubbo.metrics.model.sample;
 
+import org.apache.dubbo.common.logger.ErrorTypeAwareLogger;
+import org.apache.dubbo.common.logger.LoggerFactory;
 import org.apache.dubbo.metrics.model.MetricsCategory;
 import org.apache.dubbo.metrics.model.key.MetricsKey;
 import org.apache.dubbo.metrics.model.key.MetricsKeyWrapper;
 
 import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.ToDoubleFunction;
 
+import static org.apache.dubbo.common.constants.LoggerCodeConstants.COMMON_METRICS_COLLECTOR_EXCEPTION;
+
 /**
  * GaugeMetricSample.
  */
 public class GaugeMetricSample<T> extends MetricSample {
+    private static final ErrorTypeAwareLogger logger = LoggerFactory.getErrorTypeAwareLogger(GaugeMetricSample.class);
+    private final AtomicBoolean warned = new AtomicBoolean(false);
 
     private final T value;
 
@@ -71,7 +78,22 @@ public GaugeMetricSample(
             ToDoubleFunction<T> apply) {
         super(name, description, tags, Type.GAUGE, category, baseUnit);
         this.value = Objects.requireNonNull(value, "The GaugeMetricSample value cannot be null");
-        this.apply = Objects.requireNonNull(apply, "The GaugeMetricSample apply cannot be null");
+        Objects.requireNonNull(apply, "The GaugeMetricSample apply cannot be null");
+        this.apply = (e) -> {
+            try {
+                return apply.applyAsDouble(e);
+            } catch (Throwable t) {
+                if (warned.compareAndSet(false, true)) {
+                    logger.error(
+                            COMMON_METRICS_COLLECTOR_EXCEPTION,
+                            "",
+                            "",
+                            "Unexpected error occurred when applying the GaugeMetricSample",
+                            t);
+                }
+                return 0;
+            }
+        };
     }
 
     public T getValue() {