[Storage] [DataMovement] Adding fix for awaiting completion on a DataTransfer to prevent delay (#33621)

* Adding fix for awaiting completion on a DataTransfer to prevent delay

* Remove SpinWait, utilize TaskCompletionSource instead

* PR Comments

* WIP - Change to run tasks on factory

* Changelog

* Revert change to EnsureCompleted

* Run export-api

* Adding lock surrounding the job part status change

* Moving up status setting above calling the completion status

Author: amnguye

sdk/storage/Azure.Storage.DataMovement.Blobs/CHANGELOG.md

@@ -1,14 +1,7 @@
 # Release History
 
 ## 12.0.0-beta.2 (Unreleased)
-
-### Features Added
-
-### Breaking Changes
-
-### Bugs Fixed
-
-### Other Changes
+- This release contains bug fixes to improve quality.
 
 ## 12.0.0-beta.1 (2022-12-15)
 - This preview is the first release of a ground-up rewrite of our client data movement

sdk/storage/Azure.Storage.DataMovement/CHANGELOG.md

@@ -1,14 +1,7 @@
 # Release History
 
 ## 12.0.0-beta.2 (Unreleased)
-
-### Features Added
-
-### Breaking Changes
-
-### Bugs Fixed
-
-### Other Changes
+- Fix to prevent thread starvation on the DataTransfer.AwaitCompletion
 
 ## 12.0.0-beta.1 (2022-12-15)
 - This preview is the first release of a ground-up rewrite of our client data movement

sdk/storage/Azure.Storage.DataMovement/api/Azure.Storage.DataMovement.netstandard2.0.cs

@@ -1,13 +1,12 @@
 namespace Azure.Storage.DataMovement
 {
-    public partial class DataTransfer : System.IAsyncDisposable
+    public partial class DataTransfer
     {
         internal DataTransfer() { }
         public bool HasCompleted { get { throw null; } }
         public string Id { get { throw null; } }
         public Azure.Storage.DataMovement.StorageTransferStatus TransferStatus { get { throw null; } }
         public System.Threading.Tasks.Task AwaitCompletion(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
-        public System.Threading.Tasks.ValueTask DisposeAsync() { throw null; }
         public void EnsureCompleted(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { }
     }
     [System.FlagsAttribute]

sdk/storage/Azure.Storage.DataMovement/src/JobPartInternal.cs

@@ -85,6 +85,7 @@ internal abstract class JobPartInternal
         /// The current status of each job part.
         /// </summary>
         public StorageTransferStatus JobPartStatus { get; set; }
+        private object _statusLock = new object();
 
         /// <summary>
         /// Optional. If the length is known, we log it instead of doing a GetProperties call on the
@@ -206,10 +207,18 @@ internal async Task TriggerCancellation(StorageTransferStatus status)
         /// <param name="transferStatus"></param>
         internal async Task OnTransferStatusChanged(StorageTransferStatus transferStatus)
         {
-            if (transferStatus != StorageTransferStatus.None
-                && JobPartStatus != transferStatus)
+            bool statusChanged = false;
+            lock (_statusLock)
+            {
+                if (transferStatus != StorageTransferStatus.None
+                    && JobPartStatus != transferStatus)
+                {
+                    statusChanged = true;
+                    JobPartStatus = transferStatus;
+                }
+            }
+            if (statusChanged)
             {
-                JobPartStatus = transferStatus;
                 if (JobPartStatus == StorageTransferStatus.Completed)
                 {
                     await InvokeSingleCompletedArg().ConfigureAwait(false);

sdk/storage/Azure.Storage.DataMovement/src/Shared/DataTransfer.cs

@@ -8,13 +8,14 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core;
+using Azure.Core.Pipeline;
 
 namespace Azure.Storage.DataMovement
 {
     /// <summary>
     /// Holds transfer information
     /// </summary>
-    public class DataTransfer : IAsyncDisposable
+    public class DataTransfer
     {
         /// <summary>
         /// Defines whether the DataTransfer has completed.
@@ -62,56 +63,24 @@ internal DataTransfer(string id, long bytesTransferred)
             _state = new DataTransferState(id, bytesTransferred);
         }
 
-        /// <summary>
-        /// Disposes the DataTransfer object.
-        /// </summary>
-        /// <returns></returns>
-        public async ValueTask DisposeAsync()
-        {
-            await _state.DisposeAsync().ConfigureAwait(false);
-            GC.SuppressFinalize(this);
-        }
-
         /// <summary>
         /// Ensures completion of the DataTransfer and attempts to get result
         /// </summary>
         public void EnsureCompleted(CancellationToken cancellationToken = default)
         {
-#if DEBUG
-            VerifyTaskCompleted(HasCompleted);
-#endif
 #pragma warning disable AZC0102 // Do not use GetAwaiter().GetResult(). Use the TaskExtensions.EnsureCompleted() extension method instead.
-            AwaitCompletion(cancellationToken);
+            AwaitCompletion(cancellationToken).GetAwaiter().GetResult();
 #pragma warning restore AZC0102 // Do not use GetAwaiter().GetResult(). Use the TaskExtensions.EnsureCompleted() extension method instead.
         }
 
-        [Conditional("DEBUG")]
-        private static void VerifyTaskCompleted(bool isCompleted)
-        {
-            if (!isCompleted)
-            {
-                if (Debugger.IsAttached)
-                {
-                    Debugger.Break();
-                }
-            }
-        }
-
         /// <summary>
         /// Waits until the data transfer itself has completed
         /// </summary>
         /// <param name="cancellationToken"></param>
-        public Task AwaitCompletion(CancellationToken cancellationToken = default)
+        public async Task AwaitCompletion(CancellationToken cancellationToken = default)
         {
-            while (!HasCompleted)
-            {
-#if DEBUG
-                VerifyTaskCompleted(HasCompleted);
-#endif
-                CancellationHelper.ThrowIfCancellationRequested(cancellationToken);
-            }
-
-            return Task.CompletedTask;
+            cancellationToken.Register(() => _state._completionSource.TrySetCanceled(cancellationToken), useSynchronizationContext: false);
+            await _state._completionSource.Task.ConfigureAwait(false);
         }
     }
 }

sdk/storage/Azure.Storage.DataMovement/src/Shared/DataTransferState.cs

@@ -3,6 +3,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Runtime.InteropServices.ComTypes;
 using System.Threading;
 using System.Threading.Tasks;
 using Azure.Core;
@@ -13,28 +14,25 @@ namespace Azure.Storage.DataMovement
     /// <summary>
     /// Defines the state of the transfer
     /// </summary>
-    internal class DataTransferState : IAsyncDisposable
+    internal class DataTransferState
     {
-        // To detect redundant calls
-        private bool _disposedValue;
-
+        private readonly object _statusLock = new object();
         private string _id;
         private StorageTransferStatus _status;
+
         private long _currentTransferredBytes;
-        private SemaphoreSlim _statusSemaphore;
+        private object _lockCurrentBytes = new object();
+
+        public TaskCompletionSource<StorageTransferStatus> _completionSource;
 
         public StorageTransferStatus Status => _status;
 
         /// <summary>
         /// constructor
         /// </summary>
         public DataTransferState()
+            : this(StorageTransferStatus.Queued)
         {
-            _disposedValue = false;
-            _statusSemaphore = new SemaphoreSlim(1, 1);
-            _id = Guid.NewGuid().ToString();
-            _status = StorageTransferStatus.Queued;
-            _currentTransferredBytes = 0;
         }
 
         /// <summary>
@@ -45,6 +43,15 @@ public DataTransferState(StorageTransferStatus status)
             _id = Guid.NewGuid().ToString();
             _status = status;
             _currentTransferredBytes = 0;
+            _completionSource = new TaskCompletionSource<StorageTransferStatus>(
+                _status,
+                TaskCreationOptions.RunContinuationsAsynchronously);
+            if (StorageTransferStatus.Completed == status ||
+                        StorageTransferStatus.CompletedWithSkippedTransfers == status ||
+                        StorageTransferStatus.CompletedWithFailedTransfers == status)
+            {
+                _completionSource.TrySetResult(status);
+            }
         }
 
         /// <summary>
@@ -55,20 +62,9 @@ public DataTransferState(string id, long bytesTransferred)
             _id = id;
             _status = StorageTransferStatus.Queued;
             _currentTransferredBytes = bytesTransferred;
-        }
-
-        public async ValueTask DisposeAsync()
-        {
-            if (!_disposedValue)
-            {
-                _disposedValue = true;
-                if (_statusSemaphore.CurrentCount == 0)
-                {
-                    await _statusSemaphore.WaitAsync().ConfigureAwait(false);
-                    _statusSemaphore.Release();
-                }
-                _statusSemaphore.Dispose();
-            }
+            _completionSource = new TaskCompletionSource<StorageTransferStatus>(
+                _status,
+                TaskCreationOptions.RunContinuationsAsynchronously);
         }
 
         /// <summary>
@@ -111,16 +107,23 @@ public void SetId(string id)
         /// Sets the completion status
         /// </summary>
         /// <param name="status"></param>
-        public async Task SetTransferStatus(StorageTransferStatus status)
+        public void SetTransferStatus(StorageTransferStatus status)
         {
-            if (!_disposedValue)
+            lock (_statusLock)
             {
-                await _statusSemaphore.WaitAsync().ConfigureAwait(false);
                 if (_status != status)
                 {
                     _status = status;
+                    if (StorageTransferStatus.Completed == status ||
+                        StorageTransferStatus.CompletedWithSkippedTransfers == status ||
+                        StorageTransferStatus.CompletedWithFailedTransfers == status)
+                    {
+                        // If the _completionSource has been cancelled or the exception
+                        // has been set, we don't need to check if TrySetResult returns false
+                        // because it's acceptable to cancel or have an error occur before then.
+                        _completionSource.TrySetResult(status);
+                    }
                 }
-                _statusSemaphore.Release();
             }
         }
 
@@ -129,7 +132,10 @@ public async Task SetTransferStatus(StorageTransferStatus status)
         /// </summary>
         public void ResetTransferredBytes()
         {
-            Volatile.Write(ref _currentTransferredBytes, 0);
+            lock (_lockCurrentBytes)
+            {
+                Volatile.Write(ref _currentTransferredBytes, 0);
+            }
         }
 
         /// <summary>
@@ -138,7 +144,10 @@ public void ResetTransferredBytes()
         /// <param name="transferredBytes"></param>
         public void UpdateTransferBytes(long transferredBytes)
         {
-            Interlocked.Add(ref _currentTransferredBytes, transferredBytes);
+            lock (_lockCurrentBytes)
+            {
+                Interlocked.Add(ref _currentTransferredBytes, transferredBytes);
+            }
         }
     }
 }

sdk/storage/Azure.Storage.DataMovement/src/TransferJobInternal.cs

@@ -99,6 +99,7 @@ internal abstract class TransferJobInternal : IDisposable
         /// Transfer Status for the job.
         /// </summary>
         internal StorageTransferStatus _transferStatus;
+        private object _statusLock = new object();
 
         /// <summary>
         /// To help set the job status when all job parts have completed.
@@ -314,10 +315,19 @@ public async Task JobPartEvent(TransferStatusEventArgs args)
 
         public async Task OnJobStatusChangedAsync(StorageTransferStatus status)
         {
-            //TODO: change to RaiseAsync after implementing ClientDiagnostics for TransferManager
-            if (_transferStatus != status)
+            bool statusChanged = false;
+            lock (_statusLock)
+            {
+                //TODO: change to RaiseAsync after implementing ClientDiagnostics for TransferManager
+                if (_transferStatus != status)
+                {
+                    statusChanged = true;
+                    _transferStatus = status;
+                }
+                _dataTransfer._state.SetTransferStatus(status);
+            }
+            if (statusChanged)
             {
-                _transferStatus = status;
                 if (TransferStatusEventHandler != null)
                 {
                     await TransferStatusEventHandler.Invoke(
@@ -327,7 +337,6 @@ await TransferStatusEventHandler.Invoke(
                             isRunningSynchronously: false,
                             cancellationToken: _cancellationTokenSource.Token)).ConfigureAwait(false);
                 }
-                await _dataTransfer._state.SetTransferStatus(status).ConfigureAwait(false);
             }
         }
 

sdk/storage/Azure.Storage.DataMovement/tests/DataTransferTests.cs

@@ -41,7 +41,7 @@ public void EnsureCompleted_CancellationToken()
 
             TestHelper.AssertExpectedException(
                 () => transfer.EnsureCompleted(cancellationTokenSource.Token),
-                new TaskCanceledException("The operation was canceled."));
+                new TaskCanceledException("A task was canceled."));
         }
 
         [Test]
@@ -63,7 +63,7 @@ public async Task AwaitCompletion_CancellationToken()
             }
             catch (TaskCanceledException exception)
             {
-                Assert.AreEqual(exception.Message, "The operation was canceled.");
+                Assert.AreEqual(exception.Message, "A task was canceled.");
             }
         }
     }