Skip to content

Commit

Permalink
Use ansi escape code to indicate download progress
Browse files Browse the repository at this point in the history
  • Loading branch information
xPaw committed Jul 28, 2024
1 parent ae72ac0 commit ee2dcc4
Show file tree
Hide file tree
Showing 4 changed files with 206 additions and 6 deletions.
46 changes: 46 additions & 0 deletions DepotDownloader/Ansi.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
using System;
using Spectre.Console;

namespace DepotDownloader;

static class Ansi
{
// https://conemu.github.io/en/AnsiEscapeCodes.html#ConEmu_specific_OSC
// https://learn.microsoft.com/en-us/windows/terminal/tutorials/progress-bar-sequences
public enum ProgressState
{
Hidden = 0,
Default = 1,
Error = 2,
Indeterminate = 3,
Warning = 4,
}

const string ESC = "\u001b";
const string BEL = "\u0007";

private static bool useProgress;

public static void Init()
{
var (supportsAnsi, legacyConsole) = AnsiDetector.Detect(stdError: false, upgrade: true);

useProgress = supportsAnsi && !legacyConsole;
}

public static void Progress(ulong downloaded, ulong total)
{
var progress = (byte)MathF.Round(downloaded / (float)total * 100.0f);
Progress(ProgressState.Default, progress);
}

public static void Progress(ProgressState state, byte progress = 0)
{
if (!useProgress)
{
return;
}

Console.Write($"{ESC}]9;4;{(byte)state};{progress}{BEL}");
}
}
133 changes: 133 additions & 0 deletions DepotDownloader/AnsiDetector.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,133 @@
// Copied from https://github.com/spectreconsole/spectre.console/blob/d79e6adc5f8e637fb35c88f987023ffda6707243/src/Spectre.Console/Internal/Backends/Ansi/AnsiDetector.cs
// which is based on https://github.com/keqingrong/supports-ansi/blob/master/index.js
// <auto-generated/>

using System;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text.RegularExpressions;

namespace Spectre.Console;

internal static class AnsiDetector
{
private static readonly Regex[] _regexes =
[
new("^xterm"), // xterm, PuTTY, Mintty
new("^rxvt"), // RXVT
new("^eterm"), // Eterm
new("^screen"), // GNU screen, tmux
new("tmux"), // tmux
new("^vt100"), // DEC VT series
new("^vt102"), // DEC VT series
new("^vt220"), // DEC VT series
new("^vt320"), // DEC VT series
new("ansi"), // ANSI
new("scoansi"), // SCO ANSI
new("cygwin"), // Cygwin, MinGW
new("linux"), // Linux console
new("konsole"), // Konsole
new("bvterm"), // Bitvise SSH Client
new("^st-256color"), // Suckless Simple Terminal, st
new("alacritty"), // Alacritty
];

public static (bool SupportsAnsi, bool LegacyConsole) Detect(bool stdError, bool upgrade)
{
// Running on Windows?
if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
{
// Running under ConEmu?
var conEmu = Environment.GetEnvironmentVariable("ConEmuANSI");
if (!string.IsNullOrEmpty(conEmu) && conEmu.Equals("On", StringComparison.OrdinalIgnoreCase))
{
return (true, false);
}

var supportsAnsi = Windows.SupportsAnsi(upgrade, stdError, out var legacyConsole);
return (supportsAnsi, legacyConsole);
}

return DetectFromTerm();
}

private static (bool SupportsAnsi, bool LegacyConsole) DetectFromTerm()
{
// Check if the terminal is of type ANSI/VT100/xterm compatible.
var term = Environment.GetEnvironmentVariable("TERM");
if (!string.IsNullOrWhiteSpace(term))
{
if (_regexes.Any(regex => regex.IsMatch(term)))
{
return (true, false);
}
}

return (false, true);
}

private static class Windows
{
private const int STD_OUTPUT_HANDLE = -11;
private const int STD_ERROR_HANDLE = -12;
private const uint ENABLE_VIRTUAL_TERMINAL_PROCESSING = 0x0004;
private const uint DISABLE_NEWLINE_AUTO_RETURN = 0x0008;

[DllImport("kernel32.dll")]
private static extern bool GetConsoleMode(IntPtr hConsoleHandle, out uint lpMode);

[DllImport("kernel32.dll")]
private static extern bool SetConsoleMode(IntPtr hConsoleHandle, uint dwMode);

[DllImport("kernel32.dll", SetLastError = true)]
private static extern IntPtr GetStdHandle(int nStdHandle);

[DllImport("kernel32.dll")]
public static extern uint GetLastError();

public static bool SupportsAnsi(bool upgrade, bool stdError, out bool isLegacy)
{
isLegacy = false;

try
{
var @out = GetStdHandle(stdError ? STD_ERROR_HANDLE : STD_OUTPUT_HANDLE);
if (!GetConsoleMode(@out, out var mode))
{
// Could not get console mode, try TERM (set in cygwin, WSL-Shell).
var (ansiFromTerm, legacyFromTerm) = DetectFromTerm();

isLegacy = ansiFromTerm ? legacyFromTerm : isLegacy;
return ansiFromTerm;
}

if ((mode & ENABLE_VIRTUAL_TERMINAL_PROCESSING) == 0)
{
isLegacy = true;

if (!upgrade)
{
return false;
}

// Try enable ANSI support.
mode |= ENABLE_VIRTUAL_TERMINAL_PROCESSING | DISABLE_NEWLINE_AUTO_RETURN;
if (!SetConsoleMode(@out, mode))
{
// Enabling failed.
return false;
}

isLegacy = false;
}

return true;
}
catch
{
// All we know here is that we don't support ANSI.
return false;
}
}
}
}
31 changes: 25 additions & 6 deletions DepotDownloader/ContentDownloader.cs
Original file line number Diff line number Diff line change
Expand Up @@ -610,6 +610,7 @@ private class FileStreamData

private class GlobalDownloadCounter
{
public ulong completeDownloadSize;
public ulong totalBytesCompressed;
public ulong totalBytesUncompressed;
}
Expand All @@ -624,6 +625,8 @@ private class DepotDownloadCounter

private static async Task DownloadSteam3Async(List<DepotDownloadInfo> depots)
{
Ansi.Progress(Ansi.ProgressState.Indeterminate);

var cts = new CancellationTokenSource();
cdnPool.ExhaustedToken = cts;

Expand All @@ -634,7 +637,7 @@ private static async Task DownloadSteam3Async(List<DepotDownloadInfo> depots)
// First, fetch all the manifests for each depot (including previous manifests) and perform the initial setup
foreach (var depot in depots)
{
var depotFileData = await ProcessDepotManifestAndFiles(cts, depot);
var depotFileData = await ProcessDepotManifestAndFiles(cts, depot, downloadCounter);

if (depotFileData != null)
{
Expand Down Expand Up @@ -665,11 +668,13 @@ private static async Task DownloadSteam3Async(List<DepotDownloadInfo> depots)
await DownloadSteam3AsyncDepotFiles(cts, downloadCounter, depotFileData, allFileNamesAllDepots);
}

Ansi.Progress(Ansi.ProgressState.Hidden);

Console.WriteLine("Total downloaded: {0} bytes ({1} bytes uncompressed) from {2} depots",
downloadCounter.totalBytesCompressed, downloadCounter.totalBytesUncompressed, depots.Count);
}

private static async Task<DepotFilesData> ProcessDepotManifestAndFiles(CancellationTokenSource cts, DepotDownloadInfo depot)
private static async Task<DepotFilesData> ProcessDepotManifestAndFiles(CancellationTokenSource cts, DepotDownloadInfo depot, GlobalDownloadCounter downloadCounter)
{
var depotCounter = new DepotDownloadCounter();

Expand Down Expand Up @@ -751,7 +756,7 @@ private static async Task<DepotFilesData> ProcessDepotManifestAndFiles(Cancellat
}
else
{
Console.Write("Downloading depot manifest...");
Console.Write("Downloading depot manifest... ");

DepotManifest depotManifest = null;
ulong manifestRequestCode = 0;
Expand Down Expand Up @@ -814,7 +819,7 @@ private static async Task<DepotFilesData> ProcessDepotManifestAndFiles(Cancellat

if (e.StatusCode == HttpStatusCode.Unauthorized || e.StatusCode == HttpStatusCode.Forbidden)
{
Console.WriteLine("Encountered 401 for depot manifest {0} {1}. Aborting.", depot.DepotId, depot.ManifestId);
Console.WriteLine("Encountered {2} for depot manifest {0} {1}. Aborting.", depot.DepotId, depot.ManifestId, (int)e.StatusCode);
break;
}

Expand Down Expand Up @@ -889,6 +894,7 @@ private static async Task<DepotFilesData> ProcessDepotManifestAndFiles(Cancellat
Directory.CreateDirectory(Path.GetDirectoryName(fileFinalPath));
Directory.CreateDirectory(Path.GetDirectoryName(fileStagingPath));

downloadCounter.completeDownloadSize += file.TotalSize;
depotCounter.completeDownloadSize += file.TotalSize;
}
});
Expand Down Expand Up @@ -918,7 +924,7 @@ private static async Task DownloadSteam3AsyncDepotFiles(CancellationTokenSource

await Util.InvokeAsync(
files.Select(file => new Func<Task>(async () =>
await Task.Run(() => DownloadSteam3AsyncDepotFile(cts, depotFilesData, file, networkChunkQueue)))),
await Task.Run(() => DownloadSteam3AsyncDepotFile(cts, downloadCounter, depotFilesData, file, networkChunkQueue)))),
maxDegreeOfParallelism: Config.MaxDownloads
);

Expand Down Expand Up @@ -966,6 +972,7 @@ await Task.Run(() => DownloadSteam3AsyncDepotFileChunk(cts, downloadCounter, dep

private static void DownloadSteam3AsyncDepotFile(
CancellationTokenSource cts,
GlobalDownloadCounter downloadCounter,
DepotFilesData depotFilesData,
ProtoManifest.FileData file,
ConcurrentQueue<(FileStreamData, ProtoManifest.FileData, ProtoManifest.ChunkData)> networkChunkQueue)
Expand Down Expand Up @@ -1128,6 +1135,11 @@ private static void DownloadSteam3AsyncDepotFile(
Console.WriteLine("{0,6:#00.00}% {1}", (depotDownloadCounter.sizeDownloaded / (float)depotDownloadCounter.completeDownloadSize) * 100.0f, fileFinalPath);
}

lock (downloadCounter)
{
downloadCounter.completeDownloadSize -= file.TotalSize;
}

return;
}

Expand All @@ -1136,6 +1148,11 @@ private static void DownloadSteam3AsyncDepotFile(
{
depotDownloadCounter.sizeDownloaded += sizeOnDisk;
}

lock (downloadCounter)
{
downloadCounter.completeDownloadSize -= sizeOnDisk;
}
}

var fileIsExecutable = file.Flags.HasFlag(EDepotFileFlag.Executable);
Expand Down Expand Up @@ -1217,7 +1234,7 @@ private static async Task DownloadSteam3AsyncDepotFileChunk(

if (e.StatusCode == HttpStatusCode.Unauthorized || e.StatusCode == HttpStatusCode.Forbidden)
{
Console.WriteLine("Encountered 401 for chunk {0}. Aborting.", chunkID);
Console.WriteLine("Encountered {1} for chunk {0}. Aborting.", chunkID, (int)e.StatusCode);
break;
}

Expand Down Expand Up @@ -1281,6 +1298,8 @@ private static async Task DownloadSteam3AsyncDepotFileChunk(
{
downloadCounter.totalBytesCompressed += chunk.CompressedLength;
downloadCounter.totalBytesUncompressed += chunk.UncompressedLength;

Ansi.Progress(downloadCounter.totalBytesUncompressed, downloadCounter.completeDownloadSize);
}

if (remainingChunks == 0)
Expand Down
2 changes: 2 additions & 0 deletions DepotDownloader/Program.cs
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@ static async Task<int> MainAsync(string[] args)
return 1;
}

Ansi.Init();

DebugLog.Enabled = false;

AccountSettingsStore.LoadFromFile("account.config");
Expand Down

0 comments on commit ee2dcc4

Please sign in to comment.