This is an automated email from the ASF dual-hosted git repository.
morningman pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/master by this push:
new 46e1c671220 [refactor](arrow-flight) One connection pool for MySQL
connections and Arrow Flight SQL sessions (#68101)
46e1c671220 is described below
commit 46e1c67122085ad624c715a96a38899fe1d8e9c6
Author: Mingyu Chen (Rayner) <[email protected]>
AuthorDate: Sun Sep 20 15:46:53 2026 +0800
[refactor](arrow-flight) One connection pool for MySQL connections and
Arrow Flight SQL sessions (#68101)
### What problem does this PR solve?
Issue Number: #67577
Related PR: #67966 (the Flight session teardown contract this keeps,
D-13)
Problem Summary:
Stage 3, item 1 of #67577 (connection governance): one connection pool
for both protocols.
Arrow Flight SQL sessions lived in a pool of their own,
`FlightSqlConnectPoolMgr`, held to `arrow_flight_max_connections` alone:
not to `qe_max_connection`, not to the user's `max_user_connections`,
refused with a sentence of their own ("the number of arrow flight bearer
tokens should be equal to arrow flight sql max connections"), and summed
into every scheduler-wide answer — the processlist, `KILL`, the timeout
checker, the connection metrics — by `ConnectScheduler` walking two
pools.
There is one pool now:
- `ConnectPoolMgr` registers a connection of either protocol against
`qe_max_connection` and the user's `max_user_connections`, and a Flight
session additionally against its sub-quota. It indexes Flight sessions
by their peer identity (the bearer token), since that is how Flight
requests name their session. `FlightSqlConnectPoolMgr` is deleted;
`ConnectScheduler` delegates to the one pool.
- Admission is one atomic transition:
`ConnectPoolMgr.registerConnection` checks and reserves the pool, Flight
sub-quota and per-user quotas under a single lock (`admissionLock`), so
a doomed attempt cannot hold one counter's reservation across another
attempt's checks. Session creation is serialized too
(`sessionCreationLock` in `FlightSessionsWithTokenManager`): exactly one
`ConnectContext` is published per bearer token, and after registering it
the token is re-validated under the same lock and the session
unregistered if the token was invalidated meanwhile, so no session
lingers in the pool without its token.
- The bearer token cache size floors only the effective sub-quota at 1
(a legal sub-quota of 0 at `qe_max_connection = 1` still needs one token
so the first request reaches the pool and is refused with
`RESOURCE_EXHAUSTED`); an illegal `arrow_flight_token_cache_size` (<= 0)
is left as-is, keeping the existing startup failure.
- `unregisterConnection` is where every teardown path of a connection
meets — a MySQL channel closing, a Flight token expiring or being
evicted, `CloseSession`, `KILL`, the timeout checker — so it is where
the protocol releases what it still holds for the session, through the
new `ProtocolAdapter.releaseSession`: for Flight the channel-cached
results and the deferred query coordinators
(`FlightProtocolAdapter.tearDown`, unchanged from #67966), nothing for
MySQL. The transitional `ProtocolAdapter.connectPool` goes.
- A Flight session refused for a limit is refused in the words a MySQL
client is refused in (`Reach limit of connections. Total: %d, User: %d,
Current: %d`, naming the Flight sub-quota only when it is the tighter
one), as the `RESOURCE_EXHAUSTED` status of the request that would have
opened it, and its bearer token is invalidated with the refusal.
- `arrow_flight_max_connections` becomes the sub-quota of Flight
sessions within the pool: `-1`, the new default, is half of
`qe_max_connection`; an explicit value never exceeds it.
`connection_max` reports the pool's limit,
`arrow_flight_connection_total` / `arrow_flight_connection_max` the
Flight share of it. The token cache is sized from the effective
sub-quota and is otherwise untouched: the token's lifecycle is the next
PR's.
---
.../main/java/org/apache/doris/common/Config.java | 34 ++-
.../doris/arrowflight/DorisFlightSqlService.java | 42 ++-
.../protocol/FlightProtocolAdapter.java | 27 +-
.../sessions/FlightSessionsWithTokenManager.java | 63 ++++-
.../sessions/FlightSqlConnectPoolMgr.java | 99 -------
.../arrowflight/tokens/FlightTokenManagerImpl.java | 29 +-
.../java/org/apache/doris/metric/MetricRepo.java | 8 +-
.../org/apache/doris/mysql/AcceptListener.java | 9 +-
.../doris/mysql/protocol/MysqlProtocolAdapter.java | 6 +-
.../java/org/apache/doris/qe/ConnectContext.java | 17 +-
.../java/org/apache/doris/qe/ConnectPoolMgr.java | 133 ++++++++-
.../java/org/apache/doris/qe/ConnectScheduler.java | 81 ++----
.../apache/doris/qe/protocol/ProtocolAdapter.java | 11 +-
.../arrowflight/DorisFlightSqlProducerTest.java | 7 +-
.../arrowflight/DorisFlightSqlServiceTest.java | 42 +++
.../protocol/FlightProtocolAdapterTest.java | 52 +++-
.../sessions/FlightSqlConnectPoolMgrTest.java | 70 -----
.../java/org/apache/doris/metric/MetricsTest.java | 37 ++-
.../apache/doris/mysql/ConnectionExceedTest.java | 128 +++++++--
.../mysql/protocol/MysqlProtocolAdapterTest.java | 4 +-
.../org/apache/doris/qe/ConnectPoolMgrTest.java | 314 +++++++++++++++++++++
.../apache/doris/qe/ConnectPoolTestSupport.java | 65 +++++
.../test_connection_quota.groovy | 246 ++++++++++++++++
23 files changed, 1171 insertions(+), 353 deletions(-)
diff --git a/fe/fe-common/src/main/java/org/apache/doris/common/Config.java
b/fe/fe-common/src/main/java/org/apache/doris/common/Config.java
index 7d020f42c35..c13a29dc507 100644
--- a/fe/fe-common/src/main/java/org/apache/doris/common/Config.java
+++ b/fe/fe-common/src/main/java/org/apache/doris/common/Config.java
@@ -758,7 +758,9 @@ public class Config extends ConfigBase {
+ "Set long enough to fit your tablet size.")
public static long check_consistency_default_timeout_second = 600; // 10
min
- @ConfField(description = "Maximum number of MySQL server connections per
FE.")
+ @ConfField(description = "Maximum number of connections per FE. MySQL
connections and Arrow Flight SQL "
+ + "sessions share this one pool (see arrow_flight_max_connections
for the share Flight sessions "
+ + "may take of it: half by default).")
public static int qe_max_connection = 1024;
@ConfField(mutable = true, description = "Colocate join PlanFragment
instance memory limit penalty factor. The "
@@ -2650,8 +2652,23 @@ public class Config extends ConfigBase {
+ "automatically. Set to 0 or negative value to disable " + "this
limit for user-specified buckets.")
public static int max_bucket_num_per_partition = 768;
- @ConfField(description = "Maximum number of connections for the Arrow
Flight Server per FE.")
- public static int arrow_flight_max_connections = 4096;
+ @ConfField(description = "Arrow Flight SQL sessions share the one
connection pool with MySQL connections:"
+ + " both count against qe_max_connection and the user's
max_user_connections. This is the sub-quota of"
+ + " Arrow Flight SQL sessions within that pool: -1 (the default)
is half of qe_max_connection (512 with"
+ + " the default pool of 1024), and an explicit value never exceeds
qe_max_connection (a larger one is"
+ + " capped, with a warning at startup). Mind how a Flight session
ends: with CloseSession, a KILL"
+ + " CONNECTION from another connection, wait_timeout, or the
expiry or eviction of its bearer token"
+ + " (arrow_flight_token_alive_time_second; the token cache, see
below; and per user at most"
+ + " max_user_connections / 2 tokens). Most Flight clients never
send CloseSession, so a session whose"
+ + " client has gone stays in the pool until wait_timeout (8 hours
by default) or its token's expiry"
+ + " (24 hours by default), whichever comes first; the default
leaves the other half of the pool to"
+ + " MySQL connections however many such sessions there are. Raise
it with qe_max_connection, or set"
+ + " it to qe_max_connection on an FE that serves Arrow Flight SQL
only. The bearer token cache is"
+ + " sized to this sub-quota (capped by
arrow_flight_token_cache_size), so the Flight limit shows as"
+ + " the eviction of the oldest token and its session rather than
as a refusal. -1 is accepted from"
+ + " this version on: an older FE that serves Arrow Flight SQL
exits at startup with -1 in fe.conf;"
+ + " remove the setting or set a positive value before a
downgrade.")
+ public static int arrow_flight_max_connections = -1;
@ConfField(mutable = true, description = "Arrow Flight SQL only. A query
that scans an external table in "
+ "batch mode keeps its FE coordinator alive after GetFlightInfo,
so the BE can keep fetching splits "
@@ -2679,10 +2696,13 @@ public class Config extends ConfigBase {
+ "an abnormal case and triggers an alert.")
public static double autobucket_out_of_bounds_percent_threshold = 0.5;
- @ConfField(description = "(Deprecated, replaced by
arrow_flight_max_connection) The cache limit of all user "
- + "tokens in Arrow Flight Server, which will be eliminated by LRU
rules after exceeding "
- + "the limit. Arrow Flight SQL is a stateless protocol; the
connection is usually not "
- + "actively disconnected. A bearer token evicted from the cache
will unregister its " + "ConnectContext.")
+ @ConfField(description = "The cap of the bearer token cache of the Arrow
Flight SQL server. The cache holds"
+ + " as many tokens as the Arrow Flight SQL sub-quota of the
connection pool allows"
+ + " (arrow_flight_max_connections, half of qe_max_connection by
default) but never more than"
+ + " this; beyond that the oldest token is evicted by LRU, and the
session it names is closed with it."
+ + " Arrow Flight SQL clients rarely close their session, so the
effective cache size - the sub-quota"
+ + " unless this is smaller - is what bounds their sessions in
practice, and per user"
+ + " max_user_connections / 2 tokens. The effective cache size is
logged when the server starts.")
public static int arrow_flight_token_cache_size = 4096;
@ConfField(description = "The alive time of the user token in Arrow Flight
Server (expire after write), in "
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/DorisFlightSqlService.java
b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/DorisFlightSqlService.java
index 4bc5f1e48cd..d1548bd7db0 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/DorisFlightSqlService.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/DorisFlightSqlService.java
@@ -24,8 +24,10 @@ import
org.apache.doris.arrowflight.sessions.FlightSessionsWithTokenManager;
import org.apache.doris.arrowflight.tokens.FlightTokenManager;
import org.apache.doris.arrowflight.tokens.FlightTokenManagerImpl;
import org.apache.doris.common.Config;
+import org.apache.doris.qe.ConnectPoolMgr;
import org.apache.doris.service.FrontendOptions;
+import com.google.common.annotations.VisibleForTesting;
import io.grpc.ServerBuilder;
import org.apache.arrow.flight.FlightServer;
import org.apache.arrow.flight.Location;
@@ -48,12 +50,40 @@ public class DorisFlightSqlService {
private final FlightSessionsManager flightSessionsManager;
private volatile boolean running;
+ /**
+ * The bearer token cache size: the effective Flight sub-quota (a session
opens on a token, so the
+ * sub-quota is what bounds live sessions), capped by {@code
arrow_flight_token_cache_size}. The
+ * sub-quota is floored at 1 -- a legal sub-quota of 0 ({@code
qe_max_connection = 1}) still needs
+ * one token so the first request reaches the pool and is refused with
RESOURCE_EXHAUSTED, instead
+ * of a {@code maximumSize(0)} cache evicting the freshly issued token and
answering UNAUTHENTICATED.
+ * The floor is NOT applied to {@code arrow_flight_token_cache_size}: an
illegal value there
+ * ({@literal <= 0}) stays the loud failure it is on the base (Guava
rejects a negative maximumSize;
+ * 0 evicts every token) rather than the FE silently running on a
one-token cache.
+ */
+ @VisibleForTesting
+ static int effectiveTokenCacheSize(int flightMaxConnections, int
tokenCacheConfig) {
+ return Math.min(Math.max(1, flightMaxConnections), tokenCacheConfig);
+ }
+
public DorisFlightSqlService(int port) {
BufferAllocator allocator = new RootAllocator();
// arrow flight sql is a stateless protocol, connection is usually not
actively disconnected.
// bearer token is evict from the cache will unregister ConnectContext.
- this.flightTokenManager = new FlightTokenManagerImpl(
- Math.min(Config.arrow_flight_max_connections,
Config.arrow_flight_token_cache_size),
+ int flightMaxConnections =
ConnectPoolMgr.effectiveFlightMaxConnections(
+ Config.qe_max_connection, Config.arrow_flight_max_connections);
+ if (Config.arrow_flight_max_connections > Config.qe_max_connection) {
+ // A fe.conf from before the pools were merged may still carry the
old default, 4096: capped
+ // to the whole pool, that is exactly the sharing the default of
half is there to prevent.
+ LOG.warn("arrow_flight_max_connections={} exceeds
qe_max_connection={}: Arrow Flight SQL sessions are"
+ + " connections of the one pool, so the sub-quota
is capped at {}, the whole pool."
+ + " Flight sessions, which their clients mostly
never close, can then hold every"
+ + " connection until wait_timeout and refuse MySQL
logins. On an FE that serves both"
+ + " protocols, remove the setting (the default is
half of qe_max_connection; 4096 was"
+ + " the default before the pools were merged) or
set it below qe_max_connection",
+ Config.arrow_flight_max_connections,
Config.qe_max_connection, flightMaxConnections);
+ }
+ int tokenCacheSize = effectiveTokenCacheSize(flightMaxConnections,
Config.arrow_flight_token_cache_size);
+ this.flightTokenManager = new FlightTokenManagerImpl(tokenCacheSize,
Config.arrow_flight_token_alive_time_second);
this.flightSessionsManager = new
FlightSessionsWithTokenManager(flightTokenManager);
@@ -63,9 +93,11 @@ public class DorisFlightSqlService {
.transportHint(GRPC_BUILDER_CONSUMER,
(Consumer<ServerBuilder<?>>) builder ->
builder.addStreamTracerFactory(new
FlightRemoteIpServerStreamTracer.Factory()))
.headerAuthenticator(new
FlightBearerTokenAuthenticator(flightTokenManager)).build();
- LOG.info("Arrow Flight SQL service is created, port: {},
arrow_flight_max_connections: {}，"
- + "arrow_flight_token_alive_time_second: {}", port,
Config.arrow_flight_max_connections,
- Config.arrow_flight_token_alive_time_second);
+ LOG.info("Arrow Flight SQL service is created, port: {},
arrow_flight_max_connections: {} (effective: {},"
+ + " within qe_max_connection: {}), token cache size:
{} (arrow_flight_token_cache_size: {}),"
+ + " arrow_flight_token_alive_time_second: {}", port,
+ Config.arrow_flight_max_connections, flightMaxConnections,
Config.qe_max_connection,
+ tokenCacheSize, Config.arrow_flight_token_cache_size,
Config.arrow_flight_token_alive_time_second);
}
// start Arrow Flight SQL service, return true if success, otherwise false
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/protocol/FlightProtocolAdapter.java
b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/protocol/FlightProtocolAdapter.java
index 7f8a376dba5..340e8058c43 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/protocol/FlightProtocolAdapter.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/protocol/FlightProtocolAdapter.java
@@ -24,10 +24,9 @@ import org.apache.doris.common.Config;
import org.apache.doris.common.ErrorCode;
import org.apache.doris.common.Status;
import org.apache.doris.common.util.DebugUtil;
+import org.apache.doris.common.util.TokenMasker;
import org.apache.doris.qe.ConnectContext;
import org.apache.doris.qe.ConnectContext.ConnectType;
-import org.apache.doris.qe.ConnectPoolMgr;
-import org.apache.doris.qe.ConnectScheduler;
import org.apache.doris.qe.QueryState;
import org.apache.doris.qe.ShowResultSet;
import org.apache.doris.qe.StmtExecutor;
@@ -266,9 +265,27 @@ public class FlightProtocolAdapter implements
ProtocolAdapter {
public void fillForwardRequest(ConnectContext ctx, TMasterOpRequest
request) {
}
+ /**
+ * Every Flight SQL session teardown path - the idle timeout
(wait_timeout), bearer token expiry
+ * or eviction, CloseSession, a KILL CONNECTION from another connection -
reaches here through
+ * the pool's unregisterConnection. The
+ * channel-cached Arrow results go first, then the session is closed for
good: teardown does not
+ * wait for a command that may still be running, and what that command
defers afterwards is
+ * finalized on the spot ({@link #tearDown}).
+ */
@Override
- public ConnectPoolMgr connectPool(ConnectScheduler scheduler) {
- return scheduler.getFlightSqlConnectPoolMgr();
+ public void releaseSession(ConnectContext ctx) {
+ try {
+ channel.close();
+ } catch (Throwable t) {
+ // RootAllocator.close() marks the allocator closed before it
reports outstanding
+ // bytes. The error is actionable, but session teardown must still
release the
+ // coordinator, transaction and pool/token bookkeeping. The peer
identity IS the bearer
+ // token, so it is logged as a masked id, the same one
FlightTokenManagerImpl uses.
+ LOG.warn("failed to close Flight SQL channel while unregistering
connection {}, peer identity {}",
+ ctx.getConnectionId(), TokenMasker.tokenId(peerIdentity),
t);
+ }
+ tearDown();
}
/**
@@ -345,7 +362,7 @@ public class FlightProtocolAdapter implements
ProtocolAdapter {
@Override
public void closeConnection(ConnectContext ctx) {
// Releases the channel, the deferred executors and the transaction of
the session.
- connectPool(ctx.getConnectScheduler()).unregisterConnection(ctx);
+
ctx.getConnectScheduler().getConnectPoolMgr().unregisterConnection(ctx);
}
public String getPeerIdentity() {
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/sessions/FlightSessionsWithTokenManager.java
b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/sessions/FlightSessionsWithTokenManager.java
index 0327c24b61f..14a20989e56 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/sessions/FlightSessionsWithTokenManager.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/sessions/FlightSessionsWithTokenManager.java
@@ -23,10 +23,12 @@ import org.apache.doris.common.ErrorCode;
import org.apache.doris.common.util.TokenMasker;
import org.apache.doris.common.util.Util;
import org.apache.doris.qe.ConnectContext;
+import org.apache.doris.qe.ConnectPoolMgr;
import org.apache.doris.qe.ConnectScheduler;
import org.apache.doris.service.ExecuteEnv;
import org.apache.arrow.flight.CallStatus;
+import org.apache.arrow.flight.FlightRuntimeException;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
@@ -34,6 +36,8 @@ public class FlightSessionsWithTokenManager implements
FlightSessionsManager {
private static final Logger LOG =
LogManager.getLogger(FlightSessionsWithTokenManager.class);
private final FlightTokenManager flightTokenManager;
+ // Serializes session creation so exactly one ConnectContext is published
per bearer token.
+ private final Object sessionCreationLock = new Object();
public FlightSessionsWithTokenManager(FlightTokenManager
flightTokenManager) {
this.flightTokenManager = flightTokenManager;
@@ -42,13 +46,27 @@ public class FlightSessionsWithTokenManager implements
FlightSessionsManager {
@Override
public ConnectContext getConnectContext(String peerIdentity) {
try {
- ConnectContext connectContext =
ExecuteEnv.getInstance().getScheduler().getFlightSqlConnectPoolMgr()
- .getContextWithFlightToken(peerIdentity);
- if (null == connectContext) {
- connectContext = createConnectContext(peerIdentity);
+ ConnectContext connectContext =
ExecuteEnv.getInstance().getScheduler()
+ .getContextWithPeerIdentity(peerIdentity);
+ if (null != connectContext) {
return connectContext;
}
- return connectContext;
+ // Publish exactly one context per bearer token. Two concurrent
first requests would both
+ // find none above and, without this, both build and register one:
two pool slots and quota
+ // counters for one token, with the peer-identity index naming
only the last, so evicting
+ // that one leaves the other orphaned in the pool until
wait_timeout. Serialize creation and
+ // re-check the index inside, so the loser returns the winner's
published context.
+ synchronized (sessionCreationLock) {
+ connectContext = ExecuteEnv.getInstance().getScheduler()
+ .getContextWithPeerIdentity(peerIdentity);
+ if (null != connectContext) {
+ return connectContext;
+ }
+ return createConnectContext(peerIdentity);
+ }
+ } catch (FlightRuntimeException e) {
+ // Already the status the client is meant to see (a connection
refused for its limit).
+ throw e;
} catch (Exception e) {
LOG.warn("get ConnectContext failed, " + e.getMessage(), e);
throw
CallStatus.INTERNAL.withDescription(Util.getRootCauseMessage(e)).withCause(e).toRuntimeException();
@@ -71,15 +89,34 @@ public class FlightSessionsWithTokenManager implements
FlightSessionsManager {
flightTokenDetails.getUserIdentity(),
flightTokenDetails.getRemoteIp());
ConnectScheduler connectScheduler =
ExecuteEnv.getInstance().getScheduler();
connectScheduler.submit(connectContext);
- int res =
connectScheduler.getFlightSqlConnectPoolMgr().registerConnection(connectContext);
+ // The one pool every protocol registers in: qe_max_connection, the
user's
+ // max_user_connections and the Arrow Flight SQL sub-quota, refused in
the words a MySQL
+ // client is refused in. The token goes with the refusal, so that the
client does not keep a
+ // credential that can never open a session.
+ ConnectPoolMgr pool = connectScheduler.getConnectPoolMgr();
+ int res = pool.registerConnection(connectContext);
if (res >= 0) {
- String errMsg = String.format(
- "Register arrow flight sql connection failed, Unknown
Error, the number of arrow flight "
- + "bearer tokens should be equal to arrow flight
sql max connections, "
- + "max connections: %d, used: %d.",
-
connectScheduler.getFlightSqlConnectPoolMgr().getMaxConnections(), res);
- connectContext.getState().setError(ErrorCode.ERR_UNKNOWN_ERROR,
errMsg);
- throw new IllegalArgumentException(errMsg);
+ String errMsg = pool.limitReachedMessage(connectContext, res);
+
connectContext.getState().setError(ErrorCode.ERR_TOO_MANY_USER_CONNECTIONS,
errMsg);
+ // The refused session never entered the pool, so nothing else
releases what its
+ // adapter allocated (the channel's allocator).
+ connectContext.releaseProtocolSession();
+ flightTokenManager.invalidateToken(peerIdentity);
+ LOG.warn("refuse arrow flight sql session, bearer token id: {},
user: {}: {}",
+ TokenMasker.tokenId(peerIdentity),
connectContext.getQualifiedUser(), errMsg);
+ throw
CallStatus.RESOURCE_EXHAUSTED.withDescription(errMsg).toRuntimeException();
+ }
+ // Called under sessionCreationLock (see getConnectContext). The token
can be invalidated between
+ // the validateToken above and publishing the session here -- a
concurrent CloseSession, or the
+ // same-segment LRU evicting it -- which would leave a session in the
pool whose token no longer
+ // exists, holding its pool/user/Flight slots until wait_timeout.
Re-check the token and, if it is
+ // gone, unregister the session so no orphan lingers. (The token
lifecycle moves under one owner
+ // in the next PR, which closes the residual window after this
re-check.)
+ try {
+ flightTokenManager.validateToken(peerIdentity);
+ } catch (IllegalArgumentException e) {
+
connectScheduler.getConnectPoolMgr().unregisterConnection(connectContext);
+ throw e;
}
return connectContext;
}
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/sessions/FlightSqlConnectPoolMgr.java
b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/sessions/FlightSqlConnectPoolMgr.java
deleted file mode 100644
index a65a73b23cd..00000000000
---
a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/sessions/FlightSqlConnectPoolMgr.java
+++ /dev/null
@@ -1,99 +0,0 @@
-// 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.sessions;
-
-import org.apache.doris.arrowflight.results.FlightSqlChannel;
-import org.apache.doris.common.util.TokenMasker;
-import org.apache.doris.qe.ConnectContext;
-import org.apache.doris.qe.ConnectContext.ConnectType;
-import org.apache.doris.qe.ConnectPoolMgr;
-
-import com.google.common.collect.Maps;
-import org.apache.logging.log4j.LogManager;
-import org.apache.logging.log4j.Logger;
-
-import java.util.Map;
-
-public class FlightSqlConnectPoolMgr extends ConnectPoolMgr {
- private static final Logger LOG = LogManager.getLogger(
- FlightSqlConnectPoolMgr.class);
- private final Map<String, Integer> flightToken2ConnectionId =
Maps.newConcurrentMap();
-
- public FlightSqlConnectPoolMgr(int maxConnections) {
- super(maxConnections);
- }
-
- // Register one connection with its connection id.
- // Return -1 means register OK
- // Return >=0 means register failed, and return value is current
connection num.
- @Override
- public int registerConnection(ConnectContext ctx) {
- if (numberConnection.incrementAndGet() > maxConnections) {
- numberConnection.decrementAndGet();
- return numberConnection.get();
- }
- // not check user
- connectionMap.put(ctx.getConnectionId(), ctx);
- if (ctx.getConnectType().equals(ConnectType.ARROW_FLIGHT_SQL)) {
- flightToken2ConnectionId.put(ctx.getPeerIdentity(),
ctx.getConnectionId());
- }
- return -1;
- }
-
- @Override
- public void unregisterConnection(ConnectContext ctx) {
- // All Flight SQL session teardown paths (idle/query timeout, bearer
token expiry, and
- // explicit CloseSession) reach here. Release channel-cached Arrow
results before removing
- // the context from the pool.
- FlightSqlChannel flightSqlChannel = ctx.getFlightSqlChannel();
- if (flightSqlChannel != null) {
- try {
- flightSqlChannel.close();
- } catch (Throwable t) {
- // RootAllocator.close() marks the allocator closed before it
reports outstanding
- // bytes. The error is actionable, but session teardown must
still release the
- // coordinator, transaction and pool/token bookkeeping below.
- // For an Arrow Flight SQL connection the peer identity IS the
bearer token, so it is
- // logged as a masked id, the same one FlightTokenManagerImpl
uses.
- LOG.warn("failed to close Flight SQL channel while
unregistering connection {}, peer identity {}",
- ctx.getConnectionId(),
TokenMasker.tokenId(ctx.getPeerIdentity()), t);
- }
- }
- // Finalize any Arrow Flight query whose coordinator was kept alive
across the
- // GetFlightInfo -> DoGet phases (see #62259), releasing its resources
(e.g. external-table
- // batch SplitSources and the query queue slot), and close the session
for good: teardown
- // does not wait for a command that may still be running, and what
that command defers
- // afterwards is finalized on the spot
(FlightProtocolAdapter.tearDown).
- ctx.tearDownFlightSqlSession();
- ctx.closeTxn();
- if (connectionMap.remove(ctx.getConnectionId()) != null) {
- numberConnection.decrementAndGet();
- if (ctx.getConnectType().equals(ConnectType.ARROW_FLIGHT_SQL)) {
- flightToken2ConnectionId.remove(ctx.getPeerIdentity());
- }
- }
- }
-
- public ConnectContext getContextWithFlightToken(String flightToken) {
- if (flightToken2ConnectionId.containsKey(flightToken)) {
- int connectionId = flightToken2ConnectionId.get(flightToken);
- return getContext(connectionId);
- }
- return null;
- }
-}
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/tokens/FlightTokenManagerImpl.java
b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/tokens/FlightTokenManagerImpl.java
index ca2487b8379..7d24692f52f 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/arrowflight/tokens/FlightTokenManagerImpl.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/arrowflight/tokens/FlightTokenManagerImpl.java
@@ -21,9 +21,11 @@ package org.apache.doris.arrowflight.tokens;
import org.apache.doris.arrowflight.auth2.FlightAuthResult;
import org.apache.doris.catalog.Env;
+import org.apache.doris.common.Config;
import org.apache.doris.common.CustomThreadFactory;
import org.apache.doris.common.util.TokenMasker;
import org.apache.doris.qe.ConnectContext;
+import org.apache.doris.qe.ConnectPoolMgr;
import org.apache.doris.service.ExecuteEnv;
import com.google.common.base.Preconditions;
@@ -60,10 +62,12 @@ public class FlightTokenManagerImpl implements
FlightTokenManager {
private ScheduledExecutorService cleanupExecutor;
public FlightTokenManagerImpl(final int cacheSize, final int
cacheExpiration) {
- // The cache size of all user tokens in Arrow Flight Server. which
will be eliminated by
- // LRU rules after exceeding the limit, the default value is
arrow_flight_max_connections,
- // arrow flight sql is a stateless protocol, the connection is usually
not actively
- // disconnected, bearer token is evict from the cache will unregister
ConnectContext.
+ // The cache size of all user tokens in Arrow Flight Server, which
will be eliminated by
+ // LRU rules after exceeding the limit. The size is the Arrow Flight
SQL sub-quota of the
+ // connection pool (arrow_flight_max_connections, half of
qe_max_connection by default),
+ // capped by arrow_flight_token_cache_size - see
DorisFlightSqlService. Arrow flight sql is a
+ // stateless protocol, the connection is usually not actively
disconnected; a bearer token
+ // evicted from the cache unregisters its ConnectContext.
this.cacheSize = cacheSize;
this.cacheExpiration = cacheExpiration;
@@ -75,10 +79,10 @@ public class FlightTokenManagerImpl implements
FlightTokenManager {
// TODO: broadcast this message to other FE
String token = notification.getKey();
FlightTokenDetails tokenDetails =
notification.getValue();
- ConnectContext context =
ExecuteEnv.getInstance().getScheduler().getFlightSqlConnectPoolMgr()
- .getContextWithFlightToken(token);
+ ConnectContext context =
ExecuteEnv.getInstance().getScheduler()
+ .getContextWithPeerIdentity(token);
if (context != null) {
-
ExecuteEnv.getInstance().getScheduler().getFlightSqlConnectPoolMgr()
+
ExecuteEnv.getInstance().getScheduler().getConnectPoolMgr()
.unregisterConnection(context);
LOG.info("evict bearer token: " +
TokenMasker.tokenId(token) + " from tokenCache, "
+ "reason: " + notification.getCause()
@@ -155,9 +159,14 @@ public class FlightTokenManagerImpl implements
FlightTokenManager {
if (value.getToken().equals("")) {
throw new IllegalArgumentException("invalid bearer token, token
id: " + TokenMasker.tokenId(token)
+ ", try reconnect, bearer token may not be created, or
may have been evict, search for this "
- + "token id in fe.log to see the evict reason. currently
in fe.conf, "
- + "`arrow_flight_max_connections`=" + this.cacheSize
- + ", `arrow_flight_token_alive_time_second`=" +
this.cacheExpiration);
+ + "token id in fe.log to see the evict reason. the token
cache is sized to " + this.cacheSize
+ + " tokens = min(max(1, effective Arrow Flight SQL
sub-quota of the connection pool "
+ +
ConnectPoolMgr.effectiveFlightMaxConnections(Config.qe_max_connection,
+ Config.arrow_flight_max_connections)
+ + " [fe.conf `arrow_flight_max_connections`=" +
Config.arrow_flight_max_connections
+ + " within `qe_max_connection`=" +
Config.qe_max_connection + "]), `arrow_flight_token_cache_size`="
+ + Config.arrow_flight_token_cache_size + "),
`arrow_flight_token_alive_time_second`="
+ + this.cacheExpiration);
}
if (System.currentTimeMillis() >= value.getExpiresAt()) {
tokenCache.invalidate(token);
diff --git a/fe/fe-core/src/main/java/org/apache/doris/metric/MetricRepo.java
b/fe/fe-core/src/main/java/org/apache/doris/metric/MetricRepo.java
index 7c7d5a54e9e..38328622119 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/metric/MetricRepo.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/metric/MetricRepo.java
@@ -457,11 +457,13 @@ public final class MetricRepo {
}
};
DORIS_METRIC_REGISTER.addMetrics(connections);
+ // Arrow Flight SQL sessions are connections of the one pool: they are
counted in
+ // connection_total and held to connection_max, and these two report
their share of it.
GAUGE_ARROW_FLIGHT_CONNECTIONS = new
GaugeMetric<Integer>("arrow_flight_connection_total",
MetricUnit.CONNECTIONS, "total arrow flight connections") {
@Override
public Integer getValue() {
- return
ExecuteEnv.getInstance().getScheduler().getFlightSqlConnectPoolMgr().getConnectionNum();
+ return
ExecuteEnv.getInstance().getScheduler().getConnectPoolMgr().getFlightConnectionNum();
}
};
DORIS_METRIC_REGISTER.addMetrics(GAUGE_ARROW_FLIGHT_CONNECTIONS);
@@ -469,7 +471,7 @@ public final class MetricRepo {
MetricUnit.CONNECTIONS, "max connections") {
@Override
public Integer getValue() {
- return Config.qe_max_connection +
Config.arrow_flight_max_connections;
+ return
ExecuteEnv.getInstance().getScheduler().getConnectPoolMgr().getMaxConnections();
}
};
DORIS_METRIC_REGISTER.addMetrics(GAUGE_CONNECTION_MAX);
@@ -477,7 +479,7 @@ public final class MetricRepo {
MetricUnit.CONNECTIONS, "max arrow flight connections") {
@Override
public Integer getValue() {
- return Config.arrow_flight_max_connections;
+ return
ExecuteEnv.getInstance().getScheduler().getConnectPoolMgr().getFlightMaxConnections();
}
};
DORIS_METRIC_REGISTER.addMetrics(GAUGE_ARROW_FLIGHT_CONNECTION_MAX);
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/mysql/AcceptListener.java
b/fe/fe-core/src/main/java/org/apache/doris/mysql/AcceptListener.java
index 0a796cb9367..b6aae649dab 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/mysql/AcceptListener.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/mysql/AcceptListener.java
@@ -115,10 +115,11 @@ public class AcceptListener implements
ChannelListener<AcceptingChannel<StreamCo
connection.setCloseListener(
streamConnection ->
connectScheduler.getConnectPoolMgr().unregisterConnection(context));
} else {
- long userConnLimit =
context.getEnv().getAuth().getMaxConn(context.getQualifiedUser());
- String errMsg = String.format(
- "Reach limit of connections. Total: %d, User: %d,
Current: %d",
-
connectScheduler.getConnectPoolMgr().getMaxConnections(), userConnLimit, res);
+ String errMsg =
connectScheduler.getConnectPoolMgr().limitReachedMessage(context, res);
+ // The refused client sees the message; the operator finds it
here, since the login
+ // that would have shown the pool's state in SHOW PROCESSLIST
is the one refused.
+ LOG.warn("refused MySQL connection of user {} from {}: {}",
context.getQualifiedUser(),
+ context.getRemoteHostPortString(), errMsg);
context.getState().setError(ErrorCode.ERR_TOO_MANY_USER_CONNECTIONS, errMsg);
MysqlProto.sendResponsePacket(context);
throw new AfterConnectedException(errMsg);
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/mysql/protocol/MysqlProtocolAdapter.java
b/fe/fe-core/src/main/java/org/apache/doris/mysql/protocol/MysqlProtocolAdapter.java
index d7f997ffc3e..5f76d55772f 100644
---
a/fe/fe-core/src/main/java/org/apache/doris/mysql/protocol/MysqlProtocolAdapter.java
+++
b/fe/fe-core/src/main/java/org/apache/doris/mysql/protocol/MysqlProtocolAdapter.java
@@ -33,9 +33,7 @@ import org.apache.doris.nereids.glue.LogicalPlanAdapter;
import org.apache.doris.nereids.stats.StatsErrorEstimator;
import org.apache.doris.qe.ConnectContext;
import org.apache.doris.qe.ConnectContext.ConnectType;
-import org.apache.doris.qe.ConnectPoolMgr;
import org.apache.doris.qe.ConnectProcessor;
-import org.apache.doris.qe.ConnectScheduler;
import org.apache.doris.qe.QueryState.MysqlStateType;
import org.apache.doris.qe.ShowResultSet;
import org.apache.doris.qe.StmtExecutor;
@@ -216,8 +214,8 @@ public class MysqlProtocolAdapter implements
ProtocolAdapter {
}
@Override
- public ConnectPoolMgr connectPool(ConnectScheduler scheduler) {
- return scheduler.getConnectPoolMgr();
+ public void releaseSession(ConnectContext ctx) {
+ // Nothing is held for a MySQL session beyond its channel, which
closeConnection closes.
}
/**
diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java
b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java
index 8d96b1b2f7a..dbce88cd992 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectContext.java
@@ -1032,13 +1032,14 @@ public class ConnectContext {
}
}
- // The session is over (see FlightProtocolAdapter.tearDown): its deferred
executors are
- // finalized, and it keeps none and runs no command from now on. Nothing
to do for a
- // connection of any other protocol.
- public void tearDownFlightSqlSession() {
- if (protocolAdapter instanceof FlightProtocolAdapter) {
- ((FlightProtocolAdapter) protocolAdapter).tearDown();
- }
+ /**
+ * Releases what the protocol still holds for this session, when the
connection leaves the pool
+ * (see {@link ConnectPoolMgr#unregisterConnection}): for Arrow Flight SQL
the channel-cached
+ * results and the deferred executors, after which the session keeps none
and runs no command
+ * (see {@code FlightProtocolAdapter.tearDown}); nothing for a MySQL
connection. Idempotent.
+ */
+ public void releaseProtocolSession() {
+ protocolAdapter.releaseSession(this);
}
// A snapshot; empty for a connection of any other protocol.
@@ -1182,7 +1183,7 @@ public class ConnectContext {
}
this.queryId = queryId;
if (connectScheduler != null && !Strings.isNullOrEmpty(traceId)) {
-
protocolAdapter.connectPool(connectScheduler).putTraceId2QueryId(traceId,
queryId);
+ connectScheduler.getConnectPoolMgr().putTraceId2QueryId(traceId,
queryId);
}
}
diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectPoolMgr.java
b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectPoolMgr.java
index 643908e52c1..e4e614bdf41 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectPoolMgr.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectPoolMgr.java
@@ -22,6 +22,7 @@ import org.apache.doris.catalog.Env;
import org.apache.doris.common.Status;
import org.apache.doris.common.util.DebugUtil;
import org.apache.doris.mysql.privilege.PrivPredicate;
+import org.apache.doris.qe.ConnectContext.ConnectType;
import org.apache.doris.qe.ConnectContext.ThreadInfo;
import org.apache.doris.thrift.TUniqueId;
@@ -36,21 +37,66 @@ import java.util.Map;
import java.util.Optional;
import java.util.concurrent.atomic.AtomicInteger;
+/**
+ * The one pool of every connection this frontend serves, whatever protocol it
speaks: MySQL
+ * connections and Arrow Flight SQL sessions share {@code qe_max_connection},
the per-user
+ * {@code max_user_connections}, the processlist, KILL, the timeout checker
and the connection
+ * metrics. Arrow Flight SQL sessions additionally count against their own
sub-quota
+ * ({@code arrow_flight_max_connections}, half of the pool's limit unless set)
and are indexed by
+ * their peer identity, the bearer token, since that is how Flight requests
name their session.
+ *
+ * <p>{@link #unregisterConnection} is where every teardown path of a
connection meets - a MySQL
+ * channel closing, a Flight bearer token expiring or being evicted,
CloseSession, a KILL CONNECTION
+ * from another connection, the timeout checker past wait_timeout - so it is
where the protocol
+ * releases what it still holds for the session.
+ */
public class ConnectPoolMgr {
private static final Logger LOG =
LogManager.getLogger(ConnectPoolMgr.class);
protected final int maxConnections;
+ private final int flightMaxConnections;
protected final AtomicInteger numberConnection;
+ private final AtomicInteger numberFlightConnection = new AtomicInteger(0);
protected final Map<Integer, ConnectContext> connectionMap =
Maps.newConcurrentMap();
protected final Map<String, AtomicInteger> connByUser =
Maps.newConcurrentMap();
+ // Arrow Flight SQL: peer identity (bearer token) -> connection id
+ private final Map<String, Integer> peerIdentity2ConnectionId =
Maps.newConcurrentMap();
+ // Serializes the check-and-reserve of the pool/Flight/user quotas in
registerConnection, so a
+ // doomed attempt cannot hold one counter's reservation across another
attempt's checks.
+ private final Object admissionLock = new Object();
// valid trace id -> query id
protected final Map<String, TUniqueId> traceId2QueryId =
Maps.newConcurrentMap();
public ConnectPoolMgr(int maxConnections) {
+ this(maxConnections, -1);
+ }
+
+ /**
+ * @param maxConnections the pool's limit, {@code qe_max_connection}
+ * @param flightMaxConnections the Arrow Flight SQL sub-quota, {@code
arrow_flight_max_connections};
+ * negative is half of {@code maxConnections}
+ */
+ public ConnectPoolMgr(int maxConnections, int flightMaxConnections) {
this.maxConnections = maxConnections;
+ this.flightMaxConnections =
effectiveFlightMaxConnections(maxConnections, flightMaxConnections);
numberConnection = new AtomicInteger(0);
}
+ /**
+ * The Arrow Flight SQL sub-quota as enforced: a negative setting (the
default) is half of the
+ * pool's limit, so that Flight sessions - which their clients mostly
never close, and which
+ * therefore stay until a timeout - cannot take the half MySQL clients
connect through; an
+ * explicit one can never exceed the pool's limit, since every Flight
session is a connection of
+ * the pool too.
+ */
+ public static int effectiveFlightMaxConnections(int maxConnections, int
flightMaxConnections) {
+ return flightMaxConnections < 0 ? maxConnections / 2 :
Math.min(maxConnections, flightMaxConnections);
+ }
+
+ private static boolean isFlight(ConnectContext ctx) {
+ return ctx.getConnectType() == ConnectType.ARROW_FLIGHT_SQL;
+ }
+
public void timeoutChecker(long now) {
for (ConnectContext connectContext : connectionMap.values()) {
try {
@@ -66,23 +112,64 @@ public class ConnectPoolMgr {
// Return -1 means register OK
// Return >=0 means register failed, and return value is current
connection num.
public int registerConnection(ConnectContext ctx) {
- if (numberConnection.incrementAndGet() > maxConnections) {
- numberConnection.decrementAndGet();
- return numberConnection.get();
- }
- // Check user
- connByUser.putIfAbsent(ctx.getQualifiedUser(), new AtomicInteger(0));
- AtomicInteger conns = connByUser.get(ctx.getQualifiedUser());
- if (conns.incrementAndGet() >
ctx.getEnv().getAuth().getMaxConn(ctx.getQualifiedUser())) {
- conns.decrementAndGet();
- numberConnection.decrementAndGet();
- return numberConnection.get();
+ boolean flight = isFlight(ctx);
+ String qualifiedUser = ctx.getQualifiedUser();
+ long userLimit = ctx.getEnv().getAuth().getMaxConn(qualifiedUser);
+ // Check and reserve the pool, Flight sub-quota and per-user quotas as
one transition. Doing it
+ // under a lock (rather than the earlier increment-check-rollback per
counter) keeps a doomed
+ // attempt from reserving one counter across another attempt's checks:
a second Flight attempt
+ // that will fail its sub-quota could otherwise briefly hold the last
pool slot and make a
+ // concurrent MySQL connection - which fits in every serial ordering -
be refused. Registration
+ // is once per connection, not per query, so this is not a hot path.
+ synchronized (admissionLock) {
+ if (numberConnection.get() >= maxConnections) {
+ return numberConnection.get();
+ }
+ if (flight && numberFlightConnection.get() >=
flightMaxConnections) {
+ return numberConnection.get();
+ }
+ AtomicInteger conns = connByUser.computeIfAbsent(qualifiedUser, u
-> new AtomicInteger(0));
+ if (conns.get() >= userLimit) {
+ return numberConnection.get();
+ }
+ numberConnection.incrementAndGet();
+ if (flight) {
+ numberFlightConnection.incrementAndGet();
+ }
+ conns.incrementAndGet();
+ connectionMap.put(ctx.getConnectionId(), ctx);
+ if (flight) {
+ peerIdentity2ConnectionId.put(ctx.getPeerIdentity(),
ctx.getConnectionId());
+ }
}
- connectionMap.put(ctx.getConnectionId(), ctx);
return -1;
}
+ /**
+ * The refusal a client is told when {@link #registerConnection} returned
a count instead of -1:
+ * the same sentence for every protocol, naming the limits a connection is
held to. The Flight
+ * sub-quota and its usage are named whenever Flight sessions hold part of
the pool - a MySQL
+ * client refused by a pool that Flight sessions filled is told where the
connections went -
+ * and to a Flight client also when the sub-quota is tighter than the
pool's limit, since then
+ * it can be the one that was reached.
+ */
+ public String limitReachedMessage(ConnectContext ctx, int current) {
+ long userLimit =
ctx.getEnv().getAuth().getMaxConn(ctx.getQualifiedUser());
+ String message = String.format("Reach limit of connections. Total: %d,
User: %d, Current: %d",
+ maxConnections, userLimit, current);
+ int flightCurrent = numberFlightConnection.get();
+ if (flightCurrent > 0 || (isFlight(ctx) && flightMaxConnections <
maxConnections)) {
+ message += String.format(", Arrow Flight SQL: %d (current: %d)",
flightMaxConnections, flightCurrent);
+ }
+ return message;
+ }
+
public void unregisterConnection(ConnectContext ctx) {
+ // First, before any bookkeeping that could fail: what the protocol
holds for the session -
+ // for Arrow Flight SQL the channel-cached results and the deferred
query coordinators
+ // (see FlightProtocolAdapter.tearDown) - is released whether or not
the connection is
+ // still in the pool, so an abandoned session is cleaned up rather
than leaked.
+ ctx.releaseProtocolSession();
ctx.closeTxn();
if (connectionMap.remove(ctx.getConnectionId()) != null) {
AtomicInteger conns = connByUser.get(ctx.getQualifiedUser());
@@ -92,6 +179,12 @@ public class ConnectPoolMgr {
if (ctx.traceId() != null) {
traceId2QueryId.remove(ctx.traceId());
}
+ if (isFlight(ctx)) {
+ // Only this connection's own entry: a second session created
under the same token by
+ // a concurrent first request must not lose its index to the
first one's teardown.
+ peerIdentity2ConnectionId.remove(ctx.getPeerIdentity(),
ctx.getConnectionId());
+ numberFlightConnection.decrementAndGet();
+ }
numberConnection.decrementAndGet();
}
}
@@ -100,6 +193,12 @@ public class ConnectPoolMgr {
return connectionMap.get(connectionId);
}
+ /** The Arrow Flight SQL session registered under this peer identity
(bearer token), or null. */
+ public ConnectContext getContextWithPeerIdentity(String peerIdentity) {
+ Integer connectionId = peerIdentity2ConnectionId.get(peerIdentity);
+ return connectionId == null ? null : getContext(connectionId);
+ }
+
public ConnectContext getContextWithQueryId(String queryId) {
for (ConnectContext context : connectionMap.values()) {
if (queryId.equals(DebugUtil.printId(context.queryId)) ||
queryId.equals(context.traceId())) {
@@ -124,6 +223,11 @@ public class ConnectPoolMgr {
return numberConnection.get();
}
+ /** How many of the pool's connections are Arrow Flight SQL sessions. */
+ public int getFlightConnectionNum() {
+ return numberFlightConnection.get();
+ }
+
public List<ThreadInfo> listConnection(String user, boolean isFull) {
List<ConnectContext.ThreadInfo> infos = Lists.newArrayList();
for (ConnectContext ctx : connectionMap.values()) {
@@ -180,4 +284,9 @@ public class ConnectPoolMgr {
public int getMaxConnections() {
return maxConnections;
}
+
+ /** The Arrow Flight SQL sub-quota as enforced (see {@link
#effectiveFlightMaxConnections}). */
+ public int getFlightMaxConnections() {
+ return flightMaxConnections;
+ }
}
diff --git a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectScheduler.java
b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectScheduler.java
index 06e544971f6..7a2b99dcd4c 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectScheduler.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/qe/ConnectScheduler.java
@@ -18,19 +18,14 @@
package org.apache.doris.qe;
import org.apache.doris.analysis.UserIdentity;
-import org.apache.doris.arrowflight.sessions.FlightSqlConnectPoolMgr;
import org.apache.doris.common.Config;
import org.apache.doris.common.Status;
import org.apache.doris.common.ThreadPoolManager;
import org.apache.doris.qe.ConnectContext.ThreadInfo;
-import com.google.common.base.Strings;
-import com.google.common.collect.Lists;
-import com.google.common.collect.Maps;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
-import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Optional;
@@ -45,8 +40,8 @@ import java.util.concurrent.atomic.AtomicInteger;
public class ConnectScheduler {
private static final Logger LOG =
LogManager.getLogger(ConnectScheduler.class);
private final AtomicInteger nextConnectionId;
+ // Every connection of every protocol: see ConnectPoolMgr.
private final ConnectPoolMgr connectPoolMgr;
- private final FlightSqlConnectPoolMgr flightSqlConnectPoolMgr;
// Use a thread to check whether connection is timeout. Because
// 1. If use a scheduler, the task maybe a huge number when query is messy.
@@ -55,25 +50,25 @@ public class ConnectScheduler {
private final ScheduledExecutorService checkTimer =
ThreadPoolManager.newDaemonScheduledThreadPool(1,
"connect-scheduler-check-timer", true);
- public ConnectScheduler(int commonMaxConnections, int
flightSqlMaxConnections) {
+ /**
+ * @param maxConnections the pool's limit, {@code qe_max_connection}
+ * @param flightMaxConnections the Arrow Flight SQL sub-quota, {@code
arrow_flight_max_connections};
+ * negative is half of {@code maxConnections}
+ */
+ public ConnectScheduler(int maxConnections, int flightMaxConnections) {
nextConnectionId = new AtomicInteger(0);
- this.connectPoolMgr = new ConnectPoolMgr(commonMaxConnections);
- this.flightSqlConnectPoolMgr = new
FlightSqlConnectPoolMgr(flightSqlMaxConnections);
+ this.connectPoolMgr = new ConnectPoolMgr(maxConnections,
flightMaxConnections);
checkTimer.scheduleAtFixedRate(new TimeoutChecker(), 0, 1000L,
TimeUnit.MILLISECONDS);
}
- public ConnectScheduler(int commonMaxConnections) {
- this(commonMaxConnections, Config.arrow_flight_max_connections);
+ public ConnectScheduler(int maxConnections) {
+ this(maxConnections, Config.arrow_flight_max_connections);
}
public ConnectPoolMgr getConnectPoolMgr() {
return connectPoolMgr;
}
- public FlightSqlConnectPoolMgr getFlightSqlConnectPoolMgr() {
- return flightSqlConnectPoolMgr;
- }
-
// submit one MysqlContext to this scheduler.
// return true, if this connection has been successfully submitted,
otherwise return false.
// Caller should close ConnectContext if return false.
@@ -87,83 +82,57 @@ public class ConnectScheduler {
}
public ConnectContext getContext(int connectionId) {
- ConnectContext ctx = connectPoolMgr.getContext(connectionId);
- if (ctx == null) {
- ctx = flightSqlConnectPoolMgr.getContext(connectionId);
- }
- return ctx;
+ return connectPoolMgr.getContext(connectionId);
+ }
+
+ /** The Arrow Flight SQL session registered under this peer identity
(bearer token), or null. */
+ public ConnectContext getContextWithPeerIdentity(String peerIdentity) {
+ return connectPoolMgr.getContextWithPeerIdentity(peerIdentity);
}
public ConnectContext getContextWithQueryId(String queryId) {
- ConnectContext ctx = connectPoolMgr.getContextWithQueryId(queryId);
- if (ctx == null) {
- ctx = flightSqlConnectPoolMgr.getContextWithQueryId(queryId);
- }
- return ctx;
+ return connectPoolMgr.getContextWithQueryId(queryId);
}
public boolean cancelQuery(String queryId, Status cancelReason) {
- boolean ret = connectPoolMgr.cancelQuery(queryId, cancelReason);
- if (!ret) {
- ret = flightSqlConnectPoolMgr.cancelQuery(queryId, cancelReason);
- }
- return ret;
+ return connectPoolMgr.cancelQuery(queryId, cancelReason);
}
public int getConnectionNum() {
- return connectPoolMgr.getConnectionNum() +
flightSqlConnectPoolMgr.getConnectionNum();
+ return connectPoolMgr.getConnectionNum();
}
public List<ThreadInfo> listConnection(String user, boolean isFull) {
- List<ConnectContext.ThreadInfo> infos = Lists.newArrayList();
- infos.addAll(connectPoolMgr.listConnection(user, isFull));
- infos.addAll(flightSqlConnectPoolMgr.listConnection(user, isFull));
- return infos;
+ return connectPoolMgr.listConnection(user, isFull);
}
// used for thrift
public List<List<String>> listConnectionForRpc(UserIdentity userIdentity,
boolean isShowFullSql,
Optional<String> timeZone) {
- List<List<String>> list = new ArrayList<>();
- list.addAll(connectPoolMgr.listConnectionForRpc(userIdentity,
isShowFullSql, timeZone));
- list.addAll(flightSqlConnectPoolMgr.listConnectionForRpc(userIdentity,
isShowFullSql, timeZone));
- return list;
+ return connectPoolMgr.listConnectionForRpc(userIdentity,
isShowFullSql, timeZone);
}
public String getQueryIdByTraceId(String traceId) {
- String queryId = connectPoolMgr.getQueryIdByTraceId(traceId);
- if (Strings.isNullOrEmpty(queryId)) {
- queryId = flightSqlConnectPoolMgr.getQueryIdByTraceId(traceId);
- }
- return queryId;
+ return connectPoolMgr.getQueryIdByTraceId(traceId);
}
public void removeOldTraceId(String traceId) {
connectPoolMgr.removeTraceId(traceId);
- flightSqlConnectPoolMgr.removeTraceId(traceId);
}
public Map<Integer, ConnectContext> getConnectionMap() {
- Map<Integer, ConnectContext> map = Maps.newConcurrentMap();
- map.putAll(connectPoolMgr.getConnectionMap());
- map.putAll(flightSqlConnectPoolMgr.getConnectionMap());
- return map;
+ return connectPoolMgr.getConnectionMap();
}
public Map<String, AtomicInteger> getUserConnectionMap() {
- Map<String, AtomicInteger> map = Maps.newConcurrentMap();
- map.putAll(connectPoolMgr.getUserConnectionMap());
- map.putAll(flightSqlConnectPoolMgr.getUserConnectionMap());
- return map;
+ return connectPoolMgr.getUserConnectionMap();
}
private class TimeoutChecker extends TimerTask {
@Override
public void run() {
try {
- long now = System.currentTimeMillis();
- connectPoolMgr.timeoutChecker(now);
- flightSqlConnectPoolMgr.timeoutChecker(now);
+ connectPoolMgr.timeoutChecker(System.currentTimeMillis());
} catch (Throwable t) {
LOG.warn("failed to check connection timeout", t);
}
diff --git
a/fe/fe-core/src/main/java/org/apache/doris/qe/protocol/ProtocolAdapter.java
b/fe/fe-core/src/main/java/org/apache/doris/qe/protocol/ProtocolAdapter.java
index 94c418ee5a2..0d77ecd1a1b 100644
--- a/fe/fe-core/src/main/java/org/apache/doris/qe/protocol/ProtocolAdapter.java
+++ b/fe/fe-core/src/main/java/org/apache/doris/qe/protocol/ProtocolAdapter.java
@@ -19,8 +19,6 @@ package org.apache.doris.qe.protocol;
import org.apache.doris.qe.ConnectContext;
import org.apache.doris.qe.ConnectContext.ConnectType;
-import org.apache.doris.qe.ConnectPoolMgr;
-import org.apache.doris.qe.ConnectScheduler;
import org.apache.doris.qe.StmtExecutor;
import org.apache.doris.thrift.TMasterOpRequest;
import org.apache.doris.thrift.TResultSinkType;
@@ -148,10 +146,13 @@ public interface ProtocolAdapter {
throws IOException;
/**
- * The pool this connection is registered in. Each protocol still keeps
its own pool; this
- * goes away when they are merged.
+ * Releases what the protocol still holds for the session when the
connection leaves the pool
+ * ({@code ConnectPoolMgr.unregisterConnection}, where every teardown path
of a connection
+ * meets): for Arrow Flight SQL the channel-cached results and the
deferred query coordinators,
+ * nothing for MySQL, whose channel {@link #closeConnection} closes. Must
be idempotent, and
+ * must not depend on the connection being in the pool.
*/
- ConnectPoolMgr connectPool(ConnectScheduler scheduler);
+ void releaseSession(ConnectContext ctx);
/**
* Called from {@link ConnectContext#clear()} once the response of a
statement has been sent,
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/arrowflight/DorisFlightSqlProducerTest.java
b/fe/fe-core/src/test/java/org/apache/doris/arrowflight/DorisFlightSqlProducerTest.java
index 6fb16d62fbf..59bcec4f768 100644
---
a/fe/fe-core/src/test/java/org/apache/doris/arrowflight/DorisFlightSqlProducerTest.java
+++
b/fe/fe-core/src/test/java/org/apache/doris/arrowflight/DorisFlightSqlProducerTest.java
@@ -17,6 +17,7 @@
package org.apache.doris.arrowflight;
+import org.apache.doris.analysis.UserIdentity;
import org.apache.doris.arrowflight.protocol.FlightProtocolAdapter;
import org.apache.doris.arrowflight.results.FlightSqlChannel;
import org.apache.doris.arrowflight.sessions.FlightSessionsManager;
@@ -25,6 +26,7 @@ import org.apache.doris.common.FeConstants;
import org.apache.doris.common.IncrWindowNotReadyException;
import org.apache.doris.mysql.MysqlCommand;
import org.apache.doris.qe.ConnectContext;
+import org.apache.doris.qe.ConnectPoolTestSupport;
import org.apache.doris.qe.ConnectScheduler;
import org.apache.doris.qe.QueryState;
import org.apache.doris.qe.StmtExecutor;
@@ -434,11 +436,12 @@ public class DorisFlightSqlProducerTest {
// statement runs, all start wait_timeout over.
@Test
public void testSessionOptionActionsKeepTheSessionAlive() throws Exception
{
- ConnectContext ctx = ConnectContext.forFlight("token");
+ ConnectContext ctx =
ConnectPoolTestSupport.flightSession(ConnectPoolTestSupport.envAllowing(100),
+ UserIdentity.ROOT, "token");
ConnectScheduler scheduler = new ConnectScheduler(10, 10);
ctx.setConnectScheduler(scheduler);
scheduler.submit(ctx);
- Assertions.assertEquals(-1,
scheduler.getFlightSqlConnectPoolMgr().registerConnection(ctx));
+ Assertions.assertEquals(-1,
scheduler.getConnectPoolMgr().registerConnection(ctx));
ctx.setCommand(MysqlCommand.COM_SLEEP);
ctx.setStartTime();
long waitTimeoutMs = ctx.getSessionVariable().getWaitTimeoutS() *
1000L;
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/arrowflight/DorisFlightSqlServiceTest.java
b/fe/fe-core/src/test/java/org/apache/doris/arrowflight/DorisFlightSqlServiceTest.java
new file mode 100644
index 00000000000..ea615cfe391
--- /dev/null
+++
b/fe/fe-core/src/test/java/org/apache/doris/arrowflight/DorisFlightSqlServiceTest.java
@@ -0,0 +1,42 @@
+// 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.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class DorisFlightSqlServiceTest {
+
+ // The token cache size floors the effective Flight sub-quota at 1, but
never the configured cap:
+ // a legal sub-quota of 0 (qe_max_connection = 1) keeps one token so the
first request reaches the
+ // pool and is refused with RESOURCE_EXHAUSTED, while an illegal
arrow_flight_token_cache_size (<= 0)
+ // is passed through so the base's loud failure (Guava rejects a negative
maximumSize; 0 evicts every
+ // token) is preserved rather than the FE silently running on a one-token
cache.
+ @Test
+ public void testEffectiveTokenCacheSizeFloorsOnlyTheSubQuota() {
+ // Legal sub-quota of 0 -> floored to 1.
+ Assertions.assertEquals(1,
DorisFlightSqlService.effectiveTokenCacheSize(0, 4096));
+ // Normal: the smaller of the sub-quota and the cap.
+ Assertions.assertEquals(512,
DorisFlightSqlService.effectiveTokenCacheSize(512, 4096));
+ Assertions.assertEquals(4096,
DorisFlightSqlService.effectiveTokenCacheSize(8192, 4096));
+ Assertions.assertEquals(10,
DorisFlightSqlService.effectiveTokenCacheSize(512, 10));
+ // An illegal cap is not floored: it stays <= 0 for the base's loud
failure.
+ Assertions.assertEquals(0,
DorisFlightSqlService.effectiveTokenCacheSize(0, 0));
+ Assertions.assertEquals(-1,
DorisFlightSqlService.effectiveTokenCacheSize(512, -1));
+ }
+}
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/arrowflight/protocol/FlightProtocolAdapterTest.java
b/fe/fe-core/src/test/java/org/apache/doris/arrowflight/protocol/FlightProtocolAdapterTest.java
index 02630b9c22e..0028dcd4d3f 100644
---
a/fe/fe-core/src/test/java/org/apache/doris/arrowflight/protocol/FlightProtocolAdapterTest.java
+++
b/fe/fe-core/src/test/java/org/apache/doris/arrowflight/protocol/FlightProtocolAdapterTest.java
@@ -17,6 +17,7 @@
package org.apache.doris.arrowflight.protocol;
+import org.apache.doris.analysis.UserIdentity;
import org.apache.doris.arrowflight.results.FlightSqlEndpointsLocation;
import org.apache.doris.catalog.Column;
import org.apache.doris.catalog.ScalarType;
@@ -27,6 +28,7 @@ import org.apache.doris.common.util.DebugUtil;
import org.apache.doris.mysql.protocol.MysqlProtocolAdapter;
import org.apache.doris.qe.ConnectContext;
import org.apache.doris.qe.ConnectContext.ConnectType;
+import org.apache.doris.qe.ConnectPoolTestSupport;
import org.apache.doris.qe.ConnectScheduler;
import org.apache.doris.qe.QueryState.MysqlStateType;
import org.apache.doris.qe.ShowResultSet;
@@ -63,8 +65,9 @@ import java.util.concurrent.TimeUnit;
/**
* An Arrow Flight SQL session is a ConnectContext bound to a
FlightProtocolAdapter. The adapter
- * owns what only that protocol has (result cache, endpoints, deferred
executors, the pool the
- * session is registered in) and serializes the session's commands, which gRPC
does not do.
+ * owns what only that protocol has (result cache, endpoints, deferred
executors) and serializes
+ * the session's commands, which gRPC does not do; the session is registered
in the one connection
+ * pool every protocol shares, which asks the adapter to release those on
teardown.
*/
public class FlightProtocolAdapterTest {
private boolean savedRunningUnitTest;
@@ -83,7 +86,10 @@ public class FlightProtocolAdapterTest {
}
private static ConnectContext flightSession() {
- return ConnectContext.forFlight("test-peer-identity");
+ // Registrable in the one pool: the pool files a session under its
user and asks its Env
+ // for the user's connection limit.
+ return
ConnectPoolTestSupport.flightSession(ConnectPoolTestSupport.envAllowing(100),
UserIdentity.ROOT,
+ "test-peer-identity");
}
// A command that blocks on a latch, run from a plain thread.
@@ -129,7 +135,7 @@ public class FlightProtocolAdapterTest {
}
@Test
- public void testSessionRegistersItsTraceIdInTheFlightPool() {
+ public void testSessionRegistersItsTraceIdInThePool() {
ConnectScheduler scheduler = new ConnectScheduler(10, 10);
ConnectContext ctx = flightSession();
ctx.setConnectScheduler(scheduler);
@@ -138,26 +144,52 @@ public class FlightProtocolAdapterTest {
ctx.setQueryId(queryId);
- Assertions.assertEquals(DebugUtil.printId(queryId),
-
scheduler.getFlightSqlConnectPoolMgr().getQueryIdByTraceId("trace-1"));
- Assertions.assertEquals("",
scheduler.getConnectPoolMgr().getQueryIdByTraceId("trace-1"));
+ Assertions.assertEquals(DebugUtil.printId(queryId),
scheduler.getQueryIdByTraceId("trace-1"));
}
@Test
- public void testKillUnregistersTheSessionFromTheFlightPool() {
+ public void testKillUnregistersTheSessionFromThePool() {
ConnectScheduler scheduler = new ConnectScheduler(10, 10);
ConnectContext ctx = flightSession();
ctx.setConnectScheduler(scheduler);
scheduler.submit(ctx);
- Assertions.assertEquals(-1,
scheduler.getFlightSqlConnectPoolMgr().registerConnection(ctx));
+ Assertions.assertEquals(-1,
scheduler.getConnectPoolMgr().registerConnection(ctx));
Assertions.assertSame(ctx,
scheduler.getContext(ctx.getConnectionId()));
+ Assertions.assertSame(ctx,
scheduler.getContextWithPeerIdentity(ctx.getPeerIdentity()));
ctx.kill(true);
Assertions.assertTrue(ctx.isKilled());
Assertions.assertNull(scheduler.getContext(ctx.getConnectionId()));
+
Assertions.assertNull(scheduler.getContextWithPeerIdentity(ctx.getPeerIdentity()));
}
+ // Every Flight session teardown path meets in the pool's
unregisterConnection, which asks the
+ // protocol to release what it holds: the channel-cached results first,
then the deferred
+ // executors (tearDown), whether or not the session was ever registered.
+ @Test
+ public void testReleaseSessionClosesTheChannelAndTearsTheSessionDown() {
+ ConnectContext ctx = flightSession();
+ FlightProtocolAdapter adapter = FlightProtocolAdapter.of(ctx);
+ StmtExecutor deferred = Mockito.mock(StmtExecutor.class);
+ ctx.addFlightSqlDeferredExecutor(deferred);
+ // A result the client never pulled: off-heap Arrow buffers the
channel holds.
+ adapter.getChannel().addOKResult("query-1", "SELECT 1");
+ Assertions.assertEquals(1, adapter.getChannel().resultNum());
+ Assertions.assertTrue(adapter.getChannel().getAllocatedMemory() > 0);
+
+ ctx.releaseProtocolSession();
+
+ // The channel's results are released with their buffers, and the
deferred query finalized.
+ Assertions.assertEquals(0, adapter.getChannel().resultNum());
+ Assertions.assertEquals(0, adapter.getChannel().getAllocatedMemory());
+ Assertions.assertTrue(ctx.getFlightSqlDeferredExecutors().isEmpty());
+ // Idempotent: the second release finds nothing to do, throws nothing,
and finalizes nothing twice.
+ Assertions.assertDoesNotThrow(ctx::releaseProtocolSession);
+ Mockito.verify(deferred, Mockito.times(1)).finalizeArrowFlightQuery();
+ }
+
+
@Test
public void testCommandsOfOneSessionRunOneAtATime() throws Exception {
ConnectContext ctx = flightSession();
@@ -468,7 +500,7 @@ public class FlightProtocolAdapterTest {
ConnectContext ctx = flightSession();
ctx.setConnectScheduler(scheduler);
scheduler.submit(ctx);
- Assertions.assertEquals(-1,
scheduler.getFlightSqlConnectPoolMgr().registerConnection(ctx));
+ Assertions.assertEquals(-1,
scheduler.getConnectPoolMgr().registerConnection(ctx));
FlightProtocolAdapter adapter = FlightProtocolAdapter.of(ctx);
StmtExecutor before = Mockito.mock(StmtExecutor.class);
StmtExecutor late = Mockito.mock(StmtExecutor.class);
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/arrowflight/sessions/FlightSqlConnectPoolMgrTest.java
b/fe/fe-core/src/test/java/org/apache/doris/arrowflight/sessions/FlightSqlConnectPoolMgrTest.java
deleted file mode 100644
index 3dd5b94b4f1..00000000000
---
a/fe/fe-core/src/test/java/org/apache/doris/arrowflight/sessions/FlightSqlConnectPoolMgrTest.java
+++ /dev/null
@@ -1,70 +0,0 @@
-// 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.sessions;
-
-import org.apache.doris.arrowflight.results.FlightSqlChannel;
-import org.apache.doris.qe.ConnectContext;
-
-import org.junit.jupiter.api.Assertions;
-import org.junit.jupiter.api.Test;
-import org.mockito.Mockito;
-
-public class FlightSqlConnectPoolMgrTest {
-
- // Arrow Flight SQL keeps a query's coordinator alive across GetFlightInfo
-> DoGet (see #62259).
- // unregisterConnection() is the catch-all teardown path: idle/query
timeout, bearer token expiry
- // and explicit CloseSession all reach here. It must tear the Flight
session down -- finalize the
- // deferred coordinators, so an abandoned connection cannot leak them (the
external-table batch
- // SplitSource and the query queue slot the coordinator holds), and close
the session to what a
- // still-running command defers afterwards.
- @Test
- public void testUnregisterConnectionFinalizesDeferredExecutors() {
- FlightSqlConnectPoolMgr poolMgr = new FlightSqlConnectPoolMgr(100);
- ConnectContext ctx = Mockito.mock(ConnectContext.class);
- FlightSqlChannel channel = Mockito.mock(FlightSqlChannel.class);
- Mockito.when(ctx.getFlightSqlChannel()).thenReturn(channel);
-
- poolMgr.unregisterConnection(ctx);
-
- // The deferred coordinators must be released on teardown even though
this connection was
- // never registered in the pool (an abandoned connection is still
cleaned up, not leaked).
- Mockito.verify(channel).close();
- Mockito.verify(ctx).tearDownFlightSqlSession();
- }
-
- // Cleanup must run before the connection bookkeeping (closeTxn / map
removal), so that a failure
- // there cannot strand the deferred coordinators. Verify the deferred
executors are finalized
- // even when the context is the one stored in the pool.
- @Test
- public void testUnregisterRegisteredConnectionFinalizesDeferredExecutors()
{
- FlightSqlConnectPoolMgr poolMgr = new FlightSqlConnectPoolMgr(100);
- ConnectContext ctx = Mockito.mock(ConnectContext.class);
- FlightSqlChannel channel = Mockito.mock(FlightSqlChannel.class);
- Mockito.when(ctx.getFlightSqlChannel()).thenReturn(channel);
- Mockito.when(ctx.getConnectionId()).thenReturn(7);
-
Mockito.when(ctx.getConnectType()).thenReturn(ConnectContext.ConnectType.ARROW_FLIGHT_SQL);
- Mockito.when(ctx.getPeerIdentity()).thenReturn("token-7");
- poolMgr.getConnectionMap().put(7, ctx);
-
- poolMgr.unregisterConnection(ctx);
-
- Mockito.verify(channel).close();
- Mockito.verify(ctx).tearDownFlightSqlSession();
- Assertions.assertNull(poolMgr.getConnectionMap().get(7));
- }
-}
diff --git a/fe/fe-core/src/test/java/org/apache/doris/metric/MetricsTest.java
b/fe/fe-core/src/test/java/org/apache/doris/metric/MetricsTest.java
index ddb71021794..448bc5430ab 100644
--- a/fe/fe-core/src/test/java/org/apache/doris/metric/MetricsTest.java
+++ b/fe/fe-core/src/test/java/org/apache/doris/metric/MetricsTest.java
@@ -17,6 +17,7 @@
package org.apache.doris.metric;
+import org.apache.doris.analysis.UserIdentity;
import org.apache.doris.catalog.Env;
import org.apache.doris.cloud.CloudWarmUpJob;
import org.apache.doris.cloud.JobWarmUpStats;
@@ -37,6 +38,10 @@ import org.apache.doris.monitor.jvm.JvmService;
import org.apache.doris.monitor.jvm.JvmStats;
import org.apache.doris.mysql.privilege.Auth;
import org.apache.doris.mysql.privilege.UserProperty;
+import org.apache.doris.qe.ConnectContext;
+import org.apache.doris.qe.ConnectPoolTestSupport;
+import org.apache.doris.qe.ConnectScheduler;
+import org.apache.doris.service.ExecuteEnv;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.google.common.collect.Lists;
@@ -44,6 +49,8 @@ import lombok.extern.slf4j.Slf4j;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
import java.lang.management.GarbageCollectorMXBean;
import java.lang.management.ManagementFactory;
@@ -86,22 +93,34 @@ public class MetricsTest {
@Test
public void testConnectionMaxMetrics() throws Exception {
- int originQeMaxConnection = Config.qe_max_connection;
- int originArrowFlightMaxConnections =
Config.arrow_flight_max_connections;
- try {
- Config.qe_max_connection = 4321;
- Config.arrow_flight_max_connections = 8765;
+ // One pool for every protocol: connection_max is its limit,
connection_total every
+ // connection in it, and the two arrow_flight gauges the Flight share
of each. A scheduler
+ // of known numbers, with one connection of each protocol registered,
so that each gauge is
+ // pinned to a literal and not to the getter it reads.
+ ConnectScheduler scheduler = new ConnectScheduler(4321, 300);
+ Env env = ConnectPoolTestSupport.envAllowing(100);
+ ConnectContext mysql = ConnectPoolTestSupport.mysqlConnection(env,
UserIdentity.ROOT);
+ ConnectContext flight = ConnectPoolTestSupport.flightSession(env,
UserIdentity.ROOT, "metric-token");
+ scheduler.submit(mysql);
+ scheduler.submit(flight);
+ Assertions.assertEquals(-1,
scheduler.getConnectPoolMgr().registerConnection(mysql));
+ Assertions.assertEquals(-1,
scheduler.getConnectPoolMgr().registerConnection(flight));
+ ExecuteEnv executeEnv = Mockito.mock(ExecuteEnv.class);
+ Mockito.when(executeEnv.getScheduler()).thenReturn(scheduler);
+ try (MockedStatic<ExecuteEnv> mockedExecuteEnv =
Mockito.mockStatic(ExecuteEnv.class)) {
+
mockedExecuteEnv.when(ExecuteEnv::getInstance).thenReturn(executeEnv);
MetricRepo.updateUserConnectionMaxMetric("metric_user", 321L);
MetricVisitor visitor = new PrometheusMetricVisitor();
MetricRepo.DORIS_METRIC_REGISTER.accept(visitor);
String metricResult = visitor.finish();
Assertions.assertTrue(metricResult.contains("# TYPE
doris_fe_connection_max gauge"));
-
Assertions.assertTrue(metricResult.contains("doris_fe_connection_max 13086"));
+
Assertions.assertTrue(metricResult.contains("doris_fe_connection_max 4321\n"),
metricResult);
+
Assertions.assertTrue(metricResult.contains("doris_fe_connection_total 2\n"),
metricResult);
Assertions.assertTrue(metricResult.contains("# TYPE
doris_fe_arrow_flight_connection_total gauge"));
-
Assertions.assertTrue(metricResult.contains("doris_fe_arrow_flight_connection_total
0"));
+
Assertions.assertTrue(metricResult.contains("doris_fe_arrow_flight_connection_total
1\n"), metricResult);
Assertions.assertTrue(metricResult.contains("# TYPE
doris_fe_arrow_flight_connection_max gauge"));
-
Assertions.assertTrue(metricResult.contains("doris_fe_arrow_flight_connection_max
8765"));
+
Assertions.assertTrue(metricResult.contains("doris_fe_arrow_flight_connection_max
300\n"), metricResult);
Assertions.assertTrue(metricResult.contains("# TYPE
doris_fe_user_connection_max gauge"));
Assertions.assertTrue(metricResult.contains("doris_fe_user_connection_max{user=\"metric_user\"}
321"));
@@ -123,8 +142,6 @@ public class MetricsTest {
Assertions.assertTrue(metricResult.contains("doris_fe_user_connection_max{user=\"root\"}
456"));
Assertions.assertFalse(metricResult.contains("doris_fe_user_connection_max{user=\"root\"}
789"));
} finally {
- Config.qe_max_connection = originQeMaxConnection;
- Config.arrow_flight_max_connections =
originArrowFlightMaxConnections;
MetricRepo.removeUserConnectionMaxMetric("metric_user");
Env.getServingEnv().getAuth().updateUserPropertyInternal(Auth.ROOT_USER,
Lists.newArrayList(
Pair.of(UserProperty.PROP_MAX_USER_CONNECTIONS, "100")),
true);
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/mysql/ConnectionExceedTest.java
b/fe/fe-core/src/test/java/org/apache/doris/mysql/ConnectionExceedTest.java
index 813a796c856..35345eb3fe8 100644
--- a/fe/fe-core/src/test/java/org/apache/doris/mysql/ConnectionExceedTest.java
+++ b/fe/fe-core/src/test/java/org/apache/doris/mysql/ConnectionExceedTest.java
@@ -32,8 +32,11 @@ import org.apache.doris.qe.ConnectScheduler;
import org.apache.doris.qe.QueryState;
import org.apache.doris.service.ExecuteEnv;
+import org.apache.arrow.flight.FlightRuntimeException;
+import org.apache.arrow.flight.FlightStatusCode;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
import org.mockito.InOrder;
import org.mockito.MockedStatic;
import org.mockito.Mockito;
@@ -42,6 +45,10 @@ import org.xnio.XnioIoThread;
import org.xnio.XnioWorker;
import org.xnio.conduits.ConduitStreamSourceChannel;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CyclicBarrier;
import java.util.concurrent.RejectedExecutionException;
public class ConnectionExceedTest {
@@ -128,15 +135,24 @@ public class ConnectionExceedTest {
}
}
+ // An Arrow Flight SQL session is a connection of the one pool: refused at
the pool's limit, the
+ // user's limit or the Flight sub-quota, in the words a MySQL client is
refused in, as the
+ // RESOURCE_EXHAUSTED status of the request that would have opened it -
and the bearer token
+ // goes with the refusal, since it could never open a session.
@Test
public void testFlightSessionConnectionExceed() throws Exception {
- try (MockedStatic<ExecuteEnv> mockedExecEnv =
Mockito.mockStatic(ExecuteEnv.class)) {
- // Create a scheduler with small max connections
+ try (MockedStatic<ExecuteEnv> mockedExecEnv =
Mockito.mockStatic(ExecuteEnv.class);
+ MockedStatic<Env> mockedEnvStatic =
Mockito.mockStatic(Env.class)) {
+ // A pool of 1000 with a Flight sub-quota of 2
ConnectScheduler scheduler = new ConnectScheduler(1000, 2);
// Setup expectations
Mockito.when(mockEnv.getInternalCatalog()).thenReturn(mockCatalog);
Mockito.when(mockCatalog.getName()).thenReturn("internal");
+
Mockito.when(mockAuth.getMaxConn(Mockito.anyString())).thenReturn(100L);
+ Mockito.when(mockEnv.getAuth()).thenReturn(mockAuth);
+ // The session the token manager builds takes its Env from
Env.getCurrentEnv().
+ mockedEnvStatic.when(Env::getCurrentEnv).thenReturn(mockEnv);
mockedExecEnv.when(ExecuteEnv::getInstance).thenReturn(mockExecuteEnv);
Mockito.when(mockExecuteEnv.getScheduler()).thenReturn(scheduler);
@@ -151,35 +167,91 @@ public class ConnectionExceedTest {
);
Mockito.when(mockTokenManager.validateToken("test_token")).thenReturn(tokenDetails);
- // Create first context and register
- ConnectContext context1 = new ConnectContext();
- context1.setEnv(mockEnv);
-
context1.setCurrentUserIdentity(UserIdentity.createAnalyzedUserIdentWithIp("test_user",
"%"));
- Assertions.assertTrue(scheduler.submit(context1));
- Assertions.assertEquals(-1,
scheduler.getFlightSqlConnectPoolMgr().registerConnection(context1));
-
- // Create second context and register
- ConnectContext context2 = new ConnectContext();
- context2.setEnv(mockEnv);
-
context2.setCurrentUserIdentity(UserIdentity.createAnalyzedUserIdentWithIp("test_user",
"%"));
- Assertions.assertTrue(scheduler.submit(context2));
- Assertions.assertEquals(-1,
scheduler.getFlightSqlConnectPoolMgr().registerConnection(context2));
+ // Two Flight sessions fill the sub-quota, next to a MySQL
connection of the same user: the
+ // refusal has to tell the Flight usage from the pool's count.
+ ConnectContext mysql = new ConnectContext();
+ mysql.setEnv(mockEnv);
+
mysql.setCurrentUserIdentity(UserIdentity.createAnalyzedUserIdentWithIp("test_user",
"%"));
+ Assertions.assertTrue(scheduler.submit(mysql));
+ Assertions.assertEquals(-1,
scheduler.getConnectPoolMgr().registerConnection(mysql));
+ for (String token : new String[] {"token-1", "token-2"}) {
+ ConnectContext session = ConnectContext.forFlight(token);
+ session.setEnv(mockEnv);
+
session.setCurrentUserIdentity(UserIdentity.createAnalyzedUserIdentWithIp("test_user",
"%"));
+ Assertions.assertTrue(scheduler.submit(session));
+ Assertions.assertEquals(-1,
scheduler.getConnectPoolMgr().registerConnection(session));
+ }
// Create FlightSessionsWithTokenManager and try to create a new
connection
FlightSessionsWithTokenManager manager = new
FlightSessionsWithTokenManager(mockTokenManager);
- try {
- manager.createConnectContext("test_token");
- Assertions.fail("Should throw IllegalArgumentException");
- } catch (IllegalArgumentException e) {
- // Verify error message is set correctly
- String expectedMsg = String.format(
- "Register arrow flight sql connection failed, Unknown
Error, the number of arrow flight "
- + "bearer tokens should be equal to arrow
flight sql max connections, "
- + "max connections: %d, used: %d.",
-
scheduler.getFlightSqlConnectPoolMgr().getMaxConnections(),
- scheduler.getConnectionNum());
- Assertions.assertEquals(expectedMsg, e.getMessage());
- }
+ FlightRuntimeException refused =
Assertions.assertThrows(FlightRuntimeException.class,
+ () -> manager.getConnectContext("test_token"));
+ Assertions.assertEquals(FlightStatusCode.RESOURCE_EXHAUSTED,
refused.status().code());
+ Assertions.assertEquals(
+ "Reach limit of connections. Total: 1000, User: 100,
Current: 3, Arrow Flight SQL: 2 (current: 2)",
+ refused.status().description());
+ Mockito.verify(mockTokenManager).invalidateToken("test_token");
+ Assertions.assertEquals(3, scheduler.getConnectionNum());
+ Assertions.assertEquals(2,
scheduler.getConnectPoolMgr().getFlightConnectionNum());
+ }
+ }
+
+ // Two first requests carrying the same bearer token race session
creation. getConnectContext
+ // serializes creation and re-checks the peer-identity index, so exactly
one ConnectContext is
+ // published for the token and every caller gets it. Without that (both
build and register one, or
+ // the loser hits the created-session guard) the pool would hold a second,
orphaned session, or a
+ // caller would fail. Each worker installs its own static mocks because
Mockito's mockStatic is
+ // thread-local; the ConnectScheduler is the one shared instance they
publish into.
+ @Test
+ @Timeout(60)
+ public void testConcurrentFirstRequestsPublishExactlyOneSessionPerToken()
throws Exception {
+ ConnectScheduler scheduler = new ConnectScheduler(1000, 100);
+ Mockito.when(mockEnv.getInternalCatalog()).thenReturn(mockCatalog);
+ Mockito.when(mockCatalog.getName()).thenReturn("internal");
+
Mockito.when(mockAuth.getMaxConn(Mockito.anyString())).thenReturn(100L);
+ Mockito.when(mockEnv.getAuth()).thenReturn(mockAuth);
+ Mockito.when(mockExecuteEnv.getScheduler()).thenReturn(scheduler);
+
+ UserIdentity userIdentity =
UserIdentity.createAnalyzedUserIdentWithIp("test_user", "%");
+ FlightTokenDetails tokenDetails = new FlightTokenDetails(
+ "shared_token", "test_user",
+ System.currentTimeMillis(), System.currentTimeMillis() +
3600000,
+ userIdentity, "127.0.0.1");
+
Mockito.when(mockTokenManager.validateToken("shared_token")).thenReturn(tokenDetails);
+
+ FlightSessionsWithTokenManager manager = new
FlightSessionsWithTokenManager(mockTokenManager);
+
+ int threadCount = 6;
+ CyclicBarrier barrier = new CyclicBarrier(threadCount);
+ List<ConnectContext> results = Collections.synchronizedList(new
ArrayList<>());
+ List<Throwable> errors = Collections.synchronizedList(new
ArrayList<>());
+ List<Thread> threads = new ArrayList<>();
+ for (int i = 0; i < threadCount; i++) {
+ Thread t = new Thread(() -> {
+ try (MockedStatic<ExecuteEnv> execEnv =
Mockito.mockStatic(ExecuteEnv.class);
+ MockedStatic<Env> env = Mockito.mockStatic(Env.class))
{
+
execEnv.when(ExecuteEnv::getInstance).thenReturn(mockExecuteEnv);
+ env.when(Env::getCurrentEnv).thenReturn(mockEnv);
+ barrier.await();
+ results.add(manager.getConnectContext("shared_token"));
+ } catch (Throwable e) {
+ errors.add(e);
+ }
+ });
+ threads.add(t);
+ t.start();
+ }
+ for (Thread t : threads) {
+ t.join();
+ }
+
+ Assertions.assertTrue(errors.isEmpty(), "getConnectContext threw: " +
errors);
+ Assertions.assertEquals(1, scheduler.getConnectionNum(),
+ "concurrent first requests on one token must publish exactly
one session");
+ ConnectContext published =
scheduler.getContextWithPeerIdentity("shared_token");
+ Assertions.assertNotNull(published);
+ for (ConnectContext c : results) {
+ Assertions.assertSame(published, c, "every caller must get the one
published session");
}
}
}
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/mysql/protocol/MysqlProtocolAdapterTest.java
b/fe/fe-core/src/test/java/org/apache/doris/mysql/protocol/MysqlProtocolAdapterTest.java
index a9f05fedb63..bc4fcb7ea71 100644
---
a/fe/fe-core/src/test/java/org/apache/doris/mysql/protocol/MysqlProtocolAdapterTest.java
+++
b/fe/fe-core/src/test/java/org/apache/doris/mysql/protocol/MysqlProtocolAdapterTest.java
@@ -118,7 +118,7 @@ public class MysqlProtocolAdapterTest {
}
@Test
- public void testConnectionRegistersItsTraceIdInTheMysqlPool() {
+ public void testConnectionRegistersItsTraceIdInThePool() {
ConnectScheduler scheduler = new ConnectScheduler(10, 10);
ConnectContext ctx = new ConnectContext();
ctx.setConnectScheduler(scheduler);
@@ -128,7 +128,7 @@ public class MysqlProtocolAdapterTest {
ctx.setQueryId(queryId);
Assertions.assertEquals(DebugUtil.printId(queryId),
scheduler.getConnectPoolMgr().getQueryIdByTraceId("trace-1"));
- Assertions.assertEquals("",
scheduler.getFlightSqlConnectPoolMgr().getQueryIdByTraceId("trace-1"));
+ Assertions.assertEquals(DebugUtil.printId(queryId),
scheduler.getQueryIdByTraceId("trace-1"));
}
@Test
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectPoolMgrTest.java
b/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectPoolMgrTest.java
new file mode 100644
index 00000000000..9cc903b46d3
--- /dev/null
+++ b/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectPoolMgrTest.java
@@ -0,0 +1,314 @@
+// 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.qe;
+
+import org.apache.doris.analysis.UserIdentity;
+import org.apache.doris.catalog.Env;
+import org.apache.doris.datasource.InternalCatalog;
+import org.apache.doris.mysql.privilege.Auth;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+import org.mockito.InOrder;
+import org.mockito.Mockito;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * One pool for every protocol: MySQL connections and Arrow Flight SQL
sessions share the pool's
+ * limit and the user's limit, Flight sessions additionally their sub-quota,
and every teardown
+ * path meets in unregisterConnection.
+ */
+public class ConnectPoolMgrTest {
+
+ private static final UserIdentity ALICE =
UserIdentity.createAnalyzedUserIdentWithIp("alice", "%");
+ private static final UserIdentity BOB =
UserIdentity.createAnalyzedUserIdentWithIp("bob", "%");
+ private static final UserIdentity CAROL =
UserIdentity.createAnalyzedUserIdentWithIp("carol", "%");
+
+ private static ConnectContext registered(ConnectPoolMgr pool,
ConnectContext ctx, int connectionId) {
+ ctx.setConnectionId(connectionId);
+ Assertions.assertEquals(-1, pool.registerConnection(ctx));
+ return ctx;
+ }
+
+ @Test
+ public void testBothProtocolsShareThePoolsLimit() {
+ Env env = ConnectPoolTestSupport.envAllowing(100);
+ // The sub-quota is the whole pool here: only the pool's limit is in
play.
+ ConnectPoolMgr pool = new ConnectPoolMgr(2, 2);
+
+ ConnectContext mysql = registered(pool,
ConnectPoolTestSupport.mysqlConnection(env, ALICE), 1);
+ ConnectContext flight = registered(pool,
ConnectPoolTestSupport.flightSession(env, BOB, "token-2"), 2);
+ Assertions.assertEquals(2, pool.getConnectionNum());
+ Assertions.assertEquals(1, pool.getFlightConnectionNum());
+
+ // The third connection is refused whichever protocol it speaks, with
the count it was refused at,
+ // and a refused registration changes no count: not the pool's, not
the Flight share, not the user's.
+ Assertions.assertEquals(2,
pool.registerConnection(ConnectPoolTestSupport.mysqlConnection(env, ALICE)));
+ Assertions.assertEquals(2,
+
pool.registerConnection(ConnectPoolTestSupport.flightSession(env, ALICE,
"token-3")));
+ Assertions.assertEquals(2, pool.getConnectionNum());
+ Assertions.assertEquals(1, pool.getFlightConnectionNum());
+ Assertions.assertEquals(1,
pool.getUserConnectionMap().get(ALICE.getQualifiedUser()).get());
+ Assertions.assertEquals(1,
pool.getUserConnectionMap().get(BOB.getQualifiedUser()).get());
+
+ Assertions.assertSame(flight,
pool.getContextWithPeerIdentity("token-2"));
+ Assertions.assertNull(pool.getContextWithPeerIdentity("token-3"));
+ Assertions.assertSame(mysql, pool.getContext(1));
+
+ // Unregistering the MySQL connection frees its slot and touches
nothing of the Flight share.
+ pool.unregisterConnection(mysql);
+ Assertions.assertEquals(1, pool.getConnectionNum());
+ Assertions.assertEquals(1, pool.getFlightConnectionNum());
+ Assertions.assertEquals(-1,
+
pool.registerConnection(ConnectPoolTestSupport.flightSession(env, ALICE,
"token-3")));
+ Assertions.assertEquals(2, pool.getFlightConnectionNum());
+ }
+
+ @Test
+ public void testAUsersLimitCountsBothProtocols() {
+ Env env = ConnectPoolTestSupport.envAllowing(1);
+ ConnectPoolMgr pool = new ConnectPoolMgr(10);
+
+ registered(pool, ConnectPoolTestSupport.mysqlConnection(env, ALICE),
1);
+ // Alice's one connection is taken by MySQL: her Flight session is
refused...
+ Assertions.assertEquals(1,
pool.registerConnection(ConnectPoolTestSupport.flightSession(env, ALICE, "a")));
+ // ...and nothing of the refused registration lingers: not in the
pool, not counted anywhere.
+ Assertions.assertEquals(1, pool.getConnectionNum());
+ Assertions.assertEquals(0, pool.getFlightConnectionNum());
+ Assertions.assertEquals(1,
pool.getUserConnectionMap().get(ALICE.getQualifiedUser()).get());
+ Assertions.assertNull(pool.getContextWithPeerIdentity("a"));
+ // Bob is not affected.
+ registered(pool, ConnectPoolTestSupport.flightSession(env, BOB, "b"),
2);
+ }
+
+ @Test
+ public void testTheFlightSubQuotaIsEnforcedWithinThePoolsLimit() {
+ Env env = ConnectPoolTestSupport.envAllowing(100);
+ ConnectPoolMgr pool = new ConnectPoolMgr(10, 1);
+ Assertions.assertEquals(1, pool.getFlightMaxConnections());
+
+ registered(pool, ConnectPoolTestSupport.flightSession(env, ALICE,
"a"), 1);
+ Assertions.assertEquals(1,
pool.registerConnection(ConnectPoolTestSupport.flightSession(env, BOB, "b")));
+ Assertions.assertEquals(1, pool.getConnectionNum());
+ Assertions.assertEquals(1, pool.getFlightConnectionNum());
+ // MySQL connections are not held to the Flight sub-quota.
+ registered(pool, ConnectPoolTestSupport.mysqlConnection(env, BOB), 2);
+ Assertions.assertEquals(2, pool.getConnectionNum());
+ }
+
+ // Unset, the sub-quota is half of the pool's limit: Flight sessions,
which their clients mostly
+ // never close, cannot take the half MySQL clients connect through. Set,
it never exceeds the limit.
+ @Test
+ public void
testTheFlightSubQuotaIsHalfThePoolsLimitUnlessSetAndNeverExceedsIt() {
+ Assertions.assertEquals(5,
ConnectPoolMgr.effectiveFlightMaxConnections(10, -1));
+ Assertions.assertEquals(5, new
ConnectPoolMgr(10).getFlightMaxConnections());
+ Assertions.assertEquals(512,
ConnectPoolMgr.effectiveFlightMaxConnections(1024, -1));
+ Assertions.assertEquals(0,
ConnectPoolMgr.effectiveFlightMaxConnections(1, -1));
+ Assertions.assertEquals(3,
ConnectPoolMgr.effectiveFlightMaxConnections(10, 3));
+ Assertions.assertEquals(10,
ConnectPoolMgr.effectiveFlightMaxConnections(10, 10));
+ Assertions.assertEquals(10,
ConnectPoolMgr.effectiveFlightMaxConnections(10, 4096));
+
+ // The default half is enforced like any sub-quota: the sixth Flight
session of a pool of ten is
+ // refused while MySQL connections still get in.
+ Env env = ConnectPoolTestSupport.envAllowing(100);
+ ConnectPoolMgr pool = new ConnectPoolMgr(10);
+ for (int i = 1; i <= 5; i++) {
+ registered(pool, ConnectPoolTestSupport.flightSession(env, ALICE,
"token-" + i), i);
+ }
+ Assertions.assertEquals(5,
pool.registerConnection(ConnectPoolTestSupport.flightSession(env, BOB, "b")));
+ registered(pool, ConnectPoolTestSupport.mysqlConnection(env, BOB), 6);
+ Assertions.assertEquals(6, pool.getConnectionNum());
+ Assertions.assertEquals(5, pool.getFlightConnectionNum());
+ }
+
+ @Test
+ public void testTheRefusalReadsTheSameForEveryProtocol() {
+ Env env = ConnectPoolTestSupport.envAllowing(5);
+ // A sub-quota equal to the pool's limit and no Flight session in the
pool: the plain sentence.
+ ConnectPoolMgr pool = new ConnectPoolMgr(10, 10);
+ Assertions.assertEquals("Reach limit of connections. Total: 10, User:
5, Current: 10",
+
pool.limitReachedMessage(ConnectPoolTestSupport.mysqlConnection(env, ALICE),
10));
+ Assertions.assertEquals("Reach limit of connections. Total: 10, User:
5, Current: 10",
+
pool.limitReachedMessage(ConnectPoolTestSupport.flightSession(env, ALICE, "a"),
10));
+
+ // Once Flight sessions hold part of the pool, both protocols are told
the sub-quota and its
+ // usage - a MySQL client refused by a pool that Flight sessions
filled reads where they went.
+ // One connection of each protocol, so that the Flight usage cannot
pass for the pool's count.
+ ConnectPoolMgr quota = new ConnectPoolMgr(10, 2);
+ registered(quota, ConnectPoolTestSupport.flightSession(env, BOB, "b"),
1);
+ registered(quota, ConnectPoolTestSupport.mysqlConnection(env, BOB), 2);
+ Assertions.assertEquals(
+ "Reach limit of connections. Total: 10, User: 5, Current: 2,
Arrow Flight SQL: 2 (current: 1)",
+
quota.limitReachedMessage(ConnectPoolTestSupport.flightSession(env, ALICE,
"a"), 2));
+ Assertions.assertEquals(
+ "Reach limit of connections. Total: 10, User: 5, Current: 2,
Arrow Flight SQL: 2 (current: 1)",
+
quota.limitReachedMessage(ConnectPoolTestSupport.mysqlConnection(env, ALICE),
2));
+ ConnectPoolMgr halved = new ConnectPoolMgr(10);
+ registered(halved, ConnectPoolTestSupport.flightSession(env, BOB,
"b"), 1);
+ registered(halved, ConnectPoolTestSupport.mysqlConnection(env, BOB),
2);
+ Assertions.assertEquals(
+ "Reach limit of connections. Total: 10, User: 5, Current: 10,
Arrow Flight SQL: 5 (current: 1)",
+
halved.limitReachedMessage(ConnectPoolTestSupport.mysqlConnection(env, ALICE),
10));
+
+ // A sub-quota tighter than the pool's limit is named to a Flight
client even while no
+ // Flight session is in the pool: it is the limit that refused it.
+ ConnectPoolMgr none = new ConnectPoolMgr(10, 0);
+ Assertions.assertEquals(
+ "Reach limit of connections. Total: 10, User: 5, Current: 0,
Arrow Flight SQL: 0 (current: 0)",
+
none.limitReachedMessage(ConnectPoolTestSupport.flightSession(env, ALICE, "a"),
0));
+ Assertions.assertEquals("Reach limit of connections. Total: 10, User:
5, Current: 0",
+
none.limitReachedMessage(ConnectPoolTestSupport.mysqlConnection(env, ALICE),
0));
+ }
+
+ // Arrow Flight SQL keeps a query's coordinator alive across GetFlightInfo
-> DoGet (see #62259).
+ // unregisterConnection() is the catch-all teardown path: the idle timeout
(wait_timeout), bearer
+ // token expiry or eviction, CloseSession and a KILL CONNECTION from
another connection all reach
+ // here (a query timeout only cancels the query). The protocol must
release what it holds for the
+ // session -- for Flight the channel-cached results and the deferred
coordinators -- even for a
+ // connection that was never registered (an abandoned connection is
cleaned up, not leaked), and
+ // before the bookkeeping, so that a failure there cannot strand the
coordinators.
+ @Test
+ public void
testUnregisterReleasesTheProtocolSessionFirstEvenWhenNotRegistered() {
+ ConnectPoolMgr pool = new ConnectPoolMgr(100);
+ ConnectContext ctx = Mockito.mock(ConnectContext.class);
+
+ pool.unregisterConnection(ctx);
+
+ InOrder inOrder = Mockito.inOrder(ctx);
+ inOrder.verify(ctx).releaseProtocolSession();
+ inOrder.verify(ctx).closeTxn();
+ }
+
+ @Test
+ public void testUnregisterRemovesAFlightSessionAndItsPeerIdentity() {
+ Env env = ConnectPoolTestSupport.envAllowing(100);
+ ConnectPoolMgr pool = new ConnectPoolMgr(100, 1);
+ ConnectContext ctx = registered(pool,
ConnectPoolTestSupport.flightSession(env, ALICE, "token-7"), 7);
+ Assertions.assertSame(ctx, pool.getContextWithPeerIdentity("token-7"));
+
+ pool.unregisterConnection(ctx);
+
+ Assertions.assertNull(pool.getContext(7));
+ Assertions.assertNull(pool.getContextWithPeerIdentity("token-7"));
+ Assertions.assertEquals(0, pool.getConnectionNum());
+ Assertions.assertEquals(0, pool.getFlightConnectionNum());
+ Assertions.assertEquals(0,
pool.getUserConnectionMap().get(ALICE.getQualifiedUser()).get());
+ // The sub-quota slot is free again.
+ registered(pool, ConnectPoolTestSupport.flightSession(env, ALICE,
"token-8"), 8);
+ // Unregistering twice is harmless.
+ pool.unregisterConnection(ctx);
+ Assertions.assertEquals(1, pool.getConnectionNum());
+ }
+
+ // Two sessions under one bearer token, as two concurrent first requests
of a token can open:
+ // the index names the later one, and the earlier one's teardown must not
take that entry away.
+ @Test
+ public void testUnregisterRemovesOnlyTheConnectionsOwnPeerIdentityEntry() {
+ Env env = ConnectPoolTestSupport.envAllowing(100);
+ ConnectPoolMgr pool = new ConnectPoolMgr(100, 2);
+ ConnectContext earlier = registered(pool,
ConnectPoolTestSupport.flightSession(env, ALICE, "token-x"), 1);
+ ConnectContext later = registered(pool,
ConnectPoolTestSupport.flightSession(env, ALICE, "token-x"), 2);
+ Assertions.assertSame(later,
pool.getContextWithPeerIdentity("token-x"));
+
+ pool.unregisterConnection(earlier);
+
+ Assertions.assertNull(pool.getContext(1));
+ Assertions.assertSame(later,
pool.getContextWithPeerIdentity("token-x"));
+ Assertions.assertEquals(1, pool.getConnectionNum());
+ Assertions.assertEquals(1, pool.getFlightConnectionNum());
+
+ pool.unregisterConnection(later);
+ Assertions.assertNull(pool.getContextWithPeerIdentity("token-x"));
+ Assertions.assertEquals(0, pool.getFlightConnectionNum());
+ }
+
+ // qe_max_connection = 1 makes the default Flight sub-quota 0 (half,
floored). No Flight session may
+ // open: the first one is refused at the pool (returns the count, not -1)
while a MySQL connection
+ // still fits. The service floors the token cache at 1 so this refusal is
what the client meets as
+ // RESOURCE_EXHAUSTED, instead of the freshly issued token being evicted
first (see
+ // DorisFlightSqlService); here we cover the pool half of that boundary.
+ @Test
+ public void testAZeroFlightSubQuotaRefusesEveryFlightSessionButNotMysql() {
+ Assertions.assertEquals(0,
ConnectPoolMgr.effectiveFlightMaxConnections(1, -1));
+ Env env = ConnectPoolTestSupport.envAllowing(100);
+ ConnectPoolMgr pool = new ConnectPoolMgr(1);
+ Assertions.assertEquals(0, pool.getFlightMaxConnections());
+
+
Assertions.assertTrue(pool.registerConnection(ConnectPoolTestSupport.flightSession(env,
ALICE, "t")) >= 0,
+ "a Flight session must be refused when the sub-quota is 0");
+ Assertions.assertEquals(0, pool.getFlightConnectionNum());
+ Assertions.assertNull(pool.getContextWithPeerIdentity("t"));
+
+ // The one pool slot is still open to a MySQL connection.
+ Assertions.assertEquals(-1,
pool.registerConnection(ConnectPoolTestSupport.mysqlConnection(env, BOB)));
+ Assertions.assertEquals(1, pool.getConnectionNum());
+ }
+
+ // One attempt must not hold the last pool slot across another attempt's
checks. Freeze one
+ // registration inside getMaxConn (the user check) -- where the earlier
increment-check-rollback code
+ // was already holding the pool increment -- and register a connection
that fits on another thread:
+ // it must be admitted. The staged code refused it (the frozen attempt's
pool increment made the pool
+ // look full); the single admission critical section admits it, because
the pool is reserved only
+ // inside the lock and getMaxConn is read before the lock. Without the
synchronized block this fails.
+ @Test
+ @Timeout(30)
+ public void
testConcurrentAdmissionDoesNotRefuseAFittingConnectionWhileAnotherIsInFlight()
throws Exception {
+ CountDownLatch frozenInGetMaxConn = new CountDownLatch(1);
+ CountDownLatch releaseFrozen = new CountDownLatch(1);
+ Auth auth = Mockito.mock(Auth.class);
+ Mockito.when(auth.getMaxConn(Mockito.anyString())).thenAnswer(inv -> {
+ if (ALICE.getQualifiedUser().equals(inv.getArgument(0))) {
+ frozenInGetMaxConn.countDown();
+ releaseFrozen.await();
+ }
+ return 100L;
+ });
+ Env env = Mockito.mock(Env.class);
+ Mockito.when(env.getAuth()).thenReturn(auth);
+ InternalCatalog internalCatalog = Mockito.mock(InternalCatalog.class);
+
Mockito.when(internalCatalog.getName()).thenReturn(InternalCatalog.INTERNAL_CATALOG_NAME);
+ Mockito.when(env.getInternalCatalog()).thenReturn(internalCatalog);
+
+ ConnectPoolMgr pool = new ConnectPoolMgr(2);
+ registered(pool, ConnectPoolTestSupport.mysqlConnection(env, CAROL),
1); // one slot left
+
+ ConnectContext frozen = ConnectPoolTestSupport.mysqlConnection(env,
ALICE);
+ frozen.setConnectionId(2);
+ ConnectContext fitting = ConnectPoolTestSupport.mysqlConnection(env,
BOB);
+ fitting.setConnectionId(3);
+
+ Thread frozenThread = new Thread(() ->
pool.registerConnection(frozen));
+ frozenThread.start();
+ Assertions.assertTrue(frozenInGetMaxConn.await(10, TimeUnit.SECONDS),
+ "the frozen attempt never reached getMaxConn");
+
+ AtomicInteger fittingResult = new AtomicInteger(Integer.MIN_VALUE);
+ Thread fittingThread = new Thread(() ->
fittingResult.set(pool.registerConnection(fitting)));
+ fittingThread.start();
+ fittingThread.join();
+ Assertions.assertEquals(-1, fittingResult.get(),
+ "a connection that fits the pool was refused while another
attempt's admission was in flight");
+
+ releaseFrozen.countDown();
+ frozenThread.join();
+ }
+}
diff --git
a/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectPoolTestSupport.java
b/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectPoolTestSupport.java
new file mode 100644
index 00000000000..7928b3198d3
--- /dev/null
+++ b/fe/fe-core/src/test/java/org/apache/doris/qe/ConnectPoolTestSupport.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.doris.qe;
+
+import org.apache.doris.analysis.UserIdentity;
+import org.apache.doris.catalog.Env;
+import org.apache.doris.datasource.InternalCatalog;
+import org.apache.doris.mysql.privilege.Auth;
+
+import org.mockito.Mockito;
+
+/**
+ * Connections as the pool sees them, for tests that register real contexts:
the pool asks the
+ * connection's Env for the user's connection limit and files the connection
under its user, so a
+ * context needs both before it can be registered.
+ */
+public final class ConnectPoolTestSupport {
+
+ private ConnectPoolTestSupport() {
+ }
+
+ /** An Env whose Auth allows every user {@code maxUserConnections}
connections. */
+ public static Env envAllowing(long maxUserConnections) {
+ Env env = Mockito.mock(Env.class);
+ Auth auth = Mockito.mock(Auth.class);
+
Mockito.when(auth.getMaxConn(Mockito.anyString())).thenReturn(maxUserConnections);
+ Mockito.when(env.getAuth()).thenReturn(auth);
+ // ConnectContext.setEnv starts the connection in the internal catalog.
+ InternalCatalog internalCatalog = Mockito.mock(InternalCatalog.class);
+
Mockito.when(internalCatalog.getName()).thenReturn(InternalCatalog.INTERNAL_CATALOG_NAME);
+ Mockito.when(env.getInternalCatalog()).thenReturn(internalCatalog);
+ return env;
+ }
+
+ /** A MySQL connection of {@code user}, registrable in a pool. */
+ public static ConnectContext mysqlConnection(Env env, UserIdentity user) {
+ ConnectContext ctx = new ConnectContext();
+ ctx.setEnv(env);
+ ctx.setCurrentUserIdentity(user);
+ return ctx;
+ }
+
+ /** An Arrow Flight SQL session of {@code user} under the bearer token
{@code peerIdentity}, registrable in a pool. */
+ public static ConnectContext flightSession(Env env, UserIdentity user,
String peerIdentity) {
+ ConnectContext ctx = ConnectContext.forFlight(peerIdentity);
+ ctx.setEnv(env);
+ ctx.setCurrentUserIdentity(user);
+ return ctx;
+ }
+}
diff --git
a/regression-test/suites/arrow_flight_sql_p0/test_connection_quota.groovy
b/regression-test/suites/arrow_flight_sql_p0/test_connection_quota.groovy
new file mode 100644
index 00000000000..3067a060a4b
--- /dev/null
+++ b/regression-test/suites/arrow_flight_sql_p0/test_connection_quota.groovy
@@ -0,0 +1,246 @@
+// 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.
+
+import java.sql.DriverManager
+import java.sql.SQLException
+import java.util.regex.Pattern
+import java.util.concurrent.TimeUnit
+
+import org.awaitility.Awaitility
+
+// The Flight SQL JDBC driver on the classpath shades Arrow Flight; its
FlightSqlClient is the
+// one a test can drive directly (see test_session_options).
+import
org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.CloseSessionRequest
+import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.FlightClient
+import
org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.FlightRuntimeException
+import
org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.FlightStatusCode
+import org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.Location
+import
org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.flight.sql.FlightSqlClient
+import
org.apache.arrow.driver.jdbc.shaded.org.apache.arrow.memory.RootAllocator
+
+// An Arrow Flight SQL session is a connection of the same pool a MySQL
connection is in: the
+// user's max_user_connections counts both, whichever came first, and a
refusal reads the same
+// over either protocol - over Flight as the RESOURCE_EXHAUSTED status of the
request that would
+// have opened the session, since a session opens on its first request, not
when the token is
+// issued.
+//
+// The limit is 4 rather than 1 because the token manager still keeps at most
+// max_user_connections / 2 bearer tokens per user and evicts the oldest -
session included - when
+// one more is issued; three MySQL connections take the slots a token cannot,
so that the second
+// Flight session is the one the pool refuses. (The per-user token cache goes
with PR-3.2.)
+//
+// Not in the 'arrow_flight_sql' group on purpose: `sql` stays the MySQL
control connection, and
+// the sessions under test are raw Flight SQL clients and plain MySQL
connections of their own.
+suite("test_connection_quota") {
+ String host = context.config.otherConfigs.get("extArrowFlightSqlHost")
+ int port = context.config.otherConfigs.get("extArrowFlightSqlPort") as int
+
+ String user = "flight_quota_user"
+ String password = "Quota_12345"
+ int limit = 4
+ // What both protocols say when the user's limit is reached; only the
count at the refusal varies.
+ Pattern refusal = Pattern.compile("^Reach limit of connections\\. Total:
(\\d+), User: ${limit}, Current: \\d+")
+
+ sql "DROP USER IF EXISTS '${user}'"
+ sql "CREATE USER '${user}' IDENTIFIED BY '${password}'"
+ sql "GRANT SELECT_PRIV ON *.* TO '${user}'"
+ // In cloud mode a query runs on a compute group; a user with no
USAGE_PRIV on one is refused with
+ // an INTERNAL error when it runs a statement (the Flight SELECT 1 below),
before the pool's quota
+ // is ever reached. Grant it the way the other cloud Flight suites do (see
test_auth_remote_ip).
+ if (isCloudMode()) {
+ def computeGroups = sql "SHOW COMPUTE GROUPS"
+ assertTrue(!computeGroups.isEmpty(), "cloud mode but SHOW COMPUTE
GROUPS returned nothing")
+ sql "GRANT USAGE_PRIV ON COMPUTE GROUP '${computeGroups[0][0]}' TO
'${user}'"
+ }
+ sql "SET PROPERTY FOR '${user}' 'max_user_connections' = '${limit}'"
+
+ // The connection quota is enforced per-FE, and
information_schema.processlist is a cluster-wide
+ // view when fetch_all_fe_for_system_table is on (the default): on a
multi-FE cluster the count then
+ // depends on which FE answers and can include connections on another FE,
so a per-FE quota test
+ // cannot rely on it. Anchor everything to the one FE that serves the
configured Flight endpoint --
+ // open the MySQL connections on that FE (its Flight host, the query port
and options from jdbcUrl)
+ // and count only that FE's own connections (fetch_all_fe_for_system_table
= false). Without this the
+ // MySQL connections, the Flight sessions and the counting query need not
land on the same FE and the
+ // count never settles (build 1049514).
+ def jdbcMatcher = (context.config.jdbcUrl =~
/^(jdbc:mysql:\/\/)[^\/:@]+(:\d+.*)$/)
+ assertTrue(jdbcMatcher.matches(),
+ "cannot derive the Flight FE's MySQL url from jdbcUrl:
${context.config.jdbcUrl}")
+ def feMysqlUrl = jdbcMatcher.replaceFirst("\$1${host}\$2")
+ def allocator = new RootAllocator()
+ def client = FlightClient.builder(allocator,
Location.forGrpcInsecure(host, port)).build()
+ def flight = new FlightSqlClient(client)
+ def mysqlConnections = []
+ def openTokens = []
+ def countConn = null
+ try {
+ // A dedicated connection on the Flight FE that reports only that FE's
connections, so the count
+ // is exactly the MySQL connections and Flight sessions this suite
opened there.
+ countConn = DriverManager.getConnection(feMysqlUrl,
context.config.jdbcUser, context.config.jdbcPassword)
+ countConn.createStatement().withCloseable { it.execute("SET
fetch_all_fe_for_system_table = false") }
+ // The anchoring above assumes the FE reached at feMysqlUrl is the one
serving the configured
+ // Flight endpoint (its host, jdbcUrl's query port). Assert it here --
otherwise the suite would
+ // fail later at the quota checks with a confusing pool-pointing error
-- by comparing this FE's
+ // own arrow_flight_sql_port to the configured Flight port.
forward_to_master keeps SHOW FRONTEND
+ // CONFIG reporting this FE rather than the master.
+ countConn.createStatement().withCloseable { it.execute("SET
forward_to_master = false") }
+ countConn.createStatement().withCloseable { st ->
+ def rs = st.executeQuery("SHOW FRONTEND CONFIG LIKE
'arrow_flight_sql_port'")
+ assertTrue(rs.next(), "SHOW FRONTEND CONFIG returned no
arrow_flight_sql_port")
+ assertEquals(port as String, rs.getString("Value"),
+ "the FE reached at ${feMysqlUrl} does not serve the Flight
endpoint on port ${port}; "
+ + "this suite needs MySQL, Flight and the count
query on the same FE")
+ }
+ // The user's connections as that FE's pool sees them: MySQL
connections and Flight sessions alike.
+ def connectionsOf = {
+ def st = countConn.createStatement()
+ try {
+ def rs = st.executeQuery(
+ "SELECT COUNT(*) FROM information_schema.processlist
WHERE User = '${user}'")
+ rs.next()
+ return rs.getInt(1)
+ } finally {
+ st.close()
+ }
+ }
+ // A closed MySQL connection leaves the pool on the frontend's nio
thread after the client's
+ // COM_QUIT, and Connector/J does not wait for that; so wait here
before the next connection
+ // is opened against the count.
+ def awaitConnections = { int expected ->
+ Awaitility.await().atMost(30, TimeUnit.SECONDS).pollInterval(200,
TimeUnit.MILLISECONDS)
+ .until { connectionsOf() == expected }
+ }
+ // Closes a session the way the drivers do, once; a token whose
session is already gone is
+ // UNAUTHENTICATED and needs nothing.
+ def closeSession = { cred ->
+ try {
+ flight.closeSession(new CloseSessionRequest(), cred)
+ } catch (FlightRuntimeException e) {
+ assertEquals(FlightStatusCode.UNAUTHENTICATED,
e.status().code())
+ }
+ openTokens.remove(cred)
+ }
+ // Whether the pool refused this token's session: its
RESOURCE_EXHAUSTED refusal message, or null
+ // when the session was admitted. The quota is checked when the
session is registered, before the
+ // probe query runs (DorisFlightSqlProducer.getFlightInfoStatement
gets the ConnectContext, which
+ // registers it, and only then executes the statement). So a
RESOURCE_EXHAUSTED is the refusal --
+ // the session never entered the pool and the frontend invalidated the
token -- while any other
+ // outcome (success, or the statement failing for a reason outside
this suite's scope, e.g. an
+ // environment-specific query error) means the session was admitted
and is in the pool; the
+ // connectionsOf() checks below verify that count. The token is
tracked before the request so a
+ // session that opened is closed at the end.
+ def refusalOf = { cred ->
+ openTokens << cred
+ try {
+ flight.execute("SELECT 1", cred).getEndpoints()
+ return null
+ } catch (FlightRuntimeException e) {
+ def code = e.status().code()
+ if (code == FlightStatusCode.RESOURCE_EXHAUSTED) {
+ openTokens.remove(cred)
+ return e.status().description()
+ }
+ if (code == FlightStatusCode.INTERNAL) {
+ // Admitted: the pool check passed and the session was
registered before the statement
+ // ran, so it counts (the connectionsOf() checks confirm
it). The probe query then
+ // failed for a reason outside a connection-quota test's
scope -- an environment-
+ // specific execution error -- which is logged here so it
stays diagnosable.
+ logger.warn("test_connection_quota: an admitted Flight
session's probe query failed "
+ + "with ${code}: ${e.status().description()}")
+ return null
+ }
+ // Neither a quota refusal nor an admitted session (e.g.
UNAUTHENTICATED); let it surface.
+ throw e
+ }
+ }
+ // The pool's limit named by a refusal, asserting the refusal is
worded as MySQL's is.
+ def totalOf = { String message ->
+ def m = refusal.matcher(message)
+ assertTrue(m.find(), "expected the MySQL wording of the refusal,
got: ${message}")
+ return m.group(1)
+ }
+ // A MySQL connection of the user that could not open: the refusal, or
null when it opened.
+ def mysqlRefusal = {
+ try {
+ mysqlConnections << DriverManager.getConnection(feMysqlUrl,
user, password)
+ return null
+ } catch (SQLException e) {
+ return e.getMessage()
+ }
+ }
+
+ // 1. Three MySQL connections, then the first Flight session: the
user's four.
+ (1..limit - 1).each { assertNull(mysqlRefusal(), "MySQL connection
${it} of ${limit - 1} was refused") }
+ awaitConnections(limit - 1)
+ def first = client.authenticateBasicToken(user, password).get()
+ assertNull(refusalOf(first), "the Flight session should open as the
user's last connection")
+ assertEquals(limit, connectionsOf(), "the Flight session is a
connection of the user's like the others")
+
+ // 2. The second Flight session is refused, in MySQL's words...
+ def second = client.authenticateBasicToken(user, password).get()
+ String flightRefused = refusalOf(second)
+ assertNotNull(flightRefused, "the second Flight session opened
although the user's limit is reached")
+ String total = totalOf(flightRefused)
+
+ // 3. ...and so is a MySQL connection, in the same words.
+ String mysqlRefused = mysqlRefusal()
+ assertNotNull(mysqlRefused, "the MySQL connection opened although the
user's limit is reached")
+ assertEquals(total, totalOf(mysqlRefused))
+
+ // 4. CloseSession releases the Flight session's connection: a MySQL
connection opens now, and
+ // once it is closed again a Flight session does.
+ assertEquals("CLOSED", flight.closeSession(new CloseSessionRequest(),
first).getStatus().name())
+ openTokens.remove(first)
+ awaitConnections(limit - 1)
+ assertNull(mysqlRefusal(), "the MySQL connection should open once the
Flight session is closed")
+ mysqlConnections.remove(mysqlConnections.size() - 1).close()
+ awaitConnections(limit - 1)
+ def third = client.authenticateBasicToken(user, password).get()
+ assertNull(refusalOf(third), "the Flight session should open once the
MySQL connection is closed")
+ assertEquals(limit, connectionsOf(), "the reopened Flight session must
be the user's ${limit}th connection")
+ closeSession(third)
+ awaitConnections(limit - 1)
+ } finally {
+ // Sessions outlive the client: close the ones still open, or a rerun
on the same frontend
+ // starts against a user whose slots they hold until wait_timeout
(DROP USER does not end them).
+ // Nothing here asserts or throws: the failure that brought the suite
here, if any, is the one
+ // reported, and every step of the cleanup runs.
+ def quietly = { String what, Closure step ->
+ try {
+ step()
+ } catch (Exception e) {
+ logger.warn("cleanup of test_connection_quota: ${what} failed:
${e.message}")
+ }
+ }
+ openTokens.each { cred ->
+ quietly("closing a session left open") {
+ try {
+ flight.closeSession(new CloseSessionRequest(), cred)
+ } catch (FlightRuntimeException e) {
+ // A token whose session is already gone is
UNAUTHENTICATED and needs nothing.
+ if (e.status().code() != FlightStatusCode.UNAUTHENTICATED)
{
+ throw e
+ }
+ }
+ }
+ }
+ mysqlConnections.each { conn -> quietly("closing a MySQL connection")
{ conn.close() } }
+ quietly("closing the count connection") { if (countConn != null)
countConn.close() }
+ quietly("closing the Flight client") { client.close() }
+ quietly("closing the allocator") { allocator.close() }
+ sql "DROP USER IF EXISTS '${user}'"
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
