[fix](remote-doris) End a remote Doris scan's Flight SQL session with the query; stop the regression framework leaking connections

fe/fe-core/src/main/java/org/apache/doris/datasource/doris/source/RemoteDorisFlightSession.java

@@ -0,0 +1,153 @@
+// 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.datasource.doris.source;
+
+import org.apache.doris.common.Pair;
+import org.apache.doris.common.UserException;
+
+import com.google.common.annotations.VisibleForTesting;
+import org.apache.arrow.flight.CallOptions;
+import org.apache.arrow.flight.CloseSessionRequest;
+import org.apache.arrow.flight.FlightClient;
+import org.apache.arrow.flight.FlightInfo;
+import org.apache.arrow.flight.Location;
+import org.apache.arrow.flight.grpc.CredentialCallOption;
+import org.apache.arrow.flight.sql.FlightSqlClient;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+
+import java.io.Closeable;
+import java.net.URI;
+import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * The Flight SQL session a {@link RemoteDorisScanNode} opens on a remote Doris frontend for one
+ * scan, and ends with a CloseSession once the scan is over.
+ *
+ * <p>The handshake ({@code authenticateBasicToken}) opens a session on the remote frontend: a
+ * connection in its pool, counted against {@code qe_max_connection}, the Arrow Flight SQL sub-quota
+ * and the catalog user's {@code max_user_connections}, that only a CloseSession, a KILL or
+ * {@code wait_timeout} ends. Closing the gRPC channel does not. A session opened per scan and never
+ * closed therefore stays for hours and, one scan at a time, exhausts the catalog user's connection
+ * quota on the remote frontend - refusing that user's MySQL connections there as well.
+ *
+ * <p>The session outlives GetFlightInfo on purpose: the query it ran serves the BE's DoGet of the
+ * endpoints, and the remote frontend cancels whatever a closed session was still running - the
+ * query itself, when the remote table is an external table scanned in batch mode and the query is
+ * therefore deferred there. So {@link #close()} is called from {@link RemoteDorisScanNode#stop()},
+ * when the coordinator of the local query closes, and that coordinator is kept alive until the BE
+ * has finished scanning ({@link RemoteDorisScanNode#coordinatorMustOutliveDispatch()}).
+ */
+class RemoteDorisFlightSession implements Closeable {
+    private static final Logger LOG = LogManager.getLogger(RemoteDorisFlightSession.class);
+
+    // A bound on the CloseSession round trip. close() runs on the local query's teardown path,
+    // which must not hang on a remote frontend that has stopped answering; the session is then left
+    // to the remote frontend's wait_timeout, as every session was before this class existed.
+    @VisibleForTesting
+    static final int CLOSE_SESSION_TIMEOUT_SECONDS = 5;
+
+    private final Pair<String, Integer> hostAndPort;
+    private final BufferAllocator allocator;
+    private final FlightSqlClient client;
+    private final CredentialCallOption credential;
+    private boolean closed = false;
+
+    private RemoteDorisFlightSession(Pair<String, Integer> hostAndPort, BufferAllocator allocator,
+            FlightSqlClient client, CredentialCallOption credential) {
+        this.hostAndPort = hostAndPort;
+        this.allocator = allocator;
+        this.client = client;
+        this.credential = credential;
+    }
+
+    /**
+     * Opens a session on the remote frontend at {@code hostAndPort} with the catalog's credentials.
+     * Nothing is left behind when this fails: a handshake that was refused opened no session, and
+     * the channel and allocator are released before the exception propagates.
+     */
+    static RemoteDorisFlightSession open(Pair<String, Integer> hostAndPort, String user, String password)
+            throws Exception {
+        BufferAllocator allocator = new RootAllocator();
+        FlightClient flightClient = null;
+        try {
+            URI uri = new URI("grpc", null, hostAndPort.first, hostAndPort.second, null, null, null);
+            flightClient = FlightClient.builder(allocator, new Location(uri)).build();
+            Optional<CredentialCallOption> credential = flightClient.authenticateBasicToken(user, password);
+            if (!credential.isPresent()) {
+                throw new UserException("Authenticates with a username and password failure");
+            }
+            return new RemoteDorisFlightSession(hostAndPort, allocator, new FlightSqlClient(flightClient),
+                    credential.get());
+        } catch (Throwable t) {
+            closeQuietly(flightClient, allocator, hostAndPort);
+            throw t;
+        }
+    }
+
+    /** Runs {@code sql} on the remote frontend; the endpoints of the result are where the BE reads it. */
+    FlightInfo execute(String sql, int timeoutSec) {
+        return client.execute(sql, credential, CallOptions.timeout(timeoutSec, TimeUnit.SECONDS));
+    }
+
+    Pair<String, Integer> getHostAndPort() {
+        return hostAndPort;
+    }
+
+    /**
+     * Ends the session on the remote frontend (CloseSession), then releases the channel and the
+     * allocator. Never throws: this runs on the local query's teardown path, and a session the
+     * remote frontend could not be told to close is only left to its wait_timeout. Idempotent.
+     */
+    @Override
+    public synchronized void close() {
+        if (closed) {
+            return;
+        }
+        closed = true;
+        try {
+            client.closeSession(new CloseSessionRequest(), credential,
+                    CallOptions.timeout(CLOSE_SESSION_TIMEOUT_SECONDS, TimeUnit.SECONDS));
+        } catch (Throwable t) {
+            LOG.warn("failed to close the Arrow Flight SQL session on remote Doris frontend {}:{}, it stays open there"
+                    + " until its wait_timeout", hostAndPort.first, hostAndPort.second, t);
+        }
+        closeQuietly(client, allocator, hostAndPort);
+    }
+
+    private static void closeQuietly(AutoCloseable client, BufferAllocator allocator,
+            Pair<String, Integer> hostAndPort) {
+        try {
+            if (client != null) {
+                client.close();
+            }
+        } catch (Throwable t) {
+            LOG.warn("failed to close the Arrow Flight client to remote Doris frontend {}:{}",
+                    hostAndPort.first, hostAndPort.second, t);
+        }
+        try {
+            allocator.close();
+        } catch (Throwable t) {
+            LOG.warn("failed to close the Arrow allocator of the Flight client to remote Doris frontend {}:{}",
+                    hostAndPort.first, hostAndPort.second, t);
+        }
+    }
+}

fe/fe-core/src/main/java/org/apache/doris/datasource/doris/source/RemoteDorisScanNode.java

@@ -44,29 +44,21 @@
 import org.apache.doris.thrift.TRemoteDorisFileDesc;
 import org.apache.doris.thrift.TTableFormatFileDesc;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Joiner;
 import com.google.common.collect.Lists;
-import org.apache.arrow.flight.CallOptions;
-import org.apache.arrow.flight.FlightClient;
 import org.apache.arrow.flight.FlightEndpoint;
 import org.apache.arrow.flight.FlightInfo;
 import org.apache.arrow.flight.Location;
-import org.apache.arrow.flight.grpc.CredentialCallOption;
-import org.apache.arrow.flight.sql.FlightSqlClient;
-import org.apache.arrow.memory.BufferAllocator;
-import org.apache.arrow.memory.RootAllocator;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
-import java.net.URI;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Optional;
 import java.util.Set;
-import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
 public class RemoteDorisScanNode extends FileQueryScanNode {
@@ -79,6 +71,18 @@ public class RemoteDorisScanNode extends FileQueryScanNode {
 
     private RemoteDorisSource source;
 
+    // The Flight SQL session this scan opened on the remote frontend, from getSplits until stop()
+    // closes it (see RemoteDorisFlightSession for why it must live that long and no longer). All
+    // three guarded by this: stop() may run on another thread than the one that planned the query -
+    // a KILL, the timeout checker - and more than once (cancel, then close).
+    private RemoteDorisFlightSession flightSession;
+    private boolean stopped;
+    // Whether stop() ended a session this scan had opened: the endpoints handed to the backend
+    // belong to that session's query, and a plan dispatched again with them (the same-plan retry
+    // of StmtExecutor.handleQueryWithRetry) would read what the remote frontend may have torn
+    // down with the session.
+    private boolean sessionClosedByStop;
+
     public RemoteDorisScanNode(PlanNodeId id, TupleDescriptor desc, boolean needCheckColumnPriv,
                                SessionVariable sv, ScanContext scanContext) {
         super(id, desc, "REMOTE_DORIS_SCAN_NODE", scanContext, needCheckColumnPriv, sv);
@@ -177,7 +181,8 @@ private List<Pair<String, ByteBuffer>> executeQuery() {
                     source.nextHostAndArrowPort(),
                     source.getCatalog().getUsername(),
                     source.getCatalog().getPassword(),
-                    queryStr
+                    queryStr,
+                    source.getCatalog().getQueryTimeoutSec()
                 );
             } catch (Exception e) {
                 LOG.warn("arrow request node [{}] failures {}, try next nodes",
@@ -189,17 +194,98 @@ private List<Pair<String, ByteBuffer>> executeQuery() {
         throw new RuntimeException("Failed to execute query: " + queryStr, lastException);
     }
 
-    private List<Pair<String, ByteBuffer>> executeFlightSqlQuery(Pair<String, Integer> hostAndPort,
-                     String user, String psw, String sql) throws Exception {
-        try (
-                BufferAllocator allocatorFE = new RootAllocator();
-                FlightClient clientFE = createFlightClient(allocatorFE, hostAndPort);
-                FlightSqlClient sqlClientFE = new FlightSqlClient(clientFE)
-        ) {
-            CredentialCallOption credentialCallOption = authenticate(clientFE, user, psw);
-            FlightInfo info = executeSqlWithTimeout(sqlClientFE, sql, credentialCallOption);
-
-            return processFlightEndpoints(info.getEndpoints());
+    // Opens a Flight SQL session on the remote frontend and runs the query in it. The session is
+    // kept until stop(): its query serves the BE's DoGet of the endpoints returned here. A session
+    // whose query failed is closed right away, so a retry on the next node leaves nothing behind.
+    @VisibleForTesting
+    List<Pair<String, ByteBuffer>> executeFlightSqlQuery(Pair<String, Integer> hostAndPort,
+                     String user, String psw, String sql, int timeoutSec) throws Exception {
+        RemoteDorisFlightSession session = RemoteDorisFlightSession.open(hostAndPort, user, psw);
+        FlightInfo info;
+        try {
+            info = session.execute(sql, timeoutSec);
+        } catch (Throwable t) {
+            session.close();
+            throw t;
+        }
+        keepFlightSession(session);
+        return processFlightEndpoints(info.getEndpoints());
+    }
+
+    /**
+     * Holds {@code session} until {@link #stop()}. A session handed over after stop() already ran,
+     * or on top of one still held, is closed at once instead: this scan owns one session at most,
+     * and none once stopped. The statement registers the node as well: stop() is the coordinator's
+     * to call, but a plan that never gets one, or whose coordinator nobody closes, is stopped when
+     * the statement ends instead ({@link org.apache.doris.nereids.StatementContext#stopScanNodeAtClose}).
+     */
+    @VisibleForTesting
+    void keepFlightSession(RemoteDorisFlightSession session) {
+        RemoteDorisFlightSession toClose;
+        synchronized (this) {
+            if (stopped) {
+                toClose = session;
+            } else {
+                toClose = flightSession;
+                flightSession = session;
+            }
+        }
+        if (toClose != null) {
+            toClose.close();
+        }
+        if (toClose != session) {
+            ConnectContext.get().getStatementContext().stopScanNodeAtClose(this);
+        }
+    }
+
+    /**
+     * True once {@link #stop()} ended the session this scan opened: the endpoints in its scan
+     * ranges belong to that session's query on the remote frontend, so the same plan must not be
+     * dispatched again (see {@link ScanNode#cannotBeRedispatched()}).
+     */
+    @Override
+    public boolean cannotBeRedispatched() {
+        synchronized (this) {
+            return sessionClosedByStop;
+        }
+    }
+
+    /**
+     * True while this scan holds a Flight SQL session on the remote frontend: the local coordinator
+     * has to stay alive until the BE has finished reading the remote query, since closing it is what
+     * ends the session ({@link #stop()}) - and the remote frontend cancels what a closed session was
+     * still running. Without this, an Arrow Flight SQL query on this frontend would close its
+     * coordinator right after dispatch (#67503), while its BE may still be reading.
+     */
+    @Override
+    public boolean coordinatorMustOutliveDispatch() {
+        if (super.coordinatorMustOutliveDispatch()) {
+            return true;
+        }
+        synchronized (this) {
+            return flightSession != null;
+        }
+    }
+
+    /**
+     * Ends the Flight SQL session on the remote frontend, in addition to what {@code FileQueryScanNode}
+     * releases. Called by the coordinator when the local query closes or is cancelled, i.e. when the
+     * BE is done with (or gave up on) the remote query's endpoints.
+     */
+    @Override
+    public void stop() {
+        super.stop();
+        RemoteDorisFlightSession session;
+        synchronized (this) {
+            stopped = true;
+            session = flightSession;
+            flightSession = null;
+            if (session != null) {
+                sessionClosedByStop = true;
+            }
+        }
+        if (session != null) {
+            session.close();
         }
     }
 
@@ -290,27 +376,6 @@ private boolean isExplainStatement() {
             .trim().toLowerCase().startsWith("explain");
     }
 
-    private FlightClient createFlightClient(BufferAllocator allocator,
-                                            Pair<String, Integer> hostAndPort) throws Exception {
-        URI uri = new URI("grpc", null, hostAndPort.first, hostAndPort.second, null, null, null);
-        return FlightClient.builder(allocator, new Location(uri)).build();
-    }
-
-    private CredentialCallOption authenticate(FlightClient client, String user, String psw) throws UserException {
-        Optional<CredentialCallOption> credentialCallOption = client.authenticateBasicToken(user, psw);
-        if (!credentialCallOption.isPresent()) {
-            throw new UserException("Authenticates with a username and password failure");
-        }
-        return credentialCallOption.get();
-    }
-
-    private FlightInfo executeSqlWithTimeout(FlightSqlClient sqlClient, String sql,
-                                             CredentialCallOption credentialCallOption) {
-        int timeoutSec = source.getCatalog().getQueryTimeoutSec();
-        return sqlClient.execute(sql, credentialCallOption,
-            CallOptions.timeout(timeoutSec, TimeUnit.SECONDS));
-    }
-
     private List<Pair<String, ByteBuffer>> processFlightEndpoints(List<FlightEndpoint> endpoints) {
         List<Pair<String, ByteBuffer>> uniquePairs = new ArrayList<>();
         Set<String> seenPairs = new HashSet<>();