Ensure extension RPC endpoints ready before processing gRPC messages (#10255)

* Ensure extension RPC endpoints ready before processing gRPC messages

* Add timeout and tests to waiting on RPC extensions.

Author: jviau

release_notes.md

@@ -17,4 +17,5 @@
 - Ordered invocations are now the default (#10201)
 - Skip worker description if none of the profile conditions are met (#9932)
 - Fixed incorrect function count in the log message.(#10220)
+- Fix race condition on startup with extension RPC endpoints not being available. (#10255)
 - Adding a timeout when retrieving function metadata from metadata providers (#10219)

src/WebJobs.Script.Grpc/Server/ExtensionsCompositeEndpointDataSource.cs

@@ -6,10 +6,12 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Linq;
 using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Routing;
 using Microsoft.Azure.WebJobs.Rpc.Core.Internal;
 using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
 using Microsoft.Extensions.Primitives;
 
 namespace Microsoft.Azure.WebJobs.Script.Grpc
@@ -26,6 +28,7 @@ internal sealed class ExtensionsCompositeEndpointDataSource : EndpointDataSource
         private readonly object _lock = new();
         private readonly List<EndpointDataSource> _dataSources = new();
         private readonly IScriptHostManager _scriptHostManager;
+        private readonly TaskCompletionSource _initialized = new();
 
         private IServiceProvider _extensionServices;
         private List<Endpoint> _endpoints;
@@ -191,6 +194,7 @@ private void OnHostChanged(object sender, ActiveHostChangedEventArgs args)
                         .GetService<IEnumerable<WebJobsRpcEndpointDataSource>>()
                         ?? Enumerable.Empty<WebJobsRpcEndpointDataSource>();
                     _dataSources.AddRange(sources);
+                    _initialized.TrySetResult(); // signal we have first initialized.
                 }
                 else
                 {
@@ -301,5 +305,49 @@ private void ThrowIfDisposed()
                 throw new ObjectDisposedException(nameof(ExtensionsCompositeEndpointDataSource));
             }
         }
+
+        /// <summary>
+        /// Middleware to ensure <see cref="ExtensionsCompositeEndpointDataSource"/> is initialized before routing for the first time.
+        /// Must be registered as a singleton service.
+        /// </summary>
+        /// <param name="dataSource">The <see cref="ExtensionsCompositeEndpointDataSource"/> to ensure is initialized.</param>
+        /// <param name="logger">The logger.</param>
+        public sealed class EnsureInitializedMiddleware(ExtensionsCompositeEndpointDataSource dataSource, ILogger<EnsureInitializedMiddleware> logger) : IMiddleware
+        {
+            private TaskCompletionSource _initialized = new();
+            private bool _firstRun = true;
+
+            // used for testing to verify initialization success.
+            internal Task Initialized => _initialized.Task;
+
+            // settable only for testing purposes.
+            internal TimeSpan Timeout { get; init; } = TimeSpan.FromSeconds(2);
+
+            public Task InvokeAsync(HttpContext context, RequestDelegate next)
+            {
+                return _firstRun ? InvokeCoreAsync(context, next) : next(context);
+            }
+
+            private async Task InvokeCoreAsync(HttpContext context, RequestDelegate next)
+            {
+                try
+                {
+                    await dataSource._initialized.Task.WaitAsync(Timeout);
+                }
+                catch (TimeoutException ex)
+                {
+                    // In case of deadlock we don't want to block all gRPC requests.
+                    // Log an error and continue.
+                    logger.LogError(ex, "Error initializing extension endpoints.");
+                    _initialized.TrySetException(ex);
+                }
+
+                // Even in case of timeout we don't want to continually test for initialization on subsequent requests.
+                // That would be a serious performance degredation.
+                _firstRun = false;
+                _initialized.TrySetResult();
+                await next(context);
+            }
+        }
     }
 }

src/WebJobs.Script.Grpc/Server/Startup.cs

@@ -16,6 +16,7 @@ internal class Startup
         public void ConfigureServices(IServiceCollection services)
         {
             services.AddSingleton<ExtensionsCompositeEndpointDataSource>();
+            services.AddSingleton<ExtensionsCompositeEndpointDataSource.EnsureInitializedMiddleware>();
             services.AddGrpc(options =>
             {
                 options.MaxReceiveMessageSize = MaxMessageLengthBytes;
@@ -30,12 +31,16 @@ public void Configure(IApplicationBuilder app, IWebHostEnvironment env)
                 app.UseDeveloperExceptionPage();
             }
 
+            // This must occur before 'UseRouting'. This ensures extension endpoints are registered before the
+            // endpoints are collected by the routing middleware.
+            app.UseMiddleware<ExtensionsCompositeEndpointDataSource.EnsureInitializedMiddleware>();
             app.UseRouting();
 
             app.UseEndpoints(endpoints =>
             {
                 endpoints.MapGrpcService<FunctionRpc.FunctionRpcBase>();
-                endpoints.DataSources.Add(endpoints.ServiceProvider.GetRequiredService<ExtensionsCompositeEndpointDataSource>());
+                endpoints.DataSources.Add(
+                    endpoints.ServiceProvider.GetRequiredService<ExtensionsCompositeEndpointDataSource>());
             });
         }
     }

test/WebJobs.Script.Tests/Workers/Rpc/ExtensionsCompositeEndpointDataSourceTests.cs

@@ -3,12 +3,16 @@
 
 using System;
 using System.Collections.Generic;
+using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
 using Microsoft.Azure.WebJobs.Rpc.Core.Internal;
 using Microsoft.Azure.WebJobs.Script.Grpc;
 using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.FileProviders;
 using Microsoft.Extensions.Hosting;
+using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Logging.Abstractions;
 using Microsoft.Extensions.Primitives;
 using Moq;
 using Xunit;
@@ -17,6 +21,9 @@ namespace Microsoft.Azure.WebJobs.Script.Tests.Workers.Rpc
 {
     public class ExtensionsCompositeEndpointDataSourceTests
     {
+        private static readonly ILogger<ExtensionsCompositeEndpointDataSource.EnsureInitializedMiddleware> _logger
+            = NullLogger<ExtensionsCompositeEndpointDataSource.EnsureInitializedMiddleware>.Instance;
+
         [Fact]
         public void NoActiveHost_NoEndpoints()
         {
@@ -41,6 +48,7 @@ public void ActiveHostChanged_NullHost_NoEndpoints()
         public void ActiveHostChanged_NoExtensions_NoEndpoints()
         {
             Mock<IScriptHostManager> manager = new();
+
             ExtensionsCompositeEndpointDataSource dataSource = new(manager.Object);
 
             IChangeToken token = dataSource.GetChangeToken();
@@ -67,6 +75,45 @@ public void ActiveHostChanged_NewExtensions_NewEndpoints()
                 endpoint => Assert.Equal("Test2", endpoint.DisplayName));
         }
 
+        [Fact]
+        public async Task ActiveHostChanged_MiddlewareWaits_Success()
+        {
+            Mock<IScriptHostManager> manager = new();
+
+            ExtensionsCompositeEndpointDataSource dataSource = new(manager.Object);
+            ExtensionsCompositeEndpointDataSource.EnsureInitializedMiddleware middleware =
+                new(dataSource, _logger) { Timeout = Timeout.InfiniteTimeSpan };
+            TestDelegate next = new();
+
+            Task waiter = middleware.InvokeAsync(null, next.InvokeAsync);
+            Assert.False(waiter.IsCompleted); // should be blocked until we raise the event.
+
+            manager.Raise(x => x.ActiveHostChanged += null, new ActiveHostChangedEventArgs(null, GetHost()));
+            await waiter.WaitAsync(TimeSpan.FromSeconds(5));
+            await middleware.Initialized;
+            await next.Invoked;
+        }
+
+        [Fact]
+        public async Task NoActiveHostChanged_MiddlewareWaits_Timeout()
+        {
+            Mock<IScriptHostManager> manager = new();
+
+            ExtensionsCompositeEndpointDataSource dataSource = new(manager.Object);
+            ExtensionsCompositeEndpointDataSource.EnsureInitializedMiddleware middleware =
+                new(dataSource, _logger) { Timeout = TimeSpan.Zero };
+            TestDelegate next = new();
+
+            await middleware.InvokeAsync(null, next.InvokeAsync).WaitAsync(TimeSpan.FromSeconds(5)); // should not throw
+            await Assert.ThrowsAsync<TimeoutException>(() => middleware.Initialized);
+            await next.Invoked;
+
+            // invoke again to verify it processes the next request.
+            next = new();
+            await middleware.InvokeAsync(null, next.InvokeAsync);
+            await next.Invoked;
+        }
+
         [Fact]
         public void Dispose_GetThrows()
         {
@@ -105,5 +152,18 @@ public TestEndpoints(params Endpoint[] endpoints)
 
             public override IChangeToken GetChangeToken() => NullChangeToken.Singleton;
         }
+
+        private class TestDelegate
+        {
+            private readonly TaskCompletionSource _invoked = new();
+
+            public Task Invoked => _invoked.Task;
+
+            public Task InvokeAsync(HttpContext context)
+            {
+                _invoked.TrySetResult();
+                return Task.CompletedTask;
+            }
+        }
     }
 }