Add bootstrap mode for PinotServiceManager to avoid glitch for health…

… check (#7880)

* Add bootstrap mode for PinotServiceManager to avoid glitch for health check

* Adding healthcheck with instance id

pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerHealthCheck.java

@@ -23,12 +23,14 @@
 import io.swagger.annotations.ApiResponse;
 import io.swagger.annotations.ApiResponses;
 import javax.inject.Inject;
+import javax.inject.Named;
 import javax.ws.rs.GET;
 import javax.ws.rs.Path;
 import javax.ws.rs.Produces;
 import javax.ws.rs.WebApplicationException;
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.Response;
+import org.apache.pinot.broker.broker.BrokerAdminApiApplication;
 import org.apache.pinot.common.metrics.BrokerMeter;
 import org.apache.pinot.common.metrics.BrokerMetrics;
 import org.apache.pinot.common.utils.ServiceStatus;
@@ -37,6 +39,9 @@
 @Api(tags = "Health")
 @Path("/")
 public class PinotBrokerHealthCheck {
+  @Inject
+  @Named(BrokerAdminApiApplication.BROKER_INSTANCE_ID)
+  private String _instanceId;
 
   @Inject
   private BrokerMetrics _brokerMetrics;
@@ -50,7 +55,7 @@ public class PinotBrokerHealthCheck {
       @ApiResponse(code = 503, message = "Broker is not healthy")
   })
   public String getBrokerHealth() {
-    ServiceStatus.Status status = ServiceStatus.getServiceStatus();
+    ServiceStatus.Status status = ServiceStatus.getServiceStatus(_instanceId);
     if (status == ServiceStatus.Status.GOOD) {
       _brokerMetrics.addMeteredGlobalValue(BrokerMeter.HEALTHCHECK_OK_CALLS, 1);
       return "OK";

pinot-broker/src/main/java/org/apache/pinot/broker/broker/BrokerAdminApiApplication.java

@@ -42,6 +42,7 @@
 public class BrokerAdminApiApplication extends ResourceConfig {
   private static final String RESOURCE_PACKAGE = "org.apache.pinot.broker.api.resources";
   public static final String PINOT_CONFIGURATION = "pinotConfiguration";
+  public static final String BROKER_INSTANCE_ID = "brokerInstanceId";
 
   private HttpServer _httpServer;
 
@@ -56,6 +57,7 @@ protected void configure() {
         bind(routingManager).to(RoutingManager.class);
         bind(brokerRequestHandler).to(BrokerRequestHandler.class);
         bind(brokerMetrics).to(BrokerMetrics.class);
+        bind(brokerConf.getProperty(CommonConstants.Broker.CONFIG_OF_BROKER_ID)).named(BROKER_INSTANCE_ID);
       }
     });
     register(JacksonFeature.class);

pinot-controller/src/main/java/org/apache/pinot/controller/BaseControllerStarter.java

@@ -110,6 +110,7 @@
 public abstract class BaseControllerStarter implements ServiceStartable {
   private static final Logger LOGGER = LoggerFactory.getLogger(BaseControllerStarter.class);
 
+  public static final String CONTROLLER_INSTANCE_ID = "controllerInstanceId";
   private static final String METRICS_REGISTRY_NAME = "pinot.controller.metrics";
   private static final Long DATA_DIRECTORY_MISSING_VALUE = 1000000L;
   private static final Long DATA_DIRECTORY_EXCEPTION_VALUE = 1100000L;
@@ -445,6 +446,7 @@ private void setUpPinotController() {
       @Override
       protected void configure() {
         bind(_config).to(ControllerConf.class);
+        bind(_helixParticipantInstanceId).named(CONTROLLER_INSTANCE_ID);
         bind(_helixResourceManager).to(PinotHelixResourceManager.class);
         bind(_helixTaskResourceManager).to(PinotHelixTaskResourceManager.class);
         bind(_segmentCompletionManager).to(SegmentCompletionManager.class);

pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotControllerHealthCheck.java

@@ -23,6 +23,7 @@
 import io.swagger.annotations.ApiResponse;
 import io.swagger.annotations.ApiResponses;
 import javax.inject.Inject;
+import javax.inject.Named;
 import javax.ws.rs.GET;
 import javax.ws.rs.Path;
 import javax.ws.rs.Produces;
@@ -33,13 +34,18 @@
 import org.apache.pinot.common.metrics.ControllerMeter;
 import org.apache.pinot.common.metrics.ControllerMetrics;
 import org.apache.pinot.common.utils.ServiceStatus;
+import org.apache.pinot.controller.BaseControllerStarter;
 import org.apache.pinot.controller.ControllerConf;
 
 
 @Api(tags = Constants.HEALTH_TAG)
 @Path("/")
 public class PinotControllerHealthCheck {
 
+  @Inject
+  @Named(BaseControllerStarter.CONTROLLER_INSTANCE_ID)
+  private String _instanceId;
+
   @Inject
   ControllerConf _controllerConf;
 
@@ -64,7 +70,7 @@ public String checkHealthLegacy() {
   @ApiResponses(value = {@ApiResponse(code = 200, message = "Good")})
   @Produces(MediaType.TEXT_PLAIN)
   public String checkHealth() {
-    ServiceStatus.Status status = ServiceStatus.getServiceStatus();
+    ServiceStatus.Status status = ServiceStatus.getServiceStatus(_instanceId);
     if (status == ServiceStatus.Status.GOOD) {
       _controllerMetrics.addMeteredGlobalValue(ControllerMeter.HEALTHCHECK_OK_CALLS, 1);
       return "OK";

pinot-server/src/main/java/org/apache/pinot/server/api/resources/HealthCheckResource.java

@@ -22,6 +22,8 @@
 import io.swagger.annotations.ApiOperation;
 import io.swagger.annotations.ApiResponse;
 import io.swagger.annotations.ApiResponses;
+import javax.inject.Inject;
+import javax.inject.Named;
 import javax.ws.rs.GET;
 import javax.ws.rs.Path;
 import javax.ws.rs.Produces;
@@ -30,6 +32,7 @@
 import javax.ws.rs.core.Response;
 import org.apache.pinot.common.utils.ServiceStatus;
 import org.apache.pinot.common.utils.ServiceStatus.Status;
+import org.apache.pinot.server.starter.helix.AdminApiApplication;
 
 
 /**
@@ -39,6 +42,10 @@
 @Path("/")
 public class HealthCheckResource {
 
+  @Inject
+  @Named(AdminApiApplication.SERVER_INSTANCE_ID)
+  private String _instanceId;
+
   @GET
   @Path("/health")
   @Produces(MediaType.TEXT_PLAIN)
@@ -48,7 +55,7 @@ public class HealthCheckResource {
       @ApiResponse(code = 503, message = "Server is not healthy")
   })
   public String checkHealth() {
-    Status status = ServiceStatus.getServiceStatus();
+    Status status = ServiceStatus.getServiceStatus(_instanceId);
     if (status == Status.GOOD) {
       return "OK";
     }

pinot-server/src/main/java/org/apache/pinot/server/starter/helix/AdminApiApplication.java

@@ -51,6 +51,7 @@ public class AdminApiApplication extends ResourceConfig {
   private static final Logger LOGGER = LoggerFactory.getLogger(AdminApiApplication.class);
   public static final String PINOT_CONFIGURATION = "pinotConfiguration";
   public static final String RESOURCE_PACKAGE = "org.apache.pinot.server.api.resources";
+  public static final String SERVER_INSTANCE_ID = "serverInstanceId";
 
   private final ServerInstance _serverInstance;
   private final AccessControlFactory _accessControlFactory;
@@ -69,6 +70,7 @@ public AdminApiApplication(ServerInstance instance, AccessControlFactory accessC
       protected void configure() {
         bind(_serverInstance).to(ServerInstance.class);
         bind(accessControlFactory).to(AccessControlFactory.class);
+        bind(serverConf.getProperty(CommonConstants.Server.CONFIG_OF_INSTANCE_ID)).named(SERVER_INSTANCE_ID);
       }
     });
 

pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java

@@ -337,7 +337,7 @@ private void startupServiceStatusCheck(long endTimeMs) {
         Server.DEFAULT_STARTUP_SERVICE_STATUS_CHECK_INTERVAL_MS);
 
     while (System.currentTimeMillis() < endTimeMs) {
-      Status serviceStatus = ServiceStatus.getServiceStatus();
+      Status serviceStatus = ServiceStatus.getServiceStatus(_instanceId);
       long currentTimeMs = System.currentTimeMillis();
       if (serviceStatus == Status.GOOD) {
         LOGGER.info("Service status is GOOD after {}ms", currentTimeMs - startTimeMs);

pinot-server/src/main/java/org/apache/pinot/server/starter/helix/DefaultHelixStarterServerConfig.java

@@ -68,7 +68,6 @@ public static PinotConfiguration loadDefaultServerConf() {
     // netty port
     serverConf
         .addProperty(CommonConstants.Server.CONFIG_OF_NETTY_PORT, CommonConstants.Helix.DEFAULT_SERVER_NETTY_PORT);
-
     return serverConf;
   }
 }

pinot-server/src/test/java/org/apache/pinot/server/api/AccessControlTest.java

@@ -66,6 +66,13 @@ public void setUp()
     ServerInstance serverInstance = mock(ServerInstance.class);
 
     PinotConfiguration serverConf = DefaultHelixStarterServerConfig.loadDefaultServerConf();
+    String hostname = serverConf.getProperty(CommonConstants.Helix.KEY_OF_SERVER_NETTY_HOST,
+        serverConf.getProperty(CommonConstants.Helix.SET_INSTANCE_ID_TO_HOSTNAME_KEY, false)
+            ? NetUtils.getHostnameOrAddress() : NetUtils.getHostAddress());
+    int port = serverConf.getProperty(CommonConstants.Helix.KEY_OF_SERVER_NETTY_PORT,
+        CommonConstants.Helix.DEFAULT_SERVER_NETTY_PORT);
+    serverConf.setProperty(CommonConstants.Server.CONFIG_OF_INSTANCE_ID,
+        CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE + hostname + "_" + port);
     _adminApiApplication = new AdminApiApplication(serverInstance, new DenyAllAccessFactory(), serverConf);
     _adminApiApplication.start(Collections.singletonList(
         new ListenerConfig(CommonConstants.HTTP_PROTOCOL, "0.0.0.0", CommonConstants.Server.DEFAULT_ADMIN_API_PORT,

pinot-server/src/test/java/org/apache/pinot/server/api/BaseResourceTest.java

@@ -127,6 +127,13 @@ public void setUp()
     setUpSegment(offlineTableName, null, "default", _offlineIndexSegments);
 
     PinotConfiguration serverConf = DefaultHelixStarterServerConfig.loadDefaultServerConf();
+    String hostname = serverConf.getProperty(CommonConstants.Helix.KEY_OF_SERVER_NETTY_HOST,
+        serverConf.getProperty(CommonConstants.Helix.SET_INSTANCE_ID_TO_HOSTNAME_KEY, false)
+            ? NetUtils.getHostnameOrAddress() : NetUtils.getHostAddress());
+    int port = serverConf.getProperty(CommonConstants.Helix.KEY_OF_SERVER_NETTY_PORT,
+        CommonConstants.Helix.DEFAULT_SERVER_NETTY_PORT);
+    serverConf.setProperty(CommonConstants.Server.CONFIG_OF_INSTANCE_ID,
+        CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE + hostname + "_" + port);
     _adminApiApplication = new AdminApiApplication(serverInstance, new AllowAllAccessFactory(), serverConf);
     _adminApiApplication.start(Collections.singletonList(
         new ListenerConfig(CommonConstants.HTTP_PROTOCOL, "0.0.0.0", CommonConstants.Server.DEFAULT_ADMIN_API_PORT,

pinot-server/src/test/java/org/apache/pinot/server/api/PinotServerAppConfigsTest.java

@@ -20,10 +20,14 @@
 
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import java.net.SocketException;
+import java.net.UnknownHostException;
 import javax.ws.rs.core.Response;
 import org.apache.pinot.common.utils.PinotAppConfigs;
 import org.apache.pinot.server.starter.helix.DefaultHelixStarterServerConfig;
 import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.apache.pinot.spi.utils.NetUtils;
 import org.apache.pinot.spi.utils.Obfuscator;
 import org.testng.Assert;
 import org.testng.annotations.Test;
@@ -41,8 +45,15 @@ public class PinotServerAppConfigsTest extends BaseResourceTest {
    */
   @Test
   public void testAppConfigs()
-      throws JsonProcessingException {
+      throws JsonProcessingException, SocketException, UnknownHostException {
     PinotConfiguration expectedServerConf = DefaultHelixStarterServerConfig.loadDefaultServerConf();
+    String hostname = expectedServerConf.getProperty(CommonConstants.Helix.KEY_OF_SERVER_NETTY_HOST,
+        expectedServerConf.getProperty(CommonConstants.Helix.SET_INSTANCE_ID_TO_HOSTNAME_KEY, false)
+            ? NetUtils.getHostnameOrAddress() : NetUtils.getHostAddress());
+    int port = expectedServerConf.getProperty(CommonConstants.Helix.KEY_OF_SERVER_NETTY_PORT,
+        CommonConstants.Helix.DEFAULT_SERVER_NETTY_PORT);
+    expectedServerConf.setProperty(CommonConstants.Server.CONFIG_OF_INSTANCE_ID,
+        CommonConstants.Helix.PREFIX_OF_SERVER_INSTANCE + hostname + "_" + port);
     PinotAppConfigs expected = new PinotAppConfigs(expectedServerConf);
 
     Response response = _webTarget.path("/appconfigs").request().get(Response.class);

pinot-tools/src/main/java/org/apache/pinot/tools/service/PinotServiceManager.java

@@ -54,6 +54,7 @@ public class PinotServiceManager {
   private final String _instanceId;
   private PinotServiceManagerAdminApiApplication _pinotServiceManagerAdminApplication;
   private boolean _isStarted = false;
+  private boolean _isShuttingDown = false;
 
   public PinotServiceManager(String zkAddress, String clusterName) {
     this(zkAddress, clusterName, 0);
@@ -76,11 +77,6 @@ public PinotServiceManager(String zkAddress, String clusterName, String hostname
     _instanceId = String.format("ServiceManager_%s_%d", hostname, port);
   }
 
-  public static void main(String[] args) {
-    PinotServiceManager pinotServiceManager = new PinotServiceManager("localhost:2181", "pinot-demo", 8085);
-    pinotServiceManager.start();
-  }
-
   public String startRole(ServiceRole role, Map<String, Object> properties)
       throws Exception {
     switch (role) {
@@ -214,7 +210,7 @@ public void start() {
     LOGGER.info("Registering service status handler");
     ServiceStatus.setServiceStatusCallback(_instanceId, new PinotServiceManagerStatusCallback(this));
 
-    if (_port < 0) {
+    if (_port <= 0) {
       LOGGER.info("Skip Starting Pinot Service Manager admin application");
     } else {
       LOGGER.info("Starting Pinot Service Manager admin application on port: {}", _port);
@@ -235,6 +231,7 @@ public void stop() {
 
   public void stopAll() {
     LOGGER.info("Shutting down Pinot Service Manager with all running Pinot instances...");
+    _isShuttingDown = true;
     for (String instanceId : _runningInstanceMap.keySet()) {
       stopPinotInstanceById(instanceId);
     }
@@ -245,6 +242,10 @@ public boolean isStarted() {
     return _isStarted;
   }
 
+  public boolean isShuttingDown() {
+    return _isShuttingDown;
+  }
+
   public String getInstanceId() {
     return _instanceId;
   }

pinot-tools/src/main/java/org/apache/pinot/tools/service/PinotServiceManagerStatusCallback.java

@@ -30,19 +30,23 @@ public PinotServiceManagerStatusCallback(PinotServiceManager pinotServiceManager
 
   @Override
   public ServiceStatus.Status getServiceStatus() {
+    if (_pinotServiceManager.isShuttingDown()) {
+      return ServiceStatus.Status.SHUTTING_DOWN;
+    }
     if (_pinotServiceManager.isStarted()) {
       return ServiceStatus.Status.GOOD;
-    } else {
-      return ServiceStatus.Status.STARTING;
     }
+    return ServiceStatus.Status.STARTING;
   }
 
   @Override
   public String getStatusDescription() {
+    if (_pinotServiceManager.isShuttingDown()) {
+      return ServiceStatus.STATUS_DESCRIPTION_SHUTTING_DOWN;
+    }
     if (_pinotServiceManager.isStarted()) {
       return ServiceStatus.STATUS_DESCRIPTION_STARTED;
-    } else {
-      return ServiceStatus.STATUS_DESCRIPTION_INIT;
     }
+    return ServiceStatus.STATUS_DESCRIPTION_INIT;
   }
 }