rework download progress

Author: deansheather

Tests.Vpn.Service/DownloaderTest.cs

@@ -2,7 +2,6 @@
 using System.Security.Cryptography;
 using System.Security.Cryptography.X509Certificates;
 using System.Text;
-using System.Threading.Channels;
 using Coder.Desktop.Vpn.Service;
 using Microsoft.Extensions.Logging.Abstractions;
 
@@ -278,7 +277,7 @@ public async Task Download(CancellationToken ct)
         var dlTask = await manager.StartDownloadAsync(new HttpRequestMessage(HttpMethod.Get, url), destPath,
             NullDownloadValidator.Instance, ct);
         await dlTask.Task;
-        Assert.That(dlTask.TotalBytes, Is.EqualTo(4));
+        Assert.That(dlTask.BytesTotal, Is.EqualTo(4));
         Assert.That(dlTask.BytesWritten, Is.EqualTo(4));
         Assert.That(dlTask.Progress, Is.EqualTo(1));
         Assert.That(dlTask.IsCompleted, Is.True);
@@ -301,13 +300,13 @@ public async Task DownloadSameDest(CancellationToken ct)
             NullDownloadValidator.Instance, ct);
         var dlTask0 = await startTask0;
         await dlTask0.Task;
-        Assert.That(dlTask0.TotalBytes, Is.EqualTo(5));
+        Assert.That(dlTask0.BytesTotal, Is.EqualTo(5));
         Assert.That(dlTask0.BytesWritten, Is.EqualTo(5));
         Assert.That(dlTask0.Progress, Is.EqualTo(1));
         Assert.That(dlTask0.IsCompleted, Is.True);
         var dlTask1 = await startTask1;
         await dlTask1.Task;
-        Assert.That(dlTask1.TotalBytes, Is.EqualTo(5));
+        Assert.That(dlTask1.BytesTotal, Is.EqualTo(5));
         Assert.That(dlTask1.BytesWritten, Is.EqualTo(5));
         Assert.That(dlTask1.Progress, Is.EqualTo(1));
         Assert.That(dlTask1.IsCompleted, Is.True);
@@ -320,9 +319,9 @@ public async Task DownloadWithXOriginalContentLength(CancellationToken ct)
         using var httpServer = new TestHttpServer(async ctx =>
         {
             ctx.Response.StatusCode = 200;
-            ctx.Response.Headers.Add("X-Original-Content-Length", "6"); // wrong but should be used until complete
+            ctx.Response.Headers.Add("X-Original-Content-Length", "4");
             ctx.Response.ContentType = "text/plain";
-            ctx.Response.ContentLength64 = 4; // This should be ignored.
+            // Don't set Content-Length.
             await ctx.Response.OutputStream.WriteAsync("test"u8.ToArray(), ct);
         });
         var url = new Uri(httpServer.BaseUrl + "/test");
@@ -331,25 +330,30 @@ public async Task DownloadWithXOriginalContentLength(CancellationToken ct)
         var req = new HttpRequestMessage(HttpMethod.Get, url);
         var dlTask = await manager.StartDownloadAsync(req, destPath, NullDownloadValidator.Instance, ct);
 
-        var progressChannel = Channel.CreateUnbounded<DownloadProgressEvent>();
-        dlTask.ProgressChanged += (_, args) =>
-            Assert.That(progressChannel.Writer.TryWrite(args), Is.True);
-
         await dlTask.Task;
-        Assert.That(dlTask.TotalBytes, Is.EqualTo(4)); // should equal BytesWritten after completion
+        Assert.That(dlTask.BytesTotal, Is.EqualTo(4));
         Assert.That(dlTask.BytesWritten, Is.EqualTo(4));
-        progressChannel.Writer.Complete();
-
-        var list = progressChannel.Reader.ReadAllAsync(ct).ToBlockingEnumerable(ct).ToList();
-        Assert.That(list.Count, Is.GreaterThanOrEqualTo(2)); // there may be an item in the middle
-        // The first item should be the initial progress with 0 bytes written.
-        Assert.That(list[0].BytesWritten, Is.EqualTo(0));
-        Assert.That(list[0].BytesTotal, Is.EqualTo(6)); // from X-Original-Content-Length
-        Assert.That(list[0].Progress, Is.EqualTo(0.0d));
-        // The last item should be final progress with the actual total bytes.
-        Assert.That(list[^1].BytesWritten, Is.EqualTo(4));
-        Assert.That(list[^1].BytesTotal, Is.EqualTo(4)); // from the actual bytes written
-        Assert.That(list[^1].Progress, Is.EqualTo(1.0d));
+    }
+
+    [Test(Description = "Download with mismatched Content-Length")]
+    [CancelAfter(30_000)]
+    public async Task DownloadWithMismatchedContentLength(CancellationToken ct)
+    {
+        using var httpServer = new TestHttpServer(async ctx =>
+        {
+            ctx.Response.StatusCode = 200;
+            ctx.Response.Headers.Add("X-Original-Content-Length", "5"); // incorrect
+            ctx.Response.ContentType = "text/plain";
+            await ctx.Response.OutputStream.WriteAsync("test"u8.ToArray(), ct);
+        });
+        var url = new Uri(httpServer.BaseUrl + "/test");
+        var destPath = Path.Combine(_tempDir, "test");
+        var manager = new Downloader(NullLogger<Downloader>.Instance);
+        var req = new HttpRequestMessage(HttpMethod.Get, url);
+        var dlTask = await manager.StartDownloadAsync(req, destPath, NullDownloadValidator.Instance, ct);
+
+        var ex = Assert.ThrowsAsync<IOException>(() => dlTask.Task);
+        Assert.That(ex.Message, Is.EqualTo("Downloaded file size does not match expected response content length: Expected=5, BytesWritten=4"));
     }
 
     [Test(Description = "Download with custom headers")]

Vpn.Service/Downloader.cs

@@ -338,51 +338,14 @@ internal static async Task TaskOrCancellation(Task task, CancellationToken cance
     }
 }
 
-public class DownloadProgressEvent
-{
-    // TODO: speed calculation would be nice
-    public ulong BytesWritten { get; init; }
-    public ulong? BytesTotal { get; init; } // null if unknown
-
-    public double? Progress => BytesTotal == null ? null : (double)BytesWritten / BytesTotal.Value;
-
-    public override string ToString()
-    {
-        var s = FriendlyBytes(BytesWritten);
-        if (BytesTotal != null)
-            s += $" of {FriendlyBytes(BytesTotal.Value)}";
-        else
-            s += " of unknown";
-        if (Progress != null)
-            s += $" ({Progress:0%})";
-        return s;
-    }
-
-    private static readonly string[] ByteSuffixes = ["B", "KB", "MB", "GB", "TB", "PB", "EB"];
-
-    // Unfortunately this is copied from FriendlyByteConverter in App. Ideally
-    // it should go into some shared utilities project, but it's overkill to do
-    // that for a single tiny function until we have more shared code.
-    private static string FriendlyBytes(ulong bytes)
-    {
-        if (bytes == 0)
-            return $"0 {ByteSuffixes[0]}";
-
-        var place = Convert.ToInt32(Math.Floor(Math.Log(bytes, 1024)));
-        var num = Math.Round(bytes / Math.Pow(1024, place), 1);
-        return $"{num} {ByteSuffixes[place]}";
-    }
-}
-
 /// <summary>
 ///     Downloads a Url to a file on disk. The download will be written to a temporary file first, then moved to the final
 ///     destination. The SHA1 of any existing file will be calculated and used as an ETag to avoid downloading the file if
 ///     it hasn't changed.
 /// </summary>
 public class DownloadTask
 {
-    private const int BufferSize = 4096;
-    private const int ProgressUpdateDelayMs = 50;
+    private const int BufferSize = 64 * 1024;
     private const string XOriginalContentLengthHeader = "X-Original-Content-Length"; // overrides Content-Length if available
 
     private static readonly HttpClient HttpClient = new(new HttpClientHandler
@@ -398,22 +361,13 @@ public class DownloadTask
     private readonly string _destinationPath;
     private readonly string _tempDestinationPath;
 
-    // ProgressChanged events are always delayed by up to 50ms to avoid
-    // flooding.
-    //
-    // This will be called:
-    // - once after the request succeeds but before the read/write routine
-    //   begins
-    // - occasionally while the file is being downloaded (at least 50ms apart)
-    // - once when the download is complete
-    public EventHandler<DownloadProgressEvent>? ProgressChanged;
-
     public readonly HttpRequestMessage Request;
 
     public Task Task { get; private set; } = null!; // Set in EnsureStartedAsync
+    public bool DownloadStarted { get; private set; } // Whether we've received headers yet and started the actual download
     public ulong BytesWritten { get; private set; }
-    public ulong? TotalBytes { get; private set; }
-    public double? Progress => TotalBytes == null ? null : (double)BytesWritten / TotalBytes.Value;
+    public ulong? BytesTotal { get; private set; }
+    public double? Progress => BytesTotal == null ? null : (double)BytesWritten / BytesTotal.Value;
     public bool IsCompleted => Task.IsCompleted;
 
     internal DownloadTask(ILogger logger, HttpRequestMessage req, string destinationPath, IDownloadValidator validator)
@@ -496,32 +450,27 @@ private async Task Start(CancellationToken ct = default)
         }
 
         if (res.Content.Headers.ContentLength >= 0)
-            TotalBytes = (ulong)res.Content.Headers.ContentLength;
+            BytesTotal = (ulong)res.Content.Headers.ContentLength;
 
         // X-Original-Content-Length overrules Content-Length if set.
         if (res.Headers.TryGetValues(XOriginalContentLengthHeader, out var headerValues))
         {
             // If there are multiple we only look at the first one.
             var headerValue = headerValues.ToList().FirstOrDefault();
             if (!string.IsNullOrEmpty(headerValue) && ulong.TryParse(headerValue, out var originalContentLength))
-                TotalBytes = originalContentLength;
+                BytesTotal = originalContentLength;
             else
                 _logger.LogWarning(
                     "Failed to parse {XOriginalContentLengthHeader} header value '{HeaderValue}'",
                     XOriginalContentLengthHeader, headerValue);
         }
 
-        SendProgressUpdate(new DownloadProgressEvent
-        {
-            BytesWritten = 0,
-            BytesTotal = TotalBytes,
-        });
-
         await Download(res, ct);
     }
 
     private async Task Download(HttpResponseMessage res, CancellationToken ct)
     {
+        DownloadStarted = true;
         try
         {
             var sha1 = res.Headers.Contains("ETag") ? SHA1.Create() : null;
@@ -546,28 +495,13 @@ private async Task Download(HttpResponseMessage res, CancellationToken ct)
                     await tempFile.WriteAsync(buffer.AsMemory(0, n), ct);
                     sha1?.TransformBlock(buffer, 0, n, null, 0);
                     BytesWritten += (ulong)n;
-                    await QueueProgressUpdate(new DownloadProgressEvent
-                    {
-                        BytesWritten = BytesWritten,
-                        BytesTotal = TotalBytes,
-                    }, ct);
                 }
             }
 
-            // Clear any pending progress updates to ensure they won't be sent
-            // after the final update.
-            await ClearQueuedProgressUpdate(ct);
-            // Then write the final status update.
-            TotalBytes = BytesWritten;
-            SendProgressUpdate(new DownloadProgressEvent
-            {
-                BytesWritten = BytesWritten,
-                BytesTotal = BytesWritten,
-            });
-
-            if (TotalBytes != null && BytesWritten != TotalBytes)
+            BytesTotal ??= BytesWritten;
+            if (BytesWritten != BytesTotal)
                 throw new IOException(
-                    $"Downloaded file size does not match response Content-Length: Content-Length={TotalBytes}, BytesRead={BytesWritten}");
+                    $"Downloaded file size does not match expected response content length: Expected={BytesTotal}, BytesWritten={BytesWritten}");
 
             // Verify the ETag if it was sent by the server.
             if (res.Headers.Contains("ETag") && sha1 != null)
@@ -612,69 +546,4 @@ await QueueProgressUpdate(new DownloadProgressEvent
             throw;
         }
     }
-
-    // _progressEventLock protects _progressUpdateTask and _pendingProgressEvent.
-    private readonly RaiiSemaphoreSlim _progressEventLock = new(1, 1);
-    private readonly CancellationTokenSource _progressUpdateCts = new();
-    private Task? _progressUpdateTask;
-    private DownloadProgressEvent? _pendingProgressEvent;
-
-    // Can be called multiple times, but must not be called or in progress while
-    // SendQueuedProgressUpdateNow is called.
-    private async Task QueueProgressUpdate(DownloadProgressEvent e, CancellationToken ct)
-    {
-        using var _1 = await _progressEventLock.LockAsync(ct);
-        _pendingProgressEvent = e;
-
-        if (_progressUpdateCts.IsCancellationRequested)
-            throw new InvalidOperationException("Progress update task was cancelled, cannot queue new progress update");
-
-        // Start a task with a 50ms delay unless one is already running.
-        var cts = CancellationTokenSource.CreateLinkedTokenSource(ct, _progressUpdateCts.Token);
-        cts.CancelAfter(TimeSpan.FromSeconds(5));
-        _progressUpdateTask ??= Task.Delay(ProgressUpdateDelayMs, cts.Token)
-            .ContinueWith(t =>
-            {
-                cts.Cancel();
-                using var _2 = _progressEventLock.Lock();
-                _progressUpdateTask = null;
-                if (t.IsFaulted || t.IsCanceled) return;
-
-                var ev = _pendingProgressEvent;
-                if (ev != null) SendProgressUpdate(ev);
-            }, cts.Token);
-    }
-
-    // Must only be called after all QueueProgressUpdate calls have completed.
-    private async Task ClearQueuedProgressUpdate(CancellationToken ct)
-    {
-        Task? t;
-        using (var _ = _progressEventLock.LockAsync(ct))
-        {
-            await _progressUpdateCts.CancelAsync();
-            t = _progressUpdateTask;
-        }
-
-        // We can't continue to hold the lock here because the continuation
-        // grabs a lock. We don't need to worry about a new task spawning after
-        // this because the token is cancelled.
-        if (t == null) return;
-        try
-        {
-            await t.WaitAsync(ct);
-        }
-        catch (TaskCanceledException)
-        {
-            // Ignore
-        }
-    }
-
-    private void SendProgressUpdate(DownloadProgressEvent e)
-    {
-        var handler = ProgressChanged;
-        if (handler == null)
-            return;
-        // Start a new task in the background to invoke the event.
-        _ = Task.Run(() => handler.Invoke(this, e));
-    }
 }

Vpn.Service/Manager.cs

@@ -450,34 +450,46 @@ private async Task DownloadTunnelBinaryAsync(string baseUrl, SemVersion expected
             _logger.LogDebug("Skipping tunnel binary version validation");
         }
 
+        // Note: all ETag, signature and version validation is performed by the
+        // DownloadTask.
         var downloadTask = await _downloader.StartDownloadAsync(req, _config.TunnelBinaryPath, validators, ct);
 
-        var progressLock = new RaiiSemaphoreSlim(1, 1);
-        var progressBroadcastCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
-        downloadTask.ProgressChanged += (sender, ev) =>
+        // Wait for the download to complete, sending progress updates every
+        // 50ms.
+        while (true)
         {
-            using var _ = progressLock.Lock();
-            if (progressBroadcastCts.IsCancellationRequested) return;
-            _logger.LogInformation("Download progress: {ev}", ev);
+            // Wait for the download to complete, or for a short delay before
+            // we send a progress update.
+            var delayTask = Task.Delay(TimeSpan.FromMilliseconds(50), ct);
+            var winner = await Task.WhenAny([
+                downloadTask.Task,
+                delayTask,
+            ]);
+            if (winner == downloadTask.Task)
+                break;
+
+            // Task.WhenAny will not throw if the winner was cancelled, so
+            // check CT afterward and not beforehand.
+            ct.ThrowIfCancellationRequested();
+
+            if (!downloadTask.DownloadStarted)
+                // Don't send progress updates if we don't know what the
+                // progress is yet.
+                continue;
 
             var progress = new StartProgressDownloadProgress
             {
-                BytesWritten = ev.BytesWritten,
+                BytesWritten = downloadTask.BytesWritten,
             };
-            if (ev.BytesTotal != null)
-                progress.BytesTotal = ev.BytesTotal.Value;
-            BroadcastStartProgress(StartProgressStage.Downloading, progress, progressBroadcastCts.Token)
-                .Wait(progressBroadcastCts.Token);
-        };
+            if (downloadTask.BytesTotal != null)
+                progress.BytesTotal = downloadTask.BytesTotal.Value;
 
-        // Awaiting this will check the checksum (via the ETag) if the file
-        // exists, and will also validate the signature and version.
-        await downloadTask.Task;
+            await BroadcastStartProgress(StartProgressStage.Downloading, progress, ct);
+        }
 
-        // Prevent any lagging progress events from being sent.
-        // ReSharper disable once PossiblyMistakenUseOfCancellationToken
-        using (await progressLock.LockAsync(ct))
-            await progressBroadcastCts.CancelAsync();
+        // Await again to re-throw any exceptions that occurred during the
+        // download.
+        await downloadTask.Task;
 
         // We don't send a broadcast here as we immediately send one in the
         // parent routine.
@@ -486,7 +498,6 @@ private async Task DownloadTunnelBinaryAsync(string baseUrl, SemVersion expected
 
     private async Task BroadcastStartProgress(StartProgressStage stage, StartProgressDownloadProgress? downloadProgress = null, CancellationToken cancellationToken = default)
     {
-        _logger.LogInformation("Start progress: {stage}", stage);
         await FallibleBroadcast(new ServiceMessage
         {
             StartProgress = new StartProgress