From aca59e8d864fb3c15c7b48c5c0533c6a65fc50a7 Mon Sep 17 00:00:00 2001
From: Anairkoen Schno <aaron.cirr.com@gmail.com>
Date: Sat, 7 Dec 2019 21:58:00 -0600
Subject: [PATCH] Added a SingleThreadTaskScheduler

---
 IPA.Loader/IPA.Loader.csproj                  |   1 +
 .../Async/SingleThreadTaskScheduler.cs        | 175 ++++++++++++++++++
 2 files changed, 176 insertions(+)
 create mode 100644 IPA.Loader/Utilities/Async/SingleThreadTaskScheduler.cs

diff --git a/IPA.Loader/IPA.Loader.csproj b/IPA.Loader/IPA.Loader.csproj
index a73ed13e..9e0a7a0e 100644
--- a/IPA.Loader/IPA.Loader.csproj
+++ b/IPA.Loader/IPA.Loader.csproj
@@ -127,6 +127,7 @@
     <Compile Include="Config\ModPrefs.cs" />
     <Compile Include="JsonConverters\SemverRangeConverter.cs" />
     <Compile Include="JsonConverters\SemverVersionConverter.cs" />
+    <Compile Include="Utilities\Async\SingleThreadTaskScheduler.cs" />
     <Compile Include="Utilities\BeatSaber.cs" />
     <Compile Include="Utilities\AlmostVersion.cs" />
     <Compile Include="Utilities\CriticalSection.cs" />
diff --git a/IPA.Loader/Utilities/Async/SingleThreadTaskScheduler.cs b/IPA.Loader/Utilities/Async/SingleThreadTaskScheduler.cs
new file mode 100644
index 00000000..605e1dc6
--- /dev/null
+++ b/IPA.Loader/Utilities/Async/SingleThreadTaskScheduler.cs
@@ -0,0 +1,175 @@
+using System;
+using System.Collections.Generic;
+using System.Collections.Concurrent;
+using System.Linq;
+using System.Text;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace IPA.Utilities.Async
+{
+    /// <summary>
+    /// A single-threaded task scheduler that runs all of its tasks on the same thread.
+    /// </summary>
+    public class SingleThreadTaskScheduler : TaskScheduler, IDisposable
+    {
+        private readonly Thread runThread = new Thread(ExecuteTasksS);
+        private readonly BlockingCollection<Task> tasks = new BlockingCollection<Task>();
+        private readonly CancellationTokenSource exitTokenSource = new CancellationTokenSource();
+
+        /// <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>
+        public void StartThread()
+        {
+            ThrowIfDisposed();
+
+            runThread.Start();
+        }
+
+        /// <summary>
+        /// Terminates the runner thread, and waits for the currently running task to complete.
+        /// </summary>
+        /// <remarks>
+        /// 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>
+        public IEnumerable<Task> Exit()
+        {
+            ThrowIfDisposed();
+
+            tasks.CompleteAdding();
+            exitTokenSource.Cancel();
+            runThread.Join();
+
+            var retTasks = new List<Task>();
+            retTasks.AddRange(tasks);
+
+            Dispose(true);
+            return retTasks;
+        }
+
+        /// <summary>
+        /// Waits for the runner thread to complete all tasks in the queue, then exits.
+        /// </summary>
+        /// <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>
+        public void Join()
+        {
+            ThrowIfDisposed();
+
+            tasks.CompleteAdding();
+            runThread.Join();
+            Dispose(true);
+        }
+
+        /// <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();
+        }
+
+        /// <summary>
+        /// Queues a given <see cref="Task"/> to this scheduler. The <see cref="Task"/> <i>must></i> be
+        /// 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>
+        protected override void QueueTask(Task task)
+        {
+            ThrowIfDisposed();
+
+            tasks.Add(task);
+        }
+
+        /// <summary>
+        /// Rejects any attempts to execute a task inline.
+        /// </summary>
+        /// <remarks>
+        /// This task scheduler <i>always</i> runs its tasks on the thread that it manages, therefore it doesn't
+        /// make sense to run it inline.
+        /// </remarks>
+        /// <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>
+        protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued)
+        {
+            ThrowIfDisposed();
+
+            return false;
+        }
+
+        private void ThrowIfDisposed()
+        {
+            if (disposedValue)
+                throw new InvalidOperationException("Object already disposed");
+        }
+
+        private void ExecuteTasks()
+        {
+            ThrowIfDisposed();
+
+            var token = exitTokenSource.Token;
+
+            try
+            {
+                // while we are still accepting tasks, and we can pull out a task with an infinite wait duration
+                while (!tasks.IsCompleted && tasks.TryTake(out var task, -1, token))
+                {
+                    TryExecuteTask(task);
+                }
+            }
+            catch (OperationCanceledException)
+            {
+                // TryTake was cancelled, we'll just leave
+            }
+        }
+
+        private static void ExecuteTasksS(object param)
+        {
+            var self = param as SingleThreadTaskScheduler;
+            self.ExecuteTasks();
+        }
+
+        #region IDisposable Support
+        private bool disposedValue = false; // To detect redundant calls
+
+        /// <summary>
+        /// Disposes this object.
+        /// </summary>
+        /// <param name="disposing">whether or not to dispose managed objects</param>
+        protected virtual void Dispose(bool disposing)
+        {
+            if (!disposedValue)
+            {
+                if (disposing) 
+                {
+                    exitTokenSource.Dispose();
+                    tasks.Dispose();
+                }
+
+                disposedValue = true;
+            }
+        }
+
+        /// <summary>
+        /// Disposes this object. This puts the object into an unusable state.
+        /// </summary>
+        public void Dispose()
+        {
+            // Do not change this code. Put cleanup code in Dispose(bool disposing) above.
+            Dispose(true);
+        }
+        #endregion
+    }
+}