Add serverStats to BrokerQueryEventListener

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java

@@ -859,7 +859,7 @@ protected BrokerResponse handleRequest(long requestId, String query, SqlNodeAndO
       // Log query and stats
       _queryLogger.log(
           new QueryLogger.QueryLogParams(requestContext, tableName, brokerResponse, requesterIdentity, serverStats));
-
+      requestContext.setServerStatsMap(serverStats.getServerStatsMap());
       return brokerResponse;
     } finally {
       Tracing.ThreadAccountantOps.clear();
@@ -1898,14 +1898,23 @@ private String getGlobalQueryId(long requestId) {
    */
   public static class ServerStats {
     private String _serverStats;
+    private Map<String, Map<String, Integer>> _serverStatsMap;
 
     public String getServerStats() {
       return _serverStats;
     }
 
+    public Map getServerStatsMap() {
+      return _serverStatsMap;
+    }
+
     public void setServerStats(String serverStats) {
       _serverStats = serverStats;
     }
+
+    public void setServerStatsMap(Map<String, Map<String, Integer>> serverStatsMap) {
+      _serverStatsMap = serverStatsMap;
+    }
   }
 
   /**

pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java

@@ -126,6 +126,7 @@ protected BrokerResponseNative processBrokerRequest(long requestId, BrokerReques
         System.nanoTime() - scatterGatherStartTimeNs);
     // TODO Use scatterGatherStats as serverStats
     serverStats.setServerStats(asyncQueryResponse.getServerStats());
+    serverStats.setServerStatsMap(asyncQueryResponse.getServerStatsMap());
 
     int numServersQueried = finalResponses.size();
     long totalResponseSize = 0;

pinot-core/src/main/java/org/apache/pinot/core/transport/AsyncQueryResponse.java

@@ -25,6 +25,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import javax.annotation.concurrent.ThreadSafe;
 import org.apache.pinot.common.datatable.DataTable;
@@ -145,6 +146,21 @@ public String getServerStats() {
     return stringBuilder.toString();
   }
 
+  @Override
+  public Map<String, Map<String, Integer>> getServerStatsMap() {
+    return _responseMap.entrySet().stream()
+        .collect(Collectors.toMap(
+            entry -> entry.getKey().getShortName(),
+            entry -> Map.of(
+                "SubmitDelayMs", entry.getValue().getSubmitDelayMs(),
+                "ResponseDelayMs", entry.getValue().getResponseDelayMs(),
+                "ResponseSize", entry.getValue().getResponseSize(),
+                "DeserializationTimeMs", entry.getValue().getDeserializationTimeMs(),
+                "RequestSentDelayMs", entry.getValue().getRequestSentDelayMs()
+            )
+        ));
+  }
+
   @Override
   public long getServerResponseDelayMs(ServerRoutingInstance serverRoutingInstance) {
     return _responseMap.get(serverRoutingInstance).getResponseDelayMs();

pinot-core/src/main/java/org/apache/pinot/core/transport/QueryResponse.java

@@ -62,6 +62,8 @@ Map<ServerRoutingInstance, ServerResponse> getFinalResponses()
    */
   String getServerStats();
 
+  Map<String, Map<String, Integer>> getServerStatsMap();
+
   /**
    * Returns the time taken for the server to respond to the query.
    * @param serverRoutingInstance

pinot-spi/src/main/java/org/apache/pinot/spi/trace/DefaultRequestContext.java

@@ -81,6 +81,7 @@ public class DefaultRequestContext implements RequestScope {
   private Map<String, String> _traceInfo = new HashMap<>();
   private List<String> _processingExceptions = new ArrayList<>();
   private Map<String, List<String>> _requestHttpHeaders = new HashMap<>();
+  private Map<String, Map<String, Integer>> _serverStatsMap = new HashMap<>();
 
   public DefaultRequestContext() {
   }
@@ -573,6 +574,16 @@ public void setRequestHttpHeaders(Map<String, List<String>> requestHttpHeaders)
     _requestHttpHeaders.putAll(requestHttpHeaders);
   }
 
+  @Override
+  public Map<String, Map<String, Integer>> getServerStatsMap() {
+    return _serverStatsMap;
+  }
+
+  @Override
+  public void setServerStatsMap(Map<String, Map<String, Integer>> serverStatsMap) {
+    _serverStatsMap.putAll(serverStatsMap);
+  }
+
   @Override
   public void close() {
   }

pinot-spi/src/main/java/org/apache/pinot/spi/trace/RequestContext.java

@@ -229,6 +229,10 @@ default boolean isSampledRequest() {
    */
   void setRequestHttpHeaders(Map<String, List<String>> requestHttpHeaders);
 
+  Map<String, Map<String, Integer>> getServerStatsMap();
+
+  void setServerStatsMap(Map<String, Map<String, Integer>> serverStatsMap);
+
   enum FanoutType {
     OFFLINE, REALTIME, HYBRID
   }