use ContinueWith rather than finally for coroutine result types

microsoft · Jul 23, 2024 · 9210fbd · 9210fbd
1 parent 8d1cd07
commit 9210fbd
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/src/PowerShell/CommonFiles/PowerShellCmdlet.cs b/src/PowerShell/CommonFiles/PowerShellCmdlet.cs
@@ -97,11 +97,14 @@ internal Task RunOnMTA(Func<Task> func)
                 this.Write(StreamType.Verbose, "Already running on MTA");
                 try
                 {
-                    return func();
+                    Task result = func();
+                    result.ContinueWith((task) => this.Complete(), TaskContinuationOptions.ExecuteSynchronously);
+                    return result;
                 }
-                finally
+                catch
                 {
                     this.Complete();
+                    throw;
                 }
             }
 
@@ -150,11 +153,14 @@ internal Task<TResult> RunOnMTA<TResult>(Func<Task<TResult>> func)
                 this.Write(StreamType.Verbose, "Already running on MTA");
                 try
                 {
-                    return func();
+                    Task<TResult> result = func();
+                    result.ContinueWith((task) => this.Complete(), TaskContinuationOptions.ExecuteSynchronously);
+                    return result;
                 }
-                finally
+                catch
                 {
                     this.Complete();
+                    throw;
                 }
             }