handle internal errors like OOM in entity processing by aborting the … (#2234)

* handle internal errors like OOM in entity processing by aborting the session

* update release_notes

* fix whitespace that breaks build

Author: sebastianburckhardt

release_notes.md

@@ -1,6 +1,7 @@
 ## Bug Fixes
 - Correctly Serialize HostStoppingEvent in ActivityShim (https://github.com/Azure/azure-functions-durable-extension/pull/2178)
 - Fix NotImplementedException for management API calls from Java client (https://github.com/Azure/azure-functions-durable-extension/pull/2193)
+- Handle OOM and other exceptions in entity shim by aborting the session (https://github.com/Azure/azure-functions-durable-extension/pull/2234)
 
 ## Enhancements
 - add optional 'instanceIdPrefix' query parameter to the HTTP API for instance queries

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

@@ -8,6 +8,8 @@
 using System.Runtime.ExceptionServices;
 using System.Threading.Tasks;
 using DurableTask.Core;
+using DurableTask.Core.Common;
+using DurableTask.Core.Exceptions;
 using Microsoft.Azure.WebJobs.Host.Bindings;
 using Newtonsoft.Json;
 
@@ -132,6 +134,14 @@ public void CaptureApplicationError(Exception e)
             this.ApplicationErrors.Add(ExceptionDispatchInfo.Capture(e));
         }
 
+        public void AbortOnInternalError()
+        {
+            if (this.InternalError != null)
+            {
+                throw new SessionAbortedException($"Session aborted because of {this.InternalError.SourceException.GetType().Name}", this.InternalError.SourceException);
+            }
+        }
+
         public void ThrowInternalExceptionIfAny()
         {
             if (this.InternalError != null)

src/WebJobs.Extensions.DurableTask/DurableTaskExtension.cs

@@ -1128,11 +1128,11 @@ private async Task EntityMiddleware(DispatchMiddlewareContext dispatchContext, F
                 }
             }
 
-            await entityContext.RunDeferredTasks();
-
-            // If there were internal errors, do not commit the batch, but instead rethrow
+            // If there were internal errors, throw a SessionAbortedException
             // here so DTFx can abort the batch and back off the work item
-            entityContext.ThrowInternalExceptionIfAny();
+            entityContext.AbortOnInternalError();
+
+            await entityContext.RunDeferredTasks();
         }
 
         internal string GetDefaultConnectionName()

src/WebJobs.Extensions.DurableTask/Listener/TaskEntityShim.cs

@@ -8,6 +8,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using DurableTask.Core;
+using DurableTask.Core.Exceptions;
 using Newtonsoft.Json;
 using Newtonsoft.Json.Linq;
 
@@ -185,70 +186,78 @@ public override async Task<string> Execute(OrchestrationContext innerContext, st
 
             DurableTaskExtension.TagActivityWithOrchestrationStatus(status, this.context.InstanceId, true);
 #endif
-
-            if (this.operationBatch.Count == 0
-                && this.lockRequest == null
-                && (this.toBeRescheduled == null || this.toBeRescheduled.Count == 0)
-                && !this.suspendAndContinueWithDelay)
-            {
-                // we are idle after a ContinueAsNew - the batch is empty.
-                // Wait for more messages to get here (via extended sessions)
-                await this.doneProcessingMessages.Task;
-            }
-
-            if (!this.messageDataConverter.IsDefault)
-            {
-                innerContext.MessageDataConverter = this.messageDataConverter;
-            }
-
-            if (!this.errorDataConverter.IsDefault)
-            {
-                innerContext.ErrorDataConverter = this.errorDataConverter;
-            }
-
-            if (this.NumberEventsToReceive > 0)
+            try
             {
-                await this.doneProcessingMessages.Task;
-            }
+                if (this.operationBatch.Count == 0
+                    && this.lockRequest == null
+                    && (this.toBeRescheduled == null || this.toBeRescheduled.Count == 0)
+                    && !this.suspendAndContinueWithDelay)
+                {
+                    // we are idle after a ContinueAsNew - the batch is empty.
+                    // Wait for more messages to get here (via extended sessions)
+                    await this.doneProcessingMessages.Task;
+                }
 
-            // Commit the effects of this batch, if
-            // we have not already run into an internal error
-            // (in which case we will abort the batch instead of committing it)
-            if (this.context.InternalError == null)
-            {
-                bool writeBackSuccessful = true;
-                ResponseMessage serializationErrorMessage = null;
+                if (!this.messageDataConverter.IsDefault)
+                {
+                    innerContext.MessageDataConverter = this.messageDataConverter;
+                }
 
-                if (this.RollbackFailedOperations)
+                if (!this.errorDataConverter.IsDefault)
                 {
-                    // the state has already been written back, since it is
-                    // done right after each operation.
+                    innerContext.ErrorDataConverter = this.errorDataConverter;
                 }
-                else
+
+                if (this.NumberEventsToReceive > 0)
                 {
-                    // we are writing back the state here, after executing
-                    // the entire batch of operations.
-                    writeBackSuccessful = this.context.TryWriteback(out serializationErrorMessage, out var _);
+                    await this.doneProcessingMessages.Task;
                 }
 
-                // Reschedule all signals that were received before their time
-                this.context.RescheduleMessages(innerContext, this.toBeRescheduled);
+                // Commit the effects of this batch, if
+                // we have not already run into an internal error
+                // (in which case we will abort the batch instead of committing it)
+                if (this.context.InternalError == null)
+                {
+                    bool writeBackSuccessful = true;
+                    ResponseMessage serializationErrorMessage = null;
 
-                // Send all buffered outgoing messages
-                this.context.SendOutbox(innerContext, writeBackSuccessful, serializationErrorMessage);
+                    if (this.RollbackFailedOperations)
+                    {
+                        // the state has already been written back, since it is
+                        // done right after each operation.
+                    }
+                    else
+                    {
+                        // we are writing back the state here, after executing
+                        // the entire batch of operations.
+                        writeBackSuccessful = this.context.TryWriteback(out serializationErrorMessage, out var _);
+                    }
 
-                // send a continue signal
-                if (this.suspendAndContinueWithDelay)
-                {
-                    this.context.SendContinue(innerContext);
-                    this.suspendAndContinueWithDelay = false;
-                    this.context.State.Suspended = true;
-                }
+                    // Reschedule all signals that were received before their time
+                    this.context.RescheduleMessages(innerContext, this.toBeRescheduled);
 
-                var jstate = JToken.FromObject(this.context.State);
+                    // Send all buffered outgoing messages
+                    this.context.SendOutbox(innerContext, writeBackSuccessful, serializationErrorMessage);
 
-                // continue as new
-                innerContext.ContinueAsNew(jstate);
+                    // send a continue signal
+                    if (this.suspendAndContinueWithDelay)
+                    {
+                        this.context.SendContinue(innerContext);
+                        this.suspendAndContinueWithDelay = false;
+                        this.context.State.Suspended = true;
+                    }
+
+                    var jstate = JToken.FromObject(this.context.State);
+
+                    // continue as new
+                    innerContext.ContinueAsNew(jstate);
+                }
+            }
+            catch (Exception e)
+            {
+                // we must catch unexpected exceptions here, otherwise entity goes into permanent failed state
+                // for example, there can be an OOM thrown during serialization https://github.com/Azure/azure-functions-durable-extension/issues/2166
+                this.context.CaptureInternalError(e);
             }
 
             // The return value is not used.