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 Detour? _detour;
    private bool _installed;

    /// <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)
    {
        if (!_installed)
        {
            throw new InvalidOperationException(
                "The main-thread pump is not installed. Call Install() first.");
        }

        var tcs = new TaskCompletionSource<object?>();
        _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)
    {
        if (!_installed)
        {
            throw new InvalidOperationException(
                "The main-thread pump is not installed. Call Install() first.");
        }

        var tcs = new TaskCompletionSource<TResult>();
        object? Box() => work()!;
        _queue.Enqueue(new WorkItem(Box, r => tcs.SetResult((TResult)r!), ex => tcs.SetException(ex)));
        return tcs.Task;
    }

    /// <summary>Removes the frame-function detour if it is installed.</summary>
    public void Dispose()
    {
        if (_installed && _detour is not null)
        {
            _detour.Remove();
            _installed = false;
        }
    }

    private int PumpHook()
    {
        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;
    }

    [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.SetResult(r);
            _setException = ex => tcs.SetException(ex);
        }

        public Func<object?> Work { get; }

        public void SetResult(object? result)
        {
            _setResult?.Invoke(result);
        }

        public void SetException(Exception exception)
        {
            _setException?.Invoke(exception);
        }
    }
}