SOLR-17337: Show proper distributed stage id (#2594)

Co-authored-by: Torsten Bøgh Köster <[email protected]>

Author: cpoerschke

solr/CHANGES.txt

@@ -212,6 +212,8 @@ Bug Fixes
 
 * SOLR-17369: Fix "flags" usage in FunctionQParser that caused some issues in vectorSimilarity() with BYTE vector constants (hossman)
 
+* SOLR-17337: Display all custom distributed stages in debug output. (Torsten Bøgh Köster, Christine Poerschke)
+
 Dependency Upgrades
 ---------------------
 (No changes)

solr/core/src/java/org/apache/solr/handler/component/DebugComponent.java

@@ -19,6 +19,7 @@
 import static org.apache.solr.common.params.CommonParams.FQ;
 import static org.apache.solr.common.params.CommonParams.JSON;
 
+import com.google.common.annotations.VisibleForTesting;
 import java.io.IOException;
 import java.lang.reflect.Array;
 import java.util.ArrayList;
@@ -181,16 +182,28 @@ public void modifyRequest(ResponseBuilder rb, SearchComponent who, ShardRequest
     }
   }
 
+  @VisibleForTesting
+  protected String getDistributedStageName(int stage) {
+    String stageName = stages.get(stage);
+
+    if (stageName == null) {
+      stageName = "STAGE_" + Integer.toString(stage);
+    }
+
+    return stageName;
+  }
+
   @Override
   public void handleResponses(ResponseBuilder rb, ShardRequest sreq) {
     if (rb.isDebugTrack() && rb.isDistrib && !rb.finished.isEmpty()) {
       @SuppressWarnings("unchecked")
       NamedList<Object> stageList =
           (NamedList<Object>)
-              ((NamedList<Object>) rb.getDebugInfo().get("track")).get(stages.get(rb.stage));
+              ((NamedList<Object>) rb.getDebugInfo().get("track"))
+                  .get(getDistributedStageName(rb.stage));
       if (stageList == null) {
         stageList = new SimpleOrderedMap<>();
-        rb.addDebug(stageList, "track", stages.get(rb.stage));
+        rb.addDebug(stageList, "track", getDistributedStageName(rb.stage));
       }
       for (ShardResponse response : sreq.responses) {
         stageList.add(response.getShard(), getTrackResponse(response));

solr/core/src/test/org/apache/solr/handler/component/DebugComponentTest.java

@@ -297,4 +297,14 @@ private void addRequestId(ResponseBuilder rb, String requestId) {
     params.add(CommonParams.REQUEST_ID, requestId);
     rb.req.setParams(params);
   }
+
+  @Test
+  public void testDistributedStageResolution() throws IOException {
+    final DebugComponent debugComponent = new DebugComponent();
+    assertEquals(
+        "PARSE_QUERY", debugComponent.getDistributedStageName(ResponseBuilder.STAGE_PARSE_QUERY));
+    assertEquals("DONE", debugComponent.getDistributedStageName(ResponseBuilder.STAGE_DONE));
+    assertEquals("STAGE_400", debugComponent.getDistributedStageName(400));
+    assertEquals("STAGE_10000", debugComponent.getDistributedStageName(10000));
+  }
 }