Make more things async

playground/CustomResources/CustomResources.AppHost/TestResource.cs

@@ -44,7 +44,7 @@ public Task BeforeStartAsync(DistributedApplicationModel appModel, CancellationT
                 Task.Run(async () =>
                 {
                     // Simulate custom resource state changes
-                    var state = resourceUpdates.GetInitialSnapshot();
+                    var state = await resourceUpdates.GetInitialSnapshotAsync(_tokenSource.Token);
                     var seconds = Random.Shared.Next(2, 12);
 
                     state = state with

src/Aspire.Hosting.Dapr/DaprDistributedApplicationLifecycleHook.cs

@@ -180,7 +180,7 @@ public async Task BeforeStartAsync(DistributedApplicationModel appModel, Cancell
             }
 
             daprCli.Annotations.Add(
-                new ExecutableArgsCallbackAnnotation(
+                new CommandLineArgsCallbackAnnotation(
                     updatedArgs =>
                     {
                         AllocatedEndpointAnnotation? httpEndPoint = null;

src/Aspire.Hosting/ApplicationModel/CommandLineArgsCallbackAnnotation.cs

@@ -0,0 +1,55 @@
+// 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 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
+{
+    /// <summary>
+    /// Initializes a new instance of the <see cref="CommandLineArgsCallbackAnnotation"/> class with the specified callback action.
+    /// </summary>
+    /// <param name="callback"> The callback action to be executed.</param>
+    public CommandLineArgsCallbackAnnotation(Func<CommandLineArgsCallbackContext, Task> callback)
+    {
+        Callback = callback;
+    }
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="CommandLineArgsCallbackAnnotation"/> class with the specified callback action.
+    /// </summary>
+    /// <param name="callback"> The callback action to be executed.</param>
+    public CommandLineArgsCallbackAnnotation(Action<IList<string>> callback)
+    {
+        Callback = (c) =>
+        {
+            callback(c.Args);
+            return Task.CompletedTask;
+        };
+    }
+
+    /// <summary>
+    /// Gets the callback action to be executed when the executable arguments are parsed.
+    /// </summary>
+    public Func<CommandLineArgsCallbackContext, Task> Callback { get; }
+}
+
+/// <summary>
+/// Represents a callback context for the list of command-line arguments associated with an executable resource.
+/// </summary>
+/// <param name="args"> The list of command-line arguments.</param>
+/// <param name="cancellationToken"> The cancellation token associated with this execution.</param>
+public sealed class CommandLineArgsCallbackContext(IList<string> args, CancellationToken cancellationToken = default)
+{
+    /// <summary>
+    /// Gets the list of command-line arguments.
+    /// </summary>
+    public IList<string> Args { get; } = args;
+
+    /// <summary>
+    /// Gets the cancellation token associated with the callback context.
+    /// </summary>
+    public CancellationToken CancellationToken { get; } = cancellationToken;
+}

src/Aspire.Hosting/ApplicationModel/CustomResourceExtensions.cs

@@ -17,14 +17,27 @@ public static class CustomResourceExtensions
     /// <param name="builder">The resource builder.</param>
     /// <param name="initialSnapshotFactory">The factory to create the initial <see cref="CustomResourceSnapshot"/> for this resource.</param>
     /// <returns>The resource builder.</returns>
-    public static IResourceBuilder<TResource> WithResourceUpdates<TResource>(this IResourceBuilder<TResource> builder, Func<CustomResourceSnapshot>? initialSnapshotFactory = null)
+    public static IResourceBuilder<TResource> WithResourceUpdates<TResource>(this IResourceBuilder<TResource> builder, Func<CancellationToken, ValueTask<CustomResourceSnapshot>>? initialSnapshotFactory = null)
         where TResource : IResource
     {
-        initialSnapshotFactory ??= () => CustomResourceSnapshot.Create(builder.Resource);
+        initialSnapshotFactory ??= cancellationToken => CustomResourceSnapshot.CreateAsync(builder.Resource, cancellationToken);
 
         return builder.WithAnnotation(new ResourceUpdatesAnnotation(initialSnapshotFactory), ResourceAnnotationMutationBehavior.Replace);
     }
 
+    /// <summary>
+    /// Initializes the resource with a <see cref="ResourceUpdatesAnnotation"/> that allows publishing and subscribing to changes in the state of this resource.
+    /// </summary>
+    /// <typeparam name="TResource">The resource.</typeparam>
+    /// <param name="builder">The resource builder.</param>
+    /// <param name="initialSnapshotFactory">The factory to create the initial <see cref="CustomResourceSnapshot"/> for this resource.</param>
+    /// <returns>The resource builder.</returns>
+    public static IResourceBuilder<TResource> WithResourceUpdates<TResource>(this IResourceBuilder<TResource> builder, Func<CustomResourceSnapshot> initialSnapshotFactory)
+        where TResource : IResource
+    {
+        return builder.WithAnnotation(new ResourceUpdatesAnnotation(_ => ValueTask.FromResult(initialSnapshotFactory())), ResourceAnnotationMutationBehavior.Replace);
+    }
+
     /// <summary>
     /// Initializes the resource with a logger that writes to the log stream for the resource.
     /// </summary>

src/Aspire.Hosting/ApplicationModel/EnvironmentCallbackAnnotation.cs

@@ -21,7 +21,11 @@ public class EnvironmentCallbackAnnotation : IResourceAnnotation
     public EnvironmentCallbackAnnotation(string name, Func<string> callback)
     {
         _name = name;
-        Callback = (c) => c.EnvironmentVariables[name] = callback();
+        Callback = (c) =>
+        {
+            c.EnvironmentVariables[name] = callback();
+            return Task.CompletedTask;
+        };
     }
 
     /// <summary>
@@ -30,22 +34,39 @@ public EnvironmentCallbackAnnotation(string name, Func<string> callback)
     /// <param name="callback">The callback action to be executed.</param>
     public EnvironmentCallbackAnnotation(Action<Dictionary<string, string>> callback)
     {
-        Callback = (c) => callback(c.EnvironmentVariables);
+        Callback = (c) =>
+        {
+            callback(c.EnvironmentVariables);
+            return Task.CompletedTask;
+        };
     }
 
     /// <summary>
     /// Initializes a new instance of the <see cref="EnvironmentCallbackAnnotation"/> class with the specified callback.
     /// </summary>
     /// <param name="callback">The callback to be invoked.</param>
     public EnvironmentCallbackAnnotation(Action<EnvironmentCallbackContext> callback)
+    {
+        Callback = c =>
+        {
+            callback(c);
+            return Task.CompletedTask;
+        };
+    }
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="EnvironmentCallbackAnnotation"/> class with the specified callback.
+    /// </summary>
+    /// <param name="callback">The callback to be invoked.</param>
+    public EnvironmentCallbackAnnotation(Func<EnvironmentCallbackContext, Task> callback)
     {
         Callback = callback;
     }
 
     /// <summary>
     /// Gets or sets the callback action to be executed when the environment is being built.
     /// </summary>
-    public Action<EnvironmentCallbackContext> Callback { get; private set; }
+    public Func<EnvironmentCallbackContext, Task> Callback { get; private set; }
 
     private string DebuggerToString()
     {

src/Aspire.Hosting/ApplicationModel/EnvironmentCallbackContext.cs

@@ -8,7 +8,8 @@ namespace Aspire.Hosting.ApplicationModel;
 /// </summary>
 /// <param name="executionContext">The execution context for this invocation of the AppHost.</param>
 /// <param name="environmentVariables">The environment variables associated with this execution.</param>
-public class EnvironmentCallbackContext(DistributedApplicationExecutionContext executionContext, Dictionary<string, string>? environmentVariables = null)
+/// <param name="cancellationToken">The cancellation token associated with this execution.</param>
+public class EnvironmentCallbackContext(DistributedApplicationExecutionContext executionContext, Dictionary<string, string>? environmentVariables = null, CancellationToken cancellationToken = default)
 {
     /// <summary>
     /// Obsolete. Use ExecutionContext instead. Will be removed in next preview.
@@ -21,6 +22,11 @@ public class EnvironmentCallbackContext(DistributedApplicationExecutionContext e
     /// </summary>
     public Dictionary<string, string> EnvironmentVariables { get; } = environmentVariables ?? new();
 
+    /// <summary>
+    /// Gets the CancellationToken associated with the callback context.
+    /// </summary>
+    public CancellationToken CancellationToken { get; } = cancellationToken;
+
     /// <summary>
     /// Gets the execution context associated with this invocation of the AppHost.
     /// </summary>

src/Aspire.Hosting/ApplicationModel/IResourceWithConnectionString.cs

@@ -14,6 +14,13 @@ public interface IResourceWithConnectionString : IResource
     /// <returns>The connection string associated with the resource, when one is available.</returns>
     public string? GetConnectionString();
 
+    /// <summary>
+    /// Gets the connection string associated with the resource.
+    /// </summary>
+    /// <param name="cancellationToken"> A <see cref="CancellationToken"/> to observe while waiting for the task to complete.</param>
+    /// <returns>The connection string associated with the resource, when one is available.</returns>
+    public ValueTask<string?> GetConnectionStringAsync(CancellationToken cancellationToken = default) => new(GetConnectionString());
+
     /// <summary>
     /// Describes the connection string format string used for this resource in the manifest.
     /// </summary>

src/Aspire.Hosting/ApplicationModel/ManifestPublishingCallbackAnnotation.cs

@@ -8,13 +8,38 @@ namespace Aspire.Hosting.ApplicationModel;
 /// <summary>
 /// Represents an annotation that provides a callback to be executed during manifest publishing.
 /// </summary>
-public class ManifestPublishingCallbackAnnotation(Action<ManifestPublishingContext>? callback) : IResourceAnnotation
+public class ManifestPublishingCallbackAnnotation : IResourceAnnotation
 {
+    /// <summary>
+    /// Initializes a new instance of the <see cref="ManifestPublishingCallbackAnnotation"/> class with the specified callback.
+    /// </summary>
+    /// <param name="callback"></param>
+    public ManifestPublishingCallbackAnnotation(Action<ManifestPublishingContext>? callback)
+    {
+        if (callback is not null)
+        {
+            Callback = context =>
+            {
+                callback(context);
+                return Task.CompletedTask;
+            };
+        }
+    }
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="ManifestPublishingCallbackAnnotation"/> class with the specified callback.
+    /// </summary>
+    /// <param name="callback"></param>
+    public ManifestPublishingCallbackAnnotation(Func<ManifestPublishingContext, Task>? callback)
+    {
+        Callback = callback;
+    }
+
     /// <summary>
     /// Gets the callback action for publishing the manifest.
     /// </summary>
-    public Action<ManifestPublishingContext>? Callback { get; } = callback;
-    
+    public Func<ManifestPublishingContext, Task>? Callback { get; }
+
     /// <summary>
     /// Represents a <see langword="null"/>-based callback annotation for manifest 
     /// publishing used in scenarios where it's ignored.

src/Aspire.Hosting/ApplicationModel/ResourceUpdatesAnnotation.cs

@@ -10,7 +10,7 @@ namespace Aspire.Hosting.ApplicationModel;
 /// <summary>
 /// The annotation that allows publishing and subscribing to changes in the state of a resource.
 /// </summary>
-public sealed class ResourceUpdatesAnnotation(Func<CustomResourceSnapshot> initialSnapshotFactory) : IResourceAnnotation
+public sealed class ResourceUpdatesAnnotation(Func<CancellationToken, ValueTask<CustomResourceSnapshot>> initialSnapshotFactory) : IResourceAnnotation
 {
     private readonly CancellationTokenSource _streamClosedCts = new();
 
@@ -24,7 +24,7 @@ public sealed class ResourceUpdatesAnnotation(Func<CustomResourceSnapshot> initi
     /// <summary>
     /// Gets the initial snapshot of the dashboard state for this resource.
     /// </summary>
-    public CustomResourceSnapshot GetInitialSnapshot() => initialSnapshotFactory();
+    public ValueTask<CustomResourceSnapshot> GetInitialSnapshotAsync(CancellationToken cancellationToken = default) => initialSnapshotFactory(cancellationToken);
 
     /// <summary>
     /// Updates the snapshot of the <see cref="CustomResourceSnapshot"/> for a resource.
@@ -114,8 +114,9 @@ public sealed record CustomResourceSnapshot
     /// Creates a new <see cref="CustomResourceSnapshot"/> for a resource using the well known annotations.
     /// </summary>
     /// <param name="resource">The resource.</param>
+    /// <param name="cancellationToken">The cancellation token.</param>
     /// <returns>The new <see cref="CustomResourceSnapshot"/>.</returns>
-    public static CustomResourceSnapshot Create(IResource resource)
+    public static async ValueTask<CustomResourceSnapshot> CreateAsync(IResource resource, CancellationToken cancellationToken = default)
     {
         ImmutableArray<string> urls = [];
 
@@ -131,10 +132,10 @@ static string GetUrl(EndpointAnnotation e) =>
 
         if (resource.TryGetAnnotationsOfType<EnvironmentCallbackAnnotation>(out var environmentCallbacks))
         {
-            var envContext = new EnvironmentCallbackContext(new DistributedApplicationExecutionContext(DistributedApplicationOperation.Run));
+            var envContext = new EnvironmentCallbackContext(new DistributedApplicationExecutionContext(DistributedApplicationOperation.Run), cancellationToken: cancellationToken);
             foreach (var annotation in environmentCallbacks)
             {
-                annotation.Callback(envContext);
+                await annotation.Callback(envContext).ConfigureAwait(false);
             }
 
             environmentVariables = [.. envContext.EnvironmentVariables.Select(e => (e.Key, e.Value))];

src/Aspire.Hosting/Dashboard/DcpDataSource.cs

@@ -201,7 +201,7 @@ private async Task ProcessInitialResourceAsync(IResource resource, CancellationT
         {
             // We have a dashboard annotation, so we want to create a snapshot for the resource
             // and update data immediately. We also want to watch for changes to the dashboard state.
-            var state = resourceUpdates.GetInitialSnapshot();
+            var state = await resourceUpdates.GetInitialSnapshotAsync(cancellationToken).ConfigureAwait(false);
             var creationTimestamp = DateTime.UtcNow;
 
             var snapshot = CreateResourceSnapshot(resource, creationTimestamp, state);

src/Aspire.Hosting/Dcp/ApplicationExecutor.cs

@@ -529,16 +529,18 @@ private async Task CreateExecutablesAsync(IEnumerable<AppResource> executableRes
 
                 spec.Args ??= new();
 
-                if (er.ModelResource.TryGetAnnotationsOfType<ExecutableArgsCallbackAnnotation>(out var exeArgsCallbacks))
+                if (er.ModelResource.TryGetAnnotationsOfType<CommandLineArgsCallbackAnnotation>(out var exeArgsCallbacks))
                 {
+                    var commandLineContext = new CommandLineArgsCallbackContext(spec.Args, cancellationToken);
+
                     foreach (var exeArgsCallback in exeArgsCallbacks)
                     {
-                        exeArgsCallback.Callback(spec.Args);
+                        await exeArgsCallback.Callback(commandLineContext).ConfigureAwait(false);
                     }
                 }
 
                 var config = new Dictionary<string, string>();
-                var context = new EnvironmentCallbackContext(_executionContext, config);
+                var context = new EnvironmentCallbackContext(_executionContext, config, cancellationToken);
 
                 // Need to apply configuration settings manually; see PrepareExecutables() for details.
                 if (er.ModelResource is ProjectResource project && project.SelectLaunchProfileName() is { } launchProfileName && project.GetLaunchSettings() is { } launchSettings)
@@ -570,7 +572,7 @@ private async Task CreateExecutablesAsync(IEnumerable<AppResource> executableRes
                 {
                     foreach (var ann in envVarAnnotations)
                     {
-                        ann.Callback(context);
+                        await ann.Callback(context).ConfigureAwait(false);
                     }
                 }
 
@@ -789,11 +791,11 @@ private async Task CreateContainersAsync(IEnumerable<AppResource> containerResou
 
                 if (modelContainerResource.TryGetEnvironmentVariables(out var containerEnvironmentVariables))
                 {
-                    var context = new EnvironmentCallbackContext(_executionContext, config);
+                    var context = new EnvironmentCallbackContext(_executionContext, config, cancellationToken);
 
                     foreach (var v in containerEnvironmentVariables)
                     {
-                        v.Callback(context);
+                        await v.Callback(context).ConfigureAwait(false);
                     }
                 }
 
@@ -802,12 +804,15 @@ private async Task CreateContainersAsync(IEnumerable<AppResource> containerResou
                     dcpContainerResource.Spec.Env.Add(new EnvVar { Name = kvp.Key, Value = kvp.Value });
                 }
 
-                if (modelContainerResource.TryGetAnnotationsOfType<ExecutableArgsCallbackAnnotation>(out var argsCallback))
+                if (modelContainerResource.TryGetAnnotationsOfType<CommandLineArgsCallbackAnnotation>(out var argsCallback))
                 {
                     dcpContainerResource.Spec.Args ??= [];
+
+                    var commandLineArgsContext = new CommandLineArgsCallbackContext(dcpContainerResource.Spec.Args, cancellationToken);
+
                     foreach (var callback in argsCallback)
                     {
-                        callback.Callback(dcpContainerResource.Spec.Args);
+                        await callback.Callback(commandLineArgsContext).ConfigureAwait(false);
                     }
                 }
 

src/Aspire.Hosting/Extensions/ContainerResourceBuilderExtensions.cs

@@ -76,7 +76,7 @@ public static IResourceBuilder<T> WithBindMount<T>(this IResourceBuilder<T> buil
     /// <returns>The <see cref="IResourceBuilder{T}"/>.</returns>    
     public static IResourceBuilder<T> WithArgs<T>(this IResourceBuilder<T> builder, params string[] args) where T : ContainerResource
     {
-        var annotation = new ExecutableArgsCallbackAnnotation(updatedArgs =>
+        var annotation = new CommandLineArgsCallbackAnnotation(updatedArgs =>
         {
             updatedArgs.AddRange(args);
         });
@@ -159,7 +159,7 @@ public static IResourceBuilder<T> WithImageSHA256<T>(this IResourceBuilder<T> bu
     /// <returns>A reference to the <see cref="IResourceBuilder{T}"/>.</returns>
     public static IResourceBuilder<T> PublishAsContainer<T>(this IResourceBuilder<T> builder) where T : ContainerResource
     {
-        return builder.WithManifestPublishingCallback(context => context.WriteContainer(builder.Resource));
+        return builder.WithManifestPublishingCallback(context => context.WriteContainerAsync(builder.Resource));
     }
 }