JST-252 | JST-250: batch end and endStream error handling

src/activity/activity.ts

@@ -93,7 +93,7 @@ export class Activity {
    * @param timeout - execution timeout
    */
   public async execute(script: ExeScriptRequest, stream?: boolean, timeout?: number): Promise<Readable> {
-    let batchId, batchSize;
+    let batchId: string, batchSize: number;
     let startTime = new Date();
     try {
       batchId = await this.send(script);
@@ -169,7 +169,7 @@ export class Activity {
 
   private async pollingBatch(batchId, startTime, timeout): Promise<Readable> {
     let isBatchFinished = false;
-    let lastIndex;
+    let lastIndex: number;
     let retryCount = 0;
     const maxRetries = 5;
     const { id: activityId, agreementId } = this;

src/task/batch.spec.ts

@@ -112,8 +112,7 @@ describe("Batch", () => {
       expect(spy).toHaveBeenCalled();
     });
 
-    // FIXME: Not working due to bug: JST-250
-    xit("should call script.after() on execute error", async () => {
+    it("should call script.after() on execute error", async () => {
       const spy = jest.spyOn(batch["script"], "after");
       jest.spyOn(activity, "execute").mockRejectedValue(new Error("ERROR"));
 
@@ -168,8 +167,7 @@ describe("Batch", () => {
       expect(spy).toHaveBeenCalled();
     });
 
-    // FIXME: Not working due to bug: JST-252
-    xit("should call script.after() on result stream error", async () => {
+    it("should call script.after() on result stream error", async () => {
       const spy = jest.spyOn(batch["script"], "after");
       activity.mockResultFailure("FAILURE");
 
@@ -187,8 +185,7 @@ describe("Batch", () => {
       expect(spy).toHaveBeenCalled();
     });
 
-    // FIXME: Not working due to bug: JST-250
-    xit("should call script.after() on execute error", async () => {
+    it("should call script.after() on execute error", async () => {
       const spy = jest.spyOn(batch["script"], "after");
       jest.spyOn(activity, "execute").mockRejectedValue(new Error("ERROR"));
 
@@ -201,5 +198,20 @@ describe("Batch", () => {
 
       expect(spy).toHaveBeenCalled();
     });
+
+    it("should destroy the stream on result stream error", async () => {
+      activity.mockResultFailure("FAILURE");
+      const stream = await batch.endStream();
+      try {
+        // eslint-disable-next-line @typescript-eslint/no-unused-vars
+        for await (const r of stream) {
+          /* empty */
+        }
+        fail("Expected to throw");
+      } catch (e) {
+        /* empty */
+      }
+      expect(stream.destroyed).toBe(true);
+    });
   });
 });

src/task/batch.ts

@@ -2,7 +2,7 @@ import { DownloadFile, Run, Script, UploadFile } from "../script";
 import { Activity, Result } from "../activity";
 import { StorageProvider } from "../storage/provider";
 import { Logger, sleep } from "../utils";
-import { Readable, Transform } from "stream";
+import { Readable, Transform, pipeline } from "stream";
 import { UploadData } from "../script/command";
 
 export class Batch {
@@ -73,7 +73,14 @@ export class Batch {
   async end(): Promise<Result[]> {
     await this.script.before();
     await sleep(100, true);
-    const results = await this.activity.execute(this.script.getExeScriptRequest());
+    let results: Readable;
+    try {
+      results = await this.activity.execute(this.script.getExeScriptRequest());
+    } catch (error) {
+      // the original error is more important than the one from after()
+      await this.script.after([]).catch();
+      throw error;
+    }
     const allResults: Result[] = [];
     return new Promise((resolve, reject) => {
       results.on("data", (res) => {
@@ -99,7 +106,14 @@ export class Batch {
   async endStream(): Promise<Readable> {
     const script = this.script;
     await script.before();
-    const results = await this.activity.execute(this.script.getExeScriptRequest());
+    let results: Readable;
+    try {
+      results = await this.activity.execute(this.script.getExeScriptRequest());
+    } catch (error) {
+      // the original error is more important than the one from after()
+      await script.after([]).catch();
+      throw error;
+    }
     const decodedResults: Result[] = [];
     const errorResultHandler = new Transform({
       objectMode: true,
@@ -118,11 +132,9 @@ export class Batch {
         }
       },
     });
-    results.on("end", () => this.script.after(decodedResults).catch());
-    results.on("error", (error) => {
+    const resultsWithErrorHandling = pipeline(results, errorResultHandler, () => {
       script.after(decodedResults).catch();
-      results.destroy(error);
     });
-    return results.pipe(errorResultHandler);
+    return resultsWithErrorHandling;
   }
 }