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);
}
}
}