Detect host shutdown using CancellationToken. Fixes Azure#4251

TayoA · Jun 5, 2019 · 18244b3 · 18244b3
1 parent 81003ef
commit 18244b3
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/src/WebJobs.Script.WebHost/Middleware/FunctionInvocationMiddleware.cs b/src/WebJobs.Script.WebHost/Middleware/FunctionInvocationMiddleware.cs
@@ -9,6 +9,7 @@
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Authorization;
 using Microsoft.AspNetCore.Authorization.Policy;
+using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Features;
 using Microsoft.AspNetCore.Mvc;
@@ -28,6 +29,7 @@ namespace Microsoft.Azure.WebJobs.Script.WebHost.Middleware
     public class FunctionInvocationMiddleware
     {
         private readonly RequestDelegate _next;
+        private IApplicationLifetime _applicationLifetime;
 
         public FunctionInvocationMiddleware(RequestDelegate next)
         {
@@ -50,6 +52,8 @@ public async Task Invoke(HttpContext context)
                 await _next(context);
             }
 
+            _applicationLifetime = context.RequestServices.GetService<IApplicationLifetime>();
+
             IFunctionExecutionFeature functionExecution = context.Features.Get<IFunctionExecutionFeature>();
             if (functionExecution != null && !context.Response.HasStarted())
             {
@@ -139,8 +143,8 @@ private async Task<IActionResult> GetResultAsync(HttpContext context, IFunctionE
 
                 using (logger.BeginScope(scopeState))
                 {
-                    // TODO: Flow cancellation token from caller
-                    await functionExecution.ExecuteAsync(context.Request, CancellationToken.None);
+                    CancellationToken cancellationToken = _applicationLifetime != null ? _applicationLifetime.ApplicationStopping : CancellationToken.None;
+                    await functionExecution.ExecuteAsync(context.Request, cancellationToken);
                 }
             }
 

diff --git a/test/WebJobs.Script.Tests.Integration/WebHostEndToEnd/SamplesEndToEndTests_CSharp.cs b/test/WebJobs.Script.Tests.Integration/WebHostEndToEnd/SamplesEndToEndTests_CSharp.cs
@@ -327,6 +327,9 @@ await TestHelpers.RunWithTimeoutAsync(async () =>
 
             await AwaitHostStateAsync(ScriptHostState.Running);
 
+            // need to reinitialize TestFunctionHost to reset IApplicationLifetime
+            await _fixture.InitializeAsync();
+
             // verify functions can be invoked
             await InvokeAndValidateHttpTrigger(functionName);