(#9) Add Download Progress

This adds download progress to the console when ChocolateyProgressInfo.ShouldDisplayDownloadProgress is set to true. It works for both v2 and v3 feeds, but does not display for local feeds.
chocolatey · Dec 19, 2022 · f311ec5 · f311ec5
1 parent 6ab0442
commit f311ec5
Show file tree

Hide file tree

Showing 9 changed files with 369 additions and 13 deletions.
diff --git a/src/NuGet.Core/NuGet.Protocol/ChocolateyProgressInfo.cs b/src/NuGet.Core/NuGet.Protocol/ChocolateyProgressInfo.cs
@@ -0,0 +1,28 @@
+// Copyright (c) 2022-Present Chocolatey Software, Inc.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Text;
+using System.Threading.Tasks;
+using NuGet.Packaging.Core;
+
+namespace NuGet.Protocol
+{
+    public class ChocolateyProgressInfo
+    {
+        public ChocolateyProgressInfo(PackageIdentity identity, long? length = null, string operation = "")
+        {
+            Operation = operation;
+            Length = length;
+            Identity = identity;
+        }
+
+        public string Operation { get; set; }
+        public PackageIdentity Identity { get; set; }
+        public long? Length { get; set; }
+        public static bool ShouldDisplayDownloadProgress { get; set; }
+        public bool Completed { get; set; }
+    }
+}
diff --git a/src/NuGet.Core/NuGet.Protocol/ChocolateyProgressStream.cs b/src/NuGet.Core/NuGet.Protocol/ChocolateyProgressStream.cs
@@ -0,0 +1,55 @@
+// Copyright (c) 2022-Present Chocolatey Software, Inc.
+// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
+
+using System;
+using System.IO;
+
+namespace NuGet.Protocol
+{
+    public delegate void StreamUpdate(Stream sender, long progress, long totalProgress);
+
+    // Based on https://www.thomasbogholm.net/2021/07/15/extend-streams-with-progress-reporting-progressreportingstream/
+    public class ChocolateyProgressStream : Stream
+    {
+        public event StreamUpdate WriteProgress;
+        public event StreamUpdate ReadProgress;
+        long _totalBytesRead;
+        long _totalBytesWritten;
+        private void UpdateRead(long read)
+        {
+            _totalBytesRead += read;
+            ReadProgress?.Invoke(this, read, _totalBytesRead);
+        }
+        private void UpdateWritten(long written)
+        {
+            _totalBytesWritten += written;
+            WriteProgress?.Invoke(this, written, _totalBytesWritten);
+        }
+
+        Stream InnerStream { get; }
+
+        public override int Read(byte[] buffer, int offset, int count)
+        {
+            var result = InnerStream.Read(buffer, offset, count);
+            UpdateRead(result);
+            return result;
+        }
+        public override void Write(byte[] buffer, int offset, int count)
+        {
+            InnerStream.Write(buffer, offset, count);
+            UpdateWritten(count);
+        }
+        public override bool CanRead => InnerStream.CanRead;
+        public override bool CanSeek => InnerStream.CanSeek;
+        public override bool CanWrite => InnerStream.CanWrite;
+        public override long Length => InnerStream.Length;
+        public override long Position { get => InnerStream.Position; set => InnerStream.Position = value; }
+        public ChocolateyProgressStream(Stream s)
+        {
+            InnerStream = s;
+        }
+        public override void Flush() => InnerStream.Flush();
+        public override long Seek(long offset, SeekOrigin origin) => InnerStream.Seek(offset, origin);
+        public override void SetLength(long value) => InnerStream.SetLength(value);
+    }
+}
diff --git a/src/NuGet.Core/NuGet.Protocol/HttpSource/HttpSource.cs b/src/NuGet.Core/NuGet.Protocol/HttpSource/HttpSource.cs
@@ -264,6 +264,43 @@ public async Task<T> ProcessStreamAsync<T>(
                 token);
         }
 
+        //////////////////////////////////////////////////////////
+        // Start - Chocolatey Specific Modification
+        //////////////////////////////////////////////////////////
+        public async Task<T> ProcessStreamAsync<T>(
+            HttpSourceRequest request,
+            Func<Stream, ChocolateyProgressInfo, Task<T>> processAsync,
+            SourceCacheContext cacheContext,
+            ILogger log,
+            ChocolateyProgressInfo progressInfo,
+            CancellationToken token)
+        {
+            return await ProcessResponseAsync(
+                request,
+                async response =>
+                {
+                    if ((request.IgnoreNotFounds && response.StatusCode == HttpStatusCode.NotFound) ||
+                         response.StatusCode == HttpStatusCode.NoContent)
+                    {
+                        return await processAsync(null, null);
+                    }
+
+                    response.EnsureSuccessStatusCode();
+
+                    progressInfo.Length = response.Content.Headers.ContentLength;
+                    progressInfo.Operation = "Downloading";
+
+                    var networkStream = await response.Content.ReadAsStreamAsync();
+                    return await processAsync(networkStream, progressInfo);
+                },
+                cacheContext,
+                log,
+                token);
+        }
+        //////////////////////////////////////////////////////////
+        // End - Chocolatey Specific Modification
+        //////////////////////////////////////////////////////////
+
         public Task<T> ProcessResponseAsync<T>(
             HttpSourceRequest request,
             Func<HttpResponseMessage, Task<T>> processAsync,

diff --git a/src/NuGet.Core/NuGet.Protocol/HttpSource/IHttpSource.cs b/src/NuGet.Core/NuGet.Protocol/HttpSource/IHttpSource.cs
@@ -47,6 +47,14 @@ Task<T> ProcessStreamAsync<T>(
             ILogger log,
             CancellationToken token);
 
+        Task<T> ProcessStreamAsync<T>(
+            HttpSourceRequest request,
+            Func<Stream, ChocolateyProgressInfo, Task<T>> processAsync,
+            SourceCacheContext cacheContext,
+            ILogger log,
+            ChocolateyProgressInfo progressInfo,
+            CancellationToken token);
+
         Task<T> ProcessResponseAsync<T>(
             HttpSourceRequest request,
             Func<HttpResponseMessage, Task<T>> processAsync,

diff --git a/src/NuGet.Core/NuGet.Protocol/PublicAPI/net472/PublicAPI.Unshipped.txt b/src/NuGet.Core/NuGet.Protocol/PublicAPI/net472/PublicAPI.Unshipped.txt
@@ -16,6 +16,41 @@ NuGet.Protocol.Core.Types.DownloadCache.FileName.get -> string
 NuGet.Protocol.Core.Types.DownloadCache.FileName.set -> void
 NuGet.Protocol.Core.Types.DownloadCache.OriginalUrl.get -> string
 NuGet.Protocol.Core.Types.DownloadCache.OriginalUrl.set -> void
+NuGet.Protocol.ChocolateyProgressInfo
+NuGet.Protocol.ChocolateyProgressInfo.ChocolateyProgressInfo(NuGet.Packaging.Core.PackageIdentity identity, long? length = null, string operation = "") -> void
+NuGet.Protocol.ChocolateyProgressInfo.Completed.get -> bool
+NuGet.Protocol.ChocolateyProgressInfo.Completed.set -> void
+NuGet.Protocol.ChocolateyProgressInfo.Identity.get -> NuGet.Packaging.Core.PackageIdentity
+NuGet.Protocol.ChocolateyProgressInfo.Identity.set -> void
+NuGet.Protocol.ChocolateyProgressInfo.Length.get -> long?
+NuGet.Protocol.ChocolateyProgressInfo.Length.set -> void
+NuGet.Protocol.ChocolateyProgressInfo.Operation.get -> string
+NuGet.Protocol.ChocolateyProgressInfo.Operation.set -> void
+NuGet.Protocol.ChocolateyProgressStream
+NuGet.Protocol.ChocolateyProgressStream.ChocolateyProgressStream(System.IO.Stream s) -> void
+NuGet.Protocol.ChocolateyProgressStream.ReadProgress -> NuGet.Protocol.StreamUpdate
+NuGet.Protocol.ChocolateyProgressStream.WriteProgress -> NuGet.Protocol.StreamUpdate
+NuGet.Protocol.Core.Types.FatalProtocolException.FatalProtocolException(System.Runtime.Serialization.SerializationInfo info, System.Runtime.Serialization.StreamingContext context) -> void
+NuGet.Protocol.Core.Types.NuGetProtocolException.NuGetProtocolException(System.Runtime.Serialization.SerializationInfo info, System.Runtime.Serialization.StreamingContext context) -> void
+NuGet.Protocol.Core.Types.RetriableProtocolException.RetriableProtocolException(System.Runtime.Serialization.SerializationInfo info, System.Runtime.Serialization.StreamingContext context) -> void
+NuGet.Protocol.HttpSource.ProcessHttpStreamAsync<T>(NuGet.Protocol.HttpSourceRequest request, System.Func<System.Net.Http.HttpResponseMessage, System.Threading.Tasks.Task<T>> processAsync, NuGet.Common.ILogger log, System.Threading.CancellationToken token) -> System.Threading.Tasks.Task<T>
+NuGet.Protocol.HttpSource.ProcessStreamAsync<T>(NuGet.Protocol.HttpSourceRequest request, System.Func<System.IO.Stream, NuGet.Protocol.ChocolateyProgressInfo, System.Threading.Tasks.Task<T>> processAsync, NuGet.Protocol.Core.Types.SourceCacheContext cacheContext, NuGet.Common.ILogger log, NuGet.Protocol.ChocolateyProgressInfo progressInfo, System.Threading.CancellationToken token) -> System.Threading.Tasks.Task<T>
+NuGet.Protocol.HttpSourceResource.OverrideHttpSource(NuGet.Protocol.IHttpSource source) -> void
+NuGet.Protocol.IHttpSource
+NuGet.Protocol.IHttpSource.GetAsync<T>(NuGet.Protocol.HttpSourceCachedRequest request, System.Func<NuGet.Protocol.HttpSourceResult, System.Threading.Tasks.Task<T>> processAsync, NuGet.Common.ILogger log, System.Threading.CancellationToken token) -> System.Threading.Tasks.Task<T>
+NuGet.Protocol.IHttpSource.GetJObjectAsync(NuGet.Protocol.HttpSourceRequest request, NuGet.Common.ILogger log, System.Threading.CancellationToken token) -> System.Threading.Tasks.Task<Newtonsoft.Json.Linq.JObject>
+NuGet.Protocol.IHttpSource.HttpCacheDirectory.get -> string
+NuGet.Protocol.IHttpSource.HttpCacheDirectory.set -> void
+NuGet.Protocol.IHttpSource.PackageSource.get -> string
+NuGet.Protocol.IHttpSource.ProcessHttpStreamAsync<T>(NuGet.Protocol.HttpSourceRequest request, System.Func<System.Net.Http.HttpResponseMessage, System.Threading.Tasks.Task<T>> processAsync, NuGet.Common.ILogger log, System.Threading.CancellationToken token) -> System.Threading.Tasks.Task<T>
+NuGet.Protocol.IHttpSource.ProcessResponseAsync<T>(NuGet.Protocol.HttpSourceRequest request, System.Func<System.Net.Http.HttpResponseMessage, System.Threading.Tasks.Task<T>> processAsync, NuGet.Common.ILogger log, System.Threading.CancellationToken token) -> System.Threading.Tasks.Task<T>
+NuGet.Protocol.IHttpSource.ProcessResponseAsync<T>(NuGet.Protocol.HttpSourceRequest request, System.Func<System.Net.Http.HttpResponseMessage, System.Threading.Tasks.Task<T>> processAsync, NuGet.Protocol.Core.Types.SourceCacheContext cacheContext, NuGet.Common.ILogger log, System.Threading.CancellationToken token) -> System.Threading.Tasks.Task<T>
+NuGet.Protocol.IHttpSource.ProcessStreamAsync<T>(NuGet.Protocol.HttpSourceRequest request, System.Func<System.IO.Stream, NuGet.Protocol.ChocolateyProgressInfo, System.Threading.Tasks.Task<T>> processAsync, NuGet.Protocol.Core.Types.SourceCacheContext cacheContext, NuGet.Common.ILogger log, NuGet.Protocol.ChocolateyProgressInfo progressInfo, System.Threading.CancellationToken token) -> System.Threading.Tasks.Task<T>
+NuGet.Protocol.IHttpSource.ProcessStreamAsync<T>(NuGet.Protocol.HttpSourceRequest request, System.Func<System.IO.Stream, System.Threading.Tasks.Task<T>> processAsync, NuGet.Common.ILogger log, System.Threading.CancellationToken token) -> System.Threading.Tasks.Task<T>
+NuGet.Protocol.IHttpSource.ProcessStreamAsync<T>(NuGet.Protocol.HttpSourceRequest request, System.Func<System.IO.Stream, System.Threading.Tasks.Task<T>> processAsync, NuGet.Protocol.Core.Types.SourceCacheContext cacheContext, NuGet.Common.ILogger log, System.Threading.CancellationToken token) -> System.Threading.Tasks.Task<T>
+NuGet.Protocol.IHttpSource.RetryHandler.get -> NuGet.Protocol.IHttpRetryHandler
+NuGet.Protocol.IHttpSource.RetryHandler.set -> void
+NuGet.Protocol.InvalidCacheProtocolException.InvalidCacheProtocolException(System.Runtime.Serialization.SerializationInfo info, System.Runtime.Serialization.StreamingContext context) -> void
 NuGet.Protocol.Core.Types.IPackageSearchMetadata.DownloadCache.get -> System.Collections.Generic.IEnumerable<NuGet.Protocol.Core.Types.DownloadCache>
 NuGet.Protocol.Core.Types.IPackageSearchMetadata.DownloadCacheDate.get -> System.DateTime?
 NuGet.Protocol.Core.Types.IPackageSearchMetadata.IsApproved.get -> bool
@@ -161,6 +196,7 @@ NuGet.Protocol.PackageSearchMetadataV2Feed.PackageTestResultStatusDate.get -> Sy
 NuGet.Protocol.PackageSearchMetadataV2Feed.PackageValidationResultDate.get -> System.DateTime?
 NuGet.Protocol.PackageSearchMetadataV2Feed.PackageValidationResultStatus.get -> string
 NuGet.Protocol.PackageSearchMetadataV2Feed.VersionDownloadCount.get -> int?
+NuGet.Protocol.StreamUpdate
 NuGet.Protocol.V2FeedPackageInfo.DownloadCacheDate.get -> System.DateTime?
 NuGet.Protocol.V2FeedPackageInfo.DownloadCacheString.get -> string
 NuGet.Protocol.V2FeedPackageInfo.IsApproved.get -> bool
@@ -188,3 +224,17 @@ override NuGet.Protocol.LocalDependencyInfoResource.ResolvePackage(NuGet.Packagi
 override NuGet.Protocol.LocalDependencyInfoResource.ResolvePackages(string packageId, Chocolatey.NuGet.Frameworks.NuGetFramework projectFramework, NuGet.Protocol.Core.Types.SourceCacheContext sourceCacheContext, NuGet.Common.ILogger log, System.Threading.CancellationToken token) -> System.Threading.Tasks.Task<System.Collections.Generic.IEnumerable<NuGet.Protocol.Core.Types.SourcePackageDependencyInfo>>
 override NuGet.Protocol.Plugins.PluginPackageReader.GetSupportedFrameworks() -> System.Collections.Generic.IEnumerable<Chocolatey.NuGet.Frameworks.NuGetFramework>
 override NuGet.Protocol.Plugins.PluginPackageReader.GetSupportedFrameworksAsync(System.Threading.CancellationToken cancellationToken) -> System.Threading.Tasks.Task<System.Collections.Generic.IEnumerable<Chocolatey.NuGet.Frameworks.NuGetFramework>>
+override NuGet.Protocol.ChocolateyProgressStream.CanRead.get -> bool
+override NuGet.Protocol.ChocolateyProgressStream.CanSeek.get -> bool
+override NuGet.Protocol.ChocolateyProgressStream.CanWrite.get -> bool
+override NuGet.Protocol.ChocolateyProgressStream.Flush() -> void
+override NuGet.Protocol.ChocolateyProgressStream.Length.get -> long
+override NuGet.Protocol.ChocolateyProgressStream.Position.get -> long
+override NuGet.Protocol.ChocolateyProgressStream.Position.set -> void
+override NuGet.Protocol.ChocolateyProgressStream.Read(byte[] buffer, int offset, int count) -> int
+override NuGet.Protocol.ChocolateyProgressStream.Seek(long offset, System.IO.SeekOrigin origin) -> long
+override NuGet.Protocol.ChocolateyProgressStream.SetLength(long value) -> void
+override NuGet.Protocol.ChocolateyProgressStream.Write(byte[] buffer, int offset, int count) -> void
+static NuGet.Protocol.ChocolateyProgressInfo.ShouldDisplayDownloadProgress.get -> bool
+static NuGet.Protocol.ChocolateyProgressInfo.ShouldDisplayDownloadProgress.set -> void
+static NuGet.Protocol.GlobalPackagesFolderUtility.AddPackageAsync(string source, NuGet.Packaging.Core.PackageIdentity packageIdentity, System.IO.Stream packageStream, string globalPackagesFolder, System.Guid parentId, NuGet.Packaging.Signing.ClientPolicyContext clientPolicyContext, NuGet.Common.ILogger logger, System.Threading.CancellationToken token, NuGet.Protocol.ChocolateyProgressInfo progressInfo) -> System.Threading.Tasks.Task<NuGet.Protocol.Core.Types.DownloadResourceResult>