[DataMovement] Pause and Resume (Azure#35160)

* WIP - make APIs public

* WIP - Adding conversions header to stream

* WIP - added back api to link to job plan

* WIP

* WIP

* WIP - changed from marshal type to unsafe to get struct size in mmf

* WIP

* WIP

* WIP - convert JobPartPlanHeader to manage bytes manually

* WIP

* WIP - Working on deserializing

* WIP - serialize and deserialize tests

* WIP - fixed missing values deserialize and serialize methods

* WIP - move test infra; more checkpointer tests and fixes

* WIP - correct transfer manager to checkpointer changes

* Added Pause Resume Tests WIP: pause then resume tests, directory pause, sync tests

* ExportAPI

* Undo unrelated changes

* WIP - fix null object error and adding in transfer to track failures to end of test

* Deleted unused models; PR Comment regarding API exposure of checkpointer; small test fix

* Addressed PR comments as of 4/4/2023, waiting on replies

* Export API

* WIP

* WIP - implemented wait on complete pause

* PR comments; fix to pause resume then pause; WIP - completedwithfailures tests

* Fixed directory failed test issue

* Record tests

* Remove use of mapnames as it's unsupported in linux and mac; fix tests slash delimiter failure

* more path combine change for mac/linux

* Add check for CheckAndUpdateCancellationStatus on completed and skipped events

* Revert last commit

* Ignore flakey tests

Author: amnguye

sdk/storage/Azure.Storage.DataMovement.Blobs/src/Azure.Storage.DataMovement.Blobs.csproj

@@ -99,8 +99,6 @@
     <Compile Include="$(AzureStorageDataMovementSharedSources)ProgressMultipleParititonedStream.cs" LinkBase="Shared\DataMovement" />
     <Compile Include="$(AzureStorageDataMovementSharedSources)PathScanner.cs" LinkBase="Shared\DataMovement" />
     <Compile Include="$(AzureStorageDataMovementSharedSources)PathScannerFactory.cs" LinkBase="Shared\DataMovement" />
-    <Compile Include="$(AzureStorageDataMovementSharedSources)PlanJobWriter.cs" LinkBase="Shared\DataMovement" />
-    <Compile Include="$(AzureStorageDataMovementSharedSources)PlanJobWriterFactory.cs" LinkBase="Shared\DataMovement" />
     <Compile Include="$(AzureStorageDataMovementSharedSources)StorageManagerTransferStatus.cs" LinkBase="Shared\DataMovement" />
   </ItemGroup>
   <ItemGroup>

sdk/storage/Azure.Storage.DataMovement.Blobs/tests/Azure.Storage.DataMovement.Blobs.Tests.csproj

@@ -7,12 +7,18 @@
   </PropertyGroup>
   <ItemGroup>
     <ProjectReference Include="$(MSBuildThisFileDirectory)..\src\Azure.Storage.DataMovement.Blobs.csproj" />
+    <ProjectReference Include="$(MSBuildThisFileDirectory)..\..\Azure.Storage.DataMovement\src\Azure.Storage.DataMovement.csproj" />
   </ItemGroup>
   <ItemGroup>
     <Compile Include="$(AzureStorageDataMovementTestSharedSources)DataMovementTestBase.cs" LinkBase="Shared\DataMovement" />
+    <Compile Include="$(AzureStorageDataMovementTestSharedSources)DisposingLocalDirectory.cs" LinkBase="Shared\DataMovement" />
   </ItemGroup>
   <ItemGroup>
-    <Compile Include="$(AzureStorageSharedTestSources)\**\*.cs" LinkBase="Shared" />
+    <Compile Include="$(AzureStorageDataMovementSharedSources)..\JobPlanModels\*.cs" LinkBase="Shared\DataMovement\JobPlanModels" />
+    <Compile Include="$(AzureStorageDataMovementSharedSources)DataMovementExtensions.cs" LinkBase="Shared\DataMovement" />
+  </ItemGroup>
+  <ItemGroup>
+    <Compile Include="$(AzureStorageSharedTestSources)\**\*.cs" LinkBase="Shared\Storage" />
     <Compile Remove="$(AzureStorageSharedTestSources)\AzuriteFixture.cs" />
     <Compile Remove="$(AzureStorageSharedTestSources)\AzuriteNUnitFixture.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)..\..\Azure.Storage.Blobs\tests\BlobsClientTestFixtureAttribute.cs" LinkBase="Shared" />

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

@@ -8,6 +8,7 @@ internal DataTransfer() { }
         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 void EnsureCompleted(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { }
+        public System.Threading.Tasks.Task<bool> TryPauseAsync(System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
     }
     [System.FlagsAttribute]
     public enum ErrorHandlingOptions
@@ -106,6 +107,8 @@ public enum StorageTransferStatus
         Completed = 4,
         CompletedWithSkippedTransfers = 5,
         CompletedWithFailedTransfers = 6,
+        PauseInProgress = 7,
+        CancellationInProgress = 8,
     }
     public partial class TransferManager : System.IAsyncDisposable
     {
@@ -114,10 +117,13 @@ public TransferManager(Azure.Storage.DataMovement.TransferManagerOptions options
         public virtual System.Threading.Tasks.Task<Azure.Storage.DataMovement.DataTransfer> StartTransferAsync(Azure.Storage.DataMovement.StorageResource sourceResource, Azure.Storage.DataMovement.StorageResource destinationResource, Azure.Storage.DataMovement.Models.SingleTransferOptions transferOptions = null) { throw null; }
         public virtual System.Threading.Tasks.Task<Azure.Storage.DataMovement.DataTransfer> StartTransferAsync(Azure.Storage.DataMovement.StorageResourceContainer sourceResource, Azure.Storage.DataMovement.StorageResourceContainer destinationResource, Azure.Storage.DataMovement.Models.ContainerTransferOptions transferOptions = null) { throw null; }
         System.Threading.Tasks.ValueTask System.IAsyncDisposable.DisposeAsync() { throw null; }
+        public virtual System.Threading.Tasks.Task<bool> TryPauseTransferAsync(Azure.Storage.DataMovement.DataTransfer transfer, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
+        public virtual System.Threading.Tasks.Task<bool> TryPauseTransferAsync(string transferId, System.Threading.CancellationToken cancellationToken = default(System.Threading.CancellationToken)) { throw null; }
     }
     public partial class TransferManagerOptions
     {
         public TransferManagerOptions() { }
+        public Azure.Storage.DataMovement.Models.TransferCheckpointerOptions CheckpointerOptions { get { throw null; } set { } }
         public Azure.Storage.DataMovement.ErrorHandlingOptions ErrorHandling { get { throw null; } set { } }
         public int? MaximumConcurrency { get { throw null; } set { } }
     }
@@ -135,6 +141,7 @@ public ContainerTransferOptions() { }
         public Azure.Storage.DataMovement.StorageResourceCreateMode CreateMode { get { throw null; } set { } }
         public long? InitialTransferSize { get { throw null; } set { } }
         public long? MaximumTransferChunkSize { get { throw null; } set { } }
+        public string ResumeFromCheckpointId { get { throw null; } set { } }
         public event Azure.Core.SyncAsyncEventHandler<Azure.Storage.DataMovement.Models.SingleTransferCompletedEventArgs> SingleTransferCompleted { add { } remove { } }
         public event Azure.Core.SyncAsyncEventHandler<Azure.Storage.DataMovement.Models.TransferFailedEventArgs> TransferFailed { add { } remove { } }
         public event Azure.Core.SyncAsyncEventHandler<Azure.Storage.DataMovement.Models.TransferSkippedEventArgs> TransferSkipped { add { } remove { } }
@@ -174,6 +181,7 @@ public SingleTransferOptions() { }
         public Azure.Storage.DataMovement.StorageResourceCreateMode CreateMode { get { throw null; } set { } }
         public long? InitialTransferSize { get { throw null; } set { } }
         public long? MaximumTransferChunkSize { get { throw null; } set { } }
+        public string ResumeFromCheckpointId { get { throw null; } set { } }
         public event Azure.Core.SyncAsyncEventHandler<Azure.Storage.DataMovement.Models.TransferFailedEventArgs> TransferFailed { add { } remove { } }
         public event Azure.Core.SyncAsyncEventHandler<Azure.Storage.DataMovement.Models.TransferSkippedEventArgs> TransferSkipped { add { } remove { } }
         public event Azure.Core.SyncAsyncEventHandler<Azure.Storage.DataMovement.Models.TransferStatusEventArgs> TransferStatus { add { } remove { } }
@@ -206,6 +214,10 @@ public partial class StorageResourceWriteToOffsetOptions
         public StorageResourceWriteToOffsetOptions() { }
         public string BlockId { get { throw null; } }
     }
+    public partial class TransferCheckpointerOptions
+    {
+        public TransferCheckpointerOptions(string localCheckpointerPath) { }
+    }
     public enum TransferCopyMethod
     {
         None = 0,

sdk/storage/Azure.Storage.DataMovement/src/Azure.Storage.DataMovement.csproj

@@ -1,4 +1,4 @@
-<Project Sdk="Microsoft.NET.Sdk">
+<Project Sdk="Microsoft.NET.Sdk">
   <PropertyGroup>
     <RequiredTargetFrameworks>netstandard2.0</RequiredTargetFrameworks>
     <TargetFrameworks>$(RequiredTargetFrameworks)</TargetFrameworks>

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

@@ -5,12 +5,12 @@
 using System.Collections.Generic;
 using System.Globalization;
 using System.IO;
-using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
-using System.Xml.Linq;
 using Azure.Core;
 using Azure.Storage.DataMovement.Models;
+using Azure.Storage.DataMovement.JobPlanModels;
+using System.Linq;
 
 namespace Azure.Storage.DataMovement
 {
@@ -36,7 +36,7 @@ internal abstract class JobPartInternal
         ///
         /// Will be disposed of once all tasks of the job have completed or have been cancelled.
         /// </summary>
-        internal CancellationTokenSource _cancellationTokenSource { get; set; }
+        internal CancellationToken _cancellationToken { get; set; }
 
         /// <summary>
         /// Plan file writer for the respective job.
@@ -119,6 +119,9 @@ internal abstract class JobPartInternal
         /// </summary>
         public SyncAsyncEventHandler<SingleTransferCompletedEventArgs> SingleTransferCompletedEventHandler { get; internal set; }
 
+        private List<Task<bool>> _chunkTasks;
+        private List<TaskCompletionSource<bool>> _chunkTaskSources;
+
         /// <summary>
         /// Array pools for reading from streams to upload
         /// </summary>
@@ -143,18 +146,19 @@ internal JobPartInternal(
             SyncAsyncEventHandler<TransferFailedEventArgs> failedEventHandler,
             SyncAsyncEventHandler<TransferSkippedEventArgs> skippedEventHandler,
             SyncAsyncEventHandler<SingleTransferCompletedEventArgs> singleTransferEventHandler,
-            CancellationTokenSource cancellationTokenSource,
+            CancellationToken cancellationToken,
+            StorageTransferStatus jobPartStatus = StorageTransferStatus.Queued,
             long? length = default)
         {
-            JobPartStatus = StorageTransferStatus.Queued;
+            JobPartStatus = jobPartStatus;
             PartNumber = partNumber;
             _dataTransfer = dataTransfer;
             _sourceResource = sourceResource;
             _destinationResource = destinationResource;
             _errorHandling = errorHandling;
             _createMode = createMode;
             _checkpointer = checkpointer;
-            _cancellationTokenSource = cancellationTokenSource;
+            _cancellationToken = cancellationToken;
             _arrayPool = arrayPool;
             PartTransferStatusEventHandler = jobPartEventHandler;
             TransferStatusEventHandler = statusEventHandler;
@@ -178,26 +182,75 @@ internal JobPartInternal(
             }
 
             Length = length;
+            _chunkTasks = new List<Task<bool>>();
+            _chunkTaskSources = new List<TaskCompletionSource<bool>>();
         }
 
         public void SetQueueChunkDelegate(QueueChunkDelegate chunkDelegate)
         {
             QueueChunk = chunkDelegate;
         }
 
+        /// <summary>
+        /// Queues the task to the main chunk channel and also appends the tracking
+        /// completion source to the task. So we know the state of each chunk especially
+        /// when we're looking to stop/pause the job part.
+        /// </summary>
+        /// <returns></returns>
+        public async Task QueueChunkToChannelAsync(Task chunkTask)
+        {
+            // Attach TaskCompletionSource
+            TaskCompletionSource<bool> chunkCompleted = new TaskCompletionSource<bool>(
+                false,
+                TaskCreationOptions.RunContinuationsAsynchronously);
+            _chunkTaskSources.Add(chunkCompleted);
+            _chunkTasks.Add(chunkCompleted.Task);
+
+            await QueueChunk(
+                async () =>
+                {
+                    await chunkTask.ConfigureAwait(false);
+                    chunkCompleted.SetResult(true);
+                    await CheckAndUpdateCancellationStatusAsync().ConfigureAwait(false);
+                }).ConfigureAwait(false);
+        }
+
         /// <summary>
         /// Processes the job to job parts
         /// </summary>
         /// <returns>An IEnumerable that contains the job chunks</returns>
         public abstract Task ProcessPartToChunkAsync();
 
-        internal async Task TriggerCancellation(StorageTransferStatus status)
+        /// <summary>
+        /// Triggers the cancellation for the Job Part.
+        ///
+        /// If the status is set to <see cref="StorageTransferStatus.Paused"/>
+        /// and any chunks is still processing to be cancelled is will be set to <see cref="StorageTransferStatus.PauseInProgress"/>
+        /// until the chunks finish then it will be set to <see cref="StorageTransferStatus.Paused"/>.
+        ///
+        /// If the status is set to <see cref="StorageTransferStatus.CompletedWithFailedTransfers"/>
+        /// and any chunks is still processing to be cancelled is will be set to <see cref="StorageTransferStatus.CancellationInProgress"/>
+        /// until the chunks finish then it will be set to <see cref="StorageTransferStatus.CompletedWithFailedTransfers"/>.
+        /// </summary>
+        /// <returns></returns>
+        internal async Task TriggerCancellationAsync()
         {
-            if (!_cancellationTokenSource.IsCancellationRequested)
+            if (!_cancellationToken.IsCancellationRequested)
             {
-                _cancellationTokenSource.Cancel();
+                _dataTransfer._state.TriggerCancellation();
+            }
+            // Set the status to Pause/CancellationInProgress
+            if (StorageTransferStatus.PauseInProgress == _dataTransfer.TransferStatus)
+            {
+                // It's possible that the status hasn't propagated down to the job part
+                // status yet here since we pause from the data transfer object.
+                await OnTransferStatusChanged(StorageTransferStatus.PauseInProgress).ConfigureAwait(false);
+            }
+            else
+            {
+                // It's a cancellation if a pause wasn't called.
+                await OnTransferStatusChanged(StorageTransferStatus.CancellationInProgress).ConfigureAwait(false);
             }
-            await OnTransferStatusChanged(status).ConfigureAwait(false);
             _dataTransfer._state.ResetTransferredBytes();
         }
 
@@ -223,12 +276,15 @@ internal async Task OnTransferStatusChanged(StorageTransferStatus transferStatus
                 {
                     await InvokeSingleCompletedArg().ConfigureAwait(false);
                 }
+                // Set the status in the checkpointer
+                await SetCheckpointerStatus(transferStatus).ConfigureAwait(false);
+
                 // TODO: change to RaiseAsync
                 await PartTransferStatusEventHandler.Invoke(new TransferStatusEventArgs(
                     _dataTransfer.Id,
                     transferStatus,
                     false,
-                    _cancellationTokenSource.Token)).ConfigureAwait(false);
+                    _cancellationToken)).ConfigureAwait(false);
             }
         }
 
@@ -251,12 +307,12 @@ await SingleTransferCompletedEventHandler.Invoke(
                         _sourceResource,
                         _destinationResource,
                         false,
-                        _cancellationTokenSource.Token)).ConfigureAwait(false);
+                        _cancellationToken)).ConfigureAwait(false);
             }
         }
 
         /// <summary>
-        /// Invokes Failed Argument
+        /// Invokes Skipped Argument Event.
         /// </summary>
         public async virtual Task InvokeSkippedArg()
         {
@@ -268,13 +324,13 @@ await TransferSkippedEventHandler.Invoke(new TransferSkippedEventArgs(
                     _sourceResource,
                     _destinationResource,
                     false,
-                    _cancellationTokenSource.Token)).ConfigureAwait(false);
+                    _cancellationToken)).ConfigureAwait(false);
             }
             await OnTransferStatusChanged(StorageTransferStatus.CompletedWithSkippedTransfers).ConfigureAwait(false);
         }
 
         /// <summary>
-        /// Invokes Failed Argument
+        /// Invokes Failed Argument Event.
         /// </summary>
         public async virtual Task InvokeFailedArg(Exception ex)
         {
@@ -287,10 +343,34 @@ await TransferFailedEventHandler.Invoke(new TransferFailedEventArgs(
                     _destinationResource,
                     ex,
                     false,
-                    _cancellationTokenSource.Token)).ConfigureAwait(false);
+                    _cancellationToken)).ConfigureAwait(false);
             }
             // Trigger job cancellation if the failed handler is enabled
-            await TriggerCancellation(StorageTransferStatus.CompletedWithFailedTransfers).ConfigureAwait(false);
+            await TriggerCancellationAsync().ConfigureAwait(false);
+            await CheckAndUpdateCancellationStatusAsync().ConfigureAwait(false);
+        }
+
+        public async virtual Task AddJobPartToCheckpointer(int chunksTotal)
+        {
+            JobPartPlanHeader header = this.ToJobPartPlanHeader(StorageTransferStatus.InProgress);
+            using (Stream stream = new MemoryStream())
+            {
+                header.Serialize(stream);
+                await _checkpointer.AddNewJobPartAsync(
+                        transferId: _dataTransfer.Id,
+                        partNumber: PartNumber,
+                        chunksTotal: chunksTotal,
+                        headerStream: stream,
+                        cancellationToken: _cancellationToken).ConfigureAwait(false);
+            }
+        }
+
+        internal async virtual Task SetCheckpointerStatus(StorageTransferStatus status)
+        {
+            await _checkpointer.SetJobPartTransferStatusAsync(
+                transferId: _dataTransfer.Id,
+                partNumber: PartNumber,
+                status: status).ConfigureAwait(false);
         }
 
         internal long CalculateBlockSize(long length)
@@ -381,5 +461,20 @@ internal static long ParseRangeTotalLength(string range)
                 absolutePosition += blockLength;
             }
         }
+
+        internal async Task CheckAndUpdateCancellationStatusAsync()
+        {
+            if (_chunkTasks.All((Task task) => (task.IsCompleted)))
+            {
+                if (_dataTransfer.TransferStatus == StorageTransferStatus.PauseInProgress)
+                {
+                    await OnTransferStatusChanged(StorageTransferStatus.Paused).ConfigureAwait(false);
+                }
+                else if (_dataTransfer.TransferStatus == StorageTransferStatus.CancellationInProgress)
+                {
+                    await OnTransferStatusChanged(StorageTransferStatus.CompletedWithFailedTransfers).ConfigureAwait(false);
+                }
+            }
+        }
     }
 }