Do not unregister extension endpoints on shutdown (#10522)

* Do not unregister extension endpoints on shutdown

* Add comment about known RPC issue

* Fix unit test build

* Update unit test condition

* Remove NotEmpty check

Author: jviau

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

@@ -26,21 +26,25 @@ namespace Microsoft.Azure.WebJobs.Script.Grpc
     internal sealed class ExtensionsCompositeEndpointDataSource : EndpointDataSource, IDisposable
     {
         private readonly object _lock = new();
-        private readonly List<EndpointDataSource> _dataSources = new();
         private readonly IScriptHostManager _scriptHostManager;
         private readonly TaskCompletionSource _initialized = new();
+        private readonly ILogger _logger;
 
+        private IList<EndpointDataSource> _dataSources = Array.Empty<EndpointDataSource>();
         private IServiceProvider _extensionServices;
         private List<Endpoint> _endpoints;
         private IChangeToken _consumerChangeToken;
         private CancellationTokenSource _cts;
         private List<IDisposable> _changeTokenRegistrations;
         private bool _disposed;
 
-        public ExtensionsCompositeEndpointDataSource(IScriptHostManager scriptHostManager)
+        public ExtensionsCompositeEndpointDataSource(
+            IScriptHostManager scriptHostManager,
+            ILogger<ExtensionsCompositeEndpointDataSource> logger)
         {
             _scriptHostManager = scriptHostManager;
             _scriptHostManager.ActiveHostChanged += OnHostChanged;
+            _logger = logger;
         }
 
         /// <inheritdoc />
@@ -186,23 +190,21 @@ private void OnHostChanged(object sender, ActiveHostChangedEventArgs args)
         {
             lock (_lock)
             {
-                _dataSources.Clear();
+                // Do not clear out endpoints when host changes to null, as functions may still be running.
+                // TODO: there are still edge cases where we will switch active hosts, leading to potential
+                // issues with draining invocations.
                 if (args?.NewHost?.Services is { } services)
                 {
                     _extensionServices = services;
-                    IEnumerable<WebJobsRpcEndpointDataSource> sources = services
-                        .GetService<IEnumerable<WebJobsRpcEndpointDataSource>>()
-                        ?? Enumerable.Empty<WebJobsRpcEndpointDataSource>();
-                    _dataSources.AddRange(sources);
+                    IEnumerable<EndpointDataSource> sources = services.GetServices<WebJobsRpcEndpointDataSource>();
+                    _dataSources = sources.ToList();
+
+                    int totalEndpoints = _dataSources.Sum(x => x.Endpoints.Count);
+                    _logger.LogDebug("Host changed. Registering {ExtensionCount} extension data sources and {EndpointCount} total endpoints.", _dataSources.Count, totalEndpoints);
                     _initialized.TrySetResult(); // signal we have first initialized.
-                }
-                else
-                {
-                    _extensionServices = null;
+                    OnEndpointsChange(collectionChanged: true);
                 }
             }
-
-            OnEndpointsChange(collectionChanged: true);
         }
 
         private void OnEndpointsChange(bool collectionChanged)

src/WebJobs.Script/Host/IScriptHostManager.cs

@@ -15,7 +15,7 @@ public interface IScriptHostManager
         event EventHandler HostInitializing;
 
         /// <summary>
-        /// Evewnt raised when the active host managed by this instance changes.
+        /// Event raised when the active host managed by this instance changes.
         /// </summary>
         event EventHandler<ActiveHostChangedEventArgs> ActiveHostChanged;
 

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

@@ -21,35 +21,37 @@ namespace Microsoft.Azure.WebJobs.Script.Tests.Workers.Rpc
 {
     public class ExtensionsCompositeEndpointDataSourceTests
     {
-        private static readonly ILogger<ExtensionsCompositeEndpointDataSource.EnsureInitializedMiddleware> _logger
+        private static readonly ILogger<ExtensionsCompositeEndpointDataSource> _dataSourceLogger
+            = NullLogger<ExtensionsCompositeEndpointDataSource>.Instance;
+
+        private static readonly ILogger<ExtensionsCompositeEndpointDataSource.EnsureInitializedMiddleware> _middlewareLogger
             = NullLogger<ExtensionsCompositeEndpointDataSource.EnsureInitializedMiddleware>.Instance;
 
         [Fact]
         public void NoActiveHost_NoEndpoints()
         {
-            ExtensionsCompositeEndpointDataSource dataSource = new(Mock.Of<IScriptHostManager>());
+            ExtensionsCompositeEndpointDataSource dataSource = new(Mock.Of<IScriptHostManager>(), _dataSourceLogger);
             Assert.Empty(dataSource.Endpoints);
         }
 
         [Fact]
-        public void ActiveHostChanged_NullHost_NoEndpoints()
+        public void ActiveHostChanged_NullHost_EndpointsRemain()
         {
             Mock<IScriptHostManager> manager = new();
-            ExtensionsCompositeEndpointDataSource dataSource = new(manager.Object);
+            ExtensionsCompositeEndpointDataSource dataSource = new(manager.Object, _dataSourceLogger);
 
             IChangeToken token = dataSource.GetChangeToken();
             Assert.False(token.HasChanged);
             manager.Raise(x => x.ActiveHostChanged += null, new ActiveHostChangedEventArgs(null, null));
-            Assert.True(token.HasChanged);
-            Assert.Empty(dataSource.Endpoints);
+            Assert.False(token.HasChanged);
         }
 
         [Fact]
         public void ActiveHostChanged_NoExtensions_NoEndpoints()
         {
             Mock<IScriptHostManager> manager = new();
 
-            ExtensionsCompositeEndpointDataSource dataSource = new(manager.Object);
+            ExtensionsCompositeEndpointDataSource dataSource = new(manager.Object, _dataSourceLogger);
 
             IChangeToken token = dataSource.GetChangeToken();
             Assert.False(token.HasChanged);
@@ -62,7 +64,7 @@ public void ActiveHostChanged_NoExtensions_NoEndpoints()
         public void ActiveHostChanged_NewExtensions_NewEndpoints()
         {
             Mock<IScriptHostManager> manager = new();
-            ExtensionsCompositeEndpointDataSource dataSource = new(manager.Object);
+            ExtensionsCompositeEndpointDataSource dataSource = new(manager.Object, _dataSourceLogger);
             IHost host = GetHost(new TestEndpoints(new Endpoint(null, null, "Test1"), new Endpoint(null, null, "Test2")));
 
             IChangeToken token = dataSource.GetChangeToken();
@@ -80,9 +82,9 @@ public async Task ActiveHostChanged_MiddlewareWaits_Success()
         {
             Mock<IScriptHostManager> manager = new();
 
-            ExtensionsCompositeEndpointDataSource dataSource = new(manager.Object);
+            ExtensionsCompositeEndpointDataSource dataSource = new(manager.Object, _dataSourceLogger);
             ExtensionsCompositeEndpointDataSource.EnsureInitializedMiddleware middleware =
-                new(dataSource, _logger) { Timeout = Timeout.InfiniteTimeSpan };
+                new(dataSource, _middlewareLogger) { Timeout = Timeout.InfiniteTimeSpan };
             TestDelegate next = new();
 
             Task waiter = middleware.InvokeAsync(null, next.InvokeAsync);
@@ -99,9 +101,9 @@ public async Task NoActiveHostChanged_MiddlewareWaits_Timeout()
         {
             Mock<IScriptHostManager> manager = new();
 
-            ExtensionsCompositeEndpointDataSource dataSource = new(manager.Object);
+            ExtensionsCompositeEndpointDataSource dataSource = new(manager.Object, _dataSourceLogger);
             ExtensionsCompositeEndpointDataSource.EnsureInitializedMiddleware middleware =
-                new(dataSource, _logger) { Timeout = TimeSpan.Zero };
+                new(dataSource, _middlewareLogger) { Timeout = TimeSpan.Zero };
             TestDelegate next = new();
 
             await middleware.InvokeAsync(null, next.InvokeAsync).WaitAsync(TimeSpan.FromSeconds(5)); // should not throw
@@ -118,7 +120,7 @@ public async Task NoActiveHostChanged_MiddlewareWaits_Timeout()
         public void Dispose_GetThrows()
         {
             Mock<IScriptHostManager> manager = new();
-            ExtensionsCompositeEndpointDataSource dataSource = new(manager.Object);
+            ExtensionsCompositeEndpointDataSource dataSource = new(manager.Object, _dataSourceLogger);
             IHost host = GetHost(new TestEndpoints(new Endpoint(null, null, "Test1"), new Endpoint(null, null, "Test2")));
             manager.Raise(x => x.ActiveHostChanged += null, new ActiveHostChangedEventArgs(null, host));
             dataSource.Dispose();