fix: bound Thrift CloudFetch downloads by retriesTimeout end to end

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Thrift CloudFetch downloads are now bounded end to end by `retriesTimeout` (default 15 minutes), including every attempt and retry backoff. Previously a stalled download, where the response headers or body never arrived, blocked the fetch forever. A download still running when the deadline expires is aborted and fails with `RetryError` (`TIMEOUT_EXCEEDED`, message `CloudFetch result download failed: Retry timeout exceeded`). (#520)
 - Thrift connections made through `proxy` now apply `checkServerCertificate`, `customCaCert` and `clientCert`/`clientKey`. Previously the proxied path ignored them and verified against Node's process-wide trust store, so a certificate outside the configured trust set could be accepted. **Compatibility:** proxied connections now verify unless `checkServerCertificate: false` is set; as on direct connections, `NODE_TLS_REJECT_UNAUTHORIZED=0` no longer turns verification off. For an `https` proxy these options also apply to the proxy's own certificate. (#485)
 - Honor the client `useCloudFetch: false` default on the kernel backend so complete results can be fetched without cloud-storage downloads. (#574)
 - Thrift HTTP errors (`THTTPException`) now append server detail from the error headers (at most 500 characters). Kernel-backend errors now carry a non-enumerable `cause` summarizing the kernel error code and message (also at most 500 characters). Reyden fallback preserves KP001 after any existing kernel cause. (#467)

lib/connection/connections/HttpRetryPolicy.ts

@@ -3,9 +3,24 @@ import { HttpTransactionDetails } from '../contracts/IConnectionProvider';
 import IClientContext from '../../contracts/IClientContext';
 import RetryError, { RetryErrorCode } from '../../errors/RetryError';
 
-function delay(milliseconds: number): Promise<void> {
-  return new Promise<void>((resolve) => {
-    setTimeout(() => resolve(), milliseconds);
+// Sleeps for `milliseconds`; aborting `signal` clears the timer and rejects with
+// `signal.reason`, so a cancelled backoff does not keep the process alive.
+function delay(milliseconds: number, signal?: AbortSignal): Promise<void> {
+  return new Promise<void>((resolve, reject) => {
+    if (signal?.aborted) {
+      reject(signal.reason);
+      return;
+    }
+    let timer: NodeJS.Timeout | undefined;
+    const onAbort = () => {
+      clearTimeout(timer);
+      reject(signal?.reason);
+    };
+    timer = setTimeout(() => {
+      signal?.removeEventListener('abort', onAbort);
+      resolve();
+    }, milliseconds);
+    signal?.addEventListener('abort', onAbort, { once: true });
   });
 }
 
@@ -52,7 +67,10 @@ export default class HttpRetryPolicy implements IRetryPolicy<HttpTransactionDeta
     return { shouldRetry: false };
   }
 
-  public async invokeWithRetry(operation: RetryableOperation<HttpTransactionDetails>): Promise<HttpTransactionDetails> {
+  public async invokeWithRetry(
+    operation: RetryableOperation<HttpTransactionDetails>,
+    signal?: AbortSignal,
+  ): Promise<HttpTransactionDetails> {
     for (;;) {
       // Capture either the resolved response or the thrown error so the
       // retry-decision logic below can flow without an early `continue` and
@@ -73,7 +91,7 @@ export default class HttpRetryPolicy implements IRetryPolicy<HttpTransactionDeta
           return outcome.details;
         }
         // eslint-disable-next-line no-await-in-loop
-        await delay(status.retryAfter);
+        await delay(status.retryAfter, signal);
       } else {
         // The operation threw before producing a response. This is typically a
         // transient network failure (stale keep-alive socket reset by a load
@@ -91,7 +109,7 @@ export default class HttpRetryPolicy implements IRetryPolicy<HttpTransactionDeta
           throw outcome.error;
         }
         // eslint-disable-next-line no-await-in-loop
-        await delay(status.retryAfter);
+        await delay(status.retryAfter, signal);
       }
     }
   }

lib/connection/contracts/IRetryPolicy.ts

@@ -12,5 +12,7 @@ export type RetryableOperation<R> = () => Promise<R>;
 export default interface IRetryPolicy<R> {
   shouldRetry(details: R): Promise<ShouldRetryResult>;
 
-  invokeWithRetry(operation: RetryableOperation<R>): Promise<R>;
+  // `signal` cancels a pending backoff sleep: aborting it rejects with `signal.reason`
+  // instead of waiting for the sleep to finish (the operation itself must honour the signal).
+  invokeWithRetry(operation: RetryableOperation<R>, signal?: AbortSignal): Promise<R>;
 }

lib/contracts/IDBSQLClient.ts

@@ -141,7 +141,10 @@ export type ConnectionOptions = {
    *     attempt with no retry. Both backends honour the same total-attempt
    *     semantics (the kernel converts it to its after-initial retry count).
    *   • `retriesTimeout`   — maximum total wallclock spent retrying, in
-   *     milliseconds. Default 900000 (15 minutes).
+   *     milliseconds. Default 900000 (15 minutes). On the Thrift backend it is
+   *     also the end-to-end deadline of each CloudFetch result download (all
+   *     attempts plus backoff): a download still running when it expires is
+   *     aborted and fails with a `RetryError` (`TIMEOUT_EXCEEDED`).
    *   • `retryDelayMin`    — minimum backoff between attempts, in milliseconds.
    *     Default 1000.
    *   • `retryDelayMax`    — maximum backoff between attempts, in milliseconds.

lib/result/CloudFetchResultHandler.ts

@@ -1,7 +1,7 @@
 import fetch, { RequestInfo, RequestInit, Request, Response } from 'node-fetch';
 import { TGetResultSetMetadataResp, TRowSet, TSparkArrowResultLink } from '../../thrift/TCLIService_types';
 import HiveDriverError from '../errors/HiveDriverError';
-import RetryError from '../errors/RetryError';
+import RetryError, { RetryErrorCode } from '../errors/RetryError';
 import IClientContext from '../contracts/IClientContext';
 import IResultsProvider, { ResultsProviderFetchNextOptions } from './IResultsProvider';
 import { ArrowBatch } from './utils';
@@ -17,6 +17,9 @@ import { safeEmit } from '../telemetry/telemetryUtils';
 // unchanged while the unhandled-rejection leak is eliminated.
 type SettledDownload = { ok: true; batch: ArrowBatch } | { ok: false; error: unknown };
 
+// Largest delay `setTimeout` accepts; longer values fire immediately, so they are clamped.
+const MAX_TIMER_DELAY_MS = 2 ** 31 - 1;
+
 export default class CloudFetchResultHandler implements IResultsProvider<ArrowBatch> {
   private readonly context: IClientContext;
 
@@ -150,8 +153,21 @@ export default class CloudFetchResultHandler implements IResultsProvider<ArrowBa
     const connectionProvider = await this.context.getConnectionProvider();
     const agent = await connectionProvider.getAgent();
     const retryPolicy = await connectionProvider.getRetryPolicy();
+    const { retriesTimeout } = this.context.getConfig();
 
     const requestConfig: RequestInit = { agent, ...init };
+    // The retry policy checks `retriesTimeout` only between attempts, so a stalled attempt
+    // (headers or body never arrive) would block forever. Enforce it as a deadline over the
+    // whole download - every attempt plus the backoff between them - and abort the in-flight
+    // request and any pending backoff sleep when it expires (#520). Non-positive values keep
+    // the deadline off.
+    const deadlineMs = retriesTimeout > 0 ? Math.min(retriesTimeout, MAX_TIMER_DELAY_MS) : undefined;
+    // Aborted only once, by the deadline, so a single controller serves every attempt and backoff.
+    const controller = deadlineMs === undefined ? undefined : new AbortController();
+    let timeoutError: RetryError | undefined;
+    let lastFailure: unknown;
+    let timer: NodeJS.Timeout | undefined;
+
     // Read the body inside the retried block. CloudFetch downloads are large
     // GETs against pre-signed cloud-storage URLs that frequently surface
     // `socket hang up` / "Premature close" once the stream is mid-transfer.
@@ -160,17 +176,51 @@ export default class CloudFetchResultHandler implements IResultsProvider<ArrowBa
     let downloaded: ArrayBuffer | undefined;
     let result: { request: Request; response: Response };
     try {
-      result = await retryPolicy.invokeWithRetry(async () => {
+      const download = retryPolicy.invokeWithRetry(async () => {
+        // The timeout error is not retryable, so this also ends the retry loop when the
+        // deadline expired during a backoff sleep or aborted the attempt.
+        if (timeoutError) {
+          throw timeoutError;
+        }
         downloaded = undefined;
-        const request = new Request(url, requestConfig);
-        const response = await fetch(request);
-        if (response.ok) {
-          downloaded = await response.arrayBuffer();
+        const request = new Request(url, {
+          ...requestConfig,
+          signal: controller?.signal as unknown as RequestInit['signal'],
+        });
+        try {
+          const response = await fetch(request);
+          if (response.ok) {
+            downloaded = await response.arrayBuffer();
+          } else {
+            lastFailure = { request, response };
+          }
+          return { request, response };
+        } catch (error) {
+          if (timeoutError) {
+            throw timeoutError;
+          }
+          lastFailure = error;
+          throw error;
         }
-        return { request, response };
-      });
+      }, controller?.signal);
+
+      if (deadlineMs === undefined) {
+        result = await download;
+      } else {
+        const deadline = new Promise<never>((_, reject) => {
+          timer = setTimeout(() => {
+            timeoutError = new RetryError(RetryErrorCode.TimeoutExceeded, lastFailure);
+            // Aborts the in-flight request, and rejects a pending backoff sleep with `timeoutError`
+            controller?.abort(timeoutError);
+            reject(timeoutError);
+          }, deadlineMs);
+        });
+        result = await Promise.race([download, deadline]);
+      }
     } catch (error) {
       throw CloudFetchResultHandler.describeDownloadError(error);
+    } finally {
+      clearTimeout(timer);
     }
     // Fall back to reading the body here if the retry policy returned a
     // response without consuming it via our operation (e.g. unit-test stubs

tests/e2e/cloudfetch.test.ts

@@ -234,4 +234,82 @@ describe('CloudFetch', () => {
       await session.close();
     }
   });
+
+  it('should fail a stalled CloudFetch download with a timeout error within retriesTimeout', async () => {
+    const retriesTimeout = 3000;
+    const session = await openSession({
+      cloudFetchConcurrentDownloads: 1,
+      useLZ4Compression: false,
+      retriesTimeout,
+    });
+
+    // Accepts every download and never responds
+    let serverHits = 0;
+    const server = http.createServer(() => {
+      serverHits += 1;
+    });
+    await new Promise<void>((resolve) => {
+      server.listen(0, '127.0.0.1', () => resolve());
+    });
+    const { port } = server.address() as AddressInfo;
+
+    try {
+      const queriedRowsCount = 10000000; // result has to be quite big to enable CloudFetch
+      const operation = await session.executeStatement(
+        `
+          SELECT *
+          FROM range(0, ${queriedRowsCount}) AS t1
+          LEFT JOIN (SELECT 1) AS t2
+        `,
+        {
+          maxRows: null, // disable DirectResults
+          useCloudFetch: true,
+        },
+      );
+
+      await operation.finished();
+
+      // @ts-expect-error TS2339: Property getResultHandler does not exist on type IOperation
+      const resultHandler = await operation.getResultHandler();
+      expect(resultHandler.source.source).to.be.instanceOf(CloudFetchResultHandler);
+      const cfResultHandler = resultHandler.source.source;
+
+      const connectionProvider = await cfResultHandler.context.getConnectionProvider();
+      const getAgentStub = sinon.stub(connectionProvider, 'getAgent').resolves(undefined);
+
+      const rowSetProvider = cfResultHandler.source;
+      const originalFetchNext = rowSetProvider.fetchNext.bind(rowSetProvider);
+      sinon.stub(rowSetProvider, 'fetchNext').callsFake(async (...args: Array<any>) => {
+        const rowSet = await originalFetchNext(...args);
+        rowSet?.resultLinks?.forEach((link: any) => {
+          link.fileLink = `http://127.0.0.1:${port}/result`;
+          link.httpHeaders = {};
+        });
+        return rowSet;
+      });
+
+      const startedAt = Date.now();
+      let thrown: unknown;
+      try {
+        await operation.fetchChunk({ maxRows: 100000, disableBuffering: true });
+      } catch (error) {
+        thrown = error;
+      } finally {
+        getAgentStub.restore();
+      }
+      const elapsedMs = Date.now() - startedAt;
+
+      expect(thrown).to.be.instanceOf(RetryError);
+      expect((thrown as RetryError).errorCode).to.be.equal(RetryErrorCode.TimeoutExceeded);
+      expect((thrown as RetryError).message).to.contain('CloudFetch result download failed: Retry timeout exceeded');
+      expect(serverHits).to.be.equal(1);
+      expect(elapsedMs).to.be.lt(retriesTimeout + 5000);
+    } finally {
+      server.closeAllConnections();
+      await new Promise<void>((resolve) => {
+        server.close(() => resolve());
+      });
+      await session.close();
+    }
+  });
 });

tests/unit/connection/connections/HttpRetryPolicy.test.ts

@@ -400,5 +400,78 @@ describe('HttpRetryPolicy', () => {
         expect(operation.callCount).to.equal(clientConfig.retryMaxAttempts);
       }
     });
+
+    it('should cancel a pending backoff sleep when the signal is aborted', async () => {
+      const context = new ClientContextStub({
+        retryDelayMin: 10000,
+        retryDelayMax: 60000,
+        retryMaxAttempts: 20,
+      });
+      const policy = new HttpRetryPolicy(context);
+      const operation = sinon.stub().resolves({
+        request: new Request('http://localhost'),
+        response: new Response(undefined, { status: 503 }),
+      });
+      const controller = new AbortController();
+      const reason = new Error('deadline');
+      const clearTimeoutSpy = sinon.spy(global, 'clearTimeout');
+      const setTimeoutSpy = sinon.spy(global, 'setTimeout');
+
+      try {
+        const invocation = policy.invokeWithRetry(operation, controller.signal);
+        // Let the first attempt complete and the 20s backoff sleep start
+        await new Promise((resolve) => {
+          setImmediate(resolve);
+        });
+        const backoffCall = setTimeoutSpy.getCalls().find((call) => call.args[1] === 20000);
+        expect(backoffCall).to.not.be.undefined;
+
+        controller.abort(reason);
+
+        try {
+          await invocation;
+          expect.fail('It should reject with the abort reason');
+        } catch (error) {
+          if (error instanceof AssertionError) {
+            throw error;
+          }
+          expect(error).to.equal(reason);
+        }
+        expect(clearTimeoutSpy.calledWith(backoffCall!.returnValue)).to.be.true;
+        expect(operation.callCount).to.equal(1);
+      } finally {
+        setTimeoutSpy.restore();
+        clearTimeoutSpy.restore();
+      }
+    });
+
+    it('should not start another attempt when the signal is already aborted before the backoff', async () => {
+      const context = new ClientContextStub({
+        retryDelayMin: 10000,
+        retryDelayMax: 60000,
+        retryMaxAttempts: 20,
+      });
+      const policy = new HttpRetryPolicy(context);
+      const controller = new AbortController();
+      const reason = new Error('deadline');
+      const operation = sinon.stub().callsFake(async () => {
+        controller.abort(reason);
+        return {
+          request: new Request('http://localhost'),
+          response: new Response(undefined, { status: 503 }),
+        };
+      });
+
+      try {
+        await policy.invokeWithRetry(operation, controller.signal);
+        expect.fail('It should reject with the abort reason');
+      } catch (error) {
+        if (error instanceof AssertionError) {
+          throw error;
+        }
+        expect(error).to.equal(reason);
+      }
+      expect(operation.callCount).to.equal(1);
+    });
   });
 });