add timeout param

justcoding121 · justcoding121 · commit e8522a5e66aa · 2018-08-09T22:44:51.000-04:00
diff --git a/Advanced.Algorithms/Distributed/AsyncQueue.cs b/Advanced.Algorithms/Distributed/AsyncQueue.cs
@@ -16,18 +16,20 @@ public class AsyncQueue<T>
         private readonly Queue<TaskCompletionSource<T>> consumerQueue = new Queue<TaskCompletionSource<T>>();
         private SemaphoreSlim consumerQueueLock = new SemaphoreSlim(1);
 
+        public int Count => queue.Count;
+
         /// <summary>
         ///     Supports multi-threaded producers.
         ///     Time complexity: O(1).
         /// </summary>
-        public async Task EnqueueAsync(T value, CancellationToken taskCancellationToken = default(CancellationToken))
+        public async Task EnqueueAsync(T value, int millisecondsTimeout = int.MaxValue, CancellationToken taskCancellationToken = default(CancellationToken))
         {
-            await consumerQueueLock.WaitAsync(taskCancellationToken);
+            await consumerQueueLock.WaitAsync(millisecondsTimeout, taskCancellationToken);
 
             if(consumerQueue.Count > 0)
             {
                 var consumer = consumerQueue.Dequeue();
-                consumer.SetResult(value);
+                consumer.TrySetResult(value);
             }
             else
             {
@@ -41,22 +43,29 @@ public class AsyncQueue<T>
         ///      Supports multi-threaded consumers.
         ///      Time complexity: O(1).
         /// </summary>
-        public async Task<T> DequeueAsync(CancellationToken taskCancellationToken = default(CancellationToken))
+        public async Task<T> DequeueAsync(int millisecondsTimeout = int.MaxValue, CancellationToken taskCancellationToken = default(CancellationToken))
         {
-            await consumerQueueLock.WaitAsync(taskCancellationToken);
+            await consumerQueueLock.WaitAsync(millisecondsTimeout, taskCancellationToken);
+
+            TaskCompletionSource<T> consumer;
+
+            try
+            {
+                if (queue.Count > 0)
+                {
+                    var result = queue.Dequeue();
+                    consumerQueueLock.Release();
+                    return result;
+                }
 
-            if (queue.Count > 0)
+                consumer = new TaskCompletionSource<T>();
+                taskCancellationToken.Register(() => consumer.TrySetCanceled());
+                consumerQueue.Enqueue(consumer);
+            }
+            finally
             {
-                var result = queue.Dequeue();
                 consumerQueueLock.Release();
-                return result;
             }
-          
-            var consumer = new TaskCompletionSource<T>();
-            taskCancellationToken.Register(() => consumer.TrySetCanceled());
-            consumerQueue.Enqueue(consumer);
-            
-            consumerQueueLock.Release();
 
             return await consumer.Task;
 
