Update DownloadDepotChunkAsync to use array pool

DepotDownloader/ContentDownloader.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Buffers;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.IO;
@@ -1199,80 +1200,91 @@ namespace DepotDownloader
                 UncompressedLength = chunk.UncompressedLength
             };
 
-            DepotChunk chunkData = null;
+            var written = 0;
+            var chunkBuffer = ArrayPool<byte>.Shared.Rent((int)data.UncompressedLength);
 
-            do
+            try
             {
-                cts.Token.ThrowIfCancellationRequested();
+                do
+                {
+                    cts.Token.ThrowIfCancellationRequested();
 
-                Server connection = null;
+                    Server connection = null;
 
-                try
+                    try
-                {
+                    {
-                    connection = cdnPool.GetConnection(cts.Token);
+                        connection = cdnPool.GetConnection(cts.Token);
 
-                    DebugLog.WriteLine("ContentDownloader", "Downloading chunk {0} from {1} with {2}", chunkID, connection, cdnPool.ProxyServer != null ? cdnPool.ProxyServer : "no proxy");
+                        DebugLog.WriteLine("ContentDownloader", "Downloading chunk {0} from {1} with {2}", chunkID, connection, cdnPool.ProxyServer != null ? cdnPool.ProxyServer : "no proxy");
-                    chunkData = await cdnPool.CDNClient.DownloadDepotChunkAsync(
+                        written = await cdnPool.CDNClient.DownloadDepotChunkAsync(
-                        depot.DepotId,
+                            depot.DepotId,
-                        data,
+                            data,
-                        connection,
+                            connection,
-                        depot.DepotKey,
+                            chunkBuffer,
-                        cdnPool.ProxyServer).ConfigureAwait(false);
+                            depot.DepotKey,
+                            cdnPool.ProxyServer).ConfigureAwait(false);
 
-                    cdnPool.ReturnConnection(connection);
+                        cdnPool.ReturnConnection(connection);
-                }
+
-                catch (TaskCanceledException)
+                        break;
-                {
+                    }
-                    Console.WriteLine("Connection timeout downloading chunk {0}", chunkID);
+                    catch (TaskCanceledException)
-                }
+                    {
-                catch (SteamKitWebRequestException e)
+                        Console.WriteLine("Connection timeout downloading chunk {0}", chunkID);
-                {
+                    }
-                    cdnPool.ReturnBrokenConnection(connection);
+                    catch (SteamKitWebRequestException e)
+                    {
+                        cdnPool.ReturnBrokenConnection(connection);
+
+                        if (e.StatusCode == HttpStatusCode.Unauthorized || e.StatusCode == HttpStatusCode.Forbidden)
+                        {
+                            Console.WriteLine("Encountered {1} for chunk {0}. Aborting.", chunkID, (int)e.StatusCode);
+                            break;
+                        }
 
-                    if (e.StatusCode == HttpStatusCode.Unauthorized || e.StatusCode == HttpStatusCode.Forbidden)
+                        Console.WriteLine("Encountered error downloading chunk {0}: {1}", chunkID, e.StatusCode);
+                    }
+                    catch (OperationCanceledException)
                     {
-                        Console.WriteLine("Encountered {1} for chunk {0}. Aborting.", chunkID, (int)e.StatusCode);
                         break;
                     }
+                    catch (Exception e)
+                    {
+                        cdnPool.ReturnBrokenConnection(connection);
+                        Console.WriteLine("Encountered unexpected error downloading chunk {0}: {1}", chunkID, e.Message);
+                    }
+                } while (written == 0);
 
-                    Console.WriteLine("Encountered error downloading chunk {0}: {1}", chunkID, e.StatusCode);
+                if (written == 0)
-                }
-                catch (OperationCanceledException)
-                {
-                    break;
-                }
-                catch (Exception e)
                 {
-                    cdnPool.ReturnBrokenConnection(connection);
+                    Console.WriteLine("Failed to find any server with chunk {0} for depot {1}. Aborting.", chunkID, depot.DepotId);
-                    Console.WriteLine("Encountered unexpected error downloading chunk {0}: {1}", chunkID, e.Message);
+                    cts.Cancel();
                 }
-            } while (chunkData == null);
 
-            if (chunkData == null)
+                // Throw the cancellation exception if requested so that this task is marked failed
-            {
+                cts.Token.ThrowIfCancellationRequested();
-                Console.WriteLine("Failed to find any server with chunk {0} for depot {1}. Aborting.", chunkID, depot.DepotId);
-                cts.Cancel();
-            }
 
-            // Throw the cancellation exception if requested so that this task is marked failed
+                try
-            cts.Token.ThrowIfCancellationRequested();
+                {
+                    await fileStreamData.fileLock.WaitAsync().ConfigureAwait(false);
 
-            try
+                    if (fileStreamData.fileStream == null)
-            {
+                    {
-                await fileStreamData.fileLock.WaitAsync().ConfigureAwait(false);
+                        var fileFinalPath = Path.Combine(depot.InstallDir, file.FileName);
+                        fileStreamData.fileStream = File.Open(fileFinalPath, FileMode.Open);
+                    }
 
-                if (fileStreamData.fileStream == null)
+                    fileStreamData.fileStream.Seek((long)data.Offset, SeekOrigin.Begin);
+                    await fileStreamData.fileStream.WriteAsync(chunkBuffer.AsMemory(0, written), cts.Token);
+                }
+                finally
                 {
-                    var fileFinalPath = Path.Combine(depot.InstallDir, file.FileName);
+                    fileStreamData.fileLock.Release();
-                    fileStreamData.fileStream = File.Open(fileFinalPath, FileMode.Open);
                 }
-
-                fileStreamData.fileStream.Seek((long)chunkData.ChunkInfo.Offset, SeekOrigin.Begin);
-                await fileStreamData.fileStream.WriteAsync(chunkData.Data.AsMemory(0, chunkData.Data.Length), cts.Token);
             }
             finally
             {
-                fileStreamData.fileLock.Release();
+                ArrayPool<byte>.Shared.Return(chunkBuffer);
             }
 
             var remainingChunks = Interlocked.Decrement(ref fileStreamData.chunksToDownload);
@@ -1285,7 +1297,7 @@ namespace DepotDownloader
             ulong sizeDownloaded = 0;
             lock (depotDownloadCounter)
             {
-                sizeDownloaded = depotDownloadCounter.sizeDownloaded + (ulong)chunkData.Data.Length;
+                sizeDownloaded = depotDownloadCounter.sizeDownloaded + (ulong)written;
                 depotDownloadCounter.sizeDownloaded = sizeDownloaded;
                 depotDownloadCounter.depotBytesCompressed += chunk.CompressedLength;
                 depotDownloadCounter.depotBytesUncompressed += chunk.UncompressedLength;

DepotDownloader/DepotDownloader.csproj

@@ -15,6 +15,6 @@
   <ItemGroup>
     <PackageReference Include="protobuf-net" Version="3.2.30" />
     <PackageReference Include="QRCoder" Version="1.6.0" />
-    <PackageReference Include="SteamKit2" Version="3.0.0-Beta.1" />
+    <PackageReference Include="SteamKit2" Version="3.0.0-Beta.2" />
   </ItemGroup>
 </Project>