Warmup new executor before replacing old one (#8068)

Author: tobias-tengler

src/HotChocolate/AspNetCore/src/AspNetCore/Extensions/HotChocolateAspNetCoreServiceCollectionExtensions.Warmup.cs

@@ -34,7 +34,7 @@ public static IRequestExecutorBuilder InitializeOnStartup(
             throw new ArgumentNullException(nameof(builder));
         }
 
-        builder.Services.AddHostedService<ExecutorWarmupService>();
+        builder.Services.AddHostedService<RequestExecutorWarmupService>();
         builder.Services.AddSingleton(new WarmupSchemaTask(builder.Name, keepWarm, warmup));
         return builder;
     }

src/HotChocolate/AspNetCore/src/AspNetCore/Warmup/ExecutorWarmupService.cs

@@ -0,0 +1,14 @@
+using Microsoft.Extensions.Hosting;
+
+namespace HotChocolate.AspNetCore.Warmup;
+
+internal sealed class RequestExecutorWarmupService(
+    IRequestExecutorWarmup executorWarmup)
+    : IHostedService
+{
+    public async Task StartAsync(CancellationToken cancellationToken)
+        => await executorWarmup.WarmupAsync(cancellationToken).ConfigureAwait(false);
+
+    public Task StopAsync(CancellationToken cancellationToken)
+        => Task.CompletedTask;
+}

src/HotChocolate/AspNetCore/src/AspNetCore/Warmup/RequestExecutorWarmupService.cs

@@ -1,26 +1,34 @@
 using HotChocolate.AspNetCore.Tests.Utilities;
+using HotChocolate.Execution;
+using Microsoft.Extensions.DependencyInjection;
 
 namespace HotChocolate.AspNetCore;
 
-public class EvictSchemaTests : ServerTestBase
+public class EvictSchemaTests(TestServerFactory serverFactory) : ServerTestBase(serverFactory)
 {
-    public EvictSchemaTests(TestServerFactory serverFactory)
-        : base(serverFactory)
-    {
-    }
-
     [Fact]
     public async Task Evict_Default_Schema()
     {
         // arrange
+        var newExecutorCreatedResetEvent = new AutoResetEvent(false);
         var server = CreateStarWarsServer();
 
         var time1 = await server.GetAsync(
             new ClientQueryRequest { Query = "{ time }", });
 
+        var resolver = server.Services.GetRequiredService<IRequestExecutorResolver>();
+        resolver.Events.Subscribe(new RequestExecutorEventObserver(@event =>
+        {
+            if (@event.Type == RequestExecutorEventType.Created)
+            {
+                newExecutorCreatedResetEvent.Set();
+            }
+        }));
+
         // act
         await server.GetAsync(
             new ClientQueryRequest { Query = "{ evict }", });
+        newExecutorCreatedResetEvent.WaitOne(5000);
 
         // assert
         var time2 = await server.GetAsync(
@@ -32,16 +40,27 @@ await server.GetAsync(
     public async Task Evict_Named_Schema()
     {
         // arrange
+        var newExecutorCreatedResetEvent = new AutoResetEvent(false);
         var server = CreateStarWarsServer();
 
         var time1 = await server.GetAsync(
             new ClientQueryRequest { Query = "{ time }", },
             "/evict");
 
+        var resolver = server.Services.GetRequiredService<IRequestExecutorResolver>();
+        resolver.Events.Subscribe(new RequestExecutorEventObserver(@event =>
+        {
+            if (@event.Type == RequestExecutorEventType.Created)
+            {
+                newExecutorCreatedResetEvent.Set();
+            }
+        }));
+
         // act
         await server.GetAsync(
             new ClientQueryRequest { Query = "{ evict }", },
             "/evict");
+        newExecutorCreatedResetEvent.WaitOne(5000);
 
         // assert
         var time2 = await server.GetAsync(

src/HotChocolate/AspNetCore/src/AspNetCore/Warmup/WarmupSchemaTask.cs

@@ -21,9 +21,7 @@ private AutoUpdateRequestExecutorProxy(
         _executorProxy = requestExecutorProxy;
         _executor = initialExecutor;
 
-        _executorProxy.ExecutorEvicted += (_, _) => BeginUpdateExecutor();
-
-        BeginUpdateExecutor();
+        _executorProxy.ExecutorUpdated += (_, args) => _executor = args.Executor;
     }
 
     /// <summary>
@@ -144,26 +142,6 @@ public Task<IResponseStream> ExecuteBatchAsync(
         CancellationToken cancellationToken = default)
         => _executor.ExecuteBatchAsync(requestBatch, cancellationToken);
 
-    private void BeginUpdateExecutor()
-        => UpdateExecutorAsync().FireAndForget();
-
-    private async ValueTask UpdateExecutorAsync()
-    {
-        await _semaphore.WaitAsync().ConfigureAwait(false);
-
-        try
-        {
-            var executor = await _executorProxy
-                .GetRequestExecutorAsync(CancellationToken.None)
-                .ConfigureAwait(false);
-            _executor = executor;
-        }
-        finally
-        {
-            _semaphore.Release();
-        }
-    }
-
     /// <inheritdoc cref="IDisposable" />
     public void Dispose()
     {

src/HotChocolate/AspNetCore/test/AspNetCore.Tests/EvictSchemaTests.cs

@@ -154,10 +154,8 @@ internal static IServiceCollection TryAddRequestExecutorResolver(
         this IServiceCollection services)
     {
         services.TryAddSingleton<RequestExecutorResolver>();
-        services.TryAddSingleton<IRequestExecutorResolver>(
-            sp => sp.GetRequiredService<RequestExecutorResolver>());
-        services.TryAddSingleton<IInternalRequestExecutorResolver>(
-            sp => sp.GetRequiredService<RequestExecutorResolver>());
+        services.TryAddSingleton<IRequestExecutorResolver>(sp => sp.GetRequiredService<RequestExecutorResolver>());
+        services.TryAddSingleton<IRequestExecutorWarmup>(sp => sp.GetRequiredService<RequestExecutorResolver>());
         return services;
     }
 

src/HotChocolate/Core/src/Execution/AutoUpdateRequestExecutorProxy.cs

@@ -12,6 +12,7 @@
 
   <ItemGroup>
     <InternalsVisibleTo Include="DynamicProxyGenAssembly2" />
+    <InternalsVisibleTo Include="HotChocolate.AspNetCore" />
     <InternalsVisibleTo Include="HotChocolate.AspNetCore.Tests" />
     <InternalsVisibleTo Include="HotChocolate.Caching" />
     <InternalsVisibleTo Include="HotChocolate.CostAnalysis" />

src/HotChocolate/Core/src/Execution/DependencyInjection/InternalServiceCollectionExtensions.cs

@@ -0,0 +1,18 @@
+namespace HotChocolate.Execution;
+
+/// <summary>
+/// Allows to run the initial warmup for registered <see cref="IRequestExecutor"/>s.
+/// </summary>
+internal interface IRequestExecutorWarmup
+{
+    /// <summary>
+    /// Runs the initial warmup tasks.
+    /// </summary>
+    /// <param name="cancellationToken">
+    /// The cancellation token.
+    /// </param>
+    /// <returns>
+    /// Returns a task that completes once the warmup is done.
+    /// </returns>
+    Task WarmupAsync(CancellationToken cancellationToken);
+}

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

@@ -29,7 +29,7 @@ public RequestExecutorProxy(IRequestExecutorResolver executorResolver, string sc
         _schemaName = schemaName;
         _eventSubscription =
             _executorResolver.Events.Subscribe(
-                new ExecutorObserver(EvictRequestExecutor));
+                new RequestExecutorEventObserver(OnRequestExecutorEvent));
     }
 
     public IRequestExecutor? CurrentExecutor => _executor;
@@ -178,22 +178,40 @@ public async ValueTask<IRequestExecutor> GetRequestExecutorAsync(
         return executor;
     }
 
-    private void EvictRequestExecutor(string schemaName)
+    private void OnRequestExecutorEvent(RequestExecutorEvent @event)
     {
-        if (!_disposed && schemaName.Equals(_schemaName))
+        if (_disposed || !@event.Name.Equals(_schemaName) || _executor is null)
+        {
+            return;
+        }
+
+        if (@event.Type is RequestExecutorEventType.Evicted)
         {
             _semaphore.Wait();
 
             try
             {
-                _executor = null;
                 ExecutorEvicted?.Invoke(this, EventArgs.Empty);
             }
             finally
             {
                 _semaphore.Release();
             }
         }
+        else if (@event.Type is RequestExecutorEventType.Created)
+        {
+            _semaphore.Wait();
+
+            try
+            {
+                _executor = @event.Executor;
+                ExecutorUpdated?.Invoke(this, new RequestExecutorUpdatedEventArgs(@event.Executor));
+            }
+            finally
+            {
+                _semaphore.Release();
+            }
+        }
     }
 
     public void Dispose()
@@ -206,19 +224,4 @@ public void Dispose()
             _disposed = true;
         }
     }
-
-    private sealed class ExecutorObserver(Action<string> evicted) : IObserver<RequestExecutorEvent>
-    {
-        public void OnNext(RequestExecutorEvent value)
-        {
-            if (value.Type is RequestExecutorEventType.Evicted)
-            {
-                evicted(value.Name);
-            }
-        }
-
-        public void OnError(Exception error) { }
-
-        public void OnCompleted() { }
-    }
 }