Skip to content

Add dataSource and taskId dimensions to GroupByStatsMonitor for peons #18709

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
abhishekrb19 merged 6 commits into from
Nov 6, 2025
Merged

Add dataSource and taskId dimensions to GroupByStatsMonitor for peons #18709

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
72ed511
Add taskID and dataSource to GroupByStatsMonitor.
abhishekrb19 Oct 31, 2025
be53bda
Update docs
abhishekrb19 Oct 31, 2025
d788744
Add new method for taskId to dismabiguate id
abhishekrb19 Oct 31, 2025
e9d29bd
Review comments: add javadoc, dimensions constructor
abhishekrb19 Nov 4, 2025
32b6da6
Merge branch 'master' into group_by_stats_datasource_taskId & address…
abhishekrb19 Nov 4, 2025
1f5492a
Merge branch 'master' into group_by_stats_datasource_taskId
abhishekrb19 Nov 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion docs/ingestion/native-batch.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -458,7 +458,7 @@ Range partitioning is not possible on multi-value dimensions. If one of the prov
|`type`|Set the value to `range`.|none|yes|
|`partitionDimensions`|An array of dimensions to partition on. Order the dimensions from most frequently queried to least frequently queried. For best results, limit your number of dimensions to between three and five dimensions.|none|yes|
|`targetRowsPerSegment`|Target number of rows to include in a partition, should be a number that targets segments of 500MB\~1GB.|none|either this or `maxRowsPerSegment`|
|maxRowsPerSegment|Soft max for the number of rows to include in a partition.|none|either this or `targetRowsPerSegment`|
|`maxRowsPerSegment`|Soft max for the number of rows to include in a partition.|none|either this or `targetRowsPerSegment`|
|`assumeGrouped`|Assume that input data has already been grouped on time and dimensions. Ingestion will run faster, but may choose sub-optimal partitions if this assumption is violated.|false|no|

#### Benefits of range partitioning
Expand Down
19 changes: 12 additions & 7 deletions docs/operations/metrics.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,11 @@ Most metric values reset each emission period, as specified in `druid.monitoring

### Real-time

:::info
Monitors on peons that previously emitted the `id` dimension from `JettyMonitor`, `OshiSysMonitor`, `JvmMonitor`, `JvmCpuMonitor`, `JvmThreadsMonitor` and `SysMonitor`
to represent the task ID are deprecated and will be removed in a future release. Use the `taskId` dimension instead.
Comment on lines +123 to +124

@kfaraz kfaraz Nov 6, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
Monitors on peons that previously emitted the `id` dimension from `JettyMonitor`, `OshiSysMonitor`, `JvmMonitor`, `JvmCpuMonitor`, `JvmThreadsMonitor` and `SysMonitor`
to represent the task ID are deprecated and will be removed in a future release. Use the `taskId` dimension instead.
The dimension `id` emitted by peons to represent the task ID (by monitors such as `JettyMonitor`, `OshiSysMonitor`, `JvmMonitor`, `JvmCpuMonitor`, `JvmThreadsMonitor` and `SysMonitor`)
is now deprecated and will be removed in a future release. Use the `taskId` dimension instead.

:::

|Metric|Description|Dimensions|Normal value|
|------|-----------|----------|------------|
|`query/time`|Milliseconds taken to complete a query.|<p>Common: `dataSource`, `type`, `interval`, `hasFilters`, `duration`, `context`, `remoteAddress`, `id`, `statusCode`.</p><p> Aggregation Queries: `numMetrics`, `numComplexMetrics`.</p><p> GroupBy: `numDimensions`.</p><p> TopN: `threshold`, `dimension`.</p>|< 1s|
Expand All @@ -131,13 +136,13 @@ Most metric values reset each emission period, as specified in `druid.monitoring
|`query/failed/count`|Number of failed queries.|This metric is only available if the `QueryCountStatsMonitor` module is included.||
|`query/interrupted/count`|Number of queries interrupted due to cancellation.|This metric is only available if the `QueryCountStatsMonitor` module is included.||
|`query/timeout/count`|Number of timed out queries.|This metric is only available if the `QueryCountStatsMonitor` module is included.||
|`mergeBuffer/pendingRequests`|Number of requests waiting to acquire a batch of buffers from the merge buffer pool.|This metric is only available if the `GroupByStatsMonitor` module is included.|Should be ideally 0, though a higher number isn't representative of a problem.|
|`mergeBuffer/used`|Number of merge buffers used from the merge buffer pool.|This metric is only available if the `GroupByStatsMonitor` module is included.|Depends on the number of groupBy queries needing merge buffers.|
|`mergeBuffer/queries`|Number of groupBy queries that acquired a batch of buffers from the merge buffer pool.|This metric is only available if the `GroupByStatsMonitor` module is included.|Depends on the number of groupBy queries needing merge buffers.|
|`mergeBuffer/acquisitionTimeNs`|Total time in nanoseconds to acquire merge buffer for groupBy queries.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
|`groupBy/spilledQueries`|Number of groupBy queries that have spilled onto the disk.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
|`groupBy/spilledBytes`|Number of bytes spilled on the disk by the groupBy queries.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
|`groupBy/mergeDictionarySize`|Size of on-heap merge dictionary in bytes.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
|`mergeBuffer/pendingRequests`|Number of requests waiting to acquire a batch of buffers from the merge buffer pool. This metric is only available if the `GroupByStatsMonitor` module is included.|`dataSource`, `taskId`|Should be ideally 0, though a higher number isn't representative of a problem.|
|`mergeBuffer/used`|Number of merge buffers used from the merge buffer pool. This metric is only available if the `GroupByStatsMonitor` module is included.|`dataSource`, `taskId`|Depends on the number of groupBy queries needing merge buffers.|
|`mergeBuffer/queries`|Number of groupBy queries that acquired a batch of buffers from the merge buffer pool. This metric is only available if the `GroupByStatsMonitor` module is included.|`dataSource`, `taskId`|Depends on the number of groupBy queries needing merge buffers.|
|`mergeBuffer/acquisitionTimeNs`|Total time in nanoseconds to acquire merge buffer for groupBy queries. This metric is only available if the `GroupByStatsMonitor` module is included.|`dataSource`, `taskId`|Varies|
|`groupBy/spilledQueries`|Number of groupBy queries that have spilled onto the disk. This metric is only available if the `GroupByStatsMonitor` module is included.|`dataSource`, `taskId`|Varies|
|`groupBy/spilledBytes`|Number of bytes spilled on the disk by the groupBy queries. This metric is only available if the `GroupByStatsMonitor` module is included.|`dataSource`, `taskId`|Varies|
|`groupBy/mergeDictionarySize`|Size of on-heap merge dictionary in bytes. This metric is only available if the `GroupByStatsMonitor` module is included.|`dataSource`, `taskId`|Varies|

### Jetty

Expand Down
10 changes: 7 additions & 3 deletions server/src/main/java/org/apache/druid/server/initialization/jetty/JettyServerModule.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -519,16 +519,20 @@ private static int getTCPAcceptQueueSize()
@LazySingleton
public JettyMonitor getJettyMonitor(DataSourceTaskIdHolder dataSourceTaskIdHolder)
{
return new JettyMonitor(dataSourceTaskIdHolder.getDataSource(), dataSourceTaskIdHolder.getTaskId());
return new JettyMonitor(
MonitorsConfig.mapOfDatasourceAndTaskID(
dataSourceTaskIdHolder.getDataSource(), dataSourceTaskIdHolder.getTaskId()
)
);
}

public static class JettyMonitor extends AbstractMonitor
{
private final Map<String, String[]> dimensions;

public JettyMonitor(String dataSource, String taskId)
public JettyMonitor(Map<String, String[]> dimensions)
{
this.dimensions = MonitorsConfig.mapOfDatasourceAndTaskID(dataSource, taskId);
this.dimensions = dimensions;
}

@Override
Expand Down
14 changes: 14 additions & 0 deletions server/src/main/java/org/apache/druid/server/metrics/DataSourceTaskIdHolder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,11 @@
import org.apache.druid.server.coordination.BroadcastDatasourceLoadingSpec;
import org.apache.druid.server.lookup.cache.LookupLoadingSpec;

import javax.annotation.Nullable;

/**
* This holder is only applicable to {@code CliPeon} servers.
*/
public class DataSourceTaskIdHolder
{
public static final String DATA_SOURCE_BINDING = "druidDataSource";
Expand All @@ -34,6 +39,7 @@ public class DataSourceTaskIdHolder
@Named(DATA_SOURCE_BINDING)
@Inject(optional = true)
String dataSource = null;

@Named(TASK_ID_BINDING)
@Inject(optional = true)
String taskId = null;
Expand All @@ -46,11 +52,19 @@ public class DataSourceTaskIdHolder
@Inject(optional = true)
BroadcastDatasourceLoadingSpec broadcastDatasourceLoadingSpec = BroadcastDatasourceLoadingSpec.ALL;

/**
* @return the taskId for CliPeon servers; {@code null} for all other servers.
*/
@Nullable

@kfaraz kfaraz Nov 2, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When can this be null?

@abhishekrb19 abhishekrb19 Nov 2, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For peons, the CliPeon would provide the values of taskId and dataSource.

For all other servers, this would be null as these fields are initialized to null::

  @Named(DATA_SOURCE_BINDING)
  @Inject(optional = true)
  String dataSource = null;
  @Named(TASK_ID_BINDING)
  @Inject(optional = true)
  String taskId = null;

@kfaraz kfaraz Nov 3, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please add a short javadoc on these getters mentioning that these are non-null only for CliPeon servers.

In a follow up PR, we should also fix up the way that DatasourceTaskIdHolder is wired and bind that class itself only for peons. There is no meaning in having this instance bound but contain null values for datasource and task ID.

@abhishekrb19 abhishekrb19 Nov 4, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added javadoc!

In a follow up PR, we should also fix up the way that DatasourceTaskIdHolder is wired and bind that class itself only for peons. There is no meaning in having this instance bound but contain null values for datasource and task ID.

Yeah, there's some commentary on why it's setup the way it is. I think we'll need to revisit this Guice bindings for this.

@kfaraz kfaraz Nov 4, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, given that we have embedded tests now, it should be easy to fix as we can iterate much faster.

public String getDataSource()
{
return dataSource;
}

/**
* @return the dataSource for CliPeon servers; {@code null} for all other servers.
*/
@Nullable
public String getTaskId()
{
return taskId;
Expand Down
11 changes: 10 additions & 1 deletion server/src/main/java/org/apache/druid/server/metrics/GroupByStatsMonitor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,11 @@
import org.apache.druid.java.util.emitter.service.ServiceEmitter;
import org.apache.druid.java.util.emitter.service.ServiceMetricEvent;
import org.apache.druid.java.util.metrics.AbstractMonitor;
import org.apache.druid.java.util.metrics.MonitorUtils;
import org.apache.druid.query.groupby.GroupByStatsProvider;

import java.nio.ByteBuffer;
import java.util.Map;

@LoadScope(roles = {
NodeRole.BROKER_JSON_NAME,
Expand All @@ -41,21 +43,28 @@ public class GroupByStatsMonitor extends AbstractMonitor
{
private final GroupByStatsProvider groupByStatsProvider;
private final BlockingPool<ByteBuffer> mergeBufferPool;
private final Map<String, String[]> dimensions;

@kfaraz kfaraz Nov 3, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe rename to serviceDimensions?

In fact, I wonder if it wouldn't be cleaner to just bind the datasource and taskID as @ExtraServiceDimensions so that they are used in EmitterModule and are emitted for all metrics emitted by a peon.

@abhishekrb19 abhishekrb19 Nov 4, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can rename, but I left this field as dimensions to be consistent with the other monitors.

I wonder if it wouldn't be cleaner to just bind the datasource and taskID as @ExtraServiceDimensions so that they are used in EmitterModule and are emitted for all metrics emitted by a peon.

This is a good idea! I will try to take a stab at it separately and put up a PR as it affects all monitors - we'll need to rework the Guice bindings and likely need to reconcile how the extra service dimensions will be bound to monitors in a backwards-compatible manner too (id vs taskId is something to consider).

@kfaraz kfaraz Nov 4, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can rename, but I left this field as dimensions to be consistent with the other monitors.

I see, thanks for the clarification.


@Inject
public GroupByStatsMonitor(
GroupByStatsProvider groupByStatsProvider,
@Merging BlockingPool<ByteBuffer> mergeBufferPool
@Merging BlockingPool<ByteBuffer> mergeBufferPool,
DataSourceTaskIdHolder dataSourceTaskIdHolder

@kfaraz kfaraz Nov 3, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please pass a dimensions map here instead of this holder since the holder class has meaning only for peons.

@abhishekrb19 abhishekrb19 Nov 4, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since this monitor is injected, I don’t think it can accept a Map<String, String[]> in the constructor unless it’s bound with a specific name or qualifier?

@kfaraz kfaraz Nov 4, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, we can leave that for now. When we fix up the binding of DatasourceTaskIdHolder, we can figure that out.

)
{
this.groupByStatsProvider = groupByStatsProvider;
this.mergeBufferPool = mergeBufferPool;
this.dimensions = MonitorsConfig.mapOfDatasourceAndTaskID(
dataSourceTaskIdHolder.getDataSource(),
dataSourceTaskIdHolder.getTaskId()
);
}

@Override
public boolean doMonitor(ServiceEmitter emitter)
{
final ServiceMetricEvent.Builder builder = new ServiceMetricEvent.Builder();
MonitorUtils.addDimensionsToBuilder(builder, dimensions);

emitter.emit(builder.setMetric("mergeBuffer/pendingRequests", mergeBufferPool.getPendingRequests()));

Expand Down
13 changes: 12 additions & 1 deletion server/src/main/java/org/apache/druid/server/metrics/MonitorsConfig.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
import org.apache.druid.java.util.metrics.Monitor;
import org.apache.druid.query.DruidMetrics;

import javax.annotation.Nullable;
import javax.validation.constraints.NotNull;
import java.util.ArrayList;
import java.util.HashMap;
Expand Down Expand Up @@ -79,13 +80,23 @@ public String toString()
'}';
}

public static Map<String, String[]> mapOfDatasourceAndTaskID(final String datasource, final String taskId)

@kfaraz kfaraz Nov 4, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of adding a new method, let's do the following instead:

  • use the old method itself to emit the task ID under both the id dimension and taskId dimension
  • call out the id dimension as deprecated in the docs and release notes and remove it in D37
  • update the docs for the existing metrics to indicate that they emit taskId dimension too.

@abhishekrb19 abhishekrb19 Nov 4, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.

call out the id dimension as deprecated in the docs and release notes and remove it in D37

It looks like the docs for the six monitors that use this method don’t specify their dimensions clearly. I’ve added a general info note under the Real-time metrics seciton. Please let me know if I missed anything.

/**
* @return a map of {@code datasource} and {@code taskId} dimensions if provided; otherwise, returns an empty map.
* When {@code taskId} is provided, both {@link DruidMetrics#ID} and {@link DruidMetrics#TASK_ID} dimensions are added
* to the map for backward compatibility. {@link DruidMetrics#ID} is deprecated because it's ambiguous and will be
* removed in a future release.
*/
public static Map<String, String[]> mapOfDatasourceAndTaskID(
@Nullable final String datasource,
@Nullable final String taskId
)
{
final ImmutableMap.Builder<String, String[]> builder = ImmutableMap.builder();
if (datasource != null) {
builder.put(DruidMetrics.DATASOURCE, new String[]{datasource});
}
if (taskId != null) {
builder.put(DruidMetrics.TASK_ID, new String[]{taskId});
builder.put(DruidMetrics.ID, new String[]{taskId});
}
return builder.build();
Expand Down
5 changes: 4 additions & 1 deletion server/src/test/java/org/apache/druid/server/initialization/jetty/JettyServerModuleTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@
import org.apache.druid.server.DruidNode;
import org.apache.druid.server.initialization.ServerConfig;
import org.apache.druid.server.initialization.TLSServerConfig;
import org.apache.druid.server.metrics.MonitorsConfig;
import org.apache.druid.server.security.TLSCertificateChecker;
import org.eclipse.jetty.server.Server;
import org.eclipse.jetty.util.ssl.SslContextFactory;
Expand All @@ -54,7 +55,9 @@ public void testJettyServerModule()
Mockito.when(jettyServerThreadPool.getQueueSize()).thenReturn(50);
Mockito.when(jettyServerThreadPool.getBusyThreads()).thenReturn(60);

JettyServerModule.JettyMonitor jettyMonitor = new JettyServerModule.JettyMonitor("ds", "t0");
JettyServerModule.JettyMonitor jettyMonitor = new JettyServerModule.JettyMonitor(
MonitorsConfig.mapOfDatasourceAndTaskID("ds", "t0")
);

final StubServiceEmitter serviceEmitter = new StubServiceEmitter("service", "host");
jettyMonitor.doMonitor(serviceEmitter);
Expand Down
70 changes: 61 additions & 9 deletions server/src/test/java/org/apache/druid/server/metrics/GroupByStatsMonitorTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,10 +19,18 @@

package org.apache.druid.server.metrics;

import com.google.inject.Injector;
import com.google.inject.Key;
import com.google.inject.name.Names;
import org.apache.druid.collections.BlockingPool;
import org.apache.druid.collections.DefaultBlockingPool;
import org.apache.druid.guice.GuiceInjectors;
import org.apache.druid.guice.JsonConfigProvider;
import org.apache.druid.guice.annotations.Self;
import org.apache.druid.initialization.Initialization;
import org.apache.druid.java.util.metrics.StubServiceEmitter;
import org.apache.druid.query.groupby.GroupByStatsProvider;
import org.apache.druid.server.DruidNode;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
Expand All @@ -31,6 +39,7 @@
import java.nio.ByteBuffer;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
Expand All @@ -52,12 +61,12 @@
public synchronized AggregateStats getStatsSince()
{
return new AggregateStats(
1L,
100L,
2L,
200L,
300L
);
1L,
100L,
2L,
200L,
300L
);
}
};

Expand All @@ -75,7 +84,7 @@
public void testMonitor()
{
final GroupByStatsMonitor monitor =
new GroupByStatsMonitor(groupByStatsProvider, mergeBufferPool);
new GroupByStatsMonitor(groupByStatsProvider, mergeBufferPool, new DataSourceTaskIdHolder());
final StubServiceEmitter emitter = new StubServiceEmitter("service", "host");
monitor.doMonitor(emitter);
emitter.flush();
Expand All @@ -92,6 +101,49 @@
emitter.verifyValue("groupBy/mergeDictionarySize", 300L);
}

@Test
public void testMonitorWithServiceDimensions()
{
final String dataSource = "fooDs";
final String taskId = "taskId1";
final Injector injector = Initialization.makeInjectorWithModules(
GuiceInjectors.makeStartupInjector(),
List.of(binder -> {
JsonConfigProvider.bindInstance(
binder,
Key.get(DruidNode.class, Self.class),
new DruidNode("test-inject", null, false, null, null, true, false)
);
binder.bind(Key.get(String.class, Names.named(DataSourceTaskIdHolder.DATA_SOURCE_BINDING)))
.toInstance(dataSource);
binder.bind(Key.get(String.class, Names.named(DataSourceTaskIdHolder.TASK_ID_BINDING)))
.toInstance(taskId);
})
);
Comment on lines +109 to +122

Check notice

Code scanning / CodeQL

Deprecated method or constructor invocation Note test

Invoking
Initialization.makeInjectorWithModules
Loading
should be avoided because it has been deprecated.

Copilot Autofix

AI 9 months ago

The goal is to avoid using the deprecated Initialization.makeInjectorWithModules method. This method is a Druid utility that previously created a Guice Injector using a base injector and additional modules. In recent Druid releases, its use is discouraged, with the typical replacement being the use of Guice's own APIs (namely, Guice.createInjector or Guice.createChildInjector).

The best fix is to replace the deprecated call with a direct use of Guice's createInjector method, which receives a list of modules (in this case, the module lambda currently provided in the list argument). If the intent here is to use more than one module, just pass them as arguments; if only one, just use that module directly.

Changes required:

  • Edit line(s) 109-122 to directly call Guice.createInjector, in place of Initialization.makeInjectorWithModules.
  • Add an import for com.google.inject.Guice if it isn't already imported.

Only the code snippet shown in the question may be changed, so only changes in this file and import area are permitted.

Suggested changeset 1
server/src/test/java/org/apache/druid/server/metrics/GroupByStatsMonitorTest.java

Autofix patch

Autofix patch
Run the following command in your local git repository to apply this patch
cat << 'EOF' | git apply
diff --git a/server/src/test/java/org/apache/druid/server/metrics/GroupByStatsMonitorTest.java b/server/src/test/java/org/apache/druid/server/metrics/GroupByStatsMonitorTest.java
--- a/server/src/test/java/org/apache/druid/server/metrics/GroupByStatsMonitorTest.java
+++ b/server/src/test/java/org/apache/druid/server/metrics/GroupByStatsMonitorTest.java
@@ -22,6 +22,7 @@
 import com.google.inject.Injector;
 import com.google.inject.Key;
 import com.google.inject.name.Names;
+import com.google.inject.Guice;
 import org.apache.druid.collections.BlockingPool;
 import org.apache.druid.collections.DefaultBlockingPool;
 import org.apache.druid.guice.GuiceInjectors;
@@ -106,9 +107,8 @@
   {
     final String dataSource = "fooDs";
     final String taskId = "taskId1";
-    final Injector injector = Initialization.makeInjectorWithModules(
-        GuiceInjectors.makeStartupInjector(),
-        List.of(binder -> {
+    final Injector injector = Guice.createInjector(
+        binder -> {
           JsonConfigProvider.bindInstance(
               binder,
               Key.get(DruidNode.class, Self.class),
@@ -118,7 +118,7 @@
                 .toInstance(dataSource);
           binder.bind(Key.get(String.class, Names.named(DataSourceTaskIdHolder.TASK_ID_BINDING)))
                 .toInstance(taskId);
-        })
+        }
     );
     final DataSourceTaskIdHolder dsTaskIdHolder = new DataSourceTaskIdHolder();
     injector.injectMembers(dsTaskIdHolder);
EOF
@@ -22,6 +22,7 @@
import com.google.inject.Injector;
import com.google.inject.Key;
import com.google.inject.name.Names;
import com.google.inject.Guice;
import org.apache.druid.collections.BlockingPool;
import org.apache.druid.collections.DefaultBlockingPool;
import org.apache.druid.guice.GuiceInjectors;
@@ -106,9 +107,8 @@
{
final String dataSource = "fooDs";
final String taskId = "taskId1";
final Injector injector = Initialization.makeInjectorWithModules(
GuiceInjectors.makeStartupInjector(),
List.of(binder -> {
final Injector injector = Guice.createInjector(
binder -> {
JsonConfigProvider.bindInstance(
binder,
Key.get(DruidNode.class, Self.class),
@@ -118,7 +118,7 @@
.toInstance(dataSource);
binder.bind(Key.get(String.class, Names.named(DataSourceTaskIdHolder.TASK_ID_BINDING)))
.toInstance(taskId);
})
}
);
final DataSourceTaskIdHolder dsTaskIdHolder = new DataSourceTaskIdHolder();
injector.injectMembers(dsTaskIdHolder);
Copilot is powered by AI and may make mistakes. Always verify output.
final DataSourceTaskIdHolder dsTaskIdHolder = new DataSourceTaskIdHolder();
injector.injectMembers(dsTaskIdHolder);
final GroupByStatsMonitor monitor =
new GroupByStatsMonitor(groupByStatsProvider, mergeBufferPool, dsTaskIdHolder);
final StubServiceEmitter emitter = new StubServiceEmitter("service", "host");
monitor.doMonitor(emitter);
emitter.flush();
// Trigger metric emission
monitor.doMonitor(emitter);

final Map<String, Object> dimFilters = Map.of(
"taskId", List.of(taskId), "dataSource", List.of(dataSource), "id", List.of(taskId)
);
Assert.assertEquals(7, emitter.getNumEmittedEvents());
emitter.verifyValue("mergeBuffer/pendingRequests", dimFilters, 0L);
emitter.verifyValue("mergeBuffer/used", dimFilters, 0L);
emitter.verifyValue("mergeBuffer/queries", dimFilters, 1L);
emitter.verifyValue("mergeBuffer/acquisitionTimeNs", dimFilters, 100L);
emitter.verifyValue("groupBy/spilledQueries", dimFilters, 2L);
emitter.verifyValue("groupBy/spilledBytes", dimFilters, 200L);
emitter.verifyValue("groupBy/mergeDictionarySize", dimFilters, 300L);
}


@Test
public void testMonitoringMergeBuffer_acquiredCount()
throws ExecutionException, InterruptedException, TimeoutException
Expand All @@ -101,7 +153,7 @@
}).get(20, TimeUnit.SECONDS);

final GroupByStatsMonitor monitor =
new GroupByStatsMonitor(groupByStatsProvider, mergeBufferPool);
new GroupByStatsMonitor(groupByStatsProvider, mergeBufferPool, new DataSourceTaskIdHolder());
final StubServiceEmitter emitter = new StubServiceEmitter("DummyService", "DummyHost");
boolean ret = monitor.doMonitor(emitter);
Assert.assertTrue(ret);
Expand Down Expand Up @@ -130,7 +182,7 @@
}

final GroupByStatsMonitor monitor =
new GroupByStatsMonitor(groupByStatsProvider, mergeBufferPool);
new GroupByStatsMonitor(groupByStatsProvider, mergeBufferPool, new DataSourceTaskIdHolder());
final StubServiceEmitter emitter = new StubServiceEmitter("DummyService", "DummyHost");
boolean ret = monitor.doMonitor(emitter);
Assert.assertTrue(ret);
Expand Down