Copilot commented on code in PR #17375:
URL: https://github.com/apache/pinot/pull/17375#discussion_r2622896777
##########
pinot-controller/src/test/java/org/apache/pinot/controller/api/PinotInstanceRestletResourceTest.java:
##########
@@ -296,6 +297,184 @@ private List<String> getInstanceTags(String instance)
return tags;
}
+ @Test
+ public void testDrainMinionInstance()
+ throws Exception {
+ // Create a minion instance with minion_untagged tag
+ String createInstanceUrl = _urlBuilder.forInstanceCreate();
+ Instance minionInstance =
+ new Instance("minion1.test.com", 9514, InstanceType.MINION,
+ Collections.singletonList(Helix.UNTAGGED_MINION_INSTANCE),
+ null, 0, 0, 0, 0, false);
+ sendPostRequest(createInstanceUrl, minionInstance.toJsonString());
+ String minionInstanceId = "Minion_minion1.test.com_9514";
+
+ // Verify the minion was created with minion_untagged tag
+ checkInstanceInfo(minionInstanceId, "minion1.test.com", 9514, new
String[]{Helix.UNTAGGED_MINION_INSTANCE},
+ null, -1, -1, -1, -1, false);
+
+ // Drain the minion instance
+ String drainUrl = _urlBuilder.forInstanceState(minionInstanceId);
+ sendPutRequest(drainUrl + "?state=DRAIN", "");
+
+ // Verify the minion now has minion_drained tag instead of minion_untagged
+ JsonNode response =
JsonUtils.stringToJsonNode(sendGetRequest(_urlBuilder.forInstance(minionInstanceId)));
+ assertEquals(response.get("instanceName").asText(), minionInstanceId);
+ JsonNode tags = response.get("tags");
+ assertEquals(tags.size(), 1);
+ assertEquals(tags.get(0).asText(), Helix.DRAINED_MINION_INSTANCE);
+
+ // Cleanup - delete the minion instance
+ sendDeleteRequest(_urlBuilder.forInstance(minionInstanceId));
+ }
+
+ @Test
+ public void testDrainMinionInstanceWithMultipleTags()
+ throws Exception {
+ // Create a minion instance with multiple tags including minion_untagged
+ String createInstanceUrl = _urlBuilder.forInstanceCreate();
+ List<String> tags = Arrays.asList(Helix.UNTAGGED_MINION_INSTANCE,
"custom_tag1", "custom_tag2");
+ Instance minionInstance =
+ new Instance("minion2.test.com", 9515, InstanceType.MINION, tags,
null, 0, 0, 0, 0, false);
+ sendPostRequest(createInstanceUrl, minionInstance.toJsonString());
+ String minionInstanceId = "Minion_minion2.test.com_9515";
+
+ // Drain the minion instance
+ String drainUrl = _urlBuilder.forInstanceState(minionInstanceId);
+ sendPutRequest(drainUrl + "?state=DRAIN", "");
+
+ // Verify minion_untagged was replaced with minion_drained, other tags
remain
+ JsonNode response =
JsonUtils.stringToJsonNode(sendGetRequest(_urlBuilder.forInstance(minionInstanceId)));
+ JsonNode responseTags = response.get("tags");
+ assertEquals(responseTags.size(), 3);
+ List<String> actualTags = new ArrayList<>();
+ for (JsonNode tag : responseTags) {
+ actualTags.add(tag.asText());
+ }
+ assertTrue(actualTags.contains(Helix.DRAINED_MINION_INSTANCE));
+ assertTrue(actualTags.contains("custom_tag1"));
+ assertTrue(actualTags.contains("custom_tag2"));
+ assertFalse(actualTags.contains(Helix.UNTAGGED_MINION_INSTANCE));
+
+ // Cleanup
+ sendDeleteRequest(_urlBuilder.forInstance(minionInstanceId));
+ }
+
+ @Test
+ public void testDrainMinionInstanceWithoutUntaggedTag()
+ throws Exception {
+ // Create a minion instance without minion_untagged tag
+ String createInstanceUrl = _urlBuilder.forInstanceCreate();
+ List<String> tags = List.of("custom_tag");
Review Comment:
Using List.of() creates an immutable list. While this works for the test,
consider using Arrays.asList() or Collections.singletonList() for consistency
with other test methods in this file (see lines 336, 398).
```suggestion
List<String> tags = Arrays.asList("custom_tag");
```
##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java:
##########
@@ -3624,6 +3624,50 @@ public PinotResourceManagerResponse dropInstance(String
instanceName) {
return PinotResourceManagerResponse.success("Instance " + instanceName + "
dropped");
}
+ /**
+ * Drains a minion instance by preventing new task assignments while
allowing existing tasks to complete.
+ * This is achieved by changing the instance tag from minion_untagged to
minion_drained, which prevents
+ * Helix from assigning new tasks to this instance.
+ *
+ * @param instanceName Name of the minion instance to drain
+ * @return Response indicating success or failure
+ */
+ public synchronized PinotResourceManagerResponse drainMinionInstance(String
instanceName) {
+ InstanceConfig instanceConfig = getHelixInstanceConfig(instanceName);
+ if (instanceConfig == null) {
+ return PinotResourceManagerResponse.failure("Instance " + instanceName +
" not found");
+ }
+
+ // Replace minion_untagged with minion_drained tag
+ List<String> updatedTags = replaceMinionTag(instanceConfig.getTags());
+ instanceConfig.getRecord().setListField(
+ InstanceConfig.InstanceConfigProperty.TAG_LIST.name(), updatedTags);
+
+ // Save to Helix
+ if
(!_helixDataAccessor.setProperty(_keyBuilder.instanceConfig(instanceName),
instanceConfig)) {
+ return PinotResourceManagerResponse.failure("Failed to set instance
config for instance: " + instanceName);
+ }
+
+ LOGGER.info("Successfully drained minion instance: {}", instanceName);
+ return PinotResourceManagerResponse.success(
+ "Successfully drained minion instance: " + instanceName);
+ }
+
+ /**
+ * Replaces minion_untagged tag with minion_drained tag in the given tag
list.
+ *
+ * @param currentTags Current list of instance tags
+ * @return Updated list with a minion_drained tag instead of minion_untagged
+ */
+ private List<String> replaceMinionTag(List<String> currentTags) {
Review Comment:
The method doesn't handle null currentTags. If currentTags is null, this
will throw a NullPointerException when creating the ArrayList. Add a null check
at the beginning of the method and return a list containing only
DRAINED_MINION_INSTANCE if currentTags is null.
```suggestion
private List<String> replaceMinionTag(List<String> currentTags) {
if (currentTags == null) {
// If currentTags is null, return a list containing only the drained
minion tag
return Collections.singletonList(Helix.DRAINED_MINION_INSTANCE);
}
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
