Enable replicaSet support for MongoDb

src/Aspire.Hosting.MongoDB/MongoDBBuilderExtensions.cs

@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Text;
 using Aspire.Hosting.ApplicationModel;
 using Aspire.Hosting.MongoDB;
 using Aspire.Hosting.Utils;
@@ -94,7 +95,8 @@ public static IResourceBuilder<T> WithMongoExpress<T>(this IResourceBuilder<T> b
                                                         .WithImageRegistry(MongoDBContainerImageTags.MongoExpressRegistry)
                                                         .WithEnvironment(context => ConfigureMongoExpressContainer(context, builder.Resource))
                                                         .WithHttpEndpoint(targetPort: 8081, name: "http")
-                                                        .ExcludeFromManifest();
+                                                        .ExcludeFromManifest()
+                                                        .WaitFor(builder);
 
         configureContainer?.Invoke(resourceBuilder);
 
@@ -161,11 +163,96 @@ public static IResourceBuilder<MongoDBServerResource> WithInitBindMount(this IRe
         return builder.WithBindMount(source, "/docker-entrypoint-initdb.d", isReadOnly);
     }
 
+    /// <summary>
+    /// Adds a replica set to the MongoDB server resource.
+    /// </summary>
+    /// <param name="builder">The MongoDB server resource.</param>
+    /// <param name="replicaSetName">The name of the replica set. If not provided, defaults to <c>rs0</c>.</param>
+    /// <returns>A reference to the <see cref="IResourceBuilder{T}"/>.</returns>
+    public static IResourceBuilder<MongoDBServerResource> WithReplicaSet(this IResourceBuilder<MongoDBServerResource> builder, string? replicaSetName = null)
+    {
+        if (builder.Resource.TryGetLastAnnotation<MongoDbReplicaSetAnnotation>(out _))
+        {
+            throw new InvalidOperationException("A replica set has already been added to the MongoDB server resource.");
+        }
+
+        replicaSetName ??= "rs0";
+
+        var port = SetPortAndTargetToBeSame(builder);
+
+        // Add a container that initializes the replica set
+        var init = builder.ApplicationBuilder
+            .AddDockerfile("replicaset-init", GetReplicaSetInitDockerfileDir(replicaSetName, builder.Resource.Name, port))
+
+            // We don't want to wait for the healthchecks to be successful since the initialization is required for that. However, we also don't want this to start
+            // up until the database itself is ready
+            .WaitFor(builder, includeHealthChecks: false);
+
+        return builder
+            .WithAnnotation(new MongoDbReplicaSetAnnotation(replicaSetName, init))
+            .WithArgs("--replSet", replicaSetName, "--bind_ip_all", "--port", $"{port}");
+
+        static int SetPortAndTargetToBeSame(IResourceBuilder<MongoDBServerResource> builder)
+        {
+            foreach (var endpoint in builder.Resource.Annotations.OfType<EndpointAnnotation>())
+            {
+                if (endpoint.Name == MongoDBServerResource.PrimaryEndpointName)
+                {
+                    if (endpoint.Port is { } port)
+                    {
+                        endpoint.TargetPort = port;
+                    }
+
+                    if (endpoint.TargetPort is not { } targetPort)
+                    {
+                        throw new InvalidOperationException("Target port is not set.");
+                    }
+
+                    // In the case of replica sets, the port and target port should be the same and is not proxied
+                    endpoint.IsProxied = false;
+
+                    return targetPort;
+                }
+            }
+
+            throw new InvalidOperationException("No endpoint found for the MongoDB server resource.");
+        }
+
+        // See the conversation about setting up replica sets in Docker here: https://github.com/docker-library/mongo/issues/246
+        static string GetReplicaSetInitDockerfileDir(string replicaSet, string host, int port)
+        {
+            var dir = Directory.CreateTempSubdirectory("aspire.mongo").FullName;
+
+            var rsInitContents = $$"""rs.initiate({ _id:'{{replicaSet}}', members:[{_id:0,host:'localhost:{{port}}'}]})""";
+            var init = Path.Combine(dir, "rs.js");
+            File.WriteAllText(init, rsInitContents);
+
+            var dockerfile = Path.Combine(dir, "Dockerfile");
+            File.WriteAllText(dockerfile, $"""
+                FROM {MongoDBContainerImageTags.Image}:{MongoDBContainerImageTags.Tag}
+                WORKDIR /rsinit
+                ADD rs.js rs.js
+                ENTRYPOINT ["mongosh", "--port", "{port}", "--host", "{host}", "rs.js"]
+                """);
+            return dir;
+        }
+    }
+
     private static void ConfigureMongoExpressContainer(EnvironmentCallbackContext context, MongoDBServerResource resource)
     {
-        // Mongo Exporess assumes Mongo is being accessed over a default Aspire container network and hardcodes the resource address
+        var sb = new StringBuilder($"mongodb://{resource.Name}:{resource.PrimaryEndpoint.TargetPort}/?directConnection=true");
+
+        if (resource.TryGetLastAnnotation<MongoDbReplicaSetAnnotation>(out var replica))
+        {
+            sb.Append('&');
+            sb.Append(MongoDbReplicaSetAnnotation.QueryName);
+            sb.Append('=');
+            sb.Append(replica.ReplicaSetName);
+        }
+
+        // Mongo Express assumes Mongo is being accessed over a default Aspire container network and hardcodes the resource address
         // This will need to be refactored once updated service discovery APIs are available
-        context.EnvironmentVariables.Add("ME_CONFIG_MONGODB_URL", $"mongodb://{resource.Name}:{resource.PrimaryEndpoint.TargetPort}/?directConnection=true");
+        context.EnvironmentVariables.Add("ME_CONFIG_MONGODB_URL", sb.ToString());
         context.EnvironmentVariables.Add("ME_CONFIG_BASICAUTH", "false");
     }
 }

src/Aspire.Hosting.MongoDB/MongoDBDatabaseResource.cs

@@ -27,7 +27,17 @@ public MongoDBDatabaseResource(string name, string databaseName, MongoDBServerRe
     /// Gets the connection string expression for the MongoDB database.
     /// </summary>
     public ReferenceExpression ConnectionStringExpression
-        => ReferenceExpression.Create($"{Parent}/{DatabaseName}");
+    {
+        get
+        {
+            var builder = new ReferenceExpressionBuilder();
+
+            Parent.AppendConnectionString(builder);
+            Parent.AppendSuffix(builder, DatabaseName);
+
+            return builder.Build();
+        }
+    }
 
     /// <summary>
     /// Gets the parent MongoDB container resource.

src/Aspire.Hosting.MongoDB/MongoDBServerResource.cs

@@ -21,9 +21,59 @@ public class MongoDBServerResource(string name) : ContainerResource(name), IReso
     /// <summary>
     /// Gets the connection string for the MongoDB server.
     /// </summary>
-    public ReferenceExpression ConnectionStringExpression =>
-        ReferenceExpression.Create(
-            $"mongodb://{PrimaryEndpoint.Property(EndpointProperty.Host)}:{PrimaryEndpoint.Property(EndpointProperty.Port)}");
+    public ReferenceExpression ConnectionStringExpression
+    {
+        get
+        {
+            var builder = new ReferenceExpressionBuilder();
+
+            AppendConnectionString(builder);
+            AppendSuffix(builder);
+
+            return builder.Build();
+        }
+    }
+
+    internal void AppendConnectionString(ReferenceExpressionBuilder builder)
+    {
+        builder.AppendLiteral("mongodb://");
+        builder.AppendFormatted(PrimaryEndpoint.Property(EndpointProperty.Host));
+        builder.AppendLiteral(":");
+        builder.AppendFormatted(PrimaryEndpoint.Property(EndpointProperty.Port));
+    }
+
+    /// <summary>
+    /// Handles adding the rest of the connection string.
+    ///
+    /// - If a database name is provided, it will be appended to the connection string.
+    /// - If a replica set name is provided, it will be appended to the connection string.
+    /// - If no database but a replica set is provided, a '/' must be inserted before the '?'
+    /// </summary>
+    internal bool AppendSuffix(ReferenceExpressionBuilder builder, string? dbName = null)
+    {
+        if (dbName is { })
+        {
+            builder.AppendLiteral("/");
+            builder.AppendFormatted(dbName);
+        }
+
+        if (Annotations.OfType<MongoDbReplicaSetAnnotation>().FirstOrDefault() is { ReplicaSetName: { } replicaSetName })
+        {
+            if (dbName is null)
+            {
+                builder.AppendLiteral("/");
+            }
+
+            builder.AppendLiteral("?");
+            builder.AppendLiteral(MongoDbReplicaSetAnnotation.QueryName);
+            builder.AppendLiteral("=");
+            builder.AppendLiteral(replicaSetName);
+
+            return true;
+        }
+
+        return false;
+    }
 
     private readonly Dictionary<string, string> _databases = new Dictionary<string, string>(StringComparers.ResourceName);
 

src/Aspire.Hosting.MongoDB/MongoDbReplicaSetAnnotation.cs

@@ -0,0 +1,9 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+namespace Aspire.Hosting.ApplicationModel;
+
+internal sealed record MongoDbReplicaSetAnnotation(string ReplicaSetName, IResourceBuilder<ContainerResource> InitContainer) : IResourceAnnotation
+{
+    internal const string QueryName = "replicaSet";
+}

src/Aspire.Hosting.MongoDB/PublicAPI.Unshipped.txt

@@ -1,2 +1,2 @@
 #nullable enable
-
+static Aspire.Hosting.MongoDBBuilderExtensions.WithReplicaSet(this Aspire.Hosting.ApplicationModel.IResourceBuilder<Aspire.Hosting.ApplicationModel.MongoDBServerResource!>! builder, string? replicaSetName = null) -> Aspire.Hosting.ApplicationModel.IResourceBuilder<Aspire.Hosting.ApplicationModel.MongoDBServerResource!>!

src/Aspire.Hosting/PublicAPI.Unshipped.txt

@@ -86,6 +86,7 @@ Aspire.Hosting.DistributedApplicationExecutionContextOptions.Operation.get -> As
 Aspire.Hosting.DistributedApplicationExecutionContextOptions.ServiceProvider.get -> System.IServiceProvider?
 Aspire.Hosting.DistributedApplicationExecutionContextOptions.ServiceProvider.set -> void
 static Aspire.Hosting.ResourceBuilderExtensions.WaitFor<T>(this Aspire.Hosting.ApplicationModel.IResourceBuilder<T>! builder, Aspire.Hosting.ApplicationModel.IResourceBuilder<Aspire.Hosting.ApplicationModel.IResource!>! dependency) -> Aspire.Hosting.ApplicationModel.IResourceBuilder<T>!
+static Aspire.Hosting.ResourceBuilderExtensions.WaitFor<T>(this Aspire.Hosting.ApplicationModel.IResourceBuilder<T>! builder, Aspire.Hosting.ApplicationModel.IResourceBuilder<Aspire.Hosting.ApplicationModel.IResource!>! dependency, bool includeHealthChecks) -> Aspire.Hosting.ApplicationModel.IResourceBuilder<T>!
 static Aspire.Hosting.ResourceBuilderExtensions.WaitForCompletion<T>(this Aspire.Hosting.ApplicationModel.IResourceBuilder<T>! builder, Aspire.Hosting.ApplicationModel.IResourceBuilder<Aspire.Hosting.ApplicationModel.IResource!>! dependency, int exitCode = 0) -> Aspire.Hosting.ApplicationModel.IResourceBuilder<T>!
 static Aspire.Hosting.ResourceBuilderExtensions.WithHealthCheck<T>(this Aspire.Hosting.ApplicationModel.IResourceBuilder<T>! builder, string! key) -> Aspire.Hosting.ApplicationModel.IResourceBuilder<T>!
 static readonly Aspire.Hosting.ApplicationModel.KnownResourceStates.Exited -> string!

src/Aspire.Hosting/ResourceBuilderExtensions.cs

@@ -596,6 +596,37 @@ public static IResourceBuilder<T> ExcludeFromManifest<T>(this IResourceBuilder<T
     /// </code>
     /// </example>
     public static IResourceBuilder<T> WaitFor<T>(this IResourceBuilder<T> builder, IResourceBuilder<IResource> dependency) where T : IResource
+        => builder.WaitFor(dependency, includeHealthChecks: true);
+
+    /// <summary>
+    /// Waits for the dependency resource to enter the Running state before starting the resource.
+    /// </summary>
+    /// <typeparam name="T">The type of the resource.</typeparam>
+    /// <param name="builder">The resource builder for the resource that will be waiting.</param>
+    /// <param name="dependency">The resource builder for the dependency resource.</param>
+    /// <param name="includeHealthChecks">Optionally includes checking the health checks.</param>
+    /// <returns>The resource builder.</returns>
+    /// <remarks>
+    /// <para>This method is useful when a resource should wait until another has started running. This can help
+    /// reduce errors in logs during local development where dependency resources.</para>
+    /// <para>Some resources automatically register health checks with the application host container. For these
+    /// resources, calling <see cref="WaitFor{T}(IResourceBuilder{T}, IResourceBuilder{IResource})"/> also results
+    /// in the resource being blocked from starting until the health checks associated with the dependency resource
+    /// return <see cref="HealthStatus.Healthy"/>.</para>
+    /// <para>The <see cref="WithHealthCheck{T}(IResourceBuilder{T}, string)"/> method can be used to associate
+    /// additional health checks with a resource.</para>
+    /// </remarks>
+    /// <example>
+    /// Start message queue before starting the worker service.
+    /// <code lang="C#">
+    /// var builder = DistributedApplication.CreateBuilder(args);
+    /// var messaging = builder.AddRabbitMQ("messaging");
+    /// builder.AddProject&lt;Projects.MyApp&gt;("myapp")
+    ///        .WithReference(messaging)
+    ///        .WaitFor(messaging);
+    /// </code>
+    /// </example>
+    public static IResourceBuilder<T> WaitFor<T>(this IResourceBuilder<T> builder, IResourceBuilder<IResource> dependency, bool includeHealthChecks) where T : IResource
     {
         builder.ApplicationBuilder.Eventing.Subscribe<BeforeResourceStartedEvent>(builder.Resource, async (e, ct) =>
         {
@@ -632,7 +663,7 @@ public static IResourceBuilder<T> WaitFor<T>(this IResourceBuilder<T> builder, I
 
             // If our dependency resource has health check annotations we want to wait until they turn healthy
             // otherwise we don't care about their health status.
-            if (dependency.Resource.TryGetAnnotationsOfType<HealthCheckAnnotation>(out var _))
+            if (includeHealthChecks && dependency.Resource.TryGetAnnotationsOfType<HealthCheckAnnotation>(out var _))
             {
                 resourceLogger.LogInformation("Waiting for resource '{Name}' to become healthy.", dependency.Resource.Name);
                 await rns.WaitForResourceAsync(dependency.Resource.Name, re => re.Snapshot.HealthStatus == HealthStatus.Healthy, cancellationToken: ct).ConfigureAwait(false);