Add a warmup API for warmup triggers

Author: ankitkumarr

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

@@ -287,6 +287,17 @@ public IActionResult GetAdminToken()
             return Ok(requestHeaderToken);
         }
 
+        [HttpGet]
+        [HttpPost]
+        [Route("admin/warmup")]
+        [Authorize(Policy = PolicyNames.AdminAuthLevelOrInternal)]
+        [RequiresRunningHost]
+        public async Task<IActionResult> Warmup([FromServices] IScriptJobHost scriptHost)
+        {
+            await scriptHost.TryInvokeWarmupAsync();
+            return Ok();
+        }
+
         [AcceptVerbs("GET", "POST", "DELETE")]
         [Authorize(AuthenticationSchemes = AuthLevelAuthenticationDefaults.AuthenticationScheme)]
         [Route("runtime/webhooks/{name}/{*extra}")]

src/WebJobs.Script/Host/ScriptJobHostExtensions.cs

@@ -4,20 +4,70 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading.Tasks;
 using Microsoft.Azure.WebJobs.Script.Description;
 
 namespace Microsoft.Azure.WebJobs.Script
 {
     public static class ScriptJobHostExtensions
     {
+        private const string WarmupFunctionName = "Warmup";
+        private const string WarmupTriggerName = "WarmupTrigger";
+
         /// <summary>
         /// Lookup a function by name
         /// </summary>
         /// <param name="name">name of function</param>
         /// <returns>function or null if not found</returns>
         public static FunctionDescriptor GetFunctionOrNull(this IScriptJobHost scriptJobHost, string name)
         {
-            return scriptJobHost.Functions.FirstOrDefault(p => string.Equals(p.Name, name, StringComparison.OrdinalIgnoreCase));
+            return scriptJobHost.Functions.FirstOrDefault(f => IsFunctionNameMatch(f.Name, name));
+        }
+
+        private static bool IsFunctionNameMatch(string functionName, string comparison)
+        {
+            return string.Equals(functionName, comparison, StringComparison.OrdinalIgnoreCase);
+        }
+
+        /// <summary>
+        /// Lookup a warmup function
+        /// </summary>
+        /// <returns>Warmup function or null if not found</returns>
+        public static FunctionDescriptor GetWarmupFunctionOrNull(this IScriptJobHost scriptJobHost)
+        {
+            return scriptJobHost.Functions.FirstOrDefault(f =>
+            {
+                return IsFunctionNameMatch(f.Name, WarmupFunctionName)
+                && f.Metadata
+                    .InputBindings
+                    .Any(b => b.IsTrigger && b.Type.Equals(WarmupTriggerName, StringComparison.OrdinalIgnoreCase));
+            });
+        }
+
+        /// <summary>
+        /// Try to invoke a warmup function if available
+        /// </summary>
+        /// <returns>
+        /// A task that represents the asynchronous operation.
+        /// The task results true if a warmup function was invoked, false otherwise.
+        /// </returns>
+        public static async Task<bool> TryInvokeWarmupAsync(this IScriptJobHost scriptJobHost)
+        {
+            var warmupFunction = scriptJobHost.GetWarmupFunctionOrNull();
+            if (warmupFunction != null)
+            {
+                ParameterDescriptor inputParameter = warmupFunction.Parameters.First(p => p.IsTrigger);
+
+                var arguments = new Dictionary<string, object>()
+                {
+                    { inputParameter.Name, new WarmupContext() }
+                };
+
+                await scriptJobHost.CallAsync(warmupFunction.Name, arguments);
+                return true;
+            }
+
+            return false;
         }
     }
 }

src/WebJobs.Script/ScriptHostBuilderExtensions.cs

@@ -102,7 +102,8 @@ public static IHostBuilder AddScriptHostCore(this IHostBuilder builder, ScriptAp
                     o.SetResponse = HttpBinding.SetResponse;
                 })
                 .AddTimers()
-                .AddManualTrigger();
+                .AddManualTrigger()
+                .AddWarmup();
 
                 var extensionBundleOptions = GetExtensionBundleOptions(context);
                 var bundleManager = new ExtensionBundleManager(extensionBundleOptions, SystemEnvironment.Instance, loggerFactory);

src/WebJobs.Script/WebJobs.Script.csproj

@@ -42,7 +42,7 @@
     <PackageReference Include="Microsoft.Azure.Functions.NodeJsWorker" Version="1.1.1" />
     <PackageReference Include="Microsoft.Azure.Functions.PowerShellWorker" Version="1.0.188" />
     <PackageReference Include="Microsoft.Azure.WebJobs" Version="3.0.14-11660" />
-    <PackageReference Include="Microsoft.Azure.WebJobs.Extensions" Version="3.0.3" />
+    <PackageReference Include="Microsoft.Azure.WebJobs.Extensions" Version="3.0.4" />
     <PackageReference Include="Microsoft.Azure.WebJobs.Extensions.Http" Version="3.0.4" />
     <PackageReference Include="Microsoft.Azure.WebJobs.Logging.ApplicationInsights" Version="3.0.14-11656" />
     <PackageReference Include="Microsoft.Build" Version="15.8.166" />

test/WebJobs.Script.Tests.Integration/ApplicationInsights/ApplicationInsightsEndToEndTestsBase.cs

@@ -274,7 +274,7 @@ await TestHelpers.Await(() =>
                 !t.Message.StartsWith("Host Status")
             ).ToArray();
 
-            int expectedCount = 12;
+            int expectedCount = 13;
             Assert.True(traces.Length == expectedCount, $"Expected {expectedCount} messages, but found {traces.Length}. Actual logs:{Environment.NewLine}{string.Join(Environment.NewLine, traces.Select(t => t.Message))}");
 
             int idx = 0;
@@ -287,6 +287,7 @@ await TestHelpers.Await(() =>
             ValidateTrace(traces[idx++], "Host lock lease acquired by instance ID", ScriptConstants.LogCategoryHostGeneral);
             ValidateTrace(traces[idx++], "Host started (", LogCategories.Startup);
             ValidateTrace(traces[idx++], "Initializing Host", LogCategories.Startup);
+            ValidateTrace(traces[idx++], "Initializing Warmup Extension", LogCategories.CreateTriggerCategory("Warmup"));
             ValidateTrace(traces[idx++], "Job host started", LogCategories.Startup);
             ValidateTrace(traces[idx++], "Loading functions metadata", LogCategories.Startup);
             ValidateTrace(traces[idx++], "Starting Host (HostId=", LogCategories.Startup);

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

@@ -1,13 +1,18 @@
 // Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
+using System.Collections;
+using System.Collections.Generic;
+using System.Collections.ObjectModel;
 using System.IO;
 using System.Net;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Authentication;
 using Microsoft.AspNetCore.Authorization;
 using Microsoft.AspNetCore.Mvc;
 using Microsoft.Azure.WebJobs.Host.Scale;
+using Microsoft.Azure.WebJobs.Script.Description;
 using Microsoft.Azure.WebJobs.Script.ExtensionBundle;
 using Microsoft.Azure.WebJobs.Script.Scale;
 using Microsoft.Azure.WebJobs.Script.WebHost.Controllers;
@@ -166,5 +171,111 @@ public async Task GetScaleStatus_RuntimeScaleModeNotEnabled_ReturnsBadRequest()
             Assert.Equal(HttpStatusCode.BadRequest, (HttpStatusCode)result.StatusCode);
             Assert.Equal("Runtime scale monitoring is not enabled.", result.Value);
         }
+
+        [Theory]
+        [ClassData(typeof(WarmupTestData))]
+        public async Task TestWarmupEndpoint_Success(FunctionDescriptor[] functions, bool warmupCalled)
+        {
+            var triggerParamName = "triggerParam";
+            var scriptHostMock = new Mock<IScriptJobHost>();
+            bool functionInvoked = false;
+
+            scriptHostMock.Setup(p => p.CallAsync(It.IsAny<string>(), It.IsAny<IDictionary<string, object>>(), CancellationToken.None))
+            .Callback<string, IDictionary<string, object>, CancellationToken>((name, args, token) =>
+            {
+                Assert.Equal("warmup", name);
+                Assert.Equal(1, args.Count);
+                Assert.IsType<WarmupContext>(args[triggerParamName]);
+
+                functionInvoked = true;
+            })
+            .Returns(Task.CompletedTask);
+            scriptHostMock.SetupGet(p => p.Functions).Returns(functions);
+
+            IActionResult response = await _hostController.Warmup(scriptHostMock.Object);
+
+            Assert.Equal(warmupCalled, functionInvoked);
+            Assert.IsType<OkResult>(response);
+        }
+
+        public class WarmupTestData : IEnumerable<object[]>
+        {
+            private readonly BindingMetadata _blobInputBinding;
+            private readonly BindingMetadata _blobOutputBinding;
+            private readonly BindingMetadata _blobTriggerBinding;
+            private readonly BindingMetadata _warmupTriggerBinding;
+            private readonly BindingMetadata _manualTriggerBinding;
+
+            private readonly ParameterDescriptor _triggerParam;
+            private readonly ParameterDescriptor _nonTriggerParam;
+
+            private readonly FunctionDescriptor _warmupFunctionErrName;
+            private readonly FunctionDescriptor _warmupFunctionWarmupName;
+            private readonly FunctionDescriptor _manualFunctionWarmupName;
+            private readonly FunctionDescriptor _blobFunctionBlobName;
+
+            public WarmupTestData()
+            {
+                _triggerParam = new ParameterDescriptor("triggerParam", null)
+                {
+                    IsTrigger = true
+                };
+
+                _nonTriggerParam = new ParameterDescriptor("nonTriggerParam", null)
+                {
+                    IsTrigger = false
+                };
+
+                _blobInputBinding = GetBindingMetadata("boringBlob", "blob", BindingDirection.In);
+                _blobOutputBinding = GetBindingMetadata("bigBlob", "blob", BindingDirection.Out);
+                _blobTriggerBinding = GetBindingMetadata("beautifulBlob", "blobTrigger", BindingDirection.In);
+                _warmupTriggerBinding = GetBindingMetadata("superState", "warmupTrigger", BindingDirection.In);
+                _manualTriggerBinding = GetBindingMetadata("majesticManual", "manualTrigger", BindingDirection.In);
+
+                var warmupMetadata = new Script.Description.FunctionMetadata();
+                warmupMetadata.Bindings.Add(_warmupTriggerBinding);
+                warmupMetadata.Bindings.Add(_blobInputBinding);
+                _warmupFunctionWarmupName = new FunctionDescriptor("warmup", null, warmupMetadata,
+                    new Collection<ParameterDescriptor>() { _nonTriggerParam, _triggerParam }, null, null, null);
+
+                _warmupFunctionErrName = new FunctionDescriptor("donotwarmup", null, warmupMetadata,
+                    new Collection<ParameterDescriptor>() { _nonTriggerParam, _triggerParam }, null, null, null);
+
+                var manualMetadata = new Script.Description.FunctionMetadata();
+                manualMetadata.Bindings.Add(_manualTriggerBinding);
+                manualMetadata.Bindings.Add(_blobInputBinding);
+                _manualFunctionWarmupName = new FunctionDescriptor("warmup", null, manualMetadata,
+                    new Collection<ParameterDescriptor>() { _nonTriggerParam, _triggerParam }, null, null, null);
+
+                var blobMetadata = new Script.Description.FunctionMetadata();
+                blobMetadata.Bindings.Add(_blobTriggerBinding);
+                blobMetadata.Bindings.Add(_blobOutputBinding);
+                _blobFunctionBlobName = new FunctionDescriptor("blobFunction", null, blobMetadata,
+                    new Collection<ParameterDescriptor>() { _nonTriggerParam, _triggerParam }, null, null, null);
+            }
+
+            private BindingMetadata GetBindingMetadata(string name, string type, BindingDirection dir)
+            {
+                return new BindingMetadata()
+                {
+                    Name = name,
+                    Type = type,
+                    Direction = dir
+                };
+            }
+
+            public IEnumerator<object[]> GetEnumerator()
+            {
+                return new List<object[]>
+                {
+                    new object[] { new[] { _warmupFunctionWarmupName, _manualFunctionWarmupName },  true },
+                    new object[] { new[] { _blobFunctionBlobName, _manualFunctionWarmupName }, false },
+                    new object[] { new[] { _warmupFunctionErrName, _manualFunctionWarmupName }, false },
+                    new object[] { new[] { _warmupFunctionWarmupName, _blobFunctionBlobName }, true }
+                }.GetEnumerator();
+            }
+
+            IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
+        }
     }
 }