Ensure SingletonListener is applied before custom Listener decorators (…

…#3049)
Azure · Dec 15, 2023 · bb79ef3 · bb79ef3
1 parent 095fae1
commit bb79ef3
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 45 deletions.
diff --git a/src/Microsoft.Azure.WebJobs.Host/Hosting/WebJobsServiceCollectionExtensions.cs b/src/Microsoft.Azure.WebJobs.Host/Hosting/WebJobsServiceCollectionExtensions.cs
@@ -217,7 +217,6 @@ 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>());
         }
     }

diff --git a/src/Microsoft.Azure.WebJobs.Host/Indexers/FunctionIndexer.cs b/src/Microsoft.Azure.WebJobs.Host/Indexers/FunctionIndexer.cs
@@ -376,7 +376,7 @@ private FunctionDefinition CreateTriggeredFunctionDefinition<TTriggerValue>(
             ITriggeredFunctionBinding<TTriggerValue> functionBinding = new TriggeredFunctionBinding<TTriggerValue>(descriptor, parameterName, triggerBinding, nonTriggerBindings, _singletonManager);
             ITriggeredFunctionInstanceFactory<TTriggerValue> instanceFactory = new TriggeredFunctionInstanceFactory<TTriggerValue>(functionBinding, invoker, descriptor, _serviceScopeFactory);
             ITriggeredFunctionExecutor triggerExecutor = new TriggeredFunctionExecutor<TTriggerValue>(descriptor, _executor, instanceFactory, _loggerFactory);
-            IListenerFactory listenerFactory = new ListenerFactory(descriptor, triggerExecutor, triggerBinding, _sharedQueue);
+            IListenerFactory listenerFactory = new ListenerFactory(descriptor, triggerExecutor, triggerBinding, _sharedQueue, _singletonManager, _loggerFactory);
 
             return new FunctionDefinition(descriptor, instanceFactory, listenerFactory);
         }
@@ -470,19 +470,36 @@ private class ListenerFactory : IListenerFactory
             private readonly ITriggeredFunctionExecutor _executor;
             private readonly ITriggerBinding _binding;
             private readonly SharedQueueHandler _sharedQueue;
+            private readonly SingletonManager _singletonManager;
+            private readonly ILoggerFactory _loggerFactory;
 
-            public ListenerFactory(FunctionDescriptor descriptor, ITriggeredFunctionExecutor executor, ITriggerBinding binding, SharedQueueHandler sharedQueue)
+            public ListenerFactory(FunctionDescriptor descriptor, ITriggeredFunctionExecutor executor, ITriggerBinding binding, SharedQueueHandler sharedQueue, SingletonManager singletonManager, ILoggerFactory loggerFactory)
             {
                 _descriptor = descriptor;
                 _executor = executor;
                 _binding = binding;
                 _sharedQueue = sharedQueue;
+                _singletonManager = singletonManager;
+                _loggerFactory = loggerFactory;
             }
 
             public async Task<IListener> CreateAsync(CancellationToken cancellationToken)
             {
                 ListenerFactoryContext context = new ListenerFactoryContext(_descriptor, _executor, _sharedQueue, cancellationToken);
-                return await _binding.CreateListenerAsync(context);
+                IListener listener = await _binding.CreateListenerAsync(context);
+
+                if (listener != null)
+                {
+                    // if the listener is a Singleton, wrap it with our SingletonListener
+                    Type listenerType = listener.GetType();
+                    SingletonAttribute singletonAttribute = SingletonManager.GetListenerSingletonOrNull(listenerType, _descriptor);
+                    if (singletonAttribute != null)
+                    {
+                        listener = new SingletonListener(_descriptor, singletonAttribute, _singletonManager, listener, _loggerFactory);
+                    }
+                }
+
+                return listener;
             }
         }
 

diff --git a/src/Microsoft.Azure.WebJobs.Host/Listeners/SingletonListenerDecorator.cs b/src/Microsoft.Azure.WebJobs.Host/Listeners/SingletonListenerDecorator.cs
diff --git a/test/Microsoft.Azure.WebJobs.Host.FunctionalTests/HostListenerFactoryTests.cs b/test/Microsoft.Azure.WebJobs.Host.FunctionalTests/HostListenerFactoryTests.cs
@@ -50,15 +50,10 @@ public async Task CreateAsync_AppliesListenerDecorators()
             var innerListeners = ((IEnumerable<IListener>)listener).ToArray();
             var innerListener = innerListeners[0];
 
-            // expect the first two outer listeners to be our platform listeners
+            // expect the first listener to be our platform FunctionListener
             FunctionListener functionListener = (FunctionListener)innerListener;
             var innerListenerField = typeof(FunctionListener).GetField("_listener", BindingFlags.NonPublic | BindingFlags.Instance);
-            innerListener = (SingletonListener)innerListenerField.GetValue(functionListener);
-
-            innerListenerField = typeof(SingletonListener).GetField("_innerListener", BindingFlags.NonPublic | BindingFlags.Instance);
-            innerListener = (IListener)innerListenerField.GetValue(innerListener);
-
-            TestListenerDecorator.DecoratorListener decoratorListener = (TestListenerDecorator.DecoratorListener)innerListener;
+            TestListenerDecorator.DecoratorListener decoratorListener = (TestListenerDecorator.DecoratorListener)innerListenerField.GetValue(functionListener);
 
             // verify all decorators were consulted, resulting in a nested stack of listeners
             Assert.Equal("C", decoratorListener.Tag);
@@ -67,6 +62,14 @@ public async Task CreateAsync_AppliesListenerDecorators()
             decoratorListener = (TestListenerDecorator.DecoratorListener)decoratorListener.InnerListener;
             Assert.Equal("A", decoratorListener.Tag);
 
+            // after decorators, we expect the Singleton listener
+            SingletonListener singletonListener = (SingletonListener)decoratorListener.InnerListener;
+
+            // finally, the QueueListener
+            innerListenerField = typeof(SingletonListener).GetField("_innerListener", BindingFlags.NonPublic | BindingFlags.Instance);
+            innerListener = (IListener)innerListenerField.GetValue(singletonListener);
+            Assert.Equal("QueueListener", innerListener.GetType().Name);
+
             var logProvider = _testHost.GetTestLoggerProvider();
             var logs = logProvider.GetAllLogMessages().Where(p => p.Category == LogCategories.Startup && p.FormattedMessage.Contains(nameof(IListenerDecorator))).ToArray();
             Assert.Equal(3, logs.Length);