Support parallel upload of blobs (#210)

* parallel upload

* use concurrentdictionary instead

Author: baronfel

Microsoft.NET.Build.Containers/AuthHandshakeMessageHandler.cs

@@ -1,4 +1,5 @@
-using System.Diagnostics.CodeAnalysis;
+using System.Collections.Concurrent;
+using System.Diagnostics.CodeAnalysis;
 using System.Net;
 using System.Net.Http.Headers;
 using System.Text;
@@ -19,7 +20,7 @@ private record AuthInfo(Uri Realm, string Service, string Scope);
     /// <summary>
     /// Cache of most-recently-recieved token for each server.
     /// </summary>
-    private static Dictionary<string, string> TokenCache = new();
+    private static ConcurrentDictionary<string, string> TokenCache = new();
 
     /// <summary>
     /// the www-authenticate header must have realm, service, and scope information, so this method parses it into that shape if present
@@ -112,8 +113,10 @@ private async Task<string> GetTokenAsync(Uri realm, string service, string scope
             throw new ArgumentException("Could not deserialize token from JSON");
         }
 
-        // save the retrieved token in the cache
-        TokenCache[realm.Host] = token.ResolvedToken;
+        // save the retrieved token in the cache.
+        // if we encounter a previous token (very possible due to concurrent upload)
+        // use the more recent token.
+        TokenCache.AddOrUpdate(realm.Host, token.ResolvedToken, (previous, current) => current);
         return token.ResolvedToken;
     }
 

Microsoft.NET.Build.Containers/Registry.cs

@@ -176,19 +176,18 @@ public async Task Push(Image x, string name, string? tag, string baseName, Actio
         tag ??= "latest";
 
         using HttpClient client = GetClient();
-
-        foreach (var descriptor in x.LayerDescriptors)
-        {
+        var reg = this;
+        await Task.WhenAll(x.LayerDescriptors.Select(async descriptor => {
             string digest = descriptor.Digest;
             logProgressMessage($"Uploading layer {digest} to registry");
-            if (await BlobAlreadyUploaded(name, digest, client))
+            if (await reg.BlobAlreadyUploaded(name, digest, client))
             {
                 logProgressMessage($"Layer {digest} already existed");
-                continue;
+                return;
             }
 
             // Blob wasn't there; can we tell the server to get it from the base image?
-            HttpResponseMessage pushResponse = await client.PostAsync(new Uri(BaseUri, $"/v2/{name}/blobs/uploads/?mount={digest}&from={baseName}"), content: null);
+            HttpResponseMessage pushResponse = await client.PostAsync(new Uri(reg.BaseUri, $"/v2/{name}/blobs/uploads/?mount={digest}&from={baseName}"), content: null);
 
             if (pushResponse.StatusCode != HttpStatusCode.Created)
             {
@@ -202,10 +201,10 @@ public async Task Push(Image x, string name, string? tag, string baseName, Actio
                 // Ensure the blob is available locally
                 await x.originatingRegistry.Value.DownloadBlob(x.OriginatingName, descriptor);
                 // Then push it to the destination registry
-                await Push(Layer.FromDescriptor(descriptor), name, logProgressMessage);
+                await reg.Push(Layer.FromDescriptor(descriptor), name, logProgressMessage);
                 logProgressMessage($"Finished uploading layer {digest} to registry");
             }
-        }
+        }));
 
         using (MemoryStream stringStream = new MemoryStream(Encoding.UTF8.GetBytes(x.config.ToJsonString())))
         {
@@ -220,14 +219,12 @@ public async Task Push(Image x, string name, string? tag, string baseName, Actio
         HttpContent manifestUploadContent = new StringContent(x.manifest.ToJsonString());
         manifestUploadContent.Headers.ContentType = new MediaTypeHeaderValue(DockerManifestV2);
         var putResponse = await client.PutAsync(new Uri(BaseUri, $"/v2/{name}/manifests/{manifestDigest}"), manifestUploadContent);
-        string putresponsestr = await putResponse.Content.ReadAsStringAsync();
 
         if (!putResponse.IsSuccessStatusCode)
         {
             string jsonResponse = await putResponse.Content.ReadAsStringAsync();
             throw new ContainerHttpException("Registry push failed.", putResponse.RequestMessage?.RequestUri?.ToString(), jsonResponse);
         }
-
         logProgressMessage($"Uploaded manifest to registry");
 
         logProgressMessage($"Uploading tag {tag} to registry");