diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/StateManager.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/StateManager.cs index 3e34797bf01..b5963b2914c 100644 --- a/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/StateManager.cs +++ b/dotnet/src/Microsoft.Agents.AI.Workflows/Execution/StateManager.cs @@ -15,24 +15,53 @@ internal sealed class StateManager private readonly Dictionary _scopes = []; private readonly Dictionary _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 GetUpdatesForScopeStrict(ScopeId scopeId) + /// + /// Snapshot of the queued updates for a scope. Materialized under the lock so the caller can + /// enumerate it while other executors keep queueing. + /// + private List> 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,25 +71,28 @@ 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 keysToDelete = await scope.ReadKeysAsync().ConfigureAwait(false); + HashSet 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 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); + } } } } @@ -68,9 +100,9 @@ public async ValueTask ClearStateAsync(ScopeId scopeId) private HashSet ApplyUnpublishedUpdates(ScopeId scopeId, HashSet keys) { // Apply any queued updates for this scope - foreach (UpdateKey key in this.GetUpdatesForScopeStrict(scopeId)) + foreach (KeyValuePair 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 ReadOrInitStateAsync(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(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> queued; + lock (this._syncRoot) + { + queued = this._queuedUpdates.ToList(); + } + Dictionary>> updatesByScope = []; // Aggregate the updates for each scope - foreach (UpdateKey key in this._queuedUpdates.Keys) + foreach (KeyValuePair entry in queued) { + UpdateKey key = entry.Key; if (!updatesByScope.TryGetValue(key.ScopeId, out Dictionary>? 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 entry in queued) + { + if (this._queuedUpdates.TryGetValue(entry.Key, out StateUpdate? current) + && ReferenceEquals(current, entry.Value)) + { + this._queuedUpdates.Remove(entry.Key); + } + } + } } private static IEnumerable> ExportScope(StateScope scope) @@ -240,33 +309,44 @@ private static IEnumerable> ExportScope(St } } - internal async ValueTask> ExportStateAsync() + internal ValueTask> 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 importedState = checkpoint.StateData; + Dictionary 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; diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/StateManagerTests.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/StateManagerTests.cs index 2ddafbeb5bc..7310781a0a2 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Workflows.UnitTests/StateManagerTests.cs +++ b/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(async () => await manager.PublishUpdatesAsync(tracer: null)); + await Assert.ThrowsAsync(async () => await manager.PublishUpdatesAsync(tracer: null)); + await Assert.ThrowsAsync(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 exported = await manager.ExportStateAsync(); + Assert.Single(exported); + Assert.Equal("value1", await manager.ReadStateAsync(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(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(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 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(scope, $"e{e}_key{WritesPerExecutor - 1}")); + } + + // Nothing may be left queued after a publish, whichever thread queued it. + Dictionary exported = await manager.ExportStateAsync(); + Assert.Equal(Executors * WritesPerExecutor, exported.Count); + } }