Managing cancellation tokens and handling timeouts tends to result in duplicated logic and unnecessary complexity.
What if we could simplify this by creating a reusable solution that handles cancellation and timeouts seamlessly, without repeating the same boilerplate logic every time?
The Challenge
When performing async operations, whether it’s processing a single task or a batch of tasks, we typically want to:
Cancel all ongoing tasks immediately if one task fails or times out.
Ensure no unnecessary work is done once cancellation is triggered.
Avoid duplicating cancellation logic in each task.
Without a clean, reusable approach, these requirements can lead to a cluttered codebase with lots of repeated logic.
The Solution
By encapsulating cancellation handling and timeouts into a reusable method, we can:
Centralize Cancellation Logic: No need to repeat the same code for cancellation, timeout, or cleanup in every task.
Eliminate Boilerplate
We no longer need to manually manage callbacks or check for cancellations in every operation.
Reuse the Logics
Write it once, and reuse it everywhere, whether for a single task or a batch of tasks.
Here’s how we can solve it:
we can create a method to handle:
- Timeouts: Automatically cancel the operation if it exceeds the specified duration.
- External Cancellation: Propagate cancellation from external tokens (e.g., batch cancellation).
- Callbacks: Invoke a cleanup callback when the operation is canceled.
Reusable method
//Reusable Method
public async Task<T> RunWithTimeoutAsync<T>(
Func<CancellationToken, Task<T>> asyncFunc,
TimeSpan? timeout = null,
Action? onCanceled = null,
CancellationToken externalToken = default)
{
//Validation...
using var cts = CancellationTokenSource.CreateLinkedTokenSource(externalToken);
if (timeout.HasValue)
{
cts.CancelAfter(timeout.Value);
}
using var registration = onCanceled != null ? cts.Token.Register(() =>
{
onCanceled.Invoke();
_logger.LogInformation("Cancellation detected! Performing cleanup...");
}) : default;
try
{
return await asyncFunc(cts.Token).ConfigureAwait(false);
}
catch (OperationCanceledException) when (cts.Token.IsCancellationRequested)
{
_logger.LogInformation("Operation was canceled (either by timeout or external token).");
throw;
}
catch (Exception ex)
{
_logger.LogError(ex, "Operation failed with an exception: {Message}", ex.Message);
throw;
}
}
async Task RunExampleAsync()
{
using var cts = new CancellationTokenSource();
try
{
var myTask = _asyncOperationService.RunWithTimeoutAsync(
async (cancellationToken) => await GetDataAsync(cancellationToken),
timeout: TimeSpan.FromSeconds(3),
onCanceled: () => Console.WriteLine("Sync was canceled due to network issues."),
externalToken: cts.Token
);
var result = await myTask;
Console.WriteLine(result);
}
catch (OperationCanceledException)
{
Console.WriteLine("Synchronization was canceled.");
}
}
Usage Example
//usages
using CancellationTokenSource cts = new();
await RunExampleWithCancellationAsync();
// Simulate a delay before canceling (e.g., cancel after 3 seconds)
await Task.Delay(3000);
cts.Cancel(); // Manually cancel the operation
await RunExampleAsync();
Output
![]()
Unit Testing this reusable approach
using System;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.Logging;
using Moq;
using Xunit;
public class RunWithTimeoutAsyncTests
{
private readonly Mock<ILogger<AsyncOperationService>> _loggerMock;
private readonly AsyncOperationService _asyncOperationService;
public RunWithTimeoutAsyncTests()
{
// Mock the logger
_loggerMock = new Mock<ILogger<AsyncOperationService>>();
// Create an instance of AsyncOperationService with the mocked logger
_asyncOperationService = new AsyncOperationService(_loggerMock.Object);
}
// Test 1: Completes successfully
[Fact]
public async Task CompletesSuccessfully()
{
// Arrange
async Task<int> AsyncFunc(CancellationToken token)
{
await Task.Delay(100, token); // Simulate async work
return 42;
}
// Act
int result = await _asyncOperationService.RunWithTimeoutAsync(AsyncFunc);
// Assert
Assert.Equal(42, result);
}
// Test 2: Throws on manual cancellation
[Fact]
public async Task ThrowsOnManualCancellation()
{
// Arrange
using var cts = new CancellationTokenSource();
async Task<int> AsyncFunc(CancellationToken token)
{
await Task.Delay(1000, token); // Simulate work
return 42;
}
// Act
cts.Cancel(); // Manually cancel
var exception = await Assert.ThrowsAsync<TaskCanceledException>(() =>
_asyncOperationService.RunWithTimeoutAsync(AsyncFunc, externalToken: cts.Token)
);
// Assert
Assert.True(cts.Token.IsCancellationRequested);
Assert.True(exception.CancellationToken.IsCancellationRequested);
}
// Test 3: Throws on timeout
[Fact]
public async Task ThrowsOnTimeout()
{
// Arrange
async Task<int> AsyncFunc(CancellationToken token)
{
await Task.Delay(1000, token); // Simulate work
return 42;
}
// Act
var exception = await Assert.ThrowsAsync<TaskCanceledException>(() =>
_asyncOperationService.RunWithTimeoutAsync(AsyncFunc, timeout: TimeSpan.FromMilliseconds(500))
);
// Assert
Assert.True(exception.CancellationToken.IsCancellationRequested);
}
// Test 4: Rethrows other exceptions
[Fact]
public async Task RethrowsOtherExceptions()
{
// Arrange
async Task<int> AsyncFunc(CancellationToken token)
{
await Task.Delay(100, token); // Simulate work
throw new InvalidOperationException("Something went wrong!");
}
// Act & Assert
var exception = await Assert.ThrowsAsync<InvalidOperationException>(() =>
_asyncOperationService.RunWithTimeoutAsync(AsyncFunc)
);
Assert.Equal("Something went wrong!", exception.Message);
}
// Test 5: Handles parallel execution
[Fact]
public async Task HandlesParallelExecution()
{
// Arrange
async Task<int> AsyncFunc1(CancellationToken token)
{
await Task.Delay(100, token); // Simulate work
return 1;
}
async Task<int> AsyncFunc2(CancellationToken token)
{
await Task.Delay(100, token); // Simulate work
return 2;
}
// Act
var task1 = _asyncOperationService.RunWithTimeoutAsync(AsyncFunc1);
var task2 = _asyncOperationService.RunWithTimeoutAsync(AsyncFunc2);
var results = await Task.WhenAll(task1, task2);
// Assert
Assert.Equal(new[] { 1, 2 }, results);
}
// Test 6: Respects linked cancellation
[Fact]
public async Task RespectsLinkedCancellation()
{
// Arrange
using var externalCts = new CancellationTokenSource();
async Task<int> AsyncFunc(CancellationToken token)
{
await Task.Delay(1000, token); // Simulate work
return 42;
}
// Act
externalCts.Cancel(); // Cancel the external token
var exception = await Assert.ThrowsAsync<TaskCanceledException>(() =>
_asyncOperationService.RunWithTimeoutAsync(AsyncFunc, externalToken: externalCts.Token)
);
// Assert
Assert.True(externalCts.Token.IsCancellationRequested);
Assert.True(exception.CancellationToken.IsCancellationRequested);
}
// Test 7: Works without timeout
[Fact]
public async Task WorksWithoutTimeout()
{
// Arrange
async Task<int> AsyncFunc(CancellationToken token)
{
await Task.Delay(100, token); // Simulate work
return 42;
}
// Act
int result = await _asyncOperationService.RunWithTimeoutAsync(AsyncFunc, timeout: null);
// Assert
Assert.Equal(42, result);
}
// Test 8: Throws on null async function
[Fact]
public async Task ThrowsOnNullAsyncFunction()
{
// Act & Assert
await Assert.ThrowsAsync<NullReferenceException>(() =>
_asyncOperationService.RunWithTimeoutAsync<int>(null!)
);
}
// Test 9: Calls onCanceled callback when canceled
[Fact]
public async Task CallsOnCanceledCallback()
{
// Arrange
using var cts = new CancellationTokenSource();
bool onCanceledCalled = false;
async Task<int> AsyncFunc(CancellationToken token)
{
await Task.Delay(1000, token); // Simulate work
return 42;
}
// Act
cts.Cancel(); // Manually cancel
var exception = await Assert.ThrowsAsync<TaskCanceledException>(() =>
_asyncOperationService.RunWithTimeoutAsync(
AsyncFunc,
externalToken: cts.Token,
onCanceled: () => onCanceledCalled = true
)
);
// Assert
Assert.True(onCanceledCalled);
Assert.True(cts.Token.IsCancellationRequested);
Assert.True(exception.CancellationToken.IsCancellationRequested);
}
// Test 10: Calls onCanceled callback on timeout
[Fact]
public async Task CallsOnCanceledCallbackOnTimeout()
{
// Arrange
bool onCanceledCalled = false;
async Task<int> AsyncFunc(CancellationToken token)
{
await Task.Delay(1000, token); // Simulate work, will throw if canceled
return 42;
}
// Act
var exception = await Assert.ThrowsAsync<TaskCanceledException>(() =>
_asyncOperationService.RunWithTimeoutAsync(
AsyncFunc,
timeout: TimeSpan.FromMilliseconds(500),
onCanceled: () => onCanceledCalled = true,
externalToken: CancellationToken.None // No external token for simplicity
)
);
// Assert
Assert.True(onCanceledCalled); // Make sure callback was called
Assert.True(exception.CancellationToken.IsCancellationRequested); // Ensure cancellation was requ
}
// Test 11: Does not call onCanceled callback on success
[Fact]
public async Task DoesNotCallOnCanceledCallbackOnSuccess()
{
// Arrange
bool onCanceledCalled = false;
async Task<int> AsyncFunc(CancellationToken token)
{
await Task.Delay(100, token); // Simulate work
return 42;
}
// Act
int result = await _asyncOperationService.RunWithTimeoutAsync(
AsyncFunc,
onCanceled: () => onCanceledCalled = true
);
// Assert
Assert.False(onCanceledCalled);
Assert.Equal(42, result);
}
// Test 12: Does not call onCanceled callback on other exceptions
[Fact]
public async Task DoesNotCallOnCanceledCallbackOnOtherExceptions()
{
// Arrange
bool onCanceledCalled = false;
async Task<int> AsyncFunc(CancellationToken token)
{
await Task.Delay(100, token); // Simulate work
throw new InvalidOperationException("Something went wrong!");
}
// Act & Assert
var exception = await Assert.ThrowsAsync<InvalidOperationException>(() =>
_asyncOperationService.RunWithTimeoutAsync(
AsyncFunc,
onCanceled: () => onCanceledCalled = true
)
);
Assert.False(onCanceledCalled);
Assert.Equal("Something went wrong!", exception.Message);
}
[Fact]
public async Task UpdatesDatabaseInParallel()
{
// Arrange
var options = new DbContextOptionsBuilder<AppDbContext>()
.UseInMemoryDatabase(databaseName: "TestDatabase")
.Options;
// Seed the database with some products
using (var dbContext = new AppDbContext(options))
{
dbContext.Products.AddRange(
new Product { Id = 1, Name = "Product 1", IsProcessed = false },
new Product { Id = 2, Name = "Product 2", IsProcessed = false },
new Product { Id = 3, Name = "Product 3", IsProcessed = false }
);
await dbContext.SaveChangesAsync();
}
// Function to process a product
async Task ProcessProductAsync(int productId, CancellationToken token)
{
using var dbContext = new AppDbContext(options); // Use a new DbContext for each task
var product = await dbContext.Products.FindAsync(productId);
if (product != null)
{
await Task.Delay(100, token); // Simulate work
product.IsProcessed = true;
await dbContext.SaveChangesAsync(token);
}
}
// Act
var productIds = new[] { 1, 2, 3 };
var tasks = productIds.Select(productId =>
_asyncOperationService.RunWithTimeoutNonReturningAsync(
async token =>
{
await ProcessProductAsync(productId, token);
},
timeout: TimeSpan.FromSeconds(10) // Timeout after 5 seconds
)
);
await Task.WhenAll(tasks);
// Assert
using var finalDbContext = new AppDbContext(options); // Use a final DbContext to check the results
var processedProducts = await finalDbContext.Products.Where(p => p.IsProcessed).ToListAsync();
Assert.Equal(3, processedProducts.Count); // All products should be processed
}
[Fact]
public async Task CompareBehaviors()
{
// Arrange: In-memory collection of products
var productsWithService = new List<Product>
{
new Product { Id = 1, Name = "Product 1", IsProcessed = false },
new Product { Id = 2, Name = "Product 2", IsProcessed = false },
new Product { Id = 3, Name = "Product 3", IsProcessed = false }
};
var productsManual = new List<Product>
{
new Product { Id = 1, Name = "Product 1", IsProcessed = false },
new Product { Id = 2, Name = "Product 2", IsProcessed = false },
new Product { Id = 3, Name = "Product 3", IsProcessed = false }
};
using var batchCtsWithService = new CancellationTokenSource();
using var batchCtsManual = new CancellationTokenSource();
// Process Product Logic with Simulated Failure for Product 2
async Task ProcessProductAsync(int productId, CancellationToken token, List<Product> products)
{
var product = products.FirstOrDefault(p => p.Id == productId);
if (product == null)
{
throw new Exception($"Product with ID {productId} not found.");
}
// Exit early if cancellation is requested
token.ThrowIfCancellationRequested();
if (productId == 2)
{
// Simulate a delay to allow Product 1 to complete processing
await Task.Delay(2000); // Increased delay to ensure Product 1 completes
if (products == productsWithService)
batchCtsWithService.Cancel(); // Trigger cancellation for service test
else
batchCtsManual.Cancel(); // Trigger cancellation for manual test
throw new InvalidOperationException("Simulated failure for product 2");
}
// Simulate work with frequent cancellation checks
for (int i = 0; i < 10; i++)
{
token.ThrowIfCancellationRequested(); // Check for cancellation
await Task.Delay(50, token); // Simulate smaller chunks of work
}
// Check for cancellation again before final processing
token.ThrowIfCancellationRequested();
product.IsProcessed = true;
}
// Behavior with RunWithTimeoutAsync
var productIds = new[] { 1, 2, 3 };
var tasksWithService = productIds.Select(async productId =>
{
try
{
await _asyncOperationService.RunWithTimeoutAsync(
async token =>
{
await ProcessProductAsync(productId, token, productsWithService);
return true; // Return a dummy value
},
timeout: TimeSpan.FromSeconds(5),
externalToken: batchCtsWithService.Token
);
}
catch (Exception ex)
{
Console.WriteLine($"Error in processing product {productId} (RunWithTimeoutAsync): {ex.Message}");
if (batchCtsWithService.Token.IsCancellationRequested)
{
Console.WriteLine($"Task for product {productId} is cancelled due to batch failure.");
}
// Ensure that we re-throw to let the test capture the failure
throw;
}
}).ToArray();
// Behavior with Manual Token Handling
var tasksManual = productIds.Select(async productId =>
{
try
{
await ProcessProductAsync(productId, batchCtsManual.Token, productsManual);
}
catch (Exception ex)
{
Console.WriteLine($"Error in processing product {productId} (Manual): {ex.Message}");
if (batchCtsManual.Token.IsCancellationRequested)
{
Console.WriteLine($"Task for product {productId} is cancelled due to batch failure.");
}
// Ensure that we re-throw to let the test capture the failure
throw;
}
}).ToArray();
try
{
// Wait for all tasks and assert that cancellation happened as expected
await Task.WhenAll(tasksWithService);
}
catch (InvalidOperationException)
{
// Expected exception for product 2
}
try
{
// Wait for all tasks and assert that cancellation happened as expected
await Task.WhenAll(tasksManual);
}
catch (InvalidOperationException)
{
// Expected exception for product 2
}
// Assert: Ensure that both behaviors produce the same results
var processedProductsWithService = productsWithService.Where(p => p.IsProcessed).ToList();
var processedProductsManual = productsManual.Where(p => p.IsProcessed).ToList();
// Debugging: Print the processed products
Console.WriteLine("Processed Products (RunWithTimeoutAsync):");
foreach (var product in processedProductsWithService)
{
Console.WriteLine($"Product ID: {product.Id}, Name: {product.Name}, IsProcessed: {product.IsProcessed}");
}
Console.WriteLine("Processed Products (Manual):");
foreach (var product in processedProductsManual)
{
Console.WriteLine($"Product ID: {product.Id}, Name: {product.Name}, IsProcessed: {product.IsProcessed}");
}
// Assert that both behaviors produce the same results
Assert.Equal(processedProductsWithService.Count, processedProductsManual.Count);
Assert.All(processedProductsWithService, p => Assert.True(processedProductsManual.Any(m => m.Id == p.Id && m.IsProcessed == p.IsProcessed)));
// Ensure cancellation was correctly propagated to all tasks
Assert.True(batchCtsWithService.Token.IsCancellationRequested, "Batch cancellation was not triggered as expected (RunWithTimeoutAsync).");
Assert.True(batchCtsManual.Token.IsCancellationRequested, "Batch cancellation was not triggered as expected (Manual).");
}
}
// Sample DbContext and Product entity
public class AppDbContext : DbContext
{
public DbSet<Product> Products { get; set; }
public AppDbContext(DbContextOptions<AppDbContext> options) : base(options) { }
}
public class Product
{
public int Id { get; set; }
public string Name { get; set; }
public bool IsProcessed { get; set; }
}
Considerations
Custom Logic for Specific Scenarios:
Depending on our specific requirements, We might need to tailor the cancellation logic. For instance, if we have tasks where no result needs to be returned.
- Error Handling: The method focuses on handling cancellations gracefully, but complex workflows may require more specific error handling.
- External Token Propagation: If we’re passing external cancellation tokens, ensure that they are properly managed and not overused.
- Callback Execution: The callback is invoked during the cancellation, so ensure that cleanup logic is lightweight and doesn’t introduce additional side effects or delays.
The source code is Attached. Thanks for reading.