From 50fea67b6884ab9ec5cd0259e30d4a1e71efcf07 Mon Sep 17 00:00:00 2001 From: Anairkoen Schno Date: Thu, 16 Jan 2020 23:07:01 -0600 Subject: [PATCH] Added MainThread checks and UnityMainThreadTaskScheduler --- IPA.Injector/Injector.cs | 2 - IPA.Loader/IPA.Loader.csproj | 1 + IPA.Loader/Loader/PluginComponent.cs | 17 ++- .../Loader/StateTransitionTransaction.cs | 3 +- .../Async/SingleThreadTaskScheduler.cs | 14 +- .../Async/UnityMainThreadTaskScheduler.cs | 132 ++++++++++++++++++ IPA.Loader/Utilities/UnityGame.cs | 11 ++ 7 files changed, 167 insertions(+), 13 deletions(-) create mode 100644 IPA.Loader/Utilities/Async/UnityMainThreadTaskScheduler.cs diff --git a/IPA.Injector/Injector.cs b/IPA.Injector/Injector.cs index 4ef3b09d..6967e81f 100644 --- a/IPA.Injector/Injector.cs +++ b/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(); diff --git a/IPA.Loader/IPA.Loader.csproj b/IPA.Loader/IPA.Loader.csproj index 0b37ea80..d67fc8e4 100644 --- a/IPA.Loader/IPA.Loader.csproj +++ b/IPA.Loader/IPA.Loader.csproj @@ -140,6 +140,7 @@ + diff --git a/IPA.Loader/Loader/PluginComponent.cs b/IPA.Loader/Loader/PluginComponent.cs index 76eccfd2..3bfaa349 100644 --- a/IPA.Loader/Loader/PluginComponent.cs +++ b/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() diff --git a/IPA.Loader/Loader/StateTransitionTransaction.cs b/IPA.Loader/Loader/StateTransitionTransaction.cs index 5452cfcd..948bcd4e 100644 --- a/IPA.Loader/Loader/StateTransitionTransaction.cs +++ b/IPA.Loader/Loader/StateTransitionTransaction.cs @@ -203,11 +203,10 @@ namespace IPA.Loader /// /// /// After this completes, this transaction will be disposed. - /// It is illegal to call this anywhere but the main Unity thread. /// /// The that is returned will error if any of the mods being disabled /// error. It is up to the caller to handle these in a sane way, like logging them. If nothing else, do something like this: - /// + /// /// // get your transaction... /// var complete = transaction.Commit(); /// complete.ContinueWith(t => { diff --git a/IPA.Loader/Utilities/Async/SingleThreadTaskScheduler.cs b/IPA.Loader/Utilities/Async/SingleThreadTaskScheduler.cs index 85cced15..f092ac8a 100644 --- a/IPA.Loader/Utilities/Async/SingleThreadTaskScheduler.cs +++ b/IPA.Loader/Utilities/Async/SingleThreadTaskScheduler.cs @@ -20,7 +20,7 @@ namespace IPA.Utilities.Async /// /// Gets whether or not the underlying thread has been started. /// - /// Thrown if this object has already been disposed. + /// Thrown if this object has already been disposed. public bool IsRunning { get @@ -33,7 +33,7 @@ namespace IPA.Utilities.Async /// /// Starts the thread that executes tasks scheduled with this /// - /// Thrown if this object has already been disposed. + /// Thrown if this object has already been disposed. 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. /// /// an of s that did not execute - /// Thrown if this object has already been disposed. + /// Thrown if this object has already been disposed. public IEnumerable Exit() { ThrowIfDisposed(); @@ -69,7 +69,7 @@ namespace IPA.Utilities.Async /// /// After this method returns, this object has been disposed and is no longer in a valid state. /// - /// Thrown if this object has already been disposed. + /// Thrown if this object has already been disposed. public void Join() { ThrowIfDisposed(); @@ -94,7 +94,7 @@ namespace IPA.Utilities.Async /// scheduled for this by the runtime. /// /// the to queue - /// Thrown if this object has already been disposed. + /// Thrown if this object has already been disposed. protected override void QueueTask(Task task) { ThrowIfDisposed(); @@ -112,7 +112,7 @@ namespace IPA.Utilities.Async /// the task to attempt to execute /// whether the task was previously queued to this scheduler /// - /// Thrown if this object has already been disposed. + /// Thrown if this object has already been disposed. 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() diff --git a/IPA.Loader/Utilities/Async/UnityMainThreadTaskScheduler.cs b/IPA.Loader/Utilities/Async/UnityMainThreadTaskScheduler.cs new file mode 100644 index 00000000..f0b97184 --- /dev/null +++ b/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 tasks = new ConcurrentDictionary(); + private int queueEndPosition = 0; + private int queuePosition = 0; + + private struct QueueItem : IEquatable, IEquatable, IEquatable + { + 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; + } + } + + /// + /// Throws a . + /// + /// nothing + /// Always. + protected override IEnumerable 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 + } +} diff --git a/IPA.Loader/Utilities/UnityGame.cs b/IPA.Loader/Utilities/UnityGame.cs index cdd26f31..4b1409d3 100644 --- a/IPA.Loader/Utilities/UnityGame.cs +++ b/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; + /// + /// Checks if the currently running code is running on the Unity main thread. + /// + /// if the curent thread is the Unity main thread, otherwise + public static bool OnMainThread => Thread.CurrentThread.ManagedThreadId == mainThread?.ManagedThreadId; + + internal static void SetMainThread() + => mainThread = Thread.CurrentThread; + /// /// The different types of releases of the game. ///