xiangfu0 commented on code in PR #18709: URL: https://github.com/apache/pinot/pull/18709#discussion_r3380472122
########## pinot-broker/src/main/java/org/apache/pinot/broker/broker/BrokerDrainManager.java: ########## @@ -0,0 +1,333 @@ +/** + * 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.pinot.broker.broker; + +import com.google.common.annotations.VisibleForTesting; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Supplier; +import javax.annotation.Nullable; +import org.apache.helix.HelixManager; +import org.apache.helix.model.HelixConfigScope; +import org.apache.helix.model.builder.HelixConfigScopeBuilder; +import org.apache.pinot.common.utils.helix.HelixHelper; +import org.apache.pinot.spi.utils.CommonConstants; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/// Coordinates broker query drain. +/// +/// Drain has three effects: +/// +/// 1. New query admissions are rejected locally. +/// 2. The broker marks its Helix instance config with `shutdownInProgress=true` +/// and removes itself from `brokerResource`, so clients that discover brokers +/// through Helix stop selecting it. +/// 3. A caller can wait for in-flight queries to finish and optionally trigger +/// broker shutdown. +/// +/// The in-flight count tracks accepted broker query requests, not engine-internal +/// subqueries. That makes the wait semantics independent of single-stage, +/// multi-stage, and time-series query execution details. +public class BrokerDrainManager { + private static final Logger LOGGER = LoggerFactory.getLogger(BrokerDrainManager.class); + private static final long SHUTDOWN_CALLBACK_DELAY_MS = 1_000L; + + private final String _instanceId; + private final Supplier<HelixManager> _helixManagerSupplier; + private final Runnable _drainStartedCallback; + private final Runnable _shutdownCallback; + private final long _defaultDrainTimeoutMs; + private final boolean _enableHelixUpdates; + + private final Object _lock = new Object(); + private final AtomicBoolean _draining = new AtomicBoolean(false); + private final AtomicBoolean _shutdownTriggered = new AtomicBoolean(false); + private volatile long _drainStartTimeMs = -1L; + private volatile List<String> _tablesRemovedFromBrokerResource = Collections.emptyList(); + private int _inFlightQueries; + + public BrokerDrainManager(String instanceId, Supplier<HelixManager> helixManagerSupplier, + Runnable drainStartedCallback, Runnable shutdownCallback, long defaultDrainTimeoutMs) { + this(instanceId, helixManagerSupplier, drainStartedCallback, shutdownCallback, defaultDrainTimeoutMs, true); + } + + private BrokerDrainManager(String instanceId, Supplier<HelixManager> helixManagerSupplier, + Runnable drainStartedCallback, Runnable shutdownCallback, long defaultDrainTimeoutMs, + boolean enableHelixUpdates) { + _instanceId = instanceId; + _helixManagerSupplier = helixManagerSupplier; + _drainStartedCallback = drainStartedCallback; + _shutdownCallback = shutdownCallback; + _defaultDrainTimeoutMs = defaultDrainTimeoutMs; + _enableHelixUpdates = enableHelixUpdates; + } + + public static BrokerDrainManager noop(String instanceId) { + return new BrokerDrainManager(instanceId, () -> null, () -> { + }, () -> { + }, 0L, false); + } + + @VisibleForTesting + public static BrokerDrainManager localOnly(String instanceId, Runnable drainStartedCallback, + Runnable shutdownCallback, long defaultDrainTimeoutMs) { + return new BrokerDrainManager(instanceId, () -> null, drainStartedCallback, shutdownCallback, defaultDrainTimeoutMs, + false); + } + + @Nullable + public QueryPermit tryAcquireQuery() { + synchronized (_lock) { + if (_draining.get()) { + return null; + } + _inFlightQueries++; + return new QueryPermit(this); + } + } + + public DrainStatus drain(long timeoutMs, boolean shutdown) + throws InterruptedException { + startDrain(); + boolean drained = awaitNoInflightQueries(resolveTimeoutMs(timeoutMs)); + if (drained && shutdown) { + triggerShutdown(); + } + return getStatus(drained ? "Broker drained" : "Timed out waiting for in-flight queries to finish"); + } + + public DrainStatus getStatus() { + return getStatus(_draining.get() ? "Broker is draining" : "Broker is accepting queries"); + } + + public boolean isDraining() { + return _draining.get(); + } + + public boolean isDrainComplete() { + synchronized (_lock) { + return _draining.get() && _inFlightQueries == 0; + } + } + + public String getRejectMessage() { + return "Broker " + _instanceId + " is draining and not accepting new queries"; + } + + private void startDrain() { + if (!_draining.compareAndSet(false, true)) { Review Comment: startDrain() marks the broker draining before `getConnectedHelixManager()` and the Helix writes succeed. Because `BaseBrokerStarter` exposes the admin API before `_participantHelixManager.connect()`, a startup-window `POST /drain` (or any Helix write failure here) returns 500 but still leaves `_draining` set, so later queries are rejected until restart. Please make this transition rollback-safe or only flip the local drain state after the Helix preconditions/writes succeed. -- 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]
