Enhance purge with parallel batch deletes and partial purge timeout

src/DurableTask.AzureStorage/AzureStorageOrchestrationService.cs

@@ -1,4 +1,4 @@
-//  ----------------------------------------------------------------------------------
+//  ----------------------------------------------------------------------------------
 //  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.
@@ -1886,9 +1886,9 @@ public async Task<IList<OrchestrationState>> GetOrchestrationStateAsync(string i
         /// <summary>
         /// Gets the state of all orchestration instances that match the specified parameters.
         /// </summary>
-        /// <param name="createdTimeFrom">CreatedTime of orchestrations. Fetch status grater than this value.</param>
+        /// <param name="createdTimeFrom">CreatedTime of orchestrations. Fetch status greater than this value.</param>
         /// <param name="createdTimeTo">CreatedTime of orchestrations. Fetch status less than this value.</param>
-        /// <param name="runtimeStatus">RuntimeStatus of orchestrations. You can specify several status.</param>
+        /// <param name="runtimeStatus">RuntimeStatus of orchestrations. You can specify several statuses.</param>
         /// <param name="cancellationToken">Cancellation Token</param>
         /// <returns>List of <see cref="OrchestrationState"/></returns>
         public async Task<IList<OrchestrationState>> GetOrchestrationStateAsync(DateTime createdTimeFrom, DateTime? createdTimeTo, IEnumerable<OrchestrationStatus> runtimeStatus, CancellationToken cancellationToken = default(CancellationToken))
@@ -1900,9 +1900,9 @@ public async Task<IList<OrchestrationState>> GetOrchestrationStateAsync(string i
         /// <summary>
         /// Gets the state of all orchestration instances that match the specified parameters.
         /// </summary>
-        /// <param name="createdTimeFrom">CreatedTime of orchestrations. Fetch status grater than this value.</param>
+        /// <param name="createdTimeFrom">CreatedTime of orchestrations. Fetch status greater than this value.</param>
         /// <param name="createdTimeTo">CreatedTime of orchestrations. Fetch status less than this value.</param>
-        /// <param name="runtimeStatus">RuntimeStatus of orchestrations. You can specify several status.</param>
+        /// <param name="runtimeStatus">RuntimeStatus of orchestrations. You can specify several statuses.</param>
         /// <param name="top">Top is number of records per one request.</param>
         /// <param name="continuationToken">ContinuationToken of the pager.</param>
         /// <param name="cancellationToken">Cancellation Token</param>
@@ -2021,9 +2021,9 @@ public Task<PurgeHistoryResult> PurgeInstanceHistoryAsync(string instanceId)
         /// <summary>
         /// Purge history for orchestrations that match the specified parameters.
         /// </summary>
-        /// <param name="createdTimeFrom">CreatedTime of orchestrations. Purges history grater than this value.</param>
+        /// <param name="createdTimeFrom">CreatedTime of orchestrations. Purges history greater than this value.</param>
         /// <param name="createdTimeTo">CreatedTime of orchestrations. Purges history less than this value.</param>
-        /// <param name="runtimeStatus">RuntimeStatus of orchestrations. You can specify several status.</param>
+        /// <param name="runtimeStatus">RuntimeStatus of orchestrations. You can specify several statuses.</param>
         /// <returns>Class containing number of storage requests sent, along with instances and rows deleted/purged</returns>
         public Task<PurgeHistoryResult> PurgeInstanceHistoryAsync(DateTime createdTimeFrom, DateTime? createdTimeTo, IEnumerable<OrchestrationStatus> runtimeStatus)
         {
@@ -2040,10 +2040,28 @@ async Task<PurgeResult> IOrchestrationServicePurgeClient.PurgeInstanceStateAsync
         /// <inheritdoc />
         async Task<PurgeResult> IOrchestrationServicePurgeClient.PurgeInstanceStateAsync(PurgeInstanceFilter purgeInstanceFilter)
         {
-            PurgeHistoryResult storagePurgeHistoryResult = await this.PurgeInstanceHistoryAsync(
-                purgeInstanceFilter.CreatedTimeFrom,
-                purgeInstanceFilter.CreatedTimeTo,
-                purgeInstanceFilter.RuntimeStatus);
+            PurgeHistoryResult storagePurgeHistoryResult;
+            if (purgeInstanceFilter.Timeout.HasValue)
+            {
+                // Convert the timeout into a CancellationToken so that the tracking store
+                // only needs to observe a single cancellation mechanism.
+                using var timeoutCts = new CancellationTokenSource(purgeInstanceFilter.Timeout.Value);
+                storagePurgeHistoryResult = await this.trackingStore.PurgeInstanceHistoryAsync(
+                    purgeInstanceFilter.CreatedTimeFrom,
+                    purgeInstanceFilter.CreatedTimeTo,
+                    purgeInstanceFilter.RuntimeStatus,
+                    timeoutCts.Token);
+            }
+            else
+            {
+                // No timeout: use the original code path (no CancellationToken) to preserve
+                // backward-compatible behavior where IsComplete is null.
+                storagePurgeHistoryResult = await this.trackingStore.PurgeInstanceHistoryAsync(
+                    purgeInstanceFilter.CreatedTimeFrom,
+                    purgeInstanceFilter.CreatedTimeTo,
+                    purgeInstanceFilter.RuntimeStatus);
+            }
+
             return storagePurgeHistoryResult.ToCorePurgeHistoryResult();
         }
 #nullable enable

src/DurableTask.AzureStorage/MessageManager.cs

@@ -318,8 +318,9 @@ public string GetNewLargeMessageBlobName(MessageData message)
 
         public async Task<int> DeleteLargeMessageBlobs(string sanitizedInstanceId, CancellationToken cancellationToken = default)
         {
-            int storageOperationCount = 1;
-            if (await this.blobContainer.ExistsAsync(cancellationToken))
+            int storageOperationCount = 0;
+
+            try
             {
                 await foreach (Page<Blob> page in this.blobContainer.ListBlobsAsync(sanitizedInstanceId, cancellationToken).AsPages())
                 {
@@ -329,6 +330,22 @@ public async Task<int> DeleteLargeMessageBlobs(string sanitizedInstanceId, Cance
 
                     storageOperationCount += page.Values.Count;
                 }
+
+                // Count the list operation even if no blobs found (the initial list request still happened)
+                if (storageOperationCount == 0)
+                {
+                    storageOperationCount = 1;
+                }
+            }
+            catch (DurableTaskStorageException ex) when (ex.HttpStatusCode == 404)
+            {
+                // Container does not exist; nothing to delete.
+                storageOperationCount = 1;
+            }
+            catch (Azure.RequestFailedException ex) when (ex.Status == 404)
+            {
+                // Container does not exist; nothing to delete.
+                storageOperationCount = 1;
             }
 
             return storageOperationCount;

src/DurableTask.AzureStorage/PurgeHistoryResult.cs

@@ -33,6 +33,19 @@ public PurgeHistoryResult(int storageRequests, int instancesDeleted, int rowsDel
             this.RowsDeleted = rowsDeleted;
         }
 
+        /// <summary>
+        /// Constructor for purge history statistics with completion status.
+        /// </summary>
+        /// <param name="storageRequests">Requests sent to storage</param>
+        /// <param name="instancesDeleted">Number of instances deleted</param>
+        /// <param name="rowsDeleted">Number of rows deleted</param>
+        /// <param name="isComplete">Whether the purge operation completed all matching instances.</param>
+        public PurgeHistoryResult(int storageRequests, int instancesDeleted, int rowsDeleted, bool? isComplete)
+            : this(storageRequests, instancesDeleted, rowsDeleted)
+        {
+            this.IsComplete = isComplete;
+        }
+
         /// <summary>
         /// Number of requests sent to Storage during this execution of purge history
         /// </summary>
@@ -48,12 +61,20 @@ public PurgeHistoryResult(int storageRequests, int instancesDeleted, int rowsDel
         /// </summary>
         public int RowsDeleted { get; }
 
+        /// <summary>
+        /// Gets a value indicating whether the purge operation is complete.
+        /// <c>true</c> if all matching instances were purged;
+        /// <c>false</c> if more instances remain and purge should be called again;
+        /// <c>null</c> if completion status is unknown.
+        /// </summary>
+        public bool? IsComplete { get; }
+
         /// <summary>
         /// Converts from AzureStorage.PurgeHistoryResult to Core.PurgeResult type
         /// </summary>
         public PurgeResult ToCorePurgeHistoryResult()
         {
-            return new PurgeResult(this.InstancesDeleted);
+            return new PurgeResult(this.InstancesDeleted, this.IsComplete);
         }
     }
 }

src/DurableTask.AzureStorage/Storage/Table.cs

@@ -26,6 +26,8 @@ namespace DurableTask.AzureStorage.Storage
 
     class Table
     {
+        const int MaxParallelBatchDeletes = 10;
+
         readonly AzureStorageClient azureStorageClient;
         readonly AzureStorageOrchestrationServiceStats stats;
         readonly TableServiceClient tableServiceClient;
@@ -117,6 +119,139 @@ public async Task<TableTransactionResults> DeleteBatchAsync<T>(IEnumerable<T> en
             return await this.ExecuteBatchAsync(entityBatch, item => new TableTransactionAction(TableTransactionActionType.Delete, item), cancellationToken: cancellationToken);
         }
 
+        /// <summary>
+        /// Deletes entities in parallel batches of up to 100. Each batch is an atomic transaction,
+        /// but multiple batches are submitted concurrently for improved throughput.
+        /// Concurrency is limited to <see cref="MaxParallelBatchDeletes"/> concurrent batch transactions
+        /// to avoid overwhelming storage and starving other operations.
+        /// If a batch fails because an entity was already deleted (404/EntityNotFound),
+        /// it falls back to individual deletes for that batch, skipping already-deleted entities.
+        /// </summary>
+        public async Task<TableTransactionResults> DeleteBatchParallelAsync<T>(
+            IReadOnlyList<T> entityBatch,
+            CancellationToken cancellationToken = default) where T : ITableEntity
+        {
+            if (entityBatch.Count == 0)
+            {
+                return new TableTransactionResults(Array.Empty<Response>(), TimeSpan.Zero, 0);
+            }
+
+            const int batchSize = 100;
+            int chunkCount = (entityBatch.Count + batchSize - 1) / batchSize;
+            var chunks = new List<List<TableTransactionAction>>(chunkCount);
+
+            var currentChunk = new List<TableTransactionAction>(batchSize);
+            foreach (T entity in entityBatch)
+            {
+                currentChunk.Add(new TableTransactionAction(TableTransactionActionType.Delete, entity));
+                if (currentChunk.Count == batchSize)
+                {
+                    chunks.Add(currentChunk);
+                    currentChunk = new List<TableTransactionAction>(batchSize);
+                }
+            }
+
+            if (currentChunk.Count > 0)
+            {
+                chunks.Add(currentChunk);
+            }
+
+            var resultsBuilder = new TableTransactionResultsBuilder();
+            using var semaphore = new SemaphoreSlim(MaxParallelBatchDeletes);
+
+            var stopwatch = Stopwatch.StartNew();
+            TableTransactionResults[] allResults = await Task.WhenAll(
+                chunks.Select(async chunk =>
+                {
+                    await semaphore.WaitAsync(cancellationToken);
+                    try
+                    {
+                        return await this.ExecuteBatchWithFallbackAsync(chunk, cancellationToken);
+                    }
+                    finally
+                    {
+                        semaphore.Release();
+                    }
+                }));
+            stopwatch.Stop();
+
+            foreach (TableTransactionResults result in allResults)
+            {
+                resultsBuilder.Add(result);
+            }
+
+            TableTransactionResults aggregatedResults = resultsBuilder.ToResults();
+            return new TableTransactionResults(aggregatedResults.Responses, stopwatch.Elapsed, aggregatedResults.RequestCount);
+        }
+
+        /// <summary>
+        /// Executes a batch transaction. If it fails due to an entity not found (404),
+        /// falls back to individual delete operations, skipping entities that are already gone.
+        /// </summary>
+        async Task<TableTransactionResults> ExecuteBatchWithFallbackAsync(
+            List<TableTransactionAction> batch,
+            CancellationToken cancellationToken)
+        {
+            try
+            {
+                return await this.ExecuteBatchAsync(batch, cancellationToken);
+            }
+            catch (DurableTaskStorageException ex) when (ex.HttpStatusCode == 404)
+            {
+                // One or more entities in the batch were already deleted.
+                // Fall back to individual deletes, skipping 404s.
+                // Count the failed batch attempt as 1 storage request.
+                TableTransactionResults fallbackResults = await this.DeleteEntitiesIndividuallyAsync(batch, cancellationToken);
+                return new TableTransactionResults(
+                    fallbackResults.Responses,
+                    fallbackResults.Elapsed,
+                    fallbackResults.RequestCount + 1);
+            }
+            catch (RequestFailedException ex) when (ex.Status == 404)
+            {
+                TableTransactionResults fallbackResults = await this.DeleteEntitiesIndividuallyAsync(batch, cancellationToken);
+                return new TableTransactionResults(
+                    fallbackResults.Responses,
+                    fallbackResults.Elapsed,
+                    fallbackResults.RequestCount + 1);
+            }
+        }
+
+        async Task<TableTransactionResults> DeleteEntitiesIndividuallyAsync(
+            List<TableTransactionAction> batch,
+            CancellationToken cancellationToken)
+        {
+            var responses = new List<Response>();
+            var stopwatch = Stopwatch.StartNew();
+            int requestCount = 0;
+
+            foreach (TableTransactionAction action in batch)
+            {
+                requestCount++;
+                try
+                {
+                    Response response = await this.tableClient.DeleteEntityAsync(
+                        action.Entity.PartitionKey,
+                        action.Entity.RowKey,
+                        action.Entity.ETag,
+                        cancellationToken).DecorateFailure();
+                    responses.Add(response);
+                    this.stats.TableEntitiesWritten.Increment();
+                }
+                catch (DurableTaskStorageException ex) when (ex.HttpStatusCode == 404)
+                {
+                    // Entity already deleted; skip.
+                }
+                catch (RequestFailedException ex) when (ex.Status == 404)
+                {
+                    // Entity already deleted; skip.
+                }
+            }
+
+            stopwatch.Stop();
+            return new TableTransactionResults(responses, stopwatch.Elapsed, requestCount);
+        }
+
         public async Task<TableTransactionResults> InsertOrMergeBatchAsync<T>(IEnumerable<T> entityBatch, CancellationToken cancellationToken = default) where T : ITableEntity
         {
             TableTransactionResults results = await this.ExecuteBatchAsync(entityBatch, item => new TableTransactionAction(TableTransactionActionType.UpsertMerge, item), cancellationToken: cancellationToken);