From e084f3b73ee185abe9f5b9224cda4f123b080407 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Przemys=C5=82aw=20Grzywacz?=
 <przemyslaw.grzywacz@golem.network>
Date: Fri, 20 Oct 2023 12:22:54 +0200
Subject: [PATCH] fix: task executor end() can be called multiple times JST-524

---
 .../{executor.test.ts => executor.spec.ts}    | 23 +++++++++++++
 src/executor/executor.ts                      | 32 ++++++++++++++++---
 2 files changed, 51 insertions(+), 4 deletions(-)
 rename src/executor/{executor.test.ts => executor.spec.ts} (77%)

diff --git a/src/executor/executor.test.ts b/src/executor/executor.spec.ts
similarity index 77%
rename from src/executor/executor.test.ts
rename to src/executor/executor.spec.ts
index d247338d3..b0379077c 100644
--- a/src/executor/executor.test.ts
+++ b/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,25 @@ 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();
+      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);
+    });
+  });
 });
diff --git a/src/executor/executor.ts b/src/executor/executor.ts
index 4f5a7c728..9ac4fe71f 100644
--- a/src/executor/executor.ts
+++ b/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.
    */
-  async end() {
-    if (runtimeContextChecker.isNode) this.removeSignalHandlers();
-    if (!this.isRunning) return;
+  end(): Promise<void> {
+    if (this.isRunning) {
+      this.endPromise = this.doEnd();
+    }
+
+    return this.endPromise!;
+  }
+
+  /**
+   * Perform everything needed to cleanly shut down the executor.
+   * @private
+   */
+  private async doEnd() {
     this.isRunning = false;
+    if (runtimeContextChecker.isNode) this.removeSignalHandlers();
     clearTimeout(this.startupTimeoutId);
     if (!this.configOptions.storageProvider) await this.storageProvider?.close();
     await this.networkService?.end();