await background tasks in inline deployments (#420)

Adds an await for inline deployments

@Frassle As dicussed over slack, new pr since I used the wrong base
branch.

.changes/unreleased/bug-fixes-420.yaml

@@ -0,0 +1,6 @@
+component: sdk
+kind: bug-fixes
+body: Await background tasks during inline deployment
+time: 2024-12-04T15:32:00.3901679+01:00
+custom:
+    PR: "420"

sdk/Pulumi.Automation.Tests/LocalWorkspaceTests.cs

@@ -60,7 +60,11 @@ public LocalWorkspaceTests(ITestOutputHelper output)
             {
                 temporaryDirectory = Path.Combine(Path.GetTempPath(), "pulumi", "automation-tests");
                 Directory.CreateDirectory(temporaryDirectory);
-                Environment.SetEnvironmentVariable("PULUMI_BACKEND_URL", $"file:///{temporaryDirectory}");
+                Environment.SetEnvironmentVariable("PULUMI_BACKEND_URL",
+                    OperatingSystem.IsWindows()
+                        ? $"file://{temporaryDirectory.Replace("\\", "/")}"
+                        : $"file:///{temporaryDirectory}");
+
                 // Because we're using filestate we need to set a passphrase as well.
                 Environment.SetEnvironmentVariable("PULUMI_CONFIG_PASSPHRASE", "backup_password");
             }

sdk/Pulumi.Tests/Deployment/InlineDeploymentTests.cs

@@ -0,0 +1,90 @@
+using System;
+using System.Collections.Generic;
+using System.Collections.Immutable;
+using System.Threading.Tasks;
+using Pulumi.Testing;
+using Xunit;
+
+namespace Pulumi.Tests;
+
+public class InlineDeploymentTests
+{
+    [Fact]
+    public async Task InlineDeploymentAwaitsTasks()
+    {
+        var mocks = new AwaitingMocks();
+        var monitor = new MockMonitor(mocks);
+        var res = TryInline<Component>(monitor, async () =>
+        {
+            var stack = new Stack();
+            await Task.Delay(1);
+            return new Component("test", null, new ComponentResourceOptions
+            {
+                Parent = stack
+            });
+        });
+
+        await Task.WhenAny(res, Task.Delay(TimeSpan.FromSeconds(1)));
+        Assert.False(res.IsCompleted);
+        mocks.Complete();
+        await res;
+        Assert.Equal(4, monitor.Resources.Count);
+    }
+
+    internal static async Task<T> TryInline<T>(IMonitor monitor, Func<Task<T>> runAsync)
+    {
+        var engine = new MockEngine();
+
+        var deploymentBuilder = new MockDeploymentBuilder(engine, monitor);
+
+        var inlineDeploymentSettings = new InlineDeploymentSettings(null, "1", "1", new Dictionary<string, string>(),
+            new List<string>(), "organization", "project", "stack", 0, false);
+
+        return await Deployment.RunInlineAsyncWithResult(deploymentBuilder, inlineDeploymentSettings, runAsync)
+            .ConfigureAwait(false);
+    }
+
+    internal class Component : ComponentResource
+    {
+        public Component(string name, ResourceArgs? args, ComponentResourceOptions? options = null, bool remote = false)
+            : base("test:res:a", name, args, options, remote)
+        {
+            new ComponentResource("test:res:b", "testB", new ComponentResourceOptions
+            {
+                Parent = this
+            });
+            new ComponentResource("test:res:b", "testC", new ComponentResourceOptions
+            {
+                Parent = this
+            });
+        }
+    }
+
+    internal class AwaitingMocks : IMocks
+    {
+        private readonly TaskCompletionSource _taskCompletionSource = new();
+
+        private Task Task => _taskCompletionSource.Task;
+
+        public void Complete()
+        {
+            _taskCompletionSource.TrySetResult();
+        }
+
+        public async Task<(string? id, object state)> NewResourceAsync(MockResourceArgs args)
+        {
+            if (args.Type?.Equals("test:res:b") ?? false)
+            {
+                await Task;
+            }
+
+            string? id = args.Name;
+            return (id, ImmutableDictionary.Create<string, string>());
+        }
+
+        public Task<object> CallAsync(MockCallArgs args)
+        {
+            return Task.FromResult(new object());
+        }
+    }
+}

sdk/Pulumi/Deployment/Deployment.Runner.cs

@@ -69,6 +69,13 @@ Task<int> IRunner.RunAsync<TStack>(IServiceProvider serviceProvider)
                     ?? throw new ApplicationException($"Failed to resolve instance of type {typeof(TStack)} from service provider. Register the type with the service provider before calling {nameof(RunAsync)}."));
             }
 
+            public async Task<T> RunAsync<T>(Func<Task<T>> func)
+            {
+                var result = await func();
+                await WhileRunningAsync();
+                return result;
+            }
+
             Task<int> IRunner.RunAsync<TStack>() => RunAsync(() => (TStack)Activator.CreateInstance(typeof(TStack), BindingFlags.DoNotWrapExceptions, binder: null, args: null, culture: null)!);
 
             public Task<int> RunAsync<TStack>(Func<TStack> stackFactory) where TStack : Stack
@@ -176,8 +183,8 @@ private async Task<int> HandleExceptionsAsync(IEnumerable<Exception> exceptions)
                 {
                     // We set the exit code explicitly here in case users
                     // do not bubble up the exit code themselves to
-                    // top-level entry point of the program (for non-inline deployments). 
-                    // For example when users `await Deployment.RunAsync()` 
+                    // top-level entry point of the program (for non-inline deployments).
+                    // For example when users `await Deployment.RunAsync()`
                     // instead of `return await Deployment.RunAsync()`
                     Environment.ExitCode = exitCode;
                 }

sdk/Pulumi/Deployment/Deployment_Inline.cs

@@ -71,8 +71,13 @@ internal static async Task<InlineDeploymentResult> RunInlineAsync(IDeploymentBui
         }
 
         internal static async Task<T> RunInlineAsyncWithResult<T>(IDeploymentBuilder builder,
-            InlineDeploymentSettings settings,
-            Func<IRunner, Task<T>> runnerFunc)
+            InlineDeploymentSettings settings, Func<Task<T>> runnerFunc)
+        {
+            return await RunInlineAsyncWithResult(builder, settings, runner => runner.RunAsync(runnerFunc));
+        }
+
+        private static async Task<T> RunInlineAsyncWithResult<T>(IDeploymentBuilder builder,
+            InlineDeploymentSettings settings, Func<IRunner, Task<T>> runnerFunc)
         {
             return await CreateRunnerAndRunAsync(
                 () => new Deployment(builder, settings),

sdk/Pulumi/Deployment/IRunner.cs

@@ -15,5 +15,6 @@ internal interface IRunner
         Task<int> RunAsync<TStack>() where TStack : Stack, new();
         Task<int> RunAsync<TStack>(Func<TStack> stackFactory) where TStack : Stack;
         Task<int> RunAsync<TStack>(IServiceProvider serviceProvider) where TStack : Stack;
+        Task<T> RunAsync<T>(Func<Task<T>> func);
     }
 }

sdk/Pulumi/Provider/Provider.cs

@@ -1064,7 +1064,7 @@ public override Task<Empty> Delete(Pulumirpc.DeleteRequest request, ServerCallCo
                 var inlineDeploymentSettings = new InlineDeploymentSettings(logger, EngineAddress, request.MonitorEndpoint, request.Config,
                     request.ConfigSecretKeys, request.Organization, request.Project, request.Stack, request.Parallel, request.DryRun);
                 var domResponse = await Deployment
-                    .RunInlineAsyncWithResult(deploymentBuilder, inlineDeploymentSettings, runner => Implementation.Construct(domRequest, cts.Token))
+                    .RunInlineAsyncWithResult(deploymentBuilder, inlineDeploymentSettings, () => Implementation.Construct(domRequest, cts.Token))
                     .ConfigureAwait(false);
 
                 var state = PropertyValue.Marshal(domResponse.State);
@@ -1104,7 +1104,7 @@ public override Task<Empty> Delete(Pulumirpc.DeleteRequest request, ServerCallCo
                 var inlineDeploymentSettings = new InlineDeploymentSettings(logger, EngineAddress, request.MonitorEndpoint, request.Config,
                     request.ConfigSecretKeys, request.Organization, request.Project, request.Stack, request.Parallel, request.DryRun);
                 var domResponse = await Deployment
-                    .RunInlineAsyncWithResult(deploymentBuilder, inlineDeploymentSettings, runner => Implementation.Call(domRequest, cts.Token))
+                    .RunInlineAsyncWithResult(deploymentBuilder, inlineDeploymentSettings, () => Implementation.Call(domRequest, cts.Token))
                     .ConfigureAwait(false);
 
                 IDictionary<string, ISet<Urn>> returnDependencies = ImmutableDictionary<string, ISet<Urn>>.Empty;