morningman commented on code in PR #68101:
URL: https://github.com/apache/doris/pull/68101#discussion_r4056024198
##########
fe/fe-core/src/main/java/org/apache/doris/qe/ConnectPoolMgr.java:
##########
@@ -66,23 +109,59 @@ public void timeoutChecker(long now) {
// Return -1 means register OK
// Return >=0 means register failed, and return value is current
connection num.
public int registerConnection(ConnectContext ctx) {
+ boolean flight = isFlight(ctx);
if (numberConnection.incrementAndGet() > maxConnections) {
numberConnection.decrementAndGet();
return numberConnection.get();
}
+ if (flight && numberFlightConnection.incrementAndGet() >
flightMaxConnections) {
+ numberFlightConnection.decrementAndGet();
+ 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();
+ if (flight) {
+ numberFlightConnection.decrementAndGet();
+ }
numberConnection.decrementAndGet();
return numberConnection.get();
}
connectionMap.put(ctx.getConnectionId(), ctx);
+ if (flight) {
+ peerIdentity2ConnectionId.put(ctx.getPeerIdentity(),
ctx.getConnectionId());
Review Comment:
Fixed in b446ac89ac8 (`a435eae450b` before the rebase onto master), at the
creation rather than at this put -- once creation is serialized per token,
`registerConnection` sees one context per token and the put has nothing to
overwrite. Two changes in `FlightSessionsWithTokenManager`:
- `getConnectContext` serializes session creation under
`sessionCreationLock` and re-checks the peer-identity index inside it, so of
two concurrent first requests on one bearer token exactly one builds and
registers a `ConnectContext`; the loser returns the winner's published context
instead of registering a second one.
`ConnectionExceedTest.testConcurrentFirstRequestsPublishExactlyOneSessionPerToken`
fires six first requests on one token through a barrier and asserts one
session in the pool, returned to every caller.
- `createConnectContext` re-validates the token after `registerConnection`,
still under that lock; if the token was invalidated in between (a concurrent
CloseSession, or the same-segment LRU evicting it) the session is unregistered
right there instead of lingering in the pool until `wait_timeout`.
The residual window this PR leaves (a token invalidated after the
re-validation) is closed by the next PR of the series, which moves the token
lifecycle under one owner; the comment at the re-validation says so.
##########
fe/fe-core/src/main/java/org/apache/doris/arrowflight/DorisFlightSqlService.java:
##########
@@ -52,8 +53,24 @@ 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);
+ }
+ // The token cache holds as many tokens as the sub-quota allows,
capped by
+ // arrow_flight_token_cache_size: a session opens on a token, so this
is what bounds Flight
+ // sessions in practice (the oldest token, and its session, is evicted
first).
+ int tokenCacheSize = Math.min(flightMaxConnections,
Config.arrow_flight_token_cache_size);
Review Comment:
Fixed in b446ac89ac8 (`a435eae450b` before the rebase). The cache size is
now `DorisFlightSqlService.effectiveTokenCacheSize(subQuota,
arrow_flight_token_cache_size) = min(max(1, subQuota),
arrow_flight_token_cache_size)`: a legal sub-quota of 0 (`qe_max_connection =
1` with the `-1` default) still gets a one-token cache, so the first request
reaches the pool and is refused with the documented RESOURCE_EXHAUSTED instead
of the token being evicted at issue time and the request failing
UNAUTHENTICATED. The floor is deliberately not applied to
`arrow_flight_token_cache_size` itself: an illegal value there (<= 0) keeps the
base's loud failure rather than silently running the FE on a one-token cache.
The `invalid bearer token` diagnostic prints the formula with the floor.
Tests: `DorisFlightSqlServiceTest` (the floor applies to the sub-quota term
and not to the cap) and
`ConnectPoolMgrTest.testAZeroFlightSubQuotaRefusesEveryFlightSessionButNotMysql`
(sub-quota 0 refuses every Flight session and leaves the one pool slot to
MySQL). An end-to-end run of the zero sub-quota needs an FE started with
`qe_max_connection = 1`, which the shared regression cluster cannot be; the
end-to-end check of the same refusal path -- session refused on its first
request, RESOURCE_EXHAUSTED, MySQL wording -- is
`arrow_flight_sql_p0/test_connection_quota` through the per-user limit.
##########
fe/fe-core/src/main/java/org/apache/doris/qe/ConnectPoolMgr.java:
##########
@@ -66,23 +109,59 @@ public void timeoutChecker(long now) {
// Return -1 means register OK
// Return >=0 means register failed, and return value is current
connection num.
public int registerConnection(ConnectContext ctx) {
+ boolean flight = isFlight(ctx);
if (numberConnection.incrementAndGet() > maxConnections) {
numberConnection.decrementAndGet();
return numberConnection.get();
}
+ if (flight && numberFlightConnection.incrementAndGet() >
flightMaxConnections) {
Review Comment:
Fixed in b446ac89ac8 (`a435eae450b` before the rebase). `registerConnection`
now checks the pool, the Flight sub-quota and the user's limit and reserves all
three under `admissionLock` (check everything, then increment), instead of the
increment-check-rollback per counter that let a doomed attempt hold the last
pool slot across another attempt's checks. The user's limit is read before the
lock (`Auth.getMaxConn` takes its own locks), so the critical section is a leaf
and holds nothing that can block; registration is once per connection, not per
query. Outcome per cell is unchanged from the staged code, the refusal message
and counts included.
`ConnectPoolMgrTest.testConcurrentAdmissionDoesNotRefuseAFittingConnectionWhileAnotherIsInFlight`
freezes one registration inside `getMaxConn` and registers a connection that
fits on another thread; the staged code refused it, this code admits it.
##########
fe/fe-core/src/test/java/org/apache/doris/arrowflight/protocol/FlightProtocolAdapterTest.java:
##########
@@ -138,24 +145,78 @@ public void
testSessionRegistersItsTraceIdInTheFlightPool() {
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();
+ }
+
+ // The channel close can fail - the allocator refuses to close over a
buffer still referenced
+ // elsewhere - and the session is torn down all the same: the deferred
query is finalized and
+ // the session takes no further command.
+ @Test
+ public void
testReleaseSessionTearsTheSessionDownEvenWhenTheChannelCloseFails() {
+ ConnectContext ctx = flightSession();
+ FlightProtocolAdapter adapter = FlightProtocolAdapter.of(ctx);
+ StmtExecutor deferred = Mockito.mock(StmtExecutor.class);
+ ctx.addFlightSqlDeferredExecutor(deferred);
+ adapter.getChannel().addOKResult("query-1", "SELECT 1");
+ // A buffer of the cached result held by someone else, as a DoGet
still streaming it would:
+ // closing the result does not free it, and the allocator's close
throws over it. The
+ // reference is never given back - the failure is the point of the
test.
+ ArrowBuf held =
adapter.getChannel().getResult("query-1").getVectorSchemaRoot().getVector(0)
+ .getDataBuffer();
+ held.getReferenceManager().retain();
Review Comment:
Resolved in b446ac89ac8 (`a435eae450b` before the rebase) by dropping the
test rather than adding the `finally`: the retained buffer cannot be released
after the assertion. `BaseAllocator.close()` (arrow-memory-core 19.0.0) sets
`isClosed` before it reports the leak, so once `releaseSession` has swallowed
that failure there is no open allocator to hand `held` back to -- the only way
to keep the assertion was to keep the leaked native allocation in the test JVM,
which is exactly what you flagged. The normal release path (channel closed,
allocator closed, cached root invalidated) stays covered by the remaining
`FlightProtocolAdapterTest` cases.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
