Add CancelletaionToken.RegisterWithoutCaptureExecutionContext

neuecc · neuecc · commit 4e780c290be3 · 2018-09-12T19:02:41.000+09:00
diff --git a/Assets/Plugins/UniRx/Scripts/Async/CancellationTokenExtensions.cs b/Assets/Plugins/UniRx/Scripts/Async/CancellationTokenExtensions.cs
@@ -18,14 +18,58 @@ public static (UniTask, CancellationTokenRegistration) ToUniTask(this Cancellati
             }
 
             var promise = new UniTaskCompletionSource<AsyncUnit>();
-            return (promise.Task, cts.Register(cancellationTokenCallback, promise, false));
+            return (promise.Task, cts.RegisterWithoutCaptureExecutionContext(cancellationTokenCallback, promise));
         }
 
         static void Callback(object state)
         {
             var promise = (UniTaskCompletionSource<AsyncUnit>)state;
             promise.TrySetResult(AsyncUnit.Default);
         }
+
+        public static CancellationTokenRegistration RegisterWithoutCaptureExecutionContext(this CancellationToken cancellationToken, Action callback)
+        {
+            var restoreFlow = false;
+            if (!ExecutionContext.IsFlowSuppressed())
+            {
+                ExecutionContext.SuppressFlow();
+                restoreFlow = true;
+            }
+
+            try
+            {
+                return cancellationToken.Register(callback, false);
+            }
+            finally
+            {
+                if (restoreFlow)
+                {
+                    ExecutionContext.RestoreFlow();
+                }
+            }
+        }
+
+        public static CancellationTokenRegistration RegisterWithoutCaptureExecutionContext(this CancellationToken cancellationToken, Action<object> callback, object state)
+        {
+            var restoreFlow = false;
+            if (!ExecutionContext.IsFlowSuppressed())
+            {
+                ExecutionContext.SuppressFlow();
+                restoreFlow = true;
+            }
+
+            try
+            {
+                return cancellationToken.Register(callback, state, false);
+            }
+            finally
+            {
+                if (restoreFlow)
+                {
+                    ExecutionContext.RestoreFlow();
+                }
+            }
+        }
     }
 }
 
diff --git a/Assets/Plugins/UniRx/Scripts/Async/Triggers/AsyncTriggerBase.cs b/Assets/Plugins/UniRx/Scripts/Async/Triggers/AsyncTriggerBase.cs
@@ -188,7 +188,7 @@ protected UniTask<T> GetOrAddPromise<T>(ref AsyncTriggerPromise<T> promise, ref
                 PlayerLoopHelper.AddAction(PlayerLoopTiming.Update, new AwakeMonitor(this));
             }
 
-            var registrationToken = cancellationToken.Register(Callback, Tuple.Create((ICancellationTokenKeyDictionary)promises, (ICancelablePromise)cancellablePromise));
+            var registrationToken = cancellationToken.RegisterWithoutCaptureExecutionContext(Callback, Tuple.Create((ICancellationTokenKeyDictionary)promises, (ICancelablePromise)cancellablePromise));
             if (registeredCancellations == null)
             {
                 registeredCancellations = ArrayPool<CancellationTokenRegistration>.Shared.Rent(4);
diff --git a/Assets/Plugins/UniRx/Scripts/Async/UnityAsyncExtensions.uGUI.cs b/Assets/Plugins/UniRx/Scripts/Async/UnityAsyncExtensions.uGUI.cs
@@ -212,7 +212,7 @@ public AsyncUnityEventHandler(UnityEvent unityEvent, CancellationToken cancellat
 
             if (cancellationToken.CanBeCanceled)
             {
-                registration = cancellationToken.Register(cancellationCallback, this, false);
+                registration = cancellationToken.RegisterWithoutCaptureExecutionContext(cancellationCallback, this);
             }
 
             TaskTracker.TrackActiveTask(this, 3);
@@ -326,7 +326,7 @@ public AsyncUnityEventHandler(UnityEvent<T> unityEvent, CancellationToken cancel
 
             if (cancellationToken.CanBeCanceled)
             {
-                registration = cancellationToken.Register(cancellationCallback, this, false);
+                registration = cancellationToken.RegisterWithoutCaptureExecutionContext(cancellationCallback, this);
             }
 
             TaskTracker.TrackActiveTask(this, 3);
diff --git a/Assets/Plugins/UniRx/Scripts/Tasks/UniTaskObservableExtensions.cs b/Assets/Plugins/UniRx/Scripts/Tasks/UniTaskObservableExtensions.cs
@@ -116,7 +116,7 @@ public ToUniTaskObserver(UniTaskCompletionSource<T> promise, SingleAssignmentDis
 
                 if (this.cancellationToken.CanBeCanceled)
                 {
-                    this.registration = this.cancellationToken.Register(callback, this, false);
+                    this.registration = this.cancellationToken.RegisterWithoutCaptureExecutionContext(callback, this);
                 }
             }
 
@@ -186,7 +186,7 @@ public FirstValueToUniTaskObserver(UniTaskCompletionSource<T> promise, SingleAss
 
                 if (this.cancellationToken.CanBeCanceled)
                 {
-                    this.registration = this.cancellationToken.Register(callback, this, false);
+                    this.registration = this.cancellationToken.RegisterWithoutCaptureExecutionContext(callback, this);
                 }
             }
 
diff --git a/Assets/Plugins/UniRx/Scripts/UnityEngineBridge/ReactiveCommand.cs b/Assets/Plugins/UniRx/Scripts/UnityEngineBridge/ReactiveCommand.cs
@@ -189,7 +189,7 @@ public UniTask<T> WaitUntilExecuteAsync(CancellationToken cancellationToken)
 
             newPromise = new ReactivePropertyReusablePromise<T>(cancellationToken);
             removablePromises.Add(cancellationToken, newPromise);
-            cancellationToken.Register(Callback, Tuple.Create(this, newPromise), false);
+            cancellationToken.RegisterWithoutCaptureExecutionContext(Callback, Tuple.Create(this, newPromise));
 
             return newPromise.Task;
         }
@@ -417,7 +417,7 @@ public UniTask<T> WaitUntilExecuteAsync(CancellationToken cancellationToken)
 
             newPromise = new ReactivePropertyReusablePromise<T>(cancellationToken);
             removablePromises.Add(cancellationToken, newPromise);
-            cancellationToken.Register(Callback, Tuple.Create(this, newPromise), false);
+            cancellationToken.RegisterWithoutCaptureExecutionContext(Callback, Tuple.Create(this, newPromise));
 
             return newPromise.Task;
         }
diff --git a/Assets/Plugins/UniRx/Scripts/UnityEngineBridge/ReactiveProperty.cs b/Assets/Plugins/UniRx/Scripts/UnityEngineBridge/ReactiveProperty.cs
@@ -299,7 +299,7 @@ public UniTask<T> WaitUntilValueChangedAsync(CancellationToken cancellationToken
 
             newPromise = new ReactivePropertyReusablePromise<T>(cancellationToken);
             removablePromises.Add(cancellationToken, newPromise);
-            cancellationToken.Register(Callback, Tuple.Create(this, newPromise), false);
+            cancellationToken.RegisterWithoutCaptureExecutionContext(Callback, Tuple.Create(this, newPromise));
 
             return newPromise.Task;
         }
@@ -589,7 +589,7 @@ public UniTask<T> WaitUntilValueChangedAsync(CancellationToken cancellationToken
 
             newPromise = new ReactivePropertyReusablePromise<T>(cancellationToken);
             removablePromises.Add(cancellationToken, newPromise);
-            cancellationToken.Register(Callback, Tuple.Create(this, newPromise), false);
+            cancellationToken.RegisterWithoutCaptureExecutionContext(Callback, Tuple.Create(this, newPromise));
 
             return newPromise.Task;
         }

Original file line number	Diff line number	Diff line change
`@@ -188,7 +188,7 @@ protected UniTask<T> GetOrAddPromise<T>(ref AsyncTriggerPromise<T> promise, ref`
`188`	`188`	`PlayerLoopHelper.AddAction(PlayerLoopTiming.Update, new AwakeMonitor(this));`
`189`	`189`	`}`
`190`	`190`
`191`		`- var registrationToken = cancellationToken.Register(Callback, Tuple.Create((ICancellationTokenKeyDictionary)promises, (ICancelablePromise)cancellablePromise));`
	`191`	`+ var registrationToken = cancellationToken.RegisterWithoutCaptureExecutionContext(Callback, Tuple.Create((ICancellationTokenKeyDictionary)promises, (ICancelablePromise)cancellablePromise));`
`192`	`192`	`if (registeredCancellations == null)`
`193`	`193`	`{`
`194`	`194`	`registeredCancellations = ArrayPool<CancellationTokenRegistration>.Shared.Rent(4);`
Original file line number	Diff line number	Diff line change
`@@ -212,7 +212,7 @@ public AsyncUnityEventHandler(UnityEvent unityEvent, CancellationToken cancellat`
`212`	`212`
`213`	`213`	`if (cancellationToken.CanBeCanceled)`
`214`	`214`	`{`
`215`		`- registration = cancellationToken.Register(cancellationCallback, this, false);`
	`215`	`+ registration = cancellationToken.RegisterWithoutCaptureExecutionContext(cancellationCallback, this);`
`216`	`216`	`}`
`217`	`217`
`218`	`218`	`TaskTracker.TrackActiveTask(this, 3);`
`@@ -326,7 +326,7 @@ public AsyncUnityEventHandler(UnityEvent<T> unityEvent, CancellationToken cancel`
`326`	`326`
`327`	`327`	`if (cancellationToken.CanBeCanceled)`
`328`	`328`	`{`
`329`		`- registration = cancellationToken.Register(cancellationCallback, this, false);`
	`329`	`+ registration = cancellationToken.RegisterWithoutCaptureExecutionContext(cancellationCallback, this);`
`330`	`330`	`}`
`331`	`331`
`332`	`332`	`TaskTracker.TrackActiveTask(this, 3);`
Original file line number	Diff line number	Diff line change
`@@ -116,7 +116,7 @@ public ToUniTaskObserver(UniTaskCompletionSource<T> promise, SingleAssignmentDis`
`116`	`116`
`117`	`117`	`if (this.cancellationToken.CanBeCanceled)`
`118`	`118`	`{`
`119`		`- this.registration = this.cancellationToken.Register(callback, this, false);`
	`119`	`+ this.registration = this.cancellationToken.RegisterWithoutCaptureExecutionContext(callback, this);`
`120`	`120`	`}`
`121`	`121`	`}`
`122`	`122`
`@@ -186,7 +186,7 @@ public FirstValueToUniTaskObserver(UniTaskCompletionSource<T> promise, SingleAss`
`186`	`186`
`187`	`187`	`if (this.cancellationToken.CanBeCanceled)`
`188`	`188`	`{`
`189`		`- this.registration = this.cancellationToken.Register(callback, this, false);`
	`189`	`+ this.registration = this.cancellationToken.RegisterWithoutCaptureExecutionContext(callback, this);`
`190`	`190`	`}`
`191`	`191`	`}`
`192`	`192`
Original file line number	Diff line number	Diff line change
`@@ -189,7 +189,7 @@ public UniTask<T> WaitUntilExecuteAsync(CancellationToken cancellationToken)`
`189`	`189`
`190`	`190`	`newPromise = new ReactivePropertyReusablePromise<T>(cancellationToken);`
`191`	`191`	`removablePromises.Add(cancellationToken, newPromise);`
`192`		`- cancellationToken.Register(Callback, Tuple.Create(this, newPromise), false);`
	`192`	`+ cancellationToken.RegisterWithoutCaptureExecutionContext(Callback, Tuple.Create(this, newPromise));`
`193`	`193`
`194`	`194`	`return newPromise.Task;`
`195`	`195`	`}`
`@@ -417,7 +417,7 @@ public UniTask<T> WaitUntilExecuteAsync(CancellationToken cancellationToken)`
`417`	`417`
`418`	`418`	`newPromise = new ReactivePropertyReusablePromise<T>(cancellationToken);`
`419`	`419`	`removablePromises.Add(cancellationToken, newPromise);`
`420`		`- cancellationToken.Register(Callback, Tuple.Create(this, newPromise), false);`
	`420`	`+ cancellationToken.RegisterWithoutCaptureExecutionContext(Callback, Tuple.Create(this, newPromise));`
`421`	`421`
`422`	`422`	`return newPromise.Task;`
`423`	`423`	`}`
Original file line number	Diff line number	Diff line change
`@@ -299,7 +299,7 @@ public UniTask<T> WaitUntilValueChangedAsync(CancellationToken cancellationToken`
`299`	`299`
`300`	`300`	`newPromise = new ReactivePropertyReusablePromise<T>(cancellationToken);`
`301`	`301`	`removablePromises.Add(cancellationToken, newPromise);`
`302`		`- cancellationToken.Register(Callback, Tuple.Create(this, newPromise), false);`
	`302`	`+ cancellationToken.RegisterWithoutCaptureExecutionContext(Callback, Tuple.Create(this, newPromise));`
`303`	`303`
`304`	`304`	`return newPromise.Task;`
`305`	`305`	`}`
`@@ -589,7 +589,7 @@ public UniTask<T> WaitUntilValueChangedAsync(CancellationToken cancellationToken`
`589`	`589`
`590`	`590`	`newPromise = new ReactivePropertyReusablePromise<T>(cancellationToken);`
`591`	`591`	`removablePromises.Add(cancellationToken, newPromise);`
`592`		`- cancellationToken.Register(Callback, Tuple.Create(this, newPromise), false);`
	`592`	`+ cancellationToken.RegisterWithoutCaptureExecutionContext(Callback, Tuple.Create(this, newPromise));`
`593`	`593`
`594`	`594`	`return newPromise.Task;`
`595`	`595`	`}`