Cancel drain/resume calls on client disconnect (#10565)

* Cancel drain calls on client disconnect

* Fix test assertion

* Cancel on /admin/host/resume calls as well

Author: jviau

src/WebJobs.Script.WebHost/Controllers/HostController.cs

@@ -188,7 +188,7 @@ public async Task<IActionResult> GetWorkerProcesses([FromServices] IScriptHostMa
         [HttpPost]
         [Route("admin/host/drain")]
         [Authorize(Policy = PolicyNames.AdminAuthLevelOrInternal)]
-        public async Task<IActionResult> Drain([FromServices] IScriptHostManager scriptHostManager)
+        public async Task<IActionResult> Drain([FromServices] IScriptHostManager scriptHostManager, CancellationToken cancellation)
         {
             _logger.LogDebug("Received request to drain the host");
 
@@ -197,21 +197,30 @@ public async Task<IActionResult> Drain([FromServices] IScriptHostManager scriptH
                 return StatusCode(StatusCodes.Status503ServiceUnavailable);
             }
 
-            await _drainSemaphore.WaitAsync();
+            try
+            {
+                await _drainSemaphore.WaitAsync(cancellation);
 
-            // Stop call to some listeners gets stuck, not waiting for the stop call to complete
-            _ = drainModeManager.EnableDrainModeAsync(CancellationToken.None)
-                                .ContinueWith(
-                                    antecedent =>
-                                    {
-                                        if (antecedent.Status == TaskStatus.Faulted)
+                // Stop call to some listeners gets stuck, not waiting for the stop call to complete
+                // We do not pass in this APIs cancellation token as we do not want drain to be aborted once we have the semaphore.
+                _ = drainModeManager.EnableDrainModeAsync(CancellationToken.None)
+                                    .ContinueWith(
+                                        antecedent =>
                                         {
-                                            _logger.LogError(antecedent.Exception, "Something went wrong invoking drain mode");
-                                        }
-
-                                        _drainSemaphore.Release();
-                                    });
-            return Accepted();
+                                            if (antecedent.Status == TaskStatus.Faulted)
+                                            {
+                                                _logger.LogError(antecedent.Exception, "Something went wrong invoking drain mode");
+                                            }
+
+                                            _drainSemaphore.Release();
+                                        });
+                return Accepted();
+            }
+            catch (OperationCanceledException)
+            {
+                _logger.LogInformation("Drain request was cancelled");
+                throw;
+            }
         }
 
         [HttpGet]
@@ -251,12 +260,20 @@ public IActionResult DrainStatus([FromServices] IScriptHostManager scriptHostMan
         [HttpPost]
         [Route("admin/host/resume")]
         [Authorize(Policy = PolicyNames.AdminAuthLevelOrInternal)]
-        public async Task<IActionResult> Resume([FromServices] IScriptHostManager scriptHostManager)
+        public async Task<IActionResult> Resume([FromServices] IScriptHostManager scriptHostManager, CancellationToken cancellation)
         {
             try
             {
-                await _resumeSemaphore.WaitAsync();
+                await _resumeSemaphore.WaitAsync(cancellation);
+            }
+            catch (OperationCanceledException)
+            {
+                _logger.LogInformation("Resume request was cancelled");
+                throw;
+            }
 
+            try
+            {
                 ScriptHostState currentState = scriptHostManager.State;
 
                 _logger.LogDebug($"Received request to resume a draining host - host status: {currentState.ToString()}");

test/WebJobs.Script.Tests/Controllers/Admin/HostControllerTests.cs

@@ -275,12 +275,41 @@ public void GetDrainStatus_HostRunning_ReturnsExpected(int outstandingRetries, i
             Assert.Equal(outstandingInvocations, resultStatus.OutstandingInvocations);
         }
 
+        [Fact]
+        public async Task GetDrain_HostRunning_BeginsDrain()
+        {
+            var scriptHostManagerMock = new Mock<IScriptHostManager>(MockBehavior.Strict);
+            var drainModeManager = new Mock<IDrainModeManager>(MockBehavior.Strict);
+            drainModeManager.Setup(m => m.EnableDrainModeAsync(default)).Returns(Task.CompletedTask);
+            var serviceProviderMock = scriptHostManagerMock.As<IServiceProvider>();
+            serviceProviderMock.Setup(x => x.GetService(typeof(IDrainModeManager))).Returns(drainModeManager.Object);
+
+            var result = (AcceptedResult)await _hostController.Drain(scriptHostManagerMock.Object, default);
+            Assert.Equal(StatusCodes.Status202Accepted, result.StatusCode);
+            drainModeManager.Verify(x => x.EnableDrainModeAsync(default), Times.Once());
+        }
+
+        [Fact]
+        public async Task GetDrain_HostRunning_ClientDisconnect_DoesNotDrain()
+        {
+            var scriptHostManagerMock = new Mock<IScriptHostManager>(MockBehavior.Strict);
+            var drainModeManager = new Mock<IDrainModeManager>(MockBehavior.Strict);
+            var serviceProviderMock = scriptHostManagerMock.As<IServiceProvider>();
+            serviceProviderMock.Setup(x => x.GetService(typeof(IDrainModeManager))).Returns(drainModeManager.Object);
+
+            using var cts = new CancellationTokenSource();
+            cts.Cancel();
+
+            await Assert.ThrowsAnyAsync<OperationCanceledException>(() => _hostController.Drain(scriptHostManagerMock.Object, cts.Token));
+            drainModeManager.Verify(x => x.EnableDrainModeAsync(default), Times.Never);
+        }
+
         [Fact]
         public async Task GetDrain_HostNotRunning_ReturnsServiceUnavailable()
         {
             var scriptHostManagerMock = new Mock<IScriptHostManager>(MockBehavior.Strict);
 
-            var result = (StatusCodeResult)await _hostController.Drain(scriptHostManagerMock.Object);
+            var result = (StatusCodeResult)await _hostController.Drain(scriptHostManagerMock.Object, default);
             Assert.Equal(StatusCodes.Status503ServiceUnavailable, result.StatusCode);
         }
 
@@ -318,7 +347,7 @@ public async Task ResumeHost_HostNotRunning_ReturnsExpected(ScriptHostState host
             scriptHostManagerMock.SetupGet(p => p.State).Returns(hostStatus);
             drainModeManager.Setup(x => x.IsDrainModeEnabled).Returns(drainModeEnabled);
 
-            var result = (StatusCodeResult)await _hostController.Resume(scriptHostManagerMock.Object);
+            var result = (StatusCodeResult)await _hostController.Resume(scriptHostManagerMock.Object, default);
             Assert.Equal(expectedCode, result.StatusCode);
             scriptHostManagerMock.Verify(p => p.RestartHostAsync(It.IsAny<CancellationToken>()), Times.Never());
         }
@@ -336,7 +365,7 @@ public async Task ResumeHost_HostRunning_DrainModeEnabled_StartNewHostSuccessful
             drainModeManager.Setup(x => x.IsDrainModeEnabled).Returns(true);
 
             var expectedBody = new ResumeStatus { State = ScriptHostState.Running };
-            var result = (OkObjectResult)await _hostController.Resume(scriptHostManagerMock.Object);
+            var result = (OkObjectResult)await _hostController.Resume(scriptHostManagerMock.Object, default);
 
             Assert.Equal(StatusCodes.Status200OK, result.StatusCode);
             Assert.Equal(expectedBody.State, (result.Value as ResumeStatus).State);
@@ -355,11 +384,21 @@ public async Task ResumeHost_HostRunning_DrainModeNotEnabled_DoesNotStartNewHost
             drainModeManager.Setup(x => x.IsDrainModeEnabled).Returns(false);
 
             var expectedBody = new ResumeStatus { State = ScriptHostState.Running };
-            var result = (OkObjectResult)await _hostController.Resume(scriptHostManagerMock.Object);
+            var result = (OkObjectResult)await _hostController.Resume(scriptHostManagerMock.Object, default);
 
             Assert.Equal(StatusCodes.Status200OK, result.StatusCode);
             Assert.Equal(expectedBody.State, (result.Value as ResumeStatus).State);
             scriptHostManagerMock.Verify(p => p.RestartHostAsync(It.IsAny<CancellationToken>()), Times.Never());
         }
+
+        [Fact]
+        public async Task ResumeHost_HostRunning_ClientDisconnect_DoesNotStartNewHost()
+        {
+            var scriptHostManagerMock = new Mock<IScriptHostManager>(MockBehavior.Strict);
+            using var cts = new CancellationTokenSource();
+            cts.Cancel();
+            await Assert.ThrowsAnyAsync<OperationCanceledException>(() => _hostController.Resume(scriptHostManagerMock.Object, cts.Token));
+            scriptHostManagerMock.Verify(p => p.RestartHostAsync(It.IsAny<CancellationToken>()), Times.Never());
+        }
     }
 }