From 6c1bee4f62cb1e4f955d7cf6ea57b8ca8f5a738c Mon Sep 17 00:00:00 2001
From: Jeff Fisher <jeffish@microsoft.com>
Date: Wed, 17 Aug 2022 14:33:20 -0500
Subject: [PATCH] Fix issue with afterPhase policies getting run later than
 expected

---
 sdk/core/core-rest-pipeline/src/pipeline.ts      |  8 ++++++++
 .../core-rest-pipeline/test/pipeline.spec.ts     | 16 ++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/sdk/core/core-rest-pipeline/src/pipeline.ts b/sdk/core/core-rest-pipeline/src/pipeline.ts
index 717d13b7aa92..81ac7eef20d9 100644
--- a/sdk/core/core-rest-pipeline/src/pipeline.ts
+++ b/sdk/core/core-rest-pipeline/src/pipeline.ts
@@ -108,6 +108,7 @@ interface Phase {
   name: PipelinePhase | "None";
   policies: Set<PolicyGraphNode>;
   hasRun: boolean;
+  hasAfterPolicies: boolean;
 }
 
 /**
@@ -236,6 +237,7 @@ class HttpPipeline implements Pipeline {
         name,
         policies: new Set<PolicyGraphNode>(),
         hasRun: false,
+        hasAfterPolicies: false,
       };
     }
 
@@ -279,6 +281,7 @@ class HttpPipeline implements Pipeline {
       };
       if (options.afterPhase) {
         node.afterPhase = getPhase(options.afterPhase);
+        node.afterPhase.hasAfterPolicies = true;
       }
       policyMap.set(policyName, node);
       const phase = getPhase(options.phase);
@@ -358,6 +361,11 @@ class HttpPipeline implements Pipeline {
           // Don't proceed to the next phase until this phase finishes.
           return;
         }
+
+        if (phase.hasAfterPolicies) {
+          // Run any policies unblocked by this phase
+          walkPhase(noPhase);
+        }
       }
     }
 
diff --git a/sdk/core/core-rest-pipeline/test/pipeline.spec.ts b/sdk/core/core-rest-pipeline/test/pipeline.spec.ts
index bcf5976836d6..f9992ffdc2fb 100644
--- a/sdk/core/core-rest-pipeline/test/pipeline.spec.ts
+++ b/sdk/core/core-rest-pipeline/test/pipeline.spec.ts
@@ -423,4 +423,20 @@ describe("HttpsPipeline", function () {
     const policies = pipeline.getOrderedPolicies();
     assert.deepStrictEqual(policies, [testPolicy]);
   });
+
+  it("afterPhase respects phase ordering", function () {
+    const pipeline = createEmptyPipeline();
+    const testPolicy: PipelinePolicy = {
+      sendRequest: (request, next) => next(request),
+      name: "test",
+    };
+    const testPolicy2: PipelinePolicy = {
+      sendRequest: (request, next) => next(request),
+      name: "test2",
+    };
+    pipeline.addPolicy(testPolicy, { afterPhase: "Retry" });
+    pipeline.addPolicy(testPolicy2, { phase: "Sign" });
+    const policies = pipeline.getOrderedPolicies();
+    assert.deepStrictEqual(policies, [testPolicy, testPolicy2]);
+  });
 });