Skip to content

Instantly share code, notes, and snippets.

@brendankowitz
Last active January 24, 2024 11:24
Show Gist options
  • Save brendankowitz/5949970076952746a083054559377e56 to your computer and use it in GitHub Desktop.
Save brendankowitz/5949970076952746a083054559377e56 to your computer and use it in GitHub Desktop.
What is a good pattern for Awaiting Semaphores with a CancellationToken?
/// <summary>
/// Allows a semaphore to release with the IDisposable pattern
/// </summary>
/// <remarks>
/// Solves an issue where using the pattern:
/// <code>
/// try { await sem.WaitAsync(cancellationToken); }
/// finally { sem.Release(); }
/// </code>
/// Can result in SemaphoreFullException if the token is cancelled and the
/// the semaphore is not incremented.
/// </remarks>
public static class CancellableSemaphore
{
public static async Task<IDisposable> WaitWithCancellationAsync(this SemaphoreSlim semaphore, CancellationToken cancellationToken)
{
var task = semaphore.WaitAsync(cancellationToken);
await task.ConfigureAwait(false);
return new CancellableSemaphoreInternal(semaphore, task);
}
private class CancellableSemaphoreInternal : IDisposable
{
private readonly SemaphoreSlim _semaphoreSlim;
private Task _awaitTask;
public CancellableSemaphoreInternal(SemaphoreSlim semaphoreSlim, Task awaitTask)
{
_semaphoreSlim = semaphoreSlim;
_awaitTask = awaitTask;
}
public void Dispose()
{
if (_awaitTask?.Status == TaskStatus.RanToCompletion)
{
_semaphoreSlim.Release();
_awaitTask = null;
}
}
}
}
public class ResultProvider
{
MyResult _cachedResult;
private SemaphoreSlim _sem = new SemaphoreSlim(1, 1);
public async Task<MyResult> GetMyResultAsync(CancellationToken cancellationToken)
{
if (_cachedResult == null)
{
Task aquireSemaphoreTask = null;
try
{
aquireSemaphoreTask = _sem.WaitAsync(cancellationToken);
await aquireSemaphoreTask;
if (_cachedResult == null)
{
// ... Work ...
}
}
finally
{
// If the CancellationToken was cancelled while waiting for the Semaphore, it should not be released.
if (aquireSemaphoreTask?.Status == TaskStatus.RanToCompletion)
{
_sem.Release();
}
}
}
return _cachedResult;
}
}
public class ResultProvider
{
MyResult _cachedResult;
private SemaphoreSlim _sem = new SemaphoreSlim(1, 1);
public async Task<MyResult> GetMyResultAsync(CancellationToken cancellationToken)
{
if (_cachedResult == null)
{
try
{
await _sem.WaitAsync(cancellationToken);
if (_cachedResult == null)
{
// ... Work ...
}
}
finally
{
// This can throw System.Threading.SemaphoreFullException if the Wait() was cancelled
_sem.Release();
}
}
return _cachedResult;
}
}
@Sharkchenkos
Copy link

Nice. We've arrived to a similar pattern with checking TaskStatus.RanToCompletion in the finally block and I was just considering writing something of this sort. You can also use this with the new using syntax that doesn't require you to define an extra scope brackets. You can just add a using var lock = at the top and the whole method becomes a critical section. Can you provide context why you choose to add ConfigureAwait(false). I think it's dangerous to be put in something generic utility like this.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment