Merge branch 'experimental/non-blocking-DisposeAsync'

Yeah69 · Yeah69 · commit 9dcbb52fa8c5 · 2025-05-04T14:49:26.000+02:00
diff --git a/Main/CodeGeneration/DisposeUtility.cs b/Main/CodeGeneration/DisposeUtility.cs
@@ -26,6 +26,7 @@ internal interface IDisposeUtility
     string? DisposeExceptionHandlingAsyncFullyQualified { get; }
     string? DisposeSingularAsyncFullyQualified { get; }
     string? DisposeSingularAsyncSyncedFullyQualified { get; }
+    string? ReleaseDisposeAsyncFullyQualified { get; }
 }
 
 internal sealed class DisposeUtility : IDisposeUtility, IContainerInstance
@@ -48,6 +49,7 @@ internal sealed class DisposeUtility : IDisposeUtility, IContainerInstance
     private readonly string? _disposeExceptionHandlingSyncOnlyName;
     private readonly string? _disposeExceptionHandlingAsyncName;
     private readonly string _syncDisposalTriggeredExceptionName;
+    private readonly string? _releaseDisposeAsyncName;
 
     internal DisposeUtility(
         IReferenceGenerator referenceGenerator,
@@ -144,6 +146,12 @@ wellKnownTypes.IAsyncDisposable is not null
             : null;
         
         _syncDisposalTriggeredExceptionName = "SyncDisposalTriggeredException";
+        _releaseDisposeAsyncName = wellKnownTypes.IAsyncDisposable is not null && wellKnownTypes.ValueTask is not null
+            ? referenceGenerator.Generate("ReleaseDisposeAsync")
+            : null;
+        ReleaseDisposeAsyncFullyQualified = _releaseDisposeAsyncName is not null
+            ? $"{Constants.NamespaceForGeneratedStatics}.{ClassName}.{_releaseDisposeAsyncName}"
+            : null;
     }
 
     public DisposalUtilityInterfaceData DisposableRangeInterfaceData { get; }
@@ -160,6 +168,7 @@ wellKnownTypes.IAsyncDisposable is not null
     private string DisposeSingularFullyQualified { get; }
     public string? DisposeSingularAsyncFullyQualified { get; }
     public string? DisposeSingularAsyncSyncedFullyQualified { get; }
+    public string? ReleaseDisposeAsyncFullyQualified { get; }
     public string ClassName { get; }
 
     public string GenerateSingularDisposeFunctionsFile()
@@ -297,6 +306,23 @@ internal static class {{ClassName}}
                   """);
         }
 
+        if (_wellKnownTypes.IAsyncDisposable is not null && _wellKnownTypes.ValueTask is not null)
+        {
+            var disposedReference = _referenceGenerator.Generate("disposed");
+            var resolutionCounterReference = _referenceGenerator.Generate("resolutionCounter");
+            var awaitDisposalSourceReference = _referenceGenerator.Generate("awaitDisposalSource");
+            code.AppendLine(
+                $$"""
+                  internal static void {{_releaseDisposeAsyncName}}(ref int {{disposedReference}}, ref readonly int {{resolutionCounterReference}}, {{_wellKnownTypes.TaskCompletionSourceOfInt.FullName()}} {{awaitDisposalSourceReference}})
+                  {
+                      if ({{disposedReference}} == 0 || {{resolutionCounterReference}} > 0)
+                          return;
+                      if ({{_wellKnownTypes.Interlocked.FullName()}}.Exchange(ref {{disposedReference}}, 2) != 2)
+                          {{awaitDisposalSourceReference}}.SetResult(0);
+                  }
+                  """);
+        }
+
         code.AppendLine(
             """
                 }
diff --git a/Main/CodeGeneration/Nodes/FunctionNodeGenerator.cs b/Main/CodeGeneration/Nodes/FunctionNodeGenerator.cs
@@ -244,8 +244,13 @@ private void GenerateOneFunction(
                   finally
                   {
                   {{_wellKnownTypes.Interlocked.FullName()}}.{{nameof(Interlocked.Decrement)}}(ref {{_range.ResolutionCounterReference}});
-                  }
                   """);
+            if (_wellKnownTypes.ValueTask is not null && _wellKnownTypes.IAsyncDisposable is not null && _disposeUtility.ReleaseDisposeAsyncFullyQualified is {} releaseDisposeAsyncFullyQualified)
+            {
+                code.AppendLine(
+                    $"{releaseDisposeAsyncFullyQualified}(ref {_range.DisposalHandling.DisposedFieldReference}, ref {_range.ResolutionCounterReference}, {_range.ReleaseDisposeAsyncReference});");
+            }
+            code.AppendLine("}");
         }
 
         switch (_function)
diff --git a/Main/CodeGeneration/Nodes/RangeNodeGenerator.cs b/Main/CodeGeneration/Nodes/RangeNodeGenerator.cs
@@ -1,4 +1,5 @@
 using System.Threading;
+using System.Threading.Tasks;
 using MrMeeseeks.DIE.Configuration;
 using MrMeeseeks.DIE.Nodes.Functions;
 using MrMeeseeks.DIE.Nodes.Ranges;
@@ -63,10 +64,16 @@ public void Generate(StringBuilder code, ICodeGenerationVisitor visitor)
             code.AppendLine($"{DefaultConstructorDeclaredAccessibility}{_rangeNode.Name}() {{ }}");
         
         PreGeneralContent(code, visitor);
+
+        if (_wellKnownTypes.ValueTask is not null && _wellKnownTypes.IAsyncDisposable is not null)
+        {
+            code.AppendLine(
+                $"private readonly {_wellKnownTypes.TaskCompletionSourceOfInt.FullName()} {_rangeNode.ReleaseDisposeAsyncReference} = new {_wellKnownTypes.TaskCompletionSourceOfInt.FullName()}();");
+        }
         
         code.AppendLine(
             $$"""
-              private {{_wellKnownTypes.Int32.FullName()}} {{_rangeNode.ResolutionCounterReference}};
+              private {{_wellKnownTypes.Int32.FullName()}} {{_rangeNode.ResolutionCounterReference}} = 0;
               private {{_wellKnownTypes.ListOfListOfObject.FullName()}} {{_rangeNode.DisposalHandling.CollectionReference}} = new {{_wellKnownTypes.ListOfListOfObject.FullName()}}();
               """);
         foreach (var initializedInstance in _rangeNode.InitializedInstances)
@@ -151,6 +158,15 @@ void GenerateAddForDisposal(
                   finally
                   {
                   {{_wellKnownTypes.Interlocked.FullName()}}.{{nameof(Interlocked.Decrement)}}(ref {{_rangeNode.ResolutionCounterReference}});
+
+                  """);
+            if (_wellKnownTypes.ValueTask is not null && _wellKnownTypes.IAsyncDisposable is not null && _disposeUtility.ReleaseDisposeAsyncFullyQualified is {} releaseDisposeAsyncFullyQualified)
+            {
+                code.AppendLine(
+                    $"{releaseDisposeAsyncFullyQualified}(ref {_rangeNode.DisposalHandling.DisposedFieldReference}, ref {_rangeNode.ResolutionCounterReference}, {_rangeNode.ReleaseDisposeAsyncReference});");
+            }
+            code.AppendLine(
+                $$"""
                   }
                   }
                   """);
@@ -244,8 +260,16 @@ void GenerateDisposeFunction(bool isAsync)
                   {
                       var {{disposalHandling.DisposedLocalReference}} = {{_wellKnownTypes.Interlocked.FullName()}}.{{nameof(Interlocked.Exchange)}}(ref {{disposalHandling.DisposedFieldReference}}, 1);
                       if ({{disposalHandling.DisposedLocalReference}} != 0) return;
-                      {{_wellKnownTypes.SpinWait}}.{{nameof(SpinWait.SpinUntil)}}(() => {{_rangeNode.ResolutionCounterReference}} == 0);
                   """);
+            if (_disposeUtility.ReleaseDisposeAsyncFullyQualified is {} releaseDisposeAsyncFullyQualified)
+            {
+                code.AppendLine(
+                    $"{releaseDisposeAsyncFullyQualified}(ref {_rangeNode.DisposalHandling.DisposedFieldReference}, ref {_rangeNode.ResolutionCounterReference}, {_rangeNode.ReleaseDisposeAsyncReference});");
+            }
+
+            code.AppendLine(isAsync && _disposeUtility.ReleaseDisposeAsyncFullyQualified is not null
+                ? $"await {_rangeNode.ReleaseDisposeAsyncReference}.{nameof(TaskCompletionSource<int>.Task)};"
+                : $"{_wellKnownTypes.SpinWait}.{nameof(SpinWait.SpinUntil)}(() => {_rangeNode.ResolutionCounterReference} == 0);");
 
             switch (_rangeNode)
             {
diff --git a/Main/Nodes/Ranges/RangeNode.cs b/Main/Nodes/Ranges/RangeNode.cs
@@ -23,6 +23,7 @@ internal interface IRangeNode : INode
     string? ContainerReference { get; }
     IEnumerable<IInitializedInstanceNode> InitializedInstances { get; }
     string ResolutionCounterReference { get; }
+    string ReleaseDisposeAsyncReference { get; }
 
     IFunctionCallNode BuildCreateCall(ITypeSymbol type, IFunctionNode callingFunction, ImplementationMappingConfiguration? implementationMappingConfiguration = null);
     IWrappedAsyncFunctionCallNode BuildAsyncCreateCall(
@@ -101,6 +102,7 @@ internal abstract class RangeNode : IRangeNode
 
     public IEnumerable<IInitializedInstanceNode> InitializedInstances => InitializedInstanceNodesMap.Values;
     public string ResolutionCounterReference { get; }
+    public string ReleaseDisposeAsyncReference { get; }
 
     public IFunctionCallNode BuildEnumerableCall(INamedTypeSymbol type, IFunctionNode callingFunction,
         PassedContext passedContext) =>
@@ -278,6 +280,7 @@ internal RangeNode(
         }
         
         ResolutionCounterReference = referenceGenerator.Generate("resolutionCounter");
+        ReleaseDisposeAsyncReference = referenceGenerator.Generate("releaseDisposeAsync");
     }
     
     protected abstract IScopeManager ScopeManager { get; }
diff --git a/Main/WellKnownTypes.cs b/Main/WellKnownTypes.cs
@@ -14,6 +14,8 @@ internal sealed record WellKnownTypes(
     INamedTypeSymbol? ValueTask1, // .NET Standard 2.1
     INamedTypeSymbol Task, // .NET Standard 2.0
     INamedTypeSymbol Task1, // .NET Standard 2.0
+    INamedTypeSymbol TaskCompletionSource1, // .NET Standard 2.0
+    INamedTypeSymbol TaskCompletionSourceOfInt, // .NET Standard 2.0
     INamedTypeSymbol SpinWait, // .NET Standard 2.0
     INamedTypeSymbol Thread, // .NET Standard 2.0
     INamedTypeSymbol ObjectDisposedException, // .NET Standard 2.0
@@ -56,6 +58,7 @@ internal static WellKnownTypes Create(Compilation compilation)
         var exception = compilation.GetTypeByMetadataNameOrThrow("System.Exception");
         var task1 = compilation.GetTypeByMetadataNameOrThrow("System.Threading.Tasks.Task`1");
         var aggregateException = compilation.GetTypeByMetadataNameOrThrow("System.AggregateException");
+        var taskCompletionSource1 = compilation.GetTypeByMetadataNameOrThrow("System.Threading.Tasks.TaskCompletionSource`1");
         
         return new WellKnownTypes(
             IDisposable: iDisposable,
@@ -66,6 +69,8 @@ internal static WellKnownTypes Create(Compilation compilation)
             ValueTask1: valueTask1,
             Task: compilation.GetTypeByMetadataNameOrThrow("System.Threading.Tasks.Task"),
             Task1: task1,
+            TaskCompletionSource1: taskCompletionSource1,
+            TaskCompletionSourceOfInt: taskCompletionSource1.Construct(compilation.GetSpecialType(SpecialType.System_Int32)),
             SpinWait: compilation.GetTypeByMetadataNameOrThrow("System.Threading.SpinWait"),
             Thread: compilation.GetTypeByMetadataNameOrThrow("System.Threading.Thread"),
             ObjectDisposedException: compilation.GetTypeByMetadataNameOrThrow("System.ObjectDisposedException"),

Original file line number	Diff line number	Diff line change
`@@ -244,8 +244,13 @@ private void GenerateOneFunction(`
`244`	`244`	`finally`
`245`	`245`	`{`
`246`	`246`	`{{_wellKnownTypes.Interlocked.FullName()}}.{{nameof(Interlocked.Decrement)}}(ref {{_range.ResolutionCounterReference}});`
`247`		`- }`
`248`	`247`	`""");`
	`248`	`+ if (_wellKnownTypes.ValueTask is not null && _wellKnownTypes.IAsyncDisposable is not null && _disposeUtility.ReleaseDisposeAsyncFullyQualified is {} releaseDisposeAsyncFullyQualified)`
	`249`	`+ {`
	`250`	`+ code.AppendLine(`
	`251`	`+ $"{releaseDisposeAsyncFullyQualified}(ref {_range.DisposalHandling.DisposedFieldReference}, ref {_range.ResolutionCounterReference}, {_range.ReleaseDisposeAsyncReference});");`
	`252`	`+ }`
	`253`	`+ code.AppendLine("}");`
`249`	`254`	`}`
`250`	`255`
`251`	`256`	`switch (_function)`