ServiceBusOrchestrationService.cs

//  ----------------------------------------------------------------------------------
//  Copyright Microsoft Corporation
//  Licensed under the Apache License, Version 2.0 (the "License");
//  you may not use this file except in compliance with the License.
//  You may obtain a copy of the License at
//  http://www.apache.org/licenses/LICENSE-2.0
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
//  ----------------------------------------------------------------------------------

namespace DurableTask.ServiceBus
{
    using System;
    using System.Collections.Concurrent;
    using System.Diagnostics;
    using System.Collections.Generic;
    using System.IO;
    using System.Linq;
    using System.Threading;
    using System.Threading.Tasks;
    using System.Transactions;
    using DurableTask.Core;
    using DurableTask.Core.Common;
    using DurableTask.Core.Exceptions;
    using DurableTask.Core.History;
    using DurableTask.Core.Tracing;
    using DurableTask.Core.Tracking;
    using DurableTask.Core.Serializing;
    using DurableTask.ServiceBus.Common.Abstraction;
    using DurableTask.ServiceBus.Settings;
    using DurableTask.ServiceBus.Stats;
    using DurableTask.ServiceBus.Tracking;
    using Message = DurableTask.ServiceBus.Common.Abstraction.Message;
    using IMessageSession = DurableTask.ServiceBus.Common.Abstraction.IMessageSession;
    using RetryPolicy = DurableTask.ServiceBus.Common.Abstraction.RetryPolicy;
    using MessageSender = DurableTask.ServiceBus.Common.Abstraction.MessageSender;
    using MessageReceiver = DurableTask.ServiceBus.Common.Abstraction.MessageReceiver;
    using QueueClient = DurableTask.ServiceBus.Common.Abstraction.QueueClient;
    using SessionClient = DurableTask.ServiceBus.Common.Abstraction.SessionClient;
    using ServiceBusConnection = DurableTask.ServiceBus.Common.Abstraction.ServiceBusConnection;
    using TokenProvider = DurableTask.ServiceBus.Common.Abstraction.TokenProvider;
    using ManagementClient = DurableTask.ServiceBus.Common.Abstraction.ManagementClient;
    using ServiceBusConnectionStringBuilder = DurableTask.ServiceBus.Common.Abstraction.ServiceBusConnectionStringBuilder;
#if NETSTANDARD2_0
    using Microsoft.Azure.ServiceBus;
    using Microsoft.Azure.ServiceBus.Management;
    using Microsoft.Azure.ServiceBus.Primitives;
#else
    using Microsoft.ServiceBus.Messaging;
#endif

    /// <summary>
    /// Orchestration Service and Client implementation using Azure Service Bus
    /// Takes an optional instance store for storing state and history
    /// </summary>
    public class ServiceBusOrchestrationService : IOrchestrationService, IOrchestrationServiceClient
    {
        // This is the Max number of messages which can be processed in a single transaction.
        // Current ServiceBus limit is 100 so it has to be lower than that.  
        // This also has an impact on prefetch count as PrefetchCount cannot be greater than this value
        // as every fetched message also creates a tracking message which counts towards this limit.
        const int MaxMessageCount = 80;
        const int SessionStreamWarningSizeInBytes = 150 * 1024;
        const int StatusPollingIntervalInSeconds = 2;
        const int DuplicateDetectionWindowInHours = 4;

        /// <summary>
        /// Orchestration service settings 
        /// </summary>
        public readonly ServiceBusOrchestrationServiceSettings Settings;

        /// <summary>
        /// Instance store for state and history tracking
        /// </summary>
        public readonly IOrchestrationServiceInstanceStore InstanceStore;

        /// <summary>
        /// Blob store for oversized messages and sessions
        /// </summary>
        public readonly IOrchestrationServiceBlobStore BlobStore;

        /// <summary>
        /// Statistics for the orchestration service
        /// </summary>
        public readonly ServiceBusOrchestrationServiceStats ServiceStats;

        static readonly DataConverter DataConverter = new JsonDataConverter();

        readonly ServiceBusConnectionSettings connectionSettings;

        readonly string hubName;

        MessageSender orchestratorSender;
        readonly MessageSender orchestrationBatchMessageSender;
        QueueClient orchestratorQueueClient;
        MessageSender workerSender;
        MessageSender trackingSender;
        SessionClient orchestratorSessionClient;
        MessageReceiver workerReceiver;
        SessionClient trackingClient;
        readonly string workerEntityName;
        readonly string orchestratorEntityName;
        readonly string trackingEntityName;
        readonly WorkItemDispatcher<TrackingWorkItem> trackingDispatcher;
        readonly JumpStartManager jumpStartManager;

        ConcurrentDictionary<string, ServiceBusOrchestrationSession> orchestrationSessions;
        ConcurrentDictionary<string, Message> orchestrationMessages;
        CancellationTokenSource cancellationTokenSource;

        ServiceBusConnection serviceBusConnection;

#if NETSTANDARD2_0

        /// <summary>
        ///     Create a new ServiceBusOrchestrationService to the given service bus namespace and hub name
        /// </summary>
        /// <param name="namespaceHostName">Service Bus namespace host name</param>
        /// <param name="tokenProvider">Service Bus authentication token provider</param>
        /// <param name="hubName">Hub name to use with the Service Bus namespace</param>
        /// <param name="instanceStore">Instance store Provider, where state and history messages will be stored</param>
        /// <param name="blobStore">Blob store Provider, where oversized messages and sessions will be stored</param>
        /// <param name="settings">Settings object for service and client</param>
        public ServiceBusOrchestrationService(
            string namespaceHostName,
            ITokenProvider tokenProvider,
            string hubName,
            IOrchestrationServiceInstanceStore instanceStore,
            IOrchestrationServiceBlobStore blobStore,
            ServiceBusOrchestrationServiceSettings settings) :
                this(
                      ServiceBusConnectionSettings.Create(namespaceHostName, tokenProvider),
                      hubName,
                      instanceStore,
                      blobStore,
                      settings)
        {
        }

#endif

        /// <summary>
        ///     Create a new ServiceBusOrchestrationService to the given service bus connection string and hub name
        /// </summary>
        /// <param name="connectionString">Service Bus connection string</param>
        /// <param name="hubName">Hub name to use with the connection string</param>
        /// <param name="instanceStore">Instance store Provider, where state and history messages will be stored</param>
        /// <param name="blobStore">Blob store Provider, where oversized messages and sessions will be stored</param>
        /// <param name="settings">Settings object for service and client</param>
        public ServiceBusOrchestrationService(
            string connectionString,
            string hubName,
            IOrchestrationServiceInstanceStore instanceStore,
            IOrchestrationServiceBlobStore blobStore,
            ServiceBusOrchestrationServiceSettings settings) :
                this(
                    ServiceBusConnectionSettings.Create(connectionString),
                    hubName,
                    instanceStore,
                    blobStore,
                    settings)
        {
        }

        /// <summary>
        ///     Create a new ServiceBusOrchestrationService to the given service bus connection and hub name
        /// </summary>
        /// <param name="connectionSettings">Service Bus connection settings</param>
        /// <param name="hubName">Hub name to use with the Service Bus namespace</param>
        /// <param name="instanceStore">Instance store Provider, where state and history messages will be stored</param>
        /// <param name="blobStore">Blob store Provider, where oversized messages and sessions will be stored</param>
        /// <param name="settings">Settings object for service and client</param>
        public ServiceBusOrchestrationService(
            ServiceBusConnectionSettings connectionSettings,
            string hubName,
            IOrchestrationServiceInstanceStore instanceStore,
            IOrchestrationServiceBlobStore blobStore,
            ServiceBusOrchestrationServiceSettings settings)
        {
            this.connectionSettings = connectionSettings;
            this.hubName = hubName;
            this.ServiceStats = new ServiceBusOrchestrationServiceStats();

            this.workerEntityName = string.Format(ServiceBusConstants.WorkerEndpointFormat, this.hubName);
            this.orchestratorEntityName = string.Format(ServiceBusConstants.OrchestratorEndpointFormat, this.hubName);
            this.trackingEntityName = string.Format(ServiceBusConstants.TrackingEndpointFormat, this.hubName);

            if (!string.IsNullOrEmpty(connectionSettings.ConnectionString))
            {
                var sbConnectionStringBuilder = new ServiceBusConnectionStringBuilder(connectionSettings.ConnectionString);
                this.serviceBusConnection = new ServiceBusConnection(sbConnectionStringBuilder)
                {
                    TokenProvider = TokenProvider.CreateSharedAccessSignatureTokenProvider(sbConnectionStringBuilder.SasKeyName,
                        sbConnectionStringBuilder.SasKey, ServiceBusUtils.TokenTimeToLive)
                };
            }
#if NETSTANDARD2_0
            else if (connectionSettings.Endpoint != null && connectionSettings.TokenProvider != null)
            {
                this.serviceBusConnection = new ServiceBusConnection(connectionSettings.Endpoint.ToString(), connectionSettings.TransportType)
                {
                    TokenProvider = connectionSettings.TokenProvider
                };
            }
#endif
            else
            {
                throw new ArgumentException("Invalid Service Bus connection settings.", nameof(connectionSettings));
            }

            this.Settings = settings ?? new ServiceBusOrchestrationServiceSettings();
            this.orchestrationBatchMessageSender = new MessageSender(this.serviceBusConnection, this.orchestratorEntityName);


            this.BlobStore = blobStore;
            if (instanceStore != null)
            {
                this.InstanceStore = instanceStore;
                this.trackingDispatcher = new WorkItemDispatcher<TrackingWorkItem>(
                    "TrackingDispatcher",
                    item => item == null ? string.Empty : item.InstanceId,
                    FetchTrackingWorkItemAsync,
                    ProcessTrackingWorkItemAsync)
                {
                    GetDelayInSecondsAfterOnFetchException = GetDelayInSecondsAfterOnFetchException,
                    GetDelayInSecondsAfterOnProcessException = GetDelayInSecondsAfterOnProcessException,
                    DispatcherCount = this.Settings.TrackingDispatcherSettings.DispatcherCount,
                    MaxConcurrentWorkItems = this.Settings.TrackingDispatcherSettings.MaxConcurrentTrackingSessions
                };

                if (this.Settings.JumpStartSettings.JumpStartEnabled)
                {
                    this.jumpStartManager = new JumpStartManager(this, this.Settings.JumpStartSettings.Interval, this.Settings.JumpStartSettings.IgnoreWindow);
                }
            }
        }

        /// <summary>
        /// Starts the service initializing the required resources
        /// </summary>
        public async Task StartAsync()
        {
            this.cancellationTokenSource = new CancellationTokenSource();
            this.orchestrationSessions = new ConcurrentDictionary<string, ServiceBusOrchestrationSession>();
            this.orchestrationMessages = new ConcurrentDictionary<string, Message>();

            this.orchestratorSender = new MessageSender(this.serviceBusConnection, this.orchestratorEntityName, this.workerEntityName);
            this.workerSender = new MessageSender(this.serviceBusConnection, this.workerEntityName, this.orchestratorEntityName);
            this.trackingSender = new MessageSender(this.serviceBusConnection, this.trackingEntityName, this.orchestratorEntityName);
            this.orchestratorQueueClient = new QueueClient(this.serviceBusConnection, this.orchestratorEntityName, ReceiveMode.PeekLock, RetryPolicy.Default);
            this.workerReceiver = new MessageReceiver(serviceBusConnection, this.workerEntityName);
            this.orchestratorSessionClient = new SessionClient(serviceBusConnection, this.orchestratorEntityName, ReceiveMode.PeekLock);
            this.trackingClient = new SessionClient(serviceBusConnection, this.trackingEntityName, ReceiveMode.PeekLock);

            if (this.trackingDispatcher != null)
            {
                await this.trackingDispatcher.StartAsync();
            }

            if (this.jumpStartManager != null)
            {
                await this.jumpStartManager.StartAsync();
            }

            await Task.Factory.StartNew(() => ServiceMonitorAsync(this.cancellationTokenSource.Token), this.cancellationTokenSource.Token);
        }

        /// <summary>
        /// Stops the orchestration service gracefully
        /// </summary>
        public async Task StopAsync()
        {
            await StopAsync(false);
        }

        /// <summary>
        /// Stops the orchestration service with optional forced flag
        /// </summary>
        /// <param name="isForced">Flag when true stops resources aggressively, when false stops gracefully</param>
        public async Task StopAsync(bool isForced)
        {
            this.cancellationTokenSource?.Cancel();

            TraceHelper.Trace(TraceEventType.Information, "ServiceBusOrchestrationService-StatsFinal", "Final Service Stats: {0}", this.ServiceStats.ToString());
            // TODO : call shutdown of any remaining orchestrationSessions and orchestrationMessages

            await Task.WhenAll(
                this.workerSender.CloseAsync(),
                this.orchestratorSender.CloseAsync(),
                this.orchestrationBatchMessageSender?.CloseAsync(),
                this.trackingSender.CloseAsync(),
                this.orchestratorSessionClient.CloseAsync(),
                this.trackingClient.CloseAsync(),
                this.workerReceiver.CloseAsync()
            );
            if (this.trackingDispatcher != null)
            {
                await this.trackingDispatcher.StopAsync(isForced);
            }

            if (this.jumpStartManager != null)
            {
                await this.jumpStartManager.StopAsync();
            }
        }

        /// <summary>
        /// Deletes and creates the necessary resources for the orchestration service including the instance store
        /// </summary>
        public Task CreateAsync()
        {
            return CreateAsync(true);
        }

        /// <summary>
        /// Deletes and creates the necessary resources for the orchestration service
        /// </summary>
        /// <param name="recreateInstanceStore">Flag indicating whether to drop and create instance store</param>
        public async Task CreateAsync(bool recreateInstanceStore)
        {
            ManagementClient managementClient = this.CreateManagementClient();

            await Task.WhenAll(
                SafeDeleteAndCreateQueueAsync(managementClient, this.orchestratorEntityName, true, true, this.Settings.MaxTaskOrchestrationDeliveryCount, this.Settings.MaxQueueSizeInMegabytes),
                SafeDeleteAndCreateQueueAsync(managementClient, this.workerEntityName, false, false, this.Settings.MaxTaskActivityDeliveryCount, this.Settings.MaxQueueSizeInMegabytes)
            );

            if (this.InstanceStore != null)
            {
                await SafeDeleteAndCreateQueueAsync(managementClient, this.trackingEntityName, true, false, this.Settings.MaxTrackingDeliveryCount, this.Settings.MaxQueueSizeInMegabytes);
                await this.InstanceStore.InitializeStoreAsync(recreateInstanceStore);
            }
        }

        /// <summary>
        /// Drops and creates the necessary resources for the orchestration service and the instance store
        /// </summary>
        public async Task CreateIfNotExistsAsync()
        {
            ManagementClient managementClient = this.CreateManagementClient();

            await Task.WhenAll(
                SafeCreateQueueAsync(managementClient, this.orchestratorEntityName, true, true, this.Settings.MaxTaskOrchestrationDeliveryCount, this.Settings.MaxQueueSizeInMegabytes),
                SafeCreateQueueAsync(managementClient, this.workerEntityName, false, false, this.Settings.MaxTaskActivityDeliveryCount, this.Settings.MaxQueueSizeInMegabytes)
            );
            if (this.InstanceStore != null)
            {
                await SafeCreateQueueAsync(managementClient, this.trackingEntityName, true, false, this.Settings.MaxTrackingDeliveryCount, this.Settings.MaxQueueSizeInMegabytes);
                await this.InstanceStore.InitializeStoreAsync(false);
            }
        }

        /// <summary>
        /// Deletes the resources for the orchestration service and the instance store
        /// </summary>
        public Task DeleteAsync()
        {
            return DeleteAsync(true);
        }

        /// <summary>
        /// Deletes the resources for the orchestration service and optionally the instance store
        /// </summary>
        /// <param name="deleteInstanceStore">Flag indicating whether to drop instance store</param>
        public async Task DeleteAsync(bool deleteInstanceStore)
        {
            ManagementClient managementClient = this.CreateManagementClient();

            await Task.WhenAll(
                SafeDeleteQueueAsync(managementClient, this.orchestratorEntityName),
                SafeDeleteQueueAsync(managementClient, this.workerEntityName)
            );
            if (this.InstanceStore != null)
            {
                await SafeDeleteQueueAsync(managementClient, this.trackingEntityName);
                if (deleteInstanceStore)
                {
                    await this.InstanceStore.DeleteStoreAsync();
                }
            }

            if (this.BlobStore != null)
            {
                await this.BlobStore.DeleteStoreAsync();
            }
        }

        // Service Bus Utility methods

        /// <summary>
        /// Utility method to check if the needed resources are available for the orchestration service
        /// </summary>
        /// <returns>True if all needed queues are present, false otherwise</returns>
        public async Task<bool> HubExistsAsync()
        {
            ManagementClient managementClient = this.CreateManagementClient();

            IEnumerable<QueueDescription> queueDescriptions = (await managementClient.GetQueuesAsync()).Where(x => x.Path.StartsWith(this.hubName)).ToList();

            return queueDescriptions.Any(q => string.Equals(q.Path, this.orchestratorEntityName))
                && queueDescriptions.Any(q => string.Equals(q.Path, this.workerEntityName))
                && (this.InstanceStore == null || queueDescriptions.Any(q => string.Equals(q.Path, this.trackingEntityName)));
        }

        /// <summary>
        ///     Get the count of pending orchestrations
        /// </summary>
        /// <returns>Count of pending orchestrations</returns>
        public async Task<long> GetPendingOrchestrationsCount()
        {
            return await GetQueueCount(this.orchestratorEntityName);
        }

        /// <summary>
        ///     Get the count of pending work items (activities)
        /// </summary>
        /// <returns>Count of pending activities</returns>
        public async Task<long> GetPendingWorkItemsCount()
        {
            return await GetQueueCount(this.workerEntityName);
        }

        /// <summary>
        ///     Internal method for getting the number of items in a queue
        /// </summary>
        async Task<long> GetQueueCount(string entityName)
        {
            ManagementClient managementClient = this.CreateManagementClient();
            var queueDescription = await managementClient.GetQueueRuntimeInfoAsync(entityName);
            if (queueDescription == null)
            {
                throw TraceHelper.TraceException(
                    TraceEventType.Error,
                    "ServiceBusOrchestrationService-QueueNotFound",
                    new ArgumentException($"Queue {entityName} does not exist"));
            }

            return queueDescription.MessageCount;
        }

        /// <summary>
        ///     Internal method for getting the max delivery counts for each queue
        /// </summary>
        internal async Task<Dictionary<string, int>> GetHubQueueMaxDeliveryCountsAsync()
        {
            ManagementClient managementClient = this.CreateManagementClient();

            var result = new Dictionary<string, int>(3);

            IEnumerable<QueueDescription> queues =
                (await managementClient.GetQueuesAsync()).Where(x => x.Path.StartsWith(this.hubName)).ToList();

            result.Add("TaskOrchestration", queues.Single(q => string.Equals(q.Path, this.orchestratorEntityName))?.MaxDeliveryCount ?? -1);
            result.Add("TaskActivity", queues.Single(q => string.Equals(q.Path, this.workerEntityName))?.MaxDeliveryCount ?? -1);
            result.Add("Tracking", queues.Single(q => string.Equals(q.Path, this.trackingEntityName))?.MaxDeliveryCount ?? -1);

            return result;
        }

        /// <summary>
        ///     Checks the message count against the threshold to see if a limit is being exceeded
        /// </summary>
        /// <param name="currentMessageCount">The current message count to check</param>
        /// <param name="runtimeState">The Orchestration runtime state this message count is associated with</param>
        public bool IsMaxMessageCountExceeded(int currentMessageCount, OrchestrationRuntimeState runtimeState)
        {
            return currentMessageCount
                + ((this.InstanceStore != null) ? runtimeState.NewEvents.Count + 1 : 0) // one history message per new message + 1 for the orchestration
                > MaxMessageCount;
        }

        /// <summary>
        /// Inspects an exception to get a custom delay based on the exception (e.g. transient) properties for a process exception
        /// </summary>
        /// <param name="exception">The exception to inspect</param>
        /// <returns>Delay in seconds</returns>
        public int GetDelayInSecondsAfterOnProcessException(Exception exception)
        {
            if (IsTransientException(exception))
            {
                return this.Settings.TaskOrchestrationDispatcherSettings.TransientErrorBackOffSecs;
            }

            return 0;
        }

        /// <summary>
        /// Inspects an exception to get a custom delay based on the exception (e.g. transient) properties for a fetch exception
        /// </summary>
        /// <param name="exception">The exception to inspect</param>
        /// <returns>Delay in seconds</returns>
        public int GetDelayInSecondsAfterOnFetchException(Exception exception)
        {
            if (exception is TimeoutException)
            {
                return 0;
            }

            int delay = this.Settings.TaskOrchestrationDispatcherSettings.NonTransientErrorBackOffSecs;
            if (IsTransientException(exception))
            {
                delay = this.Settings.TaskOrchestrationDispatcherSettings.TransientErrorBackOffSecs;
            }

            return delay;
        }

        /// <summary>
        /// Gets the the number of task orchestration dispatchers
        /// </summary>
        public int TaskOrchestrationDispatcherCount => this.Settings.TaskOrchestrationDispatcherSettings.DispatcherCount;

        /// <summary>
        /// Gets the maximum number of concurrent task orchestration items
        /// </summary>
        public int MaxConcurrentTaskOrchestrationWorkItems => this.Settings.TaskOrchestrationDispatcherSettings.MaxConcurrentOrchestrations;

        /// <summary>
        ///  Should we carry over unexecuted raised events to the next iteration of an orchestration on ContinueAsNew
        /// </summary>
        public BehaviorOnContinueAsNew EventBehaviourForContinueAsNew => this.Settings.TaskOrchestrationDispatcherSettings.EventBehaviourForContinueAsNew;

        /// <summary>
        ///     Wait for the next orchestration work item and return the orchestration work item
        /// </summary>
        /// <param name="receiveTimeout">The timespan to wait for new messages before timing out</param>
        /// <param name="cancellationToken">The cancellation token to cancel execution of the task</param>
        public async Task<TaskOrchestrationWorkItem> LockNextTaskOrchestrationWorkItemAsync(TimeSpan receiveTimeout, CancellationToken cancellationToken)
        {
            var session = await this.orchestratorSessionClient.AcceptMessageSessionAsync(receiveTimeout);
            if (session == null)
            {
                return null;
            }

            this.ServiceStats.OrchestrationDispatcherStats.SessionsReceived.Increment();

            // TODO : Here and elsewhere, consider standard retry block instead of our own hand rolled version
            IList<Message> newMessages =
                (await Utils.ExecuteWithRetries(() => session.ReceiveAsync(this.Settings.PrefetchCount),
                    session.SessionId, "Receive Session Message Batch", this.Settings.MaxRetries, this.Settings.IntervalBetweenRetriesSecs)).Cast<Message>().ToList();

            this.ServiceStats.OrchestrationDispatcherStats.MessagesReceived.Increment(newMessages.Count);
            TraceHelper.TraceSession(
                TraceEventType.Information,
                "ServiceBusOrchestrationService-LockNextTaskOrchestrationWorkItem-MessageToProcess",
                session.SessionId,
                GetFormattedLog(
                    $@"{newMessages.Count} new messages to process: {
                            string.Join(",", newMessages.Select(m => m.MessageId))}, max latency: {
                            newMessages.Max(message => message.DeliveryLatency())}ms"));

            ServiceBusUtils.CheckAndLogDeliveryCount(session.SessionId, newMessages, this.Settings.MaxTaskOrchestrationDeliveryCount);

            IList<TaskMessage> newTaskMessages = await Task.WhenAll(
                newMessages.Select(async message => await ServiceBusUtils.GetObjectFromBrokeredMessageAsync<TaskMessage>(message, this.BlobStore)));

            OrchestrationRuntimeState runtimeState = await GetSessionStateAsync(session, this.BlobStore);

            long maxSequenceNumber = newMessages.Max(message => message.SystemProperties.SequenceNumber);

            Dictionary<string, Message> lockTokens = newMessages.ToDictionary(m => m.SystemProperties.LockToken.ToString(), m => m);
            var sessionState = new ServiceBusOrchestrationSession
            {
                Session = session,
                LockTokens = lockTokens,
                SequenceNumber = maxSequenceNumber
            };

            if (!this.orchestrationSessions.TryAdd(session.SessionId, sessionState))
            {
                string error = $"Duplicate orchestration session id '{session.SessionId}', id already exists in session list.";
                TraceHelper.Trace(TraceEventType.Error, "ServiceBusOrchestrationService-LockNextTaskOrchestrationWorkItem-DuplicateSessionId", error);
                throw new OrchestrationFrameworkException(error);
            }

            if (this.InstanceStore != null)
            {
                try
                {
                    TaskMessage executionStartedMessage = newTaskMessages.FirstOrDefault(m => m.Event is ExecutionStartedEvent);

                    if (executionStartedMessage != null)
                    {
                        await UpdateInstanceStoreAsync(executionStartedMessage.Event as ExecutionStartedEvent, maxSequenceNumber);
                    }
                }
                catch (Exception exception)
                {
                    this.orchestrationSessions.TryRemove(session.SessionId, out ServiceBusOrchestrationSession _);

                    string error = $"Exception while updating instance store. Session id: {session.SessionId}";
                    TraceHelper.TraceException(TraceEventType.Error, "ServiceBusOrchestrationService-LockNextTaskOrchestrationWorkItem-ErrorUpdatingInstanceStore", exception, error);

                    throw;
                }
            }

            return new TaskOrchestrationWorkItem
            {
                InstanceId = session.SessionId,
                LockedUntilUtc = session.LockedUntilUtc,
                NewMessages = newTaskMessages.ToList(),
                OrchestrationRuntimeState = runtimeState
            };
        }

        Task UpdateInstanceStoreAsync(ExecutionStartedEvent executionStartedEvent, long sequenceNumber)
        {
            // TODO: Duplicate detection: Check if the orchestration already finished

            var orchestrationState = new OrchestrationState()
            {
                Name = executionStartedEvent.Name,
                Version = executionStartedEvent.Version,
                OrchestrationInstance = executionStartedEvent.OrchestrationInstance,
                OrchestrationStatus = OrchestrationStatus.Pending,
                Input = executionStartedEvent.Input,
                Tags = executionStartedEvent.Tags,
                CreatedTime = executionStartedEvent.Timestamp,
                LastUpdatedTime = DateTime.UtcNow,
                CompletedTime = DateTimeUtils.MinDateTime,
                ParentInstance = executionStartedEvent.ParentInstance,
                ScheduledStartTime = executionStartedEvent.ScheduledStartTime
            };

            var orchestrationStateEntity = new OrchestrationStateInstanceEntity
            {
                State = orchestrationState,
                SequenceNumber = sequenceNumber
            };

            return this.InstanceStore.WriteEntitiesAsync(new[] { orchestrationStateEntity });
        }

        ServiceBusOrchestrationSession GetSessionInstanceForWorkItem(TaskOrchestrationWorkItem workItem)
        {
            if (string.IsNullOrWhiteSpace(workItem?.InstanceId))
            {
                return null;
            }

            return this.orchestrationSessions[workItem.InstanceId];
        }

        ServiceBusOrchestrationSession GetAndDeleteSessionInstanceForWorkItem(TaskOrchestrationWorkItem workItem)
        {
            if (string.IsNullOrWhiteSpace(workItem?.InstanceId))
            {
                return null;
            }

            this.orchestrationSessions.TryRemove(workItem.InstanceId, out ServiceBusOrchestrationSession sessionInstance);
            return sessionInstance;
        }

        /// <summary>
        ///     Renew the lock on an orchestration
        /// </summary>
        /// <param name="workItem">The task orchestration to renew the lock on</param>
        public async Task RenewTaskOrchestrationWorkItemLockAsync(TaskOrchestrationWorkItem workItem)
        {
            ServiceBusOrchestrationSession sessionState = GetSessionInstanceForWorkItem(workItem);
            if (sessionState?.Session == null)
            {
                return;
            }

            TraceHelper.TraceSession(TraceEventType.Information, "ServiceBusOrchestrationService-RenewTaskOrchestrationWorkItem", workItem.InstanceId, "Renew lock on orchestration session");
            await sessionState.Session.RenewSessionLockAsync();
            this.ServiceStats.OrchestrationDispatcherStats.SessionsRenewed.Increment();
            workItem.LockedUntilUtc = sessionState.Session.LockedUntilUtc;
        }

        /// <summary>
        ///     Complete an orchestration, this atomically sends any outbound messages and completes the session for all current messages
        /// </summary>
        /// <param name="workItem">The task orchestration to renew the lock on</param>
        /// <param name="newOrchestrationRuntimeState">New state of the orchestration to be persisted. Could be null if the orchestration is in completion.</param>
        /// <param name="outboundMessages">New work item messages to be processed</param>
        /// <param name="orchestratorMessages">New orchestration messages to be scheduled</param>
        /// <param name="timerMessages">Delayed execution messages to be scheduled for the orchestration</param>
        /// <param name="continuedAsNewMessage">Task Message to send to orchestrator queue to treat as new in order to rebuild state</param>
        /// <param name="orchestrationState">The prior orchestration state</param>
        public async Task CompleteTaskOrchestrationWorkItemAsync(
            TaskOrchestrationWorkItem workItem,
            OrchestrationRuntimeState newOrchestrationRuntimeState,
            IList<TaskMessage> outboundMessages,
            IList<TaskMessage> orchestratorMessages,
            IList<TaskMessage> timerMessages,
            TaskMessage continuedAsNewMessage,
            OrchestrationState orchestrationState)
        {
            OrchestrationRuntimeState runtimeState = workItem.OrchestrationRuntimeState;
            ServiceBusOrchestrationSession sessionState = GetSessionInstanceForWorkItem(workItem);
            if (sessionState == null)
            {
                // ReSharper disable once NotResolvedInText
                throw new ArgumentNullException("SessionInstance");
            }

            var session = sessionState.Session;

            using (var ts = new TransactionScope(TransactionScopeAsyncFlowOption.Enabled))
            {
                Transaction.Current.TransactionCompleted += (o, e) =>
                    TraceHelper.TraceInstance(
                        e.Transaction.TransactionInformation.Status == TransactionStatus.Committed ? TraceEventType.Information : TraceEventType.Error,
                        "ServiceBusOrchestrationService-CompleteTaskOrchestrationWorkItem-TransactionComplete",
                        runtimeState.OrchestrationInstance,
                        () => $@"Orchestration Transaction Completed {
                                e.Transaction.TransactionInformation.LocalIdentifier
                            } status: {
                                e.Transaction.TransactionInformation.Status}");

                TraceHelper.TraceInstance(
                    TraceEventType.Information,
                    "ServiceBusOrchestrationService-CompleteTaskOrchestrationWorkItem-CreateTransaction",
                    runtimeState.OrchestrationInstance,
                    () => $@"Created new Orchestration Transaction - txnid: {
                            Transaction.Current.TransactionInformation.LocalIdentifier
                        }");

                if (await TrySetSessionStateAsync(workItem, newOrchestrationRuntimeState, runtimeState, session))
                {
                    if (outboundMessages?.Count > 0)
                    {
                        MessageContainer[] outboundBrokeredMessages = await Task.WhenAll(outboundMessages.Select(async m =>
                                {
                                    Message message = await ServiceBusUtils.GetBrokeredMessageFromObjectAsync(
                                        m,
                                        this.Settings.MessageCompressionSettings,
                                        this.Settings.MessageSettings,
                                        null,
                                        "Worker outbound message",
                                        this.BlobStore,
                                        DateTimeUtils.MinDateTime);
                                    return new MessageContainer(message, m);
                                }));
                        await this.workerSender.SendAsync(outboundBrokeredMessages.Select(m => m.Message).ToList());
                        LogSentMessages(session, "Worker outbound", outboundBrokeredMessages);
                        this.ServiceStats.ActivityDispatcherStats.MessageBatchesSent.Increment();
                        this.ServiceStats.ActivityDispatcherStats.MessagesSent.Increment(outboundMessages.Count);
                    }

                    if (timerMessages?.Count > 0 && newOrchestrationRuntimeState != null)
                    {
                        MessageContainer[] timerBrokeredMessages = await Task.WhenAll(timerMessages.Select(async m =>
                                {
                                    DateTime messageFireTime = ((TimerFiredEvent) m.Event).FireAt;
                                    Message message = await ServiceBusUtils.GetBrokeredMessageFromObjectAsync(
                                        m,
                                        this.Settings.MessageCompressionSettings,
                                        this.Settings.MessageSettings,
                                        newOrchestrationRuntimeState.OrchestrationInstance,
                                        "Timer Message",
                                        this.BlobStore,
                                        messageFireTime);
                                    message.ScheduledEnqueueTimeUtc = messageFireTime;
                                    return new MessageContainer(message, m);
                                }));

                        await this.orchestratorQueueClient.SendAsync(timerBrokeredMessages.Select(m => m.Message).ToList());
                        LogSentMessages(session, "Timer Message", timerBrokeredMessages);
                        this.ServiceStats.OrchestrationDispatcherStats.MessageBatchesSent.Increment();
                        this.ServiceStats.OrchestrationDispatcherStats.MessagesSent.Increment(timerMessages.Count);
                    }

                    if (orchestratorMessages?.Count > 0)
                    {
                        MessageContainer[] orchestrationBrokeredMessages = await Task.WhenAll(orchestratorMessages.Select(async m =>
                                {
                                    Message message = await ServiceBusUtils.GetBrokeredMessageFromObjectAsync(
                                        m,
                                        this.Settings.MessageCompressionSettings,
                                        this.Settings.MessageSettings,
                                        m.OrchestrationInstance,
                                        "Sub Orchestration",
                                        this.BlobStore,
                                        DateTimeUtils.MinDateTime);
                                    return new MessageContainer(message, m);
                                }));
                        await this.orchestratorQueueClient.SendAsync(orchestrationBrokeredMessages.Select(m => m.Message).ToList());

                        LogSentMessages(session, "Sub Orchestration", orchestrationBrokeredMessages);
                        this.ServiceStats.OrchestrationDispatcherStats.MessageBatchesSent.Increment();
                        this.ServiceStats.OrchestrationDispatcherStats.MessagesSent.Increment(orchestratorMessages.Count);
                    }

                    if (continuedAsNewMessage != null)
                    {
                        Message continuedAsNewBrokeredMessage = await ServiceBusUtils.GetBrokeredMessageFromObjectAsync(
                            continuedAsNewMessage,
                            this.Settings.MessageCompressionSettings,
                            this.Settings.MessageSettings,
                            newOrchestrationRuntimeState?.OrchestrationInstance,
                            "Continue as new",
                            this.BlobStore,
                            DateTimeUtils.MinDateTime);
                        await this.orchestratorQueueClient.SendAsync(continuedAsNewBrokeredMessage);
                        LogSentMessages(session, "Continue as new", new List<MessageContainer> { new MessageContainer(continuedAsNewBrokeredMessage, null) });
                        this.ServiceStats.OrchestrationDispatcherStats.MessageBatchesSent.Increment();
                        this.ServiceStats.OrchestrationDispatcherStats.MessagesSent.Increment();
                    }

                    if (this.InstanceStore != null)
                    {
                        List<MessageContainer> trackingMessages = await CreateTrackingMessagesAsync(runtimeState, sessionState.SequenceNumber);
                        TraceHelper.TraceInstance(
                            TraceEventType.Information,
                            "ServiceBusOrchestrationService-CompleteTaskOrchestrationWorkItem-TrackingMessages",
                            runtimeState.OrchestrationInstance,
                            "Created {0} tracking messages", trackingMessages.Count);

                        if (trackingMessages.Count > 0)
                        {
                            await this.trackingSender.SendAsync(trackingMessages.Select(m => m.Message).ToList());
                            LogSentMessages(session, "Tracking messages", trackingMessages);
                            this.ServiceStats.TrackingDispatcherStats.MessageBatchesSent.Increment();
                            this.ServiceStats.TrackingDispatcherStats.MessagesSent.Increment(trackingMessages.Count);
                        }

                        if (newOrchestrationRuntimeState != null && runtimeState != newOrchestrationRuntimeState)
                        {
                            trackingMessages = await CreateTrackingMessagesAsync(newOrchestrationRuntimeState, sessionState.SequenceNumber);
                            TraceHelper.TraceInstance(
                                TraceEventType.Information,
                                "ServiceBusOrchestrationService-CompleteTaskOrchestrationWorkItem-TrackingMessages",
                                newOrchestrationRuntimeState.OrchestrationInstance,
                                "Created {0} tracking messages", trackingMessages.Count);

                            if (trackingMessages.Count > 0)
                            {
                                await this.trackingSender.SendAsync(trackingMessages.Select(m => m.Message).ToList());
                                LogSentMessages(session, "Tracking messages", trackingMessages);
                                this.ServiceStats.TrackingDispatcherStats.MessageBatchesSent.Increment();
                                this.ServiceStats.TrackingDispatcherStats.MessagesSent.Increment(trackingMessages.Count);
                            }
                        }
                    }
                }

                TraceHelper.TraceInstance(
                    TraceEventType.Information,
                    "ServiceBusOrchestrationService-CompleteTaskOrchestrationWorkItemMessages",
                    runtimeState.OrchestrationInstance,
                    () =>
                    {
                        string allIds = string.Join(" ", sessionState.LockTokens.Values.Select(m => $"[SEQ: {m.SystemProperties.SequenceNumber} LT: {m.SystemProperties.LockToken}]"));
                        return $"Completing orchestration messages sequence and lock tokens: {allIds}";
                    });

                await session.CompleteAsync(sessionState.LockTokens.Keys);
                this.ServiceStats.OrchestrationDispatcherStats.SessionBatchesCompleted.Increment();
                ts.Complete();
            }
        }

        /// <summary>
        ///     Release the lock on an orchestration, releases the session, decoupled from CompleteTaskOrchestrationWorkItemAsync to handle nested orchestrations
        /// </summary>
        /// <param name="workItem">The task orchestration to abandon</param>
        public async Task ReleaseTaskOrchestrationWorkItemAsync(TaskOrchestrationWorkItem workItem)
        {
            ServiceBusOrchestrationSession sessionState = GetAndDeleteSessionInstanceForWorkItem(workItem);
            // This is Ok, if we abandoned the message it will already be gone
            if (sessionState == null)
            {
                TraceHelper.TraceSession(
                    TraceEventType.Warning,
                    "ServiceBusOrchestrationService-ReleaseTaskOrchestrationWorkItemFailed",
                    workItem?.InstanceId,
                    "DeleteSessionInstance failed, could already be aborted");
                return;
            }

            await sessionState.Session.CloseAsync();
        }

        /// <summary>
        ///     Abandon an orchestration, this abandons ownership/locking of all messages for an orchestration and it's session
        /// </summary>
        /// <param name="workItem">The task orchestration to abandon</param>
        public async Task AbandonTaskOrchestrationWorkItemAsync(TaskOrchestrationWorkItem workItem)
        {
            ServiceBusOrchestrationSession sessionState = GetAndDeleteSessionInstanceForWorkItem(workItem);
            if (sessionState?.Session == null)
            {
                return;
            }

            TraceHelper.TraceSession(TraceEventType.Error, "ServiceBusOrchestrationService-AbandonTaskOrchestrationWorkItem", workItem.InstanceId, "Abandoning {0} messages due to work item abort", sessionState.LockTokens.Keys.Count());
            foreach (string lockToken in sessionState.LockTokens.Keys)
            {
                await sessionState.Session.AbandonAsync(lockToken);
            }

            try
            {
                await sessionState.Session.CloseAsync();
            }
            catch (Exception ex) when (!Utils.IsFatal(ex))
            {
                TraceHelper.TraceExceptionSession(TraceEventType.Warning, "ServiceBusOrchestrationService-AbandonTaskOrchestrationWorkItemError", workItem.InstanceId, ex, "Error while aborting session");
            }
        }

        /// <summary>
        /// Gets the the number of task activity dispatchers
        /// </summary>
        public int TaskActivityDispatcherCount => this.Settings.TaskActivityDispatcherSettings.DispatcherCount;

        /// <summary>
        /// Gets the maximum number of concurrent task activity items
        /// </summary>
        public int MaxConcurrentTaskActivityWorkItems => this.Settings.TaskActivityDispatcherSettings.MaxConcurrentActivities;

        /// <summary>
        ///    Wait for an lock the next task activity to be processed 
        /// </summary>
        /// <param name="receiveTimeout">The timespan to wait for new messages before timing out</param>
        /// <param name="cancellationToken">The cancellation token to cancel execution of the task</param>
        public async Task<TaskActivityWorkItem> LockNextTaskActivityWorkItem(TimeSpan receiveTimeout, CancellationToken cancellationToken)
        {
            Message receivedMessage = (Message)await this.workerReceiver.ReceiveAsync(receiveTimeout);
            if (receivedMessage == null)
            {
                return null;
            }

            this.ServiceStats.ActivityDispatcherStats.MessagesReceived.Increment();

            TraceHelper.TraceSession(
                TraceEventType.Information,
                "ServiceBusOrchestrationService-LockNextTaskActivityWorkItem-Messages",
                receivedMessage.SessionId,
                GetFormattedLog($"New message to process: {receivedMessage.MessageId} [{receivedMessage.SystemProperties.SequenceNumber}], latency: {receivedMessage.DeliveryLatency()}ms"));

            TaskMessage taskMessage = await ServiceBusUtils.GetObjectFromBrokeredMessageAsync<TaskMessage>(receivedMessage, this.BlobStore);

            ServiceBusUtils.CheckAndLogDeliveryCount(receivedMessage, this.Settings.MaxTaskActivityDeliveryCount);

            if (!this.orchestrationMessages.TryAdd(receivedMessage.MessageId, receivedMessage))
            {
                string error = $"Duplicate orchestration message id '{receivedMessage.MessageId}', id already exists in message list.";
                TraceHelper.Trace(TraceEventType.Error, "ServiceBusOrchestrationService-DuplicateOrchestration", error);
                throw new OrchestrationFrameworkException(error);
            }

            return new TaskActivityWorkItem
            {
                Id = receivedMessage.MessageId,
                LockedUntilUtc = receivedMessage.SystemProperties.LockedUntilUtc,
                TaskMessage = taskMessage
            };
        }

        Message GetBrokeredMessageForWorkItem(TaskActivityWorkItem workItem)
        {
            if (string.IsNullOrWhiteSpace(workItem?.Id))
            {
                return null;
            }

            this.orchestrationMessages.TryGetValue(workItem.Id, out Message message);
            return message;
        }

        Message GetAndDeleteBrokeredMessageForWorkItem(TaskActivityWorkItem workItem)
        {
            if (string.IsNullOrWhiteSpace(workItem?.Id))
            {
                return null;
            }

            this.orchestrationMessages.TryRemove(workItem.Id, out Message existingMessage);
            return existingMessage;
        }

        /// <summary>
        ///    Renew the lock on a still processing work item
        /// </summary>
        /// <param name="workItem">Work item to renew the lock on</param>
        public async Task<TaskActivityWorkItem> RenewTaskActivityWorkItemLockAsync(TaskActivityWorkItem workItem)
        {
            Message message = GetBrokeredMessageForWorkItem(workItem);

            if (message != null)
            {
                await this.workerReceiver.RenewLockAsync(message);
                workItem.LockedUntilUtc = message.SystemProperties.LockedUntilUtc;
                this.ServiceStats.ActivityDispatcherStats.SessionsRenewed.Increment();
            }

            return workItem;
        }