From baaa4349940a17146f22af54870bc16b6378a671 Mon Sep 17 00:00:00 2001
From: Paulo <pnmanjos@hotmail.com>
Date: Tue, 17 Nov 2020 23:51:25 -0300
Subject: [PATCH] Wait identify before opening the connection

---
 .../Net/Queue/RequestQueue.cs                 | 45 ++++++++++---------
 .../DiscordShardedClient.cs                   |  4 +-
 .../DiscordSocketClient.cs                    |  2 +
 3 files changed, 28 insertions(+), 23 deletions(-)

diff --git a/src/Discord.Net.Rest/Net/Queue/RequestQueue.cs b/src/Discord.Net.Rest/Net/Queue/RequestQueue.cs
index 8192f6a41..488e2c5c8 100644
--- a/src/Discord.Net.Rest/Net/Queue/RequestQueue.cs
+++ b/src/Discord.Net.Rest/Net/Queue/RequestQueue.cs
@@ -136,27 +136,6 @@ namespace Discord.Net.Queue
             var requestBucket = GatewayBucket.Get(request.Options.BucketId);
             if (requestBucket.Type == GatewayBucketType.Unbucketed)
                 return;
-            
-            //Identify is per-account so we won't trigger global until we can actually go for it
-            if (requestBucket.Type == GatewayBucketType.Identify)
-            {
-                if (_masterIdentifySemaphore == null)
-                    throw new InvalidOperationException("Not a RequestQueue with WebSocket data.");
-
-                if (_identifySemaphore == null)
-                    await _masterIdentifySemaphore.WaitAsync(request.CancelToken);
-                else
-                {
-                    bool master;
-                    while (!(master = _masterIdentifySemaphore.Wait(0)) && !_identifySemaphore.Wait(0)) //To not block the thread
-                        await Task.Delay(100, request.CancelToken);
-                    if (master && _identifySemaphoreMaxConcurrency > 1)
-                        _identifySemaphore.Release(_identifySemaphoreMaxConcurrency - 1);
-                }
-#if DEBUG_LIMITS
-                Debug.WriteLine($"[{id}] Acquired identify ticket");
-#endif
-            }
 
             //It's not a global request, so need to remove one from global (per-session)
             var globalBucketType = GatewayBucket.Get(GatewayBucketType.Unbucketed);
@@ -179,6 +158,30 @@ namespace Discord.Net.Queue
 #endif
         }
 
+        public async Task AcquireIdentifyTicket(CancellationToken cancellationToken)
+        {
+            try
+            {
+                if (_masterIdentifySemaphore == null)
+                    throw new InvalidOperationException("Not a RequestQueue with WebSocket data.");
+
+                if (_identifySemaphore == null)
+                    await _masterIdentifySemaphore.WaitAsync(cancellationToken);
+                else
+                {
+                    bool master;
+                    while (!(master = _masterIdentifySemaphore.Wait(0)) && !_identifySemaphore.Wait(0)) //To not block the thread
+                        await Task.Delay(100, cancellationToken);
+                    if (master && _identifySemaphoreMaxConcurrency > 1)
+                        _identifySemaphore.Release(_identifySemaphoreMaxConcurrency - 1);
+                }
+#if DEBUG_LIMITS
+                Debug.WriteLine($"[{id}] Acquired identify ticket");
+#endif
+            }
+            catch(OperationCanceledException) { }
+        }
+
         private RequestBucket GetOrCreateBucket(RequestOptions options, IRequest request)
         {
             var bucketId = options.BucketId;
diff --git a/src/Discord.Net.WebSocket/DiscordShardedClient.cs b/src/Discord.Net.WebSocket/DiscordShardedClient.cs
index 662f67e40..6c98d9e43 100644
--- a/src/Discord.Net.WebSocket/DiscordShardedClient.cs
+++ b/src/Discord.Net.WebSocket/DiscordShardedClient.cs
@@ -86,7 +86,7 @@ namespace Discord.WebSocket
                     _shardIdsToIndex.Add(_shardIds[i], i);
                     var newConfig = config.Clone();
                     newConfig.ShardId = _shardIds[i];
-                    _shards[i] = new DiscordSocketClient(newConfig, i != 0 ? _shards[0] : null, masterIdentifySemaphore, config.IdentifyMaxConcurrency > 1 ? null : identifySemaphores[i / config.IdentifyMaxConcurrency], config.IdentifyMaxConcurrency);
+                    _shards[i] = new DiscordSocketClient(newConfig, i != 0 ? _shards[0] : null, masterIdentifySemaphore, config.IdentifyMaxConcurrency == 1 ? null : identifySemaphores[i / config.IdentifyMaxConcurrency], config.IdentifyMaxConcurrency);
                     RegisterEvents(_shards[i], i == 0);
                 }
             }
@@ -120,7 +120,7 @@ namespace Discord.WebSocket
                     var newConfig = _baseConfig.Clone();
                     newConfig.ShardId = _shardIds[i];
                     newConfig.TotalShards = _totalShards;
-                    _shards[i] = new DiscordSocketClient(newConfig, i != 0 ? _shards[0] : null, masterIdentifySemaphore, maxConcurrency > 1 ? null : identifySemaphores[i / maxConcurrency], maxConcurrency);
+                    _shards[i] = new DiscordSocketClient(newConfig, i != 0 ? _shards[0] : null, masterIdentifySemaphore, maxConcurrency == 1 ? null : identifySemaphores[i / maxConcurrency], maxConcurrency);
                     RegisterEvents(_shards[i], i == 0);
                 }
             }
diff --git a/src/Discord.Net.WebSocket/DiscordSocketClient.cs b/src/Discord.Net.WebSocket/DiscordSocketClient.cs
index cd83699a7..77f051f74 100644
--- a/src/Discord.Net.WebSocket/DiscordSocketClient.cs
+++ b/src/Discord.Net.WebSocket/DiscordSocketClient.cs
@@ -228,6 +228,8 @@ namespace Discord.WebSocket
 
         private async Task OnConnectingAsync()
         {
+            await ApiClient.RequestQueue.AcquireIdentifyTicket(_connection.CancelToken);
+
             await _gatewayLogger.DebugAsync("Connecting ApiClient").ConfigureAwait(false);
             await ApiClient.ConnectAsync().ConfigureAwait(false);