BepInEx/Runtimes/Unity/BepInEx.Unity.Mono/ThreadingHelper.cs

using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Linq;
using System.Reflection;
using System.Threading;
using BepInEx.Logging;
using UnityEngine;
using Logger = BepInEx.Logging.Logger;

namespace BepInEx.Unity.Mono;

/// <summary>
///     Provides methods for running code on other threads and synchronizing with the main thread.
/// </summary>
public sealed class ThreadingHelper : MonoBehaviour, ISynchronizeInvoke
{
    private readonly object _invokeLock = new();
    private Action _invokeList;
    private Thread _mainThread;

    /// <summary>
    ///     Current instance of the helper.
    /// </summary>
    public static ThreadingHelper Instance { get; private set; }

    /// <summary>
    ///     Gives methods for invoking delegates on the main unity thread, both synchronously and asynchronously.
    ///     Can be used in many built-in framework types, for example
    ///     <see cref="System.IO.FileSystemWatcher.SynchronizingObject" />
    ///     and <see cref="System.Timers.Timer.SynchronizingObject" /> to make their events fire on the main unity thread.
    /// </summary>
    public static ISynchronizeInvoke SynchronizingObject => Instance;

    private void Update()
    {
        // The CurrentThread can change between Awake and later methods, it's safest to get it here.
        if (_mainThread == null)
            _mainThread = Thread.CurrentThread;

        // Safe to do outside of lock because nothing can remove callbacks, at worst we execute with 1 frame delay
        if (_invokeList == null) return;

        Action toRun;
        lock (_invokeLock)
        {
            toRun = _invokeList;
            _invokeList = null;
        }

        // Need to execute outside of the lock in case the callback itself calls Invoke we could deadlock
        // The invocation would also block any threads that call Invoke
        foreach (var action in toRun.GetInvocationList().Cast<Action>())
            try
            {
                action();
            }
            catch (Exception ex)
            {
                LogInvocationException(ex);
            }
    }

    internal static void Initialize()
    {
        var go = new GameObject("BepInEx_ThreadingHelper") { hideFlags = HideFlags.HideAndDontSave };
        DontDestroyOnLoad(go);
        Instance = go.AddComponent<ThreadingHelper>();
    }

    /// <summary>
    ///     Queue the delegate to be invoked on the main unity thread. Use to synchronize your threads.
    /// </summary>
    public void StartSyncInvoke(Action action)
    {
        if (action == null) throw new ArgumentNullException(nameof(action));

        lock (_invokeLock)
        {
            _invokeList += action;
        }
    }

    /// <summary>
    ///     Queue the delegate to be invoked on a background thread. Use this to run slow tasks without affecting the game.
    ///     NOTE: Most of Unity API can not be accessed while running on another thread!
    /// </summary>
    /// <param name="action">
    ///     Task to be executed on another thread. Can optionally return an Action that will be executed on the main thread.
    ///     You can use this action to return results of your work safely. Return null if this is not needed.
    /// </param>
    public void StartAsyncInvoke(Func<Action> action)
    {
        void DoWork(object _)
        {
            try
            {
                var result = action();

                if (result != null)
                    StartSyncInvoke(result);
            }
            catch (Exception ex)
            {
                LogInvocationException(ex);
            }
        }

        if (!ThreadPool.QueueUserWorkItem(DoWork))
            throw new NotSupportedException("Failed to queue the action on ThreadPool");
    }

    private static void LogInvocationException(Exception ex)
    {
        Logger.Log(LogLevel.Error, ex);
        if (ex.InnerException != null) Logger.Log(LogLevel.Error, $"INNER: {ex.InnerException}");
    }

    #region ISynchronizeInvoke

    IAsyncResult ISynchronizeInvoke.BeginInvoke(Delegate method, object[] args)
    {
        var result = new InvokeResult();

        object Invoke()
        {
            try
            {
                return method.DynamicInvoke(args);
            }
            catch (Exception ex)
            {
                result.ExceptionThrown = true;
                return ex;
            }
        }

        if (!InvokeRequired)
            result.Finish(Invoke(), true);
        else
            StartSyncInvoke(() => result.Finish(Invoke(), false));

        return result;
    }

    object ISynchronizeInvoke.EndInvoke(IAsyncResult result)
    {
        var invokeResult = (InvokeResult) result;
        invokeResult.AsyncWaitHandle.WaitOne();

        if (invokeResult.ExceptionThrown)
            throw (Exception) invokeResult.AsyncState;
        return invokeResult.AsyncState;
    }

    object ISynchronizeInvoke.Invoke(Delegate method, object[] args)
    {
        var invokeResult = ((ISynchronizeInvoke) this).BeginInvoke(method, args);
        return ((ISynchronizeInvoke) this).EndInvoke(invokeResult);
    }

    /// <summary>
    ///     False if current code is executing on the main unity thread, otherwise True.
    ///     Warning: Will return true before the first frame finishes (i.e. inside plugin Awake and Start methods).
    /// </summary>
    /// <inheritdoc />
    public bool InvokeRequired => _mainThread == null || _mainThread != Thread.CurrentThread;

    private sealed class InvokeResult : IAsyncResult
    {
        internal bool ExceptionThrown;

        public InvokeResult()
        {
            AsyncWaitHandle = new EventWaitHandle(false, EventResetMode.ManualReset);
        }

        public bool IsCompleted { get; private set; }
        public WaitHandle AsyncWaitHandle { get; }
        public object AsyncState { get; private set; }
        public bool CompletedSynchronously { get; private set; }

        public void Finish(object result, bool completedSynchronously)
        {
            AsyncState = result;
            CompletedSynchronously = completedSynchronously;
            IsCompleted = true;
            ((EventWaitHandle) AsyncWaitHandle).Set();
        }
    }

    #endregion
}

/// <summary>
///     Convenience extensions for utilizing multiple threads and using the <see cref="ThreadingHelper" />.
/// </summary>
public static class ThreadingExtensions
{
    /// <inheritdoc cref="RunParallel{TIn,TOut}(IList{TIn},Func{TIn,TOut},int)" />
    public static IEnumerable<TOut> RunParallel<TIn, TOut>(this IEnumerable<TIn> data,
                                                           Func<TIn, TOut> work,
                                                           int workerCount = -1)
    {
        foreach (var result in RunParallel(data.ToList(), work))
            yield return result;
    }

    /// <summary>
    ///     Apply a function to a collection of data by spreading the work on multiple threads.
    ///     Outputs of the functions are returned to the current thread and yielded one by one.
    /// </summary>
    /// <typeparam name="TIn">Type of the input values.</typeparam>
    /// <typeparam name="TOut">Type of the output values.</typeparam>
    /// <param name="data">Input values for the work function.</param>
    /// <param name="work">Function to apply to the data on multiple threads at once.</param>
    /// <param name="workerCount">Number of worker threads. By default Environment.ProcessorCount is used.</param>
    /// <exception cref="TargetInvocationException">
    ///     An exception was thrown inside one of the threads, and the operation was
    ///     aborted.
    /// </exception>
    /// <exception cref="ArgumentException">Need at least 1 workerCount.</exception>
    public static IEnumerable<TOut> RunParallel<TIn, TOut>(this IList<TIn> data,
                                                           Func<TIn, TOut> work,
                                                           int workerCount = -1)
    {
        if (workerCount < 0)
            workerCount = Mathf.Max(2, Environment.ProcessorCount);
        else if (workerCount == 0)
            throw new ArgumentException("Need at least 1 worker", nameof(workerCount));

        var perThreadCount = Mathf.CeilToInt(data.Count / (float) workerCount);
        var doneCount = 0;

        var lockObj = new object();
        var are = new ManualResetEvent(false);
        IEnumerable<TOut> doneItems = null;
        Exception exceptionThrown = null;

        // Start threads to process the data
        for (var i = 0; i < workerCount; i++)
        {
            var first = i * perThreadCount;
            var last = Mathf.Min(first + perThreadCount, data.Count);
            ThreadPool.QueueUserWorkItem(
                                         _ =>
                                         {
                                             var results = new List<TOut>(perThreadCount);

                                             try
                                             {
                                                 for (var dataIndex = first; dataIndex < last; dataIndex++)
                                                 {
                                                     if (exceptionThrown != null) break;
                                                     results.Add(work(data[dataIndex]));
                                                 }
                                             }
                                             catch (Exception ex)
                                             {
                                                 exceptionThrown = ex;
                                             }

                                             lock (lockObj)
                                             {
                                                 doneItems = doneItems == null
                                                                 ? results
                                                                 : results.Concat(doneItems);
                                                 doneCount++;
                                                 are.Set();
                                             }
                                         });
        }

        // Main thread waits for results and returns them until all threads finish
        while (true)
        {
            are.WaitOne();

            IEnumerable<TOut> toOutput;
            bool isDone;
            lock (lockObj)
            {
                toOutput = doneItems;
                doneItems = null;
                isDone = doneCount == workerCount;
            }

            if (toOutput != null)
                foreach (var doneItem in toOutput)
                    yield return doneItem;

            if (isDone)
                break;
        }

        if (exceptionThrown != null)
            throw new TargetInvocationException("An exception was thrown inside one of the threads",
                                                exceptionThrown);
    }

    /// <summary>
    /// Apply a function to a collection of data by spreading the work on multiple threads.
    /// Lower overhead than RunParallel but it blocks the main thread until all work is completed or an exception has been thrown.
    /// </summary>
    /// <typeparam name="T">Type of the input values.</typeparam>
    /// <param name="data">Input values for the work function.</param>
    /// <param name="work">Function to apply to the data on multiple threads at once.</param>
    /// <param name="workerCount">Number of worker threads. By default Environment.ProcessorCount is used.</param>
    /// <exception cref="TargetInvocationException">An exception was thrown inside one of the threads, and the operation was aborted.</exception>
    /// <exception cref="ArgumentException">Need at least 1 workerCount.</exception>
    public static void ForEachParallel<T>(this IList<T> data, Action<T> work, int workerCount = -1)
    {
        if (workerCount < 0)
            workerCount = Mathf.Max(2, Environment.ProcessorCount);
        else if (workerCount == 0)
            throw new ArgumentException("Need at least 1 worker", nameof(workerCount));

        var currentIndex = data.Count;

        var are = new ManualResetEvent(false);
        var runningCount = workerCount;
        Exception exceptionThrown = null;

        void DoWork(object _)
        {
            try
            {
                while (true)
                {
                    if (exceptionThrown != null)
                        return;

                    var decrementedIndex = Interlocked.Decrement(ref currentIndex);
                    if (decrementedIndex < 0)
                        return;

                    work(data[decrementedIndex]);
                }
            }
            catch (Exception ex)
            {
                exceptionThrown = ex;
            }
            finally
            {
                var decCount = Interlocked.Decrement(ref runningCount);
                if (decCount <= 0)
                    are.Set();
            }
        }

        // Start threads to process the data
        for (var i = 0; i < workerCount - 1; i++)
            ThreadPool.QueueUserWorkItem(DoWork);

        DoWork(null);

        are.WaitOne();

        if (exceptionThrown != null)
            throw new TargetInvocationException("An exception was thrown inside one of the threads", exceptionThrown);
    }
}