Fix event order and callback calling during container creation

Author: karolz-ms

src/Aspire.Hosting/ApplicationModel/ArgumentsExecutionConfigurationGatherer.cs

@@ -13,18 +13,24 @@ internal class ArgumentsExecutionConfigurationGatherer : IExecutionConfiguration
     /// <inheritdoc/>
     public async ValueTask GatherAsync(IExecutionConfigurationGathererContext context, IResource resource, ILogger resourceLogger, DistributedApplicationExecutionContext executionContext, CancellationToken cancellationToken = default)
     {
-        if (resource.TryGetAnnotationsOfType<CommandLineArgsCallbackAnnotation>(out var callbacks))
+        if (resource.TryGetAnnotationsOfType<CommandLineArgsCallbackAnnotation>(out var argumentAnnotations))
         {
-            var callbackContext = new CommandLineArgsCallbackContext(context.Arguments, resource, cancellationToken)
+            IList<object> args = [.. context.Arguments];
+            var callbackContext = new CommandLineArgsCallbackContext(args, resource, cancellationToken)
             {
                 Logger = resourceLogger,
                 ExecutionContext = executionContext
             };
 
-            foreach (var callback in callbacks)
+            foreach (var ann in argumentAnnotations)
             {
-                await callback.Callback(callbackContext).ConfigureAwait(false);
+                // Each annotation operates on a shared context.
+                args = await ann.AsCallbackAnnotation().EvaluateOnceAsync(callbackContext).ConfigureAwait(false);
             }
+
+            // Take the final result and apply to the gatherer context.
+            context.Arguments.Clear();
+            context.Arguments.AddRange(args);
         }
     }
 }

src/Aspire.Hosting/ApplicationModel/CommandLineArgsCallbackAnnotation.cs

@@ -1,16 +1,22 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Collections.Immutable;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Logging.Abstractions;
 
 namespace Aspire.Hosting.ApplicationModel;
 
+using IArgCallbackAnnotation = ICallbackResourceAnnotation<CommandLineArgsCallbackContext, IList<object>>;
+
 /// <summary>
 /// Represents an annotation that provides a callback to be executed with a list of command-line arguments when an executable resource is started.
 /// </summary>
-public class CommandLineArgsCallbackAnnotation : IResourceAnnotation
+public class CommandLineArgsCallbackAnnotation : IResourceAnnotation, IArgCallbackAnnotation
 {
+    private Task<IList<object>>? _callbackTask;
+    private readonly object _lock = new();
+
     /// <summary>
     /// Initializes a new instance of the <see cref="CommandLineArgsCallbackAnnotation"/> class with the specified callback action.
     /// </summary>
@@ -41,6 +47,35 @@ public CommandLineArgsCallbackAnnotation(Action<IList<object>> callback)
     /// Gets the callback action to be executed when the executable arguments are parsed.
     /// </summary>
     public Func<CommandLineArgsCallbackContext, Task> Callback { get; }
+
+    internal IArgCallbackAnnotation AsCallbackAnnotation() => this;
+
+    Task<IList<object>> IArgCallbackAnnotation.EvaluateOnceAsync(CommandLineArgsCallbackContext context)
+    {
+        lock(_lock)
+        {
+            if (_callbackTask is null)
+            {
+                _callbackTask = ExecuteCallbackAsync(context);
+            }
+            return _callbackTask;
+        }
+    }
+
+    void IArgCallbackAnnotation.ForgetCachedResult()
+    {
+        lock(_lock)
+        {
+            _callbackTask = null;
+        }
+    }
+
+    private async Task<IList<object>> ExecuteCallbackAsync(CommandLineArgsCallbackContext context)
+    {
+        await Callback(context).ConfigureAwait(false);
+        var result = context.Args.ToImmutableList();
+        return result;
+    }
 }
 
 /// <summary>

src/Aspire.Hosting/ApplicationModel/EnvironmentCallbackAnnotation.cs

@@ -5,13 +5,17 @@
 
 namespace Aspire.Hosting.ApplicationModel;
 
+using IEnvCallbackAnnotation = ICallbackResourceAnnotation<EnvironmentCallbackContext, Dictionary<string, object>>;
+
 /// <summary>
 /// Represents an annotation that provides a callback to modify the environment variables of an application.
 /// </summary>
 [DebuggerDisplay("{DebuggerToString(),nq}")]
-public class EnvironmentCallbackAnnotation : IResourceAnnotation
+public class EnvironmentCallbackAnnotation : IResourceAnnotation, IEnvCallbackAnnotation
 {
     private readonly string? _name;
+    private Task<Dictionary<string, object>>? _callbackTask;
+    private readonly object _lock = new();
 
     /// <summary>
     /// Initializes a new instance of the <see cref="EnvironmentCallbackAnnotation"/> class with the specified name and callback function.
@@ -77,6 +81,35 @@ public EnvironmentCallbackAnnotation(Func<EnvironmentCallbackContext, Task> call
     /// </summary>
     public Func<EnvironmentCallbackContext, Task> Callback { get; private set; }
 
+    internal IEnvCallbackAnnotation AsCallbackAnnotation() => this;
+
+    Task<Dictionary<string, object>> IEnvCallbackAnnotation.EvaluateOnceAsync(EnvironmentCallbackContext context)
+    {
+        lock(_lock)
+        {
+            if (_callbackTask is null)
+            {
+                _callbackTask = ExecuteCallbackAsync(context);
+            }
+            return _callbackTask;
+        }
+    }
+
+    void IEnvCallbackAnnotation.ForgetCachedResult()
+    {
+        lock(_lock)
+        {
+            _callbackTask = null;
+        }
+    }
+
+    private async Task<Dictionary<string, object>> ExecuteCallbackAsync(EnvironmentCallbackContext context)
+    {
+        await Callback(context).ConfigureAwait(false);
+        var result = new Dictionary<string, object>(context.EnvironmentVariables);
+        return result;
+    }
+
     private string DebuggerToString()
     {
         var text = $@"Type = {GetType().Name}";

src/Aspire.Hosting/ApplicationModel/EnvironmentVariablesConfigurationGatherer.cs

@@ -13,16 +13,24 @@ internal class EnvironmentVariablesExecutionConfigurationGatherer : IExecutionCo
     /// <inheritdoc/>
     public async ValueTask GatherAsync(IExecutionConfigurationGathererContext context, IResource resource, ILogger resourceLogger, DistributedApplicationExecutionContext executionContext, CancellationToken cancellationToken = default)
     {
-        if (resource.TryGetEnvironmentVariables(out var callbacks))
+        if (resource.TryGetEnvironmentVariables(out var envVarAnnotations))
         {
-            var callbackContext = new EnvironmentCallbackContext(executionContext, resource, context.EnvironmentVariables, cancellationToken)
+            var envVars = new Dictionary<string, object>(context.EnvironmentVariables);
+            var callbackContext = new EnvironmentCallbackContext(executionContext, resource, envVars, cancellationToken: cancellationToken)
             {
                 Logger = resourceLogger,
             };
 
-            foreach (var callback in callbacks)
+            foreach (var ann in envVarAnnotations)
             {
-                await callback.Callback(callbackContext).ConfigureAwait(false);
+                // Each annotation operates on a shared context.
+                envVars = await ann.AsCallbackAnnotation().EvaluateOnceAsync(callbackContext).ConfigureAwait(false);
+            }
+
+            // Take the final result and apply to the gatherer context.
+            foreach (var kvp in envVars)
+            {
+                context.EnvironmentVariables[kvp.Key] = kvp.Value;
             }
         }
     }

src/Aspire.Hosting/ApplicationModel/ICallbackResourceAnnotation.cs

@@ -0,0 +1,29 @@
+// 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;
+
+/// <summary>
+/// Represents a resource annotation whose callback should be evaluated at most once,
+/// with the result cached for subsequent retrievals.
+/// </summary>
+/// <typeparam name="TContext">The type of the context passed to the callback.</typeparam>
+/// <typeparam name="TResult">The type of the result produced by the callback.</typeparam>
+internal interface ICallbackResourceAnnotation<TContext, TResult>
+{
+    /// <summary>
+    /// Evaluates the callback if it has not been evaluated yet, caching the result.
+    /// Subsequent calls return the cached result regardless of the context passed.
+    /// </summary>
+    /// <param name="context">The context for the callback evaluation. Only used on the first call.</param>
+    /// <returns>The cached result of the callback evaluation.</returns>
+    Task<TResult> EvaluateOnceAsync(TContext context);
+
+    /// <summary>
+    /// Clears the cached result so that the next call to <see cref="EvaluateOnceAsync"/> will re-execute the callback. 
+    ///</summary>
+    /// <remarks>
+    /// Use <see cref="ForgetCachedResult"/> when a resource decorated with this callback annotation is restarted.
+    /// </remarks>
+    void ForgetCachedResult();
+}

src/Aspire.Hosting/ApplicationModel/ResourceDependencyDiscoveryMode.cs

@@ -4,7 +4,7 @@
 namespace Aspire.Hosting.ApplicationModel;
 
 /// <summary>
-/// Specifies how resource dependencies are discovered.
+/// Represents the mode for discovering resource dependencies.
 /// </summary>
 public enum ResourceDependencyDiscoveryMode
 {
@@ -20,5 +20,5 @@ public enum ResourceDependencyDiscoveryMode
     /// and from environment variables and command-line arguments, but does not recurse
     /// into the dependencies of those dependencies.
     /// </summary>
-    DirectOnly
+    DirectOnly,
 }

src/Aspire.Hosting/ApplicationModel/ResourceDependencyDiscoveryOptions.cs

@@ -0,0 +1,22 @@
+// 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;
+
+/// <summary>
+/// Changes how resource dependencies are discovered.
+/// </summary>
+public sealed class ResourceDependencyDiscoveryOptions
+{
+    /// <summary>
+    /// Sets the mode for discovering resource dependencies. See <see cref="ResourceDependencyDiscoveryMode"/> for details on the available modes.
+    /// 
+    /// </summary>
+    public ResourceDependencyDiscoveryMode DiscoveryMode { get; init; }
+
+    /// <summary>
+    /// When true, unresolved values from annotation callbacks will be cached and reused 
+    /// on subsequent evaluations of the same annotation, rather than re-evaluating the callback each time.
+    /// </summary>
+    public bool CacheAnnotationCallbackResults { get; init; }
+}