apache · szetszwo · Aug 31, 2022 · Aug 26, 2022
diff --git a/ratis-grpc/src/main/java/org/apache/ratis/grpc/metrics/GrpcServerMetrics.java b/ratis-grpc/src/main/java/org/apache/ratis/grpc/metrics/GrpcServerMetrics.java
@@ -17,14 +17,14 @@
  */
 package org.apache.ratis.grpc.metrics;
 
-import org.apache.ratis.thirdparty.com.codahale.metrics.Gauge;
 import org.apache.ratis.metrics.MetricRegistryInfo;
 import org.apache.ratis.metrics.RatisMetricRegistry;
 import org.apache.ratis.metrics.RatisMetrics;
-import org.apache.ratis.thirdparty.com.google.common.annotations.VisibleForTesting;
 
 import org.apache.ratis.thirdparty.com.codahale.metrics.Timer;
 
+import java.util.function.Supplier;
+
 public class GrpcServerMetrics extends RatisMetrics {
   private static final String RATIS_GRPC_METRICS_APP_NAME = "ratis_grpc";
   private static final String RATIS_GRPC_METRICS_COMP_NAME = "log_appender";
@@ -89,8 +89,7 @@ public void onRequestTimeout(String follower, boolean isHeartbeat) {
         follower)).inc();
   }
 
-  public void addPendingRequestsCount(String follower,
-      Gauge pendinglogQueueSize) {
+  public void addPendingRequestsCount(String follower, Supplier<Integer> pendinglogQueueSize) {
     registry.gauge(String.format(RATIS_GRPC_METRICS_LOG_APPENDER_PENDING_COUNT, follower), () -> pendinglogQueueSize);
   }
 
@@ -101,9 +100,4 @@ public void onInstallSnapshot() {
   public static String getHeartbeatSuffix(boolean heartbeat) {
     return heartbeat ? "_heartbeat" : "";
   }
-
-  @VisibleForTesting
-  public RatisMetricRegistry getRegistry() {
-    return registry;
-  }
 }
diff --git a/ratis-grpc/src/test/java/org/apache/ratis/grpc/server/TestGrpcServerMetrics.java b/ratis-grpc/src/test/java/org/apache/ratis/grpc/server/TestGrpcServerMetrics.java
@@ -27,16 +27,16 @@
 import static org.apache.ratis.grpc.metrics.GrpcServerMetrics.RATIS_GRPC_METRICS_REQUEST_RETRY_COUNT;
 import static org.mockito.Mockito.when;
 
-import java.util.SortedMap;
 import java.util.function.Consumer;
 
-import org.apache.ratis.thirdparty.com.codahale.metrics.Gauge;
+import org.apache.ratis.server.metrics.ServerMetricsTestUtils;
 import org.apache.ratis.grpc.metrics.GrpcServerMetrics;
 import org.apache.ratis.metrics.RatisMetricRegistry;
 import org.apache.ratis.proto.RaftProtos;
 import org.apache.ratis.protocol.RaftGroupId;
 import org.apache.ratis.protocol.RaftGroupMemberId;
 import org.apache.ratis.protocol.RaftPeerId;
+import org.apache.ratis.thirdparty.com.codahale.metrics.Gauge;
 import org.junit.Assert;
 import org.junit.BeforeClass;
 import org.junit.Test;
@@ -109,21 +109,12 @@ public void testGrpcLogPendingRequestCount() {
     GrpcLogAppender.RequestMap pendingRequest = Mockito.mock(GrpcLogAppender.RequestMap.class);
     when(pendingRequest.logRequestsSize()).thenReturn(0);
     grpcServerMetrics.addPendingRequestsCount(raftPeerId.toString(),
-        () -> pendingRequest.logRequestsSize());
-    Assert.assertEquals(0, getGuageWithName(
-            String.format(RATIS_GRPC_METRICS_LOG_APPENDER_PENDING_COUNT,
-                raftPeerId.toString())).getValue());
+        pendingRequest::logRequestsSize);
+    final String name = String.format(RATIS_GRPC_METRICS_LOG_APPENDER_PENDING_COUNT, raftPeerId);
+    final Gauge gauge = ServerMetricsTestUtils.getGaugeWithName(name, grpcServerMetrics::getRegistry);
+    Assert.assertEquals(0, gauge.getValue());
     when(pendingRequest.logRequestsSize()).thenReturn(10);
-    Assert.assertEquals(10, getGuageWithName(
-            String.format(RATIS_GRPC_METRICS_LOG_APPENDER_PENDING_COUNT,
-                raftPeerId.toString())).getValue());
-  }
-
-  private Gauge getGuageWithName(String gaugeName) {
-    SortedMap<String, Gauge> gaugeMap =
-        grpcServerMetrics.getRegistry().getGauges((s, metric) ->
-            s.contains(gaugeName));
-    return gaugeMap.get(gaugeMap.firstKey());
+    Assert.assertEquals(10, gauge.getValue());
   }
 
   @Test

diff --git a/ratis-metrics/src/main/java/org/apache/ratis/metrics/MetricRegistryInfo.java b/ratis-metrics/src/main/java/org/apache/ratis/metrics/MetricRegistryInfo.java
@@ -1,4 +1,4 @@
-/**
+/*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
@@ -18,11 +18,10 @@
 
 package org.apache.ratis.metrics;
 
+import org.apache.ratis.thirdparty.com.codahale.metrics.MetricRegistry;
 
 import java.util.Objects;
 
-import org.apache.ratis.thirdparty.com.codahale.metrics.*;
-
 /**
  *
  * This class holds the name and description and JMX related context names for such group of

diff --git a/ratis-metrics/src/main/java/org/apache/ratis/metrics/MetricsReporting.java b/ratis-metrics/src/main/java/org/apache/ratis/metrics/MetricsReporting.java
@@ -17,15 +17,15 @@
  */
 package org.apache.ratis.metrics;
 
-import java.util.concurrent.TimeUnit;
-import java.util.function.Consumer;
-
-import org.apache.ratis.thirdparty.com.codahale.metrics.*;
+import org.apache.ratis.thirdparty.com.codahale.metrics.ConsoleReporter;
 import org.apache.ratis.thirdparty.com.codahale.metrics.jmx.JmxReporter;
 import org.apache.ratis.util.TimeDuration;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+
 public final class MetricsReporting {
   public static final Logger LOG = LoggerFactory.getLogger(MetricsReporting.class);
 

diff --git a/ratis-metrics/src/main/java/org/apache/ratis/metrics/RatisMetricRegistry.java b/ratis-metrics/src/main/java/org/apache/ratis/metrics/RatisMetricRegistry.java
@@ -1,4 +1,4 @@
-/**
+/*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
@@ -17,25 +17,30 @@
  */
 package org.apache.ratis.metrics;
 
-import java.util.SortedMap;
-
-import org.apache.ratis.thirdparty.com.codahale.metrics.*;
+import org.apache.ratis.thirdparty.com.codahale.metrics.ConsoleReporter;
+import org.apache.ratis.thirdparty.com.codahale.metrics.Counter;
+import org.apache.ratis.thirdparty.com.codahale.metrics.Histogram;
+import org.apache.ratis.thirdparty.com.codahale.metrics.Meter;
+import org.apache.ratis.thirdparty.com.codahale.metrics.Metric;
+import org.apache.ratis.thirdparty.com.codahale.metrics.MetricRegistry;
+import org.apache.ratis.thirdparty.com.codahale.metrics.MetricSet;
+import org.apache.ratis.thirdparty.com.codahale.metrics.Timer;
 import org.apache.ratis.thirdparty.com.codahale.metrics.jmx.JmxReporter;
 import org.apache.ratis.thirdparty.com.google.common.annotations.VisibleForTesting;
 
+import java.util.function.Supplier;
+
 public interface RatisMetricRegistry {
   Timer timer(String name);
 
   Counter counter(String name);
 
   boolean remove(String name);
 
-  Gauge gauge(String name, MetricRegistry.MetricSupplier<Gauge> supplier);
+  <T> void gauge(String name, Supplier<Supplier<T>> gaugeSupplier);
 
   Timer timer(String name, MetricRegistry.MetricSupplier<Timer> supplier);
 
-  SortedMap<String, Gauge> getGauges(MetricFilter filter);
-
   Counter counter(String name, MetricRegistry.MetricSupplier<Counter> supplier);
 
   Histogram histogram(String name);

diff --git a/ratis-metrics/src/main/java/org/apache/ratis/metrics/impl/RatisMetricRegistryImpl.java b/ratis-metrics/src/main/java/org/apache/ratis/metrics/impl/RatisMetricRegistryImpl.java
@@ -1,4 +1,4 @@
-/**
+/*
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
@@ -17,20 +17,30 @@
  */
 package org.apache.ratis.metrics.impl;
 
-import java.util.Map;
-import java.util.SortedMap;
-
-import org.apache.ratis.thirdparty.com.codahale.metrics.*;
 import org.apache.ratis.metrics.MetricRegistryInfo;
 import org.apache.ratis.metrics.RatisMetricRegistry;
+import org.apache.ratis.thirdparty.com.codahale.metrics.ConsoleReporter;
+import org.apache.ratis.thirdparty.com.codahale.metrics.Counter;
+import org.apache.ratis.thirdparty.com.codahale.metrics.Gauge;
+import org.apache.ratis.thirdparty.com.codahale.metrics.Histogram;
+import org.apache.ratis.thirdparty.com.codahale.metrics.Meter;
+import org.apache.ratis.thirdparty.com.codahale.metrics.Metric;
+import org.apache.ratis.thirdparty.com.codahale.metrics.MetricFilter;
+import org.apache.ratis.thirdparty.com.codahale.metrics.MetricRegistry;
+import org.apache.ratis.thirdparty.com.codahale.metrics.MetricSet;
+import org.apache.ratis.thirdparty.com.codahale.metrics.Timer;
 import org.apache.ratis.thirdparty.com.codahale.metrics.jmx.JmxReporter;
 import org.apache.ratis.thirdparty.com.google.common.annotations.VisibleForTesting;
 
+import java.util.Map;
+import java.util.SortedMap;
+import java.util.function.Supplier;
+
 /**
  * Custom implementation of {@link MetricRegistry}.
  */
 public class RatisMetricRegistryImpl implements RatisMetricRegistry {
-  private MetricRegistry metricRegistry = new MetricRegistry();
+  private final MetricRegistry metricRegistry = new MetricRegistry();
 
   private final MetricRegistryInfo info;
 
@@ -57,15 +67,20 @@ public boolean remove(String name) {
     return metricRegistry.remove(getMetricName(name));
   }
 
-  @Override public Gauge gauge(String name, MetricRegistry.MetricSupplier<Gauge> supplier) {
-    return metricRegistry.gauge(getMetricName(name), supplier);
+  static <T> Gauge<T> toGauge(Supplier<T> supplier) {
+    return supplier::get;
+  }
+
+  @Override
+  public <T> void gauge(String name, Supplier<Supplier<T>> gaugeSupplier) {
+    metricRegistry.gauge(getMetricName(name), () -> toGauge(gaugeSupplier.get()));
   }
 
   @Override public Timer timer(String name, MetricRegistry.MetricSupplier<Timer> supplier) {
     return metricRegistry.timer(getMetricName(name), supplier);
   }
 
-  @Override public SortedMap<String, Gauge> getGauges(MetricFilter filter) {
+  public SortedMap<String, Gauge> getGauges(MetricFilter filter) {
     return metricRegistry.getGauges(filter);
   }
 

diff --git a/ratis-server/src/main/java/org/apache/ratis/server/impl/StateMachineMetrics.java b/ratis-server/src/main/java/org/apache/ratis/server/impl/StateMachineMetrics.java
@@ -56,10 +56,8 @@ public static StateMachineMetrics getStateMachineMetrics(
   private StateMachineMetrics(String serverId, LongSupplier getApplied,
       LongSupplier getApplyCompleted) {
     registry = getMetricRegistryForStateMachine(serverId);
-    registry.gauge(STATEMACHINE_APPLIED_INDEX_GAUGE,
-        () -> () -> getApplied.getAsLong());
-    registry.gauge(STATEMACHINE_APPLY_COMPLETED_GAUGE,
-        () -> () -> getApplyCompleted.getAsLong());
+    registry.gauge(STATEMACHINE_APPLIED_INDEX_GAUGE, () -> getApplied::getAsLong);
+    registry.gauge(STATEMACHINE_APPLY_COMPLETED_GAUGE, () -> getApplyCompleted::getAsLong);
   }
 
   private RatisMetricRegistry getMetricRegistryForStateMachine(String serverId) {

diff --git a/ratis-server/src/main/java/org/apache/ratis/server/metrics/RaftLogMetricsBase.java b/ratis-server/src/main/java/org/apache/ratis/server/metrics/RaftLogMetricsBase.java
@@ -18,12 +18,12 @@
 
 package org.apache.ratis.server.metrics;
 
-import org.apache.ratis.thirdparty.com.codahale.metrics.*;
 import org.apache.ratis.metrics.MetricRegistryInfo;
 import org.apache.ratis.metrics.RatisMetricRegistry;
 import org.apache.ratis.metrics.RatisMetrics;
 import org.apache.ratis.protocol.RaftGroupMemberId;
 import org.apache.ratis.server.raftlog.LogEntryHeader;
+import org.apache.ratis.thirdparty.com.codahale.metrics.Timer;
 
 public class RaftLogMetricsBase extends RatisMetrics implements RaftLogMetrics {
   public static final String RATIS_LOG_WORKER_METRICS_DESC = "Metrics for Log Worker";

diff --git a/ratis-server/src/main/java/org/apache/ratis/server/metrics/RaftServerMetricsImpl.java b/ratis-server/src/main/java/org/apache/ratis/server/metrics/RaftServerMetricsImpl.java
@@ -22,12 +22,10 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.HashMap;
 import java.util.Map;
-import java.util.SortedMap;
 import java.util.function.Function;
 import java.util.function.Supplier;
 
 import org.apache.ratis.thirdparty.com.codahale.metrics.Counter;
-import org.apache.ratis.thirdparty.com.codahale.metrics.Gauge;
 import org.apache.ratis.thirdparty.com.codahale.metrics.Timer;
 
 import org.apache.ratis.metrics.MetricRegistryInfo;
@@ -40,7 +38,6 @@
 import org.apache.ratis.metrics.RatisMetrics;
 import org.apache.ratis.server.RetryCache;
 import org.apache.ratis.thirdparty.com.google.common.annotations.VisibleForTesting;
-import org.apache.ratis.util.Preconditions;
 
 /**
  * Metric Registry for Raft Group Server. One instance per leader/follower.
@@ -88,7 +85,7 @@ public final class RaftServerMetricsImpl extends RatisMetrics implements RaftSer
   /** id -> key */
   private static final Map<RaftPeerId, String> PEER_COMMIT_INDEX_GAUGE_KEYS = new ConcurrentHashMap<>();
 
-  private static String getPeerCommitIndexGaugeKey(RaftPeerId serverId) {
+  static String getPeerCommitIndexGaugeKey(RaftPeerId serverId) {
     return PEER_COMMIT_INDEX_GAUGE_KEYS.computeIfAbsent(serverId,
         key -> String.format(LEADER_METRIC_PEER_COMMIT_INDEX, key));
   }
@@ -152,26 +149,9 @@ public void addPeerCommitIndexGauge(RaftPeerId peerId) {
         .orElse(0L));
   }
 
-  /**
-   * Get the commit index gauge for the given peer of the server
-   * @return Metric Gauge holding the value of commit index of the peer
-   */
   @VisibleForTesting
-  public static Gauge getPeerCommitIndexGauge(RaftGroupMemberId serverId, RaftPeerId peerId) {
-
-    final RaftServerMetricsImpl serverMetrics = METRICS.get(serverId);
-    if (serverMetrics == null) {
-      return null;
-    }
-
-    final String followerCommitIndexKey = getPeerCommitIndexGaugeKey(peerId);
-
-    SortedMap<String, Gauge> map =
-        serverMetrics.registry.getGauges((s, metric) ->
-            s.contains(followerCommitIndexKey));
-
-    Preconditions.assertTrue(map.size() <= 1);
-    return map.get(map.firstKey());
+  static RaftServerMetricsImpl getImpl(RaftGroupMemberId serverId) {
+    return METRICS.get(serverId);
   }
 
   /**
@@ -213,15 +193,15 @@ public void onRequestQueueLimitHit() {
     registry.counter(REQUEST_QUEUE_LIMIT_HIT_COUNTER).inc();
   }
 
-  public void addNumPendingRequestsGauge(Gauge queueSize) {
+  public void addNumPendingRequestsGauge(Supplier<Integer> queueSize) {
     registry.gauge(REQUEST_QUEUE_SIZE, () -> queueSize);
   }
 
   public boolean removeNumPendingRequestsGauge() {
     return registry.remove(REQUEST_QUEUE_SIZE);
   }
 
-  public void addNumPendingRequestsMegaByteSize(Gauge megabyteSize) {
+  public void addNumPendingRequestsMegaByteSize(Supplier<Integer> megabyteSize) {
     registry.gauge(REQUEST_MEGA_BYTE_SIZE, () -> megabyteSize);
   }
 

diff --git a/ratis-server/src/main/java/org/apache/ratis/server/metrics/SegmentedRaftLogMetrics.java b/ratis-server/src/main/java/org/apache/ratis/server/metrics/SegmentedRaftLogMetrics.java
@@ -18,12 +18,10 @@
 
 package org.apache.ratis.server.metrics;
 
-import org.apache.ratis.thirdparty.com.codahale.metrics.*;
 import org.apache.ratis.protocol.RaftGroupMemberId;
-import org.apache.ratis.server.raftlog.segmented.SegmentedRaftLogCache;
-import org.apache.ratis.util.DataQueue;
+import org.apache.ratis.thirdparty.com.codahale.metrics.Timer;
 
-import java.util.Queue;
+import java.util.function.Supplier;
 
 public class SegmentedRaftLogMetrics extends RaftLogMetricsBase {
   //////////////////////////////
@@ -82,34 +80,28 @@ public SegmentedRaftLogMetrics(RaftGroupMemberId serverId) {
     super(serverId);
   }
 
-  public void addDataQueueSizeGauge(DataQueue queue) {
-    registry.gauge(RAFT_LOG_DATA_QUEUE_SIZE, () -> queue::getNumElements);
+  public void addDataQueueSizeGauge(Supplier<Integer> numElements) {
+    registry.gauge(RAFT_LOG_DATA_QUEUE_SIZE, () -> numElements);
   }
 
-  public void addClosedSegmentsNum(SegmentedRaftLogCache cache) {
-    registry.gauge(RAFT_LOG_CACHE_CLOSED_SEGMENTS_NUM, () -> () -> {
-      return cache.getCachedSegmentNum();
-    });
+  public void addClosedSegmentsNum(Supplier<Long> cachedSegmentNum) {
+    registry.gauge(RAFT_LOG_CACHE_CLOSED_SEGMENTS_NUM, () -> cachedSegmentNum);
   }
 
-  public void addClosedSegmentsSizeInBytes(SegmentedRaftLogCache cache) {
-    registry.gauge(RAFT_LOG_CACHE_CLOSED_SEGMENTS_SIZE_IN_BYTES, () -> () -> {
-      return cache.getClosedSegmentsSizeInBytes();
-    });
+  public void addClosedSegmentsSizeInBytes(Supplier<Long> closedSegmentsSizeInBytes) {
+    registry.gauge(RAFT_LOG_CACHE_CLOSED_SEGMENTS_SIZE_IN_BYTES, () -> closedSegmentsSizeInBytes);
   }
 
-  public void addOpenSegmentSizeInBytes(SegmentedRaftLogCache cache) {
-    registry.gauge(RAFT_LOG_CACHE_OPEN_SEGMENT_SIZE_IN_BYTES, () -> () -> {
-      return cache.getOpenSegmentSizeInBytes();
-    });
+  public void addOpenSegmentSizeInBytes(Supplier<Long> openSegmentSizeInBytes) {
+    registry.gauge(RAFT_LOG_CACHE_OPEN_SEGMENT_SIZE_IN_BYTES, () -> openSegmentSizeInBytes);
   }
 
-  public void addLogWorkerQueueSizeGauge(Queue queue) {
-    registry.gauge(RAFT_LOG_WORKER_QUEUE_SIZE, () -> () -> queue.size());
+  public void addLogWorkerQueueSizeGauge(Supplier<Integer> queueSize) {
+    registry.gauge(RAFT_LOG_WORKER_QUEUE_SIZE, () -> queueSize);
   }
 
-  public void addFlushBatchSizeGauge(MetricRegistry.MetricSupplier<Gauge> supplier) {
-    registry.gauge(RAFT_LOG_SYNC_BATCH_SIZE, supplier);
+  public void addFlushBatchSizeGauge(Supplier<Integer> flushBatchSize) {
+    registry.gauge(RAFT_LOG_SYNC_BATCH_SIZE, () -> flushBatchSize);
   }
 
   private Timer getTimer(String timerName) {