From 99e8f664a9cc0a73ae97ea4409743b9bb479f0ec Mon Sep 17 00:00:00 2001
From: Gian Merlino <gianmerlino@gmail.com>
Date: Fri, 13 Sep 2024 15:47:51 -0700
Subject: [PATCH] Add "includeAllCounters()" to WorkerContext. (#17047)

This removes the need to read it from the query context.
---
 .../main/java/org/apache/druid/msq/exec/WorkerContext.java  | 6 ++++++
 .../src/main/java/org/apache/druid/msq/exec/WorkerImpl.java | 2 +-
 .../org/apache/druid/msq/indexing/IndexerWorkerContext.java | 6 ++++++
 .../org/apache/druid/msq/test/MSQTestWorkerContext.java     | 6 ++++++
 4 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerContext.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerContext.java
index 666115d774cf..95a4ce7c7ba5 100644
--- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerContext.java
+++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerContext.java
@@ -27,6 +27,7 @@
 import org.apache.druid.msq.kernel.FrameProcessorFactory;
 import org.apache.druid.msq.kernel.QueryDefinition;
 import org.apache.druid.msq.kernel.WorkOrder;
+import org.apache.druid.msq.util.MultiStageQueryContext;
 import org.apache.druid.server.DruidNode;
 
 import java.io.File;
@@ -98,4 +99,9 @@ public interface WorkerContext
   DruidNode selfNode();
 
   DataServerQueryHandlerFactory dataServerQueryHandlerFactory();
+
+  /**
+   * Whether to include all counters in reports. See {@link MultiStageQueryContext#CTX_INCLUDE_ALL_COUNTERS} for detail.
+   */
+  boolean includeAllCounters();
 }
diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java
index 92664feeabbb..f28d1be5e614 100644
--- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java
+++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/WorkerImpl.java
@@ -407,7 +407,7 @@ private void handleNewWorkOrder(
     kernel.startReading();
 
     final QueryContext queryContext = task != null ? QueryContext.of(task.getContext()) : QueryContext.empty();
-    final boolean includeAllCounters = MultiStageQueryContext.getIncludeAllCounters(queryContext);
+    final boolean includeAllCounters = context.includeAllCounters();
     final RunWorkOrder runWorkOrder = new RunWorkOrder(
         task.getControllerTaskId(),
         workOrder,
diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerWorkerContext.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerWorkerContext.java
index 0b3063ef48ba..c36b8e291db9 100644
--- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerWorkerContext.java
+++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/indexing/IndexerWorkerContext.java
@@ -312,6 +312,12 @@ public DataServerQueryHandlerFactory dataServerQueryHandlerFactory()
     return dataServerQueryHandlerFactory;
   }
 
+  @Override
+  public boolean includeAllCounters()
+  {
+    return includeAllCounters;
+  }
+
   private synchronized ServiceLocator makeControllerLocator(final String controllerId)
   {
     if (controllerLocator == null) {
diff --git a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerContext.java b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerContext.java
index 082429a9d7b1..1b92f468fced 100644
--- a/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerContext.java
+++ b/extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/test/MSQTestWorkerContext.java
@@ -159,6 +159,12 @@ public DataServerQueryHandlerFactory dataServerQueryHandlerFactory()
     return injector.getInstance(DataServerQueryHandlerFactory.class);
   }
 
+  @Override
+  public boolean includeAllCounters()
+  {
+    return true;
+  }
+
   class FrameContextImpl implements FrameContext
   {
     private final File tempDir;