#nullable enable using System; using System.Collections.Concurrent; using System.Threading; using System.Threading.Tasks; namespace ShrinkEventBus { internal interface IDotNetScheduler : IDisposable { bool IsOnSchedulerThread { get; } bool TryPost(Action action); ValueTask PostAsync(Func action, CancellationToken cancellationToken); } internal sealed class QueueFullException : InvalidOperationException { } internal static class DotNetSchedulerFactory { public static IDotNetScheduler Create(string name, ShrinkBusOptions options) { return options.Scheduler switch { ShrinkBusSchedulerKind.MainThread => new SynchronizationContextScheduler(options), ShrinkBusSchedulerKind.DedicatedThread => new DedicatedScheduler(name, options), ShrinkBusSchedulerKind.TaskPool => new TaskPoolScheduler(options), _ => new InlineScheduler() }; } } internal sealed class InlineScheduler : IDotNetScheduler { public bool IsOnSchedulerThread => true; public bool TryPost(Action action) { action(); return true; } public ValueTask PostAsync(Func action, CancellationToken cancellationToken) { cancellationToken.ThrowIfCancellationRequested(); return action(); } public void Dispose() { } } internal sealed class SynchronizationContextScheduler : IDotNetScheduler { private readonly SynchronizationContext _context; private readonly int _threadId; private int _disposed; public SynchronizationContextScheduler(ShrinkBusOptions options) { _context = SynchronizationContext.Current ?? throw new InvalidOperationException( "MainThread scheduler requires a current SynchronizationContext in a pure .NET host."); _threadId = Thread.CurrentThread.ManagedThreadId; } public bool IsOnSchedulerThread => Thread.CurrentThread.ManagedThreadId == _threadId; public bool TryPost(Action action) { if (Volatile.Read(ref _disposed) != 0) return false; if (IsOnSchedulerThread) action(); else _context.Post(_ => action(), null); return true; } public ValueTask PostAsync(Func action, CancellationToken cancellationToken) { if (IsOnSchedulerThread) return action(); return new ValueTask(PostCoreAsync(action, cancellationToken)); } private Task PostCoreAsync(Func action, CancellationToken cancellationToken) { var completion = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); _context.Post(async _ => { try { cancellationToken.ThrowIfCancellationRequested(); await action(); completion.SetResult(true); } catch (OperationCanceledException) { completion.SetCanceled(); } catch (Exception ex) { completion.SetException(ex); } }, null); return completion.Task; } public void Dispose() => Interlocked.Exchange(ref _disposed, 1); } internal sealed class TaskPoolScheduler : IDotNetScheduler { private readonly SemaphoreSlim _concurrency; private readonly int _capacity; private int _pending; private int _disposed; public TaskPoolScheduler(ShrinkBusOptions options) { var concurrency = options.DispatchMode == ShrinkDispatchMode.Ordered ? 1 : options.MaxConcurrency; _concurrency = new SemaphoreSlim(concurrency, concurrency); _capacity = options.QueueCapacity; } public bool IsOnSchedulerThread => false; public bool TryPost(Action action) { if (Volatile.Read(ref _disposed) != 0 || Interlocked.Increment(ref _pending) > _capacity) { Interlocked.Decrement(ref _pending); return false; } _ = Task.Run(async () => { await _concurrency.WaitAsync().ConfigureAwait(false); try { action(); } finally { _concurrency.Release(); Interlocked.Decrement(ref _pending); } }); return true; } public async ValueTask PostAsync(Func action, CancellationToken cancellationToken) { if (Volatile.Read(ref _disposed) != 0 || Interlocked.Increment(ref _pending) > _capacity) { Interlocked.Decrement(ref _pending); throw new QueueFullException(); } await _concurrency.WaitAsync(cancellationToken).ConfigureAwait(false); try { await action().ConfigureAwait(false); } finally { _concurrency.Release(); Interlocked.Decrement(ref _pending); } } public void Dispose() { Interlocked.Exchange(ref _disposed, 1); _concurrency.Dispose(); } } internal sealed class DedicatedScheduler : IDotNetScheduler { private sealed class Work { public Func Action = null!; public TaskCompletionSource? Completion; } private readonly BlockingCollection _queue; private readonly Thread _thread; private int _threadId; private int _disposed; public DedicatedScheduler(string name, ShrinkBusOptions options) { _queue = new BlockingCollection(options.QueueCapacity); _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 (IsOnSchedulerThread) { action(); return true; } return Volatile.Read(ref _disposed) == 0 && _queue.TryAdd(new Work { Action = () => { action(); return default; } }); } public ValueTask PostAsync(Func action, CancellationToken cancellationToken) { if (IsOnSchedulerThread) return action(); var completion = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); if (Volatile.Read(ref _disposed) != 0 || !_queue.TryAdd(new Work { Action = action, Completion = completion })) throw new QueueFullException(); return new ValueTask(completion.Task); } public void Dispose() { if (Interlocked.Exchange(ref _disposed, 1) != 0) return; _queue.CompleteAdding(); } private void Run() { Volatile.Write(ref _threadId, Thread.CurrentThread.ManagedThreadId); foreach (var work in _queue.GetConsumingEnumerable()) { try { work.Action().AsTask().GetAwaiter().GetResult(); work.Completion?.SetResult(true); } catch (OperationCanceledException) { work.Completion?.SetCanceled(); } catch (Exception ex) { work.Completion?.SetException(ex); } } _queue.Dispose(); } } }