diff --git a/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java b/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java index dc2ffea5e..716652397 100644 --- a/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java +++ b/agents/common/src/main/java/com/dbx/agent/AbstractJdbcAgent.java @@ -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(); diff --git a/agents/common/src/main/java/com/dbx/agent/JdbcConnectionPoolRegistry.java b/agents/common/src/main/java/com/dbx/agent/JdbcConnectionPoolRegistry.java index 57f7ca3d2..514dc1405 100644 --- a/agents/common/src/main/java/com/dbx/agent/JdbcConnectionPoolRegistry.java +++ b/agents/common/src/main/java/com/dbx/agent/JdbcConnectionPoolRegistry.java @@ -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; diff --git a/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java b/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java index 806812c33..f88dc2d37 100644 --- a/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java +++ b/agents/common/src/main/java/com/dbx/agent/JsonRpcServer.java @@ -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(); diff --git a/agents/common/src/test/java/com/dbx/agent/JdbcConnectionPoolingTest.java b/agents/common/src/test/java/com/dbx/agent/JdbcConnectionPoolingTest.java index dcc47ef8b..6a8ccaeeb 100644 --- a/agents/common/src/test/java/com/dbx/agent/JdbcConnectionPoolingTest.java +++ b/agents/common/src/test/java/com/dbx/agent/JdbcConnectionPoolingTest.java @@ -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(); diff --git a/crates/dbx-core/src/connection.rs b/crates/dbx-core/src/connection.rs index 85f61ff00..bc29f2440 100644 --- a/crates/dbx-core/src/connection.rs +++ b/crates/dbx-core/src/connection.rs @@ -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> { @@ -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)); }