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.Text;
using System.Threading;
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 TaskScheduler 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 ConcurrentQueue();
private static readonly ConditionalWeakTable itemTable = new ConditionalWeakTable();
private class QueueItem : IEquatable, IEquatable
{
public bool HasTask;
private readonly WeakReference weakTask = null;
public Task Task => weakTask.TryGetTarget(out var task) ? task : null;
public QueueItem(Task task)
{
HasTask = true;
weakTask = new WeakReference(task);
}
private bool Equals(WeakReference task)
=> weakTask.TryGetTarget(out var t1) && task.TryGetTarget(out var t2) && t1.Equals(t2);
public bool Equals(Task other) => HasTask && weakTask.TryGetTarget(out var task) && other.Equals(task);
public bool Equals(QueueItem other) => other.HasTask == HasTask && Equals(other.weakTask);
}
///
/// 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
TryExecuteTask(task.Task);
}
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()
{
// this is only for debuggers which we can't use sooooo
throw new NotSupportedException();
}
///
/// 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);
}
///
/// 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
///
/// 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
}
}