#nullable enable using System; using System.Collections; using System.Collections.Concurrent; using System.Collections.Generic; using System.Diagnostics; using System.Linq; using System.Runtime.CompilerServices; using System.Threading.Tasks; namespace IPA.Utilities.Async { /// /// A task scheduler that runs tasks on the Unity main thread via coroutines. /// public class UnityMainThreadTaskScheduler : TaskScheduler, IDisposable { /// /// Gets the default main thread scheduler that is managed by BSIPA. /// /// a scheduler that is managed by BSIPA public static new UnityMainThreadTaskScheduler Default { get; } = new UnityMainThreadTaskScheduler(); /// /// Gets a factory for creating tasks on . /// /// a factory for creating tasks on the default scheduler public static TaskFactory Factory { get; } = new TaskFactory(Default); private readonly ConcurrentQueue tasks = new(); private static readonly ConditionalWeakTable itemTable = new(); private class QueueItem : IEquatable, IEquatable { private bool hasTask; public bool HasTask { get => hasTask; set { hasTask = value; if (!hasTask) Task = null; } } public Task? Task { get; private set; } public Action? Action { get; private set; } public QueueItem(Task task) { HasTask = true; Task = task; } public QueueItem(Action action) { HasTask = true; Action = action; } public bool Equals(Task? other) => other is not null && HasTask && other.Equals(Task); public bool Equals(QueueItem other) => other.HasTask == HasTask && Equals(other.Task); } /// /// Gets whether or not this scheduler is currently executing tasks. /// /// if the scheduler is running, otherwise public bool IsRunning { get; private set; } = false; /// /// Gets whether or not this scheduler is in the process of shutting down. /// /// if the scheduler is shutting down, otherwise public bool Cancelling { get; private set; } = false; private int yieldAfterTasks = 64; /// /// Gets or sets the number of tasks to execute before yielding back to Unity. /// /// the number of tasks to execute per resume public int YieldAfterTasks { get => yieldAfterTasks; set { ThrowIfDisposed(); if (value < 1) throw new ArgumentException("Value cannot be less than 1", nameof(value)); yieldAfterTasks = value; } } private TimeSpan yieldAfterTime = TimeSpan.FromMilliseconds(.5); // auto-yield if more than half a millis has passed by default /// /// Gets or sets the amount of time to execute tasks for before yielding back to Unity. Default is 0.5ms. /// /// the amount of time to execute tasks for before yielding back to Unity public TimeSpan YieldAfterTime { get => yieldAfterTime; set { ThrowIfDisposed(); if (value <= TimeSpan.Zero) throw new ArgumentException("Value must be greater than zero", nameof(value)); yieldAfterTime = value; } } /// /// When used as a Unity coroutine, runs the scheduler. Otherwise, this is an invalid call. /// /// /// /// Do not ever call on this /// coroutine, nor on the behaviour hosting /// this coroutine. This has no way to detect this, and this object will become invalid. /// /// /// If you need to stop this coroutine, first call , then wait for it to /// exit on its own. /// /// /// a Unity coroutine /// if this scheduler is disposed /// if the scheduler is already running public IEnumerator Coroutine() { ThrowIfDisposed(); if (IsRunning) throw new InvalidOperationException("Scheduler already running"); Cancelling = false; IsRunning = true; yield return null; // yield immediately var sw = new Stopwatch(); try { while (!Cancelling) { if (!tasks.IsEmpty) { var yieldAfter = YieldAfterTasks; sw.Start(); for (int i = 0; i < yieldAfter && !tasks.IsEmpty && sw.Elapsed < YieldAfterTime; i++) { QueueItem task; do if (!tasks.TryDequeue(out task)) goto exit; // try dequeue, if we can't exit while (!task.HasTask); // if the dequeued task is empty, try again if (task.Task is not null) { _ = TryExecuteTask(task.Task); } task.Action?.Invoke(); } exit: sw.Reset(); } yield return null; } } finally { sw.Reset(); IsRunning = false; } } /// /// Cancels the scheduler. If the scheduler is currently executing tasks, that batch will finish first. /// All remaining tasks will be left in the queue. /// /// if this scheduler is disposed /// if the scheduler is not running public void Cancel() { ThrowIfDisposed(); if (!IsRunning) throw new InvalidOperationException("The scheduler is not running"); Cancelling = true; } /// /// Throws a . /// /// nothing /// Always. protected override IEnumerable GetScheduledTasks() => tasks.ToArray().Where(q => q.HasTask).Select(q => q.Task).NonNull().ToArray(); /// /// Queues a given to this scheduler. The must be /// scheduled for this by the runtime. /// /// the to queue /// Thrown if this object has already been disposed. protected override void QueueTask(Task task) { ThrowIfDisposed(); var item = new QueueItem(task); itemTable.Add(task, item); tasks.Enqueue(item); } internal void QueueAction(Action action) { ThrowIfDisposed(); tasks.Enqueue(new(action)); } /// /// Runs the task inline if the current thread is the Unity main thread. /// /// the task to attempt to execute /// whether the task was previously queued to this scheduler /// if the task could not be run, if it was /// Thrown if this object has already been disposed. protected override bool TryExecuteTaskInline(Task task, bool taskWasPreviouslyQueued) { ThrowIfDisposed(); if (!UnityGame.OnMainThread) return false; if (taskWasPreviouslyQueued) { if (itemTable.TryGetValue(task, out var item)) { if (!item.HasTask) return false; item.HasTask = false; } else 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 /// /// Disposes this object. /// /// whether or not to dispose managed objects protected virtual void Dispose(bool disposing) { if (!disposedValue) { if (disposing) { } disposedValue = true; } } /// /// Disposes this object. This puts the object into an unusable state. /// // 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 } }