From 90f7a82b5e0462605723f85d46641a31277eaabb Mon Sep 17 00:00:00 2001
From: Vic Zhang <vzhang@fb.com>
Date: Tue, 3 Sep 2019 11:11:18 -0400
Subject: [PATCH] Add query task count to statistics field of
 QueryCompletedEvent

---
 .../main/java/com/facebook/presto/event/QueryMonitor.java | 2 ++
 .../presto/spi/eventlistener/QueryStatistics.java         | 8 ++++++++
 2 files changed, 10 insertions(+)

diff --git a/presto-main/src/main/java/com/facebook/presto/event/QueryMonitor.java b/presto-main/src/main/java/com/facebook/presto/event/QueryMonitor.java
index 086701f1c22a..c56ee114c910 100644
--- a/presto-main/src/main/java/com/facebook/presto/event/QueryMonitor.java
+++ b/presto-main/src/main/java/com/facebook/presto/event/QueryMonitor.java
@@ -164,6 +164,7 @@ public void queryImmediateFailureEvent(BasicQueryInfo queryInfo, ExecutionFailur
                         0,
                         0,
                         0,
+                        0,
                         ImmutableList.of(),
                         0,
                         true,
@@ -230,6 +231,7 @@ private QueryStatistics createQueryStatistics(QueryInfo queryInfo)
                 ofMillis(queryStats.getTotalScheduledTime().toMillis()),
                 ofMillis(queryStats.getQueuedTime().toMillis()),
                 Optional.of(ofMillis(queryStats.getAnalysisTime().toMillis())),
+                queryStats.getPeakRunningTasks(),
                 queryStats.getPeakUserMemoryReservation().toBytes(),
                 queryStats.getPeakTotalMemoryReservation().toBytes(),
                 queryStats.getPeakTaskUserMemory().toBytes(),
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/eventlistener/QueryStatistics.java b/presto-spi/src/main/java/com/facebook/presto/spi/eventlistener/QueryStatistics.java
index de41835dd992..8edbf49ea494 100644
--- a/presto-spi/src/main/java/com/facebook/presto/spi/eventlistener/QueryStatistics.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/eventlistener/QueryStatistics.java
@@ -26,6 +26,7 @@ public class QueryStatistics
     private final Duration queuedTime;
     private final Optional<Duration> analysisTime;
 
+    private final int peakRunningTasks;
     private final long peakUserMemoryBytes;
     // peak of user + system memory
     private final long peakTotalNonRevocableMemoryBytes;
@@ -55,6 +56,7 @@ public QueryStatistics(
             Duration wallTime,
             Duration queuedTime,
             Optional<Duration> analysisTime,
+            int peakRunningTasks,
             long peakUserMemoryBytes,
             long peakTotalNonRevocableMemoryBytes,
             long peakTaskUserMemory,
@@ -77,6 +79,7 @@ public QueryStatistics(
         this.wallTime = requireNonNull(wallTime, "wallTime is null");
         this.queuedTime = requireNonNull(queuedTime, "queuedTime is null");
         this.analysisTime = requireNonNull(analysisTime, "analysisTime is null");
+        this.peakRunningTasks = peakRunningTasks;
         this.peakUserMemoryBytes = peakUserMemoryBytes;
         this.peakTotalNonRevocableMemoryBytes = peakTotalNonRevocableMemoryBytes;
         this.peakTaskUserMemory = peakTaskUserMemory;
@@ -116,6 +119,11 @@ public Optional<Duration> getAnalysisTime()
         return analysisTime;
     }
 
+    public int getPeakRunningTasks()
+    {
+        return peakRunningTasks;
+    }
+
     public long getPeakUserMemoryBytes()
     {
         return peakUserMemoryBytes;