diff --git a/packages/databricks/package.json b/packages/databricks/package.json index fcc0d37a..b61cadea 100644 --- a/packages/databricks/package.json +++ b/packages/databricks/package.json @@ -10,6 +10,10 @@ "types": "./dist/index.d.ts", "import": "./dist/index.js" }, + "./api": { + "types": "./dist/api/index.d.ts", + "import": "./dist/api/index.js" + }, "./apierror": { "types": "./dist/apierror/index.d.ts", "import": "./dist/apierror/index.js" diff --git a/packages/databricks/src/api/index.ts b/packages/databricks/src/api/index.ts new file mode 100644 index 00000000..5eebf785 --- /dev/null +++ b/packages/databricks/src/api/index.ts @@ -0,0 +1,8 @@ +/** + * Databricks API client utilities. + * + * @packageDocumentation + */ + +export {BackoffPolicy, retryOn} from './retrier'; +export type {BackoffPolicyOptions, Retrier} from './retrier'; diff --git a/packages/databricks/src/api/retrier.ts b/packages/databricks/src/api/retrier.ts new file mode 100644 index 00000000..0c20a736 --- /dev/null +++ b/packages/databricks/src/api/retrier.ts @@ -0,0 +1,108 @@ +/** Options for configuring a {@link BackoffPolicy}. */ +export interface BackoffPolicyOptions { + /** Initial delay in milliseconds; defaults to 1000. */ + initial?: number; + + /** Maximum delay in milliseconds; defaults to 60000. */ + maximum?: number; + + /** + * Factor by which the delay is multiplied after each retry. The value must + * be greater or equal to 1. If not, it defaults to 2. + */ + factor?: number; +} + +// Random number generation, wrapped in an object for testability. +export const rand = { + // Returns a random integer in [0, n). + int(n: number): number { + return Math.floor(Math.random() * n); + }, +}; + +/** + * BackoffPolicy implements an exponential backoff policy. The delay between + * retries is randomly computed between 0 and the "exponential delay" as + * recommended in [Exponential Backoff And Jitter]. The retry delay starts from + * initial and grows exponentially by factor at every retry. The maximum retry + * delay is capped by maximum. + * + * There is no parameter to limit the number of retries. This is intended as + * such logic should be implemented upstream (e.g. in a Retrier). + * + * [Exponential Backoff And Jitter]: https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/ + */ +export class BackoffPolicy { + /** Initial delay in milliseconds. */ + readonly initial: number; + + /** Maximum delay in milliseconds. */ + readonly maximum: number; + + /** Factor by which the delay is multiplied after each retry. */ + readonly factor: number; + + // Current delay before the next retry. + private current: number; + + constructor(options?: BackoffPolicyOptions) { + let initial = options?.initial ?? 1000; // Default initial delay of 1 second. + const maximum = options?.maximum ?? 60000; // Default maximum delay of 60 seconds. + + if (initial > maximum) { + // Initial cannot be greater than maximum. + initial = maximum; + } + + this.initial = initial; + this.maximum = maximum; + this.factor = + options?.factor !== undefined && options.factor >= 1 ? options.factor : 2; + this.current = this.initial; + } + + /** Returns a random delay in [0, current] and grows the current delay. */ + delay(): number { + // Random duration in the range [0, this.current]. + const d = rand.int(this.current + 1); + + // Grow delay for the next call. + this.current = Math.min(this.current * this.factor, this.maximum); + + return d; + } +} + +/** Retrier defines a retry behavior. */ +export interface Retrier { + /** + * Returns the delay in milliseconds before the next retry, or undefined if + * the error is not retriable. Implementations should assume that the given + * error is never undefined. + */ + isRetriable(err: Error): number | undefined; +} + +/** + * Returns a Retrier that retries based on the isRetriable predicate and relies + * on an internal backoff policy to decide how long to wait between retries. + * + * Important: the retrier has its own backoff policy which cannot be trivially + * reset by design. Users who need to reset the backoff policy should rather + * create a new retrier. + */ +export function retryOn( + options: BackoffPolicyOptions, + isRetriable: (err: Error) => boolean +): Retrier { + const bp = new BackoffPolicy(options); + return { + isRetriable(err: Error): number | undefined { + if (!isRetriable(err)) { + return undefined; + } + return bp.delay(); + }, + }; +} diff --git a/packages/databricks/tests/api/retrier.test.ts b/packages/databricks/tests/api/retrier.test.ts new file mode 100644 index 00000000..4f91bb9e --- /dev/null +++ b/packages/databricks/tests/api/retrier.test.ts @@ -0,0 +1,126 @@ +import {describe, it, expect, vi} from 'vitest'; +import {BackoffPolicy, rand, retryOn} from '../../src/api/retrier'; + +// Always returns the maximum value (n - 1) for deterministic tests. +function deterministicRand(n: number): number { + return n - 1; +} + +describe('retryOn isRetriable', () => { + const testCases: { + name: string; + fn: (err: Error) => boolean; + wantDelay: number | undefined; + }[] = [ + { + name: 'retriable returns delay', + fn: () => true, + wantDelay: 99, + }, + { + name: 'not retriable returns undefined', + fn: () => false, + wantDelay: undefined, + }, + ]; + + it.each(testCases)('$name', ({fn, wantDelay}) => { + vi.spyOn(rand, 'int').mockImplementation(deterministicRand); + const r = retryOn({initial: wantDelay ?? 0}, fn); + + const got = r.isRetriable(new Error('an error')); + if (wantDelay === undefined) { + expect(got).toBeUndefined(); + } else { + expect(got).toBe(wantDelay); + } + vi.restoreAllMocks(); + }); +}); + +describe('BackoffPolicy defaults', () => { + const testCases: { + name: string; + options?: { + initial?: number; + maximum?: number; + factor?: number; + }; + wantInitial: number; + wantMaximum: number; + wantFactor: number; + }[] = [ + { + name: 'default', + wantInitial: 1000, + wantMaximum: 60000, + wantFactor: 2, + }, + { + name: 'custom initial smaller than maximum', + options: {initial: 100}, + wantInitial: 100, + wantMaximum: 60000, + wantFactor: 2, + }, + { + name: 'custom initial greater than maximum', + options: {initial: 10000, maximum: 1000}, + wantInitial: 1000, + wantMaximum: 1000, + wantFactor: 2, + }, + { + name: 'custom factor less than 1', + options: {factor: 0.5}, + wantInitial: 1000, + wantMaximum: 60000, + wantFactor: 2, + }, + { + name: 'custom factor greater than 1', + options: {factor: 1.5}, + wantInitial: 1000, + wantMaximum: 60000, + wantFactor: 1.5, + }, + ]; + + it.each(testCases)( + '$name', + ({options, wantInitial, wantMaximum, wantFactor}) => { + const bp = new BackoffPolicy(options); + + expect(bp.initial).toBe(wantInitial); + expect(bp.maximum).toBe(wantMaximum); + expect(bp.factor).toBe(wantFactor); + } + ); +}); + +describe('BackoffPolicy exponential delay', () => { + it('should grow exponentially and cap at maximum', () => { + vi.spyOn(rand, 'int').mockImplementation(deterministicRand); + const bp = new BackoffPolicy({ + initial: 100, + maximum: 10000, + factor: 2.0, + }); + + const wantDelays = [ + 100, + 200, + 400, + 800, + 1600, + 3200, + 6400, + 10000, // Capped by maximum. + ]; + + for (const want of wantDelays) { + expect(bp.delay()).toBe(want); + } + vi.restoreAllMocks(); + }); +});