Update 'instances' API with ability to filter by tag (#12585)

Co-authored-by: Xiaotian (Jackie) Jiang <17555551+Jackie-Jiang@users.noreply.github.com>
apache · Mar 8, 2024 · 88c901a · 88c901a
1 parent c7cc821
commit 88c901a
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 1 deletion.
diff --git a/...src/main/java/org/apache/pinot/controller/api/resources/PinotInstanceRestletResource.java b/...src/main/java/org/apache/pinot/controller/api/resources/PinotInstanceRestletResource.java
@@ -104,7 +104,11 @@ public List<String> getInstances() {
       @ApiResponse(code = 200, message = "Success"),
       @ApiResponse(code = 500, message = "Internal error")
   })
-  public Instances getAllInstances() {
+  public Instances getAllInstances(
+      @ApiParam("Filter by tag") @QueryParam("tag") String tag) {
+    if (tag != null) {
+      return new Instances(_pinotHelixResourceManager.getAllInstancesWithTag(tag));
+    }
     return new Instances(_pinotHelixResourceManager.getAllInstances());
   }
 

diff --git a/...oller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java b/...oller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java
@@ -414,6 +414,15 @@ public List<String> getAllInstances() {
     return _helixAdmin.getInstancesInCluster(_helixClusterName);
   }
 
+  /**
+   * Get Ids of all instance with the given tag.
+   *
+   * @return List of instance Ids
+   */
+  public List<String> getAllInstancesWithTag(String tag) {
+    return HelixHelper.getInstancesWithTag(_helixZkManager, tag);
+  }
+
   /**
    * Get all live instance Ids.
    *

diff --git a/...t/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManagerStatelessTest.java b/...t/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManagerStatelessTest.java
@@ -147,6 +147,17 @@ private void resetServerTags() {
     assertEquals(_helixResourceManager.getOnlineUnTaggedServerInstanceList().size(), 0);
   }
 
+  @Test
+  public void testGetInstancesByTag() {
+    List<String> controllersByTag = _helixResourceManager.getAllInstancesWithTag("controller");
+    List<InstanceConfig> controllerConfigs = _helixResourceManager.getAllControllerInstanceConfigs();
+
+    assertEquals(controllersByTag.size(), controllerConfigs.size());
+    for (InstanceConfig c: controllerConfigs) {
+      assertTrue(controllersByTag.contains(c.getInstanceName()));
+    }
+  }
+
   @Test
   public void testGetDataInstanceAdminEndpoints()
       throws Exception {