github-actions[bot] commented on code in PR #68101:
URL: https://github.com/apache/doris/pull/68101#discussion_r4046064594
##########
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:
[P1] Make shared-quota admission one atomic transition
With `maxConnections=2`, `flightMaxConnections=1`, and one Flight session, a
doomed second Flight attempt can reserve the last global slot here. A
concurrent valid MySQL attempt then increments to 3 and is refused; afterward
Flight fails its sub-quota and rolls the global count back to 1. Both calls
fail although MySQL fits in every serial ordering, so saturated Flight traffic
can consume the capacity this PR intends to reserve for MySQL. Check/reserve
the global, Flight, and user quotas in one short critical transition and add a
barrier-based test.
##########
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:
[P1] Make token/session publication atomic
Two first RPCs can both pass the unlocked lookup/`createdSession` check, and
this put overwrites one connection id. Token eviction or CloseSession then
unregisters only the last context, leaving the other in `connectionMap` and all
three quota counters until `wait_timeout`; invalidation between `validateToken`
and this put creates the same orphan with only one request. Since those slots
now also block MySQL, commit exactly one live token-to-context owner atomically
(and roll back the loser/removal race), with latch-based tests for both
interleavings.
##########
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:
[P2] Define zero-quota authentication before building the cache
With `qe_max_connection=1` and the new `-1` default, the effective Flight
quota is 0 (the new pool test asserts this), so this constructs a Guava
`maximumSize(0)` token cache. Every newly issued token is evicted immediately;
basic auth can return it, but the first RPC is rejected as `UNAUTHENTICATED`
before `registerConnection` can return the documented quota error, and the
removal callback runs before that user's cache is created. Handle quota 0
explicitly (or retain a token long enough to produce `RESOURCE_EXHAUSTED`) and
add an end-to-end boundary test.
##########
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:
[P2] Release the retained Arrow buffer after this assertion
This `retain()` creates the extra native reference needed to make
`allocator.close()` fail, but the test never closes `held`. After the cached
`VectorSchemaRoot` is invalidated, that retained reference is the remaining
owner, and `releaseSession` intentionally swallows the allocator-close failure,
so the native allocation survives the test. Keep the failure assertions, then
release `held` in a `finally` block so the rest of the FE test JVM is not
contaminated.
--
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]
