dev-npgsql/npgsql: Npgsql is the .NET data provider for PostgreSQL.

src/Npgsql/TaskExtensions.cs
87 строк4 KB
Nikita Kazmin
Fixed passing a negative timeout to the cts while opening a connection (#3589)
10 мар 2021, 23:05
10 мар 2021, 23:050c05921
using System;
using System.Threading;
using System.Threading.Tasks;
using Npgsql.Util;

namespace Npgsql
{
    static class TaskExtensions
    {
        /// <summary>
        /// Utility that simplifies awaiting a task with a timeout. If the given task does not
        /// complete within <paramref name="timeout"/>, a <see cref="TimeoutException"/> is thrown.
        /// </summary>
        /// <param name="task">The task to be awaited</param>
        /// <param name="timeout">How much time to allow <paramref name="task"/> to complete before throwing a <see cref="TimeoutException"/></param>
        /// <returns>An awaitable task that represents the original task plus the timeout</returns>
        internal static async Task<T> WithTimeout<T>(this Task<T> task, NpgsqlTimeout timeout)
        {
            if (!timeout.IsSet)
                return await task;
            var timeLeft = timeout.CheckAndGetTimeLeft();
            if (task != await Task.WhenAny(task, Task.Delay(timeLeft)))
                throw new TimeoutException();
            return await task;
        }

        /// <summary>
        /// Allows you to cancel awaiting for a non-cancellable task.
        /// </summary>
        /// <remarks>
        /// Read https://blogs.msdn.com/b/pfxteam/archive/2012/10/05/how-do-i-cancel-non-cancelable-async-operations.aspx
        /// and be very careful with this.
        /// </remarks>
        internal static async Task<T> WithCancellation<T>(this Task<T> task, CancellationToken cancellationToken)
        {
            var tcs = new TaskCompletionSource<bool>();
            using (cancellationToken.Register(s => ((TaskCompletionSource<bool>)s!).TrySetResult(true), tcs))
                if (task != await Task.WhenAny(task, tcs.Task))
                    throw new TaskCanceledException(task);
            return await task;
        }

        internal static Task<T> WithCancellationAndTimeout<T>(this Task<T> task, NpgsqlTimeout timeout, CancellationToken cancellationToken)
            => task.WithCancellation(cancellationToken).WithTimeout(timeout);

#if NETSTANDARD2_0 || NETSTANDARD2_1 || NETCOREAPP3_1
        /// <summary>
        /// Utility that simplifies awaiting a task with a timeout. If the given task does not
        /// complete within <paramref name="timeout"/>, a <see cref="TimeoutException"/> is thrown.
        /// </summary>
        /// <param name="task">The task to be awaited</param>
        /// <param name="timeout">How much time to allow <paramref name="task"/> to complete before throwing a <see cref="TimeoutException"/></param>
        /// <returns>An awaitable task that represents the original task plus the timeout</returns>
        internal static async Task WithTimeout(this Task task, NpgsqlTimeout timeout)
        {
            if (!timeout.IsSet)
            {
                await task;
                return;
            }
            var timeLeft = timeout.CheckAndGetTimeLeft();
            if (task != await Task.WhenAny(task, Task.Delay(timeLeft)))
                throw new TimeoutException();
            await task;
        }

        /// <summary>
        /// Allows you to cancel awaiting for a non-cancellable task.
        /// </summary>
        /// <remarks>
        /// Read https://blogs.msdn.com/b/pfxteam/archive/2012/10/05/how-do-i-cancel-non-cancelable-async-operations.aspx
        /// and be very careful with this.
        /// </remarks>
        internal static async Task WithCancellation(this Task task, CancellationToken cancellationToken)
        {
            var tcs = new TaskCompletionSource<bool>();
            using (cancellationToken.Register(s => ((TaskCompletionSource<bool>)s!).TrySetResult(true), tcs))
                if (task != await Task.WhenAny(task, tcs.Task))
                    throw new TaskCanceledException(task);
            await task;
        }

        internal static Task WithCancellationAndTimeout(this Task task, NpgsqlTimeout timeout, CancellationToken cancellationToken)
            => task.WithCancellation(cancellationToken).WithTimeout(timeout);
#endif
    }
}