Stats API refactor; Cancellation of stream through StreamTransportResponse

Signed-off-by: Rishabh Maurya <[email protected]>

Author: rishabhmaurya

plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/stats/FlightStatsCollector.java

@@ -48,13 +48,14 @@ public class FlightStatsCollector extends AbstractLifecycleComponent {
     private final AtomicLong clientBatchTimeMax = new AtomicLong();
 
     // Shared metrics
-    private final AtomicLong bytesSentTotal = new AtomicLong();
-    private final AtomicLong bytesReceivedTotal = new AtomicLong();
-    private final AtomicLong streamErrorsTotal = new AtomicLong();
-    private final AtomicLong connectionErrorsTotal = new AtomicLong();
-    private final AtomicLong timeoutErrorsTotal = new AtomicLong();
-    private final AtomicLong streamsCompletedSuccessfully = new AtomicLong();
-    private final AtomicLong streamsFailedTotal = new AtomicLong();
+    private final AtomicLong bytesSent = new AtomicLong();
+    private final AtomicLong bytesReceived = new AtomicLong();
+    private final AtomicLong clientApplicationErrors = new AtomicLong();
+    private final AtomicLong clientTransportErrors = new AtomicLong();
+    private final AtomicLong serverApplicationErrors = new AtomicLong();
+    private final AtomicLong serverTransportErrors = new AtomicLong();
+    private final AtomicLong clientStreamsCompleted = new AtomicLong();
+    private final AtomicLong serverStreamsCompleted = new AtomicLong();
     private final long startTimeMillis = System.currentTimeMillis();
 
     private final AtomicLong channelsActive = new AtomicLong();
@@ -107,18 +108,19 @@ public FlightTransportStats collectStats() {
             totalClientBatches,
             totalClientResponses,
             totalServerBatches,
-            bytesSentTotal.get(),
-            bytesReceivedTotal.get()
+            bytesSent.get(),
+            bytesReceived.get()
         );
 
         ResourceUtilizationStats resourceUtilization = collectResourceStats();
 
         ReliabilityStats reliability = new ReliabilityStats(
-            streamErrorsTotal.get(),
-            connectionErrorsTotal.get(),
-            timeoutErrorsTotal.get(),
-            streamsCompletedSuccessfully.get(),
-            streamsFailedTotal.get(),
+            clientApplicationErrors.get(),
+            clientTransportErrors.get(),
+            serverApplicationErrors.get(),
+            serverTransportErrors.get(),
+            clientStreamsCompleted.get(),
+            serverStreamsCompleted.get(),
             System.currentTimeMillis() - startTimeMillis
         );
 
@@ -146,37 +148,44 @@ private ResourceUtilizationStats collectResourceStats() {
             directMemoryUsed = Runtime.getRuntime().totalMemory() - Runtime.getRuntime().freeMemory();
         }
 
-        int flightThreadsActive = 0;
-        int flightThreadsTotal = 0;
+        int clientThreadsActive = 0;
+        int clientThreadsTotal = 0;
+        int serverThreadsActive = 0;
+        int serverThreadsTotal = 0;
 
         if (threadPool != null) {
             try {
                 var allStats = threadPool.stats();
                 for (var stat : allStats) {
-                    if (ServerConfig.FLIGHT_SERVER_THREAD_POOL_NAME.equals(stat.getName())
-                        || ServerConfig.FLIGHT_CLIENT_THREAD_POOL_NAME.equals(stat.getName())) {
-                        flightThreadsActive += stat.getActive();
-                        flightThreadsTotal += stat.getThreads();
+                    if (ServerConfig.FLIGHT_CLIENT_THREAD_POOL_NAME.equals(stat.getName())) {
+                        clientThreadsActive += stat.getActive();
+                        clientThreadsTotal += stat.getThreads();
+                    } else if (ServerConfig.FLIGHT_SERVER_THREAD_POOL_NAME.equals(stat.getName())) {
+                        serverThreadsActive += stat.getActive();
+                        serverThreadsTotal += stat.getThreads();
                     }
                 }
             } catch (Exception e) {
                 // Ignore thread pool stats errors
             }
         }
 
+        // Add Netty event loop threads to server total
         if (bossEventLoopGroup != null && !bossEventLoopGroup.isShutdown()) {
-            flightThreadsTotal += 1;
+            serverThreadsTotal += 1;
         }
         if (workerEventLoopGroup != null && !workerEventLoopGroup.isShutdown()) {
-            flightThreadsTotal += Runtime.getRuntime().availableProcessors() * 2;
+            serverThreadsTotal += Runtime.getRuntime().availableProcessors() * 2;
         }
 
         return new ResourceUtilizationStats(
             arrowAllocatedBytes,
             arrowPeakBytes,
             directMemoryUsed,
-            flightThreadsActive,
-            flightThreadsTotal,
+            clientThreadsActive,
+            clientThreadsTotal,
+            serverThreadsActive,
+            serverThreadsTotal,
             (int) channelsActive.get(),
             (int) channelsActive.get()
         );
@@ -257,58 +266,43 @@ public void addClientBatchTime(long timeMillis) {
     /** Adds bytes sent
      * @param bytes number of bytes */
     public void addBytesSent(long bytes) {
-        bytesSentTotal.addAndGet(bytes);
+        bytesSent.addAndGet(bytes);
     }
 
     /** Adds bytes received
      * @param bytes number of bytes */
     public void addBytesReceived(long bytes) {
-        bytesReceivedTotal.addAndGet(bytes);
+        bytesReceived.addAndGet(bytes);
     }
 
-    /** Increments stream errors counter */
-    public void incrementStreamErrors() {
-        streamErrorsTotal.incrementAndGet();
+    /** Increments client application errors counter */
+    public void incrementClientApplicationErrors() {
+        clientApplicationErrors.incrementAndGet();
     }
 
-    /** Increments connection errors counter */
-    public void incrementConnectionErrors() {
-        connectionErrorsTotal.incrementAndGet();
+    /** Increments client transport errors counter */
+    public void incrementClientTransportErrors() {
+        clientTransportErrors.incrementAndGet();
     }
 
-    /** Increments timeout errors counter */
-    public void incrementTimeoutErrors() {
-        timeoutErrorsTotal.incrementAndGet();
+    /** Increments server application errors counter */
+    public void incrementServerApplicationErrors() {
+        serverApplicationErrors.incrementAndGet();
     }
 
-    /** Increments serialization errors counter */
-    public void incrementSerializationErrors() {
-        streamErrorsTotal.incrementAndGet();
+    /** Increments server transport errors counter */
+    public void incrementServerTransportErrors() {
+        serverTransportErrors.incrementAndGet();
     }
 
-    /** Increments transport errors counter */
-    public void incrementTransportErrors() {
-        streamErrorsTotal.incrementAndGet();
+    /** Increments client streams completed counter */
+    public void incrementClientStreamsCompleted() {
+        clientStreamsCompleted.incrementAndGet();
     }
 
-    /** Increments channel errors counter */
-    public void incrementChannelErrors() {
-        streamErrorsTotal.incrementAndGet();
-    }
-
-    /** Increments Flight server errors counter */
-    public void incrementFlightServerErrors() {
-        streamErrorsTotal.incrementAndGet();
-    }
-
-    /** Increments completed streams counter */
-    public void incrementStreamsCompleted() {
-        streamsCompletedSuccessfully.incrementAndGet();
-    }
-
-    /** Increments failed streams counter */
-    public void incrementStreamsFailed() {
-        streamsFailedTotal.incrementAndGet();
+    /** Increments server streams completed counter */
+    public void incrementServerStreamsCompleted() {
+        serverStreamsCompleted.incrementAndGet();
     }
 
     /** Increments active channels counter */

plugins/arrow-flight-rpc/src/main/java/org/opensearch/arrow/flight/stats/FlightStatsResponse.java

@@ -11,8 +11,10 @@
 import org.opensearch.action.FailedNodeException;
 import org.opensearch.action.support.nodes.BaseNodesResponse;
 import org.opensearch.cluster.ClusterName;
+import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.common.unit.ByteSizeValue;
 import org.opensearch.core.xcontent.ToXContentObject;
 import org.opensearch.core.xcontent.XContentBuilder;
 
@@ -72,41 +74,163 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
     }
 
     private void aggregateClusterStats(XContentBuilder builder, Params params) throws IOException {
+        // Performance aggregates
         long totalServerRequests = 0;
         long totalServerRequestsCurrent = 0;
+        long totalServerBatches = 0;
         long totalClientBatches = 0;
         long totalClientResponses = 0;
         long totalBytesSent = 0;
         long totalBytesReceived = 0;
-        long totalStreamErrors = 0;
+        long totalServerRequestTime = 0;
+        long totalServerBatchTime = 0;
+        long totalClientBatchTime = 0;
+
+        // Reliability aggregates
+        long totalClientApplicationErrors = 0;
+        long totalClientTransportErrors = 0;
+        long totalServerApplicationErrors = 0;
+        long totalServerTransportErrors = 0;
+        long totalClientStreamsCompleted = 0;
+        long totalServerStreamsCompleted = 0;
+        long totalUptime = 0;
+
+        // Resource aggregates
+        long totalArrowAllocated = 0;
+        long totalArrowPeak = 0;
+        long totalDirectMemory = 0;
+        int totalClientThreadsActive = 0;
+        int totalClientThreadsTotal = 0;
+        int totalServerThreadsActive = 0;
+        int totalServerThreadsTotal = 0;
+        int totalConnections = 0;
+        int totalChannels = 0;
 
         for (FlightNodeStats nodeStats : getNodes()) {
             FlightTransportStats stats = nodeStats.getFlightStats();
+
+            // Performance
             totalServerRequests += stats.performance.serverRequestsReceived;
             totalServerRequestsCurrent += stats.performance.serverRequestsCurrent;
+            totalServerBatches += stats.performance.serverBatchesSent;
             totalClientBatches += stats.performance.clientBatchesReceived;
             totalClientResponses += stats.performance.clientResponsesReceived;
-            totalBytesSent += stats.performance.bytesSentTotal;
-            totalBytesReceived += stats.performance.bytesReceivedTotal;
-            totalStreamErrors += stats.reliability.streamErrorsTotal;
+            totalBytesSent += stats.performance.bytesSent;
+            totalBytesReceived += stats.performance.bytesReceived;
+            totalServerRequestTime += stats.performance.serverRequestTotalMillis;
+            totalServerBatchTime += stats.performance.serverBatchTotalMillis;
+            totalClientBatchTime += stats.performance.clientBatchTotalMillis;
+
+            // Reliability
+            totalClientApplicationErrors += stats.reliability.clientApplicationErrors;
+            totalClientTransportErrors += stats.reliability.clientTransportErrors;
+            totalServerApplicationErrors += stats.reliability.serverApplicationErrors;
+            totalServerTransportErrors += stats.reliability.serverTransportErrors;
+            totalClientStreamsCompleted += stats.reliability.clientStreamsCompleted;
+            totalServerStreamsCompleted += stats.reliability.serverStreamsCompleted;
+            totalUptime = Math.max(totalUptime, stats.reliability.uptimeMillis);
+
+            // Resources
+            totalArrowAllocated += stats.resourceUtilization.arrowAllocatedBytes;
+            totalArrowPeak = Math.max(totalArrowPeak, stats.resourceUtilization.arrowPeakBytes);
+            totalDirectMemory += stats.resourceUtilization.directMemoryBytes;
+            totalClientThreadsActive += stats.resourceUtilization.clientThreadsActive;
+            totalClientThreadsTotal += stats.resourceUtilization.clientThreadsTotal;
+            totalServerThreadsActive += stats.resourceUtilization.serverThreadsActive;
+            totalServerThreadsTotal += stats.resourceUtilization.serverThreadsTotal;
+            totalConnections += stats.resourceUtilization.connectionsActive;
+            totalChannels += stats.resourceUtilization.channelsActive;
         }
 
+        // Performance stats
         builder.startObject("performance");
-        builder.field("total_server_requests", totalServerRequests);
-        builder.field("total_server_requests_current", totalServerRequestsCurrent);
-        builder.field("total_client_batches", totalClientBatches);
-        builder.field("total_client_responses", totalClientResponses);
-        builder.field("total_bytes_sent", totalBytesSent);
-        builder.field("total_bytes_received", totalBytesReceived);
+        builder.field("server_requests_total", totalServerRequests);
+        builder.field("server_requests_current", totalServerRequestsCurrent);
+        builder.field("server_batches_sent", totalServerBatches);
+        builder.field("client_batches_received", totalClientBatches);
+        builder.field("client_responses_received", totalClientResponses);
+        builder.field("bytes_sent", totalBytesSent);
+        if (params.paramAsBoolean("human", false)) {
+            builder.field("bytes_sent_human", new ByteSizeValue(totalBytesSent).toString());
+        }
+        builder.field("bytes_received", totalBytesReceived);
+        if (params.paramAsBoolean("human", false)) {
+            builder.field("bytes_received_human", new ByteSizeValue(totalBytesReceived).toString());
+        }
+        if (totalServerRequests > 0) {
+            long avgRequestTime = totalServerRequestTime / totalServerRequests;
+            builder.field("server_request_avg_millis", avgRequestTime);
+            if (params.paramAsBoolean("human", false)) {
+                builder.field("server_request_avg_time", TimeValue.timeValueMillis(avgRequestTime).toString());
+            }
+        }
+        if (totalServerBatches > 0) {
+            long avgBatchTime = totalServerBatchTime / totalServerBatches;
+            builder.field("server_batch_avg_millis", avgBatchTime);
+            if (params.paramAsBoolean("human", false)) {
+                builder.field("server_batch_avg_time", TimeValue.timeValueMillis(avgBatchTime).toString());
+            }
+        }
+        if (totalClientBatches > 0) {
+            long avgClientBatchTime = totalClientBatchTime / totalClientBatches;
+            builder.field("client_batch_avg_millis", avgClientBatchTime);
+            if (params.paramAsBoolean("human", false)) {
+                builder.field("client_batch_avg_time", TimeValue.timeValueMillis(avgClientBatchTime).toString());
+            }
+        }
         builder.endObject();
 
+        // Reliability stats
         builder.startObject("reliability");
-        builder.field("total_stream_errors", totalStreamErrors);
-        if (totalServerRequests > 0) {
-            builder.field("cluster_error_rate_percent", (totalStreamErrors * 100.0) / totalServerRequests);
+        builder.field("client_application_errors", totalClientApplicationErrors);
+        builder.field("client_transport_errors", totalClientTransportErrors);
+        builder.field("server_application_errors", totalServerApplicationErrors);
+        builder.field("server_transport_errors", totalServerTransportErrors);
+        builder.field("client_streams_completed", totalClientStreamsCompleted);
+        builder.field("server_streams_completed", totalServerStreamsCompleted);
+        builder.field("cluster_uptime_millis", totalUptime);
+        if (params.paramAsBoolean("human", false)) {
+            builder.field("cluster_uptime", TimeValue.timeValueMillis(totalUptime).toString());
+        }
+
+        long totalErrors = totalClientApplicationErrors + totalClientTransportErrors + totalServerApplicationErrors
+            + totalServerTransportErrors;
+        long totalStreams = totalClientStreamsCompleted + totalServerStreamsCompleted + totalErrors;
+        if (totalStreams > 0) {
+            builder.field("cluster_error_rate_percent", (totalErrors * 100.0) / totalStreams);
+            builder.field(
+                "cluster_success_rate_percent",
+                ((totalClientStreamsCompleted + totalServerStreamsCompleted) * 100.0) / totalStreams
+            );
         }
         builder.endObject();
 
-        // Resource utilization stats are per-node only
+        // Resource utilization stats
+        builder.startObject("resource_utilization");
+        builder.field("arrow_allocated_bytes_total", totalArrowAllocated);
+        if (params.paramAsBoolean("human", false)) {
+            builder.field("arrow_allocated_total", new ByteSizeValue(totalArrowAllocated).toString());
+        }
+        builder.field("arrow_peak_bytes_max", totalArrowPeak);
+        if (params.paramAsBoolean("human", false)) {
+            builder.field("arrow_peak_max", new ByteSizeValue(totalArrowPeak).toString());
+        }
+        builder.field("direct_memory_bytes_total", totalDirectMemory);
+        if (params.paramAsBoolean("human", false)) {
+            builder.field("direct_memory_total", new ByteSizeValue(totalDirectMemory).toString());
+        }
+        builder.field("client_threads_active", totalClientThreadsActive);
+        builder.field("client_threads_total", totalClientThreadsTotal);
+        builder.field("server_threads_active", totalServerThreadsActive);
+        builder.field("server_threads_total", totalServerThreadsTotal);
+        builder.field("connections_active", totalConnections);
+        builder.field("channels_active", totalChannels);
+        if (totalClientThreadsTotal > 0) {
+            builder.field("client_thread_utilization_percent", (totalClientThreadsActive * 100.0) / totalClientThreadsTotal);
+        }
+        if (totalServerThreadsTotal > 0) {
+            builder.field("server_thread_utilization_percent", (totalServerThreadsActive * 100.0) / totalServerThreadsTotal);
+        }
+        builder.endObject();
     }
 }