natsdotnet/tests/NATS.Server.TestUtilities/PollHelper.cs

namespace NATS.Server.TestUtilities;

/// <summary>
/// Provides bounded polling helpers for test fixtures that need to wait
/// for asynchronous conditions across independent components (e.g. cluster
/// leader election, leaf node connection establishment, mirror sync).
/// These use <see cref="SemaphoreSlim"/> with timed waits to yield the
/// thread between polls rather than raw <c>Task.Delay</c>.
/// </summary>
public static class PollHelper
{
    /// <summary>
    /// Polls <paramref name="condition"/> at <paramref name="intervalMs"/> intervals
    /// until it returns <c>true</c> or <paramref name="timeoutMs"/> elapses.
    /// Returns <c>true</c> if the condition was met, <c>false</c> on timeout.
    /// </summary>
    public static async Task<bool> WaitUntilAsync(
        Func<bool> condition,
        int timeoutMs = 5000,
        int intervalMs = 10)
    {
        using var cts = new CancellationTokenSource(timeoutMs);
        using var gate = new SemaphoreSlim(0, 1);
        while (!cts.IsCancellationRequested)
        {
            if (condition())
                return true;

            try
            {
                await gate.WaitAsync(intervalMs, cts.Token);
            }
            catch (OperationCanceledException)
            {
                break;
            }
        }

        return condition();
    }

    /// <summary>
    /// Polls <paramref name="condition"/> (async overload) at <paramref name="intervalMs"/>
    /// intervals until it returns <c>true</c> or <paramref name="timeoutMs"/> elapses.
    /// Returns <c>true</c> if the condition was met, <c>false</c> on timeout.
    /// </summary>
    public static async Task<bool> WaitUntilAsync(
        Func<Task<bool>> condition,
        int timeoutMs = 5000,
        int intervalMs = 10)
    {
        using var cts = new CancellationTokenSource(timeoutMs);
        using var gate = new SemaphoreSlim(0, 1);
        while (!cts.IsCancellationRequested)
        {
            if (await condition())
                return true;

            try
            {
                await gate.WaitAsync(intervalMs, cts.Token);
            }
            catch (OperationCanceledException)
            {
                break;
            }
        }

        return await condition();
    }

    /// <summary>
    /// Polls <paramref name="condition"/> and throws <see cref="TimeoutException"/>
    /// with <paramref name="message"/> if the condition is not met within <paramref name="timeoutMs"/>.
    /// </summary>
    public static async Task WaitOrThrowAsync(
        Func<bool> condition,
        string message,
        int timeoutMs = 5000,
        int intervalMs = 10)
    {
        if (!await WaitUntilAsync(condition, timeoutMs, intervalMs))
            throw new TimeoutException(message);
    }

    /// <summary>
    /// Polls <paramref name="condition"/> (async overload) and throws
    /// <see cref="TimeoutException"/> with <paramref name="message"/> if the condition
    /// is not met within <paramref name="timeoutMs"/>.
    /// </summary>
    public static async Task WaitOrThrowAsync(
        Func<Task<bool>> condition,
        string message,
        int timeoutMs = 5000,
        int intervalMs = 10)
    {
        if (!await WaitUntilAsync(condition, timeoutMs, intervalMs))
            throw new TimeoutException(message);
    }

    /// <summary>
    /// Yields the current task for approximately <paramref name="delayMs"/> using a
    /// semaphore-based timed wait rather than <c>Task.Delay</c>.
    /// </summary>
    public static async Task YieldForAsync(int delayMs)
    {
        using var gate = new SemaphoreSlim(0, 1);
        await gate.WaitAsync(delayMs);
    }
}