等待任务执行完成

Author: lindexi

AsyncWorkerCollection/AsyncQueue.cs

@@ -10,7 +10,7 @@ namespace dotnetCampus.Threading
     /// 提供一个异步的队列。可以使用 await 关键字异步等待出队，当有元素入队的时候，等待就会完成。
     /// </summary>
     /// <typeparam name="T">存入异步队列中的元素类型。</typeparam>
-    public class AsyncQueue<T> : IDisposable
+    public class AsyncQueue<T> : IDisposable, IAsyncDisposable
     {
         private readonly SemaphoreSlim _semaphoreSlim;
         private readonly ConcurrentQueue<T> _queue;
@@ -36,6 +36,8 @@ public AsyncQueue()
         /// <param name="item">要入队的元素。</param>
         public void Enqueue(T item)
         {
+            ThrowIfDisposing();
+
             _queue.Enqueue(item);
             _semaphoreSlim.Release();
         }
@@ -46,12 +48,15 @@ public void Enqueue(T item)
         /// <param name="source">要入队的元素序列。</param>
         public void EnqueueRange(IEnumerable<T> source)
         {
+            ThrowIfDisposing();
+
             var n = 0;
             foreach (var item in source)
             {
                 _queue.Enqueue(item);
                 n++;
             }
+
             _semaphoreSlim.Release(n);
         }
 
@@ -63,6 +68,7 @@ public void EnqueueRange(IEnumerable<T> source)
         /// 由于此方法有返回值，后续方法可能依赖于此返回值，所以如果取消将抛出 <see cref="TaskCanceledException"/>。
         /// </param>
         /// <returns>可以异步等待的队列返回的元素。</returns>
+        /// <exception cref="ObjectDisposedException"></exception>
         public async Task<T> DequeueAsync(CancellationToken cancellationToken = default)
         {
             while (!_isDisposed)
@@ -73,8 +79,19 @@ public async Task<T> DequeueAsync(CancellationToken cancellationToken = default)
                 {
                     return item;
                 }
+                else
+                {
+                    // 没有内容了，此时也准备gg了，那么就 Break 了
+                    if (_isDisposing)
+                    {
+                        _disposeTaskCompletionSource.TrySetResult(true);
+                        break;
+                    }
+                }
             }
 
+            ThrowIfDisposing();
+
             return default;
         }
 
@@ -85,13 +102,44 @@ public void Dispose()
         {
             // 当释放的时候，将通过 _queue 的 Clear 清空内容，而通过 _semaphoreSlim 的释放让 DequeueAsync 释放锁
             // 此时将会在 DequeueAsync 进入 TryDequeue 方法，也许此时依然有开发者在 _queue.Clear() 之后插入元素，但是没关系，我只是需要保证调用 Dispose 之后会让 DequeueAsync 方法返回而已
-            _isDisposed = true;
+            _isDisposing = true;
             _queue.Clear();
             // 释放 DequeueAsync 方法
             _semaphoreSlim.Release(int.MaxValue);
+            _isDisposed = true;
             _semaphoreSlim.Dispose();
         }
 
         private bool _isDisposed;
+
+        /// <summary>
+        /// 等待所有任务执行完成
+        /// </summary>
+        /// <returns></returns>
+        public async ValueTask DisposeAsync()
+        {
+            _disposeTaskCompletionSource = new TaskCompletionSource<bool>();
+            _isDisposing = true;
+
+            // 释放 DequeueAsync 方法
+            _semaphoreSlim.Release(int.MaxValue);
+
+            await _disposeTaskCompletionSource.Task;
+
+            _isDisposed = true;
+            _semaphoreSlim.Dispose();
+        }
+
+        private bool _isDisposing;
+        private TaskCompletionSource<bool> _disposeTaskCompletionSource;
+
+        // 这里忽略线程安全
+        private void ThrowIfDisposing()
+        {
+            if (_isDisposing)
+            {
+                throw new ObjectDisposedException(nameof(AsyncQueue<T>));
+            }
+        }
     }
-}
+}