From a970f6c2de7e94ee3140995fe2b154eae7371927 Mon Sep 17 00:00:00 2001 From: Lucas Capistrant Date: Tue, 21 Jul 2026 09:38:32 -0500 Subject: [PATCH] Discover Dart enabled historicals via a new DartWorkerService instead of inferring workers solely off service == historical --- .../controller/DartControllerContext.java | 30 ++++- .../DartControllerContextFactoryImpl.java | 13 ++- .../dart/controller/DartMessageRelays.java | 10 +- .../msq/dart/guice/DartWorkerModule.java | 17 +++ .../msq/dart/worker/DartWorkerService.java | 50 +++++++++ .../druid/msq/guice/MSQIndexingModule.java | 5 + .../controller/DartControllerContextTest.java | 106 +++++++++++++++--- .../dart/worker/DartWorkerServiceTest.java | 65 +++++++++++ .../TestDartControllerContextFactoryImpl.java | 74 +++++++++++- 9 files changed, 350 insertions(+), 20 deletions(-) create mode 100644 multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerService.java create mode 100644 multi-stage-query/src/test/java/org/apache/druid/msq/dart/worker/DartWorkerServiceTest.java diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java index ec748a74b7de..cd8faecf09be 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContext.java @@ -22,6 +22,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.inject.Injector; import org.apache.druid.client.TimelineServerView; +import org.apache.druid.discovery.DruidNodeDiscovery; import org.apache.druid.error.DruidException; import org.apache.druid.indexing.common.TaskLockType; import org.apache.druid.indexing.common.actions.TaskActionClient; @@ -52,6 +53,7 @@ import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.Set; import java.util.stream.Collectors; /** @@ -91,6 +93,7 @@ public class DartControllerContext implements ControllerContext private final List inputSpecSlicerProviders; private final ServiceEmitter emitter; private final QueryContext context; + private final DruidNodeDiscovery dartWorkerDiscovery; public DartControllerContext( final Injector injector, @@ -101,7 +104,8 @@ public DartControllerContext( final TimelineServerView serverView, final List inputSpecSlicerProviders, final ServiceEmitter emitter, - final QueryContext context + final QueryContext context, + final DruidNodeDiscovery dartWorkerDiscovery ) { this.injector = injector; @@ -113,6 +117,7 @@ public DartControllerContext( this.inputSpecSlicerProviders = inputSpecSlicerProviders; this.emitter = emitter; this.context = context; + this.dartWorkerDiscovery = dartWorkerDiscovery; } @Override @@ -126,17 +131,38 @@ public ControllerQueryKernelConfig queryKernelConfig(final MSQSpec querySpec) { final List servers = serverView.getDruidServerMetadatas(); + final Set dartWorkerHosts = + dartWorkerDiscovery.getAllNodes() + .stream() + .map(node -> node.getDruidNode().getHostAndPortToUse()) + .collect(Collectors.toSet()); + // Lock in the list of workers when creating the kernel config. There is a race here: the serverView itself is // allowed to float. If a segment moves to a new server that isn't part of our list after the WorkerManager is // created, we won't be able to find a valid server for certain segments. This isn't expected to be a problem, // since the serverView is referenced shortly after the worker list is created. final List workerIds = new ArrayList<>(servers.size()); for (final DruidServerMetadata server : servers) { - if (server.getType() == ServerType.HISTORICAL) { + if (server.getType() == ServerType.HISTORICAL && dartWorkerHosts.contains(server.getHost())) { workerIds.add(WorkerId.fromDruidServerMetadata(server, queryId()).toString()); } } + // Fail fast rather than running a query with no workers. + if (workerIds.isEmpty()) { + final boolean anyHistoricals = servers.stream().anyMatch(s -> s.getType() == ServerType.HISTORICAL); + throw DruidException.forPersona(DruidException.Persona.OPERATOR) + .ofCategory(DruidException.Category.RUNTIME_FAILURE) + .build( + anyHistoricals + ? "No Dart workers are available: Historicals are present but none advertise a Dart " + + "worker. Set druid.msq.dart.enabled=true on the Historicals that should run Dart " + + "queries." + : "No Dart workers are available: no Historicals are currently available to run the " + + "query." + ); + } + // Shuffle workerIds, so we don't bias towards specific servers when running multiple queries concurrently. For any // given query, lower-numbered workers tend to do more work, because the controller prefers using lower-numbered // workers when maxWorkerCount for a stage is less than the total number of workers. diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java index 260d6156c8c4..a5839d20b587 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerContextFactoryImpl.java @@ -23,6 +23,9 @@ import com.google.inject.Inject; import com.google.inject.Injector; import org.apache.druid.client.TimelineServerView; +import org.apache.druid.discovery.DruidNodeDiscovery; +import org.apache.druid.discovery.DruidNodeDiscoveryProvider; +import org.apache.druid.discovery.NodeRole; import org.apache.druid.guice.annotations.EscalatedGlobal; import org.apache.druid.guice.annotations.Json; import org.apache.druid.guice.annotations.Self; @@ -30,6 +33,7 @@ import org.apache.druid.java.util.emitter.service.ServiceEmitter; import org.apache.druid.msq.dart.Dart; import org.apache.druid.msq.dart.worker.DartWorkerClientImpl; +import org.apache.druid.msq.dart.worker.DartWorkerService; import org.apache.druid.msq.exec.ControllerContext; import org.apache.druid.msq.exec.MemoryIntrospector; import org.apache.druid.msq.input.InputSpecSlicerProvider; @@ -52,6 +56,7 @@ public class DartControllerContextFactoryImpl implements DartControllerContextFa protected final MemoryIntrospector memoryIntrospector; protected final List inputSpecSlicerProviders; protected final ServiceEmitter emitter; + protected final DruidNodeDiscovery dartWorkerDiscovery; @Inject public DartControllerContextFactoryImpl( @@ -63,7 +68,8 @@ public DartControllerContextFactoryImpl( final MemoryIntrospector memoryIntrospector, final TimelineServerView serverView, @Dart final Set inputSpecSlicerProviders, - final ServiceEmitter emitter + final ServiceEmitter emitter, + final DruidNodeDiscoveryProvider discoveryProvider ) { this.injector = injector; @@ -75,6 +81,8 @@ public DartControllerContextFactoryImpl( this.memoryIntrospector = memoryIntrospector; this.inputSpecSlicerProviders = List.copyOf(inputSpecSlicerProviders); this.emitter = emitter; + this.dartWorkerDiscovery = + discoveryProvider.getForServiceAndRoles(DartWorkerService.NAME, Set.of(NodeRole.HISTORICAL)); } @Override @@ -90,7 +98,8 @@ public ControllerContext newContext(final QueryContext context) serverView, inputSpecSlicerProviders, emitter, - context + context, + dartWorkerDiscovery ); } } diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartMessageRelays.java b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartMessageRelays.java index 23accd35ecbe..0f1ec4b5eaf5 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartMessageRelays.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartMessageRelays.java @@ -24,6 +24,9 @@ import org.apache.druid.messages.client.MessageRelayFactory; import org.apache.druid.messages.client.MessageRelays; import org.apache.druid.msq.dart.controller.messages.ControllerMessage; +import org.apache.druid.msq.dart.worker.DartWorkerService; + +import java.util.Set; /** * Specialized {@link MessageRelays} for Dart controllers. @@ -35,6 +38,11 @@ public DartMessageRelays( final MessageRelayFactory messageRelayFactory ) { - super(() -> discoveryProvider.getForNodeRole(NodeRole.HISTORICAL), messageRelayFactory); + // Only relay with Historicals that run a Dart worker (advertise DartWorkerService); Dart-disabled Historicals + // expose no outbox. This is the same discovery handle the controller uses to enroll workers. + super( + () -> discoveryProvider.getForServiceAndRoles(DartWorkerService.NAME, Set.of(NodeRole.HISTORICAL)), + messageRelayFactory + ); } } diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerModule.java b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerModule.java index 406c2acac840..dd220c4d7e28 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerModule.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerModule.java @@ -28,7 +28,10 @@ import com.google.inject.Key; import com.google.inject.Module; import com.google.inject.Provides; +import com.google.inject.multibindings.ProvidesIntoSet; +import com.google.inject.name.Named; import org.apache.druid.discovery.DruidNodeDiscoveryProvider; +import org.apache.druid.discovery.DruidService; import org.apache.druid.discovery.NodeRole; import org.apache.druid.guice.Jerseys; import org.apache.druid.guice.JsonConfigProvider; @@ -57,6 +60,7 @@ import org.apache.druid.msq.dart.worker.DartWorkerContextFactory; import org.apache.druid.msq.dart.worker.DartWorkerContextFactoryImpl; import org.apache.druid.msq.dart.worker.DartWorkerRunner; +import org.apache.druid.msq.dart.worker.DartWorkerService; import org.apache.druid.msq.dart.worker.http.DartWorkerResource; import org.apache.druid.msq.exec.MemoryIntrospector; import org.apache.druid.msq.guice.MSQBinders; @@ -118,6 +122,19 @@ public void configure(Binder binder) .in(LazySingleton.class); } + /** + * Advertise {@link DartWorkerService} in node discovery. Contributed from {@link ActualModule}, which is only + * installed when Dart is enabled, so advertisement tracks actually running a Dart worker. Merges into the same + * {@code @Named("historical")} service set announced by + * {@link org.apache.druid.guice.HistoricalServiceModule}. + */ + @ProvidesIntoSet + @Named(NodeRole.HISTORICAL_JSON_NAME) + public Class getDartWorkerService() + { + return DartWorkerService.class; + } + @Provides @ManageLifecycle public DartWorkerRunner createWorkerRunner( diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerService.java b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerService.java new file mode 100644 index 000000000000..910ce23ef854 --- /dev/null +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerService.java @@ -0,0 +1,50 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.druid.msq.dart.worker; + +import org.apache.druid.discovery.DruidService; + +/** + * No-payload {@link DruidService} advertised in node discovery by a Historical whose Dart worker runtime is + * installed (i.e. {@code druid.msq.dart.enabled=true}). The Dart controller enrolls workers only when they are + * Historicals that advertise this capability. + */ +public class DartWorkerService extends DruidService +{ + public static final String NAME = "dartWorkerService"; + + @Override + public String getName() + { + return NAME; + } + + @Override + public boolean equals(Object o) + { + return o != null && getClass() == o.getClass(); + } + + @Override + public int hashCode() + { + return DartWorkerService.class.hashCode(); + } +} diff --git a/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQIndexingModule.java b/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQIndexingModule.java index 8e472e6e7d85..a18b7cddfcb3 100644 --- a/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQIndexingModule.java +++ b/multi-stage-query/src/main/java/org/apache/druid/msq/guice/MSQIndexingModule.java @@ -40,6 +40,7 @@ import org.apache.druid.msq.counters.StorageCounters; import org.apache.druid.msq.counters.SuperSorterProgressTrackerCounter; import org.apache.druid.msq.counters.WarningCounters; +import org.apache.druid.msq.dart.worker.DartWorkerService; import org.apache.druid.msq.indexing.IndexerControllerContextFactory; import org.apache.druid.msq.indexing.IndexerSegmentsInputSliceReaderProvider; import org.apache.druid.msq.indexing.IndexerTableInputSpecSlicerProvider; @@ -229,6 +230,10 @@ public List getJacksonModules() module.registerSubtypes(new NamedType(MSQCompactionRunner.class, MSQCompactionRunner.TYPE)); + // Registered here rather than in DartWorkerModule so that every process that discovers Historicals can parse + // their DartWorkerService announcement. + module.registerSubtypes(new NamedType(DartWorkerService.class, DartWorkerService.NAME)); + FAULT_CLASSES.forEach(module::registerSubtypes); module.addSerializer(new CounterSnapshotsSerializer()); return Collections.singletonList(module); diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/DartControllerContextTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/DartControllerContextTest.java index 1ab1b55024c5..d21e181bcd71 100644 --- a/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/DartControllerContextTest.java +++ b/multi-stage-query/src/test/java/org/apache/druid/msq/dart/controller/DartControllerContextTest.java @@ -22,12 +22,17 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import org.apache.druid.client.BrokerServerView; +import org.apache.druid.discovery.DiscoveryDruidNode; +import org.apache.druid.discovery.DruidNodeDiscovery; +import org.apache.druid.error.DruidException; +import org.apache.druid.msq.dart.worker.DartWorkerService; import org.apache.druid.msq.dart.worker.WorkerId; import org.apache.druid.msq.exec.MemoryIntrospector; import org.apache.druid.msq.exec.MemoryIntrospectorImpl; import org.apache.druid.msq.indexing.LegacyMSQSpec; import org.apache.druid.msq.indexing.destination.TaskReportMSQDestination; import org.apache.druid.msq.kernel.controller.ControllerQueryKernelConfig; +import org.apache.druid.msq.test.TestDartControllerContextFactoryImpl; import org.apache.druid.msq.util.MultiStageQueryContext; import org.apache.druid.query.Query; import org.apache.druid.query.QueryContext; @@ -43,6 +48,7 @@ import org.mockito.Mockito; import org.mockito.MockitoAnnotations; +import java.util.Arrays; import java.util.List; import java.util.stream.Collectors; @@ -104,9 +110,68 @@ public void tearDown() throws Exception } @Test - public void test_queryKernelConfig() + public void test_queryKernelConfig_allHistoricalsAdvertiseDartWorker() { - final DartControllerContext controllerContext = new DartControllerContext( + // Both HISTORICALs advertise a Dart worker: enroll both, matching pre-capability behavior. + final ControllerQueryKernelConfig queryKernelConfig = + makeControllerContext(discoveryOf(SERVERS.get(0), SERVERS.get(1))).queryKernelConfig(querySpec); + + assertCommonKernelConfig(queryKernelConfig); + Assertions.assertEquals( + ImmutableList.of( + WorkerId.fromDruidServerMetadata(SERVERS.get(0), QUERY_ID).toString(), + WorkerId.fromDruidServerMetadata(SERVERS.get(1), QUERY_ID).toString() + ), + sortedWorkerIds(queryKernelConfig) + ); + } + + @Test + public void test_queryKernelConfig_onlySomeHistoricalsAdvertiseDartWorker() + { + // Only the first HISTORICAL advertises a Dart worker: the Dart-disabled one is excluded. + final ControllerQueryKernelConfig queryKernelConfig = + makeControllerContext(discoveryOf(SERVERS.get(0))).queryKernelConfig(querySpec); + + assertCommonKernelConfig(queryKernelConfig); + Assertions.assertEquals( + ImmutableList.of(WorkerId.fromDruidServerMetadata(SERVERS.get(0), QUERY_ID).toString()), + sortedWorkerIds(queryKernelConfig) + ); + } + + @Test + public void test_queryKernelConfig_noHistoricalsAdvertiseDartWorker_failsFast() + { + // Historicals are present in the server view but none advertise a Dart worker so we must fail fast + final DruidException e = Assertions.assertThrows( + DruidException.class, + () -> makeControllerContext(discoveryOf()).queryKernelConfig(querySpec) + ); + + Assertions.assertEquals(DruidException.Persona.OPERATOR, e.getTargetPersona()); + Assertions.assertTrue(e.getMessage().contains("No Dart workers are available"), e.getMessage()); + Assertions.assertTrue(e.getMessage().contains("druid.msq.dart.enabled"), e.getMessage()); + } + + @Test + public void test_queryKernelConfig_noHistoricalsAtAll_failsFast() + { + // No historical servers at all: still fail fast, but with a descriptive message about the cause + Mockito.when(serverView.getDruidServerMetadatas()) + .thenReturn(ImmutableList.of(SERVERS.get(2))); // realtime only + + final DruidException e = Assertions.assertThrows( + DruidException.class, + () -> makeControllerContext(discoveryOf()).queryKernelConfig(querySpec) + ); + + Assertions.assertTrue(e.getMessage().contains("no Historicals are currently available"), e.getMessage()); + } + + private DartControllerContext makeControllerContext(final DruidNodeDiscovery dartWorkerDiscovery) + { + return new DartControllerContext( null, null, SELF_NODE, @@ -115,24 +180,39 @@ public void test_queryKernelConfig() serverView, List.of(), null, - queryContext + queryContext, + dartWorkerDiscovery ); - final ControllerQueryKernelConfig queryKernelConfig = controllerContext.queryKernelConfig(querySpec); + } + private static void assertCommonKernelConfig(final ControllerQueryKernelConfig queryKernelConfig) + { Assertions.assertFalse(queryKernelConfig.isFaultTolerant()); Assertions.assertFalse(queryKernelConfig.isDurableStorage()); Assertions.assertEquals(3, queryKernelConfig.getMaxConcurrentStages()); Assertions.assertEquals(TaskReportMSQDestination.instance(), queryKernelConfig.getDestination()); Assertions.assertTrue(queryKernelConfig.isPipeline()); + } - // Check workerIds after sorting, because they've been shuffled. - Assertions.assertEquals( - ImmutableList.of( - // Only the HISTORICAL servers - WorkerId.fromDruidServerMetadata(SERVERS.get(0), QUERY_ID).toString(), - WorkerId.fromDruidServerMetadata(SERVERS.get(1), QUERY_ID).toString() - ), - queryKernelConfig.getWorkerIds().stream().sorted().collect(Collectors.toList()) - ); + /** + * The workerIds are shuffled by {@link DartControllerContext#queryKernelConfig}, so sort before comparing. + */ + private static List sortedWorkerIds(final ControllerQueryKernelConfig queryKernelConfig) + { + return queryKernelConfig.getWorkerIds().stream().sorted().collect(Collectors.toList()); + } + + /** + * A {@link DruidNodeDiscovery} that reports the given servers as Historicals advertising a {@link DartWorkerService}. + */ + private static DruidNodeDiscovery discoveryOf(final DruidServerMetadata... servers) + { + final List nodes = + Arrays.stream(servers) + .map(TestDartControllerContextFactoryImpl::historicalDartWorkerNode) + .collect(Collectors.toList()); + final DruidNodeDiscovery discovery = Mockito.mock(DruidNodeDiscovery.class); + Mockito.when(discovery.getAllNodes()).thenReturn(nodes); + return discovery; } } diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/dart/worker/DartWorkerServiceTest.java b/multi-stage-query/src/test/java/org/apache/druid/msq/dart/worker/DartWorkerServiceTest.java new file mode 100644 index 000000000000..88f5bf2873cd --- /dev/null +++ b/multi-stage-query/src/test/java/org/apache/druid/msq/dart/worker/DartWorkerServiceTest.java @@ -0,0 +1,65 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.druid.msq.dart.worker; + +import com.fasterxml.jackson.databind.InjectableValues; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.ImmutableMap; +import org.apache.druid.discovery.DiscoveryDruidNode; +import org.apache.druid.discovery.DruidService; +import org.apache.druid.discovery.NodeRole; +import org.apache.druid.guice.ServerModule; +import org.apache.druid.jackson.DefaultObjectMapper; +import org.apache.druid.msq.guice.MSQIndexingModule; +import org.apache.druid.server.DruidNode; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +/** + * Verifies a Historical's {@link DartWorkerService} announcement round-trips through the {@link DruidService} + * polymorphic types, so any process with {@link MSQIndexingModule} installed can read it and discover the node. + */ +public class DartWorkerServiceTest +{ + private static final DruidNode HISTORICAL_NODE = + new DruidNode("no", "localhost", false, 8083, -1, true, false); + + @Test + public void test_serde_withinDiscoveryDruidNode() throws Exception + { + final ObjectMapper mapper = createMapper(); + final DiscoveryDruidNode node = new DiscoveryDruidNode( + HISTORICAL_NODE, + NodeRole.HISTORICAL, + ImmutableMap.of(DartWorkerService.NAME, new DartWorkerService()) + ); + final String json = mapper.writeValueAsString(node); + Assertions.assertEquals(node, mapper.readValue(json, DiscoveryDruidNode.class)); + } + + private static ObjectMapper createMapper() + { + final ObjectMapper mapper = new DefaultObjectMapper(); + mapper.registerModules(new ServerModule().getJacksonModules()); + mapper.registerModules(new MSQIndexingModule().getJacksonModules()); + mapper.setInjectableValues(new InjectableValues.Std().addValue(ObjectMapper.class, mapper)); + return mapper; + } +} diff --git a/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java b/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java index d7f48f2dc419..3dd81e204d4a 100644 --- a/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java +++ b/multi-stage-query/src/test/java/org/apache/druid/msq/test/TestDartControllerContextFactoryImpl.java @@ -20,12 +20,18 @@ package org.apache.druid.msq.test; import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.ImmutableMap; +import com.google.common.net.HostAndPort; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.ListeningExecutorService; import com.google.common.util.concurrent.MoreExecutors; import com.google.inject.Inject; import com.google.inject.Injector; import org.apache.druid.client.TimelineServerView; +import org.apache.druid.discovery.DiscoveryDruidNode; +import org.apache.druid.discovery.DruidNodeDiscovery; +import org.apache.druid.discovery.DruidNodeDiscoveryProvider; +import org.apache.druid.discovery.NodeRole; import org.apache.druid.guice.annotations.EscalatedGlobal; import org.apache.druid.guice.annotations.Json; import org.apache.druid.guice.annotations.Self; @@ -38,6 +44,7 @@ import org.apache.druid.msq.dart.controller.DartControllerContext; import org.apache.druid.msq.dart.controller.DartControllerContextFactoryImpl; import org.apache.druid.msq.dart.worker.DartWorkerClient; +import org.apache.druid.msq.dart.worker.DartWorkerService; import org.apache.druid.msq.exec.Controller; import org.apache.druid.msq.exec.ControllerContext; import org.apache.druid.msq.exec.MSQMetricEventBuilder; @@ -52,10 +59,15 @@ import org.apache.druid.query.QueryContext; import org.apache.druid.rpc.ServiceClientFactory; import org.apache.druid.server.DruidNode; +import org.apache.druid.server.coordination.DruidServerMetadata; +import org.apache.druid.server.coordination.ServerType; +import java.util.Collection; +import java.util.List; import java.util.Map; import java.util.Set; import java.util.concurrent.Executors; +import java.util.stream.Collectors; public class TestDartControllerContextFactoryImpl extends DartControllerContextFactoryImpl { @@ -78,6 +90,7 @@ public TestDartControllerContextFactoryImpl( final TimelineServerView serverView, @Dart final Set inputSpecSlicerProviders, final ServiceEmitter emitter, + final DruidNodeDiscoveryProvider discoveryProvider, @Dart Map workerMap ) { @@ -90,7 +103,8 @@ public TestDartControllerContextFactoryImpl( memoryIntrospector, serverView, inputSpecSlicerProviders, - emitter + emitter, + discoveryProvider ); this.workerMap = workerMap; } @@ -107,7 +121,8 @@ public ControllerContext newContext(QueryContext context) serverView, inputSpecSlicerProviders, emitter, - context + context, + advertiseAllHistoricals(serverView) ) { @Override @@ -131,6 +146,61 @@ public boolean isDebug() }; } + /** + * Builds a {@link DruidNodeDiscovery} that reports every {@link ServerType#HISTORICAL} in the server view as a + * Dart worker (advertising {@link DartWorkerService}). {@link DartControllerContext#queryKernelConfig} reads + * only {@link DruidNodeDiscovery#getAllNodes()}, so a no-op listener registration is sufficient. + */ + private static DruidNodeDiscovery advertiseAllHistoricals(final TimelineServerView serverView) + { + final List nodes = + serverView.getDruidServerMetadatas() + .stream() + .filter(server -> server.getType() == ServerType.HISTORICAL) + .map(TestDartControllerContextFactoryImpl::historicalDartWorkerNode) + .collect(Collectors.toList()); + + return new DruidNodeDiscovery() + { + @Override + public Collection getAllNodes() + { + return nodes; + } + + @Override + public void registerListener(Listener listener) + { + // Not used by queryKernelConfig. + } + }; + } + + /** + * A Historical {@link DiscoveryDruidNode} advertising {@link DartWorkerService}, whose + * {@link DruidNode#getHostAndPortToUse()} matches {@link DruidServerMetadata#getHost()}. + */ + public static DiscoveryDruidNode historicalDartWorkerNode(final DruidServerMetadata server) + { + // Build a DruidNode whose getHostAndPortToUse() equals server.getHost() (the key queryKernelConfig matches on). + final boolean tls = server.getHostAndTlsPort() != null; + final HostAndPort hostAndPort = HostAndPort.fromString(server.getHost()); + final DruidNode druidNode = new DruidNode( + "no", + hostAndPort.getHost(), + false, + tls ? -1 : hostAndPort.getPort(), + tls ? hostAndPort.getPort() : -1, + !tls, + tls + ); + return new DiscoveryDruidNode( + druidNode, + NodeRole.HISTORICAL, + ImmutableMap.of(DartWorkerService.NAME, new DartWorkerService()) + ); + } + public class DartTestWorkerClient extends MSQTestWorkerClient implements DartWorkerClient {