Merge pull request #136 from /issues/134-CdnServersExhausted

tpill90 · Sep 18, 2022 · 2a18cec · 2a18cec
2 parents 9b4527c + 005fbea
commit 2a18cec
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 9 deletions.
diff --git a/SteamPrefill/Extensions/MiscExtensions.cs b/SteamPrefill/Extensions/MiscExtensions.cs
@@ -2,9 +2,9 @@
 {
     public static class MiscExtensions
     {
-        public static ConcurrentQueue<T> ToConcurrentBag<T>(this IOrderedEnumerable<T> list)
+        public static bool Empty<T>(this IEnumerable<T> enumerable)
         {
-            return new ConcurrentQueue<T>(list);
+            return !enumerable.Any();
         }
 
         [SuppressMessage("Microsoft.Security", "CA5350", Justification = "SHA1 is required by Steam")]

diff --git a/SteamPrefill/Handlers/DownloadHandler.cs b/SteamPrefill/Handlers/DownloadHandler.cs
@@ -83,15 +83,16 @@ private async Task<ConcurrentBag<QueuedRequest>> AttemptDownloadAsync(ProgressCo
 
             // Running as many requests as possible in parallel, evenly distributed across 3 cdns
             var cdnServers = _cdnPool.TakeConnections(3).ToList();
+            var connCount = cdnServers.Count;
             await Parallel.ForEachAsync(requestsToDownload, new ParallelOptions { MaxDegreeOfParallelism = 50 }, async (request, _) =>
             {
                 var buffer = new byte[4096];
                 try
                 {
                     var url = ZString.Format("http://{0}/depot/{1}/chunk/{2}", _lancacheAddress, request.DepotId, request.ChunkId);
                     using var requestMessage = new HttpRequestMessage(HttpMethod.Get, url);
-                    // Evenly distributes requests across CDNs
-                    requestMessage.Headers.Host = cdnServers[request.ChunkNum % 3].Host;
+                    // Evenly distributes requests the available CDNs
+                    requestMessage.Headers.Host = cdnServers[request.ChunkNum % connCount].Host;
 
                     var response = await _client.SendAsync(requestMessage, HttpCompletionOption.ResponseHeadersRead);
                     using Stream responseStream = await response.Content.ReadAsStreamAsync();

diff --git a/SteamPrefill/Handlers/Steam/CdnPool.cs b/SteamPrefill/Handlers/Steam/CdnPool.cs
@@ -51,7 +51,7 @@ await _ansiConsole.StatusSpinner().StartAsync("Getting available CDNs", async _
                 }
             });
 
-            if (!_availableServerEndpoints.Any())
+            if (_availableServerEndpoints.Empty())
             {
                 throw new CdnExhaustionException("Unable to get available CDN servers from Steam!");
             }
@@ -67,21 +67,33 @@ await _ansiConsole.StatusSpinner().StartAsync("Getting available CDNs", async _
         /// <exception cref="CdnExhaustionException">If no servers are available for use, this exception will be thrown.</exception>
         public Server TakeConnection()
         {
-            if (!_availableServerEndpoints.Any())
+            if (_availableServerEndpoints.Empty())
             {
                 throw new CdnExhaustionException("Available Steam CDN servers exhausted!  No more servers available to retry!  Try again in a few minutes");
             }
+
             var server = _availableServerEndpoints.First();
             _availableServerEndpoints.RemoveAt(0);
             return server;
         }
 
-        public IEnumerable<Server> TakeConnections(int connectionCount)
+        /// <summary>
+        /// Attempts to take multiple available connections from the pool.
+        /// If the desired number of connections is not available, then the remaining available will be returned.
+        /// Once finished with the connection, it should be returned to the pool using <seealso cref="ReturnConnection"/>
+        /// </summary>
+        /// <param name="targetCount">The desired number of servers to be returned.  If the desired amount is not available, the remaining servers will be returned</param>
+        /// <exception cref="CdnExhaustionException">If no servers are available for use, this exception will be thrown.</exception>
+        public List<Server> TakeConnections(int targetCount)
         {
-            for (int i = 0; i < connectionCount; i++)
+            if (_availableServerEndpoints.Empty())
             {
-                yield return TakeConnection();
+                throw new CdnExhaustionException("Available Steam CDN servers exhausted!  No more servers available to retry!  Try again in a few minutes");
             }
+
+            var connections = _availableServerEndpoints.Take(targetCount).ToList();
+            _availableServerEndpoints.RemoveRange(0, connections.Count);
+            return connections;
         }
 
         /// <summary>