Fixing /admin/host/resume ObjectDisposedException

src/WebJobs.Script.WebHost/DependencyInjection/JobHostScopedServiceProvider.cs

@@ -0,0 +1,71 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+using System.Collections.Concurrent;
+using System.Linq;
+using System.Threading.Tasks;
+using Microsoft.Extensions.DependencyInjection;
+
+namespace Microsoft.Azure.WebJobs.Script.WebHost.DependencyInjection
+{
+    internal class JobHostScopedServiceProvider : IServiceProvider, IServiceScopeFactory, IDisposable
+    {
+        private readonly IServiceProvider _root;
+        private readonly ConcurrentDictionary<JobHostServiceScope, object> _activeScopes = new();
+
+        public JobHostScopedServiceProvider(IServiceProvider root)
+        {
+            _root = root ?? throw new ArgumentNullException(nameof(root));
+        }
+
+        public IServiceScope CreateScope()
+        {
+            var scope = new JobHostServiceScope(_root.CreateScope());
+            _activeScopes.TryAdd(scope, null);
+            scope.DisposedTask.ContinueWith(t => { _activeScopes.TryRemove(scope, out _); });
+            return scope;
+        }
+
+        public void Dispose()
+        {
+            Task childScopeTasks = Task.WhenAll(_activeScopes.Keys.Select(s => s.DisposedTask));
+            Task.WhenAny(childScopeTasks, Task.Delay(5000))
+                .ContinueWith(t =>
+                {
+                    (_root as IDisposable)?.Dispose();
+                }, TaskContinuationOptions.ExecuteSynchronously);
+        }
+
+        public object GetService(Type serviceType)
+        {
+            if (serviceType == typeof(IServiceScopeFactory))
+            {
+                return this;
+            }
+
+            return _root.GetService(serviceType);
+        }
+
+        private class JobHostServiceScope : IServiceScope
+        {
+            private readonly TaskCompletionSource _disposalTaskSource = new();
+            private readonly IServiceScope _root;
+
+            public JobHostServiceScope(IServiceScope root)
+            {
+                _root = root;
+            }
+
+            public IServiceProvider ServiceProvider => _root.ServiceProvider;
+
+            public Task DisposedTask => _disposalTaskSource.Task;
+
+            public void Dispose()
+            {
+                _root.Dispose();
+                _disposalTaskSource.TrySetResult();
+            }
+        }
+    }
+}

src/WebJobs.Script.WebHost/DependencyInjection/JobHostScopedServiceProviderFactory.cs

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System;
@@ -73,7 +73,7 @@ public IServiceProvider CreateServiceProvider(IServiceCollection services)
                 jobHostServices.Add(service);
             }
 
-            return jobHostServices.BuildServiceProvider();
+            return new JobHostScopedServiceProvider(jobHostServices.BuildServiceProvider());
         }
 
         private static void ShimBreakingChanges(IServiceCollection services, ILogger logger)

src/WebJobs.Script.WebHost/WebJobsScriptHostService.cs

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System;
@@ -317,7 +317,7 @@ private async Task StartHostAsync(CancellationToken cancellationToken, int attem
                 }
                 finally
                 {
-                    activeOperation.Dispose();
+                    EndStartupOperation(activeOperation);
                     _hostStartSemaphore.Release();
                 }
             }

test/WebJobs.Script.Tests.Integration/WebHostEndToEnd/DrainModeResumeEndToEndTests.cs

@@ -1,9 +1,13 @@
 using System;
+using System.Collections.Generic;
 using System.IO;
+using System.Linq;
 using System.Net;
 using System.Threading.Tasks;
+using Microsoft.Azure.WebJobs.Script.Configuration;
 using Microsoft.Azure.WebJobs.Script.WebHost.Models;
 using Microsoft.Azure.WebJobs.Script.Workers.Rpc;
+using Microsoft.Extensions.Configuration;
 using Microsoft.WebJobs.Script.Tests;
 using Newtonsoft.Json;
 using Xunit;
@@ -43,9 +47,10 @@ public async Task DrainModeEnabled_RunningHost_StartsNewHost_ReturnsOk()
 
             // Validate host is "Running" after resume is called
             response = await SamplesTestHelpers.InvokeResume(this);
+            Assert.True(HttpStatusCode.OK == response.StatusCode,
+                string.Join(Environment.NewLine, Host.GetWebHostLogMessages().Where(m => m.Level == Microsoft.Extensions.Logging.LogLevel.Error).Select(m => m.ToString())));
             responseString = await response.Content.ReadAsStringAsync();
             var resumeStatus = JsonConvert.DeserializeObject<ResumeStatus>(responseString);
-
             Assert.Equal(ScriptHostState.Running, resumeStatus.State);
 
             // Validate the drain state is changed to "Disabled"
@@ -88,18 +93,21 @@ public async Task DrainModeDisabled_RunningHost_ReturnsOk()
 
     public class ResumeTestFixture : EndToEndTestFixture
     {
-        static ResumeTestFixture()
-        {
-        }
-
         public ResumeTestFixture()
             : base(Path.Combine(Environment.CurrentDirectory, "..", "..", "..", "..", "sample", "NodeResume"), "samples", RpcWorkerConstants.NodeLanguageWorkerName)
         {
         }
 
-        public override void ConfigureScriptHost(IWebJobsBuilder webJobsBuilder)
+        public override void ConfigureWebHost(IConfigurationBuilder configBuilder)
         {
-            base.ConfigureScriptHost(webJobsBuilder);
+            base.ConfigureWebHost(configBuilder);
+
+            configBuilder.AddInMemoryCollection(new Dictionary<string, string>
+            {
+                // This forces the hosts to be stopped and disposed before a new one starts.
+                // There was a bug hiding here originally, so we'll run all these tests this way.
+                { ConfigurationSectionNames.SequentialJobHostRestart, "true" }
+            });
         }
     }
 }