AsyncLock with timeout

[mikernet]

I know there has been a couple issues about this, but I wanted to chime in on the topic. The code to do this properly and efficiently isn't exactly straightforward, so I think it would be a valuable addition. I had to add this to my project:

public static IDisposable Lock(this AsyncLock asyncLock, TimeSpan timeout, CancellationToken cancellationToken = default)
{
    if (timeout == Timeout.InfiniteTimeSpan)
        return asyncLock.Lock(cancellationToken);

    if (timeout == TimeSpan.Zero)
    {
        try
        {
            return asyncLock.Lock(new CancellationToken(true));
        }
        catch (OperationCanceledException)
        {
            throw new TimeoutException("Failed to acquire lock immediately.");
        }
    }

    using var timeoutCts = new CancellationTokenSource(timeout);
    using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCts.Token);

    try
    {
        return asyncLock.Lock(linkedCts.Token);
    }
    catch (OperationCanceledException) when (timeoutCts.IsCancellationRequested)
    {
        throw new TimeoutException("Failed to acquire lock within the specified timeout.");
    }
}

public static AwaitableDisposable<IDisposable> LockAsync(this AsyncLock asyncLock, TimeSpan timeout, CancellationToken cancellationToken = default)
{
    return new(Impl(asyncLock, timeout, cancellationToken));

    static async Task<IDisposable> Impl(AsyncLock asyncLock, TimeSpan timeout, CancellationToken cancellationToken)
    {
        if (timeout == Timeout.InfiniteTimeSpan)
            return await asyncLock.LockAsync(cancellationToken);

        if (timeout == TimeSpan.Zero)
        {
            try
            {
                return await asyncLock.LockAsync(new CancellationToken(true));
            }
            catch (OperationCanceledException)
            {
                throw new TimeoutException("Failed to acquire lock immediately.");
            }
        }

        using var timeoutCts = new CancellationTokenSource(timeout);
        using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCts.Token);

        try
        {
            return await asyncLock.LockAsync(linkedCts.Token);
        }
        catch (OperationCanceledException) when (timeoutCts.IsCancellationRequested)
        {
            throw new TimeoutException("Failed to acquire lock within the specified timeout.");
        }
    }
}

[mikernet]

I'll add that without a method that takes a timeout and handles the conversion to TimeoutException like above, the following pattern is incorrect (as I'm sure you know), but the mistake is easily overlooked:

public async Task M(CancellationToken cancellationToken)
{
    using var timeoutCts = new CancellationTokenSource(timeout);
    using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCts.Token);
    
    try
    {
        using (await _asyncLock.Lock(linkedCts))
        {
            await DoStuffAsync(cancellationToken);
        }
    }
    catch (OperationCancelledException) when (timeoutCts.IsCancellationRequested)
    {
        throw new TimeoutException(...);
    }
}

I think providing a pit of success and not requiring users to have a nuanced understanding of all these concerns when dealing with the very common timeout scenario would definitely be beneficial.