Added shutdown token to CDN connection pool monitor

SteamRE · May 2, 2018 · 7a75710 · 7a75710
1 parent 5425ef7
commit 7a75710
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 2 deletions.
diff --git a/DepotDownloader/CDNClientPool.cs b/DepotDownloader/CDNClientPool.cs
@@ -24,6 +24,7 @@ class CDNClientPool
 
         private AutoResetEvent populatePoolEvent;
         private Task monitorTask;
+        private CancellationTokenSource shutdownToken;
         public CancellationTokenSource ExhaustedToken { get; set; }
 
         public CDNClientPool(Steam3Session steamSession)
@@ -35,15 +36,21 @@ public CDNClientPool(Steam3Session steamSession)
             availableServerEndpoints = new BlockingCollection<CDNClient.Server>();
 
             populatePoolEvent = new AutoResetEvent(true);
+            shutdownToken = new CancellationTokenSource();
 
             monitorTask = Task.Factory.StartNew(ConnectionPoolMonitorAsync).Unwrap();
         }
 
+        public void Shutdown()
+        {
+            shutdownToken.Cancel();
+        }
+
         private async Task<IList<CDNClient.Server>> FetchBootstrapServerListAsync()
         {
             CDNClient bootstrap = new CDNClient(steamSession.steamClient);
 
-            while (true)
+            while (!shutdownToken.IsCancellationRequested)
             {
                 try
                 {
@@ -58,13 +65,15 @@ private async Task<IList<CDNClient.Server>> FetchBootstrapServerListAsync()
                     Console.WriteLine("Failed to retrieve content server list: {0}", ex.Message);
                 }
             }
+
+            return null;
         }
 
         private async Task ConnectionPoolMonitorAsync()
         {
             bool didPopulate = false;
 
-            while(true)
+            while(!shutdownToken.IsCancellationRequested)
             {
                 populatePoolEvent.WaitOne(TimeSpan.FromSeconds(1));
 
@@ -75,6 +84,12 @@ private async Task ConnectionPoolMonitorAsync()
                 {
                     var servers = await FetchBootstrapServerListAsync().ConfigureAwait(false);
 
+                    if (servers == null)
+                    {
+                        ExhaustedToken?.Cancel();
+                        return;
+                    }
+
                     var weightedCdnServers = servers.Select(x =>
                     {
                         int penalty = 0;
@@ -95,6 +110,7 @@ private async Task ConnectionPoolMonitorAsync()
                 else if ( availableServerEndpoints.Count == 0 && !steamSession.steamClient.IsConnected && didPopulate )
                 {
                     ExhaustedToken?.Cancel();
+                    return;
                 }
             }
         }

diff --git a/DepotDownloader/ContentDownloader.cs b/DepotDownloader/ContentDownloader.cs
@@ -363,6 +363,12 @@ public static bool InitializeSteam3( string username, string password )
 
         public static void ShutdownSteam3()
         {
+            if (cdnPool != null)
+            {
+                cdnPool.Shutdown();
+                cdnPool = null;
+            }
+
             if ( steam3 == null )
                 return;