Workspace/Assets/Modules/ShrinkEventBus/Runtime/ShrinkEventBusInstance.cs

#nullable enable

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Runtime.CompilerServices;
using System.Threading;
using Cysharp.Threading.Tasks;

namespace ShrinkEventBus
{
    internal class ShrinkEventBusInstance : IShrinkEventBus
    {
        private sealed class SingleBusResolver : IShrinkBusResolver
        {
            private readonly IShrinkEventBus _bus;

            public SingleBusResolver(IShrinkEventBus bus) => _bus = bus;

            public IShrinkEventBus GetBus(ShrinkBusKey key) => key == _bus.Key
                ? _bus
                : throw new KeyNotFoundException($"Bus '{key}' is not available from '{_bus.Key}'.");

            public bool TryGetBus(ShrinkBusKey key, out IShrinkEventBus bus)
            {
                bus = _bus;
                return key == _bus.Key;
            }
        }

        private sealed class EventSubscription : IDisposable
        {
            private ShrinkEventBusInstance? _owner;
            private readonly long _subscriptionId;

            public EventSubscription(ShrinkEventBusInstance owner, long subscriptionId)
            {
                _owner = owner;
                _subscriptionId = subscriptionId;
            }

            public void Dispose()
            {
                Interlocked.Exchange(ref _owner, null)?.UnregisterSubscription(_subscriptionId);
            }
        }

        private static int _nextBusSlot;
        private readonly Dictionary<Type, IShrinkEventChannel> _channels = new();
        private readonly object _channelGate = new();
        private readonly IShrinkBusScheduler _scheduler;
        private readonly Action<IShrinkEvent, Type, ShrinkBusKey>? _postObserver;
        private readonly int _busSlot;
        private readonly bool _isInline;
        private long _nextSubscriptionId;
        private int _subscriberCount;
        private int _disposed;

        public ShrinkEventBusInstance(ShrinkEventBusBuilder builder)
        {
            if (builder == null)
                throw new ArgumentNullException(nameof(builder));

            Key = builder.Key;
            Options = builder.Options.CloneValidated();
            _busSlot = Interlocked.Increment(ref _nextBusSlot);
            _scheduler = ShrinkBusSchedulerFactory.Create(Options, Key.Id);
            _postObserver = builder.PostObserver;
            _isInline = Options.Scheduler == ShrinkBusSchedulerKind.Inline;
        }

        public ShrinkBusKey Key { get; }
        public ShrinkBusOptions Options { get; }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public virtual ShrinkPostResult Post<TEvent>(in TEvent eventData) where TEvent : IShrinkEvent
        {
            if (eventData is null)
                throw new ArgumentNullException(nameof(eventData));
            if (Volatile.Read(ref _disposed) != 0)
                return TraceWithoutDispatch(in eventData,
                    ShrinkPostResult.Rejected(ShrinkPostFailure.BusStopped), false);

            if (_isInline || _scheduler.IsOnSchedulerThread)
                return Dispatch(in eventData);

            var queuedEvent = eventData;
            if (_scheduler.TryPost(() => Dispatch(in queuedEvent)))
                return new ShrinkPostResult(true, false, false);
            return TraceWithoutDispatch(in eventData,
                ShrinkPostResult.Rejected(ShrinkPostFailure.QueueFull), false);
        }

        public virtual UniTask<ShrinkPostResult> PostAsync<TEvent>(TEvent eventData,
            CancellationToken cancellationToken = default) where TEvent : IShrinkEvent
        {
            if (eventData is null)
                throw new ArgumentNullException(nameof(eventData));
            if (Volatile.Read(ref _disposed) != 0)
                return UniTask.FromResult(TraceWithoutDispatch(in eventData,
                    ShrinkPostResult.Rejected(ShrinkPostFailure.BusStopped), true));
            if (cancellationToken.IsCancellationRequested)
                return UniTask.FromResult(TraceWithoutDispatch(in eventData,
                    new ShrinkPostResult(true, false, false, ShrinkPostFailure.Canceled), true));

            if (_isInline || _scheduler.IsOnSchedulerThread)
                return DispatchAsync(eventData, cancellationToken);

            return PostScheduledAsync(eventData, cancellationToken);
        }

        private async UniTask<ShrinkPostResult> PostScheduledAsync<TEvent>(TEvent eventData,
            CancellationToken cancellationToken) where TEvent : IShrinkEvent
        {
            var result = default(ShrinkPostResult);
            var dispatchStarted = false;
            try
            {
                await _scheduler.PostAsync(async () =>
                {
                    dispatchStarted = true;
                    result = await DispatchAsync(eventData, cancellationToken);
                }, cancellationToken);
                return result;
            }
            catch (ShrinkQueueFullException)
            {
                return TraceWithoutDispatch(in eventData,
                    ShrinkPostResult.Rejected(ShrinkPostFailure.QueueFull), true);
            }
            catch (ObjectDisposedException)
            {
                return TraceWithoutDispatch(in eventData,
                    ShrinkPostResult.Rejected(ShrinkPostFailure.BusStopped), true);
            }
            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
            {
                var canceled = new ShrinkPostResult(true, false, false, ShrinkPostFailure.Canceled);
                return dispatchStarted ? canceled : TraceWithoutDispatch(in eventData, canceled, true);
            }
        }

        public IDisposable Attach(object target)
        {
            if (target == null)
                throw new ArgumentNullException(nameof(target));
            if (target is not IShrinkGeneratedSubscriber generated)
                throw new InvalidOperationException(
                    $"Type {target.GetType().FullName} has no generated [ShrinkSubscribe] binding.");
            return generated.AttachGenerated(new SingleBusResolver(this), Key);
        }

        internal IDisposable SubscribeGenerated<TEvent>(Action<TEvent> handler,
            ShrinkSubscribeDescriptor descriptor) where TEvent : IShrinkEvent
        {
            var subscriptionId = Interlocked.Increment(ref _nextSubscriptionId);
            GetOrCreateChannel<TEvent>().Add(subscriptionId, subscriptionId, handler, descriptor);
            Interlocked.Increment(ref _subscriberCount);
            return new EventSubscription(this, subscriptionId);
        }

        internal IDisposable SubscribeGenerated<TEvent>(ShrinkAsyncEventHandler<TEvent> handler,
            ShrinkSubscribeDescriptor descriptor) where TEvent : IShrinkEvent
        {
            var subscriptionId = Interlocked.Increment(ref _nextSubscriptionId);
            GetOrCreateChannel<TEvent>().Add(subscriptionId, subscriptionId, handler, descriptor);
            Interlocked.Increment(ref _subscriberCount);
            return new EventSubscription(this, subscriptionId);
        }

        internal int SubscriberCount
        {
            get
            {
                return Math.Max(0, Volatile.Read(ref _subscriberCount));
            }
        }

        internal async UniTask ShutdownAsync()
        {
            if (Interlocked.Exchange(ref _disposed, 1) != 0)
                return;
            await _scheduler.ShutdownAsync(Options.DrainOnShutdown, Options.ShutdownTimeout);
            ClearChannels();
        }

        public void Dispose()
        {
            if (Interlocked.Exchange(ref _disposed, 1) != 0)
                return;
            _scheduler.Dispose();
            ClearChannels();
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        protected ShrinkPostResult Dispatch<TEvent>(in TEvent eventData) where TEvent : IShrinkEvent
        {
            if (_postObserver != null && EventBus.HasDetailedPostedObservers)
                return DispatchObserved(in eventData);

            var channel = ShrinkEventChannelSlots<TEvent>.Get(_busSlot);
            var result = channel == null
                ? new ShrinkPostResult(true, false, false)
                : channel.Post(in eventData);
            if (_postObserver != null && EventBus.HasPostedObservers)
                _postObserver(eventData, typeof(TEvent), Key);
            return result;
        }

        protected async UniTask<ShrinkPostResult> DispatchAsync<TEvent>(TEvent eventData,
            CancellationToken cancellationToken) where TEvent : IShrinkEvent
        {
            if (_postObserver != null && EventBus.HasDetailedPostedObservers)
                return await DispatchAsyncObserved(eventData, cancellationToken);

            var channel = ShrinkEventChannelSlots<TEvent>.Get(_busSlot);

            var result = channel == null
                ? ShrinkPostResult.Completed(false)
                : await channel.PostAsync(eventData, Options.DispatchMode, cancellationToken);
            if (_postObserver != null && EventBus.HasPostedObservers)
                _postObserver(eventData, typeof(TEvent), Key);
            return result;
        }

        private ShrinkPostResult DispatchObserved<TEvent>(in TEvent eventData) where TEvent : IShrinkEvent
        {
            var startedAt = Stopwatch.GetTimestamp();
            var threadId = Thread.CurrentThread.ManagedThreadId;
            var channel = ShrinkEventChannelSlots<TEvent>.Get(_busSlot);
            ShrinkPostResult result;
            try
            {
                result = channel == null
                    ? new ShrinkPostResult(true, false, false)
                    : channel.Post(in eventData);
            }
            catch
            {
                var failedResult = new ShrinkPostResult(true, channel != null,
                    eventData is IShrinkCancelableEvent cancelable && cancelable.IsCanceled,
                    ShrinkPostFailure.HandlerException);
                EventBus.NotifyDetailedPosted(new ShrinkEventTrace(DateTime.UtcNow, typeof(TEvent), Key,
                    Options, threadId, Stopwatch.GetTimestamp() - startedAt, failedResult, false));
                throw;
            }

            var elapsedTicks = Stopwatch.GetTimestamp() - startedAt;
            if (EventBus.HasPostedObservers)
                _postObserver!(eventData, typeof(TEvent), Key);
            EventBus.NotifyDetailedPosted(new ShrinkEventTrace(DateTime.UtcNow, typeof(TEvent), Key,
                Options, threadId, elapsedTicks, result, false));
            return result;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        protected ShrinkPostResult TraceWithoutDispatch<TEvent>(in TEvent eventData,
            ShrinkPostResult result, bool isAsync) where TEvent : IShrinkEvent
        {
            if (_postObserver != null && EventBus.HasDetailedPostedObservers)
            {
                EventBus.NotifyDetailedPosted(new ShrinkEventTrace(DateTime.UtcNow, typeof(TEvent), Key,
                    Options, Thread.CurrentThread.ManagedThreadId, 0L, result, isAsync));
            }
            return result;
        }

        private async UniTask<ShrinkPostResult> DispatchAsyncObserved<TEvent>(TEvent eventData,
            CancellationToken cancellationToken) where TEvent : IShrinkEvent
        {
            var startedAt = Stopwatch.GetTimestamp();
            var threadId = Thread.CurrentThread.ManagedThreadId;
            var channel = ShrinkEventChannelSlots<TEvent>.Get(_busSlot);
            ShrinkPostResult result;
            try
            {
                result = channel == null
                    ? ShrinkPostResult.Completed(false)
                    : await channel.PostAsync(eventData, Options.DispatchMode, cancellationToken);
            }
            catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
            {
                var canceledResult = new ShrinkPostResult(true, channel != null, false,
                    ShrinkPostFailure.Canceled);
                EventBus.NotifyDetailedPosted(new ShrinkEventTrace(DateTime.UtcNow, typeof(TEvent), Key,
                    Options, threadId, Stopwatch.GetTimestamp() - startedAt, canceledResult, true));
                throw;
            }
            catch
            {
                var failedResult = new ShrinkPostResult(true, channel != null,
                    eventData is IShrinkCancelableEvent cancelable && cancelable.IsCanceled,
                    ShrinkPostFailure.HandlerException);
                EventBus.NotifyDetailedPosted(new ShrinkEventTrace(DateTime.UtcNow, typeof(TEvent), Key,
                    Options, threadId, Stopwatch.GetTimestamp() - startedAt, failedResult, true));
                throw;
            }

            var elapsedTicks = Stopwatch.GetTimestamp() - startedAt;
            if (EventBus.HasPostedObservers)
                _postObserver!(eventData, typeof(TEvent), Key);
            EventBus.NotifyDetailedPosted(new ShrinkEventTrace(DateTime.UtcNow, typeof(TEvent), Key,
                Options, threadId, elapsedTicks, result, true));
            return result;
        }

        private ShrinkEventChannel<TEvent> GetOrCreateChannel<TEvent>() where TEvent : IShrinkEvent
        {
            var existing = ShrinkEventChannelSlots<TEvent>.Get(_busSlot);
            if (existing != null)
                return existing;

            lock (_channelGate)
            {
                existing = ShrinkEventChannelSlots<TEvent>.Get(_busSlot);
                if (existing != null)
                    return existing;
                var channel = new ShrinkEventChannel<TEvent>();
                _channels.Add(typeof(TEvent), channel);
                ShrinkEventChannelSlots<TEvent>.Set(_busSlot, channel);
                return channel;
            }
        }

        private void UnregisterSubscription(long subscriptionId)
        {
            lock (_channelGate)
            {
                foreach (var channel in _channels.Values)
                {
                    if (channel.RemoveSubscription(subscriptionId))
                    {
                        Interlocked.Decrement(ref _subscriberCount);
                        break;
                    }
                }
            }
        }

        private void ClearChannels()
        {
            lock (_channelGate)
            {
                foreach (var channel in _channels.Values)
                {
                    channel.DetachSlot(_busSlot);
                    channel.Clear();
                }
                _channels.Clear();
                Volatile.Write(ref _subscriberCount, 0);
            }
        }

        protected bool IsDisposed
        {
            [MethodImpl(MethodImplOptions.AggressiveInlining)]
            get => Volatile.Read(ref _disposed) != 0;
        }
    }

    internal sealed class ShrinkInlineEventBusInstance : ShrinkEventBusInstance
    {
        public ShrinkInlineEventBusInstance(ShrinkEventBusBuilder builder) : base(builder)
        {
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        public override ShrinkPostResult Post<TEvent>(in TEvent eventData)
        {
            if (eventData is null)
                throw new ArgumentNullException(nameof(eventData));
            if (IsDisposed)
                return TraceWithoutDispatch(in eventData,
                    ShrinkPostResult.Rejected(ShrinkPostFailure.BusStopped), false);
            return Dispatch(in eventData);
        }

        public override UniTask<ShrinkPostResult> PostAsync<TEvent>(TEvent eventData,
            CancellationToken cancellationToken = default)
        {
            if (eventData is null)
                throw new ArgumentNullException(nameof(eventData));
            if (IsDisposed)
                return UniTask.FromResult(TraceWithoutDispatch(in eventData,
                    ShrinkPostResult.Rejected(ShrinkPostFailure.BusStopped), true));
            if (cancellationToken.IsCancellationRequested)
                return UniTask.FromResult(TraceWithoutDispatch(in eventData,
                    new ShrinkPostResult(true, false, false, ShrinkPostFailure.Canceled), true));
            return DispatchAsync(eventData, cancellationToken);
        }
    }
}