Files
ShrinkEventBus/Runtime/ShrinkBusSchedulers.cs

555 lines
18 KiB
C#

#nullable enable
using System;
using System.Collections.Concurrent;
using System.Diagnostics;
using System.Threading;
using System.Threading.Tasks;
using Cysharp.Threading.Tasks;
using ShrinkSDK.Runtime;
namespace ShrinkEventBus
{
internal interface IShrinkBusScheduler : IDisposable
{
bool IsOnSchedulerThread { get; }
bool TryPost(Action action);
UniTask PostAsync(Func<UniTask> action, CancellationToken cancellationToken);
UniTask ShutdownAsync(bool drain, TimeSpan timeout);
}
internal sealed class ShrinkQueueFullException : InvalidOperationException
{
public ShrinkQueueFullException(string schedulerName)
: base($"Scheduler queue '{schedulerName}' is full.")
{
}
}
internal static class ShrinkBusSchedulerFactory
{
public static IShrinkBusScheduler Create(ShrinkBusOptions options, string name)
{
return options.Scheduler switch
{
ShrinkBusSchedulerKind.MainThread => new ShrinkMainThreadScheduler(name, options),
ShrinkBusSchedulerKind.DedicatedThread => new ShrinkDedicatedThreadScheduler(name, options),
ShrinkBusSchedulerKind.TaskPool => new ShrinkTaskPoolScheduler(name, options),
_ => new ShrinkInlineScheduler()
};
}
}
internal sealed class ShrinkSchedulerWorkItem
{
private readonly UniTaskCompletionSource? _completion;
public ShrinkSchedulerWorkItem(Func<UniTask> action, bool awaitable)
{
Action = action ?? throw new ArgumentNullException(nameof(action));
if (awaitable)
_completion = new UniTaskCompletionSource();
}
public Func<UniTask> Action { get; }
public UniTask Completion => _completion?.Task ?? UniTask.CompletedTask;
public void Complete() => _completion?.TrySetResult();
public void Fail(Exception exception) => _completion?.TrySetException(exception);
public void Cancel(CancellationToken cancellationToken) => _completion?.TrySetCanceled(cancellationToken);
public void Drop(string schedulerName) =>
_completion?.TrySetException(new ShrinkQueueFullException(schedulerName));
}
internal sealed class ShrinkSchedulerQueue : IDisposable
{
private readonly ConcurrentQueue<ShrinkSchedulerWorkItem> _queue = new();
private readonly SemaphoreSlim _slots;
private readonly ShrinkQueueOverflowPolicy _overflowPolicy;
private readonly string _name;
private int _disposed;
public ShrinkSchedulerQueue(string name, ShrinkBusOptions options)
{
_name = name;
_overflowPolicy = options.OverflowPolicy;
_slots = new SemaphoreSlim(options.QueueCapacity, options.QueueCapacity);
}
public bool IsEmpty => _queue.IsEmpty;
public bool TryEnqueue(ShrinkSchedulerWorkItem item)
{
if (Volatile.Read(ref _disposed) != 0)
return false;
if (!_slots.Wait(0))
{
if (_overflowPolicy != ShrinkQueueOverflowPolicy.DropOldest ||
!_queue.TryDequeue(out var dropped))
return false;
dropped.Drop(_name);
_slots.Release();
if (!_slots.Wait(0))
return false;
}
_queue.Enqueue(item);
return true;
}
public async UniTask EnqueueAsync(ShrinkSchedulerWorkItem item, CancellationToken cancellationToken)
{
if (Volatile.Read(ref _disposed) != 0)
throw new ObjectDisposedException(_name);
if (_overflowPolicy == ShrinkQueueOverflowPolicy.Wait)
{
await _slots.WaitAsync(cancellationToken).AsUniTask(useCurrentSynchronizationContext: false);
if (Volatile.Read(ref _disposed) != 0)
{
_slots.Release();
throw new ObjectDisposedException(_name);
}
_queue.Enqueue(item);
return;
}
if (!TryEnqueue(item))
throw new ShrinkQueueFullException(_name);
}
public bool TryDequeue(out ShrinkSchedulerWorkItem item)
{
if (!_queue.TryDequeue(out item!))
return false;
_slots.Release();
return true;
}
public void DropPending()
{
while (_queue.TryDequeue(out var item))
{
_slots.Release();
item.Drop(_name);
}
}
public void Dispose()
{
if (Interlocked.Exchange(ref _disposed, 1) != 0)
return;
DropPending();
_slots.Dispose();
}
}
internal static class ShrinkSchedulerWorkItemRunner
{
public static async UniTask RunAsync(ShrinkSchedulerWorkItem item)
{
try
{
await item.Action();
item.Complete();
}
catch (OperationCanceledException ex)
{
item.Cancel(ex.CancellationToken);
}
catch (Exception ex)
{
item.Fail(ex);
ShrinkEventDiagnostics.LogException(ex);
}
}
public static void RunBlocking(ShrinkSchedulerWorkItem item)
{
try
{
item.Action().GetAwaiter().GetResult();
item.Complete();
}
catch (OperationCanceledException ex)
{
item.Cancel(ex.CancellationToken);
}
catch (Exception ex)
{
item.Fail(ex);
ShrinkEventDiagnostics.LogException(ex);
}
}
}
internal sealed class ShrinkInlineScheduler : IShrinkBusScheduler
{
public bool IsOnSchedulerThread => true;
public bool TryPost(Action action)
{
action();
return true;
}
public UniTask PostAsync(Func<UniTask> action, CancellationToken cancellationToken)
{
cancellationToken.ThrowIfCancellationRequested();
return action();
}
public UniTask ShutdownAsync(bool drain, TimeSpan timeout) => UniTask.CompletedTask;
public void Dispose() { }
}
internal sealed class ShrinkMainThreadScheduler : IShrinkBusScheduler
{
private readonly ShrinkSchedulerQueue _queue;
private readonly IShrinkMainThreadDispatcher _dispatcher;
private int _pumpScheduled;
private int _disposed;
public ShrinkMainThreadScheduler(string name, ShrinkBusOptions options)
{
_queue = new ShrinkSchedulerQueue(name, options);
_dispatcher = ShrinkEventBusRuntime.MainThreadDispatcher;
}
public bool IsOnSchedulerThread => _dispatcher.IsMainThread;
public bool TryPost(Action action)
{
if (Volatile.Read(ref _disposed) != 0)
return false;
if (IsOnSchedulerThread)
{
action();
return true;
}
var item = new ShrinkSchedulerWorkItem(() =>
{
action();
return UniTask.CompletedTask;
}, awaitable: false);
if (!_queue.TryEnqueue(item))
return false;
SchedulePump();
return true;
}
public async UniTask PostAsync(Func<UniTask> action, CancellationToken cancellationToken)
{
if (Volatile.Read(ref _disposed) != 0)
throw new ObjectDisposedException(nameof(ShrinkMainThreadScheduler));
cancellationToken.ThrowIfCancellationRequested();
if (IsOnSchedulerThread)
{
await action();
return;
}
var item = new ShrinkSchedulerWorkItem(action, awaitable: true);
await _queue.EnqueueAsync(item, cancellationToken);
SchedulePump();
await item.Completion.AttachExternalCancellation(cancellationToken);
}
public async UniTask ShutdownAsync(bool drain, TimeSpan timeout)
{
Interlocked.Exchange(ref _disposed, 1);
if (!drain)
_queue.DropPending();
else
await ShrinkSchedulerShutdown.WaitUntilDrainedAsync(
_queue, () => Volatile.Read(ref _pumpScheduled) != 0, timeout);
_queue.Dispose();
}
public void Dispose()
{
if (Interlocked.Exchange(ref _disposed, 1) != 0)
return;
_queue.Dispose();
}
private void SchedulePump()
{
if (Interlocked.Exchange(ref _pumpScheduled, 1) != 0)
return;
if (_dispatcher.TryPost(() => PumpAsync().Forget()))
return;
Interlocked.Exchange(ref _pumpScheduled, 0);
ShrinkEventDiagnostics.LogException(new InvalidOperationException(
"The configured main-thread dispatcher rejected an EventBus pump."));
}
private async UniTask PumpAsync()
{
try
{
while (_queue.TryDequeue(out var item))
await ShrinkSchedulerWorkItemRunner.RunAsync(item);
}
finally
{
Interlocked.Exchange(ref _pumpScheduled, 0);
if (!_queue.IsEmpty && Volatile.Read(ref _disposed) == 0)
SchedulePump();
}
}
}
internal sealed class ShrinkTaskPoolScheduler : IShrinkBusScheduler
{
private readonly ShrinkSchedulerQueue _queue;
private readonly int _maxConcurrency;
private int _workers;
private int _disposed;
public ShrinkTaskPoolScheduler(string name, ShrinkBusOptions options)
{
_queue = new ShrinkSchedulerQueue(name, options);
_maxConcurrency = options.DispatchMode == ShrinkDispatchMode.Ordered ? 1 : options.MaxConcurrency;
}
public bool IsOnSchedulerThread => false;
public bool TryPost(Action action)
{
if (Volatile.Read(ref _disposed) != 0)
return false;
var item = new ShrinkSchedulerWorkItem(() =>
{
action();
return UniTask.CompletedTask;
}, awaitable: false);
if (!_queue.TryEnqueue(item))
return false;
ScheduleWorkers();
return true;
}
public async UniTask PostAsync(Func<UniTask> action, CancellationToken cancellationToken)
{
if (Volatile.Read(ref _disposed) != 0)
throw new ObjectDisposedException(nameof(ShrinkTaskPoolScheduler));
var item = new ShrinkSchedulerWorkItem(action, awaitable: true);
await _queue.EnqueueAsync(item, cancellationToken);
ScheduleWorkers();
await item.Completion.AttachExternalCancellation(cancellationToken);
}
public async UniTask ShutdownAsync(bool drain, TimeSpan timeout)
{
Interlocked.Exchange(ref _disposed, 1);
if (!drain)
_queue.DropPending();
else
await ShrinkSchedulerShutdown.WaitUntilDrainedAsync(
_queue, () => Volatile.Read(ref _workers) != 0, timeout);
_queue.Dispose();
}
public void Dispose()
{
if (Interlocked.Exchange(ref _disposed, 1) != 0)
return;
_queue.Dispose();
}
private void ScheduleWorkers()
{
while (!_queue.IsEmpty)
{
var current = Volatile.Read(ref _workers);
if (current >= _maxConcurrency ||
Interlocked.CompareExchange(ref _workers, current + 1, current) != current)
return;
UniTask.Void(WorkerAsync);
}
}
private async UniTaskVoid WorkerAsync()
{
await UniTask.SwitchToThreadPool();
try
{
while (_queue.TryDequeue(out var item))
await ShrinkSchedulerWorkItemRunner.RunAsync(item);
}
finally
{
Interlocked.Decrement(ref _workers);
if (!_queue.IsEmpty && Volatile.Read(ref _disposed) == 0)
ScheduleWorkers();
}
}
}
internal sealed class ShrinkDedicatedThreadScheduler : IShrinkBusScheduler
{
private readonly ShrinkSchedulerQueue _queue;
private readonly AutoResetEvent _signal = new(false);
private readonly ManualResetEventSlim _stopped = new(false);
private readonly Thread _thread;
private int _disposed;
private int _threadId;
private int _shutdownTimedOut;
private int _waitHandlesDisposed;
public ShrinkDedicatedThreadScheduler(string name, ShrinkBusOptions options)
{
_queue = new ShrinkSchedulerQueue(name, options);
_thread = new Thread(Run)
{
IsBackground = true,
Name = $"ShrinkBus:{name}"
};
_thread.Start();
}
public bool IsOnSchedulerThread => Thread.CurrentThread.ManagedThreadId == Volatile.Read(ref _threadId);
public bool TryPost(Action action)
{
if (Volatile.Read(ref _disposed) != 0)
return false;
if (IsOnSchedulerThread)
{
action();
return true;
}
var item = new ShrinkSchedulerWorkItem(() =>
{
action();
return UniTask.CompletedTask;
}, awaitable: false);
if (!_queue.TryEnqueue(item))
return false;
_signal.Set();
return true;
}
public async UniTask PostAsync(Func<UniTask> action, CancellationToken cancellationToken)
{
if (Volatile.Read(ref _disposed) != 0)
throw new ObjectDisposedException(nameof(ShrinkDedicatedThreadScheduler));
if (IsOnSchedulerThread)
{
await action();
return;
}
var item = new ShrinkSchedulerWorkItem(action, awaitable: true);
await _queue.EnqueueAsync(item, cancellationToken);
_signal.Set();
await item.Completion.AttachExternalCancellation(cancellationToken);
}
public async UniTask ShutdownAsync(bool drain, TimeSpan timeout)
{
Interlocked.Exchange(ref _disposed, 1);
if (!drain)
_queue.DropPending();
_signal.Set();
var waitMs = timeout == Timeout.InfiniteTimeSpan
? Timeout.Infinite
: Math.Max(0, (int)Math.Min(int.MaxValue, timeout.TotalMilliseconds));
var stopped = await Task.Run(() => _stopped.Wait(waitMs))
.AsUniTask(useCurrentSynchronizationContext: false);
if (!stopped)
{
_queue.DropPending();
Volatile.Write(ref _shutdownTimedOut, 1);
try
{
if (_stopped.IsSet)
DisposeWaitHandles();
}
catch (ObjectDisposedException)
{
}
return;
}
DisposeWaitHandles();
}
public void Dispose()
{
if (Interlocked.Exchange(ref _disposed, 1) != 0)
return;
_queue.DropPending();
Volatile.Write(ref _shutdownTimedOut, 1);
_signal.Set();
}
private void Run()
{
Volatile.Write(ref _threadId, Thread.CurrentThread.ManagedThreadId);
try
{
while (Volatile.Read(ref _disposed) == 0 || !_queue.IsEmpty)
{
if (!_queue.TryDequeue(out var item))
{
_signal.WaitOne(50);
continue;
}
ShrinkSchedulerWorkItemRunner.RunBlocking(item);
}
}
finally
{
_stopped.Set();
if (Volatile.Read(ref _shutdownTimedOut) != 0)
DisposeWaitHandles();
}
}
private void DisposeWaitHandles()
{
if (Interlocked.Exchange(ref _waitHandlesDisposed, 1) != 0)
return;
_queue.Dispose();
_signal.Dispose();
_stopped.Dispose();
}
}
internal static class ShrinkEventDiagnostics
{
public static void LogException(Exception exception)
{
#if UNITY_5_3_OR_NEWER
UnityEngine.Debug.LogException(exception);
#else
Trace.TraceError(exception.ToString());
#endif
}
}
internal static class ShrinkSchedulerShutdown
{
public static async UniTask WaitUntilDrainedAsync(ShrinkSchedulerQueue queue,
Func<bool> hasActiveWork, TimeSpan timeout)
{
var started = Stopwatch.StartNew();
while (!queue.IsEmpty || hasActiveWork())
{
if (timeout != Timeout.InfiniteTimeSpan && started.Elapsed >= timeout)
{
queue.DropPending();
return;
}
await Task.Delay(1).AsUniTask(useCurrentSynchronizationContext: false);
}
}
}
}