diff --git a/src/Microsoft.ComponentDetection.Common/AsyncExecution.cs b/src/Microsoft.ComponentDetection.Common/AsyncExecution.cs index c0728eec5..62dfaa254 100644 --- a/src/Microsoft.ComponentDetection.Common/AsyncExecution.cs +++ b/src/Microsoft.ComponentDetection.Common/AsyncExecution.cs @@ -19,19 +19,19 @@ public static class AsyncExecution /// The result of the function. /// Thrown when is null. /// Thrown when the execution does not complete within the timeout. - public static async Task ExecuteWithTimeoutAsync(Func> toExecute, TimeSpan timeout, CancellationToken cancellationToken) + public static async Task ExecuteWithTimeoutAsync(Func> toExecute, TimeSpan timeout, CancellationToken cancellationToken = default) { ArgumentNullException.ThrowIfNull(toExecute); - var work = Task.Run(toExecute); - - var completedInTime = await Task.Run(() => work.Wait(timeout)); - if (!completedInTime) + var work = toExecute(); + try { - throw new TimeoutException($"The execution did not complete in the allotted time ({timeout.TotalSeconds} seconds) and has been terminated prior to completion"); + return await work.WaitAsync(timeout, cancellationToken); + } + catch (TimeoutException ex) + { + throw new TimeoutException($"The execution did not complete in the allotted time ({timeout.TotalSeconds} seconds) and has been terminated prior to completion", ex); } - - return await work; } /// @@ -43,15 +43,19 @@ public static async Task ExecuteWithTimeoutAsync(Func> toExecute, /// A representing the asynchronous operation. /// Thrown when is null. /// Thrown when the execution does not complete within the timeout. - public static async Task ExecuteVoidWithTimeoutAsync(Action toExecute, TimeSpan timeout, CancellationToken cancellationToken) + public static async Task ExecuteVoidWithTimeoutAsync(Action toExecute, TimeSpan timeout, CancellationToken cancellationToken = default) { ArgumentNullException.ThrowIfNull(toExecute); - var work = Task.Run(toExecute, cancellationToken); - var completedInTime = await Task.Run(() => work.Wait(timeout)); - if (!completedInTime) + var work = Task.Run(toExecute, CancellationToken.None); + + try + { + await work.WaitAsync(timeout, cancellationToken); + } + catch (TimeoutException ex) { - throw new TimeoutException($"The execution did not complete in the allotted time ({timeout.TotalSeconds} seconds) and has been terminated prior to completion"); + throw new TimeoutException($"The execution did not complete in the allotted time ({timeout.TotalSeconds} seconds) and has been terminated prior to completion", ex); } } }