Adding registerExtraComponents to allow registering additional components in various services

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

@@ -162,4 +162,8 @@ public void stop() {
     LOGGER.info("Shutting down executor service");
     _executorService.shutdownNow();
   }
+
+  public HttpServer getHttpServer() {
+    return _httpServer;
+  }
 }

pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java

@@ -346,9 +346,7 @@ public void start()
       _sqlQueryExecutor = new SqlQueryExecutor(_spectatorHelixManager);
     }
     LOGGER.info("Starting broker admin application on: {}", ListenerConfigUtil.toString(_listenerConfigs));
-    _brokerAdminApplication =
-        new BrokerAdminApiApplication(_routingManager, _brokerRequestHandler, _brokerMetrics, _brokerConf,
-            _sqlQueryExecutor, _serverRoutingStatsManager, _accessControlFactory, _spectatorHelixManager);
+    _brokerAdminApplication = createBrokerAdminApp();
     registerExtraComponents(_brokerAdminApplication);
     _brokerAdminApplication.start(_listenerConfigs);
 
@@ -413,6 +411,7 @@ public void start()
   }
 
   /**
+   * @deprecated Use {@link #createBrokerAdminApp()} instead.
    * This method is called after initialization of BrokerAdminApiApplication object
    * and before calling start to allow custom broker starters to register additional
    * components.
@@ -600,4 +599,9 @@ public AccessControlFactory getAccessControlFactory() {
   public BrokerRequestHandler getBrokerRequestHandler() {
     return _brokerRequestHandler;
   }
+
+  protected BrokerAdminApiApplication createBrokerAdminApp() {
+    return new BrokerAdminApiApplication(_routingManager, _brokerRequestHandler, _brokerMetrics, _brokerConf,
+        _sqlQueryExecutor, _serverRoutingStatsManager, _accessControlFactory, _spectatorHelixManager);
+  }
 }

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

@@ -240,7 +240,7 @@ public void init(PinotConfiguration pinotConfiguration)
       // Initialize FunctionRegistry before starting the admin application (PinotQueryResource requires it to compile
       // queries)
       FunctionRegistry.init();
-      _adminApp = new ControllerAdminApiApplication(_config);
+      _adminApp = createControllerAdminApp();
       // Do not use this before the invocation of {@link PinotHelixResourceManager::start()}, which happens in {@link
       // ControllerStarter::start()}
       _helixResourceManager = new PinotHelixResourceManager(_config);
@@ -927,4 +927,8 @@ public PinotMetricsRegistry getMetricsRegistry() {
   public ControllerMetrics getControllerMetrics() {
     return _controllerMetrics;
   }
+
+  protected ControllerAdminApiApplication createControllerAdminApp() {
+    return new ControllerAdminApiApplication(_config);
+  }
 }

pinot-controller/src/main/java/org/apache/pinot/controller/api/ControllerAdminApiApplication.java

@@ -122,4 +122,8 @@ public void filter(ContainerRequestContext containerRequestContext,
       }
     }
   }
+
+  public HttpServer getHttpServer() {
+    return _httpServer;
+  }
 }

pinot-minion/src/main/java/org/apache/pinot/minion/BaseMinionStarter.java

@@ -280,7 +280,7 @@ public void start()
     minionContext.setHelixPropertyStore(_helixManager.getHelixPropertyStore());
     minionContext.setHelixManager(_helixManager);
     LOGGER.info("Starting minion admin application on: {}", ListenerConfigUtil.toString(_listenerConfigs));
-    _minionAdminApplication = new MinionAdminApiApplication(_instanceId, _config);
+    _minionAdminApplication = createMinionAdminApp();
     _minionAdminApplication.start(_listenerConfigs);
 
     // Initialize health check callback
@@ -347,4 +347,8 @@ public void stop() {
     }
     LOGGER.info("Pinot minion stopped");
   }
+
+  protected MinionAdminApiApplication createMinionAdminApp() {
+    return new MinionAdminApiApplication(_instanceId, _config);
+  }
 }

pinot-minion/src/main/java/org/apache/pinot/minion/MinionAdminApiApplication.java

@@ -95,4 +95,8 @@ public void stop() {
       _httpServer.shutdownNow();
     }
   }
+
+  public HttpServer getHttpServer() {
+    return _httpServer;
+  }
 }

pinot-server/src/main/java/org/apache/pinot/server/api/AdminApiApplication.java

@@ -141,4 +141,8 @@ public void startShuttingDown() {
   public void stop() {
     _httpServer.shutdownNow();
   }
+
+  public HttpServer getHttpServer() {
+    return _httpServer;
+  }
 }

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

@@ -139,6 +139,7 @@ public abstract class BaseServerStarter implements ServiceStartable {
   protected HelixManager _helixManager;
   protected HelixAdmin _helixAdmin;
   protected ServerInstance _serverInstance;
+  protected AccessControlFactory _accessControlFactory;
   protected AdminApiApplication _adminApiApplication;
   protected ServerQueriesDisabledTracker _serverQueriesDisabledTracker;
   protected RealtimeLuceneIndexRefreshState _realtimeLuceneIndexRefreshState;
@@ -572,9 +573,8 @@ public void start()
     String accessControlFactoryClass =
         _serverConf.getProperty(Server.ACCESS_CONTROL_FACTORY_CLASS, Server.DEFAULT_ACCESS_CONTROL_FACTORY_CLASS);
     LOGGER.info("Using class: {} as the AccessControlFactory", accessControlFactoryClass);
-    AccessControlFactory accessControlFactory;
     try {
-      accessControlFactory = PluginManager.get().createInstance(accessControlFactoryClass);
+      _accessControlFactory = PluginManager.get().createInstance(accessControlFactoryClass);
     } catch (Exception e) {
       throw new RuntimeException(
           "Caught exception while creating new AccessControlFactory instance using class '" + accessControlFactoryClass
@@ -593,7 +593,7 @@ public void start()
     ServerSegmentCompletionProtocolHandler.init(
         _serverConf.subset(SegmentCompletionProtocol.PREFIX_OF_CONFIG_OF_SEGMENT_UPLOADER));
     ServerConf serverConf = new ServerConf(_serverConf);
-    _serverInstance = new ServerInstance(serverConf, _helixManager, accessControlFactory);
+    _serverInstance = new ServerInstance(serverConf, _helixManager, _accessControlFactory);
     ServerMetrics serverMetrics = _serverInstance.getServerMetrics();
 
     InstanceDataManager instanceDataManager = _serverInstance.getInstanceDataManager();
@@ -617,7 +617,7 @@ public void start()
 
     // Start restlet server for admin API endpoint
     LOGGER.info("Starting server admin application on: {}", ListenerConfigUtil.toString(_listenerConfigs));
-    _adminApiApplication = new AdminApiApplication(_serverInstance, accessControlFactory, _serverConf);
+    _adminApiApplication = createServerAdminApp();
     _adminApiApplication.start(_listenerConfigs);
 
     // Init QueryRewriterFactory
@@ -931,4 +931,8 @@ private void initSegmentFetcher(PinotConfiguration config)
     PinotConfiguration pinotCrypterConfig = config.subset(CommonConstants.Server.PREFIX_OF_CONFIG_OF_PINOT_CRYPTER);
     PinotCrypterFactory.init(pinotCrypterConfig);
   }
+
+  protected AdminApiApplication createServerAdminApp() {
+    return new AdminApiApplication(_serverInstance, _accessControlFactory, _serverConf);
+  }
 }