Reusable Approach to Handling Cancellation Token Logic & Timeout

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.