diff --git a/packages/databricks-vscode/src/test/e2e/wdio.conf.ts b/packages/databricks-vscode/src/test/e2e/wdio.conf.ts index 8d85dc86e..6df2b84bc 100644 --- a/packages/databricks-vscode/src/test/e2e/wdio.conf.ts +++ b/packages/databricks-vscode/src/test/e2e/wdio.conf.ts @@ -31,6 +31,7 @@ import { formatRecoveredSpecsReport, } from "../retry.ts"; import {SpecRetryTracker} from "../SpecRetryTracker.ts"; +import {startCluster} from "../startCluster.ts"; // WebdriverIO v9 loads TypeScript by injecting `--import ` into // NODE_OPTIONS for every worker process. wdio-vscode-service installs the @@ -381,7 +382,10 @@ export const config: WebdriverIO.Config = { await fs.mkdir(WORKSPACE_PATH, {recursive: true}); const client = getWorkspaceClient(config); - await startCluster(client, process.env["TEST_DEFAULT_CLUSTER_ID"]); + await startCluster( + client.apiClient, + process.env["TEST_DEFAULT_CLUSTER_ID"] + ); process.env.DATABRICKS_HOST = config.host!; process.env.DATABRICKS_VSCODE_INTEGRATION_TEST = "true"; @@ -739,46 +743,3 @@ function getWorkspaceClient(config: Config) { return client; } -async function startCluster( - workspaceClient: WorkspaceClient, - clusterId: string, - attempt = 0 -) { - console.log(`Cluster ID: ${clusterId}`); - if (attempt > 100) { - throw new Error("Failed to start the cluster: too many attempts"); - } - const cluster = await workspaceClient.clusters.get({ - cluster_id: clusterId, - }); - console.log(`Cluster State: ${cluster.state}`); - switch (cluster.state) { - case "RUNNING": - console.log("Cluster is already running"); - break; - case "TERMINATED": - case "ERROR": - case "UNKNOWN": - console.log("Starting the cluster..."); - await ( - await workspaceClient.clusters.start({ - cluster_id: clusterId, - }) - ).wait({ - onProgress: async (state) => { - console.log(`Cluster state: ${state.state}`); - }, - }); - break; - case "PENDING": - case "RESIZING": - case "TERMINATING": - case "RESTARTING": - console.log("Waiting and retrying..."); - await sleep(10000); - await startCluster(workspaceClient, clusterId, attempt + 1); - break; - default: - throw new Error(`Unknown cluster state: ${cluster.state}`); - } -} diff --git a/packages/databricks-vscode/src/test/startCluster.test.ts b/packages/databricks-vscode/src/test/startCluster.test.ts new file mode 100644 index 000000000..b5fbfb38a --- /dev/null +++ b/packages/databricks-vscode/src/test/startCluster.test.ts @@ -0,0 +1,166 @@ +/* eslint-disable @typescript-eslint/naming-convention */ + +import {ApiClient, Time, TimeUnits, compute} from "@databricks/sdk-experimental"; +import * as assert from "node:assert"; +import {mock, when, instance, anything, objectContaining, verify} from "ts-mockito"; +import FakeTimers from "@sinonjs/fake-timers"; +import {startCluster, ClusterStartError} from "./startCluster"; + +describe(__filename, function () { + this.timeout(new Time(10, TimeUnits.minutes).toMillSeconds().value); + + const clusterId = "test-cluster-id"; + let mockedClient: ApiClient; + let fakeTimer: FakeTimers.Clock; + + const details = ( + state: compute.State, + extra: Partial = {} + ): compute.ClusterDetails => + ({cluster_id: clusterId, state, ...extra}) as compute.ClusterDetails; + + const whenGet = () => + when( + mockedClient.request( + objectContaining({path: "/api/2.1/clusters/get", method: "GET"}), + anything() + ) + ); + const whenStart = () => + when( + mockedClient.request( + objectContaining({ + path: "/api/2.1/clusters/start", + method: "POST", + }), + anything() + ) + ); + const verifyStarted = (times: number) => + verify( + mockedClient.request( + objectContaining({ + path: "/api/2.1/clusters/start", + method: "POST", + }), + anything() + ) + ).times(times); + + beforeEach(() => { + mockedClient = mock(ApiClient); + fakeTimer = FakeTimers.install({shouldClearNativeTimers: true}); + }); + + afterEach(() => { + fakeTimer.uninstall(); + }); + + it("returns without starting when the cluster is already RUNNING", async () => { + whenGet().thenResolve(details("RUNNING")); + + await startCluster(instance(mockedClient), clusterId); + + verifyStarted(0); + }); + + it("starts a stopped cluster and polls until RUNNING", async () => { + whenGet().thenResolve( + details("TERMINATED"), + details("PENDING"), + details("RUNNING") + ); + whenStart().thenResolve({}); + + const startPromise = startCluster(instance(mockedClient), clusterId); + await fakeTimer.runToLastAsync(); + await startPromise; + + verifyStarted(1); + }); + + it("fails fast when the cluster returns to a terminal state after start", async () => { + whenGet().thenResolve( + details("TERMINATED"), + details("TERMINATED", {state_message: "bad spark config"}) + ); + whenStart().thenResolve({}); + + const startPromise = startCluster(instance(mockedClient), clusterId); + const rejection = assert.rejects( + startPromise, + (e: Error) => + e instanceof ClusterStartError && /bad spark config/.test(e.message) + ); + await fakeTimer.runToLastAsync(); + await rejection; + }); + + it("tolerates a concurrent start race on the shared cluster", async () => { + // Initial TERMINATED -> our start() races a sibling and throws -> the + // re-check finds it already coming up (PENDING) -> RUNNING. + whenGet().thenResolve( + details("TERMINATED"), + details("PENDING"), + details("RUNNING") + ); + whenStart().thenReject( + new Error(`Cluster ${clusterId} is in unexpected state Pending.`) + ); + + const startPromise = startCluster(instance(mockedClient), clusterId); + await fakeTimer.runToLastAsync(); + await startPromise; + + // The raced start() was attempted, and polling still reached RUNNING. + verifyStarted(1); + }); + + it("propagates a non-race start error when the cluster stays stopped", async () => { + // start() fails and the re-check shows the cluster still stopped, so the + // original (actionable) error surfaces rather than being masked. + whenGet().thenResolve(details("TERMINATED"), details("TERMINATED")); + whenStart().thenReject(new Error("permission denied")); + + const startPromise = startCluster(instance(mockedClient), clusterId); + const rejection = assert.rejects( + startPromise, + (e: Error) => /permission denied/.test(e.message) + ); + await fakeTimer.runToLastAsync(); + await rejection; + }); + + it("waits for a TERMINATING cluster to stop, then starts it", async () => { + whenGet().thenResolve( + details("TERMINATING"), + details("TERMINATED"), + details("RUNNING") + ); + whenStart().thenResolve({}); + + const startPromise = startCluster(instance(mockedClient), clusterId); + await fakeTimer.runToLastAsync(); + await startPromise; + + verifyStarted(1); + }); + + it("fails fast when the cluster is UNKNOWN after start", async () => { + whenGet().thenResolve( + details("TERMINATED"), + details("UNKNOWN", {state_message: "lost the cluster"}) + ); + whenStart().thenResolve({}); + + const startPromise = startCluster(instance(mockedClient), clusterId); + const rejection = assert.rejects( + startPromise, + (e: Error) => + e instanceof ClusterStartError && + /lost the cluster/.test(e.message) + ); + await fakeTimer.runToLastAsync(); + await rejection; + }); +}); diff --git a/packages/databricks-vscode/src/test/startCluster.ts b/packages/databricks-vscode/src/test/startCluster.ts new file mode 100644 index 000000000..29ca40b5b --- /dev/null +++ b/packages/databricks-vscode/src/test/startCluster.ts @@ -0,0 +1,126 @@ +/* eslint-disable no-console */ + +import { + ApiClient, + Time, + TimeUnits, + compute, + retry, + retries, +} from "@databricks/sdk-experimental"; + +// Warms up the shared e2e test cluster from wdio's onPrepare, before any spec +// or session runs (a throw here aborts the whole shard — specFileRetries can't +// recover it). Deliberately duplicates the SDK retry() idiom rather than +// reusing the production Cluster.start() so the e2e path can own a long +// timeout (the shared cluster's cloud node placement has been seen taking ~1h) +// and tolerate the shared-cluster start race, without changing production. +export class ClusterStartError extends Error {} + +const DEFAULT_START_TIMEOUT = new Time(60, TimeUnits.minutes); +const POLL_INTERVAL = new Time(10, TimeUnits.seconds); + +export async function startCluster( + client: ApiClient, + clusterId: string, + timeout: Time = DEFAULT_START_TIMEOUT +) { + const clusterApi = new compute.ClustersService(client); + const log = (c: compute.ClusterDetails) => + console.log( + `Cluster ${clusterId} state: ${c.state}${ + c.state_message ? ` - ${c.state_message}` : "" + }` + ); + + // One deadline across both the shutdown wait and the start poll, so a slow + // TERMINATING phase can't hand the poll a fresh full timeout and let the + // total exceed the caller's bound. + const deadline = Date.now() + timeout.toMillSeconds().value; + const remaining = () => + new Time(Math.max(0, deadline - Date.now()), TimeUnits.milliseconds); + + let cluster = await clusterApi.get({cluster_id: clusterId}); + log(cluster); + if (cluster.state === "RUNNING") { + return; + } + + // If it's shutting down, wait for that to finish before restarting it. + if (cluster.state === "TERMINATING") { + await retry({ + timeout: remaining(), + retryPolicy: new retries.LinearRetryPolicy(POLL_INTERVAL), + fn: async () => { + cluster = await clusterApi.get({cluster_id: clusterId}); + log(cluster); + if (cluster.state === "TERMINATING") { + throw new retries.RetriableError(); + } + }, + }); + } + + if ( + cluster.state === "TERMINATED" || + cluster.state === "ERROR" || + cluster.state === "UNKNOWN" + ) { + try { + await clusterApi.start({cluster_id: clusterId}); + } catch (e) { + // The cluster is shared across ~40 shards, so a sibling may have + // already started it, racing this call into an error. Re-check: if + // it's now coming up we merely raced, so poll below; if it's still + // stopped the start genuinely failed (auth, permissions, bad + // request), so surface that actionable error rather than mask it. + cluster = await clusterApi.get({cluster_id: clusterId}); + log(cluster); + if ( + cluster.state === "TERMINATED" || + cluster.state === "ERROR" || + cluster.state === "UNKNOWN" + ) { + throw e; + } + } + } + + // Poll to RUNNING under one deadline. A terminal state here is a real launch + // failure (the start was already issued), so fail fast with the cloud-side + // reason instead of burning the whole timeout. + await retry({ + timeout: remaining(), + retryPolicy: new retries.LinearRetryPolicy(POLL_INTERVAL), + fn: async () => { + cluster = await clusterApi.get({cluster_id: clusterId}); + log(cluster); + switch (cluster.state) { + case "RUNNING": + return; + case "TERMINATED": + case "ERROR": + case "UNKNOWN": { + // state_message is a string; termination_reason is an + // object — stringify only the latter so a plain message + // isn't wrapped in quotes. + const reason = + cluster.state_message ?? + cluster.termination_reason ?? + "unknown reason"; + throw new ClusterStartError( + `Cluster ${clusterId} failed to start (${ + cluster.state + }): ${ + typeof reason === "string" + ? reason + : JSON.stringify(reason) + }` + ); + } + default: + throw new retries.RetriableError(); + } + }, + }); +}