Provide a hook for invoke

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

@@ -0,0 +1,26 @@
+// 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;
+using System.Threading.Tasks;
+
+namespace Microsoft.Azure.WebJobs.Host.Executors
+{
+    /// <summary>
+    /// Interface defining the contract for executing a triggered function.
+    /// Allows a hook around the underlying execution.
+    /// This should only be used by extensions that need very specific control over the invocation. 
+    /// </summary>
+    public interface ITriggeredFunctionExecutorWithHook
+    {
+        /// <summary>
+        /// Try to invoke the triggered function using the values specified.
+        /// </summary>
+        /// <param name="input">The trigger invocation details.</param>
+        /// <param name="cancellationToken">The cancellation token</param>
+        /// <param name="hook">a hook that wraps the underlying invocation</param>
+        /// <returns>A <see cref="FunctionResult"/> describing the results of the invocation.</returns>
+        Task<FunctionResult> TryExecuteAsync(TriggeredFunctionData input, CancellationToken cancellationToken, Func<Func<Task>, Task> hook);
+    }
+}

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

@@ -2,14 +2,15 @@
 // Licensed under the MIT License. See License.txt in the project root for license information.
 
 using System;
+using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Azure.WebJobs.Host.Protocols;
 using Microsoft.Azure.WebJobs.Host.Triggers;
 
 namespace Microsoft.Azure.WebJobs.Host.Executors
 {
-    internal class TriggeredFunctionExecutor<TTriggerValue> : ITriggeredFunctionExecutor
+    internal class TriggeredFunctionExecutor<TTriggerValue> : ITriggeredFunctionExecutor, ITriggeredFunctionExecutorWithHook
     {
         private FunctionDescriptor _descriptor;
         private ITriggeredFunctionInstanceFactory<TTriggerValue> _instanceFactory;
@@ -30,16 +31,46 @@ public FunctionDescriptor Function
             }
         }
 
-        public async Task<FunctionResult> TryExecuteAsync(TriggeredFunctionData input, CancellationToken cancellationToken)
+        public Task<FunctionResult> TryExecuteAsync(TriggeredFunctionData input, CancellationToken cancellationToken)
+        {
+            return TryExecuteAsync(input, cancellationToken, null);
+        }
+
+        public async Task<FunctionResult> TryExecuteAsync(TriggeredFunctionData input, CancellationToken cancellationToken, Func<Func<Task>, Task> hook)
         {
             IFunctionInstance instance = _instanceFactory.Create((TTriggerValue)input.TriggerValue, input.ParentId);
+            if (hook != null)
+            {
+                IFunctionInvoker invoker = new InvokeWrapper(instance.Invoker, hook);
+                instance = new FunctionInstance(instance.Id, instance.ParentId, instance.Reason, instance.BindingSource, invoker, instance.FunctionDescriptor);
+            }
+
             IDelayedException exception = await _executor.TryExecuteAsync(instance, cancellationToken);
 
             FunctionResult result = exception != null ?
-                new FunctionResult(exception.Exception) 
+                new FunctionResult(exception.Exception)
                 : new FunctionResult(true);
 
             return result;
         }
+
+        private class InvokeWrapper : IFunctionInvoker
+        {
+            private readonly IFunctionInvoker _inner;
+            private readonly Func<Func<Task>, Task> _hook;
+
+            public InvokeWrapper(IFunctionInvoker inner, Func<Func<Task>, Task> hook)
+            {
+                _inner = inner;
+                _hook = hook;
+            }
+            public IReadOnlyList<string> ParameterNames => _inner.ParameterNames;
+
+            public Task InvokeAsync(object[] arguments)
+            {
+                Func<Task> inner = () => _inner.InvokeAsync(arguments);
+                return _hook(inner);
+            }
+        }
     }
 }

src/Microsoft.Azure.WebJobs.Host/WebJobs.Host.csproj

@@ -490,6 +490,7 @@
     <Compile Include="Diagnostics\ExceptionFormatter.cs" />
     <Compile Include="Executors\CompositeFunctionEventCollector.cs" />
     <Compile Include="Executors\FunctionInstanceTraceWriter.cs" />
+    <Compile Include="Executors\ITriggeredFunctionExecutorWithHook.cs" />
     <Compile Include="Extensions\CloudQueueMessageExtensions.cs" />
     <Compile Include="Extensions\IJobHostMetadataProvider.cs" />
     <Compile Include="Extensions\JobHostMetadataProvider.cs" />

test/Microsoft.Azure.WebJobs.Host.UnitTests/Executors/TriggeredFunctionExecutorTests.cs

@@ -0,0 +1,72 @@
+// 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.Azure.WebJobs.Host.Executors;
+using Xunit;
+using Microsoft.Azure.WebJobs.Host.Protocols;
+using Microsoft.Azure.WebJobs.Host.Triggers;
+using System.Threading.Tasks;
+using System.Threading;
+using System.Text;
+using Moq;
+
+namespace Microsoft.Azure.WebJobs.Host.UnitTests.Executors
+{
+    public class TriggeredFunctionExecutorTests
+    {
+        // Test ITriggeredFunctionExecutorWithHook
+        [Fact]
+        public async Task TestHook()
+        {
+            StringBuilder sb = new StringBuilder();
+
+            var descr = new FunctionDescriptor();
+
+            // IFunctionExecutor just passes through to Invoker.
+            var mockExecutor = new Mock<IFunctionExecutor>();
+            mockExecutor.Setup(m => m.TryExecuteAsync(It.IsAny<IFunctionInstance>(), It.IsAny<CancellationToken>())).
+                Returns<IFunctionInstance, CancellationToken>((x, y) =>
+                {
+                    sb.Append("2>");
+                    x.Invoker.InvokeAsync(null).Wait();
+                    sb.Append("<6");
+                    return Task.FromResult<IDelayedException>(null);
+                });
+            IFunctionExecutor executor = mockExecutor.Object;
+
+            var mockInvoker = new Mock<IFunctionInvoker>();
+            mockInvoker.Setup(m => m.InvokeAsync(null)).Returns(() =>
+            {
+                sb.Append("4");
+                return Task.CompletedTask;
+            }
+                );
+            IFunctionInvoker innerInvoker = mockInvoker.Object;
+
+            IFunctionInstance inner = new FunctionInstance(Guid.NewGuid(), null, ExecutionReason.HostCall, null, innerInvoker, null);
+
+            var mockInstanceFactory = new Mock<ITriggeredFunctionInstanceFactory<int>>();
+            mockInstanceFactory.Setup(m => m.Create(It.IsAny<int>(), null)).Returns(inner);
+            ITriggeredFunctionInstanceFactory<int> instanceFactory = mockInstanceFactory.Object;
+
+            var trigger = new TriggeredFunctionExecutor<int>(descr, executor, instanceFactory);
+
+            var trigger2 = (ITriggeredFunctionExecutorWithHook)trigger;
+
+
+
+            Func<Func<Task>, Task> hook = async (x) => {
+                sb.Append("3>");
+                await x();
+                sb.Append("<5");
+            };
+
+            sb.Append("1>");
+            await trigger2.TryExecuteAsync(new TriggeredFunctionData { TriggerValue = 123 }, CancellationToken.None, hook);
+            sb.Append("<7");
+
+            Assert.Equal("1>2>3>4<5<6<7", sb.ToString());
+        }
+    }
+}

test/Microsoft.Azure.WebJobs.Host.UnitTests/PublicSurfaceTests.cs

@@ -196,6 +196,7 @@ public void WebJobsHostPublicSurface_LimitedToSpecificTypes()
                 "AmbientConnectionStringProvider",
                 "IExtensionRegistryExtensions",
                 "ITriggeredFunctionExecutor",
+                "ITriggeredFunctionExecutorWithHook",
                 "ListenerFactoryContext",
                 "BindingTemplateSource",
                 "TriggeredFunctionData",

test/Microsoft.Azure.WebJobs.Host.UnitTests/WebJobs.Host.UnitTests.csproj

@@ -256,6 +256,7 @@
     <Compile Include="Blobs\UncompletedAsyncResult.cs" />
     <Compile Include="Blobs\CompletedAsyncResult.cs" />
     <Compile Include="Blobs\WatchableReadStreamTests.cs" />
+    <Compile Include="Executors\TriggeredFunctionExecutorTests.cs" />
     <Compile Include="ExtensionConfigContextTests.cs" />
     <Compile Include="Loggers\ApplicationInsightsLoggerProviderTests.cs" />
     <Compile Include="Loggers\DefaultTelemetryClientFactoryTests.cs" />