diff --git a/src/ModularPipelines/Engine/ModuleContextProvider.cs b/src/ModularPipelines/Engine/ModuleContextProvider.cs index 82c4c866fd..bf3e58b26c 100644 --- a/src/ModularPipelines/Engine/ModuleContextProvider.cs +++ b/src/ModularPipelines/Engine/ModuleContextProvider.cs @@ -1,13 +1,16 @@ +using System.Collections.Concurrent; using Microsoft.Extensions.DependencyInjection; using ModularPipelines.Context; using ModularPipelines.Interfaces; namespace ModularPipelines.Engine; -internal class ModuleContextProvider : IPipelineContextProvider, IScopeDisposer +internal class ModuleContextProvider : IPipelineContextProvider, IScopeDisposer, IAsyncDisposable { private readonly IServiceProvider _serviceProvider; - private readonly List _scopes = new(); + private readonly ConcurrentBag _scopes = new(); + private readonly object _disposeLock = new(); + private bool _disposed; public ModuleContextProvider(IServiceProvider serviceProvider) { @@ -16,15 +19,84 @@ public ModuleContextProvider(IServiceProvider serviceProvider) public IPipelineContext GetModuleContext() { - var serviceScope = _serviceProvider.CreateAsyncScope(); + AsyncServiceScope serviceScope; - _scopes.Add(serviceScope); + lock (_disposeLock) + { + ObjectDisposedException.ThrowIf(_disposed, this); + + // Create and add scope inside lock to ensure atomicity with disposal. + // This prevents a race where disposal could take a snapshot between + // the _disposed check and the Add operation, leaving the scope undisposed. + serviceScope = _serviceProvider.CreateAsyncScope(); + _scopes.Add(serviceScope); + } return serviceScope.ServiceProvider.GetRequiredService(); } public IEnumerable GetScopes() { + // Return empty if disposed to prevent IScopeDisposer.Dispose from double-disposing + if (_disposed) + { + return []; + } + return _scopes; } -} \ No newline at end of file + + public async ValueTask DisposeAsync() + { + IServiceScope[] scopesToDispose; + + lock (_disposeLock) + { + if (_disposed) + { + return; + } + + _disposed = true; + + // Take snapshot inside lock to ensure atomicity. + // Any thread that passed the _disposed check in GetModuleContext + // before we set it will have already added to _scopes. + scopesToDispose = _scopes.ToArray(); + } + + foreach (var scope in scopesToDispose) + { + // AsyncServiceScope always implements IAsyncDisposable + await ((IAsyncDisposable)scope).DisposeAsync().ConfigureAwait(false); + } + + GC.SuppressFinalize(this); + } + + // Explicit implementation to coordinate with DisposeAsync and prevent double disposal + void IDisposable.Dispose() + { + IServiceScope[] scopesToDispose; + + lock (_disposeLock) + { + if (_disposed) + { + return; + } + + _disposed = true; + + // Take snapshot inside lock to ensure atomicity + scopesToDispose = _scopes.ToArray(); + } + + foreach (var scope in scopesToDispose) + { + scope.Dispose(); + } + + GC.SuppressFinalize(this); + } +}