[HealthChecks] Add health check middleware (#11173)

* Add health check middleware

* Fix test nullref

* Update deps.json and runtimeassemblies.json

* Update ExistingRuntimeAssemblies.txt

* Update deps.json

* Extract JsonSerializerOptions to reusable static

* Mark sealed

* Address PR comments

* Update deps.json

* Add max wait time, respond to RequestAborted

* Remove unused assembly

Author: jviau

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

@@ -69,7 +69,7 @@ public async Task<IActionResult> Get()
                 // Extensions that are webhook providers create their default system keys
                 // as part of host initialization (when those keys aren't already present).
                 // So we must delay key retrieval until host initialization is complete.
-                await _hostManager.DelayUntilHostReady();
+                await _hostManager.DelayUntilHostReadyAsync();
             }
 
             Dictionary<string, string> keys = await GetHostSecretsByScope(hostKeyScope);

src/WebJobs.Script.WebHost/Diagnostics/HealthChecks/HealthCheckResponseWriter.cs

@@ -0,0 +1,43 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+using System.Text.Json;
+using System.Threading.Tasks;
+using HealthChecks.UI.Client;
+using Microsoft.AspNetCore.Http;
+using Microsoft.Extensions.Diagnostics.HealthChecks;
+using Microsoft.Extensions.Primitives;
+
+namespace Microsoft.Azure.WebJobs.Script.WebHost.Diagnostics.HealthChecks
+{
+    public class HealthCheckResponseWriter
+    {
+        public static Task WriteResponseAsync(HttpContext httpContext, HealthReport report)
+        {
+            ArgumentNullException.ThrowIfNull(httpContext);
+            ArgumentNullException.ThrowIfNull(report);
+
+            // We will write a detailed report if ?expand=true is present.
+            if (httpContext.Request.Query.TryGetValue("expand", out StringValues value)
+                && bool.TryParse(value, out bool expand) && expand)
+            {
+                return UIResponseWriter.WriteHealthCheckUIResponse(httpContext, report);
+            }
+
+            return WriteMinimalResponseAsync(httpContext, report);
+        }
+
+        private static Task WriteMinimalResponseAsync(HttpContext httpContext, HealthReport report)
+        {
+            MinimalResponse body = new(report.Status);
+            return JsonSerializer.SerializeAsync(
+                httpContext.Response.Body, body, JsonSerializerOptionsProvider.Options, httpContext.RequestAborted);
+        }
+
+        internal readonly struct MinimalResponse(HealthStatus status)
+        {
+            public HealthStatus Status { get; } = status;
+        }
+    }
+}

src/WebJobs.Script.WebHost/Diagnostics/HealthChecks/HealthCheckWaitMiddleware.cs

@@ -0,0 +1,42 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+using System.Threading.Tasks;
+using Microsoft.AspNetCore.Http;
+using Microsoft.Azure.WebJobs.Script.WebHost.Models;
+using Microsoft.Extensions.Primitives;
+
+namespace Microsoft.Azure.WebJobs.Script.WebHost.Diagnostics.HealthChecks
+{
+    public sealed class HealthCheckWaitMiddleware(RequestDelegate next, IScriptHostManager manager)
+    {
+        private const int MaxWaitSeconds = 60;
+        private readonly RequestDelegate _next = next ?? throw new ArgumentNullException(nameof(next));
+        private readonly IScriptHostManager _manager = manager ?? throw new ArgumentNullException(nameof(manager));
+
+        public async Task InvokeAsync(HttpContext context)
+        {
+            ArgumentNullException.ThrowIfNull(context);
+
+            // If specified, the ?wait={seconds} query param will wait for an
+            // active script host for that duration. This is to avoid excessive polling
+            // when waiting for the initial readiness probe.
+            if (context.Request.Query.TryGetValue("wait", out StringValues wait))
+            {
+                if (!int.TryParse(wait.ToString(), out int waitSeconds) || waitSeconds < 0)
+                {
+                    context.Response.StatusCode = StatusCodes.Status400BadRequest;
+                    await context.Response.WriteAsJsonAsync(
+                        ErrorResponse.BadArgument("'wait' query param must be a positive integer", $"wait={wait}"));
+                    return;
+                }
+
+                waitSeconds = Math.Min(waitSeconds, MaxWaitSeconds);
+                await _manager.DelayUntilHostReadyAsync(waitSeconds).WaitAsync(context.RequestAborted);
+            }
+
+            await _next(context);
+        }
+    }
+}

src/WebJobs.Script.WebHost/Extensions/HttpContextExtensions.cs

@@ -23,7 +23,7 @@ public static async Task WaitForRunningHostAsync(this HttpContext httpContext, I
                 // If the host is not ready, we'll wait a bit for it to initialize.
                 // This might happen if http requests come in while the host is starting
                 // up for the first time, or if it is restarting.
-                bool hostReady = await hostManager.DelayUntilHostReady(timeoutSeconds, pollingIntervalMilliseconds);
+                bool hostReady = await hostManager.DelayUntilHostReadyAsync(timeoutSeconds, pollingIntervalMilliseconds);
 
                 if (!hostReady)
                 {

src/WebJobs.Script.WebHost/Extensions/ScriptHostManagerExtensions.cs

@@ -7,7 +7,7 @@ namespace Microsoft.Azure.WebJobs.Script
 {
     public static class ScriptHostManagerExtensions
     {
-        public static async Task<bool> DelayUntilHostReady(this IScriptHostManager hostManager, int timeoutSeconds = ScriptConstants.HostTimeoutSeconds, int pollingIntervalMilliseconds = ScriptConstants.HostPollingIntervalMilliseconds)
+        public static async Task<bool> DelayUntilHostReadyAsync(this IScriptHostManager hostManager, int timeoutSeconds = ScriptConstants.HostTimeoutSeconds, int pollingIntervalMilliseconds = ScriptConstants.HostPollingIntervalMilliseconds)
         {
             if (HostIsInitialized(hostManager))
             {

src/WebJobs.Script.WebHost/Middleware/HostAvailabilityCheckMiddleware.cs

@@ -54,7 +54,7 @@ private static async Task InvokeAwaitingHost(HttpContext context, RequestDelegat
             {
                 Logger.InitiatingHostAvailabilityCheck(logger);
 
-                bool hostReady = await scriptHostManager.DelayUntilHostReady();
+                bool hostReady = await scriptHostManager.DelayUntilHostReadyAsync();
                 if (!hostReady)
                 {
                     Logger.HostUnavailableAfterCheck(logger);

src/WebJobs.Script.WebHost/Middleware/HostWarmupMiddleware.cs

@@ -161,7 +161,7 @@ public async Task HostWarmupAsync(HttpRequest request)
                 await _hostManager.RestartHostAsync("Host warmup call requested a restart.", CancellationToken.None);
 
                 // This call is here for sanity, but we should be fully initialized.
-                await _hostManager.DelayUntilHostReady();
+                await _hostManager.DelayUntilHostReadyAsync();
             }
         }
 

src/WebJobs.Script.WebHost/Models/ApiErrorModel.cs

@@ -0,0 +1,63 @@
+// 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.Generic;
+using System.Text.Json.Serialization;
+using Newtonsoft.Json;
+
+namespace Microsoft.Azure.WebJobs.Script.WebHost.Models
+{
+    /// <summary>
+    /// Represents an error response.
+    /// See https://github.com/Azure/azure-resource-manager-rpc/blob/master/v1.0/common-api-details.md#error-response-content.
+    /// </summary>
+    /// <param name="Code">
+    /// The error code. This is NOT the HTTP status code.
+    /// Unlocalized string which can be used to programmatically identify the error.
+    /// The code should be Pascal-cased, and should serve to uniquely identify a particular class of error,
+    /// for example "BadArgument".
+    /// </param>
+    /// <param name="Message">
+    /// The error message. Describes the error in detail and provides debugging information.
+    /// If Accept-Language is set in the request, it must be localized to that language.
+    /// </param>]
+    public record ErrorResponse(
+        [property: JsonProperty("code")][property: JsonPropertyName("code")] string Code,
+        [property: JsonProperty("message")][property: JsonPropertyName("message")] string Message)
+    {
+        /// <summary>
+        /// Gets the target of the particular error. For example, the name of the property in error.
+        /// </summary>
+        [JsonProperty("target")]
+        [JsonPropertyName("target")]
+        public string Target { get; init; }
+
+        /// <summary>
+        /// Gets the details of this error.
+        /// </summary>
+        [JsonProperty("details")]
+        [JsonPropertyName("details")]
+        public IEnumerable<ErrorResponse> Details { get; init; } = [];
+
+        /// <summary>
+        /// Gets the additional information for this error.
+        /// </summary>
+        [JsonProperty("additionalInfo")]
+        [JsonPropertyName("additionalInfo")]
+        public IEnumerable<ErrorAdditionalInfo> AdditionalInfo { get; init; } = [];
+
+        public static ErrorResponse BadArgument(string message, string target = null)
+        {
+            return new("BadArgument", message) { Target = target };
+        }
+    }
+
+    /// <summary>
+    /// Represents additional information for an error.
+    /// </summary>
+    /// <param name="Type">The type of additional information.</param>
+    /// <param name="Info">The additional error information.</param>
+    public record ErrorAdditionalInfo(
+        [property: JsonProperty("type")][property: JsonPropertyName("type")] string Type,
+        [property: JsonProperty("info")][property: JsonPropertyName("info")] object Info);
+}

src/WebJobs.Script.WebHost/Models/ErrorResponse.cs

@@ -124,7 +124,7 @@ public async Task SpecializeHostCoreAsync()
 
             using (_metricsLogger.LatencyEvent(MetricEventNames.SpecializationDelayUntilHostReady))
             {
-                await _scriptHostManager.DelayUntilHostReady();
+                await _scriptHostManager.DelayUntilHostReadyAsync();
             }
         }