Improve concurrency handling in IDeploymentStateManager with optimistic concurrency control (#12270)

* Improve concurrency handling in IDeploymentStateManager

* Serialize SaveStateToStorageAsync calls to prevent concurrent file writes

Co-authored-by: davidfowl <[email protected]>

* Merge section locks and versions into single ConcurrentDictionary

Co-authored-by: davidfowl <[email protected]>

* Move version update outside of _saveLock using thread-safe AddOrUpdate

Co-authored-by: davidfowl <[email protected]>

* Make DeploymentStateSection an immutable snapshot using DeepClone

Co-authored-by: davidfowl <[email protected]>

* Remove synchronization from DeploymentStateSection and add thread-safety warning

Co-authored-by: davidfowl <[email protected]>

* Remove IDisposable from DeploymentStateSection and implement atomic version check

Co-authored-by: davidfowl <[email protected]>

* Update tests/Aspire.Hosting.Tests/Publishing/DeploymentStateManagerTests.cs

* Make Version settable and increment after save, remove unused Lock from SectionMetadata

Co-authored-by: davidfowl <[email protected]>

* Replace ConcurrentDictionary with lock + Dictionary for atomic version checks

Co-authored-by: davidfowl <[email protected]>

* Fix Azure tests to use new DeploymentStateSection constructor without IDisposable

Co-authored-by: davidfowl <[email protected]>

---------

Co-authored-by: Safia Abdalla <[email protected]>
Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: davidfowl <[email protected]>
Co-authored-by: David Fowler <[email protected]>

Author: 3 people

src/Aspire.Hosting.Azure/AzureEnvironmentResource.cs

@@ -72,7 +72,11 @@ public AzureEnvironmentResource(string name, ParameterResource location, Paramet
             var createContextStep = new PipelineStep
             {
                 Name = "create-provisioning-context",
-                Action = async ctx => provisioningContext = await CreateProvisioningContextAsync(ctx).ConfigureAwait(false)
+                Action = async ctx =>
+                {
+                    var provisioningContextProvider = ctx.Services.GetRequiredService<IProvisioningContextProvider>();
+                    provisioningContext = await provisioningContextProvider.CreateProvisioningContextAsync(ctx.CancellationToken).ConfigureAwait(false);
+                }
             };
             createContextStep.DependsOn(validateStep);
 
@@ -165,33 +169,9 @@ await context.ReportingStep.CompleteAsync(
         }
     }
 
-    private static async Task<ProvisioningContext> CreateProvisioningContextAsync(PipelineStepContext context)
-    {
-        var provisioningContextProvider = context.Services.GetRequiredService<IProvisioningContextProvider>();
-        var deploymentStateManager = context.Services.GetRequiredService<IDeploymentStateManager>();
-        var configuration = context.Services.GetRequiredService<IConfiguration>();
-
-        var userSecrets = await deploymentStateManager.LoadStateAsync(context.CancellationToken)
-            .ConfigureAwait(false);
-        var provisioningContext = await provisioningContextProvider
-            .CreateProvisioningContextAsync(userSecrets, context.CancellationToken)
-            .ConfigureAwait(false);
-
-        var clearCache = configuration.GetValue<bool>("Publishing:ClearCache");
-        if (!clearCache)
-        {
-            await deploymentStateManager.SaveStateAsync(
-                provisioningContext.DeploymentState,
-                context.CancellationToken).ConfigureAwait(false);
-        }
-
-        return provisioningContext;
-    }
-
     private static async Task ProvisionAzureBicepResourcesAsync(PipelineStepContext context, ProvisioningContext provisioningContext)
     {
         var bicepProvisioner = context.Services.GetRequiredService<IBicepProvisioner>();
-        var deploymentStateManager = context.Services.GetRequiredService<IDeploymentStateManager>();
         var configuration = context.Services.GetRequiredService<IConfiguration>();
 
         var bicepResources = context.Model.Resources.OfType<AzureBicepResource>()
@@ -258,14 +238,6 @@ await resourceTask.CompleteAsync(
         });
 
         await Task.WhenAll(provisioningTasks).ConfigureAwait(false);
-
-        var clearCache = configuration.GetValue<bool>("Publishing:ClearCache");
-        if (!clearCache)
-        {
-            await deploymentStateManager.SaveStateAsync(
-                provisioningContext.DeploymentState,
-                context.CancellationToken).ConfigureAwait(false);
-        }
     }
 
     private static async Task BuildContainerImagesAsync(PipelineStepContext context)

src/Aspire.Hosting.Azure/Provisioning/Internal/BaseProvisioningContextProvider.cs

@@ -1,10 +1,11 @@
 #pragma warning disable ASPIREINTERACTION001 // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
+#pragma warning disable ASPIREPUBLISHERS001 // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
 
 // 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.Json.Nodes;
 using System.Text.RegularExpressions;
+using Aspire.Hosting.Publishing;
 using Azure;
 using Azure.Core;
 using Azure.ResourceManager.Resources;
@@ -25,6 +26,7 @@ internal abstract partial class BaseProvisioningContextProvider(
     IArmClientProvider armClientProvider,
     IUserPrincipalProvider userPrincipalProvider,
     ITokenCredentialProvider tokenCredentialProvider,
+    IDeploymentStateManager deploymentStateManager,
     DistributedApplicationExecutionContext distributedApplicationExecutionContext) : IProvisioningContextProvider
 {
     internal const string LocationName = "Location";
@@ -73,7 +75,7 @@ protected static bool IsValidResourceGroupName(string? name)
         return !name.Contains("..");
     }
 
-    public virtual async Task<ProvisioningContext> CreateProvisioningContextAsync(JsonObject deploymentState, CancellationToken cancellationToken = default)
+    public virtual async Task<ProvisioningContext> CreateProvisioningContextAsync(CancellationToken cancellationToken = default)
     {
         var subscriptionId = _options.SubscriptionId ?? throw new MissingConfigurationException("An Azure subscription id is required. Set the Azure:SubscriptionId configuration value.");
 
@@ -98,6 +100,9 @@ public virtual async Task<ProvisioningContext> CreateProvisioningContextAsync(Js
             throw new MissingConfigurationException("An azure location/region is required. Set the Azure:Location configuration value.");
         }
 
+        // Acquire Azure state section for reading/writing configuration
+        var azureStateSection = await deploymentStateManager.AcquireSectionAsync("Azure", cancellationToken).ConfigureAwait(false);
+
         string resourceGroupName;
         bool createIfAbsent;
 
@@ -109,7 +114,7 @@ public virtual async Task<ProvisioningContext> CreateProvisioningContextAsync(Js
 
             createIfAbsent = true;
 
-            deploymentState.Prop("Azure")["ResourceGroup"] = resourceGroupName;
+            azureStateSection.Data["ResourceGroup"] = resourceGroupName;
         }
         else
         {
@@ -158,7 +163,7 @@ public virtual async Task<ProvisioningContext> CreateProvisioningContextAsync(Js
         var principal = await _userPrincipalProvider.GetUserPrincipalAsync(cancellationToken).ConfigureAwait(false);
 
         // Persist the provisioning options to deployment state so they can be reused in the future
-        var azureSection = deploymentState.Prop("Azure");
+        var azureSection = azureStateSection.Data;
         azureSection["Location"] = _options.Location;
         azureSection["SubscriptionId"] = _options.SubscriptionId;
         azureSection["ResourceGroup"] = resourceGroupName;
@@ -171,6 +176,8 @@ public virtual async Task<ProvisioningContext> CreateProvisioningContextAsync(Js
             azureSection["AllowResourceGroupCreation"] = _options.AllowResourceGroupCreation.Value;
         }
 
+        await deploymentStateManager.SaveSectionAsync(azureStateSection, cancellationToken).ConfigureAwait(false);
+
         return new ProvisioningContext(
                     credential,
                     armClient,
@@ -179,7 +186,6 @@ public virtual async Task<ProvisioningContext> CreateProvisioningContextAsync(Js
                     tenantResource,
                     location,
                     principal,
-                    deploymentState,
                     _distributedApplicationExecutionContext);
     }
 
@@ -200,25 +206,25 @@ public virtual async Task<ProvisioningContext> CreateProvisioningContextAsync(Js
             if (tenantList.Count > 0)
             {
                 tenantOptions = tenantList
-                    .Select(t => 
+                    .Select(t =>
                     {
                         var tenantId = t.TenantId?.ToString() ?? "";
-                        
+
                         // Build display name: prefer DisplayName, fall back to domain, then to "Unknown"
-                        var displayName = !string.IsNullOrEmpty(t.DisplayName) 
-                            ? t.DisplayName 
-                            : !string.IsNullOrEmpty(t.DefaultDomain) 
-                                ? t.DefaultDomain 
+                        var displayName = !string.IsNullOrEmpty(t.DisplayName)
+                            ? t.DisplayName
+                            : !string.IsNullOrEmpty(t.DefaultDomain)
+                                ? t.DefaultDomain
                                 : "Unknown";
-                        
+
                         // Build full description
                         var description = displayName;
                         if (!string.IsNullOrEmpty(t.DefaultDomain) && t.DisplayName != t.DefaultDomain)
                         {
                             description += $" ({t.DefaultDomain})";
                         }
                         description += $" — {tenantId}";
-                        
+
                         return KeyValuePair.Create(tenantId, description);
                     })
                     .OrderBy(kvp => kvp.Value)

src/Aspire.Hosting.Azure/Provisioning/Internal/IProvisioningServices.cs

@@ -1,7 +1,6 @@
 // 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.Json.Nodes;
 using Azure;
 using Azure.Core;
 using Azure.ResourceManager;
@@ -57,7 +56,9 @@ internal interface IProvisioningContextProvider
     /// <summary>
     /// Creates a provisioning context for Azure resource operations.
     /// </summary>
-    Task<ProvisioningContext> CreateProvisioningContextAsync(JsonObject userSecrets, CancellationToken cancellationToken = default);
+    /// <param name="cancellationToken">The cancellation token.</param>
+    /// <returns>A provisioning context.</returns>
+    Task<ProvisioningContext> CreateProvisioningContextAsync(CancellationToken cancellationToken = default);
 }
 
 /// <summary>

src/Aspire.Hosting.Azure/Provisioning/Internal/PublishModeProvisioningContextProvider.cs

@@ -4,10 +4,10 @@
 #pragma warning disable ASPIREINTERACTION001
 #pragma warning disable ASPIREPUBLISHERS001
 
-using System.Text.Json.Nodes;
 using Aspire.Hosting.Azure.Resources;
 using Aspire.Hosting.Azure.Utils;
 using Aspire.Hosting.Pipelines;
+using Aspire.Hosting.Publishing;
 using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
@@ -26,6 +26,7 @@ internal sealed class PublishModeProvisioningContextProvider(
     IArmClientProvider armClientProvider,
     IUserPrincipalProvider userPrincipalProvider,
     ITokenCredentialProvider tokenCredentialProvider,
+    IDeploymentStateManager deploymentStateManager,
     DistributedApplicationExecutionContext distributedApplicationExecutionContext,
     IPipelineActivityReporter activityReporter) : BaseProvisioningContextProvider(
         interactionService,
@@ -35,6 +36,7 @@ internal sealed class PublishModeProvisioningContextProvider(
         armClientProvider,
         userPrincipalProvider,
         tokenCredentialProvider,
+        deploymentStateManager,
         distributedApplicationExecutionContext)
 {
     protected override string GetDefaultResourceGroupName()
@@ -58,7 +60,7 @@ protected override string GetDefaultResourceGroupName()
         return $"{prefix}-{normalizedApplicationName}";
     }
 
-    public override async Task<ProvisioningContext> CreateProvisioningContextAsync(JsonObject userSecrets, CancellationToken cancellationToken = default)
+    public override async Task<ProvisioningContext> CreateProvisioningContextAsync(CancellationToken cancellationToken = default)
     {
         try
         {
@@ -74,7 +76,7 @@ public override async Task<ProvisioningContext> CreateProvisioningContextAsync(J
             _logger.LogError(ex, "Failed to retrieve Azure provisioning options.");
         }
 
-        return await base.CreateProvisioningContextAsync(userSecrets, cancellationToken).ConfigureAwait(false);
+        return await base.CreateProvisioningContextAsync(cancellationToken).ConfigureAwait(false);
     }
 
     private async Task RetrieveAzureProvisioningOptions(CancellationToken cancellationToken = default)

src/Aspire.Hosting.Azure/Provisioning/Internal/RunModeProvisioningContextProvider.cs

@@ -1,12 +1,13 @@
 #pragma warning disable ASPIREINTERACTION001 // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
+#pragma warning disable ASPIREPUBLISHERS001 // Type is for evaluation purposes only and is subject to change or removal in future updates. Suppress this diagnostic to proceed.
 
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System.Security.Cryptography;
-using System.Text.Json.Nodes;
 using Aspire.Hosting.Azure.Resources;
 using Aspire.Hosting.Azure.Utils;
+using Aspire.Hosting.Publishing;
 using Microsoft.Extensions.Hosting;
 using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Options;
@@ -24,6 +25,7 @@ internal sealed class RunModeProvisioningContextProvider(
     IArmClientProvider armClientProvider,
     IUserPrincipalProvider userPrincipalProvider,
     ITokenCredentialProvider tokenCredentialProvider,
+    IDeploymentStateManager deploymentStateManager,
     DistributedApplicationExecutionContext distributedApplicationExecutionContext) : BaseProvisioningContextProvider(
         interactionService,
         options,
@@ -32,6 +34,7 @@ internal sealed class RunModeProvisioningContextProvider(
         armClientProvider,
         userPrincipalProvider,
         tokenCredentialProvider,
+        deploymentStateManager,
         distributedApplicationExecutionContext)
 {
     private readonly TaskCompletionSource _provisioningOptionsAvailable = new(TaskCreationOptions.RunContinuationsAsynchronously);
@@ -87,13 +90,13 @@ private void EnsureProvisioningOptions()
         });
     }
 
-    public override async Task<ProvisioningContext> CreateProvisioningContextAsync(JsonObject userSecrets, CancellationToken cancellationToken = default)
+    public override async Task<ProvisioningContext> CreateProvisioningContextAsync(CancellationToken cancellationToken = default)
     {
         EnsureProvisioningOptions();
 
         await _provisioningOptionsAvailable.Task.ConfigureAwait(false);
 
-        return await base.CreateProvisioningContextAsync(userSecrets, cancellationToken).ConfigureAwait(false);
+        return await base.CreateProvisioningContextAsync(cancellationToken).ConfigureAwait(false);
     }
 
     private async Task RetrieveAzureProvisioningOptions(CancellationToken cancellationToken = default)

src/Aspire.Hosting.Azure/Provisioning/Provisioners/AzureProvisioner.cs

@@ -8,7 +8,6 @@
 using Aspire.Hosting.Azure.Provisioning.Internal;
 using Aspire.Hosting.Eventing;
 using Aspire.Hosting.Lifecycle;
-using Aspire.Hosting.Publishing;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.Logging;
 
@@ -22,8 +21,7 @@ internal sealed class AzureProvisioner(
     ResourceNotificationService notificationService,
     ResourceLoggerService loggerService,
     IDistributedApplicationEventing eventing,
-    IProvisioningContextProvider provisioningContextProvider,
-    IDeploymentStateManager deploymentStateManager
+    IProvisioningContextProvider provisioningContextProvider
     ) : IDistributedApplicationEventingSubscriber
 {
     internal const string AspireResourceNameTag = "aspire-resource-name";
@@ -164,11 +162,8 @@ private async Task ProvisionAzureResources(
         IList<(IResource Resource, IAzureResource AzureResource)> azureResources,
         CancellationToken cancellationToken)
     {
-        // Load deployment state first so it can be passed to the provisioning context
-        var deploymentState = await deploymentStateManager.LoadStateAsync(cancellationToken).ConfigureAwait(false);
-
         // Make resources wait on the same provisioning context
-        var provisioningContextLazy = new Lazy<Task<ProvisioningContext>>(() => provisioningContextProvider.CreateProvisioningContextAsync(deploymentState, cancellationToken));
+        var provisioningContextLazy = new Lazy<Task<ProvisioningContext>>(() => provisioningContextProvider.CreateProvisioningContextAsync(cancellationToken));
 
         var tasks = new List<Task>();
 
@@ -182,9 +177,6 @@ private async Task ProvisionAzureResources(
         // Suppress throwing so that we can save the deployment state even if the task fails
         await task.ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
 
-        // If we created any resources then save the deployment state
-        await deploymentStateManager.SaveStateAsync(deploymentState, cancellationToken).ConfigureAwait(false);
-
         // Set the completion source for all resources
         foreach (var resource in azureResources)
         {