morningman commented on code in PR #69996: URL: https://github.com/apache/doris/pull/69996#discussion_r4574163999
########## fe/fe-core/src/main/java/org/apache/doris/arrowflight/FlightSqlQueryCancellation.java: ########## @@ -0,0 +1,273 @@ +// 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.doris.arrowflight; + +import org.apache.doris.arrowflight.results.FlightSqlEndpointsLocation; +import org.apache.doris.common.Status; +import org.apache.doris.proto.InternalService.PCancelPlanFragmentResult; +import org.apache.doris.qe.ConnectContext; +import org.apache.doris.qe.Coordinator; +import org.apache.doris.rpc.BackendServiceProxy; +import org.apache.doris.thrift.TNetworkAddress; +import org.apache.doris.thrift.TStatus; +import org.apache.doris.thrift.TStatusCode; +import org.apache.doris.thrift.TUniqueId; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.function.LongSupplier; +import java.util.stream.Collectors; + +/** + * Routes the early close of a Flight result stream, which a BE reports through cancelFlightQuery, to every BE of + * the stream's query, and cancels the queries of a request that fails before answering with tickets for their + * results ({@link #cancelFailedRequest}). A route outlives the request and the session that registered it: the + * client reads the results from the BEs, so it can still hold a stream of the query open after its session has + * started another request or has been torn down, and close that stream early then. A route goes when its + * cancellation succeeds, when the request that registered it fails before answering with tickets for its results + * ({@link #unregister}), when its query's execution timeout (plus 5 s) has passed, and when newer routes push it + * out of the bounds: {@link #MAX_ROUTED_RESULTS} result ids and {@link #MAX_ROUTED_BACKENDS} BE addresses. Those go + * oldest first, as the oldest routes are the likeliest to belong to streams that have ended; a query whose route + * was pushed out ends at its timeout, as every Flight query did before routes existed. + */ +public final class FlightSqlQueryCancellation { + private static final Logger LOG = LogManager.getLogger(FlightSqlQueryCancellation.class); + // A route costs about 270 bytes plus about 30 bytes per address of its query's BEs, and the result ids of a + // query share one. Bounding both the result ids and the addresses keeps the registry within about 55 MB however + // many Flight queries are registered within one execution timeout, and however many BEs each of them runs on. + static final int MAX_ROUTED_RESULTS = 100_000; + static final int MAX_ROUTED_BACKENDS = 1_000_000; + public static final FlightSqlQueryCancellation INSTANCE = + new FlightSqlQueryCancellation(System::nanoTime, MAX_ROUTED_RESULTS, MAX_ROUTED_BACKENDS); + private final LongSupplier nanoClock; + private final int resultCapacity; + private final int backendCapacity; + // In registration order, so the eldest entry belongs to the oldest route. Guarded by this. + private final LinkedHashMap<TUniqueId, Route> results = new LinkedHashMap<>(); + // The addresses of the routes that some result id still maps to. Guarded by this. + private long routedBackends; + + FlightSqlQueryCancellation(LongSupplier nanoClock, int resultCapacity, int backendCapacity) { + this.nanoClock = nanoClock; + this.resultCapacity = resultCapacity; + this.backendCapacity = backendCapacity; + } + + /** + * Registers the route of a Flight query whose results stay on the BEs, from the result ids of the endpoints + * the session collected for it, before its coordinator is closed or deferred. Nothing of the session is + * kept, so neither its next request nor its teardown drops the route. + */ + public void register(ConnectContext context, Coordinator coord) { + register(coord.getQueryId(), resultIds(context), coord.getBackendBrpcAddresses(), + coord.getQueryOptions().getExecutionTimeout()); + } + + /** + * Drops the routes of the results the session collected for a request that fails before answering with + * tickets for them: no client can pull those results, so no BE can report an early close of them. Called + * only once the request has begun, so the session's endpoints are its own and not those of the previous + * request, whose streams may still be open. + */ + public void unregister(ConnectContext context) { + take(resultIds(context)); + } + + /** + * Like {@link #unregister}, and cancels on all its BEs the query of each route that goes, unless + * {@code cancelledQueries} names it. Such a query has closed its coordinator already, so nothing else + * would cancel it before the BE's result buffer timer (execution_timeout + 5 s) ends it; the named ones + * were cancelled through their deferred coordinators. The BEs' answers are not awaited, as a deferred + * coordinator's cancellation does not await them either: the routes go whatever they are, so there is no + * retry to drive, and the failed request's error should not wait for them. + */ + public void cancelFailedRequest(ConnectContext context, List<TUniqueId> cancelledQueries, Status reason) { + for (Route route : take(resultIds(context))) { + if (!cancelledQueries.contains(route.queryId)) { + sendCancellation(route, reason); + } + } + } + + public void register(TUniqueId queryId, List<TUniqueId> resultIds, + List<TNetworkAddress> backends, int timeoutSeconds) { + if (backends.isEmpty()) { + throw new IllegalArgumentException("Flight query has no cancellation backends"); + } + // A route leaves the registry, and releases its addresses, with its result ids. + if (resultIds.isEmpty()) { + throw new IllegalArgumentException("Flight query has no result ids"); + } + // Keep only cancellation addresses, not the coordinator, scan state, or query queue slot. + // Ordinary Flight coordinators are unregistered when GetFlightInfo returns. + Route route = new Route(queryId.deepCopy(), + resultIds.stream().map(TUniqueId::deepCopy).distinct().collect(Collectors.toList()), + backends.stream().map(TNetworkAddress::deepCopy).distinct().collect(Collectors.toList()), + nanoClock.getAsLong() + TimeUnit.SECONDS.toNanos(Math.max(0L, timeoutSeconds) + 5)); + synchronized (this) { + for (TUniqueId resultId : route.resultIds) { + Route previous = results.put(resultId, route); + if (previous != null) { + unrouteOne(previous); + } + } + route.routedIds = route.resultIds.size(); + routedBackends += route.backends.size(); + // Whole routes leave from the eldest end: the expired ones there, then as many as the bounds require. + long now = nanoClock.getAsLong(); + while (!results.isEmpty()) { + Route eldest = results.values().iterator().next(); + if (results.size() <= resultCapacity && routedBackends <= backendCapacity + && !eldest.isExpired(now)) { + break; + } + remove(eldest); + } + } + } + + public TStatus cancel(TUniqueId resultId) { + Route route; + synchronized (this) { + route = results.get(resultId); + if (route != null && route.isExpired(nanoClock.getAsLong())) { + remove(route); + route = null; + } + } + if (route == null) { + return new TStatus(TStatusCode.NOT_FOUND); + } + Status reason = new Status(TStatusCode.CANCELLED, "Arrow Flight stream closed before EOF"); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(1); + List<Future<PCancelPlanFragmentResult>> futures = sendCancellation(route, reason); Review Comment: Not changing this here: the blocking step is master's shared FE-to-BE client, and nothing is lost while it blocks. `BackendServiceProxy.getProxy()` serves every FE-to-BE RPC. When a cached client's IP changed (FQDN mode, after the 60 s DNS refresh) or its channel left CONNECTING/IDLE/READY, it publishes a new client and retires the old one on the caller's thread: `shutdown()` waits only while that channel has calls in flight, up to 5 s + 5 s, once per proxy instance (48) and BE, for whichever RPC comes first. `Coordinator.cancel` dispatches the same way, including the deferred-coordinator cancellation in the same failed-GetFlightInfo catch; "not awaited" means the answers, as there. Meanwhile the handler still cancels on every BE once the channel is retired, answers INTERNAL_ERROR so the route stays, and the BE keeps its own bound (2 s per attempt, one retry, then the next FE). The retry goes through another proxy instance and normally removes the route; otherwise the query is cancelled up to 10 s late and the route expires. Retiring replaced clients asynchronously in `getProxy()` fixes every caller; that belongs in its own PR, not in an executor for this one caller. The description's "at most 1 s" bounds the wait for the answers; the merge message will say so. -- 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]
