Merge pull request #229 from carlst99/m_gwResumeUrl

Support Resume Gateway URL
Remora · Aug 10, 2022 · 7f118b6 · 7f118b6
2 parents 0c5c31e + 50fbebf
commit 7f118b6
Show file tree

Hide file tree

Showing 6 changed files with 37 additions and 3 deletions.
diff --git a/Backend/Remora.Discord.API.Abstractions/API/Gateway/Events/ConnectingResuming/IReady.cs b/Backend/Remora.Discord.API.Abstractions/API/Gateway/Events/ConnectingResuming/IReady.cs
@@ -54,6 +54,11 @@ public interface IReady : IGatewayEvent
     /// </summary>
     string SessionID { get; }
 
+    /// <summary>
+    /// Gets the resume gateway URL.
+    /// </summary>
+    string ResumeGatewayUrl { get; }
+
     /// <summary>
     /// Gets the shard information associated with this session.
     /// </summary>

diff --git a/Backend/Remora.Discord.API/API/Gateway/Events/ConnectingResuming/Ready.cs b/Backend/Remora.Discord.API/API/Gateway/Events/ConnectingResuming/Ready.cs
@@ -39,6 +39,7 @@ public record Ready
     IUser User,
     IReadOnlyList<IUnavailableGuild> Guilds,
     string SessionID,
+    string ResumeGatewayUrl,
     Optional<IShardIdentification> Shard,
     IPartialApplication Application
 ) : IReady;
diff --git a/Backend/Remora.Discord.Gateway/DiscordGatewayClient.cs b/Backend/Remora.Discord.Gateway/DiscordGatewayClient.cs
@@ -109,6 +109,11 @@ public class DiscordGatewayClient : IDisposable
     /// </summary>
     private string? _sessionID;
 
+    /// <summary>
+    /// Holds the resume gateway URL.
+    /// </summary>
+    private string? _resumeGatewayUrl;
+
     /// <summary>
     /// Holds the cancellation token source for internal operations.
     /// </summary>
@@ -258,6 +263,7 @@ public async Task<Result> RunAsync(CancellationToken stopRequested)
                     if (withNewSession)
                     {
                         _sessionID = null;
+                        _resumeGatewayUrl = null;
                         _connectionStatus = GatewayConnectionStatus.Disconnected;
                     }
                     else
@@ -295,11 +301,13 @@ public async Task<Result> RunAsync(CancellationToken stopRequested)
         finally
         {
             _sessionID = null;
+            _resumeGatewayUrl = null;
             _connectionStatus = GatewayConnectionStatus.Offline;
         }
 
         // Reconnection is not allowed at this point.
         _sessionID = null;
+        _resumeGatewayUrl = null;
         _connectionStatus = GatewayConnectionStatus.Offline;
 
         return Result.FromSuccess();
@@ -516,7 +524,11 @@ private async Task<Result> RunConnectionIterationAsync(CancellationToken stopReq
                     );
                 }
 
-                var gatewayEndpoint = $"{getGatewayEndpoint.Entity.Url}?v={(int)DiscordAPIVersion.V10}&encoding=json";
+                var gatewayEndpointUrl = _resumeGatewayUrl is not null && _isSessionResumable
+                    ? _resumeGatewayUrl
+                    : getGatewayEndpoint.Entity.Url;
+
+                var gatewayEndpoint = $"{gatewayEndpointUrl}?v={(int)DiscordAPIVersion.V10}&encoding=json";
                 if (!Uri.TryCreate(gatewayEndpoint, UriKind.Absolute, out var gatewayUri))
                 {
                     return new GatewayError
@@ -756,6 +768,7 @@ private async Task<Result> CreateNewSessionAsync(CancellationToken ct = default)
                         }
 
                         _sessionID = ready.Data.SessionID;
+                        _resumeGatewayUrl = ready.Data.ResumeGatewayUrl;
 
                         return Result.FromSuccess();
                     }

diff --git a/Tests/Remora.Discord.Gateway.Tests/Constants.cs b/Tests/Remora.Discord.Gateway.Tests/Constants.cs
@@ -51,4 +51,9 @@ public static class Constants
     /// Gets the default mocked session ID.
     /// </summary>
     public static string MockSessionID => "mock-session";
+
+    /// <summary>
+    /// Gets the default mocked resume gateway URL.
+    /// </summary>
+    public static string MockResumeGatewayUrl => "wss://us-east1-b.gateway.discord.gg";
 }
diff --git a/Tests/Remora.Discord.Gateway.Tests/Tests/DiscordGatewayClientTests.cs b/Tests/Remora.Discord.Gateway.Tests/Tests/DiscordGatewayClientTests.cs
@@ -96,6 +96,7 @@ public async Task CanConnectAsync()
                             Constants.BotUser,
                             new List<IUnavailableGuild>(),
                             "mock-session",
+                            Constants.MockResumeGatewayUrl,
                             default,
                             new PartialApplication()
                         )
@@ -156,13 +157,14 @@ public async Task CanReconnectAndResumeAsync()
                             Constants.BotUser,
                             new List<IUnavailableGuild>(),
                             Constants.MockSessionID,
+                            Constants.MockResumeGatewayUrl,
                             default,
                             new PartialApplication()
                         )
                     )
                     .Send<Reconnect>()
                     .ExpectDisconnect()
-                    .ExpectConnection(new Uri($"wss://gateway.discord.gg/?v={(int)DiscordAPIVersion.V10}&encoding=json"))
+                    .ExpectConnection(new Uri($"{Constants.MockResumeGatewayUrl}?v={(int)DiscordAPIVersion.V10}&encoding=json"))
                     .Send(new Hello(TimeSpan.FromMilliseconds(200)))
                     .Expect<Resume>
                     (
@@ -229,13 +231,14 @@ public async Task CanReconnectAndCreateNewSessionAsync()
                             Constants.BotUser,
                             new List<IUnavailableGuild>(),
                             Constants.MockSessionID,
+                            Constants.MockResumeGatewayUrl,
                             default,
                             new PartialApplication()
                         )
                     )
                     .Send<Reconnect>()
                     .ExpectDisconnect()
-                    .ExpectConnection(new Uri($"wss://gateway.discord.gg/?v={(int)DiscordAPIVersion.V10}&encoding=json"))
+                    .ExpectConnection(new Uri($"{Constants.MockResumeGatewayUrl}?v={(int)DiscordAPIVersion.V10}&encoding=json"))
                     .Send(new Hello(TimeSpan.FromMilliseconds(200)))
                     .Expect<Resume>
                     (
@@ -262,6 +265,7 @@ public async Task CanReconnectAndCreateNewSessionAsync()
                             Constants.BotUser,
                             new List<IUnavailableGuild>(),
                             Constants.MockSessionID,
+                            Constants.MockResumeGatewayUrl,
                             default,
                             new PartialApplication()
                         )
@@ -337,6 +341,7 @@ public async Task CanRetrySessionCreationIfInvalidatedAsync()
                             Constants.BotUser,
                             new List<IUnavailableGuild>(),
                             Constants.MockSessionID,
+                            Constants.MockResumeGatewayUrl,
                             default,
                             new PartialApplication()
                         )
@@ -412,6 +417,7 @@ public async Task CanRetrySessionCreationIfReconnectRequestedAsync()
                             Constants.BotUser,
                             new List<IUnavailableGuild>(),
                             Constants.MockSessionID,
+                            Constants.MockResumeGatewayUrl,
                             default,
                             new PartialApplication()
                         )
@@ -472,6 +478,7 @@ public async Task CanReconnectAfterExceptionAsync()
                             Constants.BotUser,
                             Array.Empty<IUnavailableGuild>(),
                             Constants.MockSessionID,
+                            Constants.MockResumeGatewayUrl,
                             default,
                             new PartialApplication()
                         )
@@ -495,6 +502,7 @@ public async Task CanReconnectAfterExceptionAsync()
                             Constants.BotUser,
                             Array.Empty<IUnavailableGuild>(),
                             Constants.MockSessionID,
+                            Constants.MockResumeGatewayUrl,
                             default,
                             new PartialApplication()
                         )
@@ -518,6 +526,7 @@ public async Task CanReconnectAfterExceptionAsync()
                             Constants.BotUser,
                             Array.Empty<IUnavailableGuild>(),
                             Constants.MockSessionID,
+                            Constants.MockResumeGatewayUrl,
                             default,
                             new PartialApplication()
                         )

diff --git a/Tests/Remora.Discord.Tests/Samples/Gateway/Events/READY/READY.json b/Tests/Remora.Discord.Tests/Samples/Gateway/Events/READY/READY.json
@@ -14,6 +14,7 @@
       "avatar": null
     },
     "session_id": "none",
+    "resume_gateway_url": "wss://gateway.discord.gg/",
     "guilds": [],
     "application": {
       "id": "999999999999999999",