Handle disposable transformers

Author: captainsafia

src/OpenApi/src/Services/OpenApiDocumentService.cs

@@ -70,7 +70,25 @@ private async Task ApplyTransformersAsync(OpenApiDocument document, Cancellation
         for (var i = 0; i < _options.DocumentTransformers.Count; i++)
         {
             var transformer = _options.DocumentTransformers[i];
-            await transformer.TransformAsync(document, documentTransformerContext, cancellationToken);
+            try
+            {
+                await transformer.TransformAsync(document, documentTransformerContext, cancellationToken);
+            }
+            finally
+            {
+                if (transformer is IDisposable disposable)
+                {
+                    disposable.Dispose();
+                }
+                if (transformer is IAsyncDisposable asyncDisposable)
+                {
+                    await asyncDisposable.DisposeAsync();
+                }
+                if (transformer is TypeBasedOpenApiDocumentTransformer typedTransformer)
+                {
+                    await typedTransformer.DisposeAsync();
+                }
+            }
         }
     }
 

src/OpenApi/src/Transformers/TypeBasedOpenApiDocumentTransformer.cs

@@ -10,10 +10,25 @@ namespace Microsoft.AspNetCore.OpenApi;
 internal sealed class TypeBasedOpenApiDocumentTransformer(Type transformerType) : IOpenApiDocumentTransformer
 {
     private readonly ObjectFactory _transformerFactory = ActivatorUtilities.CreateFactory(transformerType, []);
+    private IOpenApiDocumentTransformer? _transformer;
+
+    public ValueTask DisposeAsync()
+    {
+        if (_transformer is IAsyncDisposable asyncDisposable)
+        {
+            return asyncDisposable.DisposeAsync();
+        }
+        if (_transformer is IDisposable disposable)
+        {
+            disposable.Dispose();
+        }
+        return ValueTask.CompletedTask;
+    }
+
     public Task TransformAsync(OpenApiDocument document, OpenApiDocumentTransformerContext context, CancellationToken cancellationToken)
     {
-        var transformer = _transformerFactory.Invoke(context.ApplicationServices, []) as IOpenApiDocumentTransformer;
-        Debug.Assert(transformer != null, $"The type {transformerType} does not implement {nameof(IOpenApiDocumentTransformer)}.");
-        return transformer.TransformAsync(document, context, cancellationToken);
+        _transformer = _transformerFactory.Invoke(context.ApplicationServices, []) as IOpenApiDocumentTransformer;
+        Debug.Assert(_transformer != null, $"The type {transformerType} does not implement {nameof(IOpenApiDocumentTransformer)}.");
+        return _transformer.TransformAsync(document, context, cancellationToken);
     }
 }

src/OpenApi/test/Transformers/DocumentTransformerTests.cs

@@ -122,6 +122,84 @@ await VerifyOpenApiDocument(builder, options, document =>
         });
     }
 
+    [Fact]
+    public async Task DocumentTransformer_SupportsDisposableInstanceTransformer()
+    {
+        var builder = CreateBuilder();
+
+        builder.MapGet("/todo", () => { });
+        builder.MapGet("/user", () => { });
+
+        var options = new OpenApiOptions();
+        var transformer = new DisposableTransformer();
+        options.UseTransformer(transformer);
+
+        Assert.False(transformer.Disposed);
+        await VerifyOpenApiDocument(builder, options, document =>
+        {
+            Assert.Equal("Info Description", document.Info.Description);
+        });
+        Assert.True(transformer.Disposed);
+    }
+
+    [Fact]
+    public async Task DocumentTransformer_SupportsAsyncDisposableInstanceTransformer()
+    {
+        var builder = CreateBuilder();
+
+        builder.MapGet("/todo", () => { });
+        builder.MapGet("/user", () => { });
+
+        var options = new OpenApiOptions();
+        var transformer = new AsyncDisposableTransformer();
+        options.UseTransformer(transformer);
+
+        Assert.False(transformer.Disposed);
+        await VerifyOpenApiDocument(builder, options, document =>
+        {
+            Assert.Equal("Info Description", document.Info.Description);
+        });
+        Assert.True(transformer.Disposed);
+    }
+
+    [Fact]
+    public async Task DocumentTransformer_SupportsDisposableActivatedTransformer()
+    {
+        var builder = CreateBuilder();
+
+        builder.MapGet("/todo", () => { });
+        builder.MapGet("/user", () => { });
+
+        var options = new OpenApiOptions();
+        options.UseTransformer<DisposableTransformer>();
+
+        DisposableTransformer.DisposeCount = 0;
+        await VerifyOpenApiDocument(builder, options, document =>
+        {
+            Assert.Equal("Info Description", document.Info.Description);
+        });
+        Assert.Equal(1, DisposableTransformer.DisposeCount);
+    }
+
+    [Fact]
+    public async Task DocumentTransformer_SupportsAsyncDisposableActivatedTransformer()
+    {
+        var builder = CreateBuilder();
+
+        builder.MapGet("/todo", () => { });
+        builder.MapGet("/user", () => { });
+
+        var options = new OpenApiOptions();
+        options.UseTransformer<AsyncDisposableTransformer>();
+
+        AsyncDisposableTransformer.DisposeCount = 0;
+        await VerifyOpenApiDocument(builder, options, document =>
+        {
+            Assert.Equal("Info Description", document.Info.Description);
+        });
+        Assert.Equal(1, AsyncDisposableTransformer.DisposeCount);
+    }
+
     private class ActivatedTransformer : IOpenApiDocumentTransformer
     {
         public Task TransformAsync(OpenApiDocument document, OpenApiDocumentTransformerContext context, CancellationToken cancellationToken)
@@ -131,6 +209,43 @@ public Task TransformAsync(OpenApiDocument document, OpenApiDocumentTransformerC
         }
     }
 
+    private class DisposableTransformer : IOpenApiDocumentTransformer, IDisposable
+    {
+        internal bool Disposed = false;
+        internal static int DisposeCount = 0;
+
+        public void Dispose()
+        {
+            Disposed = true;
+            DisposeCount += 1;
+        }
+
+        public Task TransformAsync(OpenApiDocument document, OpenApiDocumentTransformerContext context, CancellationToken cancellationToken)
+        {
+            document.Info.Description = "Info Description";
+            return Task.CompletedTask;
+        }
+    }
+
+    private class AsyncDisposableTransformer : IOpenApiDocumentTransformer, IAsyncDisposable
+    {
+        internal bool Disposed = false;
+        internal static int DisposeCount = 0;
+
+        public ValueTask DisposeAsync()
+        {
+            Disposed = true;
+            DisposeCount += 1;
+            return ValueTask.CompletedTask;
+        }
+
+        public Task TransformAsync(OpenApiDocument document, OpenApiDocumentTransformerContext context, CancellationToken cancellationToken)
+        {
+            document.Info.Description = "Info Description";
+            return Task.CompletedTask;
+        }
+    }
+
     private class ActivatedTransformerWithDependency(Dependency dependency) : IOpenApiDocumentTransformer
     {
         public Task TransformAsync(OpenApiDocument document, OpenApiDocumentTransformerContext context, CancellationToken cancellationToken)