Fix OperationCancelledException and add IAsyncEnumerable to wait without thread blocking (#1562)

* Fixed Task Scheduler operation cancelled error caused by Orphaned RunCleanup task on RequestQueue not being awaited on dispose

* Added async disposable interface to various components

* Added incorrect early marking of disposed for DiscordSocketApi client

* Added general task canceled exception catch to cleanup task

* Fix merge errors

Co-authored-by: Quin Lynch <[email protected]>
Co-authored-by: quin lynch <[email protected]>

Author: Jagreaper

src/Discord.Net.Core/IDiscordClient.cs

@@ -8,7 +8,7 @@ namespace Discord
     /// <summary>
     ///     Represents a generic Discord client.
     /// </summary>
-    public interface IDiscordClient : IDisposable
+    public interface IDiscordClient : IDisposable, IAsyncDisposable
     {
         /// <summary>
         ///     Gets the current state of connection.

src/Discord.Net.Rest/BaseDiscordClient.cs

@@ -149,9 +149,24 @@ internal virtual void Dispose(bool disposing)
                 _isDisposed = true;
             }
         }
+
+        internal virtual async ValueTask DisposeAsync(bool disposing)
+        {
+            if (!_isDisposed)
+            {
+#pragma warning disable IDISP007
+                await ApiClient.DisposeAsync().ConfigureAwait(false);
+#pragma warning restore IDISP007
+                _stateLock?.Dispose();
+                _isDisposed = true;
+            }
+        }
+
         /// <inheritdoc />
         public void Dispose() => Dispose(true);
 
+        public ValueTask DisposeAsync() => DisposeAsync(true);
+
         /// <inheritdoc />
         public Task<int> GetRecommendedShardCountAsync(RequestOptions options = null)
             => ClientHelper.GetRecommendShardCountAsync(this, options);

src/Discord.Net.Rest/DiscordRestApiClient.cs

@@ -21,7 +21,7 @@
 
 namespace Discord.API
 {
-    internal class DiscordRestApiClient : IDisposable
+    internal class DiscordRestApiClient : IDisposable, IAsyncDisposable
     {
         #region DiscordRestApiClient
         private static readonly ConcurrentDictionary<string, Func<BucketIds, BucketId>> _bucketIdGenerators = new ConcurrentDictionary<string, Func<BucketIds, BucketId>>();
@@ -97,8 +97,29 @@ internal virtual void Dispose(bool disposing)
                 _isDisposed = true;
             }
         }
+
+        internal virtual async ValueTask DisposeAsync(bool disposing)
+        {
+            if (!_isDisposed)
+            {
+                if (disposing)
+                {
+                    _loginCancelToken?.Dispose();
+                    RestClient?.Dispose();
+
+                    if (!(RequestQueue is null))
+                        await RequestQueue.DisposeAsync().ConfigureAwait(false);
+
+                    _stateLock?.Dispose();
+                }
+                _isDisposed = true;
+            }
+        }
+
         public void Dispose() => Dispose(true);
 
+        public ValueTask DisposeAsync() => DisposeAsync(true);
+
         public async Task LoginAsync(TokenType tokenType, string token, RequestOptions options = null)
         {
             await _stateLock.WaitAsync().ConfigureAwait(false);

src/Discord.Net.Rest/DiscordRestClient.cs

@@ -47,6 +47,14 @@ internal override void Dispose(bool disposing)
             base.Dispose(disposing);
         }
 
+        internal override async ValueTask DisposeAsync(bool disposing)
+        {
+            if (disposing)
+                await ApiClient.DisposeAsync().ConfigureAwait(false);
+
+            base.Dispose(disposing);
+        }
+
         /// <inheritdoc />
         internal override async Task OnLoginAsync(TokenType tokenType, string token)
         {

src/Discord.Net.Rest/Net/Queue/RequestQueue.cs

@@ -1,3 +1,4 @@
+using Newtonsoft.Json.Bson;
 using System;
 using System.Collections.Concurrent;
 #if DEBUG_LIMITS
@@ -10,7 +11,7 @@
 
 namespace Discord.Net.Queue
 {
-    internal class RequestQueue : IDisposable
+    internal class RequestQueue : IDisposable, IAsyncDisposable
     {
         public event Func<BucketId, RateLimitInfo?, string, Task> RateLimitTriggered;
 
@@ -187,13 +188,31 @@ private async Task RunCleanup()
                     await Task.Delay(60000, _cancelTokenSource.Token).ConfigureAwait(false); //Runs each minute
                 }
             }
-            catch (OperationCanceledException) { }
+            catch (TaskCanceledException) { }
             catch (ObjectDisposedException) { }
         }
 
         public void Dispose()
         {
-            _cancelTokenSource?.Dispose();
+            if (!(_cancelTokenSource is null))
+            {
+                _cancelTokenSource.Cancel();
+                _cancelTokenSource.Dispose();
+                _cleanupTask.GetAwaiter().GetResult();
+            }
+            _tokenLock?.Dispose();
+            _clearToken?.Dispose();
+            _requestCancelTokenSource?.Dispose();
+        }
+
+        public async ValueTask DisposeAsync()
+        {
+            if (!(_cancelTokenSource is null))
+            {
+                _cancelTokenSource.Cancel();
+                _cancelTokenSource.Dispose();
+                await _cleanupTask.ConfigureAwait(false);
+            }
             _tokenLock?.Dispose();
             _clearToken?.Dispose();
             _requestCancelTokenSource?.Dispose();

src/Discord.Net.WebSocket/DiscordShardedClient.cs

@@ -568,6 +568,25 @@ internal override void Dispose(bool disposing)
 
             base.Dispose(disposing);
         }
+
+        internal override ValueTask DisposeAsync(bool disposing)
+        {
+            if (!_isDisposed)
+            {
+                if (disposing)
+                {
+                    if (_shards != null)
+                    {
+                        foreach (var client in _shards)
+                            client?.Dispose();
+                    }
+                }
+
+                _isDisposed = true;
+            }
+
+            return base.DisposeAsync(disposing);
+        }
         #endregion
     }
 }

src/Discord.Net.WebSocket/DiscordSocketApiClient.cs

@@ -124,12 +124,39 @@ internal override void Dispose(bool disposing)
                     _decompressor?.Dispose();
                     _compressed?.Dispose();
                 }
-                _isDisposed = true;
             }
 
             base.Dispose(disposing);
         }
 
+#if NETSTANDARD2_1
+        internal override async ValueTask DisposeAsync(bool disposing)
+#else
+        internal override ValueTask DisposeAsync(bool disposing)
+#endif
+        {
+            if (!_isDisposed)
+            {
+                if (disposing)
+                {
+                    _connectCancelToken?.Dispose();
+                    (WebSocketClient as IDisposable)?.Dispose();
+#if NETSTANDARD2_1
+                    if (!(_decompressor is null))
+                        await _decompressor.DisposeAsync().ConfigureAwait(false);
+#else
+                    _decompressor?.Dispose();
+#endif
+                }
+            }
+
+#if NETSTANDARD2_1
+            await base.DisposeAsync(disposing).ConfigureAwait(false);
+#else
+            return base.DisposeAsync(disposing);
+#endif
+        }
+
         public async Task ConnectAsync()
         {
             await _stateLock.WaitAsync().ConfigureAwait(false);

src/Discord.Net.WebSocket/DiscordSocketClient.cs

@@ -213,6 +213,27 @@ internal override void Dispose(bool disposing)
             base.Dispose(disposing);
         }
 
+
+        internal override async ValueTask DisposeAsync(bool disposing)
+        {
+            if (!_isDisposed)
+            {
+                if (disposing)
+                {
+                    await StopAsync().ConfigureAwait(false);
+
+                    if (!(ApiClient is null))
+                        await ApiClient.DisposeAsync().ConfigureAwait(false);
+
+                    _stateLock?.Dispose();
+                }
+                _isDisposed = true;
+            }
+
+            await base.DisposeAsync(disposing).ConfigureAwait(false);
+        }
+
+        /// <inheritdoc />
         internal override async Task OnLoginAsync(TokenType tokenType, string token)
         {
             if (_shardedClient == null && _defaultStickers.Length == 0 && AlwaysDownloadDefaultStickers)