ASB transaction support

src/NServiceBus.AzureFunctions.InProcess.ServiceBus/FodyWeavers.xml

@@ -0,0 +1,4 @@
+<?xml version="1.0" encoding="utf-8"?>
+<Weavers GenerateXsd="false">
+  <Obsolete />
+</Weavers>

src/NServiceBus.AzureFunctions.InProcess.ServiceBus/FunctionEndpoint.cs

@@ -10,6 +10,8 @@
     using Extensibility;
     using Logging;
     using Microsoft.Azure.ServiceBus;
+    using Microsoft.Azure.ServiceBus.Core;
+    using Microsoft.Azure.WebJobs;
     using Microsoft.Extensions.Logging;
     using Transport;
     using ExecutionContext = Microsoft.Azure.WebJobs.ExecutionContext;
@@ -28,54 +30,118 @@ internal FunctionEndpoint(IStartableEndpointWithExternallyManagedContainer exter
             endpointFactory = _ => externallyManagedContainerEndpoint.Start(serviceProvider);
         }
 
+        /// <summary>
+        /// Processes a message received from an AzureServiceBus trigger using the NServiceBus message pipeline. This method will lookup the <see cref="ServiceBusTriggerAttribute.AutoComplete"/> setting to determine whether to use transactional or non-transactional processing.
+        /// </summary>
+        Task IFunctionEndpoint.Process(Message message, ExecutionContext executionContext, IMessageReceiver messageReceiver, ILogger functionsLogger) =>
+            ReflectionHelper.GetAutoCompleteValue()
+                ? ProcessNonTransactional(message, executionContext, messageReceiver, functionsLogger)
+                : ProcessTransactional(message, executionContext, messageReceiver, functionsLogger);
+
+        /// <summary>
+        /// Processes a message received from an AzureServiceBus trigger using the NServiceBus message pipeline. All messages are committed transactionally with the successful processing of the incoming message.
+        /// <remarks>Requires <see cref="ServiceBusTriggerAttribute.AutoComplete"/> to be set to false!</remarks>
+        /// </summary>
+        public async Task ProcessTransactional(Message message, ExecutionContext executionContext, IMessageReceiver messageReceiver, ILogger functionsLogger = null)
+        {
+            FunctionsLoggerFactory.Instance.SetCurrentLogger(functionsLogger);
+
+            var functionExecutionContext = new FunctionExecutionContext(executionContext, functionsLogger);
+
+            try
+            {
+                await InitializeEndpointIfNecessary(functionExecutionContext, CancellationToken.None)
+                    .ConfigureAwait(false);
+
+                await Process(message, new MessageReceiverTransactionStrategy(message, messageReceiver), pipeline)
+                    .ConfigureAwait(false);
+            }
+            catch (Exception)
+            {
+                // abandon message outside of a transaction scope to ensure the abandon operation can't be rolled back
+                await messageReceiver.AbandonAsync(message.SystemProperties.LockToken).ConfigureAwait(false);
+                throw;
+            }
+        }
+
+        /// <summary>
+        /// Processes a message received from an AzureServiceBus trigger using the NServiceBus message pipeline.
+        /// </summary>
+        public Task ProcessNonTransactional(Message message, ExecutionContext executionContext, IMessageReceiver messageReceiver, ILogger functionsLogger = null) => Process(message, executionContext, functionsLogger);
+
         /// <summary>
         /// Processes a message received from an AzureServiceBus trigger using the NServiceBus message pipeline.
         /// </summary>
+        [ObsoleteEx(
+            ReplacementTypeOrMember = "Process(Message, ExecutionContext, IMessageReceiver, ILogger)",
+            TreatAsErrorFromVersion = "2",
+            RemoveInVersion = "3")]
         public async Task Process(Message message, ExecutionContext executionContext, ILogger functionsLogger = null)
         {
             FunctionsLoggerFactory.Instance.SetCurrentLogger(functionsLogger);
 
-            var messageContext = CreateMessageContext(message);
             var functionExecutionContext = new FunctionExecutionContext(executionContext, functionsLogger);
 
-            await InitializeEndpointIfNecessary(functionExecutionContext,
-                messageContext.ReceiveCancellationTokenSource.Token).ConfigureAwait(false);
+            await InitializeEndpointIfNecessary(functionExecutionContext, CancellationToken.None)
+                .ConfigureAwait(false);
+
+            await Process(message, NoTransactionStrategy.Instance, pipeline)
+                .ConfigureAwait(false);
+        }
+
+        internal static async Task Process(Message message, ITransactionStrategy transactionStrategy, PipelineInvoker pipeline)
+        {
+            var messageId = message.GetMessageId();
 
             try
             {
-                await pipeline.PushMessage(messageContext).ConfigureAwait(false);
+                using (var transaction = transactionStrategy.CreateTransaction())
+                {
+                    var transportTransaction = transactionStrategy.CreateTransportTransaction(transaction);
+                    var messageContext = CreateMessageContext(transportTransaction);
+
+                    await pipeline.PushMessage(messageContext).ConfigureAwait(false);
+
+                    await transactionStrategy.Complete(transaction).ConfigureAwait(false);
+
+                    transaction?.Commit();
+                }
             }
             catch (Exception exception)
             {
-                var errorContext = new ErrorContext(
-                    exception,
-                    message.GetHeaders(),
-                    messageContext.MessageId,
-                    messageContext.Body,
-                    new TransportTransaction(),
-                    message.SystemProperties.DeliveryCount);
+                using (var transaction = transactionStrategy.CreateTransaction())
+                {
+                    var transportTransaction = transactionStrategy.CreateTransportTransaction(transaction);
+                    var errorContext = new ErrorContext(
+                        exception,
+                        message.GetHeaders(),
+                        messageId,
+                        message.Body,
+                        transportTransaction,
+                        message.SystemProperties.DeliveryCount);
+
+                    var errorHandleResult = await pipeline.PushFailedMessage(errorContext).ConfigureAwait(false);
+
+                    if (errorHandleResult == ErrorHandleResult.Handled)
+                    {
+                        await transactionStrategy.Complete(transaction).ConfigureAwait(false);
 
-                var errorHandleResult = await pipeline.PushFailedMessage(errorContext).ConfigureAwait(false);
+                        transaction?.Commit();
+                        return;
+                    }
 
-                if (errorHandleResult == ErrorHandleResult.Handled)
-                {
-                    // return to signal to the Functions host it can complete the incoming message
-                    return;
+                    throw;
                 }
-
-                throw;
             }
 
-            MessageContext CreateMessageContext(Message originalMessage)
-            {
-                return new MessageContext(
-                    originalMessage.GetMessageId(),
-                    originalMessage.GetHeaders(),
-                    originalMessage.Body,
-                    new TransportTransaction(),
+            MessageContext CreateMessageContext(TransportTransaction transportTransaction) =>
+                new MessageContext(
+                    messageId,
+                    message.GetHeaders(),
+                    message.Body,
+                    transportTransaction,
                     new CancellationTokenSource(),
                     new ContextBag());
-            }
         }
 
         /// <summary>

src/NServiceBus.AzureFunctions.InProcess.ServiceBus/IFunctionEndpoint.cs

@@ -3,6 +3,7 @@
     using System;
     using System.Threading.Tasks;
     using Microsoft.Azure.ServiceBus;
+    using Microsoft.Azure.ServiceBus.Core;
     using Microsoft.Azure.WebJobs;
     using Microsoft.Extensions.Logging;
 
@@ -12,9 +13,18 @@
     /// </summary>
     public interface IFunctionEndpoint
     {
+        /// <summary>
+        /// Processes a message received from an AzureServiceBus trigger using the NServiceBus message pipeline. This method will lookup the <see cref="ServiceBusTriggerAttribute.AutoComplete"/> setting to determine whether to use transactional or non-transactional processing.
+        /// </summary>
+        Task Process(Message message, ExecutionContext executionContext, IMessageReceiver messageReceiver, ILogger functionsLogger = null);
+
         /// <summary>
         /// Processes a message received from an AzureServiceBus trigger using the NServiceBus message pipeline.
         /// </summary>
+        [ObsoleteEx(
+            ReplacementTypeOrMember = "Process(Message, ExecutionContext, IMessageReceiver, ILogger)",
+            TreatAsErrorFromVersion = "2",
+            RemoveInVersion = "3")]
         Task Process(Message message, ExecutionContext executionContext, ILogger functionsLogger = null);
 
         /// <summary>

src/NServiceBus.AzureFunctions.InProcess.ServiceBus/NServiceBus.AzureFunctions.InProcess.ServiceBus.csproj

@@ -16,6 +16,7 @@
     <PackageReference Include="NServiceBus.Newtonsoft.Json" Version="[2.3.0, 3.0.0)" />
     <PackageReference Include="NServiceBus.Transport.AzureServiceBus" Version="[1.9.0, 2.0.0)" />
     <PackageReference Include="NServiceBus.Extensions.DependencyInjection" Version="[1.0.1, 2.0.0)" />
+    <PackageReference Include="Obsolete.Fody" Version="5.2.1" PrivateAssets="All"/>
     <PackageReference Include="Particular.Packaging" Version="1.2.1" PrivateAssets="All" />
   </ItemGroup>
 

src/NServiceBus.AzureFunctions.InProcess.ServiceBus/NServiceBusEndpointNameAttribute.cs

@@ -4,6 +4,10 @@
     /// Assembly attribute to specify NServiceBus logical endpoint name.
     /// This name is used to wire up an auto-generated service bus trigger function, responding to messages in the queue specified by the name provided.
     ///</summary>
+    [ObsoleteEx(
+        ReplacementTypeOrMember = nameof(NServiceBusTriggerFunctionAttribute),
+        TreatAsErrorFromVersion = "2",
+        RemoveInVersion = "3")]
     [System.AttributeUsage(System.AttributeTargets.Assembly)]
     public sealed class NServiceBusEndpointNameAttribute : System.Attribute
     {

src/NServiceBus.AzureFunctions.InProcess.ServiceBus/NServiceBusTriggerFunctionAttribute.cs

@@ -0,0 +1,34 @@
+namespace NServiceBus
+{
+    ///<summary>
+    /// Assembly attribute to configure generated NServiceBus Azure Function.
+    /// The attribute is used to wire up an auto-generated Service Bus trigger function, responding to messages in the queue specified by the name provided.
+    ///</summary>
+    [System.AttributeUsage(System.AttributeTargets.Assembly)]
+    public sealed class NServiceBusTriggerFunctionAttribute : System.Attribute
+    {
+        /// <summary>
+        /// Endpoint name that is the input queue name.
+        /// </summary>
+        public string EndpointName { get; }
+
+        /// <summary>
+        /// Override trigger function name.
+        /// </summary>
+        public string TriggerFunctionName { get; set; }
+
+        /// <summary>
+        /// Enable cross-entity transactions.
+        /// </summary>
+        public bool EnableCrossEntityTransactions { get; set; }
+
+        /// <summary>
+        /// Endpoint logical name.
+        /// </summary>
+        /// <param name="endpointName">Endpoint name that is the input queue name.</param>
+        public NServiceBusTriggerFunctionAttribute(string endpointName)
+        {
+            EndpointName = endpointName;
+        }
+    }
+}

src/NServiceBus.AzureFunctions.InProcess.ServiceBus/ReflectionHelper.cs

@@ -0,0 +1,50 @@
+namespace NServiceBus.AzureFunctions.InProcess.ServiceBus
+{
+    using System;
+    using System.Diagnostics;
+    using System.Reflection;
+    using Microsoft.Azure.ServiceBus;
+    using Microsoft.Azure.WebJobs;
+
+    class ReflectionHelper
+    {
+        public static bool GetAutoCompleteValue()
+        {
+            var triggerAttribute = FindTriggerAttributeInternal();
+            if (triggerAttribute != null)
+            {
+                return triggerAttribute.AutoComplete;
+            }
+
+            throw new Exception($"Could not locate {nameof(ServiceBusTriggerAttribute)} to infer the AutoComplete setting. Make sure that the function trigger contains a parameter decorated with {nameof(ServiceBusTriggerAttribute)} or use the advanced APIs exposed via the {nameof(FunctionEndpoint)} type instead.");
+        }
+
+        public static ServiceBusTriggerAttribute FindBusTriggerAttribute() => FindTriggerAttributeInternal();
+
+        static ServiceBusTriggerAttribute FindTriggerAttributeInternal()
+        {
+            var st = new StackTrace(skipFrames: 2); // skip first two frames because it is this method + the public method
+            var frames = st.GetFrames();
+            foreach (var frame in frames)
+            {
+                var method = frame?.GetMethod();
+                if (method?.GetCustomAttribute<FunctionNameAttribute>(false) != null)
+                {
+                    foreach (var parameter in method.GetParameters())
+                    {
+                        ServiceBusTriggerAttribute serviceBusTriggerAttribute;
+                        if (parameter.ParameterType == typeof(Message)
+                            && (serviceBusTriggerAttribute = parameter.GetCustomAttribute<ServiceBusTriggerAttribute>(false)) != null)
+                        {
+                            return serviceBusTriggerAttribute;
+                        }
+                    }
+
+                    return null;
+                }
+            }
+
+            return null;
+        }
+    }
+}

src/NServiceBus.AzureFunctions.InProcess.ServiceBus/ServiceBusTriggeredEndpointConfiguration.cs

@@ -4,7 +4,6 @@
     using System.Threading.Tasks;
     using AzureFunctions.InProcess.ServiceBus;
     using Logging;
-    using Microsoft.Azure.WebJobs;
     using Microsoft.Extensions.Configuration;
     using Serialization;
     using Transport;
@@ -134,10 +133,10 @@ static string GetConfiguredValueOrFallback(IConfiguration configuration, string
         /// </summary>
         public static ServiceBusTriggeredEndpointConfiguration FromAttributes()
         {
-            var configuration = TriggerDiscoverer.TryGet<ServiceBusTriggerAttribute>();
-            if (configuration != null)
+            var serviceBusTriggerAttribute = ReflectionHelper.FindBusTriggerAttribute();
+            if (serviceBusTriggerAttribute != null)
             {
-                return new ServiceBusTriggeredEndpointConfiguration(configuration.QueueName, configuration.Connection);
+                return new ServiceBusTriggeredEndpointConfiguration(serviceBusTriggerAttribute.QueueName, serviceBusTriggerAttribute.Connection);
             }
 
             throw new Exception(

src/NServiceBus.AzureFunctions.InProcess.ServiceBus/Transactions/ITransactionStrategy.cs

@@ -0,0 +1,13 @@
+namespace NServiceBus
+{
+    using System.Threading.Tasks;
+    using System.Transactions;
+    using Transport;
+
+    interface ITransactionStrategy
+    {
+        CommittableTransaction CreateTransaction();
+        TransportTransaction CreateTransportTransaction(CommittableTransaction transaction);
+        Task Complete(CommittableTransaction transaction);
+    }
+}