KnightsOfTheArk
/
DepotDownloader
mirror of https://github.com/SteamRE/DepotDownloader.git
23
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Refactored download Tasks and use InvokeAsync to throttle Tasks

Browse Source
pull/132/head
Ryan Kistner 6 years ago
parent 609a66f280
commit d02aec256c
3 changed files with 479 additions and 395 deletions
Show Stats Download Patch File Download Diff File

57
DepotDownloader/CDNClientPool.cs
Unescape Escape View File

@ -20,7 +20,7 @@ namespace DepotDownloader
public CDNClient CDNClient { get; }
private readonly ConcurrentBag<CDNClient.Server> activeConnectionPool;
private readonly ConcurrentStack<CDNClient.Server> activeConnectionPool;
private readonly BlockingCollection<CDNClient.Server> availableServerEndpoints;
private readonly AutoResetEvent populatePoolEvent;
@ -33,7 +33,7 @@ namespace DepotDownloader
this.steamSession = steamSession;
CDNClient = new CDNClient(steamSession.steamClient);
activeConnectionPool = new ConcurrentBag<CDNClient.Server>();
activeConnectionPool = new ConcurrentStack<CDNClient.Server>();
availableServerEndpoints = new BlockingCollection<CDNClient.Server>();
populatePoolEvent = new AutoResetEvent(true);
@ -122,24 +122,6 @@ namespace DepotDownloader
}
}
private async Task<string> AuthenticateConnection(uint appId, uint depotId, CDNClient.Server server)
{
var host = steamSession.ResolveCDNTopLevelHost(server.Host);
var cdnKey = $"{depotId:D}:{host}";
steamSession.RequestCDNAuthToken(appId, depotId, host, cdnKey);
if (steamSession.CDNAuthTokens.TryGetValue(cdnKey, out var authTokenCallbackPromise))
{
var result = await authTokenCallbackPromise.Task;
return result.Token;
}
else
{
throw new Exception($"Failed to retrieve CDN token for server {server.Host} depot {depotId}");
}
}
private CDNClient.Server BuildConnection(CancellationToken token)
{
if (availableServerEndpoints.Count < ServerEndpointMinimumSize)
@ -150,29 +132,42 @@ namespace DepotDownloader
return availableServerEndpoints.Take(token);
}
public async Task<Tuple<CDNClient.Server, string>> GetConnectionForDepot(uint appId, uint depotId, CancellationToken token)
public CDNClient.Server GetConnection(CancellationToken token)
{
// Take a free connection from the connection pool
// If there were no free connections, create a new one from the server list
if (!activeConnectionPool.TryTake(out var server))
if (!activeConnectionPool.TryPop(out var connection))
{
server = BuildConnection(token);
connection = BuildConnection(token);
}
// If we don't have a CDN token yet for this server and depot, fetch one now
var cdnToken = await AuthenticateConnection(appId, depotId, server);
return connection;
}
public async Task<string> AuthenticateConnection(uint appId, uint depotId, CDNClient.Server server)
{
var host = steamSession.ResolveCDNTopLevelHost(server.Host);
var cdnKey = $"{depotId:D}:{host}";
steamSession.RequestCDNAuthToken(appId, depotId, host, cdnKey);
return Tuple.Create(server, cdnToken);
if (steamSession.CDNAuthTokens.TryGetValue(cdnKey, out var authTokenCallbackPromise))
{
var result = await authTokenCallbackPromise.Task;
return result.Token;
}
else
{
throw new Exception($"Failed to retrieve CDN token for server {server.Host} depot {depotId}");
}
}
public void ReturnConnection(Tuple<CDNClient.Server, string> server)
public void ReturnConnection(CDNClient.Server server)
{
if (server == null) return;
activeConnectionPool.Add(server.Item1);
activeConnectionPool.Push(server);
}
public void ReturnBrokenConnection(Tuple<CDNClient.Server, string> server)
public void ReturnBrokenConnection(CDNClient.Server server)
{
if (server == null) return;

261
DepotDownloader/ContentDownloader.cs
Unescape Escape View File

@ -1,5 +1,6 @@
using SteamKit2;
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.IO;
using System.Linq;
@ -592,21 +593,54 @@ namespace DepotDownloader
public ProtoManifest.ChunkData NewChunk { get; private set; }
}
private static async Task DownloadSteam3Async( uint appId, List<DepotDownloadInfo> depots )
private class FileStreamData
{
ulong TotalBytesCompressed = 0;
ulong TotalBytesUncompressed = 0;
public FileStream fileStream;
public SemaphoreSlim fileLock;
public int chunksToDownload;
}
foreach ( var depot in depots )
private class GlobalDownloadCounter
{
public long TotalBytesCompressed;
public long TotalBytesUncompressed;
}
private class DepotDownloadCounter
{
ulong DepotBytesCompressed = 0;
ulong DepotBytesUncompressed = 0;
public long CompleteDownloadSize;
public long SizeDownloaded;
public long DepotBytesCompressed;
public long DepotBytesUncompressed;
Console.WriteLine( "Downloading depot {0} - {1}", depot.id, depot.contentName );
}
private static async Task DownloadSteam3Async(uint appId, List<DepotDownloadInfo> depots)
{
CancellationTokenSource cts = new CancellationTokenSource();
cdnPool.ExhaustedToken = cts;
GlobalDownloadCounter downloadCounter = new GlobalDownloadCounter();
foreach (var depot in depots)
{
await DownloadSteam3AsyncDepot(cts, downloadCounter, appId, depot);
cts.Token.ThrowIfCancellationRequested();
}
Console.WriteLine("Total downloaded: {0} bytes ({1} bytes uncompressed) from {2} depots",
downloadCounter.TotalBytesCompressed, downloadCounter.TotalBytesUncompressed, depots.Count);
}
private static async Task DownloadSteam3AsyncDepot(CancellationTokenSource cts,
GlobalDownloadCounter downloadCounter,
uint appId, DepotDownloadInfo depot)
{
DepotDownloadCounter depotCounter = new DepotDownloadCounter();
Console.WriteLine("Downloading depot {0} - {1}", depot.id, depot.contentName);
ProtoManifest oldProtoManifest = null;
ProtoManifest newProtoManifest = null;
string configDir = Path.Combine(depot.installDir, CONFIG_DIR);
@ -620,7 +654,7 @@ namespace DepotDownloader
if (lastManifestId != INVALID_MANIFEST_ID)
{
var oldManifestFileName = Path.Combine( configDir, string.Format( "{0}.bin", lastManifestId ) );
var oldManifestFileName = Path.Combine(configDir, string.Format("{0}_{1}.bin", depot.id, lastManifestId));
if (File.Exists(oldManifestFileName))
{
@ -687,18 +721,26 @@ namespace DepotDownloader
DepotManifest depotManifest = null;
while ( depotManifest == null )
do
{
Tuple<CDNClient.Server, string> connection = null;
cts.Token.ThrowIfCancellationRequested();
CDNClient.Server connection = null;
try
{
connection = await cdnPool.GetConnectionForDepot( appId, depot.id, CancellationToken.None );
connection = cdnPool.GetConnection(cts.Token);
var cdnToken = await cdnPool.AuthenticateConnection(appId, depot.id, connection);
depotManifest = await cdnPool.CDNClient.DownloadManifestAsync(depot.id, depot.manifestId,
connection.Item1, connection.Item2, depot.depotKey ).ConfigureAwait(false);
connection, cdnToken, depot.depotKey).ConfigureAwait(false);
cdnPool.ReturnConnection(connection);
}
catch (TaskCanceledException)
{
Console.WriteLine("Connection timeout downloading depot manifest {0} {1}", depot.id, depot.manifestId);
}
catch (SteamKitWebRequestException e)
{
cdnPool.ReturnBrokenConnection(connection);
@ -713,19 +755,27 @@ namespace DepotDownloader
Console.WriteLine("Encountered error downloading depot manifest {0} {1}: {2}", depot.id, depot.manifestId, e.StatusCode);
}
}
catch (OperationCanceledException)
{
break;
}
catch (Exception e)
{
cdnPool.ReturnBrokenConnection(connection);
Console.WriteLine("Encountered error downloading manifest for depot {0} {1}: {2}", depot.id, depot.manifestId, e.Message);
}
}
while (depotManifest == null);
if (depotManifest == null)
{
Console.WriteLine("\nUnable to download manifest {0} for depot {1}", depot.manifestId, depot.id);
return;
cts.Cancel();
}
// Throw the cancellation exception if requested so that this task is marked failed
cts.Token.ThrowIfCancellationRequested();
byte[] checksum;
newProtoManifest = new ProtoManifest(depotManifest, depot.manifestId);
@ -757,11 +807,9 @@ namespace DepotDownloader
}
File.WriteAllText(txtManifest, manifestBuilder.ToString());
continue;
return;
}
ulong complete_download_size = 0;
ulong size_downloaded = 0;
string stagingDir = Path.Combine(depot.installDir, STAGING_DIR);
var filesAfterExclusions = newProtoManifest.Files.AsParallel().Where(f => TestIsFileIncluded(f.FileName)).ToList();
@ -783,25 +831,43 @@ namespace DepotDownloader
Directory.CreateDirectory(Path.GetDirectoryName(fileFinalPath));
Directory.CreateDirectory(Path.GetDirectoryName(fileStagingPath));
complete_download_size += file.TotalSize;
Interlocked.Add(ref depotCounter.CompleteDownloadSize, (long)file.TotalSize);
}
});
var ioSemaphore = new SemaphoreSlim( Config.MaxDownloads );
var networkSemaphore = new SemaphoreSlim( Config.MaxDownloads );
var files = filesAfterExclusions.Where(f => !f.Flags.HasFlag(EDepotFileFlag.Directory)).ToArray();
var ioTasks = new Task[ files.Length ];
for ( var i = 0; i < files.Length; i++ )
{
var file = files[ i ];
var task = Task.Run( async () =>
{
cts.Token.ThrowIfCancellationRequested();
var networkChunkQueue = new ConcurrentQueue<Tuple<FileStreamData, ProtoManifest.FileData, ProtoManifest.ChunkData>>();
await Util.InvokeAsync(
files.Select(f => new Func<Task>(async () =>
await Task.Run(() => DownloadSteam3AsyncDepotFile(cts, downloadCounter, depotCounter,
stagingDir, oldProtoManifest, newProtoManifest,
networkChunkQueue,
appId, depot, f)))),
maxDegreeOfParallelism: Config.MaxDownloads
);
try
await Util.InvokeAsync(
networkChunkQueue.Select((x) => new Func<Task>(async () =>
await Task.Run(() => DownloadSteam3AsyncDepotFileChunk(cts, downloadCounter, depotCounter,
stagingDir, oldProtoManifest, newProtoManifest,
appId, depot, x.Item2, x.Item1, x.Item3)))),
maxDegreeOfParallelism: Config.MaxDownloads
);
DepotConfigStore.Instance.InstalledManifestIDs[depot.id] = depot.manifestId;
DepotConfigStore.Save();
Console.WriteLine("Depot {0} - Downloaded {1} bytes ({2} bytes uncompressed)", depot.id, depotCounter.DepotBytesCompressed, depotCounter.DepotBytesUncompressed);
}
private static void DownloadSteam3AsyncDepotFile(
CancellationTokenSource cts,
GlobalDownloadCounter downloadCounter, DepotDownloadCounter depotDownloadCounter,
string stagingDir, ProtoManifest oldProtoManifest, ProtoManifest newProtoManifest,
ConcurrentQueue<Tuple<FileStreamData, ProtoManifest.FileData, ProtoManifest.ChunkData>> networkChunkQueue,
uint appId, DepotDownloadInfo depot, ProtoManifest.FileData file)
{
await ioSemaphore.WaitAsync().ConfigureAwait( false );
cts.Token.ThrowIfCancellationRequested();
string fileFinalPath = Path.Combine(depot.installDir, file.FileName);
@ -839,6 +905,10 @@ namespace DepotDownloader
if (Config.VerifyAll || !oldManifestFile.FileHash.SequenceEqual(file.FileHash))
{
// we have a version of this file, but it doesn't fully match what we want
if (Config.VerifyAll)
{
Console.WriteLine("Validating {0}", fileFinalPath);
}
var matchingChunks = new List<ChunkMatch>();
@ -855,6 +925,8 @@ namespace DepotDownloader
}
}
var orderedChunks = matchingChunks.OrderBy(x => x.OldChunk.Offset);
File.Move(fileFinalPath, fileStagingPath);
fs = File.Open(fileFinalPath, FileMode.Create);
@ -862,7 +934,7 @@ namespace DepotDownloader
using (var fsOld = File.Open(fileStagingPath, FileMode.Open))
{
foreach ( var match in matchingChunks )
foreach (var match in orderedChunks)
{
fsOld.Seek((long)match.OldChunk.Offset, SeekOrigin.Begin);
@ -895,53 +967,49 @@ namespace DepotDownloader
fs.SetLength((long)file.TotalSize);
}
Console.WriteLine("Validating {0}", fileFinalPath);
neededChunks = Util.ValidateSteam3FileChecksums(fs, file.Chunks.OrderBy(x => x.Offset).ToArray());
}
if (neededChunks.Count() == 0)
{
size_downloaded += file.TotalSize;
Console.WriteLine( "{0,6:#00.00}% {1}", ( ( float )size_downloaded / ( float )complete_download_size ) * 100.0f, fileFinalPath );
var sizeDownloaded = Interlocked.Add(ref depotDownloadCounter.SizeDownloaded, (long)file.TotalSize);
Console.WriteLine("{0,6:#00.00}% {1}", ((float)sizeDownloaded / (float)depotDownloadCounter.CompleteDownloadSize) * 100.0f, fileFinalPath);
if (fs != null)
fs.Dispose();
return;
}
else
{
size_downloaded += ( file.TotalSize - ( ulong )neededChunks.Select( x => ( long )x.UncompressedLength ).Sum() );
var sizeDownloaded = ((long)file.TotalSize - (long)neededChunks.Select(x => (long)x.UncompressedLength).Sum());
Interlocked.Add(ref depotDownloadCounter.SizeDownloaded, sizeDownloaded);
}
}
var fileSemaphore = new SemaphoreSlim(1);
var downloadTasks = new Task<DepotManifest.ChunkData>[neededChunks.Count];
for (var x = 0; x < neededChunks.Count; x++)
FileStreamData fileStreamData = new FileStreamData
{
var chunk = neededChunks[x];
fileStream = fs,
fileLock = new SemaphoreSlim(1),
chunksToDownload = neededChunks.Count
};
var downloadTask = Task.Run(async () =>
foreach (var chunk in neededChunks)
{
cts.Token.ThrowIfCancellationRequested();
networkChunkQueue.Enqueue(Tuple.Create(fileStreamData, file, chunk));
}
}
try
private static async Task DownloadSteam3AsyncDepotFileChunk(
CancellationTokenSource cts,
GlobalDownloadCounter downloadCounter, DepotDownloadCounter depotDownloadCounter,
string stagingDir, ProtoManifest oldProtoManifest, ProtoManifest newProtoManifest,
uint appId, DepotDownloadInfo depot, ProtoManifest.FileData file,
FileStreamData fileStreamData, ProtoManifest.ChunkData chunk)
{
await networkSemaphore.WaitAsync().ConfigureAwait(false);
cts.Token.ThrowIfCancellationRequested();
string chunkID = Util.EncodeHexString(chunk.ChunkID);
CDNClient.DepotChunk chunkData = null;
while (!cts.IsCancellationRequested)
{
Tuple<CDNClient.Server, string> connection;
try
{
connection = await cdnPool.GetConnectionForDepot(appId, depot.id, cts.Token);
}
catch (OperationCanceledException)
{
break;
}
DepotManifest.ChunkData data = new DepotManifest.ChunkData();
data.ChunkID = chunk.ChunkID;
@ -950,13 +1018,27 @@ namespace DepotDownloader
data.CompressedLength = chunk.CompressedLength;
data.UncompressedLength = chunk.UncompressedLength;
CDNClient.DepotChunk chunkData = null;
do
{
cts.Token.ThrowIfCancellationRequested();
CDNClient.Server connection = null;
try
{
connection = cdnPool.GetConnection(cts.Token);
var cdnToken = await cdnPool.AuthenticateConnection(appId, depot.id, connection);
chunkData = await cdnPool.CDNClient.DownloadDepotChunkAsync(depot.id, data,
connection.Item1, connection.Item2, depot.depotKey).ConfigureAwait(false);
cdnPool.ReturnConnection(connection);
connection, cdnToken, depot.depotKey).ConfigureAwait(false);
break;
cdnPool.ReturnConnection(connection);
}
catch (TaskCanceledException)
{
Console.WriteLine("Connection timeout downloading chunk {0}", chunkID);
}
catch (SteamKitWebRequestException e)
{
@ -965,7 +1047,6 @@ namespace DepotDownloader
if (e.StatusCode == HttpStatusCode.Unauthorized || e.StatusCode == HttpStatusCode.Forbidden)
{
Console.WriteLine("Encountered 401 for chunk {0}. Aborting.", chunkID);
cts.Cancel();
break;
}
else
@ -973,9 +1054,9 @@ namespace DepotDownloader
Console.WriteLine("Encountered error downloading chunk {0}: {1}", chunkID, e.StatusCode);
}
}
catch (TaskCanceledException)
catch (OperationCanceledException)
{
Console.WriteLine("Connection timeout downloading chunk {0}", chunkID);
break;
}
catch (Exception e)
{
@ -983,6 +1064,7 @@ namespace DepotDownloader
Console.WriteLine("Encountered unexpected error downloading chunk {0}: {1}", chunkID, e.Message);
}
}
while (chunkData == null);
if (chunkData == null)
{
@ -995,61 +1077,34 @@ namespace DepotDownloader
try
{
await fileSemaphore.WaitAsync().ConfigureAwait(false);
fs.Seek((long)chunkData.ChunkInfo.Offset, SeekOrigin.Begin);
fs.Write(chunkData.Data, 0, chunkData.Data.Length);
return chunkData.ChunkInfo;
}
finally
{
fileSemaphore.Release();
}
}
finally
{
networkSemaphore.Release();
}
});
downloadTasks[x] = downloadTask;
}
var completedDownloads = await Task.WhenAll(downloadTasks).ConfigureAwait(false);
await fileStreamData.fileLock.WaitAsync().ConfigureAwait(false);
fs.Dispose();
foreach (var chunkInfo in completedDownloads)
{
TotalBytesCompressed += chunkInfo.CompressedLength;
DepotBytesCompressed += chunkInfo.CompressedLength;
TotalBytesUncompressed += chunkInfo.UncompressedLength;
DepotBytesUncompressed += chunkInfo.UncompressedLength;
fileStreamData.fileStream.Seek((long)chunkData.ChunkInfo.Offset, SeekOrigin.Begin);
fileStreamData.fileStream.Write(chunkData.Data, 0, chunkData.Data.Length);
size_downloaded += chunkInfo.UncompressedLength;
}
Console.WriteLine("{0,6:#00.00}% {1}", ((float)size_downloaded / (float)complete_download_size) * 100.0f, fileFinalPath);
}
finally
{
ioSemaphore.Release();
fileStreamData.fileLock.Release();
}
} );
ioTasks[ i ] = task;
}
Interlocked.Add(ref downloadCounter.TotalBytesCompressed, chunk.CompressedLength);
Interlocked.Add(ref depotDownloadCounter.DepotBytesCompressed, chunk.CompressedLength);
Interlocked.Add(ref downloadCounter.TotalBytesUncompressed, chunk.UncompressedLength);
Interlocked.Add(ref depotDownloadCounter.DepotBytesUncompressed, chunk.UncompressedLength);
await Task.WhenAll(ioTasks).ConfigureAwait(false);
var sizeDownloaded = Interlocked.Add(ref depotDownloadCounter.SizeDownloaded, chunkData.Data.Length);
DepotConfigStore.Instance.InstalledManifestIDs[ depot.id ] = depot.manifestId;
DepotConfigStore.Save();
int remainingChunks = Interlocked.Decrement(ref fileStreamData.chunksToDownload);
if (remainingChunks == 0)
{
fileStreamData.fileStream.Dispose();
Console.WriteLine( "Depot {0} - Downloaded {1} bytes ({2} bytes uncompressed)", depot.id, DepotBytesCompressed, DepotBytesUncompressed );
var fileFinalPath = Path.Combine(depot.installDir, file.FileName);
Console.WriteLine("{0,6:#00.00}% {1}", ((float)sizeDownloaded / (float)depotDownloadCounter.CompleteDownloadSize) * 100.0f, fileFinalPath);
}
Console.WriteLine( "Total downloaded: {0} bytes ({1} bytes uncompressed) from {2} depots", TotalBytesCompressed, TotalBytesUncompressed, depots.Count );
}
}
}

34
DepotDownloader/Util.cs
Unescape Escape View File

@ -5,6 +5,7 @@ using System.Linq;
using System.Runtime.InteropServices;
using System.Security.Cryptography;
using System.Text;
using System.Threading.Tasks;
namespace DepotDownloader
{
@ -132,5 +133,38 @@ namespace DepotDownloader
( sb, v ) => sb.Append( v.ToString( "x2" ) )
).ToString();
}
public static async Task InvokeAsync(IEnumerable<Func<Task>> taskFactories, int maxDegreeOfParallelism)
{
if (taskFactories == null) throw new ArgumentNullException(nameof(taskFactories));
if (maxDegreeOfParallelism <= 0) throw new ArgumentException(nameof(maxDegreeOfParallelism));
Func<Task>[] queue = taskFactories.ToArray();
if (queue.Length == 0)
{
return;
}
List<Task> tasksInFlight = new List<Task>(maxDegreeOfParallelism);
int index = 0;
do
{
while (tasksInFlight.Count < maxDegreeOfParallelism && index < queue.Length)
{
Func<Task> taskFactory = queue[index++];
tasksInFlight.Add(taskFactory());
}
Task completedTask = await Task.WhenAny(tasksInFlight).ConfigureAwait(false);
await completedTask.ConfigureAwait(false);
tasksInFlight.Remove(completedTask);
}
while (index < queue.Length || tasksInFlight.Count != 0);
}
}
}

Write Preview
Loading…
Cancel
Save