From b3e430f24bfedd94d3ceae028dce2e0053e268a6 Mon Sep 17 00:00:00 2001
From: Prashant Wason <pwason@uber.com>
Date: Wed, 11 Aug 2021 16:47:16 -0700
Subject: [PATCH] [HUDI-2017] Add API to set a metric in the registry. (#3084)

Registry.add() API adds the new value to existing metric value. For some use-cases We need a API to set/replace the existing value.

Metadata Table is synced in preWrite() and postWrite() functions of commit. As part of the sync, the current sizes and basefile/logfile counts are published as metrics. If we use the Registry.add() API, the count and sizes are incorrectly published as sum of the two values. This is corrected by using the Registry.set() API instead.
---
 .../org/apache/hudi/metrics/DistributedRegistry.java  |  5 +++++
 .../java/org/apache/hudi/common/metrics/Counter.java  |  4 ++++
 .../org/apache/hudi/common/metrics/LocalRegistry.java |  5 +++++
 .../java/org/apache/hudi/common/metrics/Registry.java | 11 +++++++++++
 .../apache/hudi/metadata/HoodieMetadataMetrics.java   | 10 +++++-----
 5 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metrics/DistributedRegistry.java b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metrics/DistributedRegistry.java
index 22b3afd10..60c32b34d 100644
--- a/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metrics/DistributedRegistry.java
+++ b/hudi-client/hudi-spark-client/src/main/java/org/apache/hudi/metrics/DistributedRegistry.java
@@ -60,6 +60,11 @@ public class DistributedRegistry extends AccumulatorV2<Map<String, Long>, Map<St
     counters.merge(name,  value, (oldValue, newValue) -> oldValue + newValue);
   }
 
+  @Override
+  public void set(String name, long value) {
+    counters.merge(name,  value, (oldValue, newValue) -> newValue);
+  }
+
   /**
    * Get all Counter type metrics.
    */
diff --git a/hudi-common/src/main/java/org/apache/hudi/common/metrics/Counter.java b/hudi-common/src/main/java/org/apache/hudi/common/metrics/Counter.java
index 546956d00..7e307e6fc 100644
--- a/hudi-common/src/main/java/org/apache/hudi/common/metrics/Counter.java
+++ b/hudi-common/src/main/java/org/apache/hudi/common/metrics/Counter.java
@@ -35,6 +35,10 @@ public class Counter implements Metric {
     this.count.addAndGet(n);
   }
 
+  public void set(long n) {
+    this.count.set(n);
+  }
+
   @Override
   public Long getValue() {
     return count.get();
diff --git a/hudi-common/src/main/java/org/apache/hudi/common/metrics/LocalRegistry.java b/hudi-common/src/main/java/org/apache/hudi/common/metrics/LocalRegistry.java
index 4fdf9f60d..330068f6d 100644
--- a/hudi-common/src/main/java/org/apache/hudi/common/metrics/LocalRegistry.java
+++ b/hudi-common/src/main/java/org/apache/hudi/common/metrics/LocalRegistry.java
@@ -48,6 +48,11 @@ public class LocalRegistry implements Registry {
     getCounter(name).add(value);
   }
 
+  @Override
+  public void set(String name, long value) {
+    getCounter(name).set(value);
+  }
+
   /**
    * Get all Counter type metrics.
    */
diff --git a/hudi-common/src/main/java/org/apache/hudi/common/metrics/Registry.java b/hudi-common/src/main/java/org/apache/hudi/common/metrics/Registry.java
index 4ac1e6178..bf5ee7e7b 100644
--- a/hudi-common/src/main/java/org/apache/hudi/common/metrics/Registry.java
+++ b/hudi-common/src/main/java/org/apache/hudi/common/metrics/Registry.java
@@ -100,6 +100,17 @@ public interface Registry extends Serializable {
    */
   void add(String name, long value);
 
+  /**
+   * Set the value to the metric.
+   *
+   * If the metric does not exist, it is added. If the metrics already exists, its value is replaced with the
+   * provided value.
+   *
+   * @param name Name of the metric.
+   * @param value The value to set for the metrics.
+   */
+  void set(String name, long value);
+
   /**
    * Get all Counter type metrics.
    */
diff --git a/hudi-common/src/main/java/org/apache/hudi/metadata/HoodieMetadataMetrics.java b/hudi-common/src/main/java/org/apache/hudi/metadata/HoodieMetadataMetrics.java
index 9f1b0a079..5924d8989 100644
--- a/hudi-common/src/main/java/org/apache/hudi/metadata/HoodieMetadataMetrics.java
+++ b/hudi-common/src/main/java/org/apache/hudi/metadata/HoodieMetadataMetrics.java
@@ -133,11 +133,11 @@ public class HoodieMetadataMetrics implements Serializable {
       return;
     }
 
-    // Update sizes and count for metadata table's data files
-    metricsRegistry.add("basefile.size", totalBaseFileSizeInBytes);
-    metricsRegistry.add("logfile.size", totalLogFileSizeInBytes);
-    metricsRegistry.add("basefile.count", baseFileCount);
-    metricsRegistry.add("logfile.count", logFileCount);
+    // Set new size and count for metadata table's data files
+    metricsRegistry.set("basefile.size", totalBaseFileSizeInBytes);
+    metricsRegistry.set("logfile.size", totalLogFileSizeInBytes);
+    metricsRegistry.set("basefile.count", baseFileCount);
+    metricsRegistry.set("logfile.count", logFileCount);
 
     LOG.info(String.format("Updating metadata size metrics (basefile.size=%d, logfile.size=%d, basefile.count=%d, "
         + "logfile.count=%d)", totalBaseFileSizeInBytes, totalLogFileSizeInBytes, baseFileCount, logFileCount));