whitemagic/WhiteMagic/Execution/MainThreadPump.cs

using System;
using System.Collections.Concurrent;
using System.Runtime.InteropServices;
using System.Threading.Tasks;
using WhiteMagic.Hooking;

namespace WhiteMagic.Execution;

/// <summary>
/// A crash-safe work queue drained on the target's own thread via a detour on a
/// per-frame function. Callers queue work and receive the result (or exception)
/// on their own thread through a completion handle.
/// </summary>
/// <remarks>
/// The pump assumes the frame function is parameterless and returns an <see cref="int"/>.
/// This matches common per-frame functions such as D3D9 <c>EndScene</c>.
/// </remarks>
public sealed class MainThreadPump : IDisposable
{
    private readonly DetourManager _detours;
    private readonly IntPtr _frameAddress;
    private readonly ConcurrentQueue<WorkItem> _queue = new();

    private readonly object _gate = new();
    private Detour? _detour;
    private bool _installed;
    private bool _disposed;

    /// <summary>
    /// Creates a pump that will hook the frame function at <paramref name="frameAddress"/>.
    /// </summary>
    public MainThreadPump(DetourManager detours, IntPtr frameAddress)
    {
        _detours = detours;
        _frameAddress = frameAddress;
    }

    /// <summary>Returns <see langword="true"/> after the frame hook has been applied.</summary>
    public bool IsInstalled => _installed;

    /// <summary>Installs the frame-function detour.</summary>
    public void Install()
    {
        if (_installed)
            return;

        _detour = _detours.Create("MainThreadPump", _frameAddress, (FrameDelegate)PumpHook);
        _detour.Apply();
        _installed = true;
    }

    /// <summary>
    /// Queues work to run on the hooked thread and blocks until it completes.
    /// </summary>
    public TResult Execute<TResult>(Func<TResult> work)
    {
        var tcs = new TaskCompletionSource<object?>();
        lock (_gate)
        {
            if (_disposed)
                ThrowDisposed();
            if (!_installed)
                throw new InvalidOperationException("The main-thread pump is not installed. Call Install() first.");

            _queue.Enqueue(new WorkItem(() => work()!, tcs));
        }

        object? result = tcs.Task.GetAwaiter().GetResult();
        return (TResult)result!;
    }

    /// <summary>
    /// Queues work to run on the hooked thread and returns a <see cref="Task{TResult}"/>.
    /// </summary>
    public Task<TResult> ExecuteAsync<TResult>(Func<TResult> work)
    {
        var tcs = new TaskCompletionSource<TResult>();
        lock (_gate)
        {
            if (_disposed)
                ThrowDisposed();
            if (!_installed)
                throw new InvalidOperationException("The main-thread pump is not installed. Call Install() first.");

            object? Box() => work()!;
            _queue.Enqueue(new WorkItem(Box, r => tcs.TrySetResult((TResult)r!), ex => tcs.TrySetException(ex)));
        }

        return tcs.Task;
    }

    /// <summary>Removes the frame-function detour if it is installed.</summary>
    public void Dispose()
    {
        lock (_gate)
        {
            if (_disposed)
                return;
            _disposed = true;

            if (_installed && _detour is not null)
                _detour.Remove();
            _installed = false;
        }

        // Fault any caller still blocked on queued work so Execute cannot hang forever.
        while (_queue.TryDequeue(out WorkItem? item))
        {
            item.SetException(new ObjectDisposedException(nameof(MainThreadPump)));
        }
    }

    private static void ThrowDisposed()
        => throw new ObjectDisposedException(nameof(MainThreadPump));

    private int PumpHook()
    {
        try
        {
            while (_queue.TryDequeue(out WorkItem? item))
            {
                try
                {
                    object? result = item.Work();
                    item.SetResult(result);
                }
                catch (Exception ex)
                {
                    item.SetException(ex);
                }
            }

            // Call the original frame function so rendering/game logic continues.
            return _detour is null ? 0 : (int?)_detour.CallOriginal() ?? 0;
        }
        catch
        {
            // Never let an exception escape back into the native frame caller.
            return 0;
        }
    }

    [UnmanagedFunctionPointer(CallingConvention.Cdecl)]
    private delegate int FrameDelegate();

    private sealed class WorkItem
    {
        private readonly Action<object?>? _setResult;
        private readonly Action<Exception>? _setException;

        public WorkItem(Func<object?> work, Action<object?> setResult, Action<Exception> setException)
        {
            Work = work;
            _setResult = setResult;
            _setException = setException;
        }

        public WorkItem(Func<object?> work, TaskCompletionSource<object?> tcs)
        {
            Work = work;
            _setResult = r => tcs.TrySetResult(r);
            _setException = ex => tcs.TrySetException(ex);
        }

        public Func<object?> Work { get; }

        public void SetResult(object? result)
        {
            try
            {
                _setResult?.Invoke(result);
            }
            catch (InvalidOperationException)
            {
                // Already completed, e.g. concurrent Dispose/PumpHook race.
            }
        }

        public void SetException(Exception exception)
        {
            try
            {
                _setException?.Invoke(exception);
            }
            catch (InvalidOperationException)
            {
                // Already completed, e.g. concurrent Dispose/PumpHook race.
            }
        }
    }
}