ShrinkEventBus/Tools~/DotNet/ShrinkEventBus.Core/Schedulers.cs

#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<ValueTask> 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<ValueTask> 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<ValueTask> action, CancellationToken cancellationToken)
        {
            if (IsOnSchedulerThread)
                return action();
            return new ValueTask(PostCoreAsync(action, cancellationToken));
        }

        private Task PostCoreAsync(Func<ValueTask> action, CancellationToken cancellationToken)
        {
            var completion = new TaskCompletionSource<bool>(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<ValueTask> 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<ValueTask> Action = null!;
            public TaskCompletionSource<bool>? Completion;
        }

        private readonly BlockingCollection<Work> _queue;
        private readonly Thread _thread;
        private int _threadId;
        private int _disposed;

        public DedicatedScheduler(string name, ShrinkBusOptions options)
        {
            _queue = new BlockingCollection<Work>(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<ValueTask> action, CancellationToken cancellationToken)
        {
            if (IsOnSchedulerThread)
                return action();
            var completion = new TaskCompletionSource<bool>(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();
        }
    }
}