Make sure that CancellationToken can unsubscribe

Author: hadashiA

Assets/UniTaskPubSub/Runtime/AsyncMessageBus.cs

@@ -46,6 +46,33 @@ public static IDisposable Subscribe<T>(
                 return UniTask.CompletedTask;
             });
         }
+
+        public static void Subscribe<T>(
+            this IAsyncSubscriber subscriber,
+            Func<T, CancellationToken, UniTask> action,
+            CancellationToken unsubscribeToken)
+        {
+            var subscription = subscriber.Subscribe(action);
+            unsubscribeToken.Register(x => ((IDisposable)x).Dispose(), subscription);
+        }
+
+        public static void Subscribe<T>(
+            this IAsyncSubscriber subscriber,
+            Func<T, UniTask> action,
+            CancellationToken unsubscribeToken)
+        {
+            var subscription = subscriber.Subscribe(action);
+            unsubscribeToken.Register(x => ((IDisposable)x).Dispose(), subscription);
+        }
+
+        public static void Subscribe<T>(
+            this IAsyncSubscriber subscriber,
+            Action<T> action,
+            CancellationToken unsubscribeToken)
+        {
+            var subscription = subscriber.Subscribe(action);
+            unsubscribeToken.Register(x => ((IDisposable)x).Dispose(), subscription);
+        }
     }
 
     public class AsyncMessageBus : IAsyncPublisher, IAsyncSubscriber, IDisposable

Assets/UniTaskPubSub/Tests/AsyncMessageBusTest.cs

@@ -138,17 +138,29 @@ public IEnumerator DisposeAll() => UniTask.ToCoroutine(async () =>
         });
 
         [UnityTest]
-        public IEnumerator SubscribeWithSync() => UniTask.ToCoroutine(async () =>
+        public IEnumerator CancelSubscription() => UniTask.ToCoroutine(async () =>
         {
             var messageBus = new AsyncMessageBus();
+            var unsubscribeTokenSource1 = new CancellationTokenSource();
+            var unsubscribeTokenSource2 = new CancellationTokenSource();
             var receives = 0;
 
-            messageBus.Subscribe<TestMessage>(msg =>
+            messageBus.Subscribe<TestMessage>(async msg =>
             {
+                await UniTask.Yield();
                 receives += 1;
-            });
+            }, unsubscribeTokenSource1.Token);
+
+            messageBus.Subscribe<TestMessage>(async msg =>
+            {
+                await UniTask.Yield();
+                receives += 1;
+            }, unsubscribeTokenSource2.Token);
+
+            unsubscribeTokenSource1.Cancel();
 
             await messageBus.PublishAsync(new TestMessage(100));
+
             Assert.That(receives, Is.EqualTo(1));
         });
     }