TimeoutController.cs

#pragma warning disable CS1591 // Missing XML comment for publicly visible type or member

using System.Threading;
using UnityEngine;
using Cysharp.Threading.Tasks.Triggers;
using System;
using Cysharp.Threading.Tasks.Internal;

namespace Cysharp.Threading.Tasks
{
    public sealed class TimeoutController : IDisposable
    {
        CancellationTokenSource source;
        readonly CancellationTokenSource linkedSource;
        StoppableDelayRealtimePromise timeoutDelay;

        public TimeoutController()
        {
            this.source = new CancellationTokenSource();
            this.linkedSource = null;
            this.timeoutDelay = null;
        }

        public TimeoutController(CancellationTokenSource linkedCancellationTokenSource)
        {
            this.source = new CancellationTokenSource();
            this.linkedSource = linkedCancellationTokenSource;
            this.timeoutDelay = null;
        }

        public CancellationToken Timeout(TimeSpan timeout)
        {
            if (source.IsCancellationRequested)
            {
                source = new CancellationTokenSource();
            }

            if (timeoutDelay == null)
            {
                RunDelayAsync(timeout).Forget(); // timeoutDelay = ... in RunDelayAsync(immediately, before await)
            }
            else
            {
                timeoutDelay.RestartStopwatch();
            }

            return source.Token;
        }

        public bool IsTimeout(OperationCanceledException exception)
        {
            return exception.CancellationToken == source.Token;
        }

        public void Reset()
        {
            if (timeoutDelay != null)
            {
                timeoutDelay.Stop(); // stop delay, will finish RunDelayAsync
                timeoutDelay = null;
            }
        }

        async UniTaskVoid RunDelayAsync(TimeSpan timeout)
        {
            timeoutDelay = StoppableDelayRealtimePromise.Create(timeout, PlayerLoopTiming.Update, (linkedSource == null) ? CancellationToken.None : linkedSource.Token, out var version);
            var linkedSourceCanceledOrStopCalled = await new UniTask(timeoutDelay, version).SuppressCancellationThrow();
            if (!linkedSourceCanceledOrStopCalled) // delay complete == timeout
            {
                source.Cancel();
                source.Dispose();
            }
            else if (linkedSource != null && linkedSource.IsCancellationRequested)
            {
                // linkedtoken canceled.
                source.Cancel();
                source.Dispose();
            }
            else
            {
                // Stop called, do nothing.
                // UnityEngine.Debug.Log("DEBUG:Stop called");
            }
        }

        public void Dispose()
        {
            if (timeoutDelay != null)
            {
                timeoutDelay.Stop();
                source.Dispose();
            }
        }

        sealed class StoppableDelayRealtimePromise : IUniTaskSource, IPlayerLoopItem, ITaskPoolNode<StoppableDelayRealtimePromise>
        {
            static OperationCanceledException ExterenalStopException = new OperationCanceledException();

            static TaskPool<StoppableDelayRealtimePromise> pool;
            StoppableDelayRealtimePromise nextNode;
            public ref StoppableDelayRealtimePromise NextNode => ref nextNode;

            static StoppableDelayRealtimePromise()
            {
                TaskPool.RegisterSizeGetter(typeof(StoppableDelayRealtimePromise), () => pool.Size);
            }

            long delayTimeSpanTicks;
            ValueStopwatch stopwatch;
            CancellationToken cancellationToken;
            bool externalStop;

            UniTaskCompletionSourceCore<AsyncUnit> core;

            StoppableDelayRealtimePromise()
            {
            }

            public static StoppableDelayRealtimePromise Create(TimeSpan delayTimeSpan, PlayerLoopTiming timing, CancellationToken cancellationToken, out short token)
            {
                if (!pool.TryPop(out var result))
                {
                    result = new StoppableDelayRealtimePromise();
                }

                result.stopwatch = ValueStopwatch.StartNew();
                result.delayTimeSpanTicks = delayTimeSpan.Ticks;
                result.cancellationToken = cancellationToken;
                result.externalStop = false;

                TaskTracker.TrackActiveTask(result, 3);

                PlayerLoopHelper.AddAction(timing, result);

                token = result.core.Version;
                return result;
            }

            public void Stop()
            {
                externalStop = true;
            }

            public void RestartStopwatch()
            {
                stopwatch = ValueStopwatch.StartNew();
            }

            public void GetResult(short token)
            {
                try
                {
                    core.GetResult(token);
                }
                finally
                {
                    TryReturn();
                }
            }

            public UniTaskStatus GetStatus(short token)
            {
                return core.GetStatus(token);
            }

            public UniTaskStatus UnsafeGetStatus()
            {
                return core.UnsafeGetStatus();
            }

            public void OnCompleted(Action<object> continuation, object state, short token)
            {
                core.OnCompleted(continuation, state, token);
            }

            public bool MoveNext()
            {
                if (cancellationToken.IsCancellationRequested)
                {
                    core.TrySetCanceled(cancellationToken);
                    return false;
                }

                if (externalStop)
                {
                    core.TrySetException(ExterenalStopException);
                    return false;
                }

                if (stopwatch.IsInvalid)
                {
                    core.TrySetResult(AsyncUnit.Default);
                    return false;
                }

                if (stopwatch.ElapsedTicks >= delayTimeSpanTicks)
                {
                    core.TrySetResult(AsyncUnit.Default);
                    return false;
                }

                return true;
            }

            bool TryReturn()
            {
                TaskTracker.RemoveTracking(this);
                core.Reset();
                stopwatch = default;
                cancellationToken = default;
                externalStop = false;
                return pool.TryPush(this);
            }
        }
    }   
}