Changed so that cancellation tokens can be specified in AsyncRegister.

tosh-coding · tosh-coding · commit 6c6f2dbe2eb8 · 2025-07-22T23:38:10.000+09:00
diff --git a/src/AsyncFiberWorks/Procedures/AsyncRegister.cs b/src/AsyncFiberWorks/Procedures/AsyncRegister.cs
@@ -17,13 +17,16 @@ public class AsyncRegister : IDisposable
         private readonly SemaphoreSlim _notifierSet = new SemaphoreSlim(0);
         private readonly SemaphoreSlim _notifierClear = new SemaphoreSlim(0);
         private bool _reading;
+        private CancellationToken _cancellationToken;
 
         /// <summary>
         /// Subscribe a task list.
         /// </summary>
         /// <param name="taskList"></param>
-        public AsyncRegister(ISequentialTaskListRegistry taskList)
+        /// <param name="cancellationToken"></param>
+        public AsyncRegister(ISequentialTaskListRegistry taskList, CancellationToken cancellationToken = default)
         {
+            _cancellationToken = cancellationToken;
             _subscription = taskList.Add(async () =>
             {
                 await SetFlagAndWaitClearing().ConfigureAwait(false);
@@ -43,6 +46,10 @@ public async Task SetFlagAndWaitClearing()
                 {
                     return;
                 }
+                if (_cancellationToken.IsCancellationRequested)
+                {
+                    return;
+                }
                 if (_hasValue)
                 {
                     throw new InvalidOperationException();
@@ -60,24 +67,34 @@ public async Task SetFlagAndWaitClearing()
                 _notifierSet.Release(1);
             }
 
-            await _notifierClear.WaitAsync().ConfigureAwait(false);
+            try
+            {
+                await _notifierClear.WaitAsync(_cancellationToken).ConfigureAwait(false);
+            }
+            catch (OperationCanceledException)
+            {
+                return;
+            }
         }
 
         /// <summary>
         /// Wait for the flag to be set.
         /// </summary>
-        /// <param name="token"></param>
         /// <returns></returns>
         /// <exception cref="ObjectDisposedException"></exception>
         /// <exception cref="OperationCanceledException"></exception>
-        public async Task WaitSetting(CancellationToken token = default)
+        public async Task WaitSetting()
         {
             lock (_lockObj)
             {
                 if (_isDisposed)
                 {
                     throw new ObjectDisposedException(GetType().FullName);
                 }
+                if (_cancellationToken.IsCancellationRequested)
+                {
+                    throw new OperationCanceledException();
+                }
                 if (_hasValue && _reading)
                 {
                     _hasValue = false;
@@ -86,7 +103,7 @@ public async Task WaitSetting(CancellationToken token = default)
                 }
             }
 
-            await _notifierSet.WaitAsync(token).ConfigureAwait(false);
+            await _notifierSet.WaitAsync(_cancellationToken).ConfigureAwait(false);
 
             lock (_lockObj)
             {
diff --git a/src/AsyncFiberWorks/Procedures/AsyncRegisterOfT.cs b/src/AsyncFiberWorks/Procedures/AsyncRegisterOfT.cs
@@ -19,13 +19,16 @@ public class AsyncRegister<T> : IDisposable
         private readonly SemaphoreSlim _notifierSet = new SemaphoreSlim(0);
         private readonly SemaphoreSlim _notifierClear = new SemaphoreSlim(0);
         private bool _reading;
+        private CancellationToken _cancellationToken;
 
         /// <summary>
         /// Subscribe a handler list.
         /// </summary>
         /// <param name="handlerList"></param>
-        public AsyncRegister(ISequentialHandlerListRegistry<T> handlerList)
+        /// <param name="cancellationToken"></param>
+        public AsyncRegister(ISequentialHandlerListRegistry<T> handlerList, CancellationToken cancellationToken = default)
         {
+            _cancellationToken = cancellationToken;
             _subscription = handlerList.Add((arg) => SetValueAndWaitClearing(arg));
         }
 
@@ -43,6 +46,10 @@ async Task<bool> SetValueAndWaitClearing(T newValue)
                 {
                     return false;
                 }
+                if (_cancellationToken.IsCancellationRequested)
+                {
+                    return false;
+                }
                 if (_hasValue)
                 {
                     throw new InvalidOperationException();
@@ -62,25 +69,35 @@ async Task<bool> SetValueAndWaitClearing(T newValue)
                 _notifierSet.Release(1);
             }
 
-            await _notifierClear.WaitAsync().ConfigureAwait(false);
+            try
+            {
+                await _notifierClear.WaitAsync(_cancellationToken).ConfigureAwait(false);
+            }
+            catch (OperationCanceledException)
+            {
+                return false;
+            }
             return _currentValue.Processed;
         }
 
         /// <summary>
         /// Wait for the value to be set.
         /// </summary>
-        /// <param name="token"></param>
         /// <returns></returns>
         /// <exception cref="ObjectDisposedException"></exception>
         /// <exception cref="OperationCanceledException"></exception>
-        public async Task<ProcessedFlagEventArgs<T>> WaitSetting(CancellationToken token = default)
+        public async Task<ProcessedFlagEventArgs<T>> WaitSetting()
         {
             lock (_lockObj)
             {
                 if (_isDisposed)
                 {
                     throw new ObjectDisposedException(GetType().FullName);
                 }
+                if (_cancellationToken.IsCancellationRequested)
+                {
+                    throw new OperationCanceledException();
+                }
                 if (_hasValue && _reading)
                 {
                     _currentValue.Arg = default;
@@ -90,7 +107,7 @@ public async Task<ProcessedFlagEventArgs<T>> WaitSetting(CancellationToken token
                 }
             }
 
-            await _notifierSet.WaitAsync(token).ConfigureAwait(false);
+            await _notifierSet.WaitAsync(_cancellationToken).ConfigureAwait(false);
 
             lock (_lockObj)
             {
diff --git a/src/AsyncFiberWorksTests/AsyncRegisterTests.cs b/src/AsyncFiberWorksTests/AsyncRegisterTests.cs
@@ -19,8 +19,7 @@ public async Task WaitingOfAsyncRegister()
 
             Func<int, Task> func = async (maxCount) =>
             {
-                var reg = new AsyncRegister(taskList);
-                try
+                using (var reg = new AsyncRegister(taskList))
                 {
                     int counter = 0;
                     while (counter < maxCount)
@@ -33,10 +32,6 @@ public async Task WaitingOfAsyncRegister()
                         counter += 1;
                     }
                 }
-                finally
-                {
-                    reg.Dispose();
-                }
             };
 
             var task1 = func(3);
@@ -52,6 +47,51 @@ public async Task WaitingOfAsyncRegister()
             Assert.AreEqual(3 + 6, resultCounter);
         }
 
+        [Test]
+        public async Task CancellationOfAsyncRegister()
+        {
+            var taskList = new FiberAndTaskPairList();
+            int resultCounter = 0;
+            int exceptionCounter = 0;
+            var lockObj = new object();
+            var cts = new CancellationTokenSource();
+            var cancellationToken = cts.Token;
+
+            var func = new Func<Task>(async () =>
+            {
+                using (var reg = new AsyncRegister(taskList, cancellationToken))
+                {
+                    try
+                    {
+                        while (true)
+                        {
+                            await reg.WaitSetting();
+                            lock (lockObj)
+                            {
+                                resultCounter += 1;
+                            }
+                        }
+                    }
+                    catch (OperationCanceledException)
+                    {
+                        exceptionCounter += 1;
+                    }
+                }
+            });
+
+            var task1 = func();
+
+            var fiber = new PoolFiber();
+            await taskList.InvokeSequentialAsync(fiber);
+            await taskList.InvokeSequentialAsync(fiber);
+            await taskList.InvokeSequentialAsync(fiber);
+            cts.Cancel();
+
+            await task1;
+            Assert.AreEqual(3, resultCounter);
+            Assert.AreEqual(1, exceptionCounter);
+        }
+
         [Test]
         public async Task WaitingOfAsyncRegisterOfT()
         {
@@ -61,8 +101,7 @@ public async Task WaitingOfAsyncRegisterOfT()
 
             Func<int, Task> func = async (maxCount) =>
             {
-                var reg = new AsyncRegister<int>(handlerList);
-                try
+                using (var reg = new AsyncRegister<int>(handlerList))
                 {
                     int counter = 0;
                     while (counter < maxCount)
@@ -75,10 +114,6 @@ public async Task WaitingOfAsyncRegisterOfT()
                         counter += 1;
                     }
                 }
-                finally
-                {
-                    reg.Dispose();
-                }
             };
 
             var task1 = func(3);
@@ -99,6 +134,54 @@ int Sigma(int n)
             return n * (n + 1) / 2;
         }
 
+        [Test]
+        public async Task CancellationOfAsyncRegisterOfT()
+        {
+            var handlerList = new FiberAndHandlerPairList<int>();
+            int resultCounter = 0;
+            int exceptionCounter = 0;
+            int totalCounter = 0;
+            var lockObj = new object();
+            var cts = new CancellationTokenSource();
+            var cancellationToken = cts.Token;
+
+            var func = new Func<Task>(async () =>
+            {
+                using (var reg = new AsyncRegister<int>(handlerList, cancellationToken))
+                {
+                    try
+                    {
+                        while (true)
+                        {
+                            var ret = await reg.WaitSetting();
+                            lock (lockObj)
+                            {
+                                totalCounter += ret.Arg;
+                                resultCounter += 1;
+                            }
+                        }
+                    }
+                    catch (OperationCanceledException)
+                    {
+                        exceptionCounter += 1;
+                    }
+                }
+            });
+
+            var task1 = func();
+
+            var fiber = new PoolFiber();
+            await handlerList.PublishSequentialAsync(1, fiber);
+            await handlerList.PublishSequentialAsync(2, fiber);
+            await handlerList.PublishSequentialAsync(3, fiber);
+            cts.Cancel();
+
+            await task1;
+            Assert.AreEqual(3, resultCounter);
+            Assert.AreEqual(6, totalCounter);
+            Assert.AreEqual(1, exceptionCounter);
+        }
+
         [Test]
         public async Task WaitingOfProcessedFlagEventArgs()
         {
@@ -109,12 +192,11 @@ public async Task WaitingOfProcessedFlagEventArgs()
             var cts = new CancellationTokenSource();
             Func<int, Task> func = async (threshold) =>
             {
-                var reg = new AsyncRegister<int>(driver);
-                try
+                using (var reg = new AsyncRegister<int>(driver, cts.Token))
                 {
                     while (true)
                     {
-                        var e = await reg.WaitSetting(cts.Token);
+                        var e = await reg.WaitSetting();
                         if (e.Arg < threshold)
                         {
                             e.Processed = false;
@@ -129,10 +211,6 @@ public async Task WaitingOfProcessedFlagEventArgs()
                         }
                     }
                 }
-                finally
-                {
-                    reg.Dispose();
-                }
             };
 
             var task1 = func(7);

Original file line number	Diff line number	Diff line change
`@@ -17,13 +17,16 @@ public class AsyncRegister : IDisposable`
`17`	`17`	`private readonly SemaphoreSlim _notifierSet = new SemaphoreSlim(0);`
`18`	`18`	`private readonly SemaphoreSlim _notifierClear = new SemaphoreSlim(0);`
`19`	`19`	`private bool _reading;`
	`20`	`+ private CancellationToken _cancellationToken;`
`20`	`21`
`21`	`22`	`/// <summary>`
`22`	`23`	`/// Subscribe a task list.`
`23`	`24`	`/// </summary>`
`24`	`25`	`/// <param name="taskList"></param>`
`25`		`- public AsyncRegister(ISequentialTaskListRegistry taskList)`
	`26`	`+ /// <param name="cancellationToken"></param>`
	`27`	`+ public AsyncRegister(ISequentialTaskListRegistry taskList, CancellationToken cancellationToken = default)`
`26`	`28`	`{`
	`29`	`+ _cancellationToken = cancellationToken;`
`27`	`30`	`_subscription = taskList.Add(async () =>`
`28`	`31`	`{`
`29`	`32`	`await SetFlagAndWaitClearing().ConfigureAwait(false);`
`@@ -43,6 +46,10 @@ public async Task SetFlagAndWaitClearing()`
`43`	`46`	`{`
`44`	`47`	`return;`
`45`	`48`	`}`
	`49`	`+ if (_cancellationToken.IsCancellationRequested)`
	`50`	`+ {`
	`51`	`+ return;`
	`52`	`+ }`
`46`	`53`	`if (_hasValue)`
`47`	`54`	`{`
`48`	`55`	`throw new InvalidOperationException();`
`@@ -60,24 +67,34 @@ public async Task SetFlagAndWaitClearing()`
`60`	`67`	`_notifierSet.Release(1);`
`61`	`68`	`}`
`62`	`69`
`63`		`- await _notifierClear.WaitAsync().ConfigureAwait(false);`
	`70`	`+ try`
	`71`	`+ {`
	`72`	`+ await _notifierClear.WaitAsync(_cancellationToken).ConfigureAwait(false);`
	`73`	`+ }`
	`74`	`+ catch (OperationCanceledException)`
	`75`	`+ {`
	`76`	`+ return;`
	`77`	`+ }`
`64`	`78`	`}`
`65`	`79`
`66`	`80`	`/// <summary>`
`67`	`81`	`/// Wait for the flag to be set.`
`68`	`82`	`/// </summary>`
`69`		`- /// <param name="token"></param>`
`70`	`83`	`/// <returns></returns>`
`71`	`84`	`/// <exception cref="ObjectDisposedException"></exception>`
`72`	`85`	`/// <exception cref="OperationCanceledException"></exception>`
`73`		`- public async Task WaitSetting(CancellationToken token = default)`
	`86`	`+ public async Task WaitSetting()`
`74`	`87`	`{`
`75`	`88`	`lock (_lockObj)`
`76`	`89`	`{`
`77`	`90`	`if (_isDisposed)`
`78`	`91`	`{`
`79`	`92`	`throw new ObjectDisposedException(GetType().FullName);`
`80`	`93`	`}`
	`94`	`+ if (_cancellationToken.IsCancellationRequested)`
	`95`	`+ {`
	`96`	`+ throw new OperationCanceledException();`
	`97`	`+ }`
`81`	`98`	`if (_hasValue && _reading)`
`82`	`99`	`{`
`83`	`100`	`_hasValue = false;`
`@@ -86,7 +103,7 @@ public async Task WaitSetting(CancellationToken token = default)`
`86`	`103`	`}`
`87`	`104`	`}`
`88`	`105`
`89`		`- await _notifierSet.WaitAsync(token).ConfigureAwait(false);`
	`106`	`+ await _notifierSet.WaitAsync(_cancellationToken).ConfigureAwait(false);`
`90`	`107`
`91`	`108`	`lock (_lockObj)`
`92`	`109`	`{`
Original file line number	Diff line number	Diff line change
`@@ -19,13 +19,16 @@ public class AsyncRegister<T> : IDisposable`
`19`	`19`	`private readonly SemaphoreSlim _notifierSet = new SemaphoreSlim(0);`
`20`	`20`	`private readonly SemaphoreSlim _notifierClear = new SemaphoreSlim(0);`
`21`	`21`	`private bool _reading;`
	`22`	`+ private CancellationToken _cancellationToken;`
`22`	`23`
`23`	`24`	`/// <summary>`
`24`	`25`	`/// Subscribe a handler list.`
`25`	`26`	`/// </summary>`
`26`	`27`	`/// <param name="handlerList"></param>`
`27`		`- public AsyncRegister(ISequentialHandlerListRegistry<T> handlerList)`
	`28`	`+ /// <param name="cancellationToken"></param>`
	`29`	`+ public AsyncRegister(ISequentialHandlerListRegistry<T> handlerList, CancellationToken cancellationToken = default)`
`28`	`30`	`{`
	`31`	`+ _cancellationToken = cancellationToken;`
`29`	`32`	`_subscription = handlerList.Add((arg) => SetValueAndWaitClearing(arg));`
`30`	`33`	`}`
`31`	`34`
`@@ -43,6 +46,10 @@ async Task<bool> SetValueAndWaitClearing(T newValue)`
`43`	`46`	`{`
`44`	`47`	`return false;`
`45`	`48`	`}`
	`49`	`+ if (_cancellationToken.IsCancellationRequested)`
	`50`	`+ {`
	`51`	`+ return false;`
	`52`	`+ }`
`46`	`53`	`if (_hasValue)`
`47`	`54`	`{`
`48`	`55`	`throw new InvalidOperationException();`
`@@ -62,25 +69,35 @@ async Task<bool> SetValueAndWaitClearing(T newValue)`
`62`	`69`	`_notifierSet.Release(1);`
`63`	`70`	`}`
`64`	`71`
`65`		`- await _notifierClear.WaitAsync().ConfigureAwait(false);`
	`72`	`+ try`
	`73`	`+ {`
	`74`	`+ await _notifierClear.WaitAsync(_cancellationToken).ConfigureAwait(false);`
	`75`	`+ }`
	`76`	`+ catch (OperationCanceledException)`
	`77`	`+ {`
	`78`	`+ return false;`
	`79`	`+ }`
`66`	`80`	`return _currentValue.Processed;`
`67`	`81`	`}`
`68`	`82`
`69`	`83`	`/// <summary>`
`70`	`84`	`/// Wait for the value to be set.`
`71`	`85`	`/// </summary>`
`72`		`- /// <param name="token"></param>`
`73`	`86`	`/// <returns></returns>`
`74`	`87`	`/// <exception cref="ObjectDisposedException"></exception>`
`75`	`88`	`/// <exception cref="OperationCanceledException"></exception>`
`76`		`- public async Task<ProcessedFlagEventArgs<T>> WaitSetting(CancellationToken token = default)`
	`89`	`+ public async Task<ProcessedFlagEventArgs<T>> WaitSetting()`
`77`	`90`	`{`
`78`	`91`	`lock (_lockObj)`
`79`	`92`	`{`
`80`	`93`	`if (_isDisposed)`
`81`	`94`	`{`
`82`	`95`	`throw new ObjectDisposedException(GetType().FullName);`
`83`	`96`	`}`
	`97`	`+ if (_cancellationToken.IsCancellationRequested)`
	`98`	`+ {`
	`99`	`+ throw new OperationCanceledException();`
	`100`	`+ }`
`84`	`101`	`if (_hasValue && _reading)`
`85`	`102`	`{`
`86`	`103`	`_currentValue.Arg = default;`
`@@ -90,7 +107,7 @@ public async Task<ProcessedFlagEventArgs<T>> WaitSetting(CancellationToken token`
`90`	`107`	`}`
`91`	`108`	`}`
`92`	`109`
`93`		`- await _notifierSet.WaitAsync(token).ConfigureAwait(false);`
	`110`	`+ await _notifierSet.WaitAsync(_cancellationToken).ConfigureAwait(false);`
`94`	`111`
`95`	`112`	`lock (_lockObj)`
`96`	`113`	`{`