Added MainThread checks and UnityMainThreadTaskScheduler

IPA.Injector/Injector.cs

@@ -331,8 +331,6 @@ namespace IPA.Injector
             pluginAsyncLoadTask.Wait();
             permissionFixTask.Wait();
 
-            UnityGame.EnsureRuntimeGameVersion();
-
             log.Debug("Plugins loaded");
             log.Debug(string.Join(", ", PluginLoader.PluginsMetadata.StrJP()));
             PluginComponent.Create();

IPA.Loader/IPA.Loader.csproj

@@ -140,6 +140,7 @@
     <Compile Include="JsonConverters\SemverVersionConverter.cs" />
     <Compile Include="Utilities\Async\SingleThreadTaskScheduler.cs" />
     <Compile Include="Utilities\Accessor.cs" />
+    <Compile Include="Utilities\Async\UnityMainThreadTaskScheduler.cs" />
     <Compile Include="Utilities\UnityGame.cs" />
     <Compile Include="Utilities\AlmostVersion.cs" />
     <Compile Include="Utilities\CriticalSection.cs" />

IPA.Loader/Loader/PluginComponent.cs

@@ -1,6 +1,8 @@
 using IPA.Config;
 using IPA.Loader.Composite;
-using System.Diagnostics.CodeAnalysis;
+using IPA.Utilities;
+using IPA.Utilities.Async;
+using System.Diagnostics.CodeAnalysis;
 using UnityEngine;
 using UnityEngine.SceneManagement;
 // ReSharper disable UnusedMember.Local
@@ -26,6 +28,9 @@ namespace IPA.Loader
 
             if (!initialized)
             {
+                UnityGame.SetMainThread();
+                UnityGame.EnsureRuntimeGameVersion();
+
                 PluginManager.Load();
 
                 bsPlugins = new CompositeBSPlugin(PluginManager.BSMetas);
@@ -42,7 +47,11 @@ namespace IPA.Loader
 
                 SceneManager.activeSceneChanged += OnActiveSceneChanged;
                 SceneManager.sceneLoaded += OnSceneLoaded;
-                SceneManager.sceneUnloaded += OnSceneUnloaded;
+                SceneManager.sceneUnloaded += OnSceneUnloaded;
+
+                var unitySched = UnityMainThreadTaskScheduler.Default as UnityMainThreadTaskScheduler;
+                if (!unitySched.IsRunning)
+                    StartCoroutine(unitySched.Coroutine());
 
                 initialized = true;
             }
@@ -52,6 +61,10 @@ namespace IPA.Loader
         {
             bsPlugins.OnUpdate();
             ipaPlugins.OnUpdate();
+
+            var unitySched = UnityMainThreadTaskScheduler.Default as UnityMainThreadTaskScheduler;
+            if (!unitySched.IsRunning)
+                StartCoroutine(unitySched.Coroutine());
         }
 
         void LateUpdate()

IPA.Loader/Loader/StateTransitionTransaction.cs

@@ -203,11 +203,10 @@ namespace IPA.Loader
         /// </summary>
         /// <remarks>
         /// <para>After this completes, this transaction will be disposed.</para>
-        /// <para>It is illegal to call this anywhere but the main Unity thread.</para>
         /// <para>
         /// The <see cref="Task"/> that is returned will error if <b>any</b> of the mods being <b>disabled</b>
         /// error. It is up to the caller to handle these in a sane way, like logging them. If nothing else, do something like this:
-        /// <code>
+        /// <code lang="csharp">
         /// // get your transaction...
         /// var complete = transaction.Commit();
         /// complete.ContinueWith(t => {

IPA.Loader/Utilities/Async/SingleThreadTaskScheduler.cs

@@ -20,7 +20,7 @@ namespace IPA.Utilities.Async
         /// <summary>
         /// Gets whether or not the underlying thread has been started.
         /// </summary>
-        /// <exception cref="InvalidOperationException">Thrown if this object has already been disposed.</exception>
+        /// <exception cref="ObjectDisposedException">Thrown if this object has already been disposed.</exception>
         public bool IsRunning
         {
             get
@@ -33,7 +33,7 @@ namespace IPA.Utilities.Async
         /// <summary>
         /// Starts the thread that executes tasks scheduled with this <see cref="TaskScheduler"/>
         /// </summary>
-        /// <exception cref="InvalidOperationException">Thrown if this object has already been disposed.</exception>
+        /// <exception cref="ObjectDisposedException">Thrown if this object has already been disposed.</exception>
         public void Start()
         {
             ThrowIfDisposed();
@@ -48,7 +48,7 @@ namespace IPA.Utilities.Async
         /// After this method returns, this object has been disposed and is no longer in a valid state.
         /// </remarks>
         /// <returns>an <see cref="IEnumerable{T}"/> of <see cref="Task"/>s that did not execute</returns>
-        /// <exception cref="InvalidOperationException">Thrown if this object has already been disposed.</exception>
+        /// <exception cref="ObjectDisposedException">Thrown if this object has already been disposed.</exception>
         public IEnumerable<Task> Exit()
         {
             ThrowIfDisposed();
@@ -69,7 +69,7 @@ namespace IPA.Utilities.Async
         /// <remarks>
         /// After this method returns, this object has been disposed and is no longer in a valid state.
         /// </remarks>
-        /// <exception cref="InvalidOperationException">Thrown if this object has already been disposed.</exception>
+        /// <exception cref="ObjectDisposedException">Thrown if this object has already been disposed.</exception>
         public void Join()
         {
             ThrowIfDisposed();
@@ -94,7 +94,7 @@ namespace IPA.Utilities.Async
         /// scheduled for this <see cref="TaskScheduler"/> by the runtime.
         /// </summary>
         /// <param name="task">the <see cref="Task"/> to queue</param>
-        /// <exception cref="InvalidOperationException">Thrown if this object has already been disposed.</exception>
+        /// <exception cref="ObjectDisposedException">Thrown if this object has already been disposed.</exception>
         protected override void QueueTask(Task task)
         {
             ThrowIfDisposed();
@@ -112,7 +112,7 @@ namespace IPA.Utilities.Async
         /// <param name="task">the task to attempt to execute</param>
         /// <param name="taskWasPreviouslyQueued">whether the task was previously queued to this scheduler</param>
         /// <returns><see langword="false"/></returns>
-        /// <exception cref="InvalidOperationException">Thrown if this object has already been disposed.</exception>
+        /// <exception cref="ObjectDisposedException">Thrown if this object has already been disposed.</exception>
         protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued)
         {
             ThrowIfDisposed();
@@ -123,7 +123,7 @@ namespace IPA.Utilities.Async
         private void ThrowIfDisposed()
         {
             if (disposedValue)
-                throw new InvalidOperationException("Object already disposed");
+                throw new ObjectDisposedException(nameof(SingleThreadTaskScheduler));
         }
 
         private void ExecuteTasks()

IPA.Loader/Utilities/Async/UnityMainThreadTaskScheduler.cs

@@ -0,0 +1,132 @@
+using System;
+using System.Collections;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace IPA.Utilities.Async
+{
+    public class UnityMainThreadTaskScheduler : TaskScheduler, IDisposable
+    {
+        public static new TaskScheduler Default { get; } = new UnityMainThreadTaskScheduler();
+
+        private readonly ConcurrentDictionary<QueueItem, Task> tasks = new ConcurrentDictionary<QueueItem, Task>();
+        private int queueEndPosition = 0;
+        private int queuePosition = 0;
+
+        private struct QueueItem : IEquatable<int>, IEquatable<Task>, IEquatable<QueueItem>
+        {
+            public int Index;
+            public Task Task;
+
+            public QueueItem(int index, Task task) : this()
+            {
+                Index = index;
+                Task = task;
+            }
+
+            public bool Equals(int other) => Index.Equals(other);
+            public bool Equals(Task other) => Task.Equals(other);
+            public bool Equals(QueueItem other) => other.Index == Index || other.Task == Task;
+        }
+
+        public bool IsRunning { get; private set; } = false;
+
+        public int YieldAfterTasks { get; set; } = 4;
+
+        public IEnumerator Coroutine()
+        {
+            ThrowIfDisposed();
+
+            IsRunning = true;
+            yield return null; // yield immediately
+
+            try
+            {
+                while (true)
+                {
+                    if (queuePosition < queueEndPosition)
+                    {
+                        var yieldAfter = YieldAfterTasks;
+                        for (int i = 0; i < yieldAfter && queuePosition < queueEndPosition; i++)
+                        {
+                            if (tasks.TryRemove(new QueueItem { Index = Interlocked.Increment(ref queuePosition) }, out var task))
+                                TryExecuteTask(task); // we succesfully removed the task
+                            else
+                                i++; // we didn't
+                        }
+                    }
+                    yield return null;
+                }
+            }
+            finally
+            {
+                IsRunning = false;
+            }
+        }
+
+        /// <summary>
+        /// Throws a <see cref="NotSupportedException"/>.
+        /// </summary>
+        /// <returns>nothing</returns>
+        /// <exception cref="NotSupportedException">Always.</exception>
+        protected override IEnumerable<Task> GetScheduledTasks()
+        {
+            // this is only for debuggers which we can't use sooooo
+            throw new NotSupportedException();
+        }
+
+        protected override void QueueTask(Task task)
+        {
+            ThrowIfDisposed();
+
+            tasks.TryAdd(new QueueItem(Interlocked.Increment(ref queueEndPosition), task), task);
+        }
+
+        protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued)
+        {
+            ThrowIfDisposed();
+
+            if (!UnityGame.OnMainThread) return false;
+
+            if (taskWasPreviouslyQueued)
+                if (!tasks.TryRemove(new QueueItem { Task = task }, out var _))
+                    return false; // if we couldn't remove it, its not in our queue, so it already ran
+
+            return TryExecuteTask(task);
+        }
+
+        private void ThrowIfDisposed()
+        {
+            if (disposedValue)
+                throw new ObjectDisposedException(nameof(SingleThreadTaskScheduler));
+        }
+
+        #region IDisposable Support
+        private bool disposedValue = false; // To detect redundant calls
+
+        protected virtual void Dispose(bool disposing)
+        {
+            if (!disposedValue)
+            {
+                if (disposing)
+                {
+                    // TODO: dispose managed state (managed objects).
+                }
+
+                disposedValue = true;
+            }
+        }
+
+        // This code added to correctly implement the disposable pattern.
+        public void Dispose()
+        {
+            // Do not change this code. Put cleanup code in Dispose(bool disposing) above.
+            Dispose(true);
+        }
+        #endregion
+    }
+}

IPA.Loader/Utilities/UnityGame.cs

@@ -2,6 +2,7 @@
 using System;
 using System.IO;
 using System.Reflection;
+using System.Threading;
 using UnityEngine;
 #if NET3
 using Path = Net3_Proxy.Path;
@@ -59,6 +60,16 @@ namespace IPA.Utilities
             SelfConfig.Instance.LastGameVersion = gameVer.ToString();
         }
 
+        private static Thread mainThread;
+        /// <summary>
+        /// Checks if the currently running code is running on the Unity main thread.
+        /// </summary>
+        /// <value><see langword="true"/> if the curent thread is the Unity main thread, <see langword="false"/> otherwise</value>
+        public static bool OnMainThread => Thread.CurrentThread.ManagedThreadId == mainThread?.ManagedThreadId;
+
+        internal static void SetMainThread() 
+            => mainThread = Thread.CurrentThread;
+
         /// <summary>
         /// The different types of releases of the game.
         /// </summary>