Improved metrics for server grpc query (apache#13177)

Author: soumitra-st

pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerMeter.java

@@ -108,6 +108,8 @@ public enum ServerMeter implements AbstractMetrics.Meter {
   GRPC_QUERIES("grpcQueries", true),
   GRPC_BYTES_RECEIVED("grpcBytesReceived", true),
   GRPC_BYTES_SENT("grpcBytesSent", true),
+  GRPC_TRANSPORT_READY("grpcTransport", true),
+  GRPC_TRANSPORT_TERMINATED("grpcTransport", true),
 
   NUM_SEGMENTS_PRUNED_INVALID("numSegmentsPrunedInvalid", false),
   NUM_SEGMENTS_PRUNED_BY_LIMIT("numSegmentsPrunedByLimit", false),

pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerTimer.java

@@ -53,6 +53,7 @@ public enum ServerTimer implements AbstractMetrics.Timer {
       "Total time taken to preload a table partition of an upsert table with upsert snapshot"),
   UPSERT_REMOVE_EXPIRED_PRIMARY_KEYS_TIME_MS("milliseconds", false,
       "Total time taken to delete expired primary keys based on metadataTTL or deletedKeysTTL"),
+  GRPC_QUERY_EXECUTION_MS("milliseconds", false, "Total execution time of a successful query over gRPC"),
   UPSERT_SNAPSHOT_TIME_MS("milliseconds", false, "Total time taken to take upsert table snapshot");
 
   private final String _timerName;

pinot-core/src/main/java/org/apache/pinot/core/transport/grpc/GrpcQueryServer.java

@@ -18,8 +18,11 @@
  */
 package org.apache.pinot.core.transport.grpc;
 
+import io.grpc.Attributes;
+import io.grpc.Grpc;
 import io.grpc.Server;
 import io.grpc.ServerBuilder;
+import io.grpc.ServerTransportFilter;
 import io.grpc.Status;
 import io.grpc.netty.shaded.io.grpc.netty.GrpcSslContexts;
 import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder;
@@ -33,12 +36,14 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 import nl.altindag.ssl.SSLFactory;
 import org.apache.pinot.common.config.GrpcConfig;
 import org.apache.pinot.common.config.TlsConfig;
 import org.apache.pinot.common.datatable.DataTable;
 import org.apache.pinot.common.metrics.ServerMeter;
 import org.apache.pinot.common.metrics.ServerMetrics;
+import org.apache.pinot.common.metrics.ServerTimer;
 import org.apache.pinot.common.proto.PinotQueryServerGrpc;
 import org.apache.pinot.common.proto.Server.ServerRequest;
 import org.apache.pinot.common.proto.Server.ServerResponse;
@@ -72,19 +77,41 @@ public class GrpcQueryServer extends PinotQueryServerGrpc.PinotQueryServerImplBa
   private final AccessControl _accessControl;
   private final ServerQueryLogger _queryLogger = ServerQueryLogger.getInstance();
 
+  // Filter to keep track of gRPC connections.
+  private class GrpcQueryTransportFilter extends ServerTransportFilter {
+    @Override
+    public Attributes transportReady(Attributes transportAttrs) {
+      LOGGER.info("gRPC transportReady: REMOTE_ADDR {}",
+          transportAttrs != null ? transportAttrs.get(Grpc.TRANSPORT_ATTR_REMOTE_ADDR) : "null");
+      _serverMetrics.addMeteredGlobalValue(ServerMeter.GRPC_TRANSPORT_READY, 1);
+      return super.transportReady(transportAttrs);
+    }
+
+    @Override
+    public void transportTerminated(Attributes transportAttrs) {
+      // transportTerminated can be called without transportReady before it, e.g. handshake fails
+      // So, don't emit metrics if transportAttrs is null
+      if (transportAttrs != null) {
+        LOGGER.info("gRPC transportTerminated: REMOTE_ADDR {}", transportAttrs.get(Grpc.TRANSPORT_ATTR_REMOTE_ADDR));
+        _serverMetrics.addMeteredGlobalValue(ServerMeter.GRPC_TRANSPORT_TERMINATED, 1);
+      }
+    }
+  }
+
   public GrpcQueryServer(int port, GrpcConfig config, TlsConfig tlsConfig, QueryExecutor queryExecutor,
       ServerMetrics serverMetrics, AccessControl accessControl) {
     _queryExecutor = queryExecutor;
     _serverMetrics = serverMetrics;
     if (tlsConfig != null) {
       try {
         _server = NettyServerBuilder.forPort(port).sslContext(buildGRpcSslContext(tlsConfig))
-            .maxInboundMessageSize(config.getMaxInboundMessageSizeBytes()).addService(this).build();
+            .maxInboundMessageSize(config.getMaxInboundMessageSizeBytes()).addService(this)
+            .addTransportFilter(new GrpcQueryTransportFilter()).build();
       } catch (Exception e) {
         throw new RuntimeException("Failed to start secure grpcQueryServer", e);
       }
     } else {
-      _server = ServerBuilder.forPort(port).addService(this).build();
+      _server = ServerBuilder.forPort(port).addService(this).addTransportFilter(new GrpcQueryTransportFilter()).build();
     }
     _accessControl = accessControl;
     LOGGER.info("Initialized GrpcQueryServer on port: {} with numWorkerThreads: {}", port,
@@ -136,6 +163,7 @@ public void shutdown() {
 
   @Override
   public void submit(ServerRequest request, StreamObserver<ServerResponse> responseObserver) {
+    long startTime = System.nanoTime();
     _serverMetrics.addMeteredGlobalValue(ServerMeter.GRPC_QUERIES, 1);
     _serverMetrics.addMeteredGlobalValue(ServerMeter.GRPC_BYTES_RECEIVED, request.getSerializedSize());
 
@@ -162,13 +190,16 @@ public void submit(ServerRequest request, StreamObserver<ServerResponse> respons
       _serverMetrics.addMeteredGlobalValue(ServerMeter.NO_TABLE_ACCESS, 1);
       responseObserver.onError(
           Status.NOT_FOUND.withDescription(exceptionMsg).withCause(unsupportedOperationException).asException());
+      return;
     }
 
     // Process the query
     InstanceResponseBlock instanceResponse;
     try {
-      instanceResponse =
-          _queryExecutor.execute(queryRequest, _executorService, new GrpcResultsBlockStreamer(responseObserver));
+      LOGGER.info("Executing gRPC query request {}: {} received from broker: {}", queryRequest.getRequestId(),
+          queryRequest.getQueryContext(), queryRequest.getBrokerId());
+      instanceResponse = _queryExecutor.execute(queryRequest, _executorService,
+          new GrpcResultsBlockStreamer(responseObserver, _serverMetrics));
     } catch (Exception e) {
       LOGGER.error("Caught exception while processing request {}: {} from broker: {}", queryRequest.getRequestId(),
           queryRequest.getQueryContext(), queryRequest.getBrokerId(), e);
@@ -192,6 +223,8 @@ public void submit(ServerRequest request, StreamObserver<ServerResponse> respons
     responseObserver.onNext(serverResponse);
     _serverMetrics.addMeteredGlobalValue(ServerMeter.GRPC_BYTES_SENT, serverResponse.getSerializedSize());
     responseObserver.onCompleted();
+    _serverMetrics.addTimedTableValue(queryRequest.getTableNameWithType(), ServerTimer.GRPC_QUERY_EXECUTION_MS,
+        System.nanoTime() - startTime, TimeUnit.NANOSECONDS);
 
     // Log the query
     if (_queryLogger != null) {

pinot-core/src/main/java/org/apache/pinot/core/transport/grpc/GrpcResultsBlockStreamer.java

@@ -23,6 +23,8 @@
 import java.io.IOException;
 import java.util.Collection;
 import org.apache.pinot.common.datatable.DataTable;
+import org.apache.pinot.common.metrics.ServerMeter;
+import org.apache.pinot.common.metrics.ServerMetrics;
 import org.apache.pinot.common.proto.Server;
 import org.apache.pinot.common.utils.DataSchema;
 import org.apache.pinot.core.operator.blocks.results.BaseResultsBlock;
@@ -32,9 +34,11 @@
 
 public class GrpcResultsBlockStreamer implements ResultsBlockStreamer {
   private final StreamObserver<Server.ServerResponse> _streamObserver;
+  private final ServerMetrics _serverMetrics;
 
-  public GrpcResultsBlockStreamer(StreamObserver<Server.ServerResponse> streamObserver) {
+  public GrpcResultsBlockStreamer(StreamObserver<Server.ServerResponse> streamObserver, ServerMetrics serverMetrics) {
     _streamObserver = streamObserver;
+    _serverMetrics = serverMetrics;
   }
 
   @Override
@@ -44,6 +48,8 @@ public void send(BaseResultsBlock block)
     Collection<Object[]> rows = block.getRows();
     Preconditions.checkState(dataSchema != null && rows != null, "Malformed data block");
     DataTable dataTable = block.getDataTable();
-    _streamObserver.onNext(StreamingResponseUtils.getDataResponse(dataTable));
+    Server.ServerResponse response = StreamingResponseUtils.getDataResponse(dataTable);
+    _streamObserver.onNext(response);
+    _serverMetrics.addMeteredGlobalValue(ServerMeter.GRPC_BYTES_SENT, response.getSerializedSize());
   }
 }