kfaraz commented on code in PR #20825: URL: https://github.com/apache/druid/pull/20825#discussion_r4452053546
########## multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartWorkerSyncer.java: ########## @@ -0,0 +1,271 @@ +/* + * 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.controller; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.base.Supplier; +import com.google.common.base.Suppliers; +import com.google.common.util.concurrent.FutureCallback; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.errorprone.annotations.concurrent.GuardedBy; +import io.netty.handler.codec.http.HttpMethod; +import org.apache.druid.discovery.DiscoveryDruidNode; +import org.apache.druid.discovery.DruidNodeDiscovery; +import org.apache.druid.error.DruidException; +import org.apache.druid.java.util.common.concurrent.Execs; +import org.apache.druid.java.util.common.concurrent.ScheduledExecutors; +import org.apache.druid.java.util.common.lifecycle.LifecycleStart; +import org.apache.druid.java.util.common.lifecycle.LifecycleStop; +import org.apache.druid.java.util.common.logger.Logger; +import org.apache.druid.java.util.http.client.response.BytesFullResponseHandler; +import org.apache.druid.java.util.http.client.response.BytesFullResponseHolder; +import org.apache.druid.msq.dart.guice.DartControllerConfig; +import org.apache.druid.msq.dart.worker.DartWorkerClientImpl; +import org.apache.druid.msq.dart.worker.DartWorkerRunner; +import org.apache.druid.msq.dart.worker.WorkerId; +import org.apache.druid.msq.dart.worker.http.DartWorkerResource; +import org.apache.druid.msq.dart.worker.http.SyncWorkersRequest; +import org.apache.druid.msq.dart.worker.http.SyncWorkersResponse; +import org.apache.druid.msq.exec.ControllerHolder; +import org.apache.druid.rpc.FixedServiceLocator; +import org.apache.druid.rpc.RequestBuilder; +import org.apache.druid.rpc.ServiceClient; +import org.apache.druid.rpc.ServiceClientFactory; +import org.apache.druid.rpc.ServiceLocation; +import org.apache.druid.rpc.ServiceLocator; +import org.apache.druid.rpc.StandardRetryPolicy; +import org.apache.druid.server.DruidNode; +import org.apache.druid.utils.CloseableUtils; +import org.joda.time.Duration; + +import javax.annotation.Nullable; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Iterator; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ScheduledExecutorService; + +/** + * Sends periodic syncs from the Broker to the servers that run Dart workers. Each sync sends a + * {@link SyncWorkersRequest} with the list of queries the Broker believes are running on that worker. The worker + * replies with a {@link SyncWorkersResponse} indicating what queries the worker is currently running. This allows + * discrepancies to be resolved. + * + * @see DartWorkerRunner#syncController the worker side + */ +public class DartWorkerSyncer +{ + private static final Logger log = new Logger(DartWorkerSyncer.class); + + /** + * Timeout for each sync request. + */ + static final long SYNC_TIMEOUT_MILLIS = 10_000; + + private final DartControllerRegistry controllerRegistry; + private final Supplier<DruidNodeDiscovery> workerDiscovery; + private final ServiceClientFactory clientFactory; + private final ObjectMapper jsonMapper; + private final String selfHost; + private final Duration syncPeriod; + + /** + * Worker host:port -> state of syncs to that server. + */ + @GuardedBy("servers") + private final Map<String, ServerHolder> servers = new HashMap<>(); + + @Nullable + private ScheduledExecutorService exec; + + public DartWorkerSyncer( + final DartControllerRegistry controllerRegistry, + final Supplier<DruidNodeDiscovery> workerDiscovery, + final ServiceClientFactory clientFactory, + final ObjectMapper jsonMapper, + final String selfHost, + final DartControllerConfig config + ) + { + this.controllerRegistry = controllerRegistry; + this.workerDiscovery = Suppliers.memoize(workerDiscovery); + this.clientFactory = clientFactory; + this.jsonMapper = jsonMapper; + this.selfHost = selfHost; + this.syncPeriod = config.getSyncPeriod().toStandardDuration(); + if (syncPeriod.getMillis() <= 0) { + throw DruidException.forPersona(DruidException.Persona.OPERATOR) + .ofCategory(DruidException.Category.INVALID_INPUT) + .build("druid.msq.dart.controller.syncPeriod[%s] must be positive", config.getSyncPeriod()); + } + } + + @LifecycleStart + public void start() + { + exec = Execs.scheduledSingleThreaded("dart-worker-syncer-%s"); + ScheduledExecutors.scheduleWithFixedDelay(exec, syncPeriod, syncPeriod, this::sync); + } + + @LifecycleStop + public void stop() + { + if (exec != null) { + exec.shutdownNow(); Review Comment: Nit: Maybe pass in a `ScheduledExecutorFactory` in the constructor for easier lifecycle management of the executor. It also allows for easier testing (if the `sync` method were not exposed). ########## multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerConfig.java: ########## @@ -59,6 +61,18 @@ public class DartWorkerConfig @Nullable private Integer segmentLoadAheadCount = null; + /** + * How long a controller can go without contacting this worker before {@link DartWorkerRunner} considers it gone. + */ + @JsonProperty("controllerTimeout") + private Period controllerTimeout = Period.minutes(3); + + /** + * How long a query can be missing from a controller's syncs before {@link DartWorkerRunner} cancels it. Review Comment: Nit: Might be nicer to have the javadocs linked on the getters as well, so that they can be easily accessed at the call sites. ########## multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/http/DartWorkerResource.java: ########## @@ -120,19 +122,31 @@ public Response httpPostWorkOrder( ) { MSQResourceUtils.authorizeAdminRequest(permissionMapper, authorizerMapper, req); - final String controllerHost = req.getHeader(HEADER_CONTROLLER_HOST); - if (controllerHost == null) { - throw DruidException.forPersona(DruidException.Persona.DEVELOPER) - .ofCategory(DruidException.Category.INVALID_INPUT) - .build("Missing controllerId[%s]", HEADER_CONTROLLER_HOST); - } - - workerRunner.startWorker(queryId, controllerHost, workOrder.getWorkerContext()) + workerRunner.startWorker(queryId, getControllerHost(req), workOrder.getWorkerContext()) .postWorkOrder(workOrder); return Response.status(Response.Status.ACCEPTED).build(); } + /** + * Process a {@link SyncWorkersRequest} from {@link DartWorkerRunner#syncController}. + */ + @POST + @Consumes({MediaType.APPLICATION_JSON, SmileMediaTypes.APPLICATION_JACKSON_SMILE}) + @Produces(MediaType.APPLICATION_JSON) + @Path("/sync") + public SyncWorkersResponse httpPostSync( + final SyncWorkersRequest request, + @Context final HttpServletRequest req + ) + { + MSQResourceUtils.authorizeAdminRequest(permissionMapper, authorizerMapper, req); + return workerRunner.syncController( + getControllerHost(req), + request != null ? request : new SyncWorkersRequest(Set.of()) Review Comment: Why would the request payload be null? Is this for tests only? I feel that under normal conditions, a null payload here should rather be treated as a bad request. ########## multi-stage-query/src/main/java/org/apache/druid/msq/dart/guice/DartWorkerConfig.java: ########## @@ -59,6 +61,18 @@ public class DartWorkerConfig @Nullable private Integer segmentLoadAheadCount = null; + /** + * How long a controller can go without contacting this worker before {@link DartWorkerRunner} considers it gone. + */ + @JsonProperty("controllerTimeout") + private Period controllerTimeout = Period.minutes(3); + + /** + * How long a query can be missing from a controller's syncs before {@link DartWorkerRunner} cancels it. + */ + @JsonProperty("syncGracePeriod") + private Period syncGracePeriod = Period.seconds(30); Review Comment: Maybe rename to `querySyncTimeout` instead? ########## multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerRunner.java: ########## @@ -47,31 +52,51 @@ import org.apache.druid.server.security.AuthorizerMapper; import org.apache.druid.utils.CloseableUtils; import org.joda.time.DateTime; +import org.joda.time.Duration; import javax.annotation.Nullable; import java.io.Closeable; import java.io.File; import java.io.IOException; +import java.time.Clock; import java.util.ArrayList; import java.util.Collection; import java.util.HashMap; import java.util.HashSet; +import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; import java.util.function.Predicate; -import java.util.stream.Collectors; @ManageLifecycle public class DartWorkerRunner { private static final Logger log = new Logger(DartWorkerRunner.class); /** - * Set of active controllers. Ignore requests from others. + * Time between {@link #reapControllersIfNeeded()} calls. + */ + private static final long REAPER_PERIOD_MILLIS = 10_000; + + /** + * How long {@link #stop()} waits for {@link #controllerReaperExec} to terminate. + */ + private static final long REAPER_SHUTDOWN_TIMEOUT_MILLIS = 30_000; + + /** + * Minimum value for {@link DartWorkerConfig#getControllerTimeout()}. Any shorter and we might end up reaping + * controllers in between heartbeats. + */ + private static final long MIN_CONTROLLER_TIMEOUT_MILLIS = 60_000; Review Comment: Shouldn't this be used and validated in the `DartWorkerConfig` itself? ########## multi-stage-query/src/main/java/org/apache/druid/msq/dart/controller/DartControllerRegistry.java: ########## @@ -123,6 +132,18 @@ public void stop() if (cleanupExec != null) { cleanupExec.shutdown(); } + + final List<ControllerHolder> holders = new ArrayList<>(controllerMap.values()); + if (!holders.isEmpty()) { + log.info("Canceling [%d] queries that are still running.", holders.size()); + + for (final ControllerHolder holder : holders) { + holder.cancel(CancellationReason.TASK_SHUTDOWN, null); Review Comment: Shouldn't this be a new reason `BROKER_SHUTDOWN`? ########## multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerRunner.java: ########## @@ -123,11 +180,7 @@ public Worker startWorker( final boolean newHolder; synchronized (this) { - if (!activeControllerHosts.contains(controllerHost)) { - throw DruidException.forPersona(DruidException.Persona.OPERATOR) - .ofCategory(DruidException.Category.RUNTIME_FAILURE) - .build("Received startWorker request for unknown controller[%s]", controllerHost); - } + lastControllerHeartbeat.put(controllerHost, clock.millis()); Review Comment: So, IIUC, we never reject any work anymore. Even if it's a new Broker from which we have never heard before. ########## multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerRunner.java: ########## @@ -180,6 +234,57 @@ public void stopWorker(final String queryId) } } + /** + * Handles a sync from a controller. + * + * @see DartWorkerSyncer the controller side + */ + public SyncWorkersResponse syncController(final String controllerHost, final SyncWorkersRequest request) + { + final Set<String> queryIds = request.queryIds(); + final List<WorkerHolder> holdersToCancel = new ArrayList<>(); + final Set<String> runningQueryIds = new HashSet<>(); + + synchronized (this) { + final long now = clock.millis(); + lastControllerHeartbeat.put(controllerHost, now); + + for (final Map.Entry<String, WorkerHolder> entry : workerMap.entrySet()) { + final WorkerHolder holder = entry.getValue(); + if (!holder.controllerHost.equals(controllerHost)) { + continue; + } + + // Report the query as running even if we cancel it below, since it has not exited yet. + runningQueryIds.add(entry.getKey()); + + if (queryIds.contains(entry.getKey())) { + // Null out missingFromSyncSince, signifying that this query is meant to be active. + holder.missingFromSyncSince = null; + } else { + // Add this query to workersToCancel if it has been missing long enough. + final Long missingFromSyncSince = holder.missingFromSyncSince; + if (missingFromSyncSince == null) { + holder.missingFromSyncSince = now; + } else if (now - missingFromSyncSince >= syncGraceMillis) { + holdersToCancel.add(holder); + } + } + } + } + + for (final WorkerHolder holder : holdersToCancel) { + log.info( + "Canceling query[%s] because controller[%s] no longer lists it as running.", + holder.workerContext.queryId(), + controllerHost + ); + holder.runRef.cancel(); + } + + return new SyncWorkersResponse(runningQueryIds); Review Comment: I wonder if the response should also include the missingQueryIds. The controller probably doesn't have any use for it right now. But it could be useful to see it in the overall Dart query table/report. ########## multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerRunner.java: ########## @@ -85,33 +110,65 @@ public class DartWorkerRunner private final Map<String, WorkerHolder> workerMap = new HashMap<>(); private final DartWorkerContextFactory workerFactory; private final ListeningExecutorService workerExec; - private final DruidNodeDiscoveryProvider discoveryProvider; + private final Outbox<ControllerMessage> outbox; private final ResourcePermissionMapper permissionMapper; private final AuthorizerMapper authorizerMapper; + private final long controllerTimeoutMillis; + + /** + * Cancel workers after they have been missing from {@link #syncController} for this long. A grace period is + * necessary because in-flight syncs may not include a worker that has just started. + */ + private final long syncGraceMillis; + private final File baseTempDir; + private final Clock clock; + private final ScheduledExecutorService controllerReaperExec; public DartWorkerRunner( final DartWorkerContextFactory workerFactory, final ListeningExecutorService workerExec, - final DruidNodeDiscoveryProvider discoveryProvider, + final Outbox<ControllerMessage> outbox, final ResourcePermissionMapper permissionMapper, final AuthorizerMapper authorizerMapper, - final File baseTempDir + final DartWorkerConfig workerConfig, + final File baseTempDir, + final Clock clock, + final ScheduledExecutorService controllerReaperExec ) { this.workerFactory = workerFactory; this.workerExec = workerExec; - this.discoveryProvider = discoveryProvider; + this.outbox = outbox; this.permissionMapper = permissionMapper; this.authorizerMapper = authorizerMapper; + this.controllerTimeoutMillis = workerConfig.getControllerTimeout().toStandardDuration().getMillis(); + if (controllerTimeoutMillis < MIN_CONTROLLER_TIMEOUT_MILLIS) { + throw DruidException.forPersona(DruidException.Persona.OPERATOR) + .ofCategory(DruidException.Category.INVALID_INPUT) + .build( + "druid.msq.dart.worker.controllerTimeout[%s] must be at least [%,d] ms", + workerConfig.getControllerTimeout(), + MIN_CONTROLLER_TIMEOUT_MILLIS + ); + } + this.syncGraceMillis = workerConfig.getSyncGracePeriod().toStandardDuration().getMillis(); + if (syncGraceMillis <= 0) { + throw DruidException.forPersona(DruidException.Persona.OPERATOR) Review Comment: Maybe move the validation to `DartWorkerConfig` itself. ########## multi-stage-query/src/main/java/org/apache/druid/msq/dart/worker/DartWorkerRunner.java: ########## @@ -180,6 +234,57 @@ public void stopWorker(final String queryId) } } + /** + * Handles a sync from a controller. + * + * @see DartWorkerSyncer the controller side + */ + public SyncWorkersResponse syncController(final String controllerHost, final SyncWorkersRequest request) + { + final Set<String> queryIds = request.queryIds(); + final List<WorkerHolder> holdersToCancel = new ArrayList<>(); + final Set<String> runningQueryIds = new HashSet<>(); + + synchronized (this) { + final long now = clock.millis(); + lastControllerHeartbeat.put(controllerHost, now); + + for (final Map.Entry<String, WorkerHolder> entry : workerMap.entrySet()) { + final WorkerHolder holder = entry.getValue(); + if (!holder.controllerHost.equals(controllerHost)) { + continue; + } + + // Report the query as running even if we cancel it below, since it has not exited yet. + runningQueryIds.add(entry.getKey()); + + if (queryIds.contains(entry.getKey())) { + // Null out missingFromSyncSince, signifying that this query is meant to be active. + holder.missingFromSyncSince = null; + } else { + // Add this query to workersToCancel if it has been missing long enough. + final Long missingFromSyncSince = holder.missingFromSyncSince; + if (missingFromSyncSince == null) { + holder.missingFromSyncSince = now; + } else if (now - missingFromSyncSince >= syncGraceMillis) { + holdersToCancel.add(holder); + } + } + } + } + + for (final WorkerHolder holder : holdersToCancel) { + log.info( + "Canceling query[%s] because controller[%s] no longer lists it as running.", + holder.workerContext.queryId(), + controllerHost + ); + holder.runRef.cancel(); Review Comment: Might be useful to emit a metric for the canceled queries with reason (missing query ID in sync or silent controller). -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
