Allow getting metrics after the operator is closed

core/trino-main/src/main/java/io/trino/operator/WorkProcessorPipelineSourceOperator.java

@@ -538,6 +538,11 @@ private void closeOperators(int lastOperatorIndex)
                             operatorContext.getDriverContext().getTaskId());
                 }
                 finally {
+                    workProcessorOperatorContext.metrics.set(operator.getMetrics());
+                    if (operator instanceof WorkProcessorSourceOperator) {
+                        WorkProcessorSourceOperator sourceOperator = (WorkProcessorSourceOperator) operator;
+                        workProcessorOperatorContext.connectorMetrics.set(sourceOperator.getConnectorMetrics());
+                    }
                     workProcessorOperatorContext.memoryTrackingContext.close();
                     workProcessorOperatorContext.finalOperatorInfo = operator.getOperatorInfo().orElse(null);
                     workProcessorOperatorContext.operator = null;

core/trino-main/src/main/java/io/trino/operator/WorkProcessorSourceOperatorAdapter.java

@@ -176,9 +176,9 @@ public boolean isFinished()
     public void close()
             throws Exception
     {
+        sourceOperator.close();
         operatorContext.setLatestMetrics(sourceOperator.getMetrics());
         operatorContext.setLatestConnectorMetrics(sourceOperator.getConnectorMetrics());
-        sourceOperator.close();
     }
 
     private void updateOperatorStats()

core/trino-main/src/test/java/io/trino/operator/TestWorkProcessorPipelineSourceOperator.java

@@ -208,8 +208,12 @@ public void testWorkProcessorPipelineSourceOperator()
         // assert source operator stats are correct
         OperatorStats sourceOperatorStats = operatorStats.get(0);
 
-        assertEquals(sourceOperatorStats.getMetrics().getMetrics(), ImmutableMap.of("testSourceMetric", new LongCount(1)));
-        assertEquals(sourceOperatorStats.getConnectorMetrics().getMetrics(), ImmutableMap.of("testSourceConnectorMetric", new LongCount(2)));
+        assertEquals(sourceOperatorStats.getMetrics().getMetrics(), ImmutableMap.of(
+                "testSourceMetric", new LongCount(1),
+                "testSourceClosed", new LongCount(1)));
+        assertEquals(sourceOperatorStats.getConnectorMetrics().getMetrics(), ImmutableMap.of(
+                "testSourceConnectorMetric", new LongCount(2),
+                "testSourceConnectorClosed", new LongCount(1)));
 
         assertEquals(sourceOperatorStats.getDynamicFilterSplitsProcessed(), 42L);
 
@@ -239,8 +243,12 @@ public void testWorkProcessorPipelineSourceOperator()
 
         // assert pipeline metrics
         List<OperatorStats> operatorSummaries = pipelineStats.getOperatorSummaries();
-        assertEquals(operatorSummaries.get(0).getMetrics().getMetrics(), ImmutableMap.of("testSourceMetric", new LongCount(1)));
-        assertEquals(operatorSummaries.get(0).getConnectorMetrics().getMetrics(), ImmutableMap.of("testSourceConnectorMetric", new LongCount(2)));
+        assertEquals(operatorSummaries.get(0).getMetrics().getMetrics(), ImmutableMap.of(
+                "testSourceMetric", new LongCount(1),
+                "testSourceClosed", new LongCount(1)));
+        assertEquals(operatorSummaries.get(0).getConnectorMetrics().getMetrics(), ImmutableMap.of(
+                "testSourceConnectorMetric", new LongCount(2),
+                "testSourceConnectorClosed", new LongCount(1)));
         assertEquals(operatorSummaries.get(1).getMetrics().getMetrics(), ImmutableMap.of("testOperatorMetric", new LongCount(1)));
     }
 
@@ -425,13 +433,18 @@ public long getDynamicFilterSplitsProcessed()
         @Override
         public Metrics getMetrics()
         {
-            return new Metrics(ImmutableMap.of("testSourceMetric", new LongCount(1)));
+            System.err.println("closed: " + closed);
+            return new Metrics(ImmutableMap.of(
+                    "testSourceMetric", new LongCount(1),
+                    "testSourceClosed", new LongCount(closed ? 1 : 0)));
         }
 
         @Override
         public Metrics getConnectorMetrics()
         {
-            return new Metrics(ImmutableMap.of("testSourceConnectorMetric", new LongCount(2)));
+            return new Metrics(ImmutableMap.of(
+                    "testSourceConnectorMetric", new LongCount(2),
+                    "testSourceConnectorClosed", new LongCount(closed ? 1 : 0)));
         }
 
         @Override

core/trino-spi/src/main/java/io/trino/spi/connector/ConnectorPageSource.java

@@ -87,6 +87,7 @@ default CompletableFuture<?> isBlocked()
      * Returns the connector's metrics, mapping a metric ID to its latest value.
      * Each call must return an immutable snapshot of available metrics.
      * Same ID metrics are merged across all tasks and exposed via OperatorStats.
+     * This method can be called after the page source is closed.
      */
     default Metrics getMetrics()
     {

plugin/trino-memory/src/main/java/io/trino/plugin/memory/MemoryPageSourceProvider.java

@@ -95,6 +95,7 @@ private static class DynamicFilteringPageSource
         private final boolean enableLazyDynamicFiltering;
         private long rows;
         private long completedPositions;
+        private boolean closed;
 
         private DynamicFilteringPageSource(FixedPageSource delegate, List<ColumnHandle> columns, DynamicFilter dynamicFilter, boolean enableLazyDynamicFiltering)
         {
@@ -171,14 +172,15 @@ public long getSystemMemoryUsage()
         public void close()
         {
             delegate.close();
+            closed = true;
         }
 
         @Override
         public Metrics getMetrics()
         {
             return new Metrics(ImmutableMap.of(
                     "rows", new LongCount(rows),
-                    "finished", new LongCount(isFinished() ? 1 : 0),
+                    "finished", new LongCount(closed ? 1 : 0),
                     "started", new LongCount(1)));
         }
     }