Merge pull request #113 from AnnulusGames/feature-cancel-behaviour

Add: CancelBehaviour to ToUniTask, ToValueTask, ToAwaitable

src/LitMotion/Assets/LitMotion/Runtime/CancelBehaviour.cs

@@ -0,0 +1,14 @@
+namespace LitMotion
+{
+    /// <summary>
+    /// Specifies the behavior when await is canceled.
+    /// </summary>
+    public enum CancelBehaviour
+    {
+        CancelAndCancelAwait,
+        CompleteAndCancelAwait,
+        CancelAwait,
+        Cancel,
+        Complete
+    }
+}

src/LitMotion/Assets/LitMotion/Runtime/External/UniTask/LitMotionUniTaskExtensions.cs

@@ -1,5 +1,4 @@
 #if LITMOTION_SUPPORT_UNITASK
-using System;
 using System.Threading;
 using Cysharp.Threading.Tasks;
 
@@ -19,7 +18,20 @@ public static class LitMotionUniTaskExtensions
         public static UniTask ToUniTask(this MotionHandle handle, CancellationToken cancellationToken = default)
         {
             if (!handle.IsActive()) return UniTask.CompletedTask;
-            return new UniTask(MotionConfiguredSource.Create(handle, cancellationToken, out var token), token);
+            return new UniTask(UniTaskMotionConfiguredSource.Create(handle, CancelBehaviour.CancelAndCancelAwait, cancellationToken, out var token), token);
+        }
+
+        /// <summary>
+        /// Convert motion handle to UniTask.
+        /// </summary>
+        /// <param name="handle">This motion handle</param>
+        /// <param name="cancelBehaviour">Behavior when canceling</param>
+        /// <param name="cancellationToken">CancellationToken</param>
+        /// <returns></returns>
+        public static UniTask ToUniTask(this MotionHandle handle, CancelBehaviour cancelBehaviour, CancellationToken cancellationToken = default)
+        {
+            if (!handle.IsActive()) return UniTask.CompletedTask;
+            return new UniTask(UniTaskMotionConfiguredSource.Create(handle, cancelBehaviour, cancellationToken, out var token), token);
         }
 
         /// <summary>
@@ -43,159 +55,5 @@ public static MotionHandle BindToAsyncReactiveProperty<TValue, TOptions, TAdapte
             });
         }
     }
-
-    internal sealed class MotionConfiguredSource : IUniTaskSource, ITaskPoolNode<MotionConfiguredSource>
-    {
-        static TaskPool<MotionConfiguredSource> pool;
-        MotionConfiguredSource nextNode;
-        public ref MotionConfiguredSource NextNode => ref nextNode;
-
-        static MotionConfiguredSource()
-        {
-            TaskPool.RegisterSizeGetter(typeof(MotionConfiguredSource), () => pool.Size);
-        }
-
-        readonly Action onCancelCallbackDelegate;
-        readonly Action onCompleteCallbackDelegate;
-
-        MotionHandle motionHandle;
-        CancellationToken cancellationToken;
-        CancellationTokenRegistration cancellationRegistration;
-        bool canceled;
-
-        Action originalCompleteAction;
-        Action originalCancelAction;
-        UniTaskCompletionSourceCore<AsyncUnit> core;
-
-        MotionConfiguredSource()
-        {
-            onCancelCallbackDelegate = new(OnCancelCallbackDelegate);
-            onCompleteCallbackDelegate = new(OnCompleteCallbackDelegate);
-        }
-
-        public static IUniTaskSource Create(MotionHandle motionHandle, CancellationToken cancellationToken, out short token)
-        {
-            if (cancellationToken.IsCancellationRequested)
-            {
-                motionHandle.Cancel();
-                return AutoResetUniTaskCompletionSource.CreateFromCanceled(cancellationToken, out token);
-            }
-
-            if (!pool.TryPop(out var result))
-            {
-                result = new MotionConfiguredSource();
-            }
-
-            result.motionHandle = motionHandle;
-            result.cancellationToken = cancellationToken;
-            result.canceled = false;
-
-            var callbacks = MotionStorageManager.GetMotionCallbacks(motionHandle);
-            result.originalCancelAction = callbacks.OnCancelAction;
-            result.originalCompleteAction = callbacks.OnCompleteAction;
-            callbacks.OnCancelAction = result.onCancelCallbackDelegate;
-            callbacks.OnCompleteAction = result.onCompleteCallbackDelegate;
-            MotionStorageManager.SetMotionCallbacks(motionHandle, callbacks);
-
-            if (result.originalCancelAction == result.onCancelCallbackDelegate)
-            {
-                result.originalCancelAction = null;
-            }
-            if (result.originalCompleteAction == result.onCompleteCallbackDelegate)
-            {
-                result.originalCompleteAction = null;
-            }
-
-            if (cancellationToken.CanBeCanceled)
-            {
-                result.cancellationRegistration = cancellationToken.RegisterWithoutCaptureExecutionContext(static x =>
-                {
-                    var source = (MotionConfiguredSource)x;
-                    var motionHandle = source.motionHandle;
-                    if (motionHandle.IsActive()) motionHandle.Cancel();
-                    source.core.TrySetCanceled(source.cancellationToken);
-                }, result);
-            }
-
-            TaskTracker.TrackActiveTask(result, 3);
-
-            token = result.core.Version;
-            return result;
-        }
-
-        void OnCompleteCallbackDelegate()
-        {
-            if (cancellationToken.IsCancellationRequested)
-            {
-                canceled = true;
-            }
-            if (canceled)
-            {
-                core.TrySetCanceled(cancellationToken);
-            }
-            else
-            {
-                originalCompleteAction?.Invoke();
-                core.TrySetResult(AsyncUnit.Default);
-            }
-        }
-
-        void OnCancelCallbackDelegate()
-        {
-            originalCancelAction?.Invoke();
-            core.TrySetCanceled(cancellationToken);
-        }
-
-        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);
-        }
-
-        bool TryReturn()
-        {
-            TaskTracker.RemoveTracking(this);
-            core.Reset();
-            cancellationRegistration.Dispose();
-
-            RestoreOriginalCallback();
-
-            motionHandle = default;
-            cancellationToken = default;
-            originalCompleteAction = default;
-            originalCancelAction = default;
-            return pool.TryPush(this);
-        }
-
-        void RestoreOriginalCallback()
-        {
-            if (!motionHandle.IsActive()) return;
-            var callbacks = MotionStorageManager.GetMotionCallbacks(motionHandle);
-            callbacks.OnCancelAction = originalCancelAction;
-            callbacks.OnCompleteAction = originalCompleteAction;
-            MotionStorageManager.SetMotionCallbacks(motionHandle, callbacks);
-        }
-    }
 }
 #endif

src/LitMotion/Assets/LitMotion/Runtime/External/UniTask/UniTaskMotionConfiguredSource.cs

@@ -0,0 +1,94 @@
+#if LITMOTION_SUPPORT_UNITASK
+using System;
+using System.Threading;
+using Cysharp.Threading.Tasks;
+
+namespace LitMotion
+{
+    internal sealed class UniTaskMotionConfiguredSource : MotionConfiguredSourceBase ,IUniTaskSource, ITaskPoolNode<UniTaskMotionConfiguredSource>
+    {
+        static UniTaskMotionConfiguredSource()
+        {
+            TaskPool.RegisterSizeGetter(typeof(UniTaskMotionConfiguredSource), () => pool.Size);
+        }
+
+        UniTaskMotionConfiguredSource() : base() { }
+
+        static TaskPool<UniTaskMotionConfiguredSource> pool;
+        UniTaskMotionConfiguredSource nextNode;
+        public ref UniTaskMotionConfiguredSource NextNode => ref nextNode;
+
+        UniTaskCompletionSourceCore<AsyncUnit> core;
+
+        public static IUniTaskSource Create(MotionHandle motionHandle, CancelBehaviour cancelBehaviour, CancellationToken cancellationToken, out short token)
+        {
+            if (cancellationToken.IsCancellationRequested)
+            {
+                OnCanceledTokenReceived(motionHandle, cancelBehaviour);
+                return AutoResetUniTaskCompletionSource.CreateFromCanceled(cancellationToken, out token);
+            }
+
+            if (!pool.TryPop(out var result))
+            {
+                result = new UniTaskMotionConfiguredSource();
+            }
+
+            result.Initialize(motionHandle, cancelBehaviour, cancellationToken);
+
+            TaskTracker.TrackActiveTask(result, 3);
+
+            token = result.core.Version;
+            return result;
+        }
+
+        protected override void SetTaskCanceled(CancellationToken cancellationToken)
+        {
+            core.TrySetCanceled(cancellationToken);
+        }
+
+        protected override void SetTaskCompleted()
+        {
+            core.TrySetResult(AsyncUnit.Default);
+        }
+
+        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);
+        }
+
+        bool TryReturn()
+        {
+            TaskTracker.RemoveTracking(this);
+            core.Reset();
+
+            DisposeRegistration();
+            RestoreOriginalCallback();
+            ResetFields();
+
+            return pool.TryPush(this);
+        }
+    }
+}
+#endif