diff --git a/src/Microsoft.AspNetCore.Hosting/Internal/RequestServicesContainerMiddleware.cs b/src/Microsoft.AspNetCore.Hosting/Internal/RequestServicesContainerMiddleware.cs index 8f30c3195c..22b034ee34 100644 --- a/src/Microsoft.AspNetCore.Hosting/Internal/RequestServicesContainerMiddleware.cs +++ b/src/Microsoft.AspNetCore.Hosting/Internal/RequestServicesContainerMiddleware.cs @@ -30,34 +30,17 @@ namespace Microsoft.AspNetCore.Hosting.Internal _scopeFactory = scopeFactory; } - public async Task Invoke(HttpContext httpContext) + public Task Invoke(HttpContext httpContext) { Debug.Assert(httpContext != null); // local cache for virtual disptach result var features = httpContext.Features; - var existingFeature = features.Get(); - // All done if RequestServices is set - if (existingFeature?.RequestServices != null) - { - await _next.Invoke(httpContext); - return; - } + var servicesFeature = new RequestServicesFeature(httpContext, _scopeFactory); - var replacementFeature = new RequestServicesFeature(_scopeFactory); - - try - { - features.Set(replacementFeature); - await _next.Invoke(httpContext); - } - finally - { - replacementFeature.Dispose(); - // Restore previous feature state - features.Set(existingFeature); - } + features.Set(servicesFeature); + return _next.Invoke(httpContext); } } } \ No newline at end of file diff --git a/src/Microsoft.AspNetCore.Hosting/Internal/RequestServicesFeature.cs b/src/Microsoft.AspNetCore.Hosting/Internal/RequestServicesFeature.cs index 8d85cec63e..a57df9bcbc 100644 --- a/src/Microsoft.AspNetCore.Hosting/Internal/RequestServicesFeature.cs +++ b/src/Microsoft.AspNetCore.Hosting/Internal/RequestServicesFeature.cs @@ -3,6 +3,7 @@ using System; using System.Diagnostics; +using Microsoft.AspNetCore.Http; using Microsoft.AspNetCore.Http.Features; using Microsoft.Extensions.DependencyInjection; @@ -14,10 +15,12 @@ namespace Microsoft.AspNetCore.Hosting.Internal private IServiceProvider _requestServices; private IServiceScope _scope; private bool _requestServicesSet; + private HttpContext _context; - public RequestServicesFeature(IServiceScopeFactory scopeFactory) + public RequestServicesFeature(HttpContext context, IServiceScopeFactory scopeFactory) { Debug.Assert(scopeFactory != null); + _context = context; _scopeFactory = scopeFactory; } @@ -27,6 +30,7 @@ namespace Microsoft.AspNetCore.Hosting.Internal { if (!_requestServicesSet) { + _context.Response.RegisterForDispose(this); _scope = _scopeFactory.CreateScope(); _requestServices = _scope.ServiceProvider; _requestServicesSet = true;