From 48871381a169888f1d29275ab25915b0d815d1c1 Mon Sep 17 00:00:00 2001
From: Charles Lyding <19598772+clydin@users.noreply.github.com>
Date: Tue, 4 Apr 2023 12:35:01 -0400
Subject: [PATCH] fix(@angular-devkit/architect): allow registered builder
 teardowns to execute

Previously, the base job handler was completing the entire job before any teardowns
could attempt to execute.
---
 .../angular_devkit/architect/index.md         |  2 ++
 .../architect/src/create-builder.ts           | 19 +++------------
 .../architect/src/jobs/create-job-handler.ts  | 23 +++++++++++++++----
 3 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/goldens/public-api/angular_devkit/architect/index.md b/goldens/public-api/angular_devkit/architect/index.md
index 70cdd28738ab..1dc0c183e10e 100644
--- a/goldens/public-api/angular_devkit/architect/index.md
+++ b/goldens/public-api/angular_devkit/architect/index.md
@@ -451,6 +451,8 @@ export function scheduleTargetAndForget(context: BuilderContext, target: Target,
 
 // @public
 interface SimpleJobHandlerContext<A extends JsonValue, I extends JsonValue, O extends JsonValue> extends JobHandlerContext<A, I, O> {
+    // (undocumented)
+    addTeardown(teardown: () => Promise<void> | void): void;
     // (undocumented)
     createChannel: (name: string) => Observer<JsonValue>;
     // (undocumented)
diff --git a/packages/angular_devkit/architect/src/create-builder.ts b/packages/angular_devkit/architect/src/create-builder.ts
index 2f7f8c85d9df..d832c7e67755 100644
--- a/packages/angular_devkit/architect/src/create-builder.ts
+++ b/packages/angular_devkit/architect/src/create-builder.ts
@@ -47,9 +47,8 @@ export function createBuilder<OptT = json.JsonObject, OutT extends BuilderOutput
     const scheduler = context.scheduler;
     const progressChannel = context.createChannel('progress');
     const logChannel = context.createChannel('log');
+    const addTeardown = context.addTeardown.bind(context);
     let currentState: BuilderProgressState = BuilderProgressState.Stopped;
-    const teardownLogics: Array<() => PromiseLike<void> | void> = [];
-    let tearingDown = false;
     let current = 0;
     let status = '';
     let total = 1;
@@ -83,18 +82,8 @@ export function createBuilder<OptT = json.JsonObject, OutT extends BuilderOutput
 
       const inputSubscription = context.inboundBus.subscribe((i) => {
         switch (i.kind) {
-          case JobInboundMessageKind.Stop:
-            // Run teardown logic then complete.
-            tearingDown = true;
-            Promise.all(teardownLogics.map((fn) => fn() || Promise.resolve())).then(
-              () => observer.complete(),
-              (err) => observer.error(err),
-            );
-            break;
           case JobInboundMessageKind.Input:
-            if (!tearingDown) {
-              onInput(i.value);
-            }
+            onInput(i.value);
             break;
         }
       });
@@ -209,9 +198,7 @@ export function createBuilder<OptT = json.JsonObject, OutT extends BuilderOutput
                 progress({ state: currentState, current, total, status }, context);
             }
           },
-          addTeardown(teardown: () => Promise<void> | void): void {
-            teardownLogics.push(teardown);
-          },
+          addTeardown,
         };
 
         context.reportRunning();
diff --git a/packages/angular_devkit/architect/src/jobs/create-job-handler.ts b/packages/angular_devkit/architect/src/jobs/create-job-handler.ts
index 36952e56a2ff..531c21e5fdc5 100644
--- a/packages/angular_devkit/architect/src/jobs/create-job-handler.ts
+++ b/packages/angular_devkit/architect/src/jobs/create-job-handler.ts
@@ -44,6 +44,7 @@ export interface SimpleJobHandlerContext<
 > extends JobHandlerContext<A, I, O> {
   createChannel: (name: string) => Observer<JsonValue>;
   input: Observable<I>;
+  addTeardown(teardown: () => Promise<void> | void): void;
 }
 
 /**
@@ -72,6 +73,8 @@ export function createJobHandler<A extends JsonValue, I extends JsonValue, O ext
     const inboundBus = context.inboundBus;
     const inputChannel = new Subject<I>();
     let subscription: Subscription;
+    const teardownLogics: Array<() => PromiseLike<void> | void> = [];
+    let tearingDown = false;
 
     return new Observable<JobOutboundMessage<O>>((subject) => {
       function complete() {
@@ -91,13 +94,22 @@ export function createJobHandler<A extends JsonValue, I extends JsonValue, O ext
             break;
 
           case JobInboundMessageKind.Stop:
-            // There's no way to cancel a promise or a synchronous function, but we do cancel
-            // observables where possible.
-            complete();
+            // Run teardown logic then complete.
+            tearingDown = true;
+            if (teardownLogics.length) {
+              Promise.all(teardownLogics.map((fn) => fn())).then(
+                () => complete(),
+                () => complete(),
+              );
+            } else {
+              complete();
+            }
             break;
 
           case JobInboundMessageKind.Input:
-            inputChannel.next(message.value);
+            if (!tearingDown) {
+              inputChannel.next(message.value);
+            }
             break;
         }
       });
@@ -108,6 +120,9 @@ export function createJobHandler<A extends JsonValue, I extends JsonValue, O ext
       const newContext: SimpleJobHandlerContext<A, I, O> = {
         ...context,
         input: inputChannel.asObservable(),
+        addTeardown(teardown: () => Promise<void> | void): void {
+          teardownLogics.push(teardown);
+        },
         createChannel(name: string) {
           if (channels.has(name)) {
             throw new ChannelAlreadyExistException(name);