Fix race condition in concurrent DeploymentState access causing intermittent AzureDeployerTests failures (#11974)

* Initial plan

* Fix race condition in JsonExtensions.Prop causing AzureDeployerTests failures

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

* Add thread-safety tests for JsonExtensions.Prop

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

* Use ConditionalWeakTable for lock objects instead of locking on JsonObject

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

* Move thread-safety from JsonExtensions.Prop to ProvisioningContext level

Address feedback to provide thread-safety at the ProvisioningContext level
rather than in the Prop method itself. This ensures all DeploymentState
access is synchronized when multiple resources are provisioned in parallel.

- Added WithDeploymentState methods to ProvisioningContext for thread-safe access
- Updated BicepProvisioner to use WithDeploymentState for all state modifications
- Reverted JsonExtensions.Prop to simpler implementation without locking
- Updated tests: removed concurrent test from JsonExtensionsTests
- Added comprehensive thread-safety tests to ProvisioningContextTests

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

* Revert JsonExtensions.cs to original implementation

As suggested by @davidfowl, revert changes to JsonExtensions.cs since
the core thread-safety issue is now properly handled by WithDeploymentState
at the ProvisioningContext level. The original TryAdd logic with retry
is restored.

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

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: radical <[email protected]>
Co-authored-by: davidfowl <[email protected]>

Author: 3 people

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

@@ -218,38 +218,41 @@ await notificationService.PublishUpdateAsync(resource, state =>
         // e.g. {  "sqlServerName": { "type": "String", "value": "<value>" }}
         var outputObj = outputs?.ToObjectFromJson<JsonObject>();
 
-        // Populate values into deployment state
-        var az = context.DeploymentState.Prop("Azure");
-        az["Tenant"] = context.Tenant.DefaultDomain;
+        // Populate values into deployment state with thread-safe synchronization
+        context.WithDeploymentState(deploymentState =>
+        {
+            var az = deploymentState.Prop("Azure");
+            az["Tenant"] = context.Tenant.DefaultDomain;
 
-        var resourceConfig = context.DeploymentState
-            .Prop("Azure")
-            .Prop("Deployments")
-            .Prop(resource.Name);
+            var resourceConfig = deploymentState
+                .Prop("Azure")
+                .Prop("Deployments")
+                .Prop(resource.Name);
 
-        // Clear the entire section
-        resourceConfig.AsObject().Clear();
+            // Clear the entire section
+            resourceConfig.AsObject().Clear();
 
-        // Save the deployment id to the configuration
-        resourceConfig["Id"] = deployment.Id.ToString();
+            // Save the deployment id to the configuration
+            resourceConfig["Id"] = deployment.Id.ToString();
 
-        // Stash all parameters as a single JSON string
-        resourceConfig["Parameters"] = parameters.ToJsonString();
+            // Stash all parameters as a single JSON string
+            resourceConfig["Parameters"] = parameters.ToJsonString();
 
-        if (outputObj is not null)
-        {
-            // Same for outputs
-            resourceConfig["Outputs"] = outputObj.ToJsonString();
-        }
+            if (outputObj is not null)
+            {
+                // Same for outputs
+                resourceConfig["Outputs"] = outputObj.ToJsonString();
+            }
 
-        // Write resource scope to config for consistent checksums
-        if (scope is not null)
-        {
-            resourceConfig["Scope"] = scope.ToJsonString();
-        }
+            // Write resource scope to config for consistent checksums
+            if (scope is not null)
+            {
+                resourceConfig["Scope"] = scope.ToJsonString();
+            }
 
-        // Save the checksum to the configuration
-        resourceConfig["CheckSum"] = BicepUtilities.GetChecksum(resource, parameters, scope);
+            // Save the checksum to the configuration
+            resourceConfig["CheckSum"] = BicepUtilities.GetChecksum(resource, parameters, scope);
+        });
 
         if (outputObj is not null)
         {

src/Aspire.Hosting.Azure/Provisioning/ProvisioningContext.cs

@@ -20,6 +20,9 @@ internal sealed class ProvisioningContext(
     JsonObject deploymentState,
     DistributedApplicationExecutionContext executionContext)
 {
+    // Lock object to protect concurrent access to DeploymentState from multiple provisioning tasks
+    private readonly object _deploymentStateLock = new();
+
     public TokenCredential Credential => credential;
     public IArmClient ArmClient => armClient;
     public ISubscriptionResource Subscription => subscription;
@@ -29,4 +32,30 @@ internal sealed class ProvisioningContext(
     public UserPrincipal Principal => principal;
     public JsonObject DeploymentState => deploymentState;
     public DistributedApplicationExecutionContext ExecutionContext => executionContext;
+
+    /// <summary>
+    /// Executes an action on the DeploymentState with thread-safe synchronization.
+    /// Use this method to perform any read or write operations on the DeploymentState
+    /// when multiple resources are being provisioned in parallel.
+    /// </summary>
+    public void WithDeploymentState(Action<JsonObject> action)
+    {
+        lock (_deploymentStateLock)
+        {
+            action(deploymentState);
+        }
+    }
+
+    /// <summary>
+    /// Executes a function on the DeploymentState with thread-safe synchronization.
+    /// Use this method to perform any read or write operations on the DeploymentState
+    /// when multiple resources are being provisioned in parallel.
+    /// </summary>
+    public T WithDeploymentState<T>(Func<JsonObject, T> func)
+    {
+        lock (_deploymentStateLock)
+        {
+            return func(deploymentState);
+        }
+    }
 }

tests/Aspire.Hosting.Azure.Tests/JsonExtensionsTests.cs

@@ -0,0 +1,60 @@
+// 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;
+
+namespace Aspire.Hosting.Azure.Tests;
+
+public class JsonExtensionsTests
+{
+    [Fact]
+    public void Prop_ReturnsExistingNode_WhenNodeAlreadyExists()
+    {
+        // Arrange
+        var rootJson = new JsonObject();
+        var azureNode = rootJson.Prop("Azure");
+        azureNode.AsObject()["TestProperty"] = "TestValue";
+
+        // Act
+        var retrievedNode = rootJson.Prop("Azure");
+
+        // Assert
+        Assert.Same(azureNode, retrievedNode);
+        Assert.Equal("TestValue", retrievedNode["TestProperty"]!.GetValue<string>());
+    }
+
+    [Fact]
+    public void Prop_CreatesNewNode_WhenNodeDoesNotExist()
+    {
+        // Arrange
+        var rootJson = new JsonObject();
+
+        // Act
+        var newNode = rootJson.Prop("NewProperty");
+
+        // Assert
+        Assert.NotNull(newNode);
+        Assert.Same(rootJson["NewProperty"], newNode);
+    }
+
+    [Fact]
+    public void Prop_NestedAccess_CreatesHierarchy()
+    {
+        // Arrange
+        var rootJson = new JsonObject();
+
+        // Act
+        var deeply = rootJson.Prop("Level1")
+                              .Prop("Level2")
+                              .Prop("Level3")
+                              .Prop("Level4");
+
+        // Assert
+        Assert.NotNull(rootJson["Level1"]);
+        Assert.NotNull(rootJson["Level1"]!["Level2"]);
+        Assert.NotNull(rootJson["Level1"]!["Level2"]!["Level3"]);
+        Assert.NotNull(rootJson["Level1"]!["Level2"]!["Level3"]!["Level4"]);
+        Assert.Same(deeply, rootJson["Level1"]!["Level2"]!["Level3"]!["Level4"]);
+    }
+}
+

tests/Aspire.Hosting.Azure.Tests/ProvisioningContextTests.cs

@@ -157,6 +157,175 @@ public void ProvisioningContext_CanBeCustomized()
         Assert.Equal("[email protected]", context.Principal.Name);
         Assert.Equal("value", context.DeploymentState["test"]?.ToString());
     }
+
+    [Fact]
+    public async Task WithDeploymentState_ConcurrentAccess_IsThreadSafe()
+    {
+        // Arrange
+        var deploymentState = new JsonObject();
+        var context = ProvisioningTestHelpers.CreateTestProvisioningContext(userSecrets: deploymentState);
+        const int threadCount = 10;
+        const int iterationsPerThread = 100;
+        var tasks = new Task[threadCount];
+
+        // Act - Multiple threads accessing the DeploymentState concurrently via WithDeploymentState
+        for (int i = 0; i < threadCount; i++)
+        {
+            int threadId = i;
+            tasks[i] = Task.Run(() =>
+            {
+                for (int j = 0; j < iterationsPerThread; j++)
+                {
+                    context.WithDeploymentState(state =>
+                    {
+                        // All threads try to get or create the same "Azure" property
+                        var azureNode = state.Prop("Azure");
+                        
+                        // Each thread creates a unique property
+                        var threadNode = azureNode.Prop($"Thread{threadId}");
+                        threadNode.AsObject()["Counter"] = j;
+                        
+                        // And a shared property under Azure
+                        var deploymentsNode = azureNode.Prop("Deployments");
+                        
+                        // Access a deeper nested property
+                        var resourceNode = deploymentsNode.Prop($"Resource{j % 5}");
+                        resourceNode.AsObject()["LastAccess"] = $"Thread{threadId}-{j}";
+                    });
+                }
+            });
+        }
+
+        // Assert - Should complete without exceptions
+        await Task.WhenAll(tasks).WaitAsync(TimeSpan.FromSeconds(10));
+
+        // Verify the structure was created correctly
+        context.WithDeploymentState(state =>
+        {
+            Assert.NotNull(state["Azure"]);
+            var azureObj = state["Azure"]!.AsObject();
+            Assert.NotNull(azureObj["Deployments"]);
+            
+            // Check that all thread-specific nodes were created
+            for (int i = 0; i < threadCount; i++)
+            {
+                Assert.NotNull(azureObj[$"Thread{i}"]);
+            }
+        });
+    }
+
+    [Fact]
+    public void WithDeploymentState_Action_ExecutesSuccessfully()
+    {
+        // Arrange
+        var deploymentState = new JsonObject();
+        var context = ProvisioningTestHelpers.CreateTestProvisioningContext(userSecrets: deploymentState);
+
+        // Act
+        var executed = false;
+        context.WithDeploymentState(state =>
+        {
+            state["TestKey"] = "TestValue";
+            executed = true;
+        });
+
+        // Assert
+        Assert.True(executed);
+        Assert.Equal("TestValue", deploymentState["TestKey"]!.GetValue<string>());
+    }
+
+    [Fact]
+    public void WithDeploymentState_Func_ReturnsValue()
+    {
+        // Arrange
+        var deploymentState = new JsonObject();
+        deploymentState["TestKey"] = "TestValue";
+        var context = ProvisioningTestHelpers.CreateTestProvisioningContext(userSecrets: deploymentState);
+
+        // Act
+        var result = context.WithDeploymentState(state =>
+        {
+            return state["TestKey"]!.GetValue<string>();
+        });
+
+        // Assert
+        Assert.Equal("TestValue", result);
+    }
+
+    [Fact]
+    public async Task WithDeploymentState_ConcurrentReadsAndWrites_MaintainsConsistency()
+    {
+        // Arrange
+        var deploymentState = new JsonObject();
+        var context = ProvisioningTestHelpers.CreateTestProvisioningContext(userSecrets: deploymentState);
+        const int writerCount = 5;
+        const int readerCount = 5;
+        const int iterations = 100;
+
+        // Initialize counter
+        context.WithDeploymentState(state =>
+        {
+            state["Counter"] = 0;
+        });
+
+        var writerTasks = new Task[writerCount];
+        var readerTasks = new Task[readerCount];
+
+        // Act - Writers increment counter
+        for (int i = 0; i < writerCount; i++)
+        {
+            writerTasks[i] = Task.Run(() =>
+            {
+                for (int j = 0; j < iterations; j++)
+                {
+                    context.WithDeploymentState(state =>
+                    {
+                        var current = state["Counter"]!.GetValue<int>();
+                        state["Counter"] = current + 1;
+                    });
+                }
+            });
+        }
+
+        // Readers read counter
+        var readValues = new List<int>[readerCount];
+        for (int i = 0; i < readerCount; i++)
+        {
+            int readerIndex = i;
+            readValues[readerIndex] = new List<int>();
+            readerTasks[i] = Task.Run(() =>
+            {
+                for (int j = 0; j < iterations; j++)
+                {
+                    var value = context.WithDeploymentState(state =>
+                    {
+                        return state["Counter"]!.GetValue<int>();
+                    });
+                    readValues[readerIndex].Add(value);
+                    Thread.Sleep(1); // Small delay to allow interleaving
+                }
+            });
+        }
+
+        await Task.WhenAll(writerTasks.Concat(readerTasks)).WaitAsync(TimeSpan.FromSeconds(15));
+
+        // Assert - Final counter value should be exactly writerCount * iterations
+        var finalValue = context.WithDeploymentState(state =>
+        {
+            return state["Counter"]!.GetValue<int>();
+        });
+
+        Assert.Equal(writerCount * iterations, finalValue);
+
+        // All read values should be in valid range (0 to finalValue)
+        foreach (var readerValues in readValues)
+        {
+            Assert.All(readerValues, value =>
+            {
+                Assert.InRange(value, 0, finalValue);
+            });
+        }
+    }
 }
 
 public class ProvisioningServicesTests