Added hooks for better schema phase outs (#8748)

Author: michaelstaib

src/HotChocolate/Core/src/Execution/Configuration/RequestExecutorSetup.cs

@@ -33,6 +33,11 @@ public sealed class RequestExecutorSetup
     /// </summary>
     public RequestExecutorOptions? RequestExecutorOptions { get; set; }
 
+    /// <summary>
+    /// Gets or sets the time that the executor manager waits to dispose the schema services.
+    /// </summary>
+    public TimeSpan EvictionTimeout { get; set; } = TimeSpan.FromSeconds(30);
+
     /// <summary>
     /// Gets the request executor options actions.
     /// This hook is invoke first in the schema creation process.
@@ -117,6 +122,7 @@ public void CopyTo(RequestExecutorSetup options)
         options._onBuildDocumentValidatorHooks.AddRange(_onBuildDocumentValidatorHooks);
         options._pipelineModifiers.AddRange(_pipelineModifiers);
         options._typeModules.AddRange(_typeModules);
+        options.EvictionTimeout = EvictionTimeout;
 
         if (DefaultPipelineFactory is not null)
         {

src/HotChocolate/Core/src/Execution/HotChocolate.Execution.csproj

@@ -24,6 +24,7 @@
     <InternalsVisibleTo Include="HotChocolate.Execution.Projections" />
     <InternalsVisibleTo Include="HotChocolate.Execution.Tests" />
     <InternalsVisibleTo Include="HotChocolate.Fusion.Tests" />
+    <InternalsVisibleTo Include="HotChocolate.Types.Tests" />
     <InternalsVisibleTo Include="HotChocolate.Fusion" />
     <InternalsVisibleTo Include="HotChocolate.Types.Mutations" />
     <InternalsVisibleTo Include="StrawberryShake.CodeGeneration" />
@@ -219,6 +220,15 @@
     <Compile Update="DependencyInjection\RequestExecutorBuilderExtensions.Services.cs">
       <DependentUpon>RequestExecutorBuilderExtensions.cs</DependentUpon>
     </Compile>
+    <Compile Update="RequestExecutorManager.Events.cs">
+      <DependentUpon>RequestExecutorManager.cs</DependentUpon>
+    </Compile>
+    <Compile Update="RequestExecutorManager.Hooks.cs">
+      <DependentUpon>RequestExecutorManager.cs</DependentUpon>
+    </Compile>
+    <Compile Update="RequestExecutorManager.Warmup.cs">
+      <DependentUpon>RequestExecutorManager.cs</DependentUpon>
+    </Compile>
   </ItemGroup>
 
   <ItemGroup>

src/HotChocolate/Core/src/Execution/RequestExecutorManager.cs

@@ -27,7 +27,7 @@ internal sealed partial class RequestExecutorManager
     : IRequestExecutorManager
     , IRequestExecutorEvents
     , IRequestExecutorWarmup
-    , IDisposable
+    , IAsyncDisposable
 {
     private readonly CancellationTokenSource _cts = new();
     private readonly ConcurrentDictionary<string, SemaphoreSlim> _semaphoreBySchema = new();
@@ -183,7 +183,8 @@ await CreateSchemaServicesAsync(context, setup, typeModuleChangeMonitor, cancell
             schemaServices,
             schemaServices.GetRequiredService<IExecutionDiagnosticEvents>(),
             setup,
-            typeModuleChangeMonitor);
+            typeModuleChangeMonitor,
+            setup.EvictionTimeout);
 
         var executor = registeredExecutor.Executor;
 
@@ -254,12 +255,12 @@ private static async Task RunEvictionEvents(RegisteredExecutor registeredExecuto
         {
             // we will give the request executor some grace period to finish all requests
             // in the pipeline.
-            await Task.Delay(TimeSpan.FromMinutes(5));
-            registeredExecutor.Dispose();
+            await Task.Delay(registeredExecutor.EvictionTimeout).ConfigureAwait(false);
+            await registeredExecutor.DisposeAsync().ConfigureAwait(false);
         }
     }
 
-    private async Task<IServiceProvider> CreateSchemaServicesAsync(
+    private async Task<ServiceProvider> CreateSchemaServicesAsync(
         ConfigurationContext context,
         RequestExecutorSetup setup,
         TypeModuleChangeMonitor typeModuleChangeMonitor,
@@ -518,16 +519,16 @@ private static RequestDelegate CreatePipeline(
         return next;
     }
 
-    public void Dispose()
+    public async ValueTask DisposeAsync()
     {
         if (!_disposed)
         {
             // this will stop the eviction processor.
-            _cts.Cancel();
+            await _cts.CancelAsync();
 
             foreach (var executor in _executors.Values)
             {
-                executor.Dispose();
+                await executor.DisposeAsync();
             }
 
             foreach (var semaphore in _semaphoreBySchema.Values)
@@ -545,33 +546,32 @@ public void Dispose()
 
     private sealed class RegisteredExecutor(
         IRequestExecutor executor,
-        IServiceProvider services,
+        ServiceProvider services,
         IExecutionDiagnosticEvents diagnosticEvents,
         RequestExecutorSetup setup,
-        TypeModuleChangeMonitor typeModuleChangeMonitor)
-        : IDisposable
+        TypeModuleChangeMonitor typeModuleChangeMonitor,
+        TimeSpan evictionTimeout)
+        : IAsyncDisposable
     {
         private bool _disposed;
 
         public IRequestExecutor Executor { get; } = executor;
 
-        public IServiceProvider Services { get; } = services;
+        public ServiceProvider Services { get; } = services;
 
         public IExecutionDiagnosticEvents DiagnosticEvents { get; } = diagnosticEvents;
 
         public RequestExecutorSetup Setup { get; } = setup;
 
         public TypeModuleChangeMonitor TypeModuleChangeMonitor { get; } = typeModuleChangeMonitor;
 
-        public void Dispose()
+        public TimeSpan EvictionTimeout { get; } = evictionTimeout;
+
+        public async ValueTask DisposeAsync()
         {
-            if (_disposed)
+            if (!_disposed)
             {
-                if (Services is IDisposable d)
-                {
-                    d.Dispose();
-                }
-
+                await Services.DisposeAsync();
                 TypeModuleChangeMonitor.Dispose();
                 _disposed = true;
             }

src/HotChocolate/Core/src/Types.Abstractions/Features/SchemaCancellationFeature.cs

@@ -0,0 +1,56 @@
+namespace HotChocolate.Features;
+
+/// <summary>
+/// Provides schema cancellation support for HotChocolate.
+/// </summary>
+/// <remarks>
+/// This feature attaches a <see cref="CancellationToken"/> to a schema, allowing
+/// long-lived operations such as subscriptions to be gracefully canceled when
+/// the schema is phased out (for example, during schema version replacement).
+/// </remarks>
+public sealed class SchemaCancellationFeature : IAsyncDisposable
+{
+    private readonly CancellationTokenSource _cts = new();
+    private bool _disposed;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="SchemaCancellationFeature"/> class.
+    /// </summary>
+    /// <remarks>
+    /// The internal <see cref="CancellationTokenSource"/> is created upon construction,
+    /// and the <see cref="CancellationToken"/> property is set to its token.
+    /// </remarks>
+    public SchemaCancellationFeature()
+    {
+        CancellationToken = _cts.Token;
+    }
+
+    /// <summary>
+    /// Gets the schema cancellation token.
+    /// </summary>
+    /// <remarks>
+    /// This token will be triggered when the schema is phased out. Components such
+    /// as subscriptions and background tasks can observe this token and terminate
+    /// gracefully when cancellation is requested.
+    /// </remarks>
+    public CancellationToken CancellationToken { get; }
+
+    /// <summary>
+    /// Disposes the <see cref="SchemaCancellationFeature"/> and requests cancellation.
+    /// </summary>
+    /// <remarks>
+    /// This method cancels the underlying <see cref="CancellationTokenSource"/>,
+    /// ensuring that all operations observing the token are notified. It should be
+    /// called when the schema is being phased out.
+    /// </remarks>
+    /// <returns>A task representing the asynchronous dispose operation.</returns>
+    public async ValueTask DisposeAsync()
+    {
+        if (!_disposed)
+        {
+            await _cts.CancelAsync().ConfigureAwait(false);
+            _cts.Dispose();
+            _disposed = true;
+        }
+    }
+}

src/HotChocolate/Core/src/Types.Abstractions/Features/SchemaCancellationSchemaDefinitionExtensions.cs

@@ -0,0 +1,27 @@
+namespace HotChocolate.Features;
+
+/// <summary>
+/// Provides extension methods for accessing the schema cancellation feature
+/// from an <see cref="ISchemaDefinition"/>.
+/// </summary>
+public static class SchemaCancellationSchemaDefinitionExtensions
+{
+    /// <summary>
+    /// Gets the <see cref="CancellationToken"/> associated with the schema.
+    /// </summary>
+    /// <param name="schema">
+    /// The <see cref="ISchemaDefinition"/> to retrieve the cancellation token from.
+    /// </param>
+    /// <returns>
+    /// The schema's cancellation token if the <see cref="SchemaCancellationFeature"/> is present;
+    /// otherwise, <see cref="CancellationToken.None"/>.
+    /// </returns>
+    /// <exception cref="ArgumentNullException">
+    /// Thrown if <paramref name="schema"/> is <c>null</c>.
+    /// </exception>
+    public static CancellationToken GetCancellationToken(this ISchemaDefinition schema)
+    {
+        ArgumentNullException.ThrowIfNull(schema);
+        return schema.Features.Get<SchemaCancellationFeature>()?.CancellationToken ?? CancellationToken.None;
+    }
+}

src/HotChocolate/Core/test/Types.Tests/CodeFirstTests.cs

@@ -252,6 +252,42 @@ public void Disallow_Implicitly_Binding_Object()
             () => SchemaBuilder.New().BindRuntimeType<object, StringType>());
     }
 
+    [Fact]
+    public async Task DisposeSchemaService()
+    {
+        // arrange
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(60));
+        var customService = new CustomService();
+        var customAsyncService = new CustomAsyncService();
+
+        var services = new ServiceCollection()
+            .AddGraphQLServer()
+            .AddQueryType(c => c.Field("a").Resolve("b"))
+            .ConfigureSchemaServices(s =>
+            {
+                s.AddSingleton(_ => customService);
+                s.AddSingleton(_ => customAsyncService);
+            })
+            .Configure(s => s.EvictionTimeout = TimeSpan.FromSeconds(3))
+            .Services
+            .BuildServiceProvider();
+
+        var manager = services.GetRequiredService<IRequestExecutorManager>();
+        var executor = await manager.GetExecutorAsync(cancellationToken: cts.Token);
+        Assert.NotNull(executor.Schema.Services.GetService<CustomService>());
+        Assert.NotNull(executor.Schema.Services.GetService<CustomAsyncService>());
+
+        // act
+        manager.EvictExecutor();
+
+        // assert
+        await customService.WaitAsync(cts.Token);
+        await customAsyncService.WaitAsync(cts.Token);
+
+        Assert.True(customService.Triggered);
+        Assert.True(customAsyncService.Triggered);
+    }
+
     public class Query
     {
         public string SayHello(string name) =>
@@ -438,4 +474,43 @@ public Example NestedClassNullableString()
             return new Example();
         }
     }
+
+    public class CustomService : IDisposable
+    {
+        private readonly TaskCompletionSource _tcs = new();
+
+        public bool Triggered { get; set; }
+
+        public async Task WaitAsync(CancellationToken cancellationToken)
+        {
+            cancellationToken.Register(() => _tcs.TrySetCanceled());
+            await _tcs.Task;
+        }
+
+        public void Dispose()
+        {
+            Triggered = true;
+            _tcs.TrySetResult();
+        }
+    }
+
+    public class CustomAsyncService : IAsyncDisposable
+    {
+        private readonly TaskCompletionSource _tcs = new();
+
+        public bool Triggered { get; set; }
+
+        public async Task WaitAsync(CancellationToken cancellationToken)
+        {
+            cancellationToken.Register(() => _tcs.TrySetCanceled());
+            await _tcs.Task;
+        }
+
+        public ValueTask DisposeAsync()
+        {
+            Triggered = true;
+            _tcs.TrySetResult();
+            return default;
+        }
+    }
 }

src/HotChocolate/Fusion-vnext/src/Fusion.Execution.Types/FusionSchemaDefinition.cs

@@ -9,10 +9,11 @@
 using HotChocolate.Language.Visitors;
 using HotChocolate.Serialization;
 using HotChocolate.Types;
+using Microsoft.Extensions.DependencyInjection;
 
 namespace HotChocolate.Fusion.Types;
 
-public sealed class FusionSchemaDefinition : ISchemaDefinition
+public sealed class FusionSchemaDefinition : ISchemaDefinition, IAsyncDisposable
 {
 #if NET9_0_OR_GREATER
     private readonly Lock _lock = new();
@@ -22,9 +23,11 @@ public sealed class FusionSchemaDefinition : ISchemaDefinition
     private readonly ConcurrentDictionary<string, ImmutableArray<FusionObjectTypeDefinition>> _possibleTypes = new();
     private readonly ConcurrentDictionary<(string, string?), ImmutableArray<Lookup>> _possibleLookups = new();
     private readonly ConcurrentDictionary<TransitionKey, Lookup> _bestDirectLookup = new();
+    private readonly IServiceProvider _services;
     private ImmutableArray<FusionUnionTypeDefinition> _unionTypes;
     private IFeatureCollection _features;
     private bool _sealed;
+    private bool _disposed;
 
     internal FusionSchemaDefinition(
         string name,
@@ -40,7 +43,7 @@ internal FusionSchemaDefinition(
     {
         Name = name;
         Description = description;
-        Services = services;
+        _services = services;
         QueryType = queryType;
         MutationType = mutationType;
         SubscriptionType = subscriptionType;
@@ -85,7 +88,7 @@ public static FusionSchemaDefinition Create(
     /// <summary>
     /// Gets the schema services.
     /// </summary>
-    public IServiceProvider Services { get; }
+    public IServiceProvider Services => _services;
 
     /// <summary>
     /// The type that query operations will be rooted at.
@@ -488,6 +491,24 @@ public DocumentNode ToSyntaxNode()
     ISyntaxNode ISyntaxNodeProvider.ToSyntaxNode()
         => SchemaFormatter.FormatAsDocument(this);
 
+    public async ValueTask DisposeAsync()
+    {
+        if (!_disposed)
+        {
+            if (Features.TryGet(out SchemaCancellationFeature? cancellation))
+            {
+                await cancellation.DisposeAsync().ConfigureAwait(false);
+            }
+
+            if (_services is IAsyncDisposable disposableServices)
+            {
+                await disposableServices.DisposeAsync().ConfigureAwait(false);
+            }
+
+            _disposed = true;
+        }
+    }
+
     private readonly record struct TransitionKey(string TypeName, string From, string To);
 }