cursor-lang/CursorLang.Tests.Shared/Pump.cs

using System.Collections.Concurrent;
using System.Runtime.InteropServices;

namespace CursorLang.Tests.Shared;

/// <summary>
/// A thread with a Win32 message loop on it — what the tests have instead of a
/// dispatcher.
/// </summary>
/// <remarks>
/// Core's timers are <c>SetTimer</c> timers, and they only tick where messages are
/// pumped. WPF used to provide that thread; Core must not depend on WPF, so the tests
/// provide it themselves. The loop is a real one, so timers fire on their own and a
/// test only has to await the consequences through <see cref="WaitFor"/>.
///
/// One thread for the whole run: starting and stopping message loops between tests
/// costs more than it proves, and the hook, the timers and the windows under test are
/// happy to share.
///
/// Waiting is allowed from the pump thread itself, and that is the delicate part. A
/// plain wait there would stop the queue and with it everything being waited for, so
/// on that thread the waiting is done by a nested loop that keeps dispatching.
/// </remarks>
public static class Pump
{
    private static readonly Lock Gate = new();
    private static readonly TimeSpan DefaultTimeout = TimeSpan.FromSeconds(5);
    private static readonly ConcurrentQueue<Action> Posted = new();

    private static uint _threadId;

    /// <summary>Runs an action on the pump thread and waits for it to finish.</summary>
    public static void Run(Action action) => Run<object?>(() =>
    {
        action();
        return null;
    });

    /// <summary>The same for an action that returns a result.</summary>
    public static TResult Run<TResult>(Func<TResult> action)
    {
        EnsureStarted();

        if (IsOnPumpThread)
        {
            return action();
        }

        var done = new TaskCompletionSource<TResult>(TaskCreationOptions.RunContinuationsAsynchronously);

        Posted.Enqueue(() =>
        {
            try
            {
                done.SetResult(action());
            }
            catch (Exception e)
            {
                done.SetException(e);
            }
        });

        Wake();

        return done.Task.GetAwaiter().GetResult();
    }

    /// <summary>
    /// Hands an action to the pump without waiting for it. What the agent gives the
    /// keyboard hook, so that a test sees the same deferral the application does.
    /// </summary>
    public static void Post(Action action)
    {
        EnsureStarted();
        Posted.Enqueue(action);
        Wake();
    }

    /// <summary>
    /// Runs an action on a separate STA thread and waits for it to finish.
    /// </summary>
    /// <remarks>
    /// Needed where the foreignness of the thread is the point: kernel objects such as
    /// a mutex let their own owner in again, so a "second instance" of the application
    /// on the same thread does not count as second.
    /// </remarks>
    public static void RunApart(Action action)
    {
        Exception? failure = null;

        var thread = new Thread(() =>
        {
            try
            {
                action();
            }
            catch (Exception e)
            {
                failure = e;
            }
        })
        {
            IsBackground = true,
            Name = "CursorLang.Tests apart",
        };

        thread.SetApartmentState(ApartmentState.STA);
        thread.Start();
        thread.Join();

        if (failure is not null)
        {
            throw new InvalidOperationException("The action on the separate thread failed", failure);
        }
    }

    /// <summary>
    /// Waits until everything posted to the pump has run. Work handed over with
    /// <see cref="Post"/> has happened by the time this returns.
    /// </summary>
    public static void Drain()
    {
        EnsureStarted();

        if (IsOnPumpThread)
        {
            DispatchPending();
            return;
        }

        Run(static () => { });
    }

    /// <summary>
    /// Waits for a condition, letting the pump get on with its work.
    /// </summary>
    /// <remarks>
    /// A timeout is reported as an exception rather than as a failed assertion so that
    /// this assembly needs no test framework. To the test the difference is the wording
    /// of the failure, and <paramref name="because"/> carries that either way.
    /// </remarks>
    /// <exception cref="TimeoutException">The condition never came true.</exception>
    public static void WaitFor(Func<bool> condition, string because, TimeSpan? timeout = null)
    {
        TimeSpan limit = timeout ?? DefaultTimeout;
        DateTime deadline = DateTime.UtcNow + limit;

        while (!condition())
        {
            if (DateTime.UtcNow >= deadline)
            {
                throw new TimeoutException($"Waited {limit.TotalMilliseconds:N0} ms in vain: {because}");
            }

            Idle(TimeSpan.FromMilliseconds(5));
        }
    }

    /// <summary>
    /// Waits for the given time while the pump keeps running: that is how
    /// "nothing happened during this time" is verified.
    /// </summary>
    public static void Pause(TimeSpan duration)
    {
        Idle(duration);
        Drain();
    }

    private static bool IsOnPumpThread => _threadId != 0 && GetCurrentThreadId() == _threadId;

    // A wait during which the queue still gets its chance to run. On any thread but the
    // pump's that is a plain sleep — the pump is elsewhere and busy on its own
    private static void Idle(TimeSpan duration)
    {
        EnsureStarted();

        if (!IsOnPumpThread)
        {
            Thread.Sleep(duration);
            return;
        }

        DateTime deadline = DateTime.UtcNow + duration;

        while (DateTime.UtcNow < deadline)
        {
            if (!DispatchPending())
            {
                // Nothing waiting: give the timers a moment to post something rather
                // than spinning through the deadline at full speed
                Thread.Sleep(1);
            }
        }

        DispatchPending();
    }

    /// <summary>
    /// Runs everything already in the queue — messages and posted work alike — and
    /// says whether there was anything. Only ever called on the pump thread.
    /// </summary>
    private static bool DispatchPending()
    {
        var any = false;

        while (Posted.TryDequeue(out Action? action))
        {
            action();
            any = true;
        }

        while (PeekMessage(out Message message, IntPtr.Zero, 0, 0, PM_REMOVE))
        {
            any = true;

            if (message.hwnd == IntPtr.Zero && message.message == WakeMessage)
            {
                while (Posted.TryDequeue(out Action? posted))
                {
                    posted();
                }

                continue;
            }

            TranslateMessage(ref message);
            DispatchMessage(ref message);
        }

        return any;
    }

    private static void EnsureStarted()
    {
        lock (Gate)
        {
            if (_threadId != 0)
            {
                return;
            }

            var ready = new TaskCompletionSource<uint>(TaskCreationOptions.RunContinuationsAsynchronously);

            var thread = new Thread(() =>
            {
                // The queue has to exist before anyone is told the thread is up:
                // PostThreadMessage to a thread without one is silently dropped
                PeekMessage(out Message _, IntPtr.Zero, WM_USER, WM_USER, PM_NOREMOVE);
                ready.SetResult(GetCurrentThreadId());

                Loop();
            })
            {
                IsBackground = true,
                Name = "CursorLang.Tests pump",
            };

            thread.SetApartmentState(ApartmentState.STA);
            thread.Start();

            _threadId = ready.Task.GetAwaiter().GetResult();
        }
    }

    private static void Loop()
    {
        while (GetMessage(out Message message, IntPtr.Zero, 0, 0) > 0)
        {
            if (message.hwnd == IntPtr.Zero && message.message == WakeMessage)
            {
                while (Posted.TryDequeue(out Action? action))
                {
                    action();
                }

                continue;
            }

            TranslateMessage(ref message);
            DispatchMessage(ref message);
        }
    }

    private static void Wake() => PostThreadMessage(_threadId, WakeMessage, IntPtr.Zero, IntPtr.Zero);

    /// <summary>WM_APP and up belong to the application, and here that is the tests.</summary>
    private const uint WakeMessage = 0x8000 + 200;

    private const uint WM_USER = 0x0400;
    private const uint PM_REMOVE = 0x0001;
    private const uint PM_NOREMOVE = 0x0000;

    [DllImport("user32.dll", CharSet = CharSet.Unicode, EntryPoint = "GetMessageW")]
    private static extern int GetMessage(out Message lpMsg, IntPtr hWnd, uint filterMin, uint filterMax);

    [DllImport("user32.dll", CharSet = CharSet.Unicode, EntryPoint = "PeekMessageW")]
    private static extern bool PeekMessage(
        out Message lpMsg, IntPtr hWnd, uint filterMin, uint filterMax, uint removeMsg);

    [DllImport("user32.dll")]
    private static extern bool TranslateMessage(ref Message lpMsg);

    [DllImport("user32.dll", CharSet = CharSet.Unicode, EntryPoint = "DispatchMessageW")]
    private static extern IntPtr DispatchMessage(ref Message lpMsg);

    [DllImport("user32.dll", CharSet = CharSet.Unicode, EntryPoint = "PostThreadMessageW")]
    private static extern bool PostThreadMessage(uint threadId, uint message, IntPtr wParam, IntPtr lParam);

    [DllImport("kernel32.dll")]
    private static extern uint GetCurrentThreadId();

    [StructLayout(LayoutKind.Sequential)]
    private struct Message
    {
        public IntPtr hwnd;
        public uint message;
        public IntPtr wParam;
        public IntPtr lParam;
        public uint time;
        public int x;
        public int y;
        public uint lPrivate;
    }
}