[DECO-167] Fix cluster stop/start (#253)

Author: fjakobs

packages/databricks-sdk-js/package.json

@@ -36,17 +36,17 @@
     },
     "devDependencies": {
         "@istanbuljs/nyc-config-typescript": "^1.0.2",
+        "@sinonjs/fake-timers": "^10.0.0",
         "@types/ini": "^1.3.31",
         "@types/mocha": "^10.0.0",
         "@types/node": "^18.11.9",
-        "@types/sinon": "^10.0.13",
+        "@types/sinonjs__fake-timers": "^8.1.2",
         "@types/tmp": "^0.2.3",
         "@types/uuid": "^8.3.4",
         "eslint": "^8.26.0",
         "mocha": "^10.1.0",
         "nyc": "^15.1.0",
         "prettier": "^2.7.1",
-        "sinon": "^14.0.2",
         "tmp-promise": "^3.0.3",
         "ts-loader": "^9.4.1",
         "ts-mocha": "^10.0.0",

packages/databricks-sdk-js/src/index.ts

@@ -45,7 +45,9 @@ export * from "./types";
 
 export {ClusterFixture, TokenFixture} from "./test/fixtures";
 
-export {RetryConfigs, default as retry} from "./retries/retries";
+export {default as retry} from "./retries/retries";
+export * as retries from "./retries/retries";
+export type {RetryPolicy} from "./retries/retries";
 export {TimeUnits, default as Time} from "./retries/Time";
 
 export * as logging from "./logging";

packages/databricks-sdk-js/src/retries/retries.test.ts

@@ -1,22 +1,21 @@
 import retry, {RetriableError, TimeoutError} from "./retries";
 import Time, {TimeUnits} from "./Time";
-import * as sinon from "sinon";
+import FakeTimers from "@sinonjs/fake-timers";
 import * as assert from "node:assert";
 
 class NonRetriableError extends Error {}
 
-describe(__filename, function () {
-    let fakeTimer: sinon.SinonFakeTimers;
+describe(__filename, () => {
+    let fakeTimer: FakeTimers.InstalledClock;
 
     beforeEach(() => {
-        fakeTimer = sinon.useFakeTimers();
+        fakeTimer = FakeTimers.install();
     });
 
     afterEach(() => {
-        fakeTimer.restore();
+        fakeTimer.uninstall();
     });
 
-    this.timeout(1000 * 3);
     it("should return result if timeout doesn't expire", async function () {
         const startTime = Date.now();
 

packages/databricks-sdk-js/src/retries/retries.ts

@@ -10,32 +10,63 @@ export class TimeoutError extends Error {
     }
 }
 
-export class RetryConfigs {
-    static maxJitter = new Time(750, TimeUnits.milliseconds);
-    static minJitter = new Time(50, TimeUnits.milliseconds);
-    static maxWaitTime = new Time(10, TimeUnits.seconds);
-    static defaultTimeout = new Time(10, TimeUnits.minutes);
-
-    static waitTime(attempt: number) {
-        const jitter = RetryConfigs.maxJitter
-            .sub(RetryConfigs.minJitter)
+export interface RetryPolicy {
+    waitTime(attempt: number): Time;
+}
+
+export class LinearRetryPolicy {
+    constructor(readonly _waitTime: Time) {}
+
+    waitTime(): Time {
+        return this._waitTime;
+    }
+}
+
+export class ExponetionalBackoffWithJitterRetryPolicy implements RetryPolicy {
+    maxJitter: Time;
+    minJitter: Time;
+    maxWaitTime: Time;
+
+    constructor(
+        options: {
+            maxJitter?: Time;
+            minJitter?: Time;
+            maxWaitTime?: Time;
+        } = {}
+    ) {
+        this.maxJitter =
+            options.maxJitter || new Time(750, TimeUnits.milliseconds);
+        this.minJitter =
+            options.minJitter || new Time(50, TimeUnits.milliseconds);
+        this.maxWaitTime =
+            options.maxWaitTime || new Time(10, TimeUnits.seconds);
+    }
+
+    waitTime(attempt: number): Time {
+        const jitter = this.maxJitter
+            .sub(this.minJitter)
             .multiply(Math.random())
-            .add(RetryConfigs.minJitter);
+            .add(this.minJitter);
         const timeout = new Time(attempt, TimeUnits.seconds).add(jitter);
 
-        return timeout.gt(RetryConfigs.maxWaitTime)
-            ? RetryConfigs.maxWaitTime
-            : timeout;
+        return timeout.gt(this.maxWaitTime) ? this.maxWaitTime : timeout;
     }
 }
 
+export const DEFAULT_RETRY_CONFIG =
+    new ExponetionalBackoffWithJitterRetryPolicy();
+
+export const DEFAULT_MAX_TIMEOUT = new Time(10, TimeUnits.minutes);
+
 interface RetryArgs<T> {
     timeout?: Time;
+    retryPolicy?: RetryPolicy;
     fn: () => Promise<T>;
 }
 
 export default async function retry<T>({
-    timeout = RetryConfigs.defaultTimeout,
+    timeout = DEFAULT_MAX_TIMEOUT,
+    retryPolicy: retryConfig = DEFAULT_RETRY_CONFIG,
     fn,
 }: RetryArgs<T>): Promise<T> {
     let attempt = 1;
@@ -67,7 +98,7 @@ export default async function retry<T>({
         await new Promise((resolve) =>
             setTimeout(
                 resolve,
-                RetryConfigs.waitTime(attempt).toMillSeconds().value
+                retryConfig.waitTime(attempt).toMillSeconds().value
             )
         );
 

packages/databricks-sdk-js/src/services/Cluster.integ.ts

@@ -35,4 +35,33 @@ describe(__filename, function () {
         assert(clusterA.id);
         assert.equal(clusterA.id, clusterB?.id);
     });
+
+    // skipping because running the test takes too long
+    it.skip("should start a stopping cluster", async () => {
+        const token = {
+            isCancellationRequested: false,
+        };
+
+        const cluster = integSetup.cluster;
+        // stop cluster
+        await Promise.race([
+            cluster.stop(token, async (info) =>
+                // eslint-disable-next-line no-console
+                console.log(`Stopping - ${info.state}`)
+            ),
+            new Promise<void>((resolve) => {
+                // cancel stop
+                setTimeout(() => {
+                    token.isCancellationRequested = true;
+                    resolve();
+                }, 500);
+            }),
+        ]);
+
+        // start cluster
+        await cluster.start(undefined, (state) =>
+            // eslint-disable-next-line no-console
+            console.log(`Starting ${state}`)
+        );
+    });
 });

packages/databricks-sdk-js/src/services/Cluster.test.ts

@@ -7,22 +7,25 @@ import Time, {TimeUnits} from "../retries/Time";
 import {getMockTestCluster} from "../test/fixtures/ClusterFixtures";
 import {ClusterInfo} from "../apis/clusters";
 import {TokenFixture} from "../test/fixtures/TokenFixtures";
-import {RetryConfigs} from "../retries/retries";
+import FakeTimers from "@sinonjs/fake-timers";
 
 describe(__filename, function () {
     this.timeout(new Time(10, TimeUnits.minutes).toMillSeconds().value);
 
     let mockedClient: ApiClient;
     let mockedCluster: Cluster;
     let testClusterDetails: ClusterInfo;
+    let fakeTimer: FakeTimers.InstalledClock;
 
     beforeEach(async () => {
         ({mockedCluster, mockedClient, testClusterDetails} =
             await getMockTestCluster());
 
-        RetryConfigs.waitTime = () => {
-            return new Time(0, TimeUnits.milliseconds);
-        };
+        fakeTimer = FakeTimers.install();
+    });
+
+    afterEach(() => {
+        fakeTimer.uninstall();
     });
 
     it("calling start on a non terminated state should not throw an error", async () => {
@@ -65,7 +68,9 @@ describe(__filename, function () {
         await mockedCluster.refresh();
         assert.notEqual(mockedCluster.state, "RUNNING");
 
-        await mockedCluster.start();
+        const startPromise = mockedCluster.start();
+        await fakeTimer.runToLastAsync();
+        await startPromise;
         assert.equal(mockedCluster.state, "RUNNING");
 
         verify(
@@ -125,7 +130,10 @@ describe(__filename, function () {
 
         assert.equal(mockedCluster.state, "RUNNING");
 
-        await mockedCluster.stop();
+        const stopPromise = mockedCluster.stop();
+        await fakeTimer.runToLastAsync();
+        await stopPromise;
+
         assert.equal(mockedCluster.state, "TERMINATED");
 
         verify(
@@ -175,7 +183,10 @@ describe(__filename, function () {
         await mockedCluster.refresh();
         assert.notEqual(mockedCluster.state, "RUNNING");
 
-        await mockedCluster.stop();
+        const stopPromise = mockedCluster.stop();
+        await fakeTimer.runToLastAsync();
+        await stopPromise;
+
         assert.equal(mockedCluster.state, "TERMINATED");
 
         verify(
@@ -243,7 +254,9 @@ describe(__filename, function () {
         await mockedCluster.refresh();
 
         assert.equal(mockedCluster.state, "PENDING");
-        await mockedCluster.start(instance(token));
+        const startPromise = mockedCluster.start(instance(token));
+        await fakeTimer.runToLastAsync();
+        await startPromise;
 
         verify(token.isCancellationRequested).thrice();
     });

packages/databricks-sdk-js/src/services/Cluster.ts

@@ -1,7 +1,7 @@
 /* eslint-disable @typescript-eslint/naming-convention */
 
 import {ApiClient} from "../api-client";
-import retry, {RetriableError} from "../retries/retries";
+import retry, {LinearRetryPolicy, RetriableError} from "../retries/retries";
 import {
     JobsService,
     RunLifeCycleState,
@@ -11,7 +11,7 @@ import {
 import {CancellationToken} from "../types";
 import {ExecutionContext} from "./ExecutionContext";
 import {WorkflowRun} from "./WorkflowRun";
-import {commands, PermissionsService} from "..";
+import {commands, PermissionsService, Time, TimeUnits} from "..";
 import {
     ClusterInfo,
     ClustersService,
@@ -189,6 +189,8 @@ export class Cluster {
         onProgress: (state: ClusterInfoState) => void = () => {}
     ) {
         await this.refresh();
+        onProgress(this.state);
+
         if (this.state === "RUNNING") {
             return;
         }
@@ -203,6 +205,30 @@ export class Cluster {
             });
         }
 
+        // wait for cluster to be stopped before re-starting
+        if (this.state === "TERMINATING") {
+            await retry<void>({
+                timeout: new Time(1, TimeUnits.minutes),
+                retryPolicy: new LinearRetryPolicy(
+                    new Time(1, TimeUnits.seconds)
+                ),
+                fn: async () => {
+                    if (token?.isCancellationRequested) {
+                        return;
+                    }
+                    await this.refresh();
+                    onProgress(this.state);
+
+                    if (this.state === "TERMINATING") {
+                        throw new RetriableError();
+                    }
+                },
+            });
+            await this.clusterApi.start({
+                cluster_id: this.id,
+            });
+        }
+
         this._canExecute = undefined;
         await retry({
             fn: async () => {

packages/databricks-vscode/package.json

@@ -564,11 +564,13 @@
     },
     "devDependencies": {
         "@istanbuljs/nyc-config-typescript": "^1.0.2",
+        "@sinonjs/fake-timers": "^10.0.0",
         "@types/fs-extra": "^9.0.13",
         "@types/glob": "^8.0.0",
         "@types/mocha": "^10.0.0",
         "@types/mock-require": "^2.0.1",
         "@types/node": "^18.11.9",
+        "@types/sinonjs__fake-timers": "^8.1.2",
         "@types/tmp": "^0.2.3",
         "@types/vscode": "^1.69.1",
         "@types/yargs": "^17.0.14",
@@ -609,4 +611,4 @@
         ],
         "report-dir": "coverage"
     }
-}
+}

packages/databricks-vscode/src/cluster/ClusterManager.test.ts

@@ -5,9 +5,9 @@ import {
     cluster,
     Cluster,
     ClusterFixture,
-    RetryConfigs,
     Time,
     TimeUnits,
+    retries,
 } from "@databricks/databricks-sdk";
 import {
     anything,
@@ -49,7 +49,7 @@ describe(__filename, async () => {
 
         resetCalls(mockedClient);
 
-        RetryConfigs.waitTime = () => {
+        retries.DEFAULT_RETRY_CONFIG.waitTime = () => {
             return new Time(0, TimeUnits.milliseconds);
         };
     });