Checkpointer testing cont. (Azure#47125)

jaschrep-msft · web-flow · commit 748c9c467ae0 · 2024-11-13T11:27:19.000-05:00
* change checkpointer API

Accept header object instead of stream

* move serialze closer to checkpoint write

* error report bug fix

* adapt tests to new interface

* exportapi
diff --git a/sdk/storage/Azure.Storage.DataMovement.Blobs/tests/Azure.Storage.DataMovement.Blobs.Tests.csproj b/sdk/storage/Azure.Storage.DataMovement.Blobs/tests/Azure.Storage.DataMovement.Blobs.Tests.csproj
@@ -41,6 +41,7 @@
     <Compile Include="$(AzureStorageDataMovementTestSharedSources)TestEventsRaised.cs" LinkBase="Shared\DataMovement" />
     <Compile Include="$(AzureStorageDataMovementTestSharedSources)TestTransferWithTimeout.cs" LinkBase="Shared\DataMovement" />
     <Compile Include="$(AzureStorageDataMovementTestSharedSources)DisposingLocalDirectory.cs" LinkBase="Shared\DataMovement" />
+    <Compile Include="$(AzureStorageDataMovementTestSharedSources)MemoryTransferCheckpointer.cs" LinkBase="Shared\DataMovement" />
     <Compile Include="$(AzureStorageDataMovementTestSharedSources)MockQueueInternalTasks.cs" LinkBase="Shared\DataMovement" />
     <Compile Include="$(AzureStorageDataMovementTestSharedSources)MockResourceCheckpointData.cs" LinkBase="Shared\DataMovement" />
     <Compile Include="$(AzureStorageDataMovementTestSharedSources)StartTransferUploadTestBase.cs" LinkBase="Shared\DataMovement" />
diff --git a/sdk/storage/Azure.Storage.DataMovement/api/Azure.Storage.DataMovement.net6.0.cs b/sdk/storage/Azure.Storage.DataMovement/api/Azure.Storage.DataMovement.net6.0.cs
@@ -237,10 +237,10 @@ public partial class TransferItemCompletedEventArgs : Azure.Storage.DataMovement
     }
     public partial class TransferItemFailedEventArgs : Azure.Storage.DataMovement.DataTransferEventArgs
     {
-        public TransferItemFailedEventArgs(string transferId, Azure.Storage.DataMovement.StorageResourceItem sourceResource, Azure.Storage.DataMovement.StorageResourceItem destinationResource, System.Exception exception, bool isRunningSynchronously, System.Threading.CancellationToken cancellationToken) : base (default(string), default(bool), default(System.Threading.CancellationToken)) { }
-        public Azure.Storage.DataMovement.StorageResourceItem DestinationResource { get { throw null; } }
+        public TransferItemFailedEventArgs(string transferId, Azure.Storage.DataMovement.StorageResource sourceResource, Azure.Storage.DataMovement.StorageResource destinationResource, System.Exception exception, bool isRunningSynchronously, System.Threading.CancellationToken cancellationToken) : base (default(string), default(bool), default(System.Threading.CancellationToken)) { }
+        public Azure.Storage.DataMovement.StorageResource DestinationResource { get { throw null; } }
         public System.Exception Exception { get { throw null; } }
-        public Azure.Storage.DataMovement.StorageResourceItem SourceResource { get { throw null; } }
+        public Azure.Storage.DataMovement.StorageResource SourceResource { get { throw null; } }
     }
     public partial class TransferItemSkippedEventArgs : Azure.Storage.DataMovement.DataTransferEventArgs
     {
diff --git a/sdk/storage/Azure.Storage.DataMovement/api/Azure.Storage.DataMovement.netstandard2.0.cs b/sdk/storage/Azure.Storage.DataMovement/api/Azure.Storage.DataMovement.netstandard2.0.cs
@@ -237,10 +237,10 @@ public partial class TransferItemCompletedEventArgs : Azure.Storage.DataMovement
     }
     public partial class TransferItemFailedEventArgs : Azure.Storage.DataMovement.DataTransferEventArgs
     {
-        public TransferItemFailedEventArgs(string transferId, Azure.Storage.DataMovement.StorageResourceItem sourceResource, Azure.Storage.DataMovement.StorageResourceItem destinationResource, System.Exception exception, bool isRunningSynchronously, System.Threading.CancellationToken cancellationToken) : base (default(string), default(bool), default(System.Threading.CancellationToken)) { }
-        public Azure.Storage.DataMovement.StorageResourceItem DestinationResource { get { throw null; } }
+        public TransferItemFailedEventArgs(string transferId, Azure.Storage.DataMovement.StorageResource sourceResource, Azure.Storage.DataMovement.StorageResource destinationResource, System.Exception exception, bool isRunningSynchronously, System.Threading.CancellationToken cancellationToken) : base (default(string), default(bool), default(System.Threading.CancellationToken)) { }
+        public Azure.Storage.DataMovement.StorageResource DestinationResource { get { throw null; } }
         public System.Exception Exception { get { throw null; } }
-        public Azure.Storage.DataMovement.StorageResourceItem SourceResource { get { throw null; } }
+        public Azure.Storage.DataMovement.StorageResource SourceResource { get { throw null; } }
     }
     public partial class TransferItemSkippedEventArgs : Azure.Storage.DataMovement.DataTransferEventArgs
     {
diff --git a/sdk/storage/Azure.Storage.DataMovement/src/DisabledTransferCheckpointer.cs b/sdk/storage/Azure.Storage.DataMovement/src/DisabledTransferCheckpointer.cs
@@ -16,7 +16,7 @@ public Task AddNewJobAsync(string transferId, StorageResource source, StorageRes
             return Task.CompletedTask;
         }
 
-        public Task AddNewJobPartAsync(string transferId, int partNumber, Stream headerStream, CancellationToken cancellationToken = default)
+        public Task AddNewJobPartAsync(string transferId, int partNumber, JobPartPlanHeader header, CancellationToken cancellationToken = default)
         {
             return Task.CompletedTask;
         }
diff --git a/sdk/storage/Azure.Storage.DataMovement/src/ITransferCheckpointer.cs b/sdk/storage/Azure.Storage.DataMovement/src/ITransferCheckpointer.cs
@@ -20,7 +20,7 @@ Task AddNewJobAsync(
         Task AddNewJobPartAsync(
             string transferId,
             int partNumber,
-            Stream headerStream,
+            JobPartPlanHeader header,
             CancellationToken cancellationToken = default);
 
         Task<bool> IsEnumerationCompleteAsync(
diff --git a/sdk/storage/Azure.Storage.DataMovement/src/JobPartInternal.cs b/sdk/storage/Azure.Storage.DataMovement/src/JobPartInternal.cs
@@ -473,16 +473,11 @@ public async virtual Task CleanupAbortedJobPartAsync()
         /// </summary>
         public async virtual Task AddJobPartToCheckpointerAsync()
         {
-            JobPartPlanHeader header = this.ToJobPartPlanHeader();
-            using (Stream stream = new MemoryStream())
-            {
-                header.Serialize(stream);
-                await _checkpointer.AddNewJobPartAsync(
-                    transferId: _dataTransfer.Id,
-                    partNumber: PartNumber,
-                    headerStream: stream,
-                    cancellationToken: _cancellationToken).ConfigureAwait(false);
-            }
+            await _checkpointer.AddNewJobPartAsync(
+                transferId: _dataTransfer.Id,
+                partNumber: PartNumber,
+                header: this.ToJobPartPlanHeader(),
+                cancellationToken: _cancellationToken).ConfigureAwait(false);
         }
 
         internal async virtual Task SetCheckpointerStatusAsync()
diff --git a/sdk/storage/Azure.Storage.DataMovement/src/JobPlan/JobPartPlanFile.cs b/sdk/storage/Azure.Storage.DataMovement/src/JobPlan/JobPartPlanFile.cs
@@ -38,21 +38,21 @@ public static async Task<JobPartPlanFile> CreateJobPartPlanFileAsync(
             string checkpointerPath,
             string id,
             int jobPart,
-            Stream headerStream,
+            JobPartPlanHeader header,
             CancellationToken cancellationToken = default)
         {
             Argument.AssertNotNullOrEmpty(checkpointerPath, nameof(checkpointerPath));
             Argument.AssertNotNullOrEmpty(id, nameof(id));
             Argument.AssertNotNull(jobPart, nameof(jobPart));
-            Argument.AssertNotNull(headerStream, nameof(headerStream));
+            Argument.AssertNotNull(header, nameof(header));
 
             JobPartPlanFileName fileName = new JobPartPlanFileName(checkpointerPath: checkpointerPath, id: id, jobPartNumber: jobPart);
-            return await CreateJobPartPlanFileAsync(fileName, headerStream, cancellationToken).ConfigureAwait(false);
+            return await CreateJobPartPlanFileAsync(fileName, header, cancellationToken).ConfigureAwait(false);
         }
 
         public static async Task<JobPartPlanFile> CreateJobPartPlanFileAsync(
             JobPartPlanFileName fileName,
-            Stream headerStream,
+            JobPartPlanHeader header,
             CancellationToken cancellationToken = default)
         {
             JobPartPlanFile result = new JobPartPlanFile()
@@ -63,8 +63,11 @@ public static async Task<JobPartPlanFile> CreateJobPartPlanFileAsync(
             try
             {
                 using (FileStream fileStream = File.Create(result.FileName.ToString()))
+                using (MemoryStream ms = new())
                 {
-                    await headerStream.CopyToAsync(
+                    header.Serialize(ms);
+                    ms.Position = 0;
+                    await ms.CopyToAsync(
                         fileStream,
                         DataMovementConstants.DefaultStreamCopyBufferSize,
                         cancellationToken).ConfigureAwait(false);
diff --git a/sdk/storage/Azure.Storage.DataMovement/src/LocalTransferCheckpointer.cs b/sdk/storage/Azure.Storage.DataMovement/src/LocalTransferCheckpointer.cs
@@ -104,14 +104,13 @@ public override async Task AddNewJobAsync(
         public override async Task AddNewJobPartAsync(
             string transferId,
             int partNumber,
-            Stream headerStream,
+            JobPartPlanHeader header,
             CancellationToken cancellationToken = default)
         {
             Argument.AssertNotNullOrEmpty(transferId, nameof(transferId));
             Argument.AssertNotNull(partNumber, nameof(partNumber));
-            Argument.AssertNotNull(headerStream, nameof(headerStream));
+            Argument.AssertNotNull(header, nameof(header));
             CancellationHelper.ThrowIfCancellationRequested(cancellationToken);
-            headerStream.Position = 0;
 
             if (!_transferStates.ContainsKey(transferId))
             {
@@ -124,7 +123,7 @@ public override async Task AddNewJobPartAsync(
                 _pathToCheckpointer,
                 transferId,
                 partNumber,
-                headerStream,
+                header,
                 cancellationToken).ConfigureAwait(false);
 
             // Add the job part into the current state
diff --git a/sdk/storage/Azure.Storage.DataMovement/src/SerializerTransferCheckpointer.cs b/sdk/storage/Azure.Storage.DataMovement/src/SerializerTransferCheckpointer.cs
@@ -45,15 +45,15 @@ public abstract Task AddNewJobAsync(
         /// </summary>
         /// <param name="transferId">The transfer ID.</param>
         /// <param name="partNumber">The job part number.</param>
-        /// <param name="headerStream">A <see cref="Stream"/> to the job part plan header.</param>
+        /// <param name="header">A <see cref="Stream"/> to the job part plan header.</param>
         /// <param name="cancellationToken">
         /// Optional <see cref="CancellationToken"/> to propagate
         /// notifications that the operation should be canceled.
         /// </param>
         public abstract Task AddNewJobPartAsync(
             string transferId,
             int partNumber,
-            Stream headerStream,
+            JobPartPlanHeader header,
             CancellationToken cancellationToken = default);
 
         /// <summary>
diff --git a/sdk/storage/Azure.Storage.DataMovement/src/TransferItemFailedEventArgs.cs b/sdk/storage/Azure.Storage.DataMovement/src/TransferItemFailedEventArgs.cs
@@ -13,14 +13,14 @@ namespace Azure.Storage.DataMovement
     public class TransferItemFailedEventArgs : DataTransferEventArgs
     {
         /// <summary>
-        /// Gets the <see cref="StorageResourceItem"/> that was the source resource for the transfer.
+        /// Gets the <see cref="StorageResource"/> that was the source resource for the transfer.
         /// </summary>
-        public StorageResourceItem SourceResource { get; }
+        public StorageResource SourceResource { get; }
 
         /// <summary>
-        /// Gets the <see cref="StorageResourceItem"/> that was the destination resource for the transfer.
+        /// Gets the <see cref="StorageResource"/> that was the destination resource for the transfer.
         /// </summary>
-        public StorageResourceItem DestinationResource { get; }
+        public StorageResource DestinationResource { get; }
 
         /// <summary>
         /// Gets the <see cref="Exception"/> that was thrown during the job.
@@ -53,8 +53,8 @@ public class TransferItemFailedEventArgs : DataTransferEventArgs
         /// </exception>
         public TransferItemFailedEventArgs(
             string transferId,
-            StorageResourceItem sourceResource,
-            StorageResourceItem destinationResource,
+            StorageResource sourceResource,
+            StorageResource destinationResource,
             Exception exception,
             bool isRunningSynchronously,
             CancellationToken cancellationToken)
diff --git a/sdk/storage/Azure.Storage.DataMovement/src/TransferJobInternal.cs b/sdk/storage/Azure.Storage.DataMovement/src/TransferJobInternal.cs
@@ -493,8 +493,8 @@ ex is not TaskCanceledException &&
                     await TransferFailedEventHandler.RaiseAsync(
                         new TransferItemFailedEventArgs(
                             _dataTransfer.Id,
-                            _sourceResource,
-                            _destinationResource,
+                            (StorageResource)_sourceResource ?? _sourceResourceContainer,
+                            (StorageResource)_destinationResource ?? _destinationResourceContainer,
                             ex,
                             false,
                             _cancellationToken),
diff --git a/sdk/storage/Azure.Storage.DataMovement/tests/JobPartPlanFileTests.cs b/sdk/storage/Azure.Storage.DataMovement/tests/JobPartPlanFileTests.cs
@@ -31,7 +31,20 @@ public async Task CreateJobPartPlanFileAsync_Base()
                     checkpointerPath: test.DirectoryPath,
                     id: transferId,
                     jobPart: jobPart,
-                    headerStream: stream);
+                    header: new(
+                        DataMovementConstants.JobPartPlanFile.SchemaVersion_b3,
+                        transferId,
+                        jobPart,
+                        System.DateTimeOffset.Now,
+                        "mock",
+                        "mock",
+                        "mock",
+                        "mock",
+                        default,
+                        default,
+                        default,
+                        default,
+                        new()));
             }
 
             JobPartPlanFileName fileName = new JobPartPlanFileName(
@@ -63,7 +76,20 @@ public async Task CreateJobPartPlanFileAsync_FileName()
             {
                 file = await JobPartPlanFile.CreateJobPartPlanFileAsync(
                    fileName: fileName,
-                   headerStream: stream);
+                   header: new(
+                        DataMovementConstants.JobPartPlanFile.SchemaVersion_b3,
+                        transferId,
+                        jobPart,
+                        System.DateTimeOffset.Now,
+                        "mock",
+                        "mock",
+                        "mock",
+                        "mock",
+                        default,
+                        default,
+                        default,
+                        default,
+                        new()));
             }
 
             Assert.NotNull(file);
diff --git a/sdk/storage/Azure.Storage.DataMovement/tests/LocalTransferCheckpointerTests.cs b/sdk/storage/Azure.Storage.DataMovement/tests/LocalTransferCheckpointerTests.cs
diff --git a/sdk/storage/Azure.Storage.DataMovement/tests/Shared/MemoryTransferCheckpointer.cs b/sdk/storage/Azure.Storage.DataMovement/tests/Shared/MemoryTransferCheckpointer.cs

Original file line number	Diff line number	Diff line change
`@@ -237,10 +237,10 @@ public partial class TransferItemCompletedEventArgs : Azure.Storage.DataMovement`
`237`	`237`	`}`
`238`	`238`	`public partial class TransferItemFailedEventArgs : Azure.Storage.DataMovement.DataTransferEventArgs`
`239`	`239`	`{`
`240`		`- public TransferItemFailedEventArgs(string transferId, Azure.Storage.DataMovement.StorageResourceItem sourceResource, Azure.Storage.DataMovement.StorageResourceItem destinationResource, System.Exception exception, bool isRunningSynchronously, System.Threading.CancellationToken cancellationToken) : base (default(string), default(bool), default(System.Threading.CancellationToken)) { }`
`241`		`- public Azure.Storage.DataMovement.StorageResourceItem DestinationResource { get { throw null; } }`
	`240`	`+ public TransferItemFailedEventArgs(string transferId, Azure.Storage.DataMovement.StorageResource sourceResource, Azure.Storage.DataMovement.StorageResource destinationResource, System.Exception exception, bool isRunningSynchronously, System.Threading.CancellationToken cancellationToken) : base (default(string), default(bool), default(System.Threading.CancellationToken)) { }`
	`241`	`+ public Azure.Storage.DataMovement.StorageResource DestinationResource { get { throw null; } }`
`242`	`242`	`public System.Exception Exception { get { throw null; } }`
`243`		`- public Azure.Storage.DataMovement.StorageResourceItem SourceResource { get { throw null; } }`
	`243`	`+ public Azure.Storage.DataMovement.StorageResource SourceResource { get { throw null; } }`
`244`	`244`	`}`
`245`	`245`	`public partial class TransferItemSkippedEventArgs : Azure.Storage.DataMovement.DataTransferEventArgs`
`246`	`246`	`{`
Original file line number	Diff line number	Diff line change
`@@ -16,7 +16,7 @@ public Task AddNewJobAsync(string transferId, StorageResource source, StorageRes`
`16`	`16`	`return Task.CompletedTask;`
`17`	`17`	`}`
`18`	`18`
`19`		`- public Task AddNewJobPartAsync(string transferId, int partNumber, Stream headerStream, CancellationToken cancellationToken = default)`
	`19`	`+ public Task AddNewJobPartAsync(string transferId, int partNumber, JobPartPlanHeader header, CancellationToken cancellationToken = default)`
`20`	`20`	`{`
`21`	`21`	`return Task.CompletedTask;`
`22`	`22`	`}`