From 0cb7808d41d82fe9f66c80f15cd852b68fbeea71 Mon Sep 17 00:00:00 2001
From: David Fowler <davidfowl@gmail.com>
Date: Thu, 28 Mar 2024 22:29:41 -0700
Subject: [PATCH] Make port allocation global.

---
 .../Publishing/ManifestPublishingContext.cs   | 12 ++---
 .../Utils/ManifestUtils.cs                    | 29 ++++++++++++
 .../Aspire.Hosting.Tests/WithEndpointTests.cs | 47 +++++++++++++++++++
 3 files changed, 82 insertions(+), 6 deletions(-)

diff --git a/src/Aspire.Hosting/Publishing/ManifestPublishingContext.cs b/src/Aspire.Hosting/Publishing/ManifestPublishingContext.cs
index 21065c87b00..feefc158219 100644
--- a/src/Aspire.Hosting/Publishing/ManifestPublishingContext.cs
+++ b/src/Aspire.Hosting/Publishing/ManifestPublishingContext.cs
@@ -31,6 +31,8 @@ public sealed class ManifestPublishingContext(DistributedApplicationExecutionCon
     /// </summary>
     public Utf8JsonWriter Writer { get; } = writer;
 
+    private PortAllocator PortAllocator { get; } = new();
+
     /// <summary>
     /// Gets cancellation token for this operation.
     /// </summary>
@@ -268,8 +270,6 @@ public void WriteBindings(IResource resource)
     {
         if (resource.TryGetEndpoints(out var endpoints))
         {
-            var allocator = new PortAllocator();
-
             Writer.WriteStartObject("bindings");
             foreach (var endpoint in endpoints)
             {
@@ -288,7 +288,7 @@ public void WriteBindings(IResource resource)
                     (ProjectResource, var scheme, null) when scheme is "http" or "https" => null,
 
                     // Allocate a dynamic port
-                    _ => allocator.AllocatePort()
+                    _ => PortAllocator.AllocatePort()
                 };
 
                 int? exposedPort = (endpoint.UriScheme, endpoint.Port, targetPort) switch
@@ -308,18 +308,18 @@ public void WriteBindings(IResource resource)
                     ("https", null, null) => null,
 
                     // Other schemes just allocate a port
-                    _ => allocator.AllocatePort()
+                    _ => PortAllocator.AllocatePort()
                 };
 
                 if (exposedPort is int ep)
                 {
-                    allocator.AddUsedPort(ep);
+                    PortAllocator.AddUsedPort(ep);
                     Writer.WriteNumber("port", ep);
                 }
 
                 if (targetPort is int tp)
                 {
-                    allocator.AddUsedPort(tp);
+                    PortAllocator.AddUsedPort(tp);
                     Writer.WriteNumber("targetPort", tp);
                 }
 
diff --git a/tests/Aspire.Hosting.Tests/Utils/ManifestUtils.cs b/tests/Aspire.Hosting.Tests/Utils/ManifestUtils.cs
index 84e290bd6c0..8dedb192320 100644
--- a/tests/Aspire.Hosting.Tests/Utils/ManifestUtils.cs
+++ b/tests/Aspire.Hosting.Tests/Utils/ManifestUtils.cs
@@ -28,6 +28,35 @@ public static async Task<JsonNode> GetManifest(IResource resource)
         return resourceNode;
     }
 
+    public static async Task<JsonNode[]> GetManifests(IResource[] resources)
+    {
+        using var ms = new MemoryStream();
+        var writer = new Utf8JsonWriter(ms);
+        var executionContext = new DistributedApplicationExecutionContext(DistributedApplicationOperation.Publish);
+        var context = new ManifestPublishingContext(executionContext, Path.Combine(Environment.CurrentDirectory, "manifest.json"), writer);
+
+        var results = new List<JsonNode>();
+
+        foreach (var r in resources)
+        {
+            writer.WriteStartObject();
+            await context.WriteResourceAsync(r);
+            writer.WriteEndObject();
+            writer.Flush();
+            ms.Position = 0;
+            var obj = JsonNode.Parse(ms);
+            Assert.NotNull(obj);
+            var resourceNode = obj[r.Name];
+            Assert.NotNull(resourceNode);
+            results.Add(resourceNode);
+
+            ms.Position = 0;
+            writer.Reset(ms);
+        }
+
+        return [.. results];
+    }
+
     public static async Task<(JsonNode ManifestNode, string BicepText)> GetManifestWithBicep(IResource resource)
     {
         var manifestNode = await GetManifest(resource);
diff --git a/tests/Aspire.Hosting.Tests/WithEndpointTests.cs b/tests/Aspire.Hosting.Tests/WithEndpointTests.cs
index 51c7c675329..4dac7720a69 100644
--- a/tests/Aspire.Hosting.Tests/WithEndpointTests.cs
+++ b/tests/Aspire.Hosting.Tests/WithEndpointTests.cs
@@ -389,6 +389,53 @@ public async Task VerifyManifestProjectWithHttpEndpointDoesNotAllocatePort()
         Assert.Equal(expectedManifest, manifest.ToString());
     }
 
+    [Fact]
+    public async Task VerifyManifestPortAllocationIsGlobal()
+    {
+        using var builder = TestDistributedApplicationBuilder.Create();
+        var container0 = builder.AddContainer("app0", "image")
+                               .WithEndpoint(name: "custom");
+
+        var container1 = builder.AddContainer("app1", "image")
+                               .WithEndpoint(name: "custom");
+
+        var manifests = await ManifestUtils.GetManifests([container0.Resource, container1.Resource]);
+        var expectedManifest0 =
+            """
+            {
+              "type": "container.v0",
+              "image": "image:latest",
+              "bindings": {
+                "custom": {
+                  "scheme": "tcp",
+                  "protocol": "tcp",
+                  "transport": "tcp",
+                  "targetPort": 8000
+                }
+              }
+            }
+            """;
+
+        var expectedManifest1 =
+            """
+            {
+              "type": "container.v0",
+              "image": "image:latest",
+              "bindings": {
+                "custom": {
+                  "scheme": "tcp",
+                  "protocol": "tcp",
+                  "transport": "tcp",
+                  "targetPort": 8001
+                }
+              }
+            }
+            """;
+
+        Assert.Equal(expectedManifest0, manifests[0].ToString());
+        Assert.Equal(expectedManifest1, manifests[1].ToString());
+    }
+
     private static TestProgram CreateTestProgram(string[]? args = null) => TestProgram.Create<WithEndpointTests>(args);
 
     sealed class TestProject : IProjectMetadata