Expose output buffer metrics in query completion event

trinodb · Dec 12, 2024 · 98d7c13 · 98d7c13
1 parent 2331e11
commit 98d7c13
Show file tree

Hide file tree

Showing 9 changed files with 73 additions and 0 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/event/QueryMonitor.java b/core/trino-main/src/main/java/io/trino/event/QueryMonitor.java
@@ -62,6 +62,7 @@
 import io.trino.spi.eventlistener.QueryOutputMetadata;
 import io.trino.spi.eventlistener.QueryStatistics;
 import io.trino.spi.eventlistener.StageCpuDistribution;
+import io.trino.spi.eventlistener.StageOutputBufferMetrics;
 import io.trino.spi.eventlistener.StageOutputBufferUtilization;
 import io.trino.spi.eventlistener.StageTaskStatistics;
 import io.trino.spi.metrics.Metrics;
@@ -238,6 +239,7 @@ public void queryImmediateFailureEvent(BasicQueryInfo queryInfo, ExecutionFailur
                         ImmutableList.of(),
                         ImmutableList.of(),
                         ImmutableList.of(),
+                        ImmutableList.of(),
                         Optional.empty()),
                 createQueryContext(
                         queryInfo.getSession(),
@@ -351,6 +353,7 @@ private QueryStatistics createQueryStatistics(QueryInfo queryInfo)
                 queryInfo.isFinalQueryInfo(),
                 getCpuDistributions(queryInfo),
                 getStageOutputBufferUtilizations(queryInfo),
+                getStageOutputBufferMetrics(queryInfo),
                 getStageTaskStatistics(queryInfo),
                 memoize(() -> operatorStats.stream().map(operatorStatsCodec::toJson).toList()),
                 ImmutableList.copyOf(queryInfo.getQueryStats().getOptimizerRulesSummaries()),
@@ -780,6 +783,29 @@ private static void populateStageOutputBufferUtilization(StageInfo stageInfo, Im
         }
     }
 
+    private static List<StageOutputBufferMetrics> getStageOutputBufferMetrics(QueryInfo queryInfo)
+    {
+        if (queryInfo.getOutputStage().isEmpty()) {
+            return ImmutableList.of();
+        }
+        ImmutableList.Builder<StageOutputBufferMetrics> builder = ImmutableList.builder();
+        populateStageOutputBufferMetrics(queryInfo.getOutputStage().get(), builder);
+
+        return builder.build();
+    }
+
+    private static void populateStageOutputBufferMetrics(StageInfo stageInfo, ImmutableList.Builder<StageOutputBufferMetrics> accumulator)
+    {
+        Metrics metrics = stageInfo.getStageStats().getOutputBufferMetrics();
+        if (!metrics.getMetrics().isEmpty()) {
+            accumulator.add(new StageOutputBufferMetrics(stageInfo.getStageId().getId(), metrics));
+        }
+
+        for (StageInfo subStage : stageInfo.getSubStages()) {
+            populateStageOutputBufferMetrics(subStage, accumulator);
+        }
+    }
+
     private List<StageTaskStatistics> getStageTaskStatistics(QueryInfo queryInfo)
     {
         if (queryInfo.getOutputStage().isEmpty()) {

diff --git a/core/trino-spi/pom.xml b/core/trino-spi/pom.xml
@@ -248,6 +248,13 @@
                                     <code>java.method.removed</code>
                                     <old>method io.trino.spi.connector.ConnectorTableHandle io.trino.spi.connector.ConnectorMetadata::makeCompatiblePartitioning(io.trino.spi.connector.ConnectorSession, io.trino.spi.connector.ConnectorTableHandle, io.trino.spi.connector.ConnectorPartitioningHandle)</old>
                                 </item>
+                                <item>
+                                    <ignore>true</ignore>
+                                    <code>java.method.numberOfParametersChanged</code>
+                                    <old>method void io.trino.spi.eventlistener.QueryStatistics::&lt;init&gt;(java.time.Duration, java.time.Duration, java.time.Duration, java.time.Duration, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, long, long, long, long, long, long, long, long, long, long, long, long, long, long, long, long, double, double, java.util.List&lt;io.trino.spi.eventlistener.StageGcStatistics&gt;, int, boolean, java.util.List&lt;io.trino.spi.eventlistener.StageCpuDistribution&gt;, java.util.List&lt;io.trino.spi.eventlistener.StageOutputBufferUtilization&gt;, java.util.List&lt;io.trino.spi.eventlistener.StageTaskStatistics&gt;, java.util.function.Supplier&lt;java.util.List&lt;java.lang.String&gt;&gt;, java.util.List&lt;io.trino.spi.eventlistener.QueryPlanOptimizerStatistics&gt;, java.util.Optional&lt;java.lang.String&gt;)</old>
+                                    <new>method void io.trino.spi.eventlistener.QueryStatistics::&lt;init&gt;(java.time.Duration, java.time.Duration, java.time.Duration, java.time.Duration, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, java.util.Optional&lt;java.time.Duration&gt;, long, long, long, long, long, long, long, long, long, long, long, long, long, long, long, long, double, double, java.util.List&lt;io.trino.spi.eventlistener.StageGcStatistics&gt;, int, boolean, java.util.List&lt;io.trino.spi.eventlistener.StageCpuDistribution&gt;, java.util.List&lt;io.trino.spi.eventlistener.StageOutputBufferUtilization&gt;, java.util.List&lt;io.trino.spi.eventlistener.StageOutputBufferMetrics&gt;, java.util.List&lt;io.trino.spi.eventlistener.StageTaskStatistics&gt;, java.util.function.Supplier&lt;java.util.List&lt;java.lang.String&gt;&gt;, java.util.List&lt;io.trino.spi.eventlistener.QueryPlanOptimizerStatistics&gt;, java.util.Optional&lt;java.lang.String&gt;)</new>
+                                    <justification>Added output buffer metrics</justification>
+                                </item>
                             </differences>
                         </revapi.differences>
                     </analysisConfiguration>

diff --git a/core/trino-spi/src/main/java/io/trino/spi/eventlistener/QueryStatistics.java b/core/trino-spi/src/main/java/io/trino/spi/eventlistener/QueryStatistics.java
@@ -75,6 +75,7 @@ public class QueryStatistics
     private final List<StageCpuDistribution> cpuTimeDistribution;
     private final List<StageOutputBufferUtilization> outputBufferUtilization;
     private final List<StageTaskStatistics> taskStatistics;
+    private final List<StageOutputBufferMetrics> outputBufferMetrics;
 
     /**
      * Operator summaries serialized to JSON. Serialization format and structure
@@ -131,6 +132,7 @@ public QueryStatistics(
             boolean complete,
             List<StageCpuDistribution> cpuTimeDistribution,
             List<StageOutputBufferUtilization> outputBufferUtilization,
+            List<StageOutputBufferMetrics> outputBufferMetrics,
             List<StageTaskStatistics> taskStatistics,
             List<String> operatorSummaries,
             List<QueryPlanOptimizerStatistics> optimizerRulesSummaries,
@@ -177,6 +179,7 @@ public QueryStatistics(
                 complete,
                 cpuTimeDistribution,
                 outputBufferUtilization,
+                outputBufferMetrics,
                 taskStatistics,
                 () -> operatorSummaries,
                 optimizerRulesSummaries,
@@ -224,6 +227,7 @@ public QueryStatistics(
             boolean complete,
             List<StageCpuDistribution> cpuTimeDistribution,
             List<StageOutputBufferUtilization> outputBufferUtilization,
+            List<StageOutputBufferMetrics> outputBufferMetrics,
             List<StageTaskStatistics> taskStatistics,
             Supplier<List<String>> operatorSummariesProvider,
             List<QueryPlanOptimizerStatistics> optimizerRulesSummaries,
@@ -269,6 +273,7 @@ public QueryStatistics(
         this.complete = complete;
         this.cpuTimeDistribution = requireNonNull(cpuTimeDistribution, "cpuTimeDistribution is null");
         this.outputBufferUtilization = requireNonNull(outputBufferUtilization, "outputBufferUtilization is null");
+        this.outputBufferMetrics = requireNonNull(outputBufferMetrics, "outputBufferMetrics is null");
         this.taskStatistics = requireNonNull(taskStatistics, "taskStatistics is null");
         this.operatorSummariesProvider = requireNonNull(operatorSummariesProvider, "operatorSummariesProvider is null");
         this.optimizerRulesSummaries = requireNonNull(optimizerRulesSummaries, "optimizerRulesSummaries is null");
@@ -515,6 +520,12 @@ public List<StageOutputBufferUtilization> getOutputBufferUtilization()
         return outputBufferUtilization;
     }
 
+    @JsonProperty
+    public List<StageOutputBufferMetrics> getOutputBufferMetrics()
+    {
+        return outputBufferMetrics;
+    }
+
     @JsonProperty
     public List<StageTaskStatistics> getTaskStatistics()
     {

diff --git a/core/trino-spi/src/main/java/io/trino/spi/eventlistener/StageOutputBufferMetrics.java b/core/trino-spi/src/main/java/io/trino/spi/eventlistener/StageOutputBufferMetrics.java
@@ -0,0 +1,21 @@
+/*
+ * Licensed 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 io.trino.spi.eventlistener;
+
+import io.trino.spi.metrics.Metrics;
+
+/**
+ * This class is JSON serializable for convenience and serialization compatibility is not guaranteed across versions.
+ */
+public record StageOutputBufferMetrics(int stageId, Metrics metrics) {}
diff --git a/...no-http-event-listener/src/test/java/io/trino/plugin/httpquery/TestHttpEventListener.java b/...no-http-event-listener/src/test/java/io/trino/plugin/httpquery/TestHttpEventListener.java
@@ -195,6 +195,7 @@ final class TestHttpEventListener
                 Collections.emptyList(),
                 Collections.emptyList(),
                 Collections.emptyList(),
+                Collections.emptyList(),
                 Optional.empty());
 
         splitCompleteEvent = new SplitCompletedEvent(

diff --git a/...r-event-listener/src/test/java/io/trino/plugin/httpquery/TestHttpServerEventListener.java b/...r-event-listener/src/test/java/io/trino/plugin/httpquery/TestHttpServerEventListener.java
@@ -166,6 +166,7 @@ final class TestHttpServerEventListener
                 Collections.emptyList(),
                 Collections.emptyList(),
                 Collections.emptyList(),
+                Collections.emptyList(),
                 Optional.empty());
 
         queryCompleteEvent = new QueryCompletedEvent(

diff --git a/...ino-kafka-event-listener/src/test/java/io/trino/plugin/eventlistener/kafka/TestUtils.java b/...ino-kafka-event-listener/src/test/java/io/trino/plugin/eventlistener/kafka/TestUtils.java
@@ -157,6 +157,7 @@ private TestUtils() {}
                 Collections.emptyList(),
                 Collections.emptyList(),
                 Collections.emptyList(),
+                Collections.emptyList(),
                 Optional.empty());
 
         queryFailureInfo = Optional.of(

diff --git a/...nt-listener/src/test/java/io/trino/plugin/eventlistener/mysql/TestMysqlEventListener.java b/...nt-listener/src/test/java/io/trino/plugin/eventlistener/mysql/TestMysqlEventListener.java
@@ -133,6 +133,8 @@ final class TestMysqlEventListener
             // not stored
             Collections.emptyList(),
             // not stored
+            Collections.emptyList(),
+            // not stored
             List.of("{operator: \"operator1\"}", "{operator: \"operator2\"}"),
             // not stored
             Collections.emptyList(),
@@ -293,6 +295,8 @@ final class TestMysqlEventListener
             Collections.emptyList(),
             // not stored
             Collections.emptyList(),
+            // not stored
+            Collections.emptyList(),
             Collections.emptyList(),
             Collections.emptyList(),
             // not stored

diff --git a/plugin/trino-openlineage/src/test/java/io/trino/plugin/openlineage/TrinoEventData.java b/plugin/trino-openlineage/src/test/java/io/trino/plugin/openlineage/TrinoEventData.java
@@ -137,6 +137,7 @@ private TrinoEventData()
                 Collections.emptyList(),
                 Collections.emptyList(),
                 Collections.emptyList(),
+                Collections.emptyList(),
                 Optional.empty());
 
         queryCompleteEvent = new QueryCompletedEvent(