test(e2e): give the shared cluster warm-up a longer, single deadline

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 <tsx loader>` 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}`);
-    }
-}

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> = {}
+    ): 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;
+    });
+});

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<void>({
+            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<void>({
+        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();
+            }
+        },
+    });
+}