diff --git a/src/WebJobs.Script.WebHost/DependencyInjection/JobHostScopedServiceProvider.cs b/src/WebJobs.Script.WebHost/DependencyInjection/JobHostScopedServiceProvider.cs new file mode 100644 index 0000000000..b3309bc3f8 --- /dev/null +++ b/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 _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(); + } + } + } +} diff --git a/src/WebJobs.Script.WebHost/DependencyInjection/JobHostScopedServiceProviderFactory.cs b/src/WebJobs.Script.WebHost/DependencyInjection/JobHostScopedServiceProviderFactory.cs index 6dcda68307..595af1071e 100644 --- a/src/WebJobs.Script.WebHost/DependencyInjection/JobHostScopedServiceProviderFactory.cs +++ b/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) diff --git a/src/WebJobs.Script.WebHost/WebJobsScriptHostService.cs b/src/WebJobs.Script.WebHost/WebJobsScriptHostService.cs index c0977defee..7b0453aebc 100644 --- a/src/WebJobs.Script.WebHost/WebJobsScriptHostService.cs +++ b/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(); } } diff --git a/test/WebJobs.Script.Tests.Integration/WebHostEndToEnd/DrainModeResumeEndToEndTests.cs b/test/WebJobs.Script.Tests.Integration/WebHostEndToEnd/DrainModeResumeEndToEndTests.cs index b680db266c..e7fee45ca8 100644 --- a/test/WebJobs.Script.Tests.Integration/WebHostEndToEnd/DrainModeResumeEndToEndTests.cs +++ b/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(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 + { + // 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" } + }); } } }