Add serverStats to BrokerQueryEventListener

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

@@ -101,6 +101,7 @@
 import org.apache.pinot.spi.exception.BadQueryRequestException;
 import org.apache.pinot.spi.exception.DatabaseConflictException;
 import org.apache.pinot.spi.trace.RequestContext;
+import org.apache.pinot.spi.trace.ServerStatsInfo;
 import org.apache.pinot.spi.trace.Tracing;
 import org.apache.pinot.spi.utils.CommonConstants;
 import org.apache.pinot.spi.utils.CommonConstants.Broker;
@@ -859,7 +860,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 +1899,23 @@ private String getGlobalQueryId(long requestId) {
    */
   public static class ServerStats {
     private String _serverStats;
+    private Map<String, ServerStatsInfo> _serverStatsMap;
 
     public String getServerStats() {
       return _serverStats;
     }
 
+    public Map getServerStatsMap() {
+      return _serverStatsMap;
+    }
+
     public void setServerStats(String serverStats) {
       _serverStats = serverStats;
     }
+
+    public void setServerStatsMap(Map<String, ServerStatsInfo> 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,12 +25,14 @@
 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;
 import org.apache.pinot.common.exception.QueryException;
 import org.apache.pinot.common.utils.HashUtil;
 import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager;
+import org.apache.pinot.spi.trace.ServerStatsInfo;
 
 
 /**
@@ -145,6 +147,20 @@ public String getServerStats() {
     return stringBuilder.toString();
   }
 
+  @Override
+  public Map<String, ServerStatsInfo> getServerStatsMap() {
+    return _responseMap.entrySet().stream()
+        .collect(Collectors.toMap(
+            entry -> entry.getKey().getInstanceId(),
+            entry -> new ServerStatsInfo(
+                entry.getValue().getSubmitDelayMs(),
+                entry.getValue().getRequestSentDelayMs(),
+                entry.getValue().getResponseDelayMs(),
+                entry.getValue().getResponseSize(),
+                entry.getValue().getDeserializationTimeMs())
+        ));
+  }
+
   @Override
   public long getServerResponseDelayMs(ServerRoutingInstance serverRoutingInstance) {
     return _responseMap.get(serverRoutingInstance).getResponseDelayMs();

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

@@ -23,6 +23,7 @@
 import javax.annotation.concurrent.ThreadSafe;
 import org.apache.pinot.spi.annotations.InterfaceAudience;
 import org.apache.pinot.spi.annotations.InterfaceStability;
+import org.apache.pinot.spi.trace.ServerStatsInfo;
 
 
 /**
@@ -62,6 +63,8 @@ Map<ServerRoutingInstance, ServerResponse> getFinalResponses()
    */
   String getServerStats();
 
+  Map<String, ServerStatsInfo> 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, ServerStatsInfo> _serverStatsMap = new HashMap<>();
 
   public DefaultRequestContext() {
   }
@@ -573,6 +574,16 @@ public void setRequestHttpHeaders(Map<String, List<String>> requestHttpHeaders)
     _requestHttpHeaders.putAll(requestHttpHeaders);
   }
 
+  @Override
+  public Map<String, ServerStatsInfo> getServerStatsMap() {
+    return _serverStatsMap;
+  }
+
+  @Override
+  public void setServerStatsMap(Map<String, ServerStatsInfo> 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, ServerStatsInfo> getServerStatsMap();
+
+  void setServerStatsMap(Map<String, ServerStatsInfo> serverStatsMap);
+
   enum FanoutType {
     OFFLINE, REALTIME, HYBRID
   }

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

@@ -0,0 +1,56 @@
+/**
+ * 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.spi.trace;
+
+public class ServerStatsInfo {
+    private final long _submitRequestTimeMs;
+    private final int _requestSentLatencyMs;
+    private final long _receiveDataTableTimeMs;
+    private final int _responseSize;
+    private final int _deserializationTimeMs;
+
+    public ServerStatsInfo(long submitRequestTimeMs, int requestSentLatencyMs,
+        long receiveDataTableTimeMs, int responseSize, int deserializationTimeMs) {
+      _submitRequestTimeMs = submitRequestTimeMs;
+      _requestSentLatencyMs = requestSentLatencyMs;
+      _receiveDataTableTimeMs = receiveDataTableTimeMs;
+      _responseSize = responseSize;
+      _deserializationTimeMs = deserializationTimeMs;
+    }
+
+    public long getSubmitRequestTimeMs() {
+      return _submitRequestTimeMs;
+    }
+
+    public int getRequestSentLatencyMs() {
+      return _requestSentLatencyMs;
+    }
+
+    public long getReceiveDataTableTimeMs() {
+      return _receiveDataTableTimeMs;
+    }
+
+    public int getResponseSize() {
+      return _responseSize;
+    }
+
+    public int getDeserializationTimeMs() {
+      return _deserializationTimeMs;
+    }
+}

pinot-spi/src/test/java/org/apache/pinot/spi/trace/ServerStatsInfoTest.java

@@ -0,0 +1,36 @@
+/**
+ * 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.spi.trace;
+
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+
+public class ServerStatsInfoTest {
+  ServerStatsInfo _serverStatsInfo = new ServerStatsInfo(1, 2, 3, 4, 5);
+
+  @Test
+  public void testServerStatsInfoValues() {
+    assertEquals(1, _serverStatsInfo.getSubmitRequestTimeMs());
+    assertEquals(2, _serverStatsInfo.getRequestSentLatencyMs());
+    assertEquals(3, _serverStatsInfo.getReceiveDataTableTimeMs());
+    assertEquals(4, _serverStatsInfo.getResponseSize());
+    assertEquals(5, _serverStatsInfo.getDeserializationTimeMs());
+  }
+}