Adding new logic for extended sessions to enable the DTS use-case (#1333)

* first commit?

* marked the return value as nullable

---------

Co-authored-by: Sophia Tevosyan <stevosyan@microsoft.com>

Author: sophiatev

src/DurableTask.AzureStorage/Messaging/OrchestrationSession.cs

@@ -19,7 +19,6 @@ namespace DurableTask.AzureStorage.Messaging
     using System.Linq;
     using System.Threading;
     using System.Threading.Tasks;
-    using Azure;
     using DurableTask.Core;
     using DurableTask.Core.History;
     using Newtonsoft.Json;
@@ -226,5 +225,11 @@ bool IsNonexistantInstance()
         {
             return this.RuntimeState.Events.Count == 0 || this.RuntimeState.ExecutionStartedEvent == null;
         }
+
+        public Task EndSessionAsync()
+        {
+            // No-op
+            return Task.CompletedTask;
+        }
     }
 }

src/DurableTask.Core/IOrchestrationSession.cs

@@ -11,6 +11,7 @@
 //  limitations under the License.
 //  ----------------------------------------------------------------------------------
 
+#nullable enable
 namespace DurableTask.Core
 {
     using System.Collections.Generic;
@@ -29,6 +30,12 @@ public interface IOrchestrationSession
         /// or until an internal wait period has expired. In either case, <c>null</c> can be returned
         /// and the dispatcher will shut down the session.
         /// </remarks>
-        Task<IList<TaskMessage>> FetchNewOrchestrationMessagesAsync(TaskOrchestrationWorkItem workItem);
+        Task<IList<TaskMessage>?> FetchNewOrchestrationMessagesAsync(TaskOrchestrationWorkItem workItem);
+
+        /// <summary>
+        /// Ends the session.
+        /// </summary>
+        /// <returns>A task that completes when the session has been ended.</returns>
+        Task EndSessionAsync();
     }
 }

src/DurableTask.Core/TaskEntityDispatcher.cs

@@ -193,6 +193,7 @@ async Task OnProcessWorkItemSessionAsync(TaskOrchestrationWorkItem workItem)
                     if (concurrencyLockAcquired)
                     {
                         this.concurrentSessionLock.Release();
+                        await workItem.Session.EndSessionAsync();
                     }
                 }
             }

src/DurableTask.Core/TaskOrchestrationDispatcher.cs

@@ -295,6 +295,7 @@ async Task OnProcessWorkItemSessionAsync(TaskOrchestrationWorkItem workItem)
                             "OnProcessWorkItemSession-Release",
                             $"Releasing extended session after {processCount} batch(es).");
                         this.concurrentSessionLock.Release();
+                        await workItem.Session.EndSessionAsync();
                     }
                 }
             }
@@ -552,10 +553,6 @@ protected async Task<bool> OnProcessWorkItemAsync(TaskOrchestrationWorkItem work
                                     orchestratorMessages.AddRange(subOrchestrationRewindMessages);
                                     workItem.OrchestrationRuntimeState = newRuntimeState;
                                     runtimeState = newRuntimeState;
-                                    // Setting this to true here will end an extended session if it is in progress.
-                                    // We don't want to save the state across executions, since we essentially manually modify
-                                    // the orchestration history here and so that stored by the extended session is incorrect.
-                                    isRewinding = true;
                                     break;
                                 default:
                                     throw TraceHelper.TraceExceptionInstance(