Allow exposing custom metrics from the connector

trinodb · Jul 2, 2021 · 38305b7 · 38305b7
1 parent d159c6c
commit 38305b7
Show file tree

Hide file tree

Showing 16 changed files with 228 additions and 1 deletion.
diff --git a/core/trino-main/src/main/java/io/trino/operator/OperatorContext.java b/core/trino-main/src/main/java/io/trino/operator/OperatorContext.java
@@ -28,6 +28,7 @@
 import io.trino.operator.OperationTimer.OperationTiming;
 import io.trino.spi.Page;
 import io.trino.spi.TrinoException;
+import io.trino.spi.metrics.Metrics;
 import io.trino.sql.planner.plan.PlanNodeId;
 
 import javax.annotation.Nullable;
@@ -83,6 +84,7 @@ public class OperatorContext
     private final CounterStat outputPositions = new CounterStat();
 
     private final AtomicLong dynamicFilterSplitsProcessed = new AtomicLong();
+    private final AtomicReference<Metrics> metrics = new AtomicReference<>(Metrics.EMPTY);  // this is not incremental, but gets overwritten by the latest value.
 
     private final AtomicLong physicalWrittenDataSize = new AtomicLong();
 
@@ -219,6 +221,16 @@ public void recordDynamicFilterSplitProcessed(long dynamicFilterSplits)
         dynamicFilterSplitsProcessed.getAndAdd(dynamicFilterSplits);
     }
 
+    /**
+     * Overwrites the metrics with the latest one.
+     *
+     * @param metrics Latest operator's metrics.
+     */
+    public void setLatestMetrics(Metrics metrics)
+    {
+        this.metrics.set(metrics);
+    }
+
     public void recordPhysicalWrittenData(long sizeInBytes)
     {
         physicalWrittenDataSize.getAndAdd(sizeInBytes);
@@ -531,6 +543,7 @@ public OperatorStats getOperatorStats()
                 outputPositions.getTotalCount(),
 
                 dynamicFilterSplitsProcessed.get(),
+                metrics.get(),
 
                 succinctBytes(physicalWrittenDataSize.get()),
 

diff --git a/core/trino-main/src/main/java/io/trino/operator/OperatorStats.java b/core/trino-main/src/main/java/io/trino/operator/OperatorStats.java
@@ -19,6 +19,7 @@
 import io.airlift.units.DataSize;
 import io.airlift.units.Duration;
 import io.trino.spi.Mergeable;
+import io.trino.spi.metrics.Metrics;
 import io.trino.sql.planner.plan.PlanNodeId;
 
 import javax.annotation.Nullable;
@@ -63,6 +64,7 @@ public class OperatorStats
     private final long outputPositions;
 
     private final long dynamicFilterSplitsProcessed;
+    private final Metrics metrics;
 
     private final DataSize physicalWrittenDataSize;
 
@@ -115,6 +117,7 @@ public OperatorStats(
             @JsonProperty("outputPositions") long outputPositions,
 
             @JsonProperty("dynamicFilterSplitsProcessed") long dynamicFilterSplitsProcessed,
+            @JsonProperty("metrics") Metrics metrics,
 
             @JsonProperty("physicalWrittenDataSize") DataSize physicalWrittenDataSize,
 
@@ -169,6 +172,7 @@ public OperatorStats(
         this.outputPositions = outputPositions;
 
         this.dynamicFilterSplitsProcessed = dynamicFilterSplitsProcessed;
+        this.metrics = requireNonNull(metrics);
 
         this.physicalWrittenDataSize = requireNonNull(physicalWrittenDataSize, "physicalWrittenDataSize is null");
 
@@ -332,6 +336,12 @@ public long getDynamicFilterSplitsProcessed()
         return dynamicFilterSplitsProcessed;
     }
 
+    @JsonProperty
+    public Metrics getMetrics()
+    {
+        return metrics;
+    }
+
     @JsonProperty
     public DataSize getPhysicalWrittenDataSize()
     {
@@ -451,6 +461,7 @@ public OperatorStats add(Iterable<OperatorStats> operators)
         long outputPositions = this.outputPositions;
 
         long dynamicFilterSplitsProcessed = this.dynamicFilterSplitsProcessed;
+        Metrics.Accumulator metricsAccumulator = Metrics.accumulator().add(this.getMetrics());
 
         long physicalWrittenDataSize = this.physicalWrittenDataSize.toBytes();
 
@@ -498,6 +509,7 @@ public OperatorStats add(Iterable<OperatorStats> operators)
             outputPositions += operator.getOutputPositions();
 
             dynamicFilterSplitsProcessed += operator.getDynamicFilterSplitsProcessed();
+            metricsAccumulator.add(operator.getMetrics());
 
             physicalWrittenDataSize += operator.getPhysicalWrittenDataSize().toBytes();
 
@@ -557,6 +569,7 @@ public OperatorStats add(Iterable<OperatorStats> operators)
                 outputPositions,
 
                 dynamicFilterSplitsProcessed,
+                metricsAccumulator.get(),
 
                 succinctBytes(physicalWrittenDataSize),
 
@@ -623,6 +636,7 @@ public OperatorStats summarize()
                 outputDataSize,
                 outputPositions,
                 dynamicFilterSplitsProcessed,
+                metrics,
                 physicalWrittenDataSize,
                 blockedWall,
                 finishCalls,

diff --git a/core/trino-main/src/main/java/io/trino/operator/ScanFilterAndProjectOperator.java b/core/trino-main/src/main/java/io/trino/operator/ScanFilterAndProjectOperator.java
@@ -39,6 +39,7 @@
 import io.trino.spi.connector.RecordCursor;
 import io.trino.spi.connector.RecordPageSource;
 import io.trino.spi.connector.UpdatablePageSource;
+import io.trino.spi.metrics.Metrics;
 import io.trino.spi.type.Type;
 import io.trino.split.EmptySplit;
 import io.trino.split.PageSourceProvider;
@@ -79,6 +80,7 @@ public class ScanFilterAndProjectOperator
     private long physicalBytes;
     private long readTimeNanos;
     private long dynamicFilterSplitsProcessed;
+    private Metrics metrics = Metrics.EMPTY;
 
     private ScanFilterAndProjectOperator(
             Session session,
@@ -160,6 +162,12 @@ public long getDynamicFilterSplitsProcessed()
         return dynamicFilterSplitsProcessed;
     }
 
+    @Override
+    public Metrics getConnectorMetrics()
+    {
+        return metrics;
+    }
+
     @Override
     public WorkProcessor<Page> getOutputPages()
     {
@@ -172,6 +180,7 @@ public void close()
         if (pageSource != null) {
             try {
                 pageSource.close();
+                metrics = pageSource.getMetrics();
             }
             catch (IOException e) {
                 throw new UncheckedIOException(e);
@@ -394,6 +403,7 @@ public ProcessState<Page> process()
             processedPositions += page.getPositionCount();
             physicalBytes = pageSource.getCompletedBytes();
             readTimeNanos = pageSource.getReadTimeNanos();
+            metrics = pageSource.getMetrics();
 
             return ProcessState.ofResult(page);
         }

diff --git a/core/trino-main/src/main/java/io/trino/operator/TableScanOperator.java b/core/trino-main/src/main/java/io/trino/operator/TableScanOperator.java
@@ -247,6 +247,7 @@ public void finish()
                 throw new UncheckedIOException(e);
             }
             systemMemoryContext.setBytes(source.getSystemMemoryUsage());
+            operatorContext.setLatestMetrics(source.getMetrics());
         }
     }
 
@@ -322,7 +323,7 @@ public Page getOutput()
 
         // updating system memory usage should happen after page is loaded.
         systemMemoryContext.setBytes(source.getSystemMemoryUsage());
-
+        operatorContext.setLatestMetrics(source.getMetrics());
         return page;
     }
 }
diff --git a/core/trino-main/src/main/java/io/trino/operator/TableScanWorkProcessorOperator.java b/core/trino-main/src/main/java/io/trino/operator/TableScanWorkProcessorOperator.java
@@ -32,6 +32,7 @@
 import io.trino.spi.connector.DynamicFilter;
 import io.trino.spi.connector.EmptyPageSource;
 import io.trino.spi.connector.UpdatablePageSource;
+import io.trino.spi.metrics.Metrics;
 import io.trino.split.EmptySplit;
 import io.trino.split.PageSourceProvider;
 
@@ -117,6 +118,12 @@ public long getDynamicFilterSplitsProcessed()
         return splitToPages.getDynamicFilterSplitsProcessed();
     }
 
+    @Override
+    public Metrics getConnectorMetrics()
+    {
+        return splitToPages.source.getMetrics();
+    }
+
     @Override
     public Duration getReadTime()
     {

diff --git a/core/trino-main/src/main/java/io/trino/operator/WorkProcessorPipelineSourceOperator.java b/core/trino-main/src/main/java/io/trino/operator/WorkProcessorPipelineSourceOperator.java
@@ -29,6 +29,7 @@
 import io.trino.operator.WorkProcessor.ProcessState;
 import io.trino.spi.Page;
 import io.trino.spi.connector.UpdatablePageSource;
+import io.trino.spi.metrics.Metrics;
 import io.trino.spi.type.Type;
 import io.trino.sql.planner.LocalExecutionPlanner.OperatorFactoryWithTypes;
 import io.trino.sql.planner.plan.PlanNodeId;
@@ -39,6 +40,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Supplier;
 
 import static com.google.common.base.Preconditions.checkState;
@@ -232,11 +234,14 @@ private void workProcessorOperatorStateMonitor(WorkProcessor.ProcessState<Page>
             long deltaReadTimeNanos = deltaAndSet(context.readTimeNanos, sourceOperator.getReadTime().roundTo(NANOSECONDS));
 
             long deltaDynamicFilterSplitsProcessed = deltaAndSet(context.dynamicFilterSplitsProcessed, sourceOperator.getDynamicFilterSplitsProcessed());
+            Metrics metrics = sourceOperator.getConnectorMetrics();
+            context.connectorMetrics.set(metrics);
 
             operatorContext.recordPhysicalInputWithTiming(deltaPhysicalInputDataSize, deltaPhysicalInputPositions, deltaReadTimeNanos);
             operatorContext.recordNetworkInput(deltaInternalNetworkInputDataSize, deltaInternalNetworkInputPositions);
             operatorContext.recordProcessedInput(deltaInputDataSize, deltaInputPositions);
             operatorContext.recordDynamicFilterSplitProcessed(deltaDynamicFilterSplitsProcessed);
+            operatorContext.setLatestMetrics(metrics);
         }
 
         if (state.getType() == FINISHED) {
@@ -337,6 +342,7 @@ private List<OperatorStats> getNestedOperatorStats()
                         context.outputPositions.get(),
 
                         context.dynamicFilterSplitsProcessed.get(),
+                        context.connectorMetrics.get(),
 
                         DataSize.ofBytes(0),
 
@@ -676,6 +682,7 @@ private static class WorkProcessorOperatorContext
         final AtomicLong outputPositions = new AtomicLong();
 
         final AtomicLong dynamicFilterSplitsProcessed = new AtomicLong();
+        final AtomicReference<Metrics> connectorMetrics = new AtomicReference<>(Metrics.EMPTY);
 
         final AtomicLong peakUserMemoryReservation = new AtomicLong();
         final AtomicLong peakSystemMemoryReservation = new AtomicLong();

diff --git a/core/trino-main/src/main/java/io/trino/operator/WorkProcessorSourceOperator.java b/core/trino-main/src/main/java/io/trino/operator/WorkProcessorSourceOperator.java
@@ -16,6 +16,7 @@
 import io.airlift.units.DataSize;
 import io.airlift.units.Duration;
 import io.trino.spi.connector.UpdatablePageSource;
+import io.trino.spi.metrics.Metrics;
 
 import java.util.Optional;
 import java.util.function.Supplier;
@@ -67,4 +68,6 @@ default long getDynamicFilterSplitsProcessed()
     {
         return 0;
     }
+
+    Metrics getConnectorMetrics();
 }
diff --git a/core/trino-main/src/main/java/io/trino/operator/WorkProcessorSourceOperatorAdapter.java b/core/trino-main/src/main/java/io/trino/operator/WorkProcessorSourceOperatorAdapter.java
@@ -21,6 +21,7 @@
 import io.trino.metadata.Split;
 import io.trino.spi.Page;
 import io.trino.spi.connector.UpdatablePageSource;
+import io.trino.spi.metrics.Metrics;
 import io.trino.sql.planner.plan.PlanNodeId;
 
 import java.util.ArrayList;
@@ -175,6 +176,7 @@ public boolean isFinished()
     public void close()
             throws Exception
     {
+        operatorContext.setLatestMetrics(sourceOperator.getConnectorMetrics());
         sourceOperator.close();
     }
 
@@ -191,6 +193,7 @@ private void updateOperatorStats()
         long currentInputPositions = sourceOperator.getInputPositions();
 
         long currentDynamicFilterSplitsProcessed = sourceOperator.getDynamicFilterSplitsProcessed();
+        Metrics currentMetrics = sourceOperator.getConnectorMetrics();
 
         if (currentPhysicalInputBytes != previousPhysicalInputBytes
                 || currentPhysicalInputPositions != previousPhysicalInputPositions
@@ -229,6 +232,8 @@ private void updateOperatorStats()
             operatorContext.recordDynamicFilterSplitProcessed(currentDynamicFilterSplitsProcessed - previousDynamicFilterSplitsProcessed);
             previousDynamicFilterSplitsProcessed = currentDynamicFilterSplitsProcessed;
         }
+
+        operatorContext.setLatestMetrics(currentMetrics);
     }
 
     private static class SplitBuffer

diff --git a/core/trino-main/src/test/java/io/trino/execution/TestQueryStats.java b/core/trino-main/src/test/java/io/trino/execution/TestQueryStats.java
@@ -22,6 +22,7 @@
 import io.trino.operator.OperatorStats;
 import io.trino.operator.TableWriterOperator;
 import io.trino.spi.eventlistener.StageGcStatistics;
+import io.trino.spi.metrics.Metrics;
 import io.trino.sql.planner.plan.PlanNodeId;
 import org.joda.time.DateTime;
 import org.testng.annotations.Test;
@@ -62,6 +63,7 @@ public class TestQueryStats
                     succinctBytes(116L),
                     117L,
                     1833,
+                    Metrics.EMPTY,
                     succinctBytes(118L),
                     new Duration(119, NANOSECONDS),
                     120L,
@@ -101,6 +103,7 @@ public class TestQueryStats
                     succinctBytes(216L),
                     217L,
                     2833,
+                    Metrics.EMPTY,
                     succinctBytes(218L),
                     new Duration(219, NANOSECONDS),
                     220L,
@@ -140,6 +143,7 @@ public class TestQueryStats
                     succinctBytes(316L),
                     317L,
                     3833,
+                    Metrics.EMPTY,
                     succinctBytes(318L),
                     new Duration(319, NANOSECONDS),
                     320L,

diff --git a/core/trino-main/src/test/java/io/trino/operator/TestOperatorStats.java b/core/trino-main/src/test/java/io/trino/operator/TestOperatorStats.java
@@ -13,11 +13,13 @@
  */
 package io.trino.operator;
 
+import com.google.common.collect.ImmutableMap;
 import io.airlift.json.JsonCodec;
 import io.airlift.units.DataSize;
 import io.airlift.units.Duration;
 import io.trino.connector.CatalogName;
 import io.trino.operator.PartitionedOutputOperator.PartitionedOutputInfo;
+import io.trino.spi.metrics.Metrics;
 import io.trino.sql.planner.plan.PlanNodeId;
 import org.testng.annotations.Test;
 
@@ -59,6 +61,7 @@ public class TestOperatorStats
             DataSize.ofBytes(12),
             13,
             533,
+            Metrics.EMPTY,
 
             DataSize.ofBytes(14),
 
@@ -106,6 +109,7 @@ public class TestOperatorStats
             DataSize.ofBytes(12),
             13,
             533,
+            Metrics.EMPTY,
 
             DataSize.ofBytes(14),
 
@@ -163,6 +167,7 @@ public static void assertExpectedOperatorStats(OperatorStats actual)
         assertEquals(actual.getOutputPositions(), 13);
 
         assertEquals(actual.getDynamicFilterSplitsProcessed(), 533);
+        assertEquals(actual.getMetrics().getMetrics(), ImmutableMap.of());
 
         assertEquals(actual.getPhysicalWrittenDataSize(), DataSize.ofBytes(14));
 

diff --git a/core/trino-main/src/test/java/io/trino/operator/TestWorkProcessorPipelineSourceOperator.java b/core/trino-main/src/test/java/io/trino/operator/TestWorkProcessorPipelineSourceOperator.java
@@ -26,6 +26,7 @@
 import io.trino.operator.WorkProcessorAssertion.Transform;
 import io.trino.spi.Page;
 import io.trino.spi.connector.UpdatablePageSource;
+import io.trino.spi.metrics.Metrics;
 import io.trino.sql.planner.LocalExecutionPlanner.OperatorFactoryWithTypes;
 import io.trino.sql.planner.plan.PlanNodeId;
 import org.testng.annotations.AfterClass;
@@ -392,6 +393,12 @@ public Duration getReadTime()
             return new Duration(7, NANOSECONDS);
         }
 
+        @Override
+        public Metrics getConnectorMetrics()
+        {
+            return Metrics.EMPTY;
+        }
+
         @Override
         public WorkProcessor<Page> getOutputPages()
         {