From 970d9b9d784abf316a8146f57d2b26cfc0691aa2 Mon Sep 17 00:00:00 2001
From: "Xiaotian (Jackie) Jiang"
 <17555551+Jackie-Jiang@users.noreply.github.com>
Date: Tue, 4 Jun 2024 13:14:49 -0700
Subject: [PATCH] Fix the NPE from IS update metrics (#13313)

---
 .../pinot/common/utils/helix/HelixHelper.java | 28 ++++++++++---------
 1 file changed, 15 insertions(+), 13 deletions(-)

diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/HelixHelper.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/HelixHelper.java
index 6dda43751b01..1910048cf686 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/HelixHelper.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/helix/HelixHelper.java
@@ -18,7 +18,6 @@
  */
 package org.apache.pinot.common.utils.helix;
 
-import com.google.common.base.Function;
 import com.google.common.base.Preconditions;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -29,6 +28,7 @@
 import java.util.Set;
 import java.util.concurrent.Callable;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
@@ -102,20 +102,21 @@ public static IdealState cloneIdealState(IdealState idealState) {
    * @param updater A function that returns an updated ideal state given an input ideal state
    * @return updated ideal state if successful, null if not
    */
-  public static IdealState updateIdealState(final HelixManager helixManager, final String resourceName,
-      final Function<IdealState, IdealState> updater, RetryPolicy policy, final boolean noChangeOk) {
+  public static IdealState updateIdealState(HelixManager helixManager, String resourceName,
+      Function<IdealState, IdealState> updater, RetryPolicy policy, boolean noChangeOk) {
+    // NOTE: ControllerMetrics could be null because this method might be invoked by Broker.
     ControllerMetrics controllerMetrics = ControllerMetrics.get();
     try {
       long startTimeMs = System.currentTimeMillis();
       IdealStateWrapper idealStateWrapper = new IdealStateWrapper();
-      int retries = policy.attempt(new Callable<Boolean>() {
+      int retries = policy.attempt(new Callable<>() {
         @Override
         public Boolean call() {
           HelixDataAccessor dataAccessor = helixManager.getHelixDataAccessor();
           PropertyKey idealStateKey = dataAccessor.keyBuilder().idealStates(resourceName);
           IdealState idealState = dataAccessor.getProperty(idealStateKey);
 
-          // Make a copy of the the idealState above to pass it to the updater
+          // Make a copy of the idealState above to pass it to the updater
           // NOTE: new IdealState(idealState.getRecord()) does not work because it's shallow copy for map fields and
           // list fields
           IdealState idealStateCopy = cloneIdealState(idealState);
@@ -197,20 +198,21 @@ private boolean shouldCompress(IdealState is) {
               numChars += entry.getValue().length();
             }
             numChars *= is.getNumPartitions();
-            if (_minNumCharsInISToTurnOnCompression > 0
-                && numChars > _minNumCharsInISToTurnOnCompression) {
-              return true;
-            }
+            return _minNumCharsInISToTurnOnCompression > 0 && numChars > _minNumCharsInISToTurnOnCompression;
           }
           return false;
         }
       });
-      controllerMetrics.addMeteredValue(resourceName, ControllerMeter.IDEAL_STATE_UPDATE_RETRY, retries);
-      controllerMetrics.addTimedValue(resourceName, ControllerTimer.IDEAL_STATE_UPDATE_TIME_MS,
-              System.currentTimeMillis() - startTimeMs, TimeUnit.MILLISECONDS);
+      if (controllerMetrics != null) {
+        controllerMetrics.addMeteredValue(resourceName, ControllerMeter.IDEAL_STATE_UPDATE_RETRY, retries);
+        controllerMetrics.addTimedValue(resourceName, ControllerTimer.IDEAL_STATE_UPDATE_TIME_MS,
+            System.currentTimeMillis() - startTimeMs, TimeUnit.MILLISECONDS);
+      }
       return idealStateWrapper._idealState;
     } catch (Exception e) {
-      controllerMetrics.addMeteredValue(resourceName, ControllerMeter.IDEAL_STATE_UPDATE_FAILURE, 1L);
+      if (controllerMetrics != null) {
+        controllerMetrics.addMeteredValue(resourceName, ControllerMeter.IDEAL_STATE_UPDATE_FAILURE, 1L);
+      }
       throw new RuntimeException("Caught exception while updating ideal state for resource: " + resourceName, e);
     }
   }