Enable service discovery in controller too

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

@@ -26,6 +26,7 @@
 import org.apache.pinot.broker.requesthandler.BrokerRequestHandler;
 import org.apache.pinot.broker.routing.RoutingManager;
 import org.apache.pinot.common.metrics.BrokerMetrics;
+import org.apache.pinot.core.api.ServiceAutoDiscoveryFeature;
 import org.apache.pinot.core.transport.ListenerConfig;
 import org.apache.pinot.core.util.ListenerConfigUtil;
 import org.apache.pinot.spi.env.PinotConfiguration;
@@ -51,7 +52,7 @@ public BrokerAdminApiApplication(RoutingManager routingManager, BrokerRequestHan
     packages(RESOURCE_PACKAGE);
     property(PINOT_CONFIGURATION, brokerConf);
     if (brokerConf.getProperty(CommonConstants.Broker.BROKER_SERVICE_AUTO_DISCOVERY, false)) {
-      register(BrokerServiceAutoDiscoveryFeature.class);
+      register(ServiceAutoDiscoveryFeature.class);
     }
     register(new AbstractBinder() {
       @Override

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

@@ -28,6 +28,7 @@
 import javax.ws.rs.container.ContainerResponseFilter;
 import org.apache.pinot.controller.ControllerConf;
 import org.apache.pinot.controller.api.access.AuthenticationFilter;
+import org.apache.pinot.core.api.ServiceAutoDiscoveryFeature;
 import org.apache.pinot.core.transport.ListenerConfig;
 import org.apache.pinot.core.util.ListenerConfigUtil;
 import org.apache.pinot.spi.utils.CommonConstants;
@@ -57,6 +58,9 @@ public ControllerAdminApiApplication(ControllerConf conf) {
     packages(_controllerResourcePackages);
     // TODO See ControllerResponseFilter
 //    register(new LoggingFeature());
+    if (conf.getProperty(CommonConstants.Controller.CONTROLLER_SERVICE_AUTO_DISCOVERY, false)) {
+      register(ServiceAutoDiscoveryFeature.class);
+    }
     register(JacksonFeature.class);
     register(MultiPartFeature.class);
     registerClasses(io.swagger.jaxrs.listing.ApiListingResource.class);

pinot-core/pom.xml

@@ -179,6 +179,10 @@
       <groupId>org.glassfish.grizzly</groupId>
       <artifactId>grizzly-http-server</artifactId>
     </dependency>
+    <dependency>
+      <groupId>org.glassfish.hk2</groupId>
+      <artifactId>hk2-locator</artifactId>
+    </dependency>
 
     <!-- test -->
     <dependency>

pinot-broker/src/main/java/org/apache/pinot/broker/broker/BrokerServiceAutoDiscoveryFeature.java → pinot-core/src/main/java/org/apache/pinot/core/api/ServiceAutoDiscoveryFeature.java

@@ -16,7 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-package org.apache.pinot.broker.broker;
+package org.apache.pinot.core.api;
 
 import java.io.IOException;
 import javax.inject.Inject;
@@ -31,6 +31,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+
 /**
  * Auto scan and discovery classes annotated with Service tags. This enables the feature similar to Spring's
  * auto-discovery, that if a class is annotated with some tags like @Service, the class will be auto-loaded
@@ -68,8 +69,8 @@
  * pinot-integration-tests/src/main/java/org/apache/pinot/broker/integration/tests/BrokerTestAutoLoadedService.java
  * </code>
  */
-public class BrokerServiceAutoDiscoveryFeature implements Feature {
-    private static final Logger LOGGER = LoggerFactory.getLogger(BrokerServiceAutoDiscoveryFeature.class);
+public class ServiceAutoDiscoveryFeature implements Feature {
+    private static final Logger LOGGER = LoggerFactory.getLogger(ServiceAutoDiscoveryFeature.class);
 
     @Inject
     ServiceLocator _serviceLocator;

pinot-integration-tests/src/main/java/org/apache/pinot/broker/api/resources/EchoWithAutoDiscovery.java → pinot-integration-tests/src/main/java/org/apache/pinot/broker/api/resources/BrokerEchoWithAutoDiscovery.java

@@ -26,7 +26,7 @@
 import javax.ws.rs.PathParam;
 import javax.ws.rs.Produces;
 import javax.ws.rs.core.MediaType;
-import org.apache.pinot.broker.integration.tests.BrokerTestAutoLoadedService;
+import org.apache.pinot.core.api.AutoLoadedServiceForTest;
 
 /**
  * This class is a typical "echo" service that will return whatever string you call GET with a path.
@@ -35,9 +35,9 @@
  */
 @Api(tags = "Test")
 @Path("/test")
-public class EchoWithAutoDiscovery {
+public class BrokerEchoWithAutoDiscovery {
     @Inject
-    public BrokerTestAutoLoadedService _injectedService;
+    public AutoLoadedServiceForTest _injectedService;
     @GET
     @Path("/echo/{table}")
     @Produces(MediaType.TEXT_PLAIN)

pinot-integration-tests/src/main/java/org/apache/pinot/controller/api/resources/ControllerEchoWithAutoDiscovery.java

@@ -0,0 +1,48 @@
+/**
+ * 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.api.resources;
+
+import io.swagger.annotations.Api;
+import javax.inject.Inject;
+import javax.ws.rs.GET;
+import javax.ws.rs.Path;
+import javax.ws.rs.PathParam;
+import javax.ws.rs.Produces;
+import javax.ws.rs.core.MediaType;
+import org.apache.pinot.core.api.AutoLoadedServiceForTest;
+
+
+/**
+ * This class is a typical "echo" service that will return whatever string you call GET with a path.
+ * It is both an integration test and a demonstration of how to dynamically add an endpoint to broker,
+ * create auto-service discovery
+ */
+@Api(tags = "Test")
+@Path("/test")
+public class ControllerEchoWithAutoDiscovery {
+    @Inject
+    public AutoLoadedServiceForTest _injectedService;
+    @GET
+    @Path("/echo/{table}")
+    @Produces(MediaType.TEXT_PLAIN)
+    public String echo(@PathParam("table") String table) {
+        return _injectedService.echo(table);
+    }
+}

pinot-integration-tests/src/main/java/org/apache/pinot/broker/integration/tests/BrokerTestAutoLoadedService.java → pinot-integration-tests/src/main/java/org/apache/pinot/core/api/AutoLoadedServiceForTest.java

@@ -16,14 +16,14 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-package org.apache.pinot.broker.integration.tests;
+package org.apache.pinot.core.api;
 
 import javax.inject.Singleton;
 import org.jvnet.hk2.annotations.Service;
 
 @Service
 @Singleton
-public class BrokerTestAutoLoadedService {
+public class AutoLoadedServiceForTest {
     public String echo(String echoText) {
         return echoText;
     }

pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/BrokerServiceDiscoveryIntegrationTest.java

@@ -28,7 +28,7 @@
 import org.testng.annotations.Test;
 
 /**
- * Integration test that converts Avro data for 12 segments and runs queries against it.
+ * Integration test that starts one broker with auto-discovered echo service and test it
  */
 public class BrokerServiceDiscoveryIntegrationTest extends BaseClusterIntegrationTestSet {
   private static final String TENANT_NAME = "TestTenant";

pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ControllerServiceDiscoveryIntegrationTest.java

@@ -0,0 +1,92 @@
+/**
+ * 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.integration.tests;
+
+import java.util.Map;
+import org.apache.commons.io.FileUtils;
+import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.apache.pinot.util.TestUtils;
+import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+
+/**
+ * Integration test that starts one broker with auto-discovered echo service and test it
+ */
+public class ControllerServiceDiscoveryIntegrationTest extends BaseClusterIntegrationTestSet {
+  private static final String TENANT_NAME = "TestTenant";
+
+  @Override
+  protected String getBrokerTenant() {
+    return TENANT_NAME;
+  }
+
+  @Override
+  protected String getServerTenant() {
+    return TENANT_NAME;
+  }
+
+  @Override
+  public Map<String, Object> getDefaultControllerConfiguration() {
+    Map<String, Object> retVal = super.getDefaultControllerConfiguration();
+    retVal.put(CommonConstants.Controller.CONTROLLER_SERVICE_AUTO_DISCOVERY, true);
+    return retVal;
+  }
+
+  @Override
+  protected PinotConfiguration getDefaultBrokerConfiguration() {
+    PinotConfiguration config = new PinotConfiguration();
+    config.setProperty(CommonConstants.Broker.BROKER_SERVICE_AUTO_DISCOVERY, true);
+    return config;
+  }
+
+  @BeforeClass
+  public void setUp()
+      throws Exception {
+    TestUtils.ensureDirectoriesExistAndEmpty(_tempDir, _segmentDir, _tarDir);
+
+    // Start the Pinot cluster
+    startZk();
+    startController();
+    startBrokers(1);
+    startServers(1);
+  }
+  @AfterClass
+  public void tearDown()
+          throws Exception {
+
+    // Brokers and servers has been stopped
+    stopBroker();
+    stopController();
+    stopZk();
+    FileUtils.deleteDirectory(_tempDir);
+  }
+
+  @Test
+  public void testControllerExtraEndpointsAutoLoaded()
+      throws Exception {
+    String response = sendGetRequest(_controllerBaseApiUrl + "/test/echo/doge");
+    Assert.assertEquals(response, "doge");
+    response = sendGetRequest(_brokerBaseApiUrl + "/test/echo/doge");
+    Assert.assertEquals(response, "doge");
+  }
+}

pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java

@@ -442,6 +442,8 @@ public static class Controller {
     public static final String CONFIG_OF_INSTANCE_ID = "pinot.controller.instance.id";
     public static final String CONFIG_OF_CONTROLLER_QUERY_REWRITER_CLASS_NAMES =
         "pinot.controller.query.rewriter.class.names";
+    //Set to true to load all services tagged and compiled with hk2-metadata-generator. Default to False
+    public static final String CONTROLLER_SERVICE_AUTO_DISCOVERY = "pinot.controller.service.auto.discovery";
   }
 
   public static class Minion {