Enhance SegmentStatusChecker to honor no-queryable servers and instance assignment config

pinot-controller/src/main/java/org/apache/pinot/controller/helix/SegmentStatusChecker.java

@@ -47,6 +47,8 @@
 import org.apache.pinot.controller.helix.core.periodictask.ControllerPeriodicTask;
 import org.apache.pinot.controller.helix.core.realtime.MissingConsumingSegmentFinder;
 import org.apache.pinot.controller.helix.core.realtime.PinotLLCRealtimeSegmentManager;
+import org.apache.pinot.controller.util.ServerInfoFetcher;
+import org.apache.pinot.controller.util.ServerInfoFetcher.ServerInfo;
 import org.apache.pinot.controller.util.TableSizeReader;
 import org.apache.pinot.spi.config.table.TableConfig;
 import org.apache.pinot.spi.config.table.TableType;
@@ -80,6 +82,8 @@ public class SegmentStatusChecker extends ControllerPeriodicTask<SegmentStatusCh
 
   private long _lastDisabledTableLogTimestamp = 0;
 
+  private ServerInfoFetcher _serverInfoFetcher;
+
   /**
    * Constructs the segment status checker.
    * @param pinotHelixResourceManager The resource checker used to interact with Helix
@@ -92,6 +96,7 @@ public SegmentStatusChecker(PinotHelixResourceManager pinotHelixResourceManager,
         config.getStatusCheckerInitialDelayInSeconds(), pinotHelixResourceManager, leadControllerManager,
         controllerMetrics);
 
+    _serverInfoFetcher = new ServerInfoFetcher(pinotHelixResourceManager);
     _waitForPushTimeSeconds = config.getStatusCheckerWaitForPushTimeInSeconds();
     _tableSizeReader = tableSizeReader;
   }
@@ -269,10 +274,12 @@ private void updateSegmentMetrics(String tableNameWithType, TableConfig tableCon
 
     ExternalView externalView = _pinotHelixResourceManager.getTableExternalView(tableNameWithType);
 
-    // Maximum number of replicas in ideal state
-    int maxISReplicas = Integer.MIN_VALUE;
-    // Minimum number of replicas in external view
-    int minEVReplicas = Integer.MAX_VALUE;
+    // Maximum number of replicas that is up (ONLINE/CONSUMING) in ideal state
+    int maxISReplicasUp = Integer.MIN_VALUE;
+    // Minimum number of replicas that is up (ONLINE/CONSUMING) in external view
+    int minEVReplicasUp = Integer.MAX_VALUE;
+    // Minimum percentage of replicas that is up (ONLINE/CONSUMING) in external view
+    int minEVReplicasUpPercent = 100;
     // Total compressed segment size in deep store
     long tableCompressedSize = 0;
     // Segments without ZK metadata
@@ -286,18 +293,19 @@ private void updateSegmentMetrics(String tableNameWithType, TableConfig tableCon
     List<String> segmentsInvalidStartTime = new ArrayList<>();
     List<String> segmentsInvalidEndTime = new ArrayList<>();
     for (String segment : segments) {
-      int numISReplicas = 0;
+      // Number of replicas in ideal state that is in ONLINE/CONSUMING state
+      int numISReplicasUp = 0;
       for (Map.Entry<String, String> entry : idealState.getInstanceStateMap(segment).entrySet()) {
         String state = entry.getValue();
         if (state.equals(SegmentStateModel.ONLINE) || state.equals(SegmentStateModel.CONSUMING)) {
-          numISReplicas++;
+          numISReplicasUp++;
         }
       }
-      // Skip segments not ONLINE/CONSUMING in ideal state
-      if (numISReplicas == 0) {
+      // Skip segments with no ONLINE/CONSUMING in ideal state
+      if (numISReplicasUp == 0) {
         continue;
       }
-      maxISReplicas = Math.max(maxISReplicas, numISReplicas);
+      maxISReplicasUp = Math.max(maxISReplicasUp, numISReplicasUp);
 
       SegmentZKMetadata segmentZKMetadata = _pinotHelixResourceManager.getSegmentZKMetadata(tableNameWithType, segment);
       // Skip the segment when it doesn't have ZK metadata. Most likely the segment is just deleted.
@@ -330,46 +338,49 @@ private void updateSegmentMetrics(String tableNameWithType, TableConfig tableCon
         }
       }
 
-      int numEVReplicas = 0;
+      int numEVReplicasUp = 0;
       if (externalView != null) {
         Map<String, String> stateMap = externalView.getStateMap(segment);
         if (stateMap != null) {
           for (Map.Entry<String, String> entry : stateMap.entrySet()) {
+            String serverInstanceId = entry.getKey();
             String state = entry.getValue();
-            if (state.equals(SegmentStateModel.ONLINE) || state.equals(SegmentStateModel.CONSUMING)) {
-              numEVReplicas++;
+            if (isServerQueryable(serverInstanceId)
+              && (state.equals(SegmentStateModel.ONLINE) || state.equals(SegmentStateModel.CONSUMING))) {
+              numEVReplicasUp++;
             }
             if (state.equals(SegmentStateModel.ERROR)) {
               errorSegments.add(Pair.of(segment, entry.getKey()));
             }
           }
         }
       }
-      if (numEVReplicas == 0) {
+      if (numEVReplicasUp == 0) {
         offlineSegments.add(segment);
-      } else if (numEVReplicas < numISReplicas) {
+      } else if (numEVReplicasUp < numISReplicasUp) {
         partialOnlineSegments.add(segment);
       } else {
-        // Do not allow nReplicasEV to be larger than nReplicasIS
-        numEVReplicas = numISReplicas;
+        // Do not allow numEVReplicasUp to be larger than numISReplicasUp
+        numEVReplicasUp = numISReplicasUp;
       }
-      minEVReplicas = Math.min(minEVReplicas, numEVReplicas);
+
+      minEVReplicasUp = Math.min(minEVReplicasUp, numEVReplicasUp);
+      // Total number of replicas in ideal state (including ERROR/OFFLINE states)
+      int numISReplicasTotal = Math.max(idealState.getInstanceStateMap(segment).entrySet().size(), 1);
+      minEVReplicasUpPercent = Math.min(minEVReplicasUpPercent, numEVReplicasUp * 100 / numISReplicasTotal);
     }
 
-    if (maxISReplicas == Integer.MIN_VALUE) {
+    if (maxISReplicasUp == Integer.MIN_VALUE) {
       try {
-        maxISReplicas = Math.max(Integer.parseInt(idealState.getReplicas()), 1);
+        maxISReplicasUp = Math.max(Integer.parseInt(idealState.getReplicas()), 1);
       } catch (NumberFormatException e) {
-        maxISReplicas = 1;
+        maxISReplicasUp = 1;
       }
     }
-    // Do not allow minEVReplicas to be larger than maxISReplicas
-    minEVReplicas = Math.min(minEVReplicas, maxISReplicas);
 
-    if (minEVReplicas < maxISReplicas) {
-      LOGGER.warn("Table {} has at least one segment running with only {} replicas, below replication threshold :{}",
-          tableNameWithType, minEVReplicas, maxISReplicas);
-    }
+    // Do not allow minEVReplicasUp to be larger than maxISReplicasUp
+    minEVReplicasUp = Math.min(minEVReplicasUp, maxISReplicasUp);
+
     int numSegmentsWithoutZKMetadata = segmentsWithoutZKMetadata.size();
     if (numSegmentsWithoutZKMetadata > 0) {
       LOGGER.warn("Table {} has {} segments without ZK metadata: {}", tableNameWithType, numSegmentsWithoutZKMetadata,
@@ -402,9 +413,9 @@ private void updateSegmentMetrics(String tableNameWithType, TableConfig tableCon
     }
 
     // Synchronization provided by Controller Gauge to make sure that only one thread updates the gauge
-    _controllerMetrics.setValueOfTableGauge(tableNameWithType, ControllerGauge.NUMBER_OF_REPLICAS, minEVReplicas);
+    _controllerMetrics.setValueOfTableGauge(tableNameWithType, ControllerGauge.NUMBER_OF_REPLICAS, minEVReplicasUp);
     _controllerMetrics.setValueOfTableGauge(tableNameWithType, ControllerGauge.PERCENT_OF_REPLICAS,
-        minEVReplicas * 100L / maxISReplicas);
+        minEVReplicasUpPercent);
     _controllerMetrics.setValueOfTableGauge(tableNameWithType, ControllerGauge.SEGMENTS_IN_ERROR_STATE,
         numErrorSegments);
     _controllerMetrics.setValueOfTableGauge(tableNameWithType, ControllerGauge.PERCENT_SEGMENTS_AVAILABLE,
@@ -426,6 +437,14 @@ private void updateSegmentMetrics(String tableNameWithType, TableConfig tableCon
     }
   }
 
+  private boolean isServerQueryable(String serverInstanceId) {
+    ServerInfo serverInfo = _serverInfoFetcher.getServerInfo(serverInstanceId);
+    return serverInfo != null
+        && serverInfo.isHelixEnabled()
+        && !serverInfo.isQueriesDisabled()
+        && !serverInfo.isShutdownInProgress();
+  }
+
   private static String logSegments(List<?> segments) {
     if (segments.size() <= MAX_SEGMENTS_TO_LOG) {
       return segments.toString();

pinot-controller/src/main/java/org/apache/pinot/controller/util/ServerInfoFetcher.java

@@ -0,0 +1,131 @@
+/**
+ * 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
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.controller.util;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.helix.model.InstanceConfig;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.pinot.controller.helix.core.PinotHelixResourceManager;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.apache.pinot.spi.utils.InstanceTypeUtils;
+
+
+/**
+ * This is a helper class that fetch server information from Helix/ZK. It caches the server information to avoid
+ * repeated ZK access. This class is NOT thread-safe.
+ */
+public class ServerInfoFetcher {
+  private PinotHelixResourceManager _pinotHelixResourceManager;
+  private Map<String, ServerInfo> _serverInfoCache;
+
+  public ServerInfoFetcher(PinotHelixResourceManager pinotHelixResourceManager) {
+    _pinotHelixResourceManager = pinotHelixResourceManager;
+    _serverInfoCache = new HashMap<>();
+  }
+
+  public ServerInfo getServerInfo(String instanceId) {
+    return _serverInfoCache.computeIfAbsent(instanceId, this::getServerInfoOndemand);
+  }
+
+  private ServerInfo getServerInfoOndemand(String instanceId) {
+    InstanceConfig instanceConfig = _pinotHelixResourceManager.getHelixInstanceConfig(instanceId);
+    if (instanceConfig == null || !InstanceTypeUtils.isServer(instanceId)) {
+      return null;
+    }
+    List<String> tags = instanceConfig.getTags();
+    ZNRecord record = instanceConfig.getRecord();
+    boolean helixEnabled = record.getBooleanField(
+        InstanceConfig.InstanceConfigProperty.HELIX_ENABLED.name(), false);
+    boolean queriesDisabled = record.getBooleanField(CommonConstants.Helix.QUERIES_DISABLED, false);
+    boolean shutdownInProgress = record.getBooleanField(CommonConstants.Helix.IS_SHUTDOWN_IN_PROGRESS, false);
+    return new ServerInfo(instanceId, tags, null, helixEnabled, queriesDisabled, shutdownInProgress);
+  }
+
+  public static class ServerInfo {
+    private String _instanceName;
+    private List<String> _tags;
+    private List<String> _tables;
+    private boolean _helixEnabled;
+    private boolean _queriesDisabled;
+    private boolean _shutdownInProgress;
+    private ServerInfo(String instanceName,
+        List<String> tags,
+        List<String> tables,
+        boolean helixEnabled,
+        boolean queriesDisabled,
+        boolean shutdownInProgress) {
+      _instanceName = instanceName;
+      _tags = tags;
+      _tables = tables;
+      _helixEnabled = helixEnabled;
+      _queriesDisabled = queriesDisabled;
+      _shutdownInProgress = shutdownInProgress;
+    }
+
+    public String getInstanceName() {
+      return _instanceName;
+    }
+
+    public void setInstanceName(String instanceName) {
+      _instanceName = instanceName;
+    }
+
+    public List<String> getTags() {
+      return _tags;
+    }
+
+    public void setTags(List<String> tags) {
+      _tags = tags;
+    }
+
+    public List<String> getTables() {
+      return _tables;
+    }
+
+    public void setTables(List<String> tables) {
+      _tables = tables;
+    }
+
+    public boolean isHelixEnabled() {
+      return _helixEnabled;
+    }
+
+    public void setHelixEnabled(boolean helixEnabled) {
+      _helixEnabled = helixEnabled;
+    }
+
+    public boolean isQueriesDisabled() {
+      return _queriesDisabled;
+    }
+
+    public void setQueriesDisabled(boolean queriesDisabled) {
+      _queriesDisabled = queriesDisabled;
+    }
+
+    public boolean isShutdownInProgress() {
+      return _shutdownInProgress;
+    }
+
+    public void setShutdownInProgress(boolean shutdownInProgress) {
+      _shutdownInProgress = shutdownInProgress;
+    }
+  }
+}