// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
using Aspire.Hosting.ApplicationModel;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;
using MongoDB.Bson;
using MongoDB.Driver;
#pragma warning disable ASPIREMONGODB001
namespace Aspire.Hosting.MongoDB;
internal sealed class MongoDBSingleMemberReplicaSetAnnotation(string name) : IResourceAnnotation
{
public string Name { get; } = name;
// Initialization and health polling run concurrently. Null means initialization and election succeeded.
public volatile string? InitializationError = "The single-member replica set has not been initialized.";
}
internal static class MongoDBSingleMemberReplicaSet
{
internal static async Task InitializeAsync(
MongoDBServerResource resource,
MongoDBSingleMemberReplicaSetAnnotation annotation,
InitializeResourceEvent evt,
CancellationToken cancellationToken)
{
using var stopping = CancellationTokenSource.CreateLinkedTokenSource(
cancellationToken, evt.Services.GetRequiredService<IHostApplicationLifetime>().ApplicationStopping);
try
{
var wasRunning = false;
// InitializeResourceEvent is dispatched non-blocking, so waiting for Running does not block container
// startup. ResourceReadyEvent would deadlock: primary election is a prerequisite for this resource's health.
// Observe start transitions too, so restarting a failed or recreated container retries initialization.
await foreach (var update in evt.Notifications.WatchAsync(stopping.Token).ConfigureAwait(false))
{
if (update.Resource != resource)
{
continue;
}
var isRunning = update.Snapshot.State?.Text == KnownResourceStates.Running;
if (isRunning && !wasRunning)
{
annotation.InitializationError = "Waiting for single-member replica set initialization and primary election.";
using var timeout = CancellationTokenSource.CreateLinkedTokenSource(stopping.Token);
timeout.CancelAfter(TimeSpan.FromSeconds(90));
try
{
var connectionString = await resource.ConnectionStringExpression.GetValueAsync(timeout.Token).ConfigureAwait(false)
?? throw new DistributedApplicationException($"The connection string for MongoDB resource '{resource.Name}' is unavailable.");
var settings = MongoClientSettings.FromConnectionString(connectionString);
settings.ServerSelectionTimeout = TimeSpan.FromSeconds(2);
settings.ConnectTimeout = TimeSpan.FromSeconds(2);
using var client = new MongoClient(settings);
var database = client.GetDatabase(MongoDBServerResource.DefaultAuthenticationDatabase);
// Loopback is stable across container and AppHost restarts, including random host-port changes.
// Clients use directConnection=true and never try to discover or connect to this internal address.
var host = $"localhost:{resource.PrimaryEndpoint.TargetPort}";
await InitializeAndWaitForPrimaryAsync(database, annotation.Name, host, timeout.Token).ConfigureAwait(false);
annotation.InitializationError = null;
evt.Logger.LogInformation("MongoDB resource '{ResourceName}' is the primary of single-member replica set '{ReplicaSetName}'.", resource.Name, annotation.Name);
}
catch (OperationCanceledException) when (!stopping.IsCancellationRequested)
{
annotation.InitializationError = $"MongoDB resource '{resource.Name}' did not initialize replica set '{annotation.Name}' and elect a primary within 90 seconds. Check the container logs, credentials, and existing data volume, then restart the resource.";
evt.Logger.LogError("{Message}", annotation.InitializationError);
}
catch (Exception ex) when (ex is MongoException or TimeoutException or DistributedApplicationException)
{
annotation.InitializationError = $"MongoDB resource '{resource.Name}' could not initialize replica set '{annotation.Name}'. Check the container logs and ensure the data volume belongs to this single-member set and uses the configured credentials.";
evt.Logger.LogError(ex, "{Message}", annotation.InitializationError);
}
}
wasRunning = isRunning;
}
}
catch (OperationCanceledException) when (stopping.IsCancellationRequested)
{
// AppHost shutdown cancels both the notification watch and any in-flight MongoDB command.
}
}
internal static async Task InitializeAndWaitForPrimaryAsync(IMongoDatabase database, string name, string host, CancellationToken cancellationToken)
{
while (true)
{
cancellationToken.ThrowIfCancellationRequested();
try
{
BsonDocument config;
try
{
var result = await database.RunCommandAsync<BsonDocument>(
new BsonDocument("replSetGetConfig", 1), ReadPreference.Nearest, cancellationToken).ConfigureAwait(false);
config = result["config"].AsBsonDocument;
}
catch (MongoCommandException ex) when (ex.CodeName == "NotYetInitialized")
{
// The official image creates the authenticated root user using a temporary standalone mongod.
// Initiation must target the final server, not an init script against that temporary process.
// https://github.com/docker-library/mongo/blob/master/docker-entrypoint.sh
await database.RunCommandAsync<BsonDocument>(
new BsonDocument("replSetInitiate", new BsonDocument
{
["_id"] = name,
["members"] = new BsonArray { new BsonDocument { ["_id"] = 0, ["host"] = host } },
}), ReadPreference.Nearest, cancellationToken).ConfigureAwait(false);
continue;
}
ValidateConfiguration(config, name, host);
var status = await database.RunCommandAsync<BsonDocument>(
new BsonDocument("replSetGetStatus", 1), ReadPreference.Nearest, cancellationToken).ConfigureAwait(false);
if (IsPrimary(status, name))
{
return;
}
}
catch (MongoCommandException ex) when (ex.CodeName == "AlreadyInitialized")
{
// Another initializer won the race. Read and validate its configuration; never force a reconfiguration.
}
catch (Exception ex) when (ex is MongoConnectionException or TimeoutException)
{
// Running means the container started, not that mongod is already accepting authenticated commands.
}
await Task.Delay(TimeSpan.FromSeconds(1), cancellationToken).ConfigureAwait(false);
}
}
internal static void ValidateConfiguration(BsonDocument config, string name, string host)
{
// replSetGetConfig returns { config: { _id: "mongo", members: [{ _id: 0, host: "localhost:27017", ... }], ... } }.
// Ignore server-populated defaults but refuse to rewrite existing identities, addresses, or multi-member data.
if (config["_id"].AsString != name ||
config["members"].AsBsonArray is not { Count: 1 } members ||
members[0]["host"].AsString != host)
{
throw new DistributedApplicationException($"The existing MongoDB replica set configuration does not match single-member set '{name}' at '{host}'. Use its original configuration and data volume; automatic migration or reconfiguration is not supported.");
}
}
internal static bool IsPrimary(BsonDocument status, string name) =>
status["set"].AsString == name && status["myState"].AsInt32 == 1 && status["members"].AsBsonArray.Count == 1;
}