addressing PR comments

Author: Sophia Tevosyan

src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs

@@ -1099,7 +1099,7 @@ async Task<string> IsExecutableInstanceAsync(
                 if (instanceStatus.State.OrchestrationInstance.ExecutionId == runtimeState.OrchestrationInstance.ExecutionId
                     && instanceStatus.State.OrchestrationStatus != runtimeState.OrchestrationStatus)
                 {
-                    await this.trackingStore.UpdateInstanceStatusAndDeleteOrphanedBlobsAsync(
+                    await this.trackingStore.UpdateInstanceStatusAndDeleteOrphanedBlobsForCompletedOrchestrationAsync(
                         runtimeState.OrchestrationInstance.InstanceId,
                         runtimeState.OrchestrationInstance.ExecutionId,
                         runtimeState,

src/DurableTask.AzureStorage/Tracking/AzureTableTrackingStore.cs

@@ -1043,13 +1043,32 @@ public override Task StartAsync(CancellationToken cancellationToken = default)
             return eTagValue;
         }
 
-        public override async Task UpdateInstanceStatusAndDeleteOrphanedBlobsAsync(
+        public override async Task UpdateInstanceStatusAndDeleteOrphanedBlobsForCompletedOrchestrationAsync(
             string instanceId,
             string executionId,
             OrchestrationRuntimeState runtimeState,
             object trackingStoreContext,
             CancellationToken cancellationToken = default)
         {
+            if (runtimeState.OrchestrationStatus != OrchestrationStatus.Completed &&
+                runtimeState.OrchestrationStatus != OrchestrationStatus.Canceled &&
+                runtimeState.OrchestrationStatus != OrchestrationStatus.Failed &&
+                runtimeState.OrchestrationStatus != OrchestrationStatus.Terminated)
+            {
+                return;
+            }
+
+            TrackingStoreContext context = (TrackingStoreContext)trackingStoreContext;
+            if (context.Blobs.Count > 0)
+            {
+                var tasks = new List<Task>(context.Blobs.Count);
+                foreach (string blobName in context.Blobs)
+                {
+                    tasks.Add(this.messageManager.DeleteBlobAsync(blobName));
+                }
+                await Task.WhenAll(tasks);
+            }
+
             string sanitizedInstanceId = KeySanitation.EscapePartitionKey(instanceId);
             var instanceEntity = new TableEntity(sanitizedInstanceId, string.Empty)
             {
@@ -1059,7 +1078,7 @@ public override async Task UpdateInstanceStatusAndDeleteOrphanedBlobsAsync(
                 ["ExecutionId"] = executionId,
                 ["LastUpdatedTime"] = runtimeState.Events.Last().Timestamp,
                 ["RuntimeStatus"] = runtimeState.OrchestrationStatus.ToString(),
-                ["CompletedTime"] = runtimeState.Events.Last().Timestamp // do we want to do this as a rough proxy or DateTime.UtcNow?
+                ["CompletedTime"] = runtimeState.CompletedTime
             };
 
             Stopwatch orchestrationInstanceUpdateStopwatch = Stopwatch.StartNew();
@@ -1073,17 +1092,6 @@ public override async Task UpdateInstanceStatusAndDeleteOrphanedBlobsAsync(
                 runtimeState.OrchestrationStatus,
                 Utils.GetEpisodeNumber(runtimeState),
                 orchestrationInstanceUpdateStopwatch.ElapsedMilliseconds);
-
-            TrackingStoreContext context = (TrackingStoreContext)trackingStoreContext;
-            if (context.Blobs.Count > 0)
-            {
-                var tasks = new List<Task>(context.Blobs.Count);
-                foreach (var blobName in context.Blobs)
-                {
-                    tasks.Add(this.messageManager.DeleteBlobAsync(blobName));
-                }
-                await Task.WhenAll(tasks);
-            }
         }
 
         static int GetEstimatedByteCount(TableEntity entity)

src/DurableTask.AzureStorage/Tracking/ITrackingStore.cs

@@ -104,17 +104,17 @@ interface ITrackingStore
         Task<InstanceStatus> FetchInstanceStatusAsync(string instanceId, CancellationToken cancellationToken = default);
 
         /// <summary>
-        /// Updates the instance status of the specified orchestration instance to match that of <paramref name="runtimeState"/>.
+        /// Updates the instance status of the specified orchestration instance to match that of <paramref name="runtimeState"/> for a completed orchestration.
         /// Also deletes any orphaned blobs of <paramref name="trackingStoreContext"/>.
-        /// This method is meant to be called in the case that there is an inconsistency between the instance and history table due to a 
-        /// failure during a call to <see cref="UpdateStateAsync"/>.
+        /// This method is meant to be called in the case that there is an inconsistency between the instance and history table due to a failure during a call to 
+        /// <see cref="UpdateStateAsync"/> for a completing orchestration. If the orchestration is not in a terminal state, the method will immediately return and do nothing.
         /// </summary>
         /// <param name="instanceId">The ID of the orchestration.</param>
         /// <param name="executionId">The execution ID of the orchestration.</param>
         /// <param name="runtimeState">The runtime state of the orchestration.</param>
         /// <param name="trackingStoreContext">Additional context for the execution that is maintained by the tracking store.</param>
         /// <param name="cancellationToken">The token to monitor for cancellation requests. The default value is <see cref="CancellationToken.None"/>.</param>
-        Task UpdateInstanceStatusAndDeleteOrphanedBlobsAsync(string instanceId, string executionId, OrchestrationRuntimeState runtimeState, object trackingStoreContext, CancellationToken cancellationToken = default);
+        Task UpdateInstanceStatusAndDeleteOrphanedBlobsForCompletedOrchestrationAsync(string instanceId, string executionId, OrchestrationRuntimeState runtimeState, object trackingStoreContext, CancellationToken cancellationToken = default);
 
         /// <summary>
         /// Get The Orchestration State for querying all orchestration instances

src/DurableTask.AzureStorage/Tracking/InstanceStoreBackedTrackingStore.cs

@@ -187,13 +187,21 @@ public override async Task UpdateStatusForTerminationAsync(string instanceId, st
             await this.instanceStore.WriteEntitiesAsync(instanceEntity);
         }
 
-        public override async Task UpdateInstanceStatusAndDeleteOrphanedBlobsAsync(
+        public override async Task UpdateInstanceStatusAndDeleteOrphanedBlobsForCompletedOrchestrationAsync(
             string instanceId,
             string executionId,
             OrchestrationRuntimeState runtimeState,
             object trackingStoreContext,
             CancellationToken cancellationToken = default)
         {
+            if (runtimeState.OrchestrationStatus != OrchestrationStatus.Completed &&
+                runtimeState.OrchestrationStatus != OrchestrationStatus.Canceled &&
+                runtimeState.OrchestrationStatus != OrchestrationStatus.Failed &&
+                runtimeState.OrchestrationStatus != OrchestrationStatus.Terminated)
+            {
+                return;
+            }
+
             // No blobs to delete for this tracking store implementation
             await instanceStore.WriteEntitiesAsync(new InstanceEntityBase[]
             {

src/DurableTask.AzureStorage/Tracking/TrackingStoreBase.cs

@@ -110,6 +110,6 @@ public virtual Task UpdateStatusForRewindAsync(string instanceId, CancellationTo
         public abstract Task<ETag?> UpdateStateAsync(OrchestrationRuntimeState newRuntimeState, OrchestrationRuntimeState oldRuntimeState, string instanceId, string executionId, ETag? eTag, object executionData, CancellationToken cancellationToken = default);
 
         /// <inheritdoc />
-        public abstract Task UpdateInstanceStatusAndDeleteOrphanedBlobsAsync(string instanceId, string executionId, OrchestrationRuntimeState runtimeState, object trackingStoreContext, CancellationToken cancellationToken = default);
+        public abstract Task UpdateInstanceStatusAndDeleteOrphanedBlobsForCompletedOrchestrationAsync(string instanceId, string executionId, OrchestrationRuntimeState runtimeState, object trackingStoreContext, CancellationToken cancellationToken = default);
     }
 }