From 24f70d2a6ff2e28ae42dcbcce328b1625a1d3c5e Mon Sep 17 00:00:00 2001 From: Diego Fanesi Date: Sun, 16 Aug 2026 13:13:29 +0200 Subject: [PATCH 1/6] fix(thrift): retry transient transport failures on status/close/cancel polls Long-running queries on the Thrift client path (UseThriftClient=1) could fail with "Query has been timed out due to inactivity" under sustained concurrency. A single stale pooled connection or TCP reset on a GetOperationStatus poll abandoned the still-running server operation, and the same failure class on CloseOperation/CancelOperation could leak completed operations until the server reaped them. These idempotent RPCs now run through a bounded, jittered exponential-backoff retry that reconnects on a fresh pooled connection before surfacing the error. Statement submission is deliberately excluded from the retry path to avoid double-execution. The SEA client path is unchanged. Co-authored-by: Isaac --- NEXT_CHANGELOG.md | 1 + .../impl/thrift/DatabricksThriftAccessor.java | 105 +++++++++++++++++- 2 files changed, 102 insertions(+), 4 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index e12c8875f..51692ef6f 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -8,6 +8,7 @@ - `DatabaseMetaData.getColumns(...)` with a `null` catalog now issues a single `SHOW COLUMNS IN ALL CATALOGS` statement (consistent with `getSchemas`/`getTables`) instead of enumerating every catalog and issuing a per-catalog `SHOW COLUMNS`. Older DBR versions that do not support the syntax transparently fall back to the previous enumerate-and-fan-out behavior. ### Fixed +- Fixed long-running queries on the Thrift client path (`UseThriftClient=1`) intermittently failing with `Query has been timed out due to inactivity` under sustained concurrency. A transient transport-level failure (stale pooled connection, connection reset, load-balancer idle drop) on a `GetOperationStatus` poll previously abandoned the still-running server operation after a single blip, and the same failure class on `CloseOperation`/`CancelOperation` could leak completed operations until the server reaped them. These idempotent RPCs now transparently reconnect and retry on a fresh connection with bounded, jittered exponential backoff before surfacing the error. Statement submission is deliberately excluded from the retry path to avoid double-execution. The SEA client path is unchanged. - Fixed `IdleConnectionEvictor` thread leak in long-running applications. Driver-side resources (HTTP client, background threads) are now always released when `Connection.close()` is called, even if statement cleanup or server-side session termination fails. - Throw `DatabricksSQLException` instead of an unchecked `ClassCastException` when a complex-type getter (`getArray`, `getStruct`, `getMap`) is called on a column of a different complex type. diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java index 18c126b2b..a1c38d416 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java @@ -31,6 +31,7 @@ import java.sql.SQLException; import java.util.Arrays; import java.util.Objects; +import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; import org.apache.http.HttpException; import org.apache.thrift.TBase; @@ -53,6 +54,16 @@ final class DatabricksThriftAccessor { TExecuteStatementResp._Fields.OPERATION_HANDLE.getThriftFieldId(); private static final short statusFieldId = TExecuteStatementResp._Fields.STATUS.getThriftFieldId(); + // Bounded, jittered retry for transient transport-level failures (stale pooled connection, + // connection reset, load-balancer idle drop) on idempotent status / close / cancel RPCs. A + // status poll is read-only, and closing or cancelling an operation is idempotent, so repeating + // any of them produces no additional server-side side effects. This is why statement submission + // is deliberately NOT routed through the retry path: re-sending an ExecuteStatement could run the + // query twice. + private static final int MAX_POLL_TRANSPORT_RETRIES = 5; + private static final long TRANSPORT_RETRY_MIN_BACKOFF_MILLIS = 1_000L; + private static final long TRANSPORT_RETRY_MAX_BACKOFF_MILLIS = 16_000L; + private DatabricksConfig databricksConfig; private final boolean enableDirectResults; private final int asyncPollIntervalMillis; @@ -163,7 +174,10 @@ TFetchResultsResp getResultSetResp(TOperationHandle operationHandle, long startR TCancelOperationResp cancelOperation(TCancelOperationReq req) throws DatabricksHttpException { try { - return getThriftClient().CancelOperation(req); + return withTransportRetry( + "CancelOperation", + loggableOperationHandle(req.getOperationHandle()), + () -> getThriftClient().CancelOperation(req)); } catch (TException e) { String errorMessage = String.format( @@ -176,7 +190,10 @@ TCancelOperationResp cancelOperation(TCancelOperationReq req) throws DatabricksH TCloseOperationResp closeOperation(TCloseOperationReq req) throws DatabricksHttpException { try { - return getThriftClient().CloseOperation(req); + return withTransportRetry( + "CloseOperation", + loggableOperationHandle(req.getOperationHandle()), + () -> getThriftClient().CloseOperation(req)); } catch (TException e) { String errorMessage = String.format( @@ -778,7 +795,11 @@ TFetchResultsResp fetchMetadataResults(TResp response, String contextDescription while (shouldContinuePolling(statusResp)) { metadataTimeoutHandler.checkTimeout(); try { - statusResp = getThriftClient().GetOperationStatus(statusReq); + statusResp = + withTransportRetry( + "GetOperationStatus", + statementId, + () -> getThriftClient().GetOperationStatus(statusReq)); } catch (TTransportException e) { throw buildTransportFailureException(statementId, e); } @@ -962,6 +983,78 @@ private DatabricksSQLException buildTransportFailureException( return new DatabricksSQLException(errorMsg, e, COMMUNICATION_LINK_FAILURE_SQLSTATE); } + /** A Thrift RPC that is safe to repeat after a transport-level failure. */ + @FunctionalInterface + private interface TransportSafeRpc { + T call() throws TException; + } + + /** Null-safe operation-handle rendering for log lines (handles may be absent). */ + private static String loggableOperationHandle(TOperationHandle operationHandle) { + return operationHandle != null + ? StatementId.loggableStatementId(operationHandle) + : "unknown"; + } + + /** + * Executes an idempotent Thrift RPC, transparently retrying transient transport-level failures on + * a fresh connection with jittered exponential backoff. + * + *

Every invocation of {@code rpc} builds a new transport, so a retry naturally leases a + * different pooled connection while the broken one is discarded. This lets a still-running + * server-side operation be re-polled — or a completed one be re-closed / re-cancelled — instead + * of being abandoned after a single stale-connection blip and left to expire on the server's + * inactivity timeout. Only RPCs that are safe to repeat may use this path (status polling, + * operation close, cancel); statement submission must not. + * + *

Retries are bounded by {@link #MAX_POLL_TRANSPORT_RETRIES}. Once exhausted, the original + * {@link TTransportException} is rethrown so the existing caller-side failure handling still + * applies. Non-transport {@link TException}s are never retried — they propagate on the first + * attempt. A thread interrupt during a backoff sleep restores the interrupt flag and aborts the + * retry loop. + */ + private T withTransportRetry(String rpcName, String statementId, TransportSafeRpc rpc) + throws TException { + int attempt = 0; + long backoffMillis = TRANSPORT_RETRY_MIN_BACKOFF_MILLIS; + while (true) { + try { + return rpc.call(); + } catch (TTransportException e) { + if (++attempt > MAX_POLL_TRANSPORT_RETRIES) { + LOGGER.error( + "Transport failure on {} for statement [{}] still failing after {} retries; giving" + + " up. Cause: {}", + rpcName, + statementId, + MAX_POLL_TRANSPORT_RETRIES, + e.getMessage()); + throw e; + } + // Full-jitter backoff around the current exponential ceiling, spreading concurrent + // reconnect attempts so they do not thunder against a recovering endpoint. + long sleepMillis = + ThreadLocalRandom.current().nextLong(backoffMillis / 2 + 1, backoffMillis + 1); + LOGGER.warn( + "Transport failure on {} for statement [{}] (attempt {}/{}); reconnecting and retrying" + + " in {} ms. Cause: {}", + rpcName, + statementId, + attempt, + MAX_POLL_TRANSPORT_RETRIES, + sleepMillis, + e.getMessage()); + try { + TimeUnit.MILLISECONDS.sleep(sleepMillis); + } catch (InterruptedException ie) { + Thread.currentThread().interrupt(); + throw e; + } + backoffMillis = Math.min(backoffMillis * 2, TRANSPORT_RETRY_MAX_BACKOFF_MILLIS); + } + } + } + private boolean shouldContinuePolling(TGetOperationStatusResp statusResp) { return statusResp == null || !statusResp.isSetOperationState() @@ -1042,7 +1135,11 @@ private TimeoutHandler getTimeoutHandler( TGetOperationStatusResp getOperationStatus( TGetOperationStatusReq statusReq, StatementId statementId) throws TException { long operationStatusStartTime = System.nanoTime(); - TGetOperationStatusResp operationStatus = getThriftClient().GetOperationStatus(statusReq); + TGetOperationStatusResp operationStatus = + withTransportRetry( + "GetOperationStatus", + statementId.toSQLExecStatementId(), + () -> getThriftClient().GetOperationStatus(statusReq)); long operationStatusEndTime = System.nanoTime(); long operationStatusLatencyMillis = (operationStatusEndTime - operationStatusStartTime) / 1_000_000; From 9c3a6d1335cbd4a164b7f7f4ecf518619b42f9e7 Mon Sep 17 00:00:00 2001 From: Diego Fanesi Date: Wed, 26 Aug 2026 10:06:41 +0200 Subject: [PATCH 2/6] fix(thrift): discriminate transient vs permanent transport-poll failures and honor queryTimeout Addresses reviewer feedback on the transport-poll retry (PR #1643). The retry now classifies each TTransportException by its (normalized) DatabricksHttpException cause instead of retrying every failure: - genuine connection-level failures (stale pooled connection, reset, socket timeout) and transient HTTP gateway codes (408/500/502/504) are retried; - permanent HTTP errors (401/403/404) surface immediately; - 429/503 already retried and exhausted by the shared DatabricksHttpRetryHandler (they arrive with a DatabricksRetryHandlerException in the cause chain) are not re-retried, avoiding a second retry storm on top of the HTTP layer's. The retry also honors the statement's queryTimeout: the deadline is checked before each backoff sleep and the sleep is capped to the remaining budget, so a failing poll can no longer overshoot its deadline by the backoff schedule. To support classification, DatabricksHttpException now carries the HTTP status code (populated in ValidationUtil.checkHTTPError), and DatabricksHttpClient preserves the original IOException cause so connection-level failures are identifiable. Co-authored-by: Isaac Signed-off-by: Diego Fanesi --- NEXT_CHANGELOG.md | 2 +- .../jdbc/common/util/ValidationUtil.java | 3 +- .../dbclient/impl/common/TimeoutHandler.java | 13 ++ .../impl/http/DatabricksHttpClient.java | 5 +- .../impl/thrift/DatabricksThriftAccessor.java | 158 ++++++++++++++-- .../exception/DatabricksHttpException.java | 23 +++ .../jdbc/common/util/ValidationUtilTest.java | 3 +- .../thrift/DatabricksThriftAccessorTest.java | 169 ++++++++++++++++++ 8 files changed, 357 insertions(+), 19 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 51692ef6f..3597cb64d 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -8,7 +8,7 @@ - `DatabaseMetaData.getColumns(...)` with a `null` catalog now issues a single `SHOW COLUMNS IN ALL CATALOGS` statement (consistent with `getSchemas`/`getTables`) instead of enumerating every catalog and issuing a per-catalog `SHOW COLUMNS`. Older DBR versions that do not support the syntax transparently fall back to the previous enumerate-and-fan-out behavior. ### Fixed -- Fixed long-running queries on the Thrift client path (`UseThriftClient=1`) intermittently failing with `Query has been timed out due to inactivity` under sustained concurrency. A transient transport-level failure (stale pooled connection, connection reset, load-balancer idle drop) on a `GetOperationStatus` poll previously abandoned the still-running server operation after a single blip, and the same failure class on `CloseOperation`/`CancelOperation` could leak completed operations until the server reaped them. These idempotent RPCs now transparently reconnect and retry on a fresh connection with bounded, jittered exponential backoff before surfacing the error. Statement submission is deliberately excluded from the retry path to avoid double-execution. The SEA client path is unchanged. +- Fixed long-running queries on the Thrift client path (`UseThriftClient=1`) intermittently failing with `Query has been timed out due to inactivity` under sustained concurrency. A transient transport-level failure (stale pooled connection, connection reset, load-balancer idle drop) on a `GetOperationStatus` poll previously abandoned the still-running server operation after a single blip, and the same failure class on `CloseOperation`/`CancelOperation` could leak completed operations until the server reaped them. These idempotent RPCs now transparently reconnect and retry on a fresh connection with bounded, jittered exponential backoff before surfacing the error. Only genuinely transient failures are retried — connection-level errors and transient gateway responses (408/500/502/504); permanent HTTP errors (401/403/404) surface immediately, and rate-limit/unavailable responses (429/503) are left to the existing HTTP-layer retry rather than retried again. The retry also honors the statement's `queryTimeout` so a failing poll cannot overshoot its deadline. Statement submission is deliberately excluded from the retry path to avoid double-execution. The SEA client path is unchanged. - Fixed `IdleConnectionEvictor` thread leak in long-running applications. Driver-side resources (HTTP client, background threads) are now always released when `Connection.close()` is called, even if statement cleanup or server-side session termination fails. - Throw `DatabricksSQLException` instead of an unchecked `ClassCastException` when a complex-type getter (`getArray`, `getStruct`, `getMap`) is called on a column of a different complex type. diff --git a/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java b/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java index 4c2995310..24547fef9 100644 --- a/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java +++ b/src/main/java/com/databricks/jdbc/common/util/ValidationUtil.java @@ -138,7 +138,8 @@ public static void checkHTTPError(HttpResponse response) return; } LOGGER.error(errorReason); - throw new DatabricksHttpException(errorReason, DEFAULT_HTTP_EXCEPTION_SQLSTATE); + int statusCode = response.getStatusLine().getStatusCode(); + throw new DatabricksHttpException(errorReason, DEFAULT_HTTP_EXCEPTION_SQLSTATE, statusCode); } /** diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java b/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java index edc85e454..398c9e1ba 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java @@ -72,6 +72,19 @@ public void checkTimeout() throws DatabricksTimeoutException { } } + /** + * Returns the time remaining before this operation's deadline, in milliseconds. Returns {@link + * Long#MAX_VALUE} when no timeout is configured ({@code timeoutSeconds <= 0}), and may return a + * negative value once the deadline has already passed. + */ + public long getRemainingMillis() { + if (timeoutSeconds <= 0) { + return Long.MAX_VALUE; + } + long elapsedMillis = System.currentTimeMillis() - startTimeMillis; + return TimeUnit.SECONDS.toMillis(timeoutSeconds) - elapsedMillis; + } + /** * Factory method to create a timeout handler for a databricks client with a statement ID. This * works with any client that implements {@link IDatabricksClient}. diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java b/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java index 966750225..de0c8e1a9 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/http/DatabricksHttpClient.java @@ -198,7 +198,10 @@ private static void throwHttpException(Exception e, HttpUriRequest request) "Caught error while executing http request: [%s]. Error Message: [%s]", RequestSanitizer.sanitizeRequest(request), e); LOGGER.error(e, errorMsg); - throw new DatabricksHttpException(errorMsg, DEFAULT_HTTP_EXCEPTION_SQLSTATE); + // Preserve the original cause (typically a connection-level IOException) so callers can + // distinguish a genuine transport failure from an HTTP-status error and decide whether the + // operation is safe to retry. + throw new DatabricksHttpException(errorMsg, e, DEFAULT_HTTP_EXCEPTION_SQLSTATE); } @VisibleForTesting diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java index a1c38d416..7fc321a2d 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java @@ -28,9 +28,12 @@ import com.databricks.jdbc.telemetry.TelemetryHelper; import com.databricks.sdk.core.DatabricksConfig; import com.databricks.sdk.service.sql.StatementState; +import com.google.common.annotations.VisibleForTesting; +import java.io.IOException; import java.sql.SQLException; import java.util.Arrays; import java.util.Objects; +import java.util.Set; import java.util.concurrent.ThreadLocalRandom; import java.util.concurrent.TimeUnit; import org.apache.http.HttpException; @@ -63,6 +66,12 @@ final class DatabricksThriftAccessor { private static final int MAX_POLL_TRANSPORT_RETRIES = 5; private static final long TRANSPORT_RETRY_MIN_BACKOFF_MILLIS = 1_000L; private static final long TRANSPORT_RETRY_MAX_BACKOFF_MILLIS = 16_000L; + // Transient HTTP gateway codes that the shared DatabricksHttpRetryHandler does NOT itself retry + // (it only retries 429/503 + configured custom codes). A poll that hits one of these received a + // real HTTP response but from a transiently-unhealthy hop, so re-polling on a fresh connection is + // safe and worthwhile. 429/503 are deliberately excluded here: they are owned by the HTTP layer, + // and re-retrying them would multiply load on a recovering endpoint. + private static final Set RETRYABLE_TRANSPORT_HTTP_CODES = Set.of(408, 500, 502, 504); private DatabricksConfig databricksConfig; private final boolean enableDirectResults; @@ -358,9 +367,10 @@ private TGetOperationStatusResp pollTillOperationFinished( timeoutHandler.checkTimeout(); // TTransportException means a transport-level failure (e.g. HTTP 502 Bad Gateway) - // after retries were exhausted. Other TException subtypes propagate unchanged. + // after retries were exhausted. Other TException subtypes propagate unchanged. The timeout + // handler is threaded in so the retry backoff cannot overshoot the statement's queryTimeout. try { - statusResp = getOperationStatus(statusReq, statementId); + statusResp = getOperationStatus(statusReq, statementId, timeoutHandler); } catch (TTransportException e) { throw buildTransportFailureException(statementId.toSQLExecStatementId(), e); } @@ -799,6 +809,7 @@ TFetchResultsResp fetchMetadataResults(TResp response, String contextDescription withTransportRetry( "GetOperationStatus", statementId, + metadataTimeoutHandler, () -> getThriftClient().GetOperationStatus(statusReq)); } catch (TTransportException e) { throw buildTransportFailureException(statementId, e); @@ -997,8 +1008,24 @@ private static String loggableOperationHandle(TOperationHandle operationHandle) } /** - * Executes an idempotent Thrift RPC, transparently retrying transient transport-level failures on - * a fresh connection with jittered exponential backoff. + * Executes an idempotent Thrift RPC with no query-timeout budget (used by {@code CloseOperation} + * / {@code CancelOperation}, which are not bounded by a statement timeout). See {@link + * #withTransportRetry(String, String, TimeoutHandler, TransportSafeRpc)} for the full contract. + */ + private T withTransportRetry(String rpcName, String statementId, TransportSafeRpc rpc) + throws TException { + try { + return withTransportRetry(rpcName, statementId, /* timeoutHandler= */ null, rpc); + } catch (DatabricksTimeoutException e) { + // Unreachable: a null timeout handler never enforces a deadline. Guard defensively so the + // checked timeout type cannot silently widen this method's contract. + throw new IllegalStateException("Unexpected timeout without an active timeout handler", e); + } + } + + /** + * Executes an idempotent Thrift RPC, transparently retrying transient transport-level + * failures on a fresh connection with jittered exponential backoff. * *

Every invocation of {@code rpc} builds a new transport, so a retry naturally leases a * different pooled connection while the broken one is discarded. This lets a still-running @@ -1007,20 +1034,36 @@ private static String loggableOperationHandle(TOperationHandle operationHandle) * inactivity timeout. Only RPCs that are safe to repeat may use this path (status polling, * operation close, cancel); statement submission must not. * - *

Retries are bounded by {@link #MAX_POLL_TRANSPORT_RETRIES}. Once exhausted, the original - * {@link TTransportException} is rethrown so the existing caller-side failure handling still - * applies. Non-transport {@link TException}s are never retried — they propagate on the first - * attempt. A thread interrupt during a backoff sleep restores the interrupt flag and aborts the - * retry loop. + *

Only failures classified as transient by {@link #isRetryableTransportFailure} are retried: + * genuine connection-level errors (stale pooled connection, reset, socket timeout) and transient + * HTTP gateway codes ({@link #RETRYABLE_TRANSPORT_HTTP_CODES}). Permanent HTTP errors (401/403/404 + * …) and anything the shared {@link + * com.databricks.jdbc.dbclient.impl.http.DatabricksHttpRetryHandler} already retried and + * exhausted (429/503/custom, which surface with a {@link DatabricksRetryHandlerException} in their + * cause chain) are rethrown on the first attempt — the latter avoids stacking a second retry + * storm on top of the HTTP layer's. + * + *

When {@code timeoutHandler} is non-null the operation's deadline is enforced before each + * backoff sleep and the sleep is capped to the remaining budget, so a failing RPC cannot overshoot + * the statement's {@code queryTimeout}. Retries are bounded by {@link #MAX_POLL_TRANSPORT_RETRIES}; + * once exhausted the original {@link TTransportException} is rethrown so existing caller-side + * handling still applies. A thread interrupt during a backoff sleep restores the interrupt flag + * and aborts the retry loop. */ - private T withTransportRetry(String rpcName, String statementId, TransportSafeRpc rpc) - throws TException { + private T withTransportRetry( + String rpcName, String statementId, TimeoutHandler timeoutHandler, TransportSafeRpc rpc) + throws TException, DatabricksTimeoutException { int attempt = 0; long backoffMillis = TRANSPORT_RETRY_MIN_BACKOFF_MILLIS; while (true) { try { return rpc.call(); } catch (TTransportException e) { + if (!isRetryableTransportFailure(e)) { + // Permanent error (e.g. 401/403/404) or one the HTTP layer already retried (429/503): + // surface immediately instead of hanging through the backoff schedule. + throw e; + } if (++attempt > MAX_POLL_TRANSPORT_RETRIES) { LOGGER.error( "Transport failure on {} for statement [{}] still failing after {} retries; giving" @@ -1031,10 +1074,19 @@ private T withTransportRetry(String rpcName, String statementId, TransportSa e.getMessage()); throw e; } + // Enforce the query deadline before sleeping so a failing RPC cannot overshoot the + // statement's queryTimeout by the backoff schedule (may run the timeout action and throw). + if (timeoutHandler != null) { + timeoutHandler.checkTimeout(); + } // Full-jitter backoff around the current exponential ceiling, spreading concurrent // reconnect attempts so they do not thunder against a recovering endpoint. long sleepMillis = ThreadLocalRandom.current().nextLong(backoffMillis / 2 + 1, backoffMillis + 1); + if (timeoutHandler != null) { + // Never sleep past the deadline; the next checkTimeout() will then fire promptly. + sleepMillis = Math.min(sleepMillis, Math.max(0L, timeoutHandler.getRemainingMillis())); + } LOGGER.warn( "Transport failure on {} for statement [{}] (attempt {}/{}); reconnecting and retrying" + " in {} ms. Cause: {}", @@ -1045,7 +1097,7 @@ private T withTransportRetry(String rpcName, String statementId, TransportSa sleepMillis, e.getMessage()); try { - TimeUnit.MILLISECONDS.sleep(sleepMillis); + backoffSleep(sleepMillis); } catch (InterruptedException ie) { Thread.currentThread().interrupt(); throw e; @@ -1055,6 +1107,61 @@ private T withTransportRetry(String rpcName, String statementId, TransportSa } } + /** + * Classifies a transport failure as transient (safe to retry on a fresh connection) or not. + * + *

Because the Thrift transport routes through {@link + * com.databricks.jdbc.dbclient.impl.http.DatabricksHttpClient}, every failure arrives wrapped as a + * {@link TTransportException} whose cause is normally a {@link DatabricksHttpException}. The + * decision: + * + *

+ */ + private static boolean isRetryableTransportFailure(TTransportException e) { + Throwable cause = e.getCause(); + if (chainContains(cause, DatabricksRetryHandlerException.class)) { + return false; + } + if (cause instanceof DatabricksHttpException) { + int statusCode = ((DatabricksHttpException) cause).getStatusCode(); + if (statusCode != 0) { + return RETRYABLE_TRANSPORT_HTTP_CODES.contains(statusCode); + } + // No HTTP response was received: retry only when a real connection-level IOException is the + // underlying cause (guards against unrelated status-less DatabricksHttpExceptions). + return chainContains(cause.getCause(), IOException.class); + } + return cause instanceof IOException; + } + + /** Sleeps for the transport-retry backoff. Extracted as a test seam to keep retry tests fast. */ + @VisibleForTesting + void backoffSleep(long millis) throws InterruptedException { + TimeUnit.MILLISECONDS.sleep(millis); + } + + /** Returns true if {@code throwable} or any exception in its cause chain is of {@code type}. */ + private static boolean chainContains(Throwable throwable, Class type) { + for (Throwable current = throwable; current != null; current = current.getCause()) { + if (type.isInstance(current)) { + return true; + } + } + return false; + } + private boolean shouldContinuePolling(TGetOperationStatusResp statusResp) { return statusResp == null || !statusResp.isSetOperationState() @@ -1140,9 +1247,30 @@ TGetOperationStatusResp getOperationStatus( "GetOperationStatus", statementId.toSQLExecStatementId(), () -> getThriftClient().GetOperationStatus(statusReq)); - long operationStatusEndTime = System.nanoTime(); - long operationStatusLatencyMillis = - (operationStatusEndTime - operationStatusStartTime) / 1_000_000; + return recordOperationStatusLatency(statementId, operationStatusStartTime, operationStatus); + } + + /** + * Timeout-aware variant used by the execution polling loop: the retry backoff is bounded by the + * statement's {@code queryTimeout} via {@code timeoutHandler} so a transient transport failure + * cannot overshoot the deadline. + */ + TGetOperationStatusResp getOperationStatus( + TGetOperationStatusReq statusReq, StatementId statementId, TimeoutHandler timeoutHandler) + throws TException, DatabricksTimeoutException { + long operationStatusStartTime = System.nanoTime(); + TGetOperationStatusResp operationStatus = + withTransportRetry( + "GetOperationStatus", + statementId.toSQLExecStatementId(), + timeoutHandler, + () -> getThriftClient().GetOperationStatus(statusReq)); + return recordOperationStatusLatency(statementId, operationStatusStartTime, operationStatus); + } + + private TGetOperationStatusResp recordOperationStatusLatency( + StatementId statementId, long startTimeNanos, TGetOperationStatusResp operationStatus) { + long operationStatusLatencyMillis = (System.nanoTime() - startTimeNanos) / 1_000_000; LOGGER.debug( "Statement [{}] Thrift operation status latency: {}ms", statementId, diff --git a/src/main/java/com/databricks/jdbc/exception/DatabricksHttpException.java b/src/main/java/com/databricks/jdbc/exception/DatabricksHttpException.java index 3bb04c8d6..a8127581d 100644 --- a/src/main/java/com/databricks/jdbc/exception/DatabricksHttpException.java +++ b/src/main/java/com/databricks/jdbc/exception/DatabricksHttpException.java @@ -5,20 +5,43 @@ /** Exception class to handle http errors while downloading chunk data from external links. */ public class DatabricksHttpException extends DatabricksSQLException { + /** + * HTTP status code that produced this exception, or {@code 0} when no HTTP response was received + * (e.g. a connection-level failure) or the status is otherwise unknown. + */ + private final int statusCode; + public DatabricksHttpException( String message, Throwable cause, DatabricksDriverErrorCode sqlCode) { super(message, cause, sqlCode); + this.statusCode = 0; } public DatabricksHttpException(String message, DatabricksDriverErrorCode internalCode) { super(message, null, internalCode.toString()); + this.statusCode = 0; } public DatabricksHttpException(String message, String sqlState) { super(message, null, sqlState); + this.statusCode = 0; } public DatabricksHttpException(String message, Throwable throwable, String sqlState) { super(message, throwable, sqlState); + this.statusCode = 0; + } + + public DatabricksHttpException(String message, String sqlState, int statusCode) { + super(message, null, sqlState); + this.statusCode = statusCode; + } + + /** + * Returns the HTTP status code associated with this exception, or {@code 0} when no HTTP response + * was received (connection-level failure) or the status is unknown. + */ + public int getStatusCode() { + return statusCode; } } diff --git a/src/test/java/com/databricks/jdbc/common/util/ValidationUtilTest.java b/src/test/java/com/databricks/jdbc/common/util/ValidationUtilTest.java index 384aa7761..2ac7313d0 100644 --- a/src/test/java/com/databricks/jdbc/common/util/ValidationUtilTest.java +++ b/src/test/java/com/databricks/jdbc/common/util/ValidationUtilTest.java @@ -92,10 +92,11 @@ void testUnsuccessfulResponseCheck() { when(response.getStatusLine()).thenReturn(statusLine); when(statusLine.getStatusCode()).thenReturn(400); when(statusLine.toString()).thenReturn("mockStatusLine"); - Throwable exception = + DatabricksHttpException exception = assertThrows(DatabricksHttpException.class, () -> ValidationUtil.checkHTTPError(response)); assertEquals( "HTTP request failed by code: 400, status line: mockStatusLine.", exception.getMessage()); + assertEquals(400, exception.getStatusCode()); when(statusLine.getStatusCode()).thenReturn(102); assertThrows(DatabricksHttpException.class, () -> ValidationUtil.checkHTTPError(response)); diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java index 58b65b5ba..8e590b658 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java @@ -14,14 +14,18 @@ import com.databricks.jdbc.common.StatementType; import com.databricks.jdbc.dbclient.impl.common.ClientConfigurator; import com.databricks.jdbc.dbclient.impl.common.StatementId; +import com.databricks.jdbc.dbclient.impl.common.TimeoutHandler; import com.databricks.jdbc.exception.DatabricksHttpException; import com.databricks.jdbc.exception.DatabricksParsingException; +import com.databricks.jdbc.exception.DatabricksRetryHandlerException; import com.databricks.jdbc.exception.DatabricksSQLException; import com.databricks.jdbc.exception.DatabricksTimeoutException; import com.databricks.jdbc.exception.DatabricksValidationException; import com.databricks.jdbc.model.client.thrift.generated.*; +import com.databricks.jdbc.model.telemetry.enums.DatabricksDriverErrorCode; import com.databricks.sdk.core.DatabricksConfig; import com.databricks.sdk.service.sql.StatementState; +import java.net.SocketException; import java.sql.SQLException; import java.sql.Statement; import java.util.ArrayList; @@ -1350,6 +1354,171 @@ void testExecute_remapsConcurrentModificationOnOperationStateBranchToSerializati assertEquals(1003, e.getErrorCode()); // EXECUTE_STATEMENT_FAILED stable code } + // --------------------------------------------------------------------------- + // Transport-retry classification and timeout behaviour (withTransportRetry). + // --------------------------------------------------------------------------- + + private static TTransportException transportError(Throwable cause) { + return new TTransportException(TTransportException.UNKNOWN, "transport failed", cause); + } + + @Test + void testTransportRetry_permanentHttpErrorSurfacesImmediately() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + // A 404 (invalid handle) carries a concrete HTTP status: permanent, must not be retried. + TTransportException permanent = + transportError( + new DatabricksHttpException("HTTP request failed by code: 404", "08000", 404)); + when(thriftClient.GetOperationStatus(operationStatusReq)).thenThrow(permanent); + StatementId statementId = StatementId.deserialize(TEST_STMT_ID); + + TTransportException thrown = + assertThrows( + TTransportException.class, + () -> accessor.getOperationStatus(operationStatusReq, statementId)); + + assertSame(permanent, thrown); + verify(thriftClient, times(1)).GetOperationStatus(operationStatusReq); + verify(accessor, never()).backoffSleep(anyLong()); + } + + @Test + void testTransportRetry_httpLayerHandledErrorNotReRetried() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + // 429/503 exhausted by the shared HTTP retry handler surface with a + // DatabricksRetryHandlerException in the cause chain — the outer loop must not pile on. + TTransportException exhausted = + transportError( + new DatabricksHttpException( + "Retry failure. HTTP response code: 429", + new DatabricksRetryHandlerException("rate limited", 429), + "08000")); + when(thriftClient.GetOperationStatus(operationStatusReq)).thenThrow(exhausted); + StatementId statementId = StatementId.deserialize(TEST_STMT_ID); + + assertThrows( + TTransportException.class, + () -> accessor.getOperationStatus(operationStatusReq, statementId)); + + verify(thriftClient, times(1)).GetOperationStatus(operationStatusReq); + verify(accessor, never()).backoffSleep(anyLong()); + } + + @Test + void testTransportRetry_transientGatewayRetriedThenSucceeds() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + // 502 is a transient gateway code the HTTP layer does not itself retry. + TTransportException gateway = + transportError( + new DatabricksHttpException("HTTP request failed by code: 502", "08000", 502)); + when(thriftClient.GetOperationStatus(operationStatusReq)) + .thenThrow(gateway) + .thenReturn(operationStatusFinishedResp); + StatementId statementId = StatementId.deserialize(TEST_STMT_ID); + + TGetOperationStatusResp resp = accessor.getOperationStatus(operationStatusReq, statementId); + + assertSame(operationStatusFinishedResp, resp); + verify(thriftClient, times(2)).GetOperationStatus(operationStatusReq); + verify(accessor, times(1)).backoffSleep(anyLong()); + } + + @Test + void testTransportRetry_connectionFailureRetriedThenSucceeds() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + // A stale pooled connection / reset surfaces as a status-less DatabricksHttpException whose + // cause is a connection-level IOException (SocketException) — the PR's core case. + TTransportException connectionFailure = + transportError( + new DatabricksHttpException( + "Caught error while executing http request", + new SocketException("Connection reset"), + "08000")); + when(thriftClient.GetOperationStatus(operationStatusReq)) + .thenThrow(connectionFailure) + .thenReturn(operationStatusFinishedResp); + StatementId statementId = StatementId.deserialize(TEST_STMT_ID); + + TGetOperationStatusResp resp = accessor.getOperationStatus(operationStatusReq, statementId); + + assertSame(operationStatusFinishedResp, resp); + verify(thriftClient, times(2)).GetOperationStatus(operationStatusReq); + } + + @Test + void testTransportRetry_exhaustsThenRethrowsOriginal() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + TTransportException gateway = + transportError( + new DatabricksHttpException("HTTP request failed by code: 502", "08000", 502)); + when(thriftClient.GetOperationStatus(operationStatusReq)).thenThrow(gateway); + StatementId statementId = StatementId.deserialize(TEST_STMT_ID); + + TTransportException thrown = + assertThrows( + TTransportException.class, + () -> accessor.getOperationStatus(operationStatusReq, statementId)); + + assertSame(gateway, thrown); + // 1 initial attempt + 5 bounded retries = 6 invocations, 5 backoff sleeps. + verify(thriftClient, times(6)).GetOperationStatus(operationStatusReq); + verify(accessor, times(5)).backoffSleep(anyLong()); + } + + @Test + void testTransportRetry_consultsTimeoutHandlerBeforeSleeping() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + TimeoutHandler timeoutHandler = mock(TimeoutHandler.class); + when(timeoutHandler.getRemainingMillis()).thenReturn(10_000L); + TTransportException gateway = + transportError( + new DatabricksHttpException("HTTP request failed by code: 502", "08000", 502)); + when(thriftClient.GetOperationStatus(operationStatusReq)) + .thenThrow(gateway) + .thenReturn(operationStatusFinishedResp); + StatementId statementId = StatementId.deserialize(TEST_STMT_ID); + + TGetOperationStatusResp resp = + accessor.getOperationStatus(operationStatusReq, statementId, timeoutHandler); + + assertSame(operationStatusFinishedResp, resp); + verify(timeoutHandler, atLeastOnce()).checkTimeout(); + } + + @Test + void testTransportRetry_timeoutAbortsRetryWithinDeadline() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + TimeoutHandler timeoutHandler = mock(TimeoutHandler.class); + doThrow( + new DatabricksTimeoutException( + "Statement execution timed-out", + /* cause= */ null, + DatabricksDriverErrorCode.STATEMENT_EXECUTION_TIMEOUT)) + .when(timeoutHandler) + .checkTimeout(); + TTransportException gateway = + transportError( + new DatabricksHttpException("HTTP request failed by code: 502", "08000", 502)); + when(thriftClient.GetOperationStatus(operationStatusReq)).thenThrow(gateway); + StatementId statementId = StatementId.deserialize(TEST_STMT_ID); + + assertThrows( + DatabricksTimeoutException.class, + () -> accessor.getOperationStatus(operationStatusReq, statementId, timeoutHandler)); + + // The first attempt fails (retryable), the deadline fires before any backoff sleep, so no + // further RPC attempts and no sleep — the retry cannot overshoot queryTimeout. + verify(thriftClient, times(1)).GetOperationStatus(operationStatusReq); + verify(accessor, never()).backoffSleep(anyLong()); + } + private TFetchResultsReq getFetchResultsRequest(boolean includeMetadata) throws DatabricksValidationException { TFetchResultsReq request = From 698007f3cf6d3b175c36e35e6ef8f6a16c036564 Mon Sep 17 00:00:00 2001 From: Diego Fanesi Date: Wed, 26 Aug 2026 10:09:19 +0200 Subject: [PATCH 3/6] fix(thrift): bound the retry budget for CloseOperation/CancelOperation Cleanup RPCs ran through the same 5-attempt poll retry as status polling, so during a transport outage closing a connection or statement (and the cancel invoked as a metadata-timeout action) could block for the full jittered backoff schedule before surfacing, stalling shutdown. Route CloseOperation/CancelOperation through a dedicated, tighter cleanup budget (MAX_CLEANUP_TRANSPORT_RETRIES = 2): a single transient blip is still retried past, but a sustained outage now surfaces quickly instead of hanging. Co-authored-by: Isaac Signed-off-by: Diego Fanesi --- NEXT_CHANGELOG.md | 2 +- .../impl/thrift/DatabricksThriftAccessor.java | 69 +++++++++++++++---- .../thrift/DatabricksThriftAccessorTest.java | 23 +++++++ 3 files changed, 79 insertions(+), 15 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 3597cb64d..e0375e5b4 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -8,7 +8,7 @@ - `DatabaseMetaData.getColumns(...)` with a `null` catalog now issues a single `SHOW COLUMNS IN ALL CATALOGS` statement (consistent with `getSchemas`/`getTables`) instead of enumerating every catalog and issuing a per-catalog `SHOW COLUMNS`. Older DBR versions that do not support the syntax transparently fall back to the previous enumerate-and-fan-out behavior. ### Fixed -- Fixed long-running queries on the Thrift client path (`UseThriftClient=1`) intermittently failing with `Query has been timed out due to inactivity` under sustained concurrency. A transient transport-level failure (stale pooled connection, connection reset, load-balancer idle drop) on a `GetOperationStatus` poll previously abandoned the still-running server operation after a single blip, and the same failure class on `CloseOperation`/`CancelOperation` could leak completed operations until the server reaped them. These idempotent RPCs now transparently reconnect and retry on a fresh connection with bounded, jittered exponential backoff before surfacing the error. Only genuinely transient failures are retried — connection-level errors and transient gateway responses (408/500/502/504); permanent HTTP errors (401/403/404) surface immediately, and rate-limit/unavailable responses (429/503) are left to the existing HTTP-layer retry rather than retried again. The retry also honors the statement's `queryTimeout` so a failing poll cannot overshoot its deadline. Statement submission is deliberately excluded from the retry path to avoid double-execution. The SEA client path is unchanged. +- Fixed long-running queries on the Thrift client path (`UseThriftClient=1`) intermittently failing with `Query has been timed out due to inactivity` under sustained concurrency. A transient transport-level failure (stale pooled connection, connection reset, load-balancer idle drop) on a `GetOperationStatus` poll previously abandoned the still-running server operation after a single blip, and the same failure class on `CloseOperation`/`CancelOperation` could leak completed operations until the server reaped them. These idempotent RPCs now transparently reconnect and retry on a fresh connection with bounded, jittered exponential backoff before surfacing the error. Only genuinely transient failures are retried — connection-level errors and transient gateway responses (408/500/502/504); permanent HTTP errors (401/403/404) surface immediately, and rate-limit/unavailable responses (429/503) are left to the existing HTTP-layer retry rather than retried again. The retry also honors the statement's `queryTimeout` so a failing poll cannot overshoot its deadline, and `CloseOperation`/`CancelOperation` use a tighter retry budget so cleanup on shutdown cannot stall. Statement submission is deliberately excluded from the retry path to avoid double-execution. The SEA client path is unchanged. - Fixed `IdleConnectionEvictor` thread leak in long-running applications. Driver-side resources (HTTP client, background threads) are now always released when `Connection.close()` is called, even if statement cleanup or server-side session termination fails. - Throw `DatabricksSQLException` instead of an unchecked `ClassCastException` when a complex-type getter (`getArray`, `getStruct`, `getMap`) is called on a column of a different complex type. diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java index 7fc321a2d..55567d90d 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java @@ -64,6 +64,10 @@ final class DatabricksThriftAccessor { // is deliberately NOT routed through the retry path: re-sending an ExecuteStatement could run the // query twice. private static final int MAX_POLL_TRANSPORT_RETRIES = 5; + // Cleanup RPCs (CloseOperation / CancelOperation) run on close / cancel / timeout paths that must + // not hang: one transient blip is worth reconnecting past, but a sustained outage should surface + // fast rather than stall shutdown. Hence a far smaller budget than the status-poll path. + private static final int MAX_CLEANUP_TRANSPORT_RETRIES = 2; private static final long TRANSPORT_RETRY_MIN_BACKOFF_MILLIS = 1_000L; private static final long TRANSPORT_RETRY_MAX_BACKOFF_MILLIS = 16_000L; // Transient HTTP gateway codes that the shared DatabricksHttpRetryHandler does NOT itself retry @@ -183,7 +187,7 @@ TFetchResultsResp getResultSetResp(TOperationHandle operationHandle, long startR TCancelOperationResp cancelOperation(TCancelOperationReq req) throws DatabricksHttpException { try { - return withTransportRetry( + return withCleanupTransportRetry( "CancelOperation", loggableOperationHandle(req.getOperationHandle()), () -> getThriftClient().CancelOperation(req)); @@ -199,7 +203,7 @@ TCancelOperationResp cancelOperation(TCancelOperationReq req) throws DatabricksH TCloseOperationResp closeOperation(TCloseOperationReq req) throws DatabricksHttpException { try { - return withTransportRetry( + return withCleanupTransportRetry( "CloseOperation", loggableOperationHandle(req.getOperationHandle()), () -> getThriftClient().CloseOperation(req)); @@ -1008,14 +1012,35 @@ private static String loggableOperationHandle(TOperationHandle operationHandle) } /** - * Executes an idempotent Thrift RPC with no query-timeout budget (used by {@code CloseOperation} - * / {@code CancelOperation}, which are not bounded by a statement timeout). See {@link - * #withTransportRetry(String, String, TimeoutHandler, TransportSafeRpc)} for the full contract. + * Executes an idempotent status-poll RPC with no query-timeout budget (used by the Thrift + * heartbeat / metadata-less status checks). Retries are bounded by {@link + * #MAX_POLL_TRANSPORT_RETRIES}. See {@link #withTransportRetry(String, String, TimeoutHandler, + * int, TransportSafeRpc)} for the full contract. */ private T withTransportRetry(String rpcName, String statementId, TransportSafeRpc rpc) throws TException { + return withoutTimeout(rpcName, statementId, MAX_POLL_TRANSPORT_RETRIES, rpc); + } + + /** + * Executes a cleanup RPC ({@code CloseOperation} / {@code CancelOperation}) with a deliberately + * small retry budget ({@link #MAX_CLEANUP_TRANSPORT_RETRIES}). Cleanup runs on close / cancel / + * timeout paths that must not hang: a single transient blip is worth reconnecting past, but during + * a sustained outage extra retries only delay shutdown, so the budget is far tighter than the + * status-poll path. + */ + private T withCleanupTransportRetry( + String rpcName, String statementId, TransportSafeRpc rpc) throws TException { + return withoutTimeout(rpcName, statementId, MAX_CLEANUP_TRANSPORT_RETRIES, rpc); + } + + /** Shared no-timeout entry point; adapts the deadline-aware core for callers with no deadline. */ + private T withoutTimeout( + String rpcName, String statementId, int maxRetries, TransportSafeRpc rpc) + throws TException { try { - return withTransportRetry(rpcName, statementId, /* timeoutHandler= */ null, rpc); + return withTransportRetry( + rpcName, statementId, /* timeoutHandler= */ null, maxRetries, rpc); } catch (DatabricksTimeoutException e) { // Unreachable: a null timeout handler never enforces a deadline. Guard defensively so the // checked timeout type cannot silently widen this method's contract. @@ -1023,6 +1048,18 @@ private T withTransportRetry(String rpcName, String statementId, TransportSa } } + /** + * Deadline-aware status-poll retry with the default poll budget ({@link + * #MAX_POLL_TRANSPORT_RETRIES}). See {@link #withTransportRetry(String, String, TimeoutHandler, + * int, TransportSafeRpc)}. + */ + private T withTransportRetry( + String rpcName, String statementId, TimeoutHandler timeoutHandler, TransportSafeRpc rpc) + throws TException, DatabricksTimeoutException { + return withTransportRetry( + rpcName, statementId, timeoutHandler, MAX_POLL_TRANSPORT_RETRIES, rpc); + } + /** * Executes an idempotent Thrift RPC, transparently retrying transient transport-level * failures on a fresh connection with jittered exponential backoff. @@ -1045,13 +1082,17 @@ private T withTransportRetry(String rpcName, String statementId, TransportSa * *

When {@code timeoutHandler} is non-null the operation's deadline is enforced before each * backoff sleep and the sleep is capped to the remaining budget, so a failing RPC cannot overshoot - * the statement's {@code queryTimeout}. Retries are bounded by {@link #MAX_POLL_TRANSPORT_RETRIES}; - * once exhausted the original {@link TTransportException} is rethrown so existing caller-side - * handling still applies. A thread interrupt during a backoff sleep restores the interrupt flag - * and aborts the retry loop. + * the statement's {@code queryTimeout}. Retries are bounded by {@code maxRetries}; once exhausted + * the original {@link TTransportException} is rethrown so existing caller-side handling still + * applies. A thread interrupt during a backoff sleep restores the interrupt flag and aborts the + * retry loop. */ private T withTransportRetry( - String rpcName, String statementId, TimeoutHandler timeoutHandler, TransportSafeRpc rpc) + String rpcName, + String statementId, + TimeoutHandler timeoutHandler, + int maxRetries, + TransportSafeRpc rpc) throws TException, DatabricksTimeoutException { int attempt = 0; long backoffMillis = TRANSPORT_RETRY_MIN_BACKOFF_MILLIS; @@ -1064,13 +1105,13 @@ private T withTransportRetry( // surface immediately instead of hanging through the backoff schedule. throw e; } - if (++attempt > MAX_POLL_TRANSPORT_RETRIES) { + if (++attempt > maxRetries) { LOGGER.error( "Transport failure on {} for statement [{}] still failing after {} retries; giving" + " up. Cause: {}", rpcName, statementId, - MAX_POLL_TRANSPORT_RETRIES, + maxRetries, e.getMessage()); throw e; } @@ -1093,7 +1134,7 @@ private T withTransportRetry( rpcName, statementId, attempt, - MAX_POLL_TRANSPORT_RETRIES, + maxRetries, sleepMillis, e.getMessage()); try { diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java index 8e590b658..345648578 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java @@ -1470,6 +1470,29 @@ void testTransportRetry_exhaustsThenRethrowsOriginal() throws Exception { verify(accessor, times(5)).backoffSleep(anyLong()); } + @Test + void testCleanupRetry_boundedForCancelOperation() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + TCancelOperationReq req = + new TCancelOperationReq() + .setOperationHandle( + new TOperationHandle() + .setOperationId(handleIdentifier) + .setOperationType(TOperationType.UNKNOWN)); + // A persistent transient transport failure on cleanup must not run the full poll budget. + TTransportException gateway = + transportError( + new DatabricksHttpException("HTTP request failed by code: 502", "08000", 502)); + when(thriftClient.CancelOperation(req)).thenThrow(gateway); + + assertThrows(DatabricksHttpException.class, () -> accessor.cancelOperation(req)); + + // 1 initial attempt + MAX_CLEANUP_TRANSPORT_RETRIES (2) = 3 total, far below the poll budget. + verify(thriftClient, times(3)).CancelOperation(req); + verify(accessor, times(2)).backoffSleep(anyLong()); + } + @Test void testTransportRetry_consultsTimeoutHandlerBeforeSleeping() throws Exception { setup(true); From 8b0d61f550f6d1a576bd03efcaf9542bd9c07661 Mon Sep 17 00:00:00 2001 From: Diego Fanesi Date: Wed, 26 Aug 2026 10:10:22 +0200 Subject: [PATCH 4/6] fix(thrift): keep the retry sleep-cap from busy-spinning near the deadline TimeoutHandler.checkTimeout() enforces on whole-second truncation (it throws once elapsed exceeds timeoutSeconds, i.e. at (timeoutSeconds+1)*1000 ms), but getRemainingMillis() reported time to the nominal timeoutSeconds*1000 boundary. In the sub-second window between the two, the retry's sleep cap (Math.max(0, remaining)) collapsed to 0, so the loop tight-spun rpc.call() with no backoff and burned the whole retry budget against a still-unhealthy endpoint. Report remaining time to checkTimeout()'s actual enforcement point so the cap stays positive until the deadline truly fires; the sleep is then bounded but never zero-length. Co-authored-by: Isaac Signed-off-by: Diego Fanesi --- .../dbclient/impl/common/TimeoutHandler.java | 14 +++++--- .../impl/thrift/DatabricksThriftAccessor.java | 8 +++-- .../impl/common/TimeoutHandlerTest.java | 32 +++++++++++++++++++ 3 files changed, 48 insertions(+), 6 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java b/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java index 398c9e1ba..0c2c1242a 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandler.java @@ -73,16 +73,22 @@ public void checkTimeout() throws DatabricksTimeoutException { } /** - * Returns the time remaining before this operation's deadline, in milliseconds. Returns {@link - * Long#MAX_VALUE} when no timeout is configured ({@code timeoutSeconds <= 0}), and may return a - * negative value once the deadline has already passed. + * Returns the time remaining, in milliseconds, before {@link #checkTimeout()} would actually + * enforce the deadline. Returns {@link Long#MAX_VALUE} when no timeout is configured ({@code + * timeoutSeconds <= 0}), and may return a negative value once the deadline has passed. + * + *

{@code checkTimeout()} compares whole (truncated) seconds and throws only once elapsed time + * exceeds {@code timeoutSeconds}, i.e. at {@code (timeoutSeconds + 1) * 1000} ms. This + * method reports the remaining time to that same enforcement point (not the nominal + * {@code timeoutSeconds * 1000}), so a caller capping a backoff sleep on it does not collapse the + * sleep to zero — and then busy-spin — during the sub-second window before enforcement fires. */ public long getRemainingMillis() { if (timeoutSeconds <= 0) { return Long.MAX_VALUE; } long elapsedMillis = System.currentTimeMillis() - startTimeMillis; - return TimeUnit.SECONDS.toMillis(timeoutSeconds) - elapsedMillis; + return TimeUnit.SECONDS.toMillis(timeoutSeconds + 1L) - elapsedMillis; } /** diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java index 55567d90d..02cacc080 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java @@ -1125,8 +1125,12 @@ private T withTransportRetry( long sleepMillis = ThreadLocalRandom.current().nextLong(backoffMillis / 2 + 1, backoffMillis + 1); if (timeoutHandler != null) { - // Never sleep past the deadline; the next checkTimeout() will then fire promptly. - sleepMillis = Math.min(sleepMillis, Math.max(0L, timeoutHandler.getRemainingMillis())); + // Cap to the time left before the deadline is actually enforced so we neither overshoot + // nor collapse to a zero-length (busy-spin) sleep in the sub-second window before it. + long remainingMillis = timeoutHandler.getRemainingMillis(); + if (remainingMillis < sleepMillis) { + sleepMillis = Math.max(0L, remainingMillis); + } } LOGGER.warn( "Transport failure on {} for statement [{}] (attempt {}/{}); reconnecting and retrying" diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandlerTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandlerTest.java index ae75478fc..a23a4ecd9 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandlerTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/common/TimeoutHandlerTest.java @@ -135,6 +135,38 @@ void testNullTimeoutAction() throws Exception { assertTrue(exception.getMessage().contains("timed-out after 2 seconds")); } + @Test + void testGetRemainingMillisStaysPositiveUntilEnforcement() throws Exception { + TimeoutHandler handler = + new TimeoutHandler( + 5, "Test operation", null, DatabricksDriverErrorCode.STATEMENT_EXECUTION_TIMEOUT); + Field startTimeField = TimeoutHandler.class.getDeclaredField("startTimeMillis"); + startTimeField.setAccessible(true); + long now = System.currentTimeMillis(); + + // Elapsed 5.5s: past the nominal 5s but before checkTimeout enforces (it fires only once elapsed + // exceeds 5 whole seconds, i.e. at 6s). getRemainingMillis() must stay POSITIVE in this window + // so a caller capping a backoff sleep on it does not collapse the sleep to zero and busy-spin. + startTimeField.set(handler, now - 5_500L); + assertDoesNotThrow(handler::checkTimeout); + assertTrue( + handler.getRemainingMillis() > 0, + "remaining should be > 0 in the sub-second window before enforcement"); + + // Well past enforcement: remaining goes negative and checkTimeout throws. + startTimeField.set(handler, now - 10_000L); + assertTrue(handler.getRemainingMillis() < 0); + assertThrows(DatabricksTimeoutException.class, handler::checkTimeout); + } + + @Test + void testGetRemainingMillisUnboundedWhenNoTimeout() { + TimeoutHandler handler = + new TimeoutHandler( + 0, "Test operation", null, DatabricksDriverErrorCode.STATEMENT_EXECUTION_TIMEOUT); + assertEquals(Long.MAX_VALUE, handler.getRemainingMillis()); + } + @Test void testForStatementFactory() throws Exception { when(mockStatementId.toString()).thenReturn("test-statement-id"); From d95e7af9545c279a976e99eb55ea65dbcde051a2 Mon Sep 17 00:00:00 2001 From: Diego Fanesi Date: Wed, 26 Aug 2026 10:11:14 +0200 Subject: [PATCH 5/6] fix(thrift): retry cause-less TTransportExceptions (truncated reads) isRetryableTransportFailure treated a TTransportException with no cause as non-retryable, but DatabricksHttpTTransport.read() raises bare, cause-less TTransportExceptions ("Response buffer is empty" / "No more data available.") on a truncated or empty response body. Those are plausibly transient and were retried before the classifier was added, so a truncated-response blip regressed to surfacing immediately. Treat a cause-less TTransportException as retryable, restoring the prior behavior for that class of failure. Co-authored-by: Isaac Signed-off-by: Diego Fanesi --- .../impl/thrift/DatabricksThriftAccessor.java | 9 ++++++++- .../thrift/DatabricksThriftAccessorTest.java | 18 ++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java index 02cacc080..aaa33fdd4 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java @@ -1171,11 +1171,18 @@ private T withTransportRetry( * connection-level failure (stale pooled connection, reset, socket timeout) surfaced as an * {@link IOException} cause; retryable. *

  • A direct {@link IOException} cause (e.g. a response-body read error) → retryable. - *
  • Anything else (unknown or absent cause) → not retryable. + *
  • No cause at all → retryable. The transport raises bare {@link TTransportException}s (e.g. + * {@code DatabricksHttpTTransport.read()} on a truncated / empty response body) that are + * plausibly transient; treating them as retryable preserves the pre-existing behavior of + * retrying every {@code TTransportException}. + *
  • A non-transient known cause (e.g. a permanent HTTP status) → not retryable. * */ private static boolean isRetryableTransportFailure(TTransportException e) { Throwable cause = e.getCause(); + if (cause == null) { + return true; + } if (chainContains(cause, DatabricksRetryHandlerException.class)) { return false; } diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java index 345648578..900a563b6 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java @@ -1470,6 +1470,24 @@ void testTransportRetry_exhaustsThenRethrowsOriginal() throws Exception { verify(accessor, times(5)).backoffSleep(anyLong()); } + @Test + void testTransportRetry_nullCauseRetriedThenSucceeds() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + // DatabricksHttpTTransport.read() raises bare TTransportExceptions (null cause) on a + // truncated/empty response body — plausibly transient, so they are retried. + TTransportException bare = new TTransportException("Response buffer is empty, no response."); + when(thriftClient.GetOperationStatus(operationStatusReq)) + .thenThrow(bare) + .thenReturn(operationStatusFinishedResp); + StatementId statementId = StatementId.deserialize(TEST_STMT_ID); + + TGetOperationStatusResp resp = accessor.getOperationStatus(operationStatusReq, statementId); + + assertSame(operationStatusFinishedResp, resp); + verify(thriftClient, times(2)).GetOperationStatus(operationStatusReq); + } + @Test void testCleanupRetry_boundedForCancelOperation() throws Exception { setup(true); From f204e97e60e5d4063d29590f3c473107579dbf6b Mon Sep 17 00:00:00 2001 From: Diego Fanesi Date: Wed, 26 Aug 2026 10:12:41 +0200 Subject: [PATCH 6/6] fix(thrift): stop treating HTTP 500 as a transient transport failure 500 Internal Server Error is usually a genuine server-side/operation error rather than a transient gateway-hop failure, so retrying it just wastes the budget (and, on the cleanup path, delays surfacing). Drop 500 from the retryable transport gateway codes, leaving 408/502/504; 500 now surfaces immediately like the other permanent errors. Co-authored-by: Isaac Signed-off-by: Diego Fanesi --- NEXT_CHANGELOG.md | 2 +- .../impl/thrift/DatabricksThriftAccessor.java | 7 ++++--- .../thrift/DatabricksThriftAccessorTest.java | 19 +++++++++++++++++++ 3 files changed, 24 insertions(+), 4 deletions(-) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index e0375e5b4..297c62e8e 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -8,7 +8,7 @@ - `DatabaseMetaData.getColumns(...)` with a `null` catalog now issues a single `SHOW COLUMNS IN ALL CATALOGS` statement (consistent with `getSchemas`/`getTables`) instead of enumerating every catalog and issuing a per-catalog `SHOW COLUMNS`. Older DBR versions that do not support the syntax transparently fall back to the previous enumerate-and-fan-out behavior. ### Fixed -- Fixed long-running queries on the Thrift client path (`UseThriftClient=1`) intermittently failing with `Query has been timed out due to inactivity` under sustained concurrency. A transient transport-level failure (stale pooled connection, connection reset, load-balancer idle drop) on a `GetOperationStatus` poll previously abandoned the still-running server operation after a single blip, and the same failure class on `CloseOperation`/`CancelOperation` could leak completed operations until the server reaped them. These idempotent RPCs now transparently reconnect and retry on a fresh connection with bounded, jittered exponential backoff before surfacing the error. Only genuinely transient failures are retried — connection-level errors and transient gateway responses (408/500/502/504); permanent HTTP errors (401/403/404) surface immediately, and rate-limit/unavailable responses (429/503) are left to the existing HTTP-layer retry rather than retried again. The retry also honors the statement's `queryTimeout` so a failing poll cannot overshoot its deadline, and `CloseOperation`/`CancelOperation` use a tighter retry budget so cleanup on shutdown cannot stall. Statement submission is deliberately excluded from the retry path to avoid double-execution. The SEA client path is unchanged. +- Fixed long-running queries on the Thrift client path (`UseThriftClient=1`) intermittently failing with `Query has been timed out due to inactivity` under sustained concurrency. A transient transport-level failure (stale pooled connection, connection reset, load-balancer idle drop) on a `GetOperationStatus` poll previously abandoned the still-running server operation after a single blip, and the same failure class on `CloseOperation`/`CancelOperation` could leak completed operations until the server reaped them. These idempotent RPCs now transparently reconnect and retry on a fresh connection with bounded, jittered exponential backoff before surfacing the error. Only genuinely transient failures are retried — connection-level errors and transient gateway responses (408/502/504); permanent HTTP errors (401/403/404, and 500) surface immediately, and rate-limit/unavailable responses (429/503) are left to the existing HTTP-layer retry rather than retried again. The retry also honors the statement's `queryTimeout` so a failing poll cannot overshoot its deadline, and `CloseOperation`/`CancelOperation` use a tighter retry budget so cleanup on shutdown cannot stall. Statement submission is deliberately excluded from the retry path to avoid double-execution. The SEA client path is unchanged. - Fixed `IdleConnectionEvictor` thread leak in long-running applications. Driver-side resources (HTTP client, background threads) are now always released when `Connection.close()` is called, even if statement cleanup or server-side session termination fails. - Throw `DatabricksSQLException` instead of an unchecked `ClassCastException` when a complex-type getter (`getArray`, `getStruct`, `getMap`) is called on a column of a different complex type. diff --git a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java index aaa33fdd4..4474fbbbb 100644 --- a/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java +++ b/src/main/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessor.java @@ -73,9 +73,10 @@ final class DatabricksThriftAccessor { // Transient HTTP gateway codes that the shared DatabricksHttpRetryHandler does NOT itself retry // (it only retries 429/503 + configured custom codes). A poll that hits one of these received a // real HTTP response but from a transiently-unhealthy hop, so re-polling on a fresh connection is - // safe and worthwhile. 429/503 are deliberately excluded here: they are owned by the HTTP layer, - // and re-retrying them would multiply load on a recovering endpoint. - private static final Set RETRYABLE_TRANSPORT_HTTP_CODES = Set.of(408, 500, 502, 504); + // safe and worthwhile. 429/503 are excluded (owned by the HTTP layer — re-retrying would multiply + // load on a recovering endpoint); 500 is excluded too, as it is typically a genuine server-side + // error rather than a transient hop failure and should surface rather than burn the retry budget. + private static final Set RETRYABLE_TRANSPORT_HTTP_CODES = Set.of(408, 502, 504); private DatabricksConfig databricksConfig; private final boolean enableDirectResults; diff --git a/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java b/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java index 900a563b6..25ea4fe11 100644 --- a/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java +++ b/src/test/java/com/databricks/jdbc/dbclient/impl/thrift/DatabricksThriftAccessorTest.java @@ -1470,6 +1470,25 @@ void testTransportRetry_exhaustsThenRethrowsOriginal() throws Exception { verify(accessor, times(5)).backoffSleep(anyLong()); } + @Test + void testTransportRetry_serverError500SurfacesImmediately() throws Exception { + setup(true); + doNothing().when(accessor).backoffSleep(anyLong()); + // 500 is treated as a (likely-permanent) server error, not a transient gateway hop failure. + TTransportException serverError = + transportError( + new DatabricksHttpException("HTTP request failed by code: 500", "08000", 500)); + when(thriftClient.GetOperationStatus(operationStatusReq)).thenThrow(serverError); + StatementId statementId = StatementId.deserialize(TEST_STMT_ID); + + assertThrows( + TTransportException.class, + () -> accessor.getOperationStatus(operationStatusReq, statementId)); + + verify(thriftClient, times(1)).GetOperationStatus(operationStatusReq); + verify(accessor, never()).backoffSleep(anyLong()); + } + @Test void testTransportRetry_nullCauseRetriedThenSucceeds() throws Exception { setup(true);