Make transient scope disposal tracking thread safe with a semaphore

Yeah69 · Yeah69 · commit f32cfd245035 · 2025-05-29T17:04:02.000+02:00
diff --git a/Main/CodeGeneration/CodeGenerationVisitor.cs b/Main/CodeGeneration/CodeGenerationVisitor.cs
@@ -164,10 +164,8 @@ public void VisitIScopeCallNode(IScopeCallNode scopeCall)
     public void VisitITransientScopeCallNode(ITransientScopeCallNode transientScopeCall)
     {
         VisitIElementNode(transientScopeCall.ScopeConstruction);
-        var owner = transientScopeCall.ContainerReference is { } containerReference
-            ? $"{containerReference}."
-            : "";
-        _code.AppendLine($"{owner}{transientScopeCall.TransientScopeDisposalReference}.{nameof(List<object>.Add)}({transientScopeCall.ScopeConstruction.Reference});");
+        if (transientScopeCall.TransientScopeDisposalParameter?.Calling.Reference is {} reference)
+            _code.AppendLine($"{reference}.{nameof(List<object>.Add)}({transientScopeCall.ScopeConstruction.Reference});");
         GenerateInitialization(transientScopeCall.Initialization, transientScopeCall.ScopeConstruction.Reference);
         VisitIFunctionCallNode(transientScopeCall);
     }
diff --git a/Main/CodeGeneration/Nodes/ContainerNodeGenerator.cs b/Main/CodeGeneration/Nodes/ContainerNodeGenerator.cs
@@ -58,6 +58,7 @@ protected override void PostGeneralContent(StringBuilder code, ICodeGenerationVi
               private interface {{_containerNode.ScopeInterface}} : {{GenerateDisposalInterfaceAssignments()}} {}
 
               private {{_wellKnownTypes.ListOfObject.FullName()}} {{_containerNode.TransientScopeDisposalReference}} = new {{_wellKnownTypes.ListOfObject.FullName()}}();
+              private {{_wellKnownTypes.SemaphoreSlim.FullName()}} {{_containerNode.TransientScopeDisposalSemaphoreReference}} = new {{_wellKnownTypes.SemaphoreSlim.FullName()}}(1);
               """);
         
         visitor.VisitITransientScopeInterfaceNode(_containerNode.TransientScopeInterface);
diff --git a/Main/CodeGeneration/Nodes/FunctionNodeGenerator.cs b/Main/CodeGeneration/Nodes/FunctionNodeGenerator.cs
@@ -175,8 +175,22 @@ private void GenerateOneFunction(
             var containerReference = _range.ContainerReference is { } reference
                 ? $"{reference}."
                 : "";
+            
+            var awaitPrefix = isAsyncAwait ? "await " : "";
+            var waitMethod = isAsyncAwait ? nameof(SemaphoreSlim.WaitAsync) : nameof(SemaphoreSlim.Wait);
+            
             code.AppendLine(
-                $"{containerReference}{_container.TransientScopeDisposalReference}.{nameof(List<object>.AddRange)}({_function.TransientScopeDisposalNode.Reference});");
+                $$"""
+                  {{awaitPrefix}}{{containerReference}}{{_container.TransientScopeDisposalSemaphoreReference}}.{{waitMethod}}();
+                  try 
+                  {
+                  {{containerReference}}{{_container.TransientScopeDisposalReference}}.{{nameof(List<object>.AddRange)}}({{_function.TransientScopeDisposalNode.Reference}});
+                  }
+                  finally
+                  {
+                  {{containerReference}}{{_container.TransientScopeDisposalSemaphoreReference}}.{{nameof(SemaphoreSlim.Release)}}();
+                  }
+                  """);
         }
         
 
diff --git a/Main/CodeGeneration/Nodes/RangeNodeGenerator.cs b/Main/CodeGeneration/Nodes/RangeNodeGenerator.cs
@@ -193,7 +193,23 @@ _wellKnownTypes.IAsyncDisposable is not null
         {
             case IContainerNode container:
                 code.AppendLine(
-                    $"{_wellKnownTypes.Object.FullName()}[] {_disposeUtility.DisposableRangeInterfaceData.InterfaceNameFullyQualified}.{_disposeUtility.DisposableRangeInterfaceData.TransientScopesPropertyName} => {_wellKnownTypesCollections.Enumerable}.{nameof(Enumerable.ToArray)}({container.TransientScopeDisposalReference});");
+                    $$"""
+                      {{_wellKnownTypes.Object.FullName()}}[] {{_disposeUtility.DisposableRangeInterfaceData.InterfaceNameFullyQualified}}.{{_disposeUtility.DisposableRangeInterfaceData.TransientScopesPropertyName}}
+                      {
+                      get
+                      {
+                      {{container.TransientScopeDisposalSemaphoreReference}}.{{nameof(SemaphoreSlim.Wait)}}();
+                      try
+                      {
+                      return {{_wellKnownTypesCollections.Enumerable}}.{{nameof(Enumerable.ToArray)}}({{container.TransientScopeDisposalReference}});
+                      }
+                      finally
+                      {
+                      {{container.TransientScopeDisposalSemaphoreReference}}.{{nameof(SemaphoreSlim.Release)}}();
+                      }
+                      }
+                      }
+                      """);
                 break;
             default:
                 code.AppendLine(
@@ -275,10 +291,19 @@ void GenerateDisposeFunction(bool isAsync)
             switch (_rangeNode)
             {
                 case ITransientScopeNode transientScope:
+                    var waitMethod = isAsync ? nameof(SemaphoreSlim.WaitAsync) : nameof(SemaphoreSlim.Wait);
                     code.AppendLine(
                         $$"""
-                          {{transientScope.ContainerReference}}.{{transientScope.TransientScopeDisposalReference}}.{{nameof(List<object>.Remove)}}(this);
-                          {{transientScope.ContainerReference}}.{{transientScope.TransientScopeDisposalReference}}.{{nameof(List<object>.TrimExcess)}}();
+                          {{awaitPrefix}}{{transientScope.ContainerReference}}.{{_containerNode.TransientScopeDisposalSemaphoreReference}}.{{waitMethod}}();
+                          try 
+                          {
+                          {{transientScope.ContainerReference}}.{{_containerNode.TransientScopeDisposalReference}}.{{nameof(List<object>.Remove)}}(this);
+                          {{transientScope.ContainerReference}}.{{_containerNode.TransientScopeDisposalReference}}.{{nameof(List<object>.TrimExcess)}}();
+                          }
+                          finally
+                          {
+                          {{transientScope.ContainerReference}}.{{_containerNode.TransientScopeDisposalSemaphoreReference}}.{{nameof(SemaphoreSlim.Release)}}();
+                          }
                           """);
                     break;
             }
diff --git a/Main/MsContainer/ExecuteLevelContainer.cs b/Main/MsContainer/ExecuteLevelContainer.cs
@@ -49,7 +49,7 @@ private IExecuteContainer DIE_Factory_IExecuteContainer(
             DisposeUtility disposeUtility,
             ReferenceGeneratorCounter referenceGeneratorCounter)
         {
-#pragma warning disable CA2000 *** Manually added for disposal
+#pragma warning disable CA2000 // Manually added for disposal
             var container = ContainerLevelContainer.DIE_CreateContainer(
                 DIE_Factory_GeneratorExecutionContext, 
                 containerInfo,
diff --git a/Main/Nodes/Ranges/ContainerNode.cs b/Main/Nodes/Ranges/ContainerNode.cs
@@ -25,6 +25,7 @@ internal interface IContainerNode : IRangeNode
     ITransientScopeInterfaceNode TransientScopeInterface { get; }
     string ScopeInterface { get; }
     string TransientScopeDisposalReference { get; }
+    string TransientScopeDisposalSemaphoreReference { get; }
     IFunctionCallNode BuildContainerInstanceCall(string? ownerReference, INamedTypeSymbol type, IFunctionNode callingFunction);
     IReadOnlyList<ICreateContainerFunctionNode> CreateContainerFunctions { get; }
     bool AsyncDisposablesPossible { get; }
@@ -60,6 +61,7 @@ internal sealed partial class ContainerNode : RangeNode, IContainerNode, IContai
     public ITransientScopeInterfaceNode TransientScopeInterface => _lazyTransientScopeInterfaceNode.Value;
     public string ScopeInterface { get; }
     public string TransientScopeDisposalReference { get; }
+    public string TransientScopeDisposalSemaphoreReference { get; }
 
     public IFunctionCallNode BuildContainerInstanceCall(
         string? ownerReference, 
@@ -141,6 +143,8 @@ internal ContainerNode(
         _containerNodeGenerator = containerNodeGenerator;
 
         TransientScopeDisposalReference = referenceGenerator.Generate("transientScopeDisposal");
+
+        TransientScopeDisposalSemaphoreReference = referenceGenerator.Generate("transientScopeDisposalLock");
         
         GenerateEmptyConstructor = !_containerInfo.ContainerType.InstanceConstructors.Any(ic => !ic.IsImplicitlyDeclared);
         
diff --git a/Main/Nodes/Ranges/TransientScopeNode.cs b/Main/Nodes/Ranges/TransientScopeNode.cs
@@ -15,7 +15,6 @@ namespace MrMeeseeks.DIE.Nodes.Ranges;
 internal interface ITransientScopeNode : IScopeNodeBase
 {
     string TransientScopeInterfaceName { get; }
-    string TransientScopeDisposalReference { get; }
     ITransientScopeCallNode BuildTransientScopeCallFunction(
         string containerParameter,
         INamedTypeSymbol type,
@@ -77,7 +76,6 @@ internal TransientScopeNode(
         _transientScopeNodeGenerator = transientScopeNodeGenerator;
         _createTransientScopeFunctionNodeFactory = createTransientScopeFunctionNodeFactory;
         TransientScopeInterfaceName = parentContainer.TransientScopeInterface.Name;
-        TransientScopeDisposalReference = parentContainer.TransientScopeDisposalReference;
     }
     protected override string ContainerParameterForScope => ContainerReference;
 
@@ -92,7 +90,6 @@ public override IFunctionCallNode BuildTransientScopeInstanceCall(INamedTypeSymb
         ParentContainer.TransientScopeInterface.BuildTransientScopeInstanceCall($"({Constants.ThisKeyword} as {ParentContainer.TransientScopeInterface.FullName})", type, callingFunction);
 
     public string TransientScopeInterfaceName { get; }
-    public string TransientScopeDisposalReference { get; }
 
     public ITransientScopeCallNode BuildTransientScopeCallFunction(
         string containerParameter,
diff --git a/Sample/Container.cs b/Sample/Container.cs
@@ -1,19 +1,18 @@
 ﻿using System;
-using System.Threading.Tasks;
 using MrMeeseeks.DIE.Configuration.Attributes;
 using MrMeeseeks.DIE.UserUtility;
 
 namespace MrMeeseeks.DIE.Sample;
 
-//*
-internal class Class : IContainerInstance, ITaskInitializer
-{
-    public async Task InitializeAsync()
-    {
-        await Task.Delay(1000);
-    }
-}
+internal sealed class Disposable : IDisposable { public void Dispose() { } }
 
-[CreateFunction(typeof(Func<Class>), "Create")]
+internal sealed class Class(ClassBelow _, Disposable __) : ITransientScopeRoot;
+
+internal sealed class ClassBelow(ClassS _, Disposable __) : ITransientScopeRoot;
+
+internal sealed class ClassS(ClassBelowS _, Disposable __) : IScopeRoot;
+
+internal sealed class ClassBelowS(Disposable __) : IScopeRoot;
+
+[CreateFunction(typeof(Class), "Create")]
 internal sealed partial class Container;
-//*/

Original file line number	Diff line number	Diff line change
`@@ -164,10 +164,8 @@ public void VisitIScopeCallNode(IScopeCallNode scopeCall)`
`164`	`164`	`public void VisitITransientScopeCallNode(ITransientScopeCallNode transientScopeCall)`
`165`	`165`	`{`
`166`	`166`	`VisitIElementNode(transientScopeCall.ScopeConstruction);`
`167`		`- var owner = transientScopeCall.ContainerReference is { } containerReference`
`168`		`- ? $"{containerReference}."`
`169`		`- : "";`
`170`		`- _code.AppendLine($"{owner}{transientScopeCall.TransientScopeDisposalReference}.{nameof(List<object>.Add)}({transientScopeCall.ScopeConstruction.Reference});");`
	`167`	`+ if (transientScopeCall.TransientScopeDisposalParameter?.Calling.Reference is {} reference)`
	`168`	`+ _code.AppendLine($"{reference}.{nameof(List<object>.Add)}({transientScopeCall.ScopeConstruction.Reference});");`
`171`	`169`	`GenerateInitialization(transientScopeCall.Initialization, transientScopeCall.ScopeConstruction.Reference);`
`172`	`170`	`VisitIFunctionCallNode(transientScopeCall);`
`173`	`171`	`}`
Original file line number	Diff line number	Diff line change
`@@ -49,7 +49,7 @@ private IExecuteContainer DIE_Factory_IExecuteContainer(`
`49`	`49`	`DisposeUtility disposeUtility,`
`50`	`50`	`ReferenceGeneratorCounter referenceGeneratorCounter)`
`51`	`51`	`{`
`52`		`-#pragma warning disable CA2000 *** Manually added for disposal`
	`52`	`+#pragma warning disable CA2000 // Manually added for disposal`
`53`	`53`	`var container = ContainerLevelContainer.DIE_CreateContainer(`
`54`	`54`	`DIE_Factory_GeneratorExecutionContext,`
`55`	`55`	`containerInfo,`