Extended sessions for entities in .NET isolated

src/WebJobs.Extensions.DurableTask/Bindings/EntityTriggerAttributeBindingProvider.cs

@@ -130,7 +130,7 @@ public Task<ITriggerData> BindAsync(object value, ValueBindingContext context)
                 {
                     // Generate a byte array which is the serialized protobuf payload
                     // https://developers.google.com/protocol-buffers/docs/csharptutorial#parsing_and_serialization
-                    var entityBatchRequest = remoteContext.Request.ToEntityBatchRequest();
+                    var entityBatchRequest = remoteContext.Request.ToEntityBatchRequest(remoteContext.Configurations);
 
                     // We convert the binary payload into a base64 string because that seems to be the most commonly supported
                     // format for Azure Functions language workers. Attempts to send unencoded byte[] payloads were unsuccessful.

src/WebJobs.Extensions.DurableTask/Bindings/OrchestrationTriggerAttributeBindingProvider.cs

@@ -170,7 +170,7 @@ public Task<ITriggerData> BindAsync(object? value, ValueBindingContext context)
                     var orchestratorRequest = new Microsoft.DurableTask.Protobuf.OrchestratorRequest()
                     {
                         InstanceId = remoteContext.InstanceId,
-                        PastEvents = { remoteContext.Configurations.IncludePastEvents ? remoteContext.PastEvents.Select(ProtobufUtils.ToHistoryEventProto) : Enumerable.Empty<Microsoft.DurableTask.Protobuf.HistoryEvent>() },
+                        PastEvents = { remoteContext.Configurations.IncludeState ? remoteContext.PastEvents.Select(ProtobufUtils.ToHistoryEventProto) : Enumerable.Empty<Microsoft.DurableTask.Protobuf.HistoryEvent>() },
                         NewEvents = { remoteContext.NewEvents.Select(ProtobufUtils.ToHistoryEventProto) },
                         EntityParameters = remoteContext.EntityParameters.ToProtobuf(),
                     };

src/WebJobs.Extensions.DurableTask/ContextImplementations/RemoteEntityContext.cs

@@ -5,19 +5,32 @@
 using System.Collections.Generic;
 using DurableTask.Core.Entities;
 using DurableTask.Core.Entities.OperationFormat;
+using Microsoft.Extensions.Options;
 using Newtonsoft.Json;
 
 namespace Microsoft.Azure.WebJobs.Extensions.DurableTask
 {
     internal class RemoteEntityContext
     {
-        public RemoteEntityContext(EntityBatchRequest batchRequest)
+        public RemoteEntityContext(EntityBatchRequest batchRequest, DurableTaskOptions options, bool isExtendedSession, bool includeEntityState)
         {
-            this.Request = batchRequest;
+            this.Request = batchRequest;
+            if (options.ExtendedSessionsEnabled)
+            {
+                this.Configurations = new RemoteInstanceConfiguration
+                {
+                    IsExtendedSession = isExtendedSession,
+                    IncludeState = includeEntityState,
+                    ExtendedSessionIdleTimeoutInSeconds = options.ExtendedSessionIdleTimeoutInSeconds,
+                };
+            }
         }
 
         [JsonProperty("request")]
-        internal EntityBatchRequest Request { get; private set; }
+        internal EntityBatchRequest Request { get; private set; }
+
+        [JsonProperty("configurations")]
+        public RemoteInstanceConfiguration? Configurations { get; private set; }
 
         [JsonIgnore]
         internal EntityBatchResult? Result { get; set; }

src/WebJobs.Extensions.DurableTask/ContextImplementations/RemoteOrchestratorContext.cs

@@ -26,14 +26,14 @@ public RemoteOrchestratorContext(OrchestrationRuntimeState runtimeState, TaskOrc
         {
             this.runtimeState = runtimeState ?? throw new ArgumentNullException(nameof(runtimeState));
             this.EntityParameters = entityParameters;
-            this.Configurations = new RemoteOrchestratorConfiguration
+            this.Configurations = new RemoteInstanceConfiguration
             {
                 HttpDefaultAsyncRequestSleepTimeMilliseconds = options.HttpSettings.DefaultAsyncRequestSleepTimeMilliseconds,
             };
             if (options.ExtendedSessionsEnabled)
             {
                 this.Configurations.IsExtendedSession = isExtendedSession;
-                this.Configurations.IncludePastEvents = includePastEvents;
+                this.Configurations.IncludeState = includePastEvents;
                 this.Configurations.ExtendedSessionIdleTimeoutInSeconds = options.ExtendedSessionIdleTimeoutInSeconds;
             }
         }
@@ -51,7 +51,7 @@ public RemoteOrchestratorContext(OrchestrationRuntimeState runtimeState, TaskOrc
         internal int UpperSchemaVersion { get; } = 4;
 
         [JsonProperty("configurations")]
-        public RemoteOrchestratorConfiguration Configurations { get; private set; }
+        public RemoteInstanceConfiguration Configurations { get; private set; }
 
         [JsonIgnore]
         internal bool OrchestratorCompleted { get; private set; }

src/WebJobs.Extensions.DurableTask/OutOfProcMiddleware.cs

@@ -111,10 +111,13 @@
             }
 
             WorkItemMetadata workItemMetadata = dispatchContext.GetProperty<WorkItemMetadata>();
-            bool isExtendedSession = workItemMetadata.IsExtendedSession;
-            bool includePastEvents = workItemMetadata.IncludePastEvents;
+            var context = new RemoteOrchestratorContext(
+                runtimeState,
+                entityParameters,
+                this.Options,
+                workItemMetadata.IsExtendedSession,
+                workItemMetadata.IncludeState);
 
-            var context = new RemoteOrchestratorContext(runtimeState, entityParameters, this.extension.Options, isExtendedSession, includePastEvents);
             bool workerRequiresHistory = false;
 
             var input = new TriggeredFunctionData
@@ -351,9 +354,16 @@
                 batchRequest.InstanceId,
                 batchRequest.EntityState,
                 functionType: FunctionType.Entity,
-                isReplay: false);
+                isReplay: false);
+
+            WorkItemMetadata workItemMetadata = dispatchContext.GetProperty<WorkItemMetadata>();
+            var context = new RemoteEntityContext(
+                batchRequest,
+                this.Options,
+                workItemMetadata.IsExtendedSession,
+                workItemMetadata.IncludeState);
 
-            var context = new RemoteEntityContext(batchRequest);
+            bool workerRequiresEntityState = false;
 
             var input = new TriggeredFunctionData
             {
@@ -381,6 +391,7 @@
 
                     byte[] triggerReturnValueBytes = Convert.FromBase64String(triggerReturnValue);
                     P.EntityBatchResult response = P.EntityBatchResult.Parser.ParseFrom(triggerReturnValueBytes);
+                    workerRequiresEntityState = response.RequiresState;
                     context.Result = response.ToEntityBatchResult();
 
                     context.ThrowIfFailed();
@@ -449,6 +460,11 @@
                 }
 
                 return;
+            }
+
+            if (workerRequiresEntityState)
+            {
+                throw new SessionAbortedException("The worker has since ended the extended session and needs an entity state to execute the request.");
             }
 
             EntityBatchResult batchResult = context.Result

src/WebJobs.Extensions.DurableTask/ProtobufUtils.cs

@@ -486,7 +486,7 @@
         /// <param name="entityBatchRequest">The operation request to convert.</param>
         /// <returns>The converted operation request.</returns>
         [return: NotNullIfNotNull("entityBatchRequest")]
-        internal static P.EntityBatchRequest? ToEntityBatchRequest(this EntityBatchRequest? entityBatchRequest)
+        internal static P.EntityBatchRequest? ToEntityBatchRequest(this EntityBatchRequest? entityBatchRequest, RemoteInstanceConfiguration? configurations)
         {
             if (entityBatchRequest == null)
             {
@@ -496,9 +496,11 @@
             var batchRequest = new P.EntityBatchRequest()
             {
                 InstanceId = entityBatchRequest.InstanceId,
-                EntityState = entityBatchRequest.EntityState,
+                EntityState = configurations?.IncludeState == false ? null : entityBatchRequest.EntityState,
             };
 
+            batchRequest.Properties.Add(ProtobufUtils.ConvertPocoToProtoMap(configurations));
+
             foreach (var operation in entityBatchRequest.Operations ?? Enumerable.Empty<OperationRequest>())
             {
                 batchRequest.Operations.Add(operation.ToOperationRequest());

src/WebJobs.Extensions.DurableTask/RemoteOrchestratorConfiguration.cs → src/WebJobs.Extensions.DurableTask/RemoteInstanceConfiguration.cs

@@ -4,30 +4,31 @@
 namespace Microsoft.Azure.WebJobs.Extensions.DurableTask
 {
     /// <summary>
-    /// Configuration settings for RemoteOrchestratorContext in out-of-process mode, transmitted via gRPC.
+    /// Configuration settings for <see cref="RemoteOrchestratorContext"> and <see cref="RemoteEntityContext">
+    /// in out-of-process mode, transmitted via gRPC.
     /// </summary>
-    public class RemoteOrchestratorConfiguration
+    internal class RemoteInstanceConfiguration
     {
         /// <summary>
         /// Gets or sets the default number of milliseconds between async HTTP status poll requests.
         /// </summary>
-        public int HttpDefaultAsyncRequestSleepTimeMilliseconds { get; set; } = 30000;
+        internal int HttpDefaultAsyncRequestSleepTimeMilliseconds { get; set; } = 30000;
 
         /// <summary>
-        /// Gets or sets whether or not to include the past history events in the orchestration request.
+        /// Gets or sets whether or not to include the instance state in the instance request.
         /// True by default.
         /// </summary>
-        public bool IncludePastEvents { get; set; } = true;
+        internal bool IncludeState { get; set; } = true;
 
         /// <summary>
         /// Gets or sets whether or not the orchestration request is within an extended session.
         /// False by default.
         /// </summary>
-        public bool IsExtendedSession { get; set; } = false;
+        internal bool IsExtendedSession { get; set; } = false;
 
         /// <summary>
         /// Gets or sets the amount of time in seconds before an idle extended session times out.
         /// </summary>
-        public int ExtendedSessionIdleTimeoutInSeconds { get; set; }
+        internal int ExtendedSessionIdleTimeoutInSeconds { get; set; }
     }
 }

src/Worker.Extensions.DurableTask/Execution/DurableFunctionExecutor.Entity.cs

@@ -37,7 +37,7 @@ private async ValueTask RunEntityAsync(FunctionContext context, BindingMetadata
             return;
         }
 
-        TaskEntityDispatcher dispatcher = new(encodedEntityBatch, context.InstanceServices);
+        TaskEntityDispatcher dispatcher = new(encodedEntityBatch, context.InstanceServices, extendedSessionsCache);
         triggerInputData.Value = dispatcher;
         await inner.ExecuteAsync(context);
         context.GetInvocationResult().Value = dispatcher.Result;

src/Worker.Extensions.DurableTask/TaskEntityDispatcher.cs

@@ -4,7 +4,9 @@
 using System;
 using System.Threading.Tasks;
 using Microsoft.DurableTask.Entities;
+using Microsoft.DurableTask.Worker;
 using Microsoft.DurableTask.Worker.Grpc;
+using Microsoft.Extensions.Caching.Memory;
 using Microsoft.Extensions.DependencyInjection;
 
 namespace Microsoft.Azure.Functions.Worker;
@@ -19,11 +21,13 @@
 {
     private readonly string request;
     private readonly IServiceProvider services;
+    private readonly ExtendedSessionsCache extendedSessionsCache;
 
-    internal TaskEntityDispatcher(string request, IServiceProvider services)
+    internal TaskEntityDispatcher(string request, IServiceProvider services, ExtendedSessionsCache extendedSessionsCache)
     {
         this.request = request;
         this.services = services;
+        this.extendedSessionsCache = extendedSessionsCache;
     }
 
     internal string Result { get; private set; } = string.Empty;
@@ -40,7 +44,7 @@
             throw new ArgumentNullException(nameof(entity));
         }
 
-        this.Result = await GrpcEntityRunner.LoadAndRunAsync(this.request, entity, this.services);
+        this.Result = await GrpcEntityRunner.LoadAndRunAsync(this.request, entity, this.extendedSessionsCache, this.services);
     }
 
     /// <summary>