diff --git a/src/WebJobs.Script.WebHost/Middleware/PlaceholderSpecializationMiddleware.cs b/src/WebJobs.Script.WebHost/Middleware/PlaceholderSpecializationMiddleware.cs index a77542e68e..607fd598d6 100644 --- a/src/WebJobs.Script.WebHost/Middleware/PlaceholderSpecializationMiddleware.cs +++ b/src/WebJobs.Script.WebHost/Middleware/PlaceholderSpecializationMiddleware.cs @@ -4,6 +4,7 @@ using System.Threading; using System.Threading.Tasks; using Microsoft.AspNetCore.Http; +using Microsoft.Azure.WebJobs.Script.Extensions; namespace Microsoft.Azure.WebJobs.Script.WebHost.Middleware { @@ -47,6 +48,13 @@ private async Task InvokeSpecializationCheck(HttpContext httpContext) if (Interlocked.CompareExchange(ref _specialized, 1, 0) == 0) { + // For Flex environments, set the cold start header during specialization + // since it's not automatically set by the platform like in Windows Consumption + if (_environment.IsFlexConsumptionSku()) + { + httpContext.Request.Headers[ScriptConstants.AntaresColdStartHeaderName] = "1"; + } + Interlocked.Exchange(ref _invoke, _next); } } diff --git a/test/WebJobs.Script.Tests/Middleware/PlaceholderSpecializationMiddlewareTests.cs b/test/WebJobs.Script.Tests/Middleware/PlaceholderSpecializationMiddlewareTests.cs new file mode 100644 index 0000000000..8165680290 --- /dev/null +++ b/test/WebJobs.Script.Tests/Middleware/PlaceholderSpecializationMiddlewareTests.cs @@ -0,0 +1,114 @@ +// Copyright (c) .NET Foundation. All rights reserved. +// Licensed under the MIT License. See License.txt in the project root for license information. + +using System.Threading.Tasks; +using Microsoft.AspNetCore.Http; +using Microsoft.Azure.WebJobs.Script.WebHost.Middleware; +using Microsoft.Extensions.Logging.Abstractions; +using Moq; +using Xunit; + +namespace Microsoft.Azure.WebJobs.Script.Tests.Middleware +{ + public class PlaceholderSpecializationMiddlewareTests + { + [Fact] + public async Task Invoke_FlexConsumption_AddsHeader_When_Specialized() + { + // Arrange + var webHostEnvironment = new Mock(); + var standbyManager = new Mock(); + var environment = new Mock(); + var httpContext = new DefaultHttpContext(); + + // Setup conditions for specialization + webHostEnvironment.Setup(e => e.InStandbyMode).Returns(false); + environment.Setup(e => e.IsContainerReady()).Returns(true); + environment.Setup(e => e.IsFlexConsumptionSku()).Returns(true); + + standbyManager.Setup(s => s.SpecializeHostAsync()).Returns(Task.CompletedTask); + + var nextCalled = false; + RequestDelegate next = (ctx) => + { + nextCalled = true; + return Task.CompletedTask; + }; + + var middleware = new PlaceholderSpecializationMiddleware(next, webHostEnvironment.Object, standbyManager.Object, environment.Object); + + // Act + await middleware.Invoke(httpContext); + + // Assert + Assert.True(nextCalled); + Assert.Equal("1", httpContext.Request.Headers[ScriptConstants.AntaresColdStartHeaderName]); + standbyManager.Verify(s => s.SpecializeHostAsync(), Times.Once); + } + + [Fact] + public async Task Invoke_NonFlexConsumption_DoesNotAddHeader_When_Specialized() + { + // Arrange + var webHostEnvironment = new Mock(); + var standbyManager = new Mock(); + var environment = new Mock(); + var httpContext = new DefaultHttpContext(); + + // Setup conditions for specialization but NOT Flex + webHostEnvironment.Setup(e => e.InStandbyMode).Returns(false); + environment.Setup(e => e.IsContainerReady()).Returns(true); + environment.Setup(e => e.IsFlexConsumptionSku()).Returns(false); + + standbyManager.Setup(s => s.SpecializeHostAsync()).Returns(Task.CompletedTask); + + var nextCalled = false; + RequestDelegate next = (ctx) => + { + nextCalled = true; + return Task.CompletedTask; + }; + + var middleware = new PlaceholderSpecializationMiddleware(next, webHostEnvironment.Object, standbyManager.Object, environment.Object); + + // Act + await middleware.Invoke(httpContext); + + // Assert + Assert.True(nextCalled); + Assert.False(httpContext.Request.Headers.ContainsKey(ScriptConstants.AntaresColdStartHeaderName)); + standbyManager.Verify(s => s.SpecializeHostAsync(), Times.Once); + } + + [Fact] + public async Task Invoke_NoSpecialization_DoesNotAddHeader() + { + // Arrange + var webHostEnvironment = new Mock(); + var standbyManager = new Mock(); + var environment = new Mock(); + var httpContext = new DefaultHttpContext(); + + // Setup conditions for NO specialization + webHostEnvironment.Setup(e => e.InStandbyMode).Returns(true); + environment.Setup(e => e.IsContainerReady()).Returns(true); + + var nextCalled = false; + RequestDelegate next = (ctx) => + { + nextCalled = true; + return Task.CompletedTask; + }; + + var middleware = new PlaceholderSpecializationMiddleware(next, webHostEnvironment.Object, standbyManager.Object, environment.Object); + + // Act + await middleware.Invoke(httpContext); + + // Assert + Assert.True(nextCalled); + Assert.False(httpContext.Request.Headers.ContainsKey(ScriptConstants.AntaresColdStartHeaderName)); + standbyManager.Verify(s => s.SpecializeHostAsync(), Times.Never); + } + } +} \ No newline at end of file