等待执行完成

Author: lindexi

AsyncWorkerCollection/AsyncQueue.cs

@@ -73,11 +73,37 @@ public async Task<T> DequeueAsync(CancellationToken cancellationToken = default)
                 {
                     return item;
                 }
+                else
+                {
+                    lock (_queue)
+                    {
+                        CurrentFinished?.Invoke(this, EventArgs.Empty);
+                    }
+                }
             }
 
             return default;
         }
 
+        private event EventHandler CurrentFinished;
+
+        public async ValueTask WaitForCurrentFinished()
+        {
+            if (_queue.Count == 0)
+            {
+                return;
+            }
+
+            using var currentFinishedTask = new CurrentFinishedTask(this);
+
+            if (_queue.Count == 0)
+            {
+                return;
+            }
+
+            await currentFinishedTask.WaitForCurrentFinished();
+        }
+
         /// <summary>
         /// 主要用来释放锁，让 DequeueAsync 方法返回，解决因为锁让此对象内存不释放
         /// </summary>
@@ -93,5 +119,41 @@ public void Dispose()
         }
 
         private bool _isDisposed;
+
+        class CurrentFinishedTask : IDisposable
+        {
+            public CurrentFinishedTask(AsyncQueue<T> asyncQueue)
+            {
+                _asyncQueue = asyncQueue;
+
+                lock (_asyncQueue)
+                {
+                    _asyncQueue.CurrentFinished += CurrentFinished;
+                }
+            }
+
+            private void CurrentFinished(object? sender, EventArgs e)
+            {
+                _currentFinishedTaskCompletionSource.TrySetResult(true);
+            }
+
+            public async ValueTask WaitForCurrentFinished()
+            {
+                await _currentFinishedTaskCompletionSource.Task;
+            }
+
+            private readonly TaskCompletionSource<bool> _currentFinishedTaskCompletionSource = new TaskCompletionSource<bool>();
+
+            private readonly AsyncQueue<T> _asyncQueue;
+
+            public void Dispose()
+            {
+                lock (_asyncQueue)
+                {
+                    _currentFinishedTaskCompletionSource.TrySetResult(true);
+                    _asyncQueue.CurrentFinished -= CurrentFinished;
+                }
+            }
+        }
     }
 }