fix: task executor end() can be called multiple times

src/executor/executor.test.ts → src/executor/executor.spec.ts

@@ -5,6 +5,8 @@ import { TaskExecutor } from "./executor";
 import { sleep } from "../utils";
 import { LoggerMock } from "../../tests/mock";
 
+/* eslint-disable @typescript-eslint/no-explicit-any */
+
 jest.mock("../market/service");
 jest.mock("../agreement/service");
 jest.mock("../network/service");
@@ -61,4 +63,26 @@ describe("Task Executor", () => {
       await executor.end();
     });
   });
+
+  describe("end()", () => {
+    it("should allow multiple calls", async () => {
+      // Implementation details: the same promise is always used, so it's safe to call end() multiple times.
+      const executor = await TaskExecutor.create({ package: "test", startupTimeout: 0, logger, yagnaOptions });
+      const p = Promise.resolve();
+      const spy = jest.spyOn(executor as any, "doEnd").mockReturnValue(p);
+
+      const r1 = executor.end();
+      expect(r1).toBeDefined();
+      expect(r1).toStrictEqual(p);
+
+      const r2 = executor.end();
+      expect(r1).toStrictEqual(r2);
+
+      await r1;
+
+      const r3 = executor.end();
+      expect(r3).toStrictEqual(r1);
+      expect(spy).toHaveBeenCalledTimes(1);
+    });
+  });
 });

src/executor/executor.ts

@@ -112,8 +112,19 @@ export class TaskExecutor {
   private startupTimeoutId?: NodeJS.Timeout;
   private yagna: Yagna;
 
+  /**
+   * Signal handler reference, needed to remove handlers on exit.
+   * @param signal
+   */
   private signalHandler = (signal: string) => this.cancel(signal);
 
+  /**
+   * End promise.
+   * This will be set by call to end() method.
+   * It will be resolved when the executor is fully stopped.
+   */
+  private endPromise?: Promise<void>;
+
   /**
    * Create a new Task Executor
    * @description Factory Method that create and initialize an instance of the TaskExecutor
@@ -243,12 +254,25 @@ export class TaskExecutor {
   }
 
   /**
-   * Stop all executor services and shut down executor instance
+   * Stop all executor services and shut down executor instance.
+   *
+   * You can call this method multiple times, it will resolve only once the executor is shutdown.
+   */
+  end(): Promise<void> {
+    if (this.isRunning) {
+      this.isRunning = false;
+      this.endPromise = this.doEnd();
+    }
+
+    return this.endPromise!;
+  }
+
+  /**
+   * Perform everything needed to cleanly shut down the executor.
+   * @private
    */
-  async end() {
+  private async doEnd() {
     if (runtimeContextChecker.isNode) this.removeSignalHandlers();
-    if (!this.isRunning) return;
-    this.isRunning = false;
     clearTimeout(this.startupTimeoutId);
     if (!this.configOptions.storageProvider) await this.storageProvider?.close();
     await this.networkService?.end();