Add keyed ILogger which forwards to ScriptHost when possible

src/WebJobs.Script.WebHost/Diagnostics/ILoggingBuilderExtensions.cs

@@ -1,19 +1,23 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using Microsoft.Azure.WebJobs.Script.WebHost.Diagnostics;
 using Microsoft.Extensions.DependencyInjection;
 
+#nullable enable
+
 namespace Microsoft.Extensions.Logging
 {
     public static class ILoggingBuilderExtensions
     {
-        public static void AddWebJobsSystem<T>(this ILoggingBuilder builder) where T : SystemLoggerProvider
+        public static ILoggingBuilder AddWebJobsSystem<T>(this ILoggingBuilder builder)
+            where T : SystemLoggerProvider
         {
             builder.Services.AddSingleton<ILoggerProvider, T>();
 
             // Log all logs to SystemLogger
             builder.AddDefaultWebJobsFilters<T>(LogLevel.Trace);
+            return builder;
         }
     }
 }

src/WebJobs.Script.WebHost/Program.cs

@@ -1,4 +1,4 @@
-// Copyright (c) .NET Foundation. All rights reserved.
+// Copyright (c) .NET Foundation. All rights reserved.
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System;
@@ -88,6 +88,7 @@ public static IWebHostBuilder CreateWebHostBuilder(string[] args = null)
 
                     loggingBuilder.AddDefaultWebJobsFilters();
                     loggingBuilder.AddWebJobsSystem<WebHostSystemLoggerProvider>();
+                    loggingBuilder.AddForwardingLogger();
                     loggingBuilder.Services.AddSingleton<DeferredLoggerProvider>();
                     loggingBuilder.Services.AddSingleton<ILoggerProvider>(s => s.GetRequiredService<DeferredLoggerProvider>());
                     loggingBuilder.Services.AddSingleton<ISystemLoggerFactory, SystemLoggerFactory>();

src/WebJobs.Script/Diagnostics/ForwardingLogger.cs

@@ -0,0 +1,110 @@
+// 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.Diagnostics;
+using System.Diagnostics.CodeAnalysis;
+using Microsoft.Azure.WebJobs.Script;
+using Microsoft.Extensions.DependencyInjection;
+
+#nullable enable
+
+namespace Microsoft.Extensions.Logging
+{
+    internal class ForwardingLogger : ILogger
+    {
+        // The service key to use for dependency injection to get forwarding loggers.
+        public const string ServiceKey = "Forwarding";
+
+        private readonly string _categoryName;
+        private readonly ILogger _fallback;
+        private readonly IScriptHostManager _manager;
+
+        // We use weak references so as to not keep a ScriptHost alive after it shuts down.
+        private readonly WeakReference<ILogger> _current = new(null!);
+        private readonly WeakReference<IServiceProvider> _services = new(null!);
+
+        public ForwardingLogger(string categoryName, ILogger inner, IScriptHostManager manager)
+        {
+            ArgumentNullException.ThrowIfNull(inner);
+            ArgumentNullException.ThrowIfNull(manager);
+            _categoryName = categoryName;
+            _fallback = inner;
+            _manager = manager;
+        }
+
+        private ILogger Current
+        {
+            get
+            {
+                if (TryGetCurrentLogger(out ILogger? logger))
+                {
+                    return logger;
+                }
+
+                // No current ScriptHost logger, or the ScriptHost is gone. Use the fallback WebHost logger.
+                return _fallback;
+            }
+        }
+
+        public IDisposable? BeginScope<TState>(TState state)
+            where TState : notnull
+            => Current.BeginScope(state);
+
+        public bool IsEnabled(LogLevel logLevel) => Current.IsEnabled(logLevel);
+
+        public void Log<TState>(LogLevel logLevel, EventId eventId, TState state, Exception? exception, Func<TState, Exception?, string> formatter)
+            => Current.Log(logLevel, eventId, state, exception, formatter);
+
+        private bool TryGetCurrentLogger([NotNullWhen(true)] out ILogger? logger)
+        {
+            if (IsLoggerCurrent(out logger))
+            {
+                return true;
+            }
+            else if (_manager.Services is { } services)
+            {
+                logger = services.GetRequiredService<ILoggerFactory>().CreateLogger(_categoryName);
+                _services.SetTarget(services);
+                _current.SetTarget(logger);
+                return true;
+            }
+
+            logger = null;
+            return false;
+        }
+
+        private bool IsLoggerCurrent([NotNullWhen(true)] out ILogger? logger)
+        {
+            // First check if the last IServiceProvider we used is still active.
+            if (_services.TryGetTarget(out IServiceProvider? services)
+                && ReferenceEquals(services, _manager.Services))
+            {
+                // Service provider is still correct, so our logger is current.
+                return _current.TryGetTarget(out logger);
+            }
+
+            logger = null;
+            return false;
+        }
+    }
+
+    [DebuggerDisplay("{_logger}")]
+    internal class ForwardingLogger<T> : ILogger<T>
+    {
+        private readonly ILogger _logger;
+
+        public ForwardingLogger([ForwardingLogger] ILoggerFactory factory)
+        {
+            ArgumentNullException.ThrowIfNull(factory);
+            _logger = factory.CreateLogger<T>();
+        }
+
+        IDisposable? ILogger.BeginScope<TState>(TState state) => _logger.BeginScope(state);
+
+        bool ILogger.IsEnabled(LogLevel logLevel) => _logger.IsEnabled(logLevel);
+
+        void ILogger.Log<TState>(LogLevel logLevel, EventId eventId, TState state, Exception? exception, Func<TState, Exception?, string> formatter) =>
+            _logger.Log(logLevel, eventId, state, exception, formatter);
+    }
+}

src/WebJobs.Script/Diagnostics/ForwardingLoggerAttribute.cs

@@ -0,0 +1,14 @@
+// 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 Microsoft.Extensions.DependencyInjection;
+
+namespace Microsoft.Extensions.Logging
+{
+    [AttributeUsage(AttributeTargets.Parameter)]
+    internal class ForwardingLoggerAttribute()
+        : FromKeyedServicesAttribute(ForwardingLogger.ServiceKey)
+    {
+    }
+}

src/WebJobs.Script/Diagnostics/ForwardingLoggerFactory.cs

@@ -0,0 +1,45 @@
+// 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.Diagnostics;
+using Microsoft.Azure.WebJobs.Script;
+
+#nullable enable
+
+namespace Microsoft.Extensions.Logging
+{
+    /// <summary>
+    /// A logger factory that creates loggers which track the current active ScriptHost (if any), falling
+    /// back to the WebHost logger if no ScriptHost is active.
+    /// </summary>
+    [DebuggerDisplay(@"InnerFactory = \{ {_inner} \}, ScriptHostState = {_manager.State}")]
+    public sealed class ForwardingLoggerFactory : ILoggerFactory
+    {
+        private readonly ILoggerFactory _inner;
+        private readonly IScriptHostManager _manager;
+
+        public ForwardingLoggerFactory(ILoggerFactory inner, IScriptHostManager manager)
+        {
+            ArgumentNullException.ThrowIfNull(inner);
+            ArgumentNullException.ThrowIfNull(manager);
+            _inner = inner;
+            _manager = manager;
+        }
+
+        /// <inheritdoc />
+        public void AddProvider(ILoggerProvider provider)
+            => throw new NotSupportedException(
+                $"{nameof(ILoggerProvider)} can not be added to the {nameof(ForwardingLoggerFactory)}.");
+
+        /// <inheritdoc />
+        public ILogger CreateLogger(string categoryName)
+            => new ForwardingLogger(categoryName, _inner.CreateLogger(categoryName), _manager);
+
+        /// <inheritdoc />
+        public void Dispose()
+        {
+            // no op.
+        }
+    }
+}

src/WebJobs.Script/Diagnostics/HealthChecks/TelemetryHealthCheckPublisher.cs

@@ -21,7 +21,7 @@ public partial class TelemetryHealthCheckPublisher : IHealthCheckPublisher
         public TelemetryHealthCheckPublisher(
             HealthCheckMetrics metrics,
             TelemetryHealthCheckPublisherOptions options,
-            ILogger<TelemetryHealthCheckPublisher> logger)
+            [ForwardingLogger] ILogger<TelemetryHealthCheckPublisher> logger)
         {
             _metrics = metrics ?? throw new ArgumentNullException(nameof(metrics));
             _options = options ?? throw new ArgumentNullException(nameof(options));

src/WebJobs.Script/Extensions/ScriptLoggingBuilderExtensions.cs

@@ -1,20 +1,32 @@
 // 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.Collections.Concurrent;
 using System.Linq;
 using Microsoft.AspNetCore.Hosting;
 using Microsoft.Azure.WebJobs.Script;
 using Microsoft.Azure.WebJobs.Script.Configuration;
 using Microsoft.Azure.WebJobs.Script.Workers;
 using Microsoft.Extensions.Configuration;
+using Microsoft.Extensions.DependencyInjection;
 using Microsoft.Extensions.Hosting;
 
 namespace Microsoft.Extensions.Logging
 {
     public static class ScriptLoggingBuilderExtensions
     {
-        private static ConcurrentDictionary<string, bool> _filteredCategoryCache = new ConcurrentDictionary<string, bool>();
+        private static readonly ConcurrentDictionary<string, bool> _filteredCategoryCache = new();
+
+        public static ILoggingBuilder AddForwardingLogger(this ILoggingBuilder builder)
+        {
+            ArgumentNullException.ThrowIfNull(builder);
+            builder.Services.AddKeyedSingleton<ILoggerFactory, ForwardingLoggerFactory>(
+                ForwardingLogger.ServiceKey);
+            builder.Services.AddKeyedSingleton(
+                typeof(ILogger<>), ForwardingLogger.ServiceKey, typeof(ForwardingLogger<>));
+            return builder;
+        }
 
         public static ILoggingBuilder AddDefaultWebJobsFilters(this ILoggingBuilder builder)
         {
@@ -23,7 +35,8 @@ public static ILoggingBuilder AddDefaultWebJobsFilters(this ILoggingBuilder buil
             return builder;
         }
 
-        public static ILoggingBuilder AddDefaultWebJobsFilters<T>(this ILoggingBuilder builder, LogLevel level) where T : ILoggerProvider
+        public static ILoggingBuilder AddDefaultWebJobsFilters<T>(this ILoggingBuilder builder, LogLevel level)
+            where T : ILoggerProvider
         {
             builder.AddFilter<T>(null, LogLevel.None);
             builder.AddFilter<T>((c, l) => Filter(c, l, level));
@@ -37,20 +50,24 @@ internal static bool Filter(string category, LogLevel actualLevel, LogLevel minL
 
         private static bool IsFiltered(string category)
         {
-            return _filteredCategoryCache.GetOrAdd(category, static cat => ScriptConstants.SystemLogCategoryPrefixes.Any(p => cat.StartsWith(p)));
+            return _filteredCategoryCache.GetOrAdd(
+                category,
+                static cat => ScriptConstants.SystemLogCategoryPrefixes.Any(p => cat.StartsWith(p)));
         }
 
         public static void AddConsoleIfEnabled(this ILoggingBuilder builder, HostBuilderContext context)
         {
             builder.AddConsoleIfEnabled(context.HostingEnvironment.IsDevelopment(), context.Configuration);
         }
 
-        private static void AddConsoleIfEnabled(this ILoggingBuilder builder, bool isDevelopment, IConfiguration configuration)
+        private static void AddConsoleIfEnabled(
+            this ILoggingBuilder builder, bool isDevelopment, IConfiguration configuration)
         {
             // console logging defaults to false, except for self host
             bool enableConsole = isDevelopment;
 
-            string consolePath = ConfigurationPath.Combine(ConfigurationSectionNames.JobHost, "Logging", "Console", "IsEnabled");
+            string consolePath = ConfigurationPath.Combine(
+                ConfigurationSectionNames.JobHost, "Logging", "Console", "IsEnabled");
             IConfigurationSection configSection = configuration.GetSection(consolePath);
 
             if (configSection.Exists())

src/WebJobs.Script/Host/IScriptHostManager.cs

@@ -5,6 +5,8 @@
 using System.Threading;
 using System.Threading.Tasks;
 
+#nullable enable
+
 namespace Microsoft.Azure.WebJobs.Script
 {
     public interface IScriptHostManager
@@ -24,7 +26,12 @@ public interface IScriptHostManager
         /// <summary>
         /// Gets the last host <see cref="Exception"/> that has occurred.
         /// </summary>
-        Exception LastError { get; }
+        Exception? LastError { get; }
+
+        /// <summary>
+        /// Gets the current <see cref="IServiceProvider"/> for the active Script Host.
+        /// </summary>
+        IServiceProvider? Services { get; }
 
         /// <summary>
         /// Restarts the current Script Job Host.

test/WebJobs.Script.Tests.Shared/TestHelpers.cs

@@ -654,6 +654,8 @@ event EventHandler IScriptHostManager.HostInitializing
 
             Exception IScriptHostManager.LastError => throw new NotImplementedException();
 
+            IServiceProvider IScriptHostManager.Services => this;
+
             public void OnActiveHostChanged()
             {
                 ActiveHostChanged?.Invoke(this, new ActiveHostChangedEventArgs(null, null));

test/WebJobs.Script.Tests/Diagnostics/ForwardingLoggerAttributeTests.cs

@@ -0,0 +1,44 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using AwesomeAssertions;
+using Microsoft.Extensions.DependencyInjection;
+using Microsoft.Extensions.Logging;
+using Xunit;
+
+namespace Microsoft.Azure.WebJobs.Script.Tests.Diagnostics
+{
+    public class ForwardingLoggerAttributeTests
+    {
+        [Fact]
+        public void Key_IsCorrect()
+        {
+            ForwardingLoggerAttribute attribute = new();
+
+            attribute.Key.Should().BeOfType<string>()
+                .Which.Should().NotBeNullOrWhiteSpace()
+                .And.Be(ForwardingLogger.ServiceKey);
+        }
+
+        [Fact]
+        public void Import_GetsService()
+        {
+            object nonKeyed = new();
+            object keyed = new();
+
+            ServiceCollection services = new();
+            services.AddSingleton(nonKeyed);
+            services.AddKeyedSingleton(ForwardingLogger.ServiceKey, keyed);
+            services.AddSingleton<TestClass>();
+
+            TestClass test = services.BuildServiceProvider().GetRequiredService<TestClass>();
+
+            test.Instance.Should().BeSameAs(keyed);
+        }
+
+        private class TestClass([ForwardingLogger] object instance)
+        {
+            public object Instance => instance;
+        }
+    }
+}