fix(agent): avoid validation contention on active JDBC pools

Closes #5567
This commit is contained in:
t8y2 2026-08-09 11:59:30 +08:00
parent 2b3bf1367d
commit d72d4beed3
No known key found for this signature in database
5 changed files with 73 additions and 4 deletions

View File

@ -278,6 +278,10 @@ public abstract class AbstractJdbcAgent extends BaseDatabaseAgent {
return poolRegistry != null;
}
final synchronized boolean hasActivePooledLeases() {
return poolRegistry != null && poolIdentity != null && poolRegistry.hasActiveLeases(poolIdentity);
}
final synchronized boolean quarantinePooledConnection() {
pooledConnectionPoisoned = true;
return requestActive && pooledLease != null && pooledLease.quarantine();

View File

@ -125,6 +125,11 @@ final class JdbcConnectionPoolRegistry implements AutoCloseable {
return physicalConnectionBudget.activeCount();
}
boolean hasActiveLeases(String identity) {
PoolEntry entry = pools.get(digest(identity));
return entry != null && entry.hasActiveLeases();
}
private PoolEntry createPoolEntry(String key, ConnectionFactory connectionFactory) {
try {
ConnectionFactoryDataSource factoryDataSource = new ConnectionFactoryDataSource(
@ -695,6 +700,10 @@ final class JdbcConnectionPoolRegistry implements AutoCloseable {
return retired;
}
private synchronized boolean hasActiveLeases() {
return activeLeases > 0;
}
private void retireAfterCheckoutFailure(OperationDeadline deadline) {
synchronized (this) {
retired = true;

View File

@ -91,6 +91,11 @@ public final class JsonRpcServer {
Object dispatchForRuntime(String method, JsonObject params) throws Exception {
return AgentExecutionContext.withJdbcExecutor(jdbcExecutor, () -> {
AbstractJdbcAgent jdbcAgent = pooledJdbcAgent();
if (AgentProtocol.METHOD_VALIDATE_CONNECTION.equals(method)
&& jdbcAgent != null
&& jdbcAgent.hasActivePooledLeases()) {
return Collections.singletonMap("ok", true);
}
boolean manageConnection = jdbcAgent != null && requiresConnectedConnection(method);
if (manageConnection) {
jdbcAgent.beginPooledRequest();

View File

@ -1141,6 +1141,55 @@ class JdbcConnectionPoolingTest {
}
}
@Test
void validationSkipsBusySharedPoolWithoutWaiting() throws Exception {
AtomicInteger physicalOpens = new AtomicInteger();
AtomicInteger requestIds = new AtomicInteger();
String url = h2Url("busy_validation");
try (MultiSessionJsonRpcServer server = server(url, physicalOpens, 1)) {
openSession(server, requestIds, "cursor-owner");
openSession(server, requestIds, "validation-session");
JsonObject pageParams = sessionParams("cursor-owner");
pageParams.addProperty("sql", "SELECT X FROM SYSTEM_RANGE(1, 3)");
pageParams.addProperty("pageSize", 1);
JsonObject firstPage = result(request(
server,
requestIds,
AgentProtocol.METHOD_EXECUTE_QUERY_PAGE,
pageParams
));
assertTrue(firstPage.get("has_more").getAsBoolean());
String querySessionId = firstPage.get("session_id").getAsString();
long startedAtNanos = System.nanoTime();
JsonObject validation = result(request(
server,
requestIds,
AgentProtocol.METHOD_VALIDATE_SESSION,
sessionParams("validation-session")
));
long elapsedMillis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startedAtNanos);
assertTrue(elapsedMillis < 200L, () -> "busy validation took " + elapsedMillis + "ms");
assertTrue(validation.get("ok").getAsBoolean());
assertEquals(1, physicalOpens.get());
JsonObject closeParams = sessionParams("cursor-owner");
closeParams.addProperty("sessionId", querySessionId);
assertTrue(request(
server,
requestIds,
AgentProtocol.METHOD_CLOSE_QUERY_SESSION,
closeParams
).get("result").getAsBoolean());
assertEquals(
2,
query(server, requestIds, "validation-session", "SELECT 2", null)
.getAsJsonArray("rows").get(0).getAsJsonArray().get(0).getAsInt()
);
}
}
@Test
void maintenanceExpiresAbandonedCursorAndReturnsItsConnection() throws Exception {
AtomicInteger physicalOpens = new AtomicInteger();

View File

@ -4726,10 +4726,10 @@ fn uses_agent_connection_pool(db_type: &DatabaseType) -> bool {
}
fn should_validate_existing_pool_before_reuse(db_type: DatabaseType) -> bool {
// PostgreSQL uses deadpool's Fast recycling and the query executor's
// ReconnectAndRetry path. An eager SELECT 1 here would add a network
// round-trip before every query without improving recovery behavior.
!matches!(db_type, DatabaseType::Postgres | DatabaseType::Etcd)
// PostgreSQL and Agent-backed databases validate connections when they are
// checked out for actual work. An eager probe here would add a database
// round-trip before every request and can compete with active Agent leases.
db_type != DatabaseType::Postgres && !matches!(db_type, agent_connection_pool_database_type!())
}
fn agent_pool_identity(pool: &PoolKind) -> Option<Arc<db::agent_driver::PooledAgentClient>> {
@ -5324,6 +5324,8 @@ mod tests {
fn drivers_with_internal_recovery_skip_eager_pool_validation() {
assert!(!super::should_validate_existing_pool_before_reuse(DatabaseType::Postgres));
assert!(!super::should_validate_existing_pool_before_reuse(DatabaseType::Etcd));
assert!(!super::should_validate_existing_pool_before_reuse(DatabaseType::Dameng));
assert!(!super::should_validate_existing_pool_before_reuse(DatabaseType::Oracle));
assert!(super::should_validate_existing_pool_before_reuse(DatabaseType::Mysql));
}