Adding IListenerDecorator support (#3011)

src/Microsoft.Azure.WebJobs.Host/Executors/JobHostContextFactory.cs

@@ -52,6 +52,7 @@ internal class JobHostContextFactory : IJobHostContextFactory
         private readonly IDrainModeManager _drainModeManager;
         private readonly IApplicationLifetime _applicationLifetime;
         private readonly ITargetScalerManager _targetScalerManager;
+        private readonly IEnumerable<IListenerDecorator> _listenerDecorators;
 
         public JobHostContextFactory(
             IDashboardLoggingSetup dashboardLoggingSetup,
@@ -75,7 +76,8 @@ public JobHostContextFactory(
             IScaleMonitorManager monitorManager,
             IDrainModeManager drainModeManager,
             IApplicationLifetime applicationLifetime,
-            ITargetScalerManager targetScalerManager)
+            ITargetScalerManager targetScalerManager,
+            IEnumerable<IListenerDecorator> listenerDecorators)
         {
             _dashboardLoggingSetup = dashboardLoggingSetup;
             _functionExecutor = functionExecutor;
@@ -99,6 +101,7 @@ public JobHostContextFactory(
             _drainModeManager = drainModeManager;
             _applicationLifetime = applicationLifetime;
             _targetScalerManager = targetScalerManager;
+            _listenerDecorators = listenerDecorators;
         }
 
         public async Task<JobHostContext> Create(JobHost host, CancellationToken shutdownToken, CancellationToken cancellationToken)
@@ -127,8 +130,7 @@ public async Task<JobHostContext> Create(JobHost host, CancellationToken shutdow
                     // they are started).
                     host.OnHostInitialized();
                 };
-                IListenerFactory functionsListenerFactory = new HostListenerFactory(functions.ReadAll(), _singletonManager, _activator, _nameResolver, _loggerFactory,
-                    _monitorManager, _targetScalerManager, listenersCreatedCallback, _jobHostOptions.Value.AllowPartialHostStartup, _drainModeManager);
+                IListenerFactory functionsListenerFactory = new HostListenerFactory(functions.ReadAll(), _loggerFactory, _monitorManager, _targetScalerManager, _listenerDecorators, listenersCreatedCallback, _drainModeManager);
 
                 string hostId = await _hostIdProvider.GetHostIdAsync(cancellationToken);
                 bool dashboardLoggingEnabled = _dashboardLoggingSetup.Setup(functions, functionsListenerFactory, out IFunctionExecutor hostCallExecutor,

src/Microsoft.Azure.WebJobs.Host/Hosting/WebJobsServiceCollectionExtensions.cs

@@ -128,6 +128,8 @@ public static IWebJobsBuilder AddWebJobs(this IServiceCollection services, Actio
             services.TryAddSingleton<ConcurrencyManager>();
             services.TryAddEnumerable(ServiceDescriptor.Singleton<IHostedService, ConcurrencyManagerService>());
 
+            AddListenerDecorators(services);
+
             services.ConfigureOptions<ConcurrencyOptionsSetup>();
             services.ConfigureOptions<PrimaryHostCoordinatorOptionsSetup>();
             services.AddOptions<ConcurrencyOptions>()
@@ -210,5 +212,13 @@ private static void AddOptionsLogging(this IServiceCollection services)
             services.AddSingleton<IHostedService, OptionsLoggingService>();
             services.AddSingleton<IOptionsFormatter<LoggerFilterOptions>, LoggerFilterOptionsFormatter>();
         }
+
+        private static void AddListenerDecorators(IServiceCollection services)
+        {
+            // Order is important for these platform decorator registrations!
+            // They will be applied in this order after any user registered decorators.
+            services.TryAddEnumerable(ServiceDescriptor.Singleton<IListenerDecorator, SingletonListenerDecorator>());
+            services.TryAddEnumerable(ServiceDescriptor.Singleton<IListenerDecorator, FunctionListenerDecorator>());
+        }
     }
 }

src/Microsoft.Azure.WebJobs.Host/Listeners/FunctionListenerDecorator.cs

@@ -0,0 +1,27 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Options;
+
+namespace Microsoft.Azure.WebJobs.Host.Listeners
+{
+    internal class FunctionListenerDecorator : IListenerDecorator
+    {
+        private readonly ILoggerFactory _loggerFactory;
+        private readonly IOptions<JobHostOptions> _jobHostOptions;
+
+        public FunctionListenerDecorator(ILoggerFactory loggerFactory, IOptions<JobHostOptions> jobHostOptions)
+        {
+            _loggerFactory = loggerFactory;
+            _jobHostOptions = jobHostOptions;
+        }
+
+        public IListener Decorate(ListenerDecoratorContext context)
+        {
+            // wrap the listener with a function listener to handle exceptions
+            bool allowPartialHostStartup = _jobHostOptions.Value.AllowPartialHostStartup;
+            return new FunctionListener(context.Listener, context.FunctionDefinition.Descriptor, _loggerFactory, allowPartialHostStartup);
+        }
+    }
+}

src/Microsoft.Azure.WebJobs.Host/Listeners/HostListenerFactory.cs

@@ -14,7 +14,6 @@
 using Microsoft.Azure.WebJobs.Logging;
 using Microsoft.Extensions.Configuration;
 using Microsoft.Extensions.Logging;
-using Microsoft.Extensions.Options;
 
 namespace Microsoft.Azure.WebJobs.Host.Listeners
 {
@@ -23,31 +22,24 @@ internal class HostListenerFactory : IListenerFactory
         private static readonly MethodInfo JobActivatorCreateMethod = typeof(IJobActivator).GetMethod("CreateInstance", BindingFlags.Public | BindingFlags.Instance).GetGenericMethodDefinition();
         private const string IsDisabledFunctionName = "IsDisabled";
         private readonly IEnumerable<IFunctionDefinition> _functionDefinitions;
-        private readonly SingletonManager _singletonManager;
-        private readonly IJobActivator _activator;
-        private readonly INameResolver _nameResolver;
         private readonly ILoggerFactory _loggerFactory;
         private readonly ILogger _logger;
-        private readonly bool _allowPartialHostStartup;
         private readonly Action _listenersCreatedCallback;
         private readonly IScaleMonitorManager _monitorManager;
         private readonly ITargetScalerManager _targetScalerManager;
         private readonly IDrainModeManager _drainModeManager;
+        private readonly IEnumerable<IListenerDecorator> _listenerDecorators;
 
-        public HostListenerFactory(IEnumerable<IFunctionDefinition> functionDefinitions, SingletonManager singletonManager, IJobActivator activator,
-            INameResolver nameResolver, ILoggerFactory loggerFactory, IScaleMonitorManager monitorManager, ITargetScalerManager targetScalerManager, Action listenersCreatedCallback, bool allowPartialHostStartup = false, IDrainModeManager drainModeManager = null)
+        public HostListenerFactory(IEnumerable<IFunctionDefinition> functionDefinitions, ILoggerFactory loggerFactory, IScaleMonitorManager monitorManager, ITargetScalerManager targetScalerManager, IEnumerable<IListenerDecorator> listenerDecorators, Action listenersCreatedCallback, IDrainModeManager drainModeManager = null)
         {
             _functionDefinitions = functionDefinitions;
-            _singletonManager = singletonManager;
-            _activator = activator;
-            _nameResolver = nameResolver;
             _loggerFactory = loggerFactory;
             _logger = _loggerFactory?.CreateLogger(LogCategories.Startup);
-            _allowPartialHostStartup = allowPartialHostStartup;
             _monitorManager = monitorManager;
             _targetScalerManager = targetScalerManager;
             _listenersCreatedCallback = listenersCreatedCallback;
             _drainModeManager = drainModeManager;
+            _listenerDecorators = listenerDecorators;
         }
 
         public async Task<IListener> CreateAsync(CancellationToken cancellationToken)
@@ -59,41 +51,39 @@ public async Task<IListener> CreateAsync(CancellationToken cancellationToken)
                 // Determine if the function is disabled
                 if (functionDefinition.Descriptor.IsDisabled)
                 {
-                    string msg = string.Format("Function '{0}' is disabled", functionDefinition.Descriptor.ShortName);
-                    _logger?.LogInformation(msg);
+                    _logger?.LogInformation($"Function '{functionDefinition.Descriptor.ShortName}' is disabled");
                     continue;
                 }
 
+                // Create the listener
                 IListenerFactory listenerFactory = functionDefinition.ListenerFactory;
                 if (listenerFactory == null)
                 {
                     continue;
                 }
-
                 IListener listener = await listenerFactory.CreateAsync(cancellationToken);
 
-                RegisterScaleMonitor(listener, _monitorManager);
-                RegisterTargetScaler(listener, _targetScalerManager);
-
-                // if the listener is a Singleton, wrap it with our SingletonListener
-                SingletonAttribute singletonAttribute = SingletonManager.GetListenerSingletonOrNull(listener.GetType(), functionDefinition.Descriptor);
-                if (singletonAttribute != null)
-                {
-                    listener = new SingletonListener(functionDefinition.Descriptor, singletonAttribute, _singletonManager, listener, _loggerFactory);
-                }
+                RegisterScalers(listener);
+
+                listener = ApplyDecorators(listener, functionDefinition);
 
-                // wrap the listener with a function listener to handle exceptions
-                listener = new FunctionListener(listener, functionDefinition.Descriptor, _loggerFactory, _allowPartialHostStartup);
                 listeners.Add(listener);
             }
 
             _listenersCreatedCallback?.Invoke();
 
             var compositeListener = new CompositeListener(listeners);
             _drainModeManager?.RegisterListener(compositeListener);
+
             return compositeListener;
         }
 
+        internal void RegisterScalers(IListener listener)
+        {
+            RegisterScaleMonitor(listener, _monitorManager);
+            RegisterTargetScaler(listener, _targetScalerManager);
+        }
+
         /// <summary>
         /// Check to see if the specified listener is an <see cref="IScaleMonitor"/> and if so
         /// register it with the <see cref="IScaleMonitorManager"/>.
@@ -229,5 +219,42 @@ internal static bool IsDisabledByProvider(Type providerType, MethodInfo jobFunct
                 return (bool)methodInfo.Invoke(instance, new object[] { jobFunction });
             }
         }
+
+        /// <summary>
+        /// Applies any user registered decorators, followed by any platform decorators.
+        /// See <see cref="WebJobsServiceCollectionExtensions.AddListenerDecorators(Extensions.DependencyInjection.IServiceCollection)"/>.
+        /// </summary>
+        /// <param name="listener">The listener to apply decorators to.</param>
+        /// <param name="functionDefinition">The function the the listener is for.</param>
+        /// <returns>The resulting listener with decorators applied.</returns>
+        private IListener ApplyDecorators(IListener listener, IFunctionDefinition functionDefinition)
+        {
+            Type rootListenerType = listener.GetType();
+            var platformDecorators = _listenerDecorators.Where(p => p.GetType().Assembly == typeof(HostListenerFactory).Assembly);
+            var userDecorators = _listenerDecorators.Except(platformDecorators);
+
+            listener = ApplyDecorators(userDecorators, listener, functionDefinition, rootListenerType, writeLog: true);
+
+            // Order is important - platform decorators must be applied AFTER any user decorators, in order.
+            listener = ApplyDecorators(platformDecorators, listener, functionDefinition, rootListenerType);
+
+            return listener;
+        }
+
+        private IListener ApplyDecorators(IEnumerable<IListenerDecorator> decorators, IListener listener, IFunctionDefinition functionDefinition, Type rootListenerType, bool writeLog = false)
+        {
+            foreach (var decorator in decorators)
+            {
+                var context = new ListenerDecoratorContext(functionDefinition, rootListenerType, listener);
+                listener = decorator.Decorate(context);
+
+                if (writeLog)
+                {
+                    _logger.LogDebug($"Applying IListenerDecorator '{decorator.GetType().FullName}'");
+                }
+            }
+
+            return listener;
+        }
     }
 }

src/Microsoft.Azure.WebJobs.Host/Listeners/IListenerDecorator.cs

@@ -0,0 +1,20 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+namespace Microsoft.Azure.WebJobs.Host.Listeners
+{
+    /// <summary>
+    /// Custom decorator interface called during <see cref="JobHost"/> listener creation to
+    /// allow function listeners to be customized.
+    /// </summary>
+    public interface IListenerDecorator
+    {
+        /// <summary>
+        /// Creates a listener.
+        /// </summary>
+        /// <param name="context">The listener context.</param>
+        /// <returns>The listener to use. This may be a new wrapped listener, or the original
+        /// listener.</returns>
+        IListener Decorate(ListenerDecoratorContext context);
+    }
+}

src/Microsoft.Azure.WebJobs.Host/Listeners/ListenerDecoratorContext.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 Microsoft.Azure.WebJobs.Host.Indexers;
+using System;
+
+namespace Microsoft.Azure.WebJobs.Host.Listeners
+{
+    /// <summary>
+    /// Context class for <see cref="IListenerDecorator.Decorate(ListenerDecoratorContext)"/>.
+    /// </summary>
+    public class ListenerDecoratorContext
+    {
+        /// <summary>
+        /// Constructs an instance.
+        /// </summary>
+        /// <param name="functionDefinition">The function the specified listener is for.</param>
+        /// <param name="rootListenerType">Gets the type of the root listener.</param>
+        /// <param name="listener">The listener to decorate.</param>
+        public ListenerDecoratorContext(IFunctionDefinition functionDefinition, Type rootListenerType, IListener listener)
+        {
+            FunctionDefinition = functionDefinition;
+            ListenerType = rootListenerType;
+            Listener = listener;
+        }
+
+        /// <summary>
+        /// Gets the <see cref="IFunctionDefinition"/> the specified listener is for.
+        /// </summary>
+        public IFunctionDefinition FunctionDefinition { get; }
+
+        /// <summary>
+        /// Gets the listener to decorate.
+        /// </summary>
+        public IListener Listener { get; }
+
+        /// <summary>
+        /// Gets the type of the root listener.
+        /// </summary>
+        public Type ListenerType { get; }
+    }
+}

src/Microsoft.Azure.WebJobs.Host/Listeners/SingletonListenerDecorator.cs

@@ -0,0 +1,34 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using Microsoft.Extensions.Logging;
+
+namespace Microsoft.Azure.WebJobs.Host.Listeners
+{
+    internal class SingletonListenerDecorator : IListenerDecorator
+    {
+        private readonly SingletonManager _singletonManager;
+        private readonly ILoggerFactory _loggerFactory;
+
+        public SingletonListenerDecorator(SingletonManager singletonManager, ILoggerFactory loggerFactory) 
+        {
+            _singletonManager = singletonManager;
+            _loggerFactory = loggerFactory;
+        }
+
+        public IListener Decorate(ListenerDecoratorContext context)
+        {
+            var functionDescriptor = context.FunctionDefinition.Descriptor;
+
+            // if the listener is a Singleton, wrap it with our SingletonListener
+            IListener listener = context.Listener;
+            SingletonAttribute singletonAttribute = SingletonManager.GetListenerSingletonOrNull(context.ListenerType, functionDescriptor);
+            if (singletonAttribute != null)
+            {
+                listener = new SingletonListener(functionDescriptor, singletonAttribute, _singletonManager, listener, _loggerFactory);
+            }
+
+            return listener;
+        }
+    }
+}