.NET: fix: synchronize StateManager for executors that run concurrently in a superstep

dotnet/src/Microsoft.Agents.AI.Workflows/Execution/StateManager.cs

@@ -15,24 +15,53 @@ internal sealed class StateManager
     private readonly Dictionary<ScopeId, StateScope> _scopes = [];
     private readonly Dictionary<UpdateKey, StateUpdate> _queuedUpdates = [];
 
+    // InProcessRunner.RunSuperstepAsync delivers a superstep's messages to every receiving executor
+    // concurrently (one DeliverMessagesAsync task per receiver, awaited with Task.WhenAll), and each
+    // of those executors reaches this instance through its bound IWorkflowContext. Two executors
+    // queueing a state update in the same superstep therefore write _queuedUpdates from two threads,
+    // and Dictionary<,> is not safe for that: it throws
+    // "Operations that change non-concurrent collections must have exclusive access" or, worse,
+    // silently corrupts. Every access to the two dictionaries takes this lock; the scope contents
+    // themselves are only written by PublishUpdatesAsync and ImportStateAsync, which run between
+    // supersteps, so the await-ing reads of a StateScope stay outside the lock.
+    private readonly object _syncRoot = new();
+
     private StateScope GetOrCreateScope(ScopeId scopeId)
     {
         Throw.IfNull(scopeId);
 
-        if (!this._scopes.TryGetValue(scopeId, out StateScope? scope))
+        lock (this._syncRoot)
         {
-            scope = new StateScope(scopeId);
-            this._scopes[scopeId] = scope;
+            if (!this._scopes.TryGetValue(scopeId, out StateScope? scope))
+            {
+                scope = new StateScope(scopeId);
+                this._scopes[scopeId] = scope;
+            }
+
+            return scope;
         }
+    }
 
-        return scope;
+    private bool TryGetScope(ScopeId scopeId, out StateScope? scope)
+    {
+        lock (this._syncRoot)
+        {
+            return this._scopes.TryGetValue(scopeId, out scope);
+        }
     }
 
-    private IEnumerable<UpdateKey> GetUpdatesForScopeStrict(ScopeId scopeId)
+    /// <summary>
+    /// Snapshot of the queued updates for a scope. Materialized under the lock so the caller can
+    /// enumerate it while other executors keep queueing.
+    /// </summary>
+    private List<KeyValuePair<UpdateKey, StateUpdate>> GetUpdatesForScopeStrict(ScopeId scopeId)
     {
         Throw.IfNull(scopeId);
 
-        return this._queuedUpdates.Keys.Where(key => key.IsMatchingScope(scopeId, strict: true));
+        lock (this._syncRoot)
+        {
+            return this._queuedUpdates.Where(kvp => kvp.Key.IsMatchingScope(scopeId, strict: true)).ToList();
+        }
     }
 
     public ValueTask ClearStateAsync(string executorId, string? scopeName)
@@ -42,35 +71,38 @@ public async ValueTask ClearStateAsync(ScopeId scopeId)
     {
         Throw.IfNull(scopeId);
 
-        if (this._scopes.TryGetValue(scopeId, out StateScope? scope))
+        if (this.TryGetScope(scopeId, out StateScope? scope))
         {
-            HashSet<string> keysToDelete = await scope.ReadKeysAsync().ConfigureAwait(false);
+            HashSet<string> keysToDelete = await scope!.ReadKeysAsync().ConfigureAwait(false);
 
-            foreach (UpdateKey updateKey in this.GetUpdatesForScopeStrict(scopeId))
+            lock (this._syncRoot)
             {
-                StateUpdate update = this._queuedUpdates[updateKey];
-                if (!update.IsDelete)
+                foreach (KeyValuePair<UpdateKey, StateUpdate> queued in this._queuedUpdates.Where(kvp => kvp.Key.IsMatchingScope(scopeId, strict: true)).ToList())
                 {
-                    this._queuedUpdates[updateKey] = StateUpdate.Delete(update.Key);
-                }
+                    StateUpdate update = queued.Value;
+                    if (!update.IsDelete)
+                    {
+                        this._queuedUpdates[queued.Key] = StateUpdate.Delete(update.Key);
+                    }
 
-                keysToDelete.Remove(update.Key);
-            }
+                    keysToDelete.Remove(update.Key);
+                }
 
-            foreach (string key in keysToDelete)
-            {
-                UpdateKey updateKey = new(scopeId, key);
-                this._queuedUpdates[updateKey] = StateUpdate.Delete(key);
+                foreach (string key in keysToDelete)
+                {
+                    UpdateKey updateKey = new(scopeId, key);
+                    this._queuedUpdates[updateKey] = StateUpdate.Delete(key);
+                }
             }
         }
     }
 
     private HashSet<string> ApplyUnpublishedUpdates(ScopeId scopeId, HashSet<string> keys)
     {
         // Apply any queued updates for this scope
-        foreach (UpdateKey key in this.GetUpdatesForScopeStrict(scopeId))
+        foreach (KeyValuePair<UpdateKey, StateUpdate> queued in this.GetUpdatesForScopeStrict(scopeId))
         {
-            StateUpdate update = this._queuedUpdates[key];
+            StateUpdate update = queued.Value;
             if (update.IsDelete)
             {
                 keys.Remove(update.Key);
@@ -117,7 +149,13 @@ public ValueTask<T> ReadOrInitStateAsync<T>(string executorId, string? scopeName
         bool needsInit = false;
 
         // If there is executor-local state (from a queued update), read it first
-        if (this._queuedUpdates.TryGetValue(stateKey, out StateUpdate? update))
+        StateUpdate? update;
+        lock (this._syncRoot)
+        {
+            this._queuedUpdates.TryGetValue(stateKey, out update);
+        }
+
+        if (update is not null)
         {
             // What's the right thing to do when we have a state object, but it is the wrong type?
             if (update.IsDelete || update.Value is null)
@@ -182,7 +220,10 @@ public ValueTask WriteStateAsync<T>(ScopeId scopeId, string key, T value)
         Throw.IfNullOrEmpty(key);
 
         UpdateKey stateKey = new(scopeId, key);
-        this._queuedUpdates[stateKey] = StateUpdate.Update(key, value);
+        lock (this._syncRoot)
+        {
+            this._queuedUpdates[stateKey] = StateUpdate.Update(key, value);
+        }
 
         return default;
     }
@@ -194,17 +235,32 @@ public ValueTask ClearStateAsync(ScopeId scopeId, string key)
     {
         Throw.IfNullOrEmpty(key);
         UpdateKey stateKey = new(scopeId, key);
-        this._queuedUpdates[stateKey] = StateUpdate.Delete(key);
+        lock (this._syncRoot)
+        {
+            this._queuedUpdates[stateKey] = StateUpdate.Delete(key);
+        }
+
         return default;
     }
 
     public async ValueTask PublishUpdatesAsync(IStepTracer? tracer)
     {
+        // Snapshot the queued updates under the lock, then publish the snapshot without holding it:
+        // StateScope.WriteStateAsync awaits, and a lock cannot span an await. The snapshot is only
+        // removed from the queue once publication has succeeded (below), so a publish that throws —
+        // a conflicting shared-scope update, say — leaves every update queued, as it always has.
+        List<KeyValuePair<UpdateKey, StateUpdate>> queued;
+        lock (this._syncRoot)
+        {
+            queued = this._queuedUpdates.ToList();
+        }
+
         Dictionary<ScopeId, Dictionary<string, List<StateUpdate>>> updatesByScope = [];
 
         // Aggregate the updates for each scope
-        foreach (UpdateKey key in this._queuedUpdates.Keys)
+        foreach (KeyValuePair<UpdateKey, StateUpdate> entry in queued)
         {
+            UpdateKey key = entry.Key;
             if (!updatesByScope.TryGetValue(key.ScopeId, out Dictionary<string, List<StateUpdate>>? scopeUpdates))
             {
                 updatesByScope[key.ScopeId] = scopeUpdates = [];
@@ -215,7 +271,7 @@ public async ValueTask PublishUpdatesAsync(IStepTracer? tracer)
                 scopeUpdates[key.Key] = stateUpdates = [];
             }
 
-            stateUpdates.Add(this._queuedUpdates[key]);
+            stateUpdates.Add(entry.Value);
         }
 
         if (tracer is not null && (updatesByScope.Count > 0))
@@ -229,7 +285,20 @@ public async ValueTask PublishUpdatesAsync(IStepTracer? tracer)
             await stateScope.WriteStateAsync(updatesByScope[scope]).ConfigureAwait(false);
         }
 
-        this._queuedUpdates.Clear();
+        // Remove only what was published. An entry whose update was replaced while the snapshot was
+        // being written (a newer write queued concurrently for the same key) is kept for the next
+        // publish: the instance check tells the two apart without comparing values.
+        lock (this._syncRoot)
+        {
+            foreach (KeyValuePair<UpdateKey, StateUpdate> entry in queued)
+            {
+                if (this._queuedUpdates.TryGetValue(entry.Key, out StateUpdate? current)
+                    && ReferenceEquals(current, entry.Value))
+                {
+                    this._queuedUpdates.Remove(entry.Key);
+                }
+            }
+        }
     }
 
     private static IEnumerable<KeyValuePair<ScopeKey, PortableValue>> ExportScope(StateScope scope)
@@ -240,33 +309,44 @@ private static IEnumerable<KeyValuePair<ScopeKey, PortableValue>> ExportScope(St
         }
     }
 
-    internal async ValueTask<Dictionary<ScopeKey, PortableValue>> ExportStateAsync()
+    internal ValueTask<Dictionary<ScopeKey, PortableValue>> ExportStateAsync()
     {
-        if (this._queuedUpdates.Count != 0)
+        lock (this._syncRoot)
         {
-            throw new InvalidOperationException("Cannot export state while there are queued updates. Call PublishUpdatesAsync() first.");
-        }
+            if (this._queuedUpdates.Count != 0)
+            {
+                throw new InvalidOperationException("Cannot export state while there are queued updates. Call PublishUpdatesAsync() first.");
+            }
 
-        return this._scopes.Values.SelectMany(ExportScope).ToDictionary(kvp => kvp.Key, kvp => kvp.Value);
+            return new(this._scopes.Values.SelectMany(ExportScope).ToDictionary(kvp => kvp.Key, kvp => kvp.Value));
+        }
     }
 
     internal ValueTask ImportStateAsync(Checkpoint checkpoint)
     {
-        // TODO: Should this be a warning instead?
-        if (this._queuedUpdates.Count != 0)
+        lock (this._syncRoot)
         {
-            throw new InvalidOperationException("Cannot import state while there are queued updates. Call PublishUpdatesAsync() first.");
-        }
+            // TODO: Should this be a warning instead?
+            if (this._queuedUpdates.Count != 0)
+            {
+                throw new InvalidOperationException("Cannot import state while there are queued updates. Call PublishUpdatesAsync() first.");
+            }
 
-        this._queuedUpdates.Clear();
-        this._scopes.Clear();
+            this._queuedUpdates.Clear();
+            this._scopes.Clear();
 
-        Dictionary<ScopeKey, PortableValue> importedState = checkpoint.StateData;
+            Dictionary<ScopeKey, PortableValue> importedState = checkpoint.StateData;
 
-        foreach (ScopeKey scopeKey in importedState.Keys)
-        {
-            StateScope scope = this.GetOrCreateScope(scopeKey.ScopeId);
-            scope.ImportState(scopeKey.Key, importedState[scopeKey]);
+            foreach (ScopeKey scopeKey in importedState.Keys)
+            {
+                if (!this._scopes.TryGetValue(scopeKey.ScopeId, out StateScope? scope))
+                {
+                    scope = new StateScope(scopeKey.ScopeId);
+                    this._scopes[scopeKey.ScopeId] = scope;
+                }
+
+                scope.ImportState(scopeKey.Key, importedState[scopeKey]);
+            }
         }
 
         return default;

dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/StateManagerTests.cs

@@ -2,6 +2,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using System.Threading.Tasks;
 using Microsoft.Agents.AI.Workflows.Checkpointing;
 using Microsoft.Agents.AI.Workflows.Execution;
@@ -340,6 +341,38 @@ public async Task Test_PrivateScope_ConflictingUpdatesAsync()
         await RunConflictingUpdatesTest_WriteVsClearAsync(ScopeName, isSharedScope: false);
     }
 
+    [Fact]
+    public async Task Test_FailedPublish_LeavesUpdatesQueuedAsync()
+    {
+        // A conflicting write to a shared scope makes PublishUpdatesAsync throw. The queued updates
+        // must survive that: the next publish sees the same conflict rather than silently finding an
+        // empty queue, and nothing half-published is dropped.
+        StateManager manager = new();
+        ScopeId selfView = new("executor1", "shared");
+        ScopeId otherView = new("executor2", "shared");
+
+        await manager.WriteStateAsync(selfView, "key1", "value1");
+        await manager.WriteStateAsync(otherView, "key1", "value2");
+
+        await Assert.ThrowsAsync<InvalidOperationException>(async () => await manager.PublishUpdatesAsync(tracer: null));
+        await Assert.ThrowsAsync<InvalidOperationException>(async () => await manager.PublishUpdatesAsync(tracer: null));
+        await Assert.ThrowsAsync<InvalidOperationException>(async () => await manager.ExportStateAsync());
+    }
+
+    [Fact]
+    public async Task Test_SuccessfulPublish_EmptiesTheQueueAsync()
+    {
+        StateManager manager = new();
+        ScopeId scope = new("executor1", "shared");
+
+        await manager.WriteStateAsync(scope, "key1", "value1");
+        await manager.PublishUpdatesAsync(tracer: null);
+
+        Dictionary<ScopeKey, PortableValue> exported = await manager.ExportStateAsync();
+        Assert.Single(exported);
+        Assert.Equal("value1", await manager.ReadStateAsync<string>(scope, "key1"));
+    }
+
     private static async Task RunConflictingUpdatesTest_WriteVsWriteAsync(string? scopeName, bool isSharedScope)
     {
         const string SelfExecutorId = "executor1";
@@ -550,4 +583,50 @@ public async Task Test_LoadPortableValueState_AfterSerializationAsync()
         // Check that we don't double-wrap stored PortableValues on the out path
         VerifyIsNot<PortableValue>(pvAsPV);
     }
+
+    [Theory]
+    [InlineData("step_results")]
+    [InlineData(null)]
+    public async Task Test_ConcurrentExecutors_QueueAndReadState_WithoutCorruptionAsync(string? scopeName)
+    {
+        // InProcessRunner.RunSuperstepAsync delivers a superstep's messages to every receiving executor
+        // concurrently (one task per receiver, awaited together), and each executor reaches the same
+        // StateManager through its bound IWorkflowContext. Before the manager synchronized access to its
+        // queued-update dictionary, executors that finished in the same superstep raced on it and the
+        // run failed with "Operations that change non-concurrent collections must have exclusive access"
+        // — observed in a production host when four executors completed within 14 ms of each other.
+        const int Executors = 64;
+        const int WritesPerExecutor = 200;
+
+        StateManager manager = new();
+
+        await Task.WhenAll(Enumerable.Range(0, Executors).Select(e => Task.Run(async () =>
+        {
+            ScopeId scope = new($"executor{e}", scopeName);
+            for (int i = 0; i < WritesPerExecutor; i++)
+            {
+                string key = $"e{e}_key{i}";
+                await manager.WriteStateAsync(scope, key, $"value{e}:{i}");
+                Assert.Equal($"value{e}:{i}", await manager.ReadStateAsync<string>(scope, key));
+                _ = await manager.ReadKeysAsync(scope);
+            }
+        })));
+
+        await manager.PublishUpdatesAsync(tracer: null);
+
+        bool isSharedScope = scopeName is not null;
+        for (int e = 0; e < Executors; e++)
+        {
+            ScopeId scope = new($"executor{e}", scopeName);
+            HashSet<string> keys = await manager.ReadKeysAsync(scope);
+
+            // A shared scope is one bag every executor writes into; a private scope holds only its owner's keys.
+            Assert.Equal(isSharedScope ? Executors * WritesPerExecutor : WritesPerExecutor, keys.Count);
+            Assert.Equal($"value{e}:{WritesPerExecutor - 1}", await manager.ReadStateAsync<string>(scope, $"e{e}_key{WritesPerExecutor - 1}"));
+        }
+
+        // Nothing may be left queued after a publish, whichever thread queued it.
+        Dictionary<ScopeKey, PortableValue> exported = await manager.ExportStateAsync();
+        Assert.Equal(Executors * WritesPerExecutor, exported.Count);
+    }
 }