[Storage][DataMovement] Adjust local file buffer size (Azure#47043)

jalauzon-msft · web-flow · commit 338bbe735822 · 2024-11-12T10:33:58.000-08:00
diff --git a/sdk/storage/Azure.Storage.DataMovement/src/JobPlan/JobPartPlanFile.cs b/sdk/storage/Azure.Storage.DataMovement/src/JobPlan/JobPartPlanFile.cs
@@ -29,8 +29,6 @@ internal class JobPartPlanFile : IDisposable
         /// </summary>
         public readonly SemaphoreSlim WriteLock;
 
-        private const int DefaultBufferSize = 81920;
-
         private JobPartPlanFile()
         {
             WriteLock = new SemaphoreSlim(1);
@@ -66,7 +64,10 @@ public static async Task<JobPartPlanFile> CreateJobPartPlanFileAsync(
             {
                 using (FileStream fileStream = File.Create(result.FileName.ToString()))
                 {
-                    await headerStream.CopyToAsync(fileStream, DefaultBufferSize, cancellationToken).ConfigureAwait(false);
+                    await headerStream.CopyToAsync(
+                        fileStream,
+                        DataMovementConstants.DefaultStreamCopyBufferSize,
+                        cancellationToken).ConfigureAwait(false);
                 }
             }
             catch (Exception)
diff --git a/sdk/storage/Azure.Storage.DataMovement/src/JobPlan/JobPlanFile.cs b/sdk/storage/Azure.Storage.DataMovement/src/JobPlan/JobPlanFile.cs
@@ -35,8 +35,6 @@ internal class JobPlanFile : IDisposable
         /// </summary>
         public readonly SemaphoreSlim WriteLock;
 
-        private const int DefaultBufferSize = 81920;
-
         private JobPlanFile(string id, string filePath)
         {
             Id = id;
@@ -63,7 +61,10 @@ public static async Task<JobPlanFile> CreateJobPlanFileAsync(
             {
                 using (FileStream fileStream = File.Create(jobPlanFile.FilePath))
                 {
-                    await headerStream.CopyToAsync(fileStream, DefaultBufferSize, cancellationToken).ConfigureAwait(false);
+                    await headerStream.CopyToAsync(
+                        fileStream,
+                        DataMovementConstants.DefaultStreamCopyBufferSize,
+                        cancellationToken).ConfigureAwait(false);
                 }
             }
             catch (Exception)
diff --git a/sdk/storage/Azure.Storage.DataMovement/src/LocalFileStorageResource.cs b/sdk/storage/Azure.Storage.DataMovement/src/LocalFileStorageResource.cs
@@ -142,11 +142,9 @@ protected internal override async Task CopyFromStreamAsync(
                     {
                         fileStream.Seek(position, SeekOrigin.Begin);
                     }
-                    await stream.CopyToAsync(
-                        fileStream,
-                        (int)streamLength,
-                        cancellationToken)
-                        .ConfigureAwait(false);
+
+                    int bufferSize = Math.Min((int)streamLength, DataMovementConstants.DefaultStreamCopyBufferSize);
+                    await stream.CopyToAsync(fileStream, bufferSize, cancellationToken).ConfigureAwait(false);
                 }
             }
         }
diff --git a/sdk/storage/Azure.Storage.DataMovement/src/LocalTransferCheckpointer.cs b/sdk/storage/Azure.Storage.DataMovement/src/LocalTransferCheckpointer.cs
@@ -152,8 +152,8 @@ public override async Task<Stream> ReadJobPlanFileAsync(
             int length,
             CancellationToken cancellationToken = default)
         {
-            int maxArraySize = length > 0 ? length : DataMovementConstants.DefaultArrayPoolArraySize;
-            Stream copiedStream = new PooledMemoryStream(ArrayPool<byte>.Shared, maxArraySize);
+            int bufferSize = length > 0 ? length : DataMovementConstants.DefaultStreamCopyBufferSize;
+            Stream copiedStream = new PooledMemoryStream(ArrayPool<byte>.Shared, bufferSize);
 
             CancellationHelper.ThrowIfCancellationRequested(cancellationToken);
             if (_transferStates.TryGetValue(transferId, out JobPlanFile jobPlanFile))
@@ -164,7 +164,7 @@ public override async Task<Stream> ReadJobPlanFileAsync(
                     using (MemoryMappedFile mmf = MemoryMappedFile.CreateFromFile(jobPlanFile.FilePath))
                     using (MemoryMappedViewStream mmfStream = mmf.CreateViewStream(offset, length, MemoryMappedFileAccess.Read))
                     {
-                        await mmfStream.CopyToAsync(copiedStream).ConfigureAwait(false);
+                        await mmfStream.CopyToAsync(copiedStream, bufferSize, cancellationToken).ConfigureAwait(false);
                     }
 
                     copiedStream.Position = 0;
@@ -193,16 +193,16 @@ public override async Task<Stream> ReadJobPartPlanFileAsync(
             {
                 if (jobPlanFile.JobParts.TryGetValue(partNumber, out JobPartPlanFile jobPartPlanFile))
                 {
-                    int maxArraySize = length > 0 ? length : DataMovementConstants.DefaultArrayPoolArraySize;
-                    Stream copiedStream = new PooledMemoryStream(ArrayPool<byte>.Shared, maxArraySize);
+                    int bufferSize = length > 0 ? length : DataMovementConstants.DefaultStreamCopyBufferSize;
+                    Stream copiedStream = new PooledMemoryStream(ArrayPool<byte>.Shared, bufferSize);
 
                     await jobPartPlanFile.WriteLock.WaitAsync(cancellationToken).ConfigureAwait(false);
                     try
                     {
                         using (MemoryMappedFile mmf = MemoryMappedFile.CreateFromFile(jobPartPlanFile.FilePath))
                         using (MemoryMappedViewStream mmfStream = mmf.CreateViewStream(offset, length, MemoryMappedFileAccess.Read))
                         {
-                            await mmfStream.CopyToAsync(copiedStream).ConfigureAwait(false);
+                            await mmfStream.CopyToAsync(copiedStream, bufferSize, cancellationToken).ConfigureAwait(false);
                         }
 
                         copiedStream.Position = 0;
diff --git a/sdk/storage/Azure.Storage.DataMovement/src/Shared/DataMovementConstants.cs b/sdk/storage/Azure.Storage.DataMovement/src/Shared/DataMovementConstants.cs
@@ -14,7 +14,7 @@ internal class DataMovementConstants
         internal const int MaxJobPartReaders = 64;
         internal const int MaxJobChunkTasks = 3000;
         internal const int StatusCheckInSec = 10;
-        internal const int DefaultArrayPoolArraySize = 4 * 1024;
+        internal const int DefaultStreamCopyBufferSize = 81920;  // Use the .NET default
 
         internal const long DefaultInitialTransferSize = 32 * Constants.MB;
         internal const long DefaultChunkSize = 4 * Constants.MB;
diff --git a/sdk/storage/Azure.Storage.DataMovement/src/UriToStreamJobPart.cs b/sdk/storage/Azure.Storage.DataMovement/src/UriToStreamJobPart.cs
@@ -472,7 +472,7 @@ public async Task WriteChunkToTempFile(string chunkFilePath, Stream source)
             {
                 await source.CopyToAsync(
                     fileStream,
-                    Constants.DefaultDownloadCopyBufferSize,
+                    DataMovementConstants.DefaultStreamCopyBufferSize,
                     _cancellationToken)
                     .ConfigureAwait(false);
             }

Original file line number	Diff line number	Diff line change
`@@ -29,8 +29,6 @@ internal class JobPartPlanFile : IDisposable`
`29`	`29`	`/// </summary>`
`30`	`30`	`public readonly SemaphoreSlim WriteLock;`
`31`	`31`
`32`		`- private const int DefaultBufferSize = 81920;`
`33`		`-`
`34`	`32`	`private JobPartPlanFile()`
`35`	`33`	`{`
`36`	`34`	`WriteLock = new SemaphoreSlim(1);`
`@@ -66,7 +64,10 @@ public static async Task<JobPartPlanFile> CreateJobPartPlanFileAsync(`
`66`	`64`	`{`
`67`	`65`	`using (FileStream fileStream = File.Create(result.FileName.ToString()))`
`68`	`66`	`{`
`69`		`- await headerStream.CopyToAsync(fileStream, DefaultBufferSize, cancellationToken).ConfigureAwait(false);`
	`67`	`+ await headerStream.CopyToAsync(`
	`68`	`+ fileStream,`
	`69`	`+ DataMovementConstants.DefaultStreamCopyBufferSize,`
	`70`	`+ cancellationToken).ConfigureAwait(false);`
`70`	`71`	`}`
`71`	`72`	`}`
`72`	`73`	`catch (Exception)`
Original file line number	Diff line number	Diff line change
`@@ -35,8 +35,6 @@ internal class JobPlanFile : IDisposable`
`35`	`35`	`/// </summary>`
`36`	`36`	`public readonly SemaphoreSlim WriteLock;`
`37`	`37`
`38`		`- private const int DefaultBufferSize = 81920;`
`39`		`-`
`40`	`38`	`private JobPlanFile(string id, string filePath)`
`41`	`39`	`{`
`42`	`40`	`Id = id;`
`@@ -63,7 +61,10 @@ public static async Task<JobPlanFile> CreateJobPlanFileAsync(`
`63`	`61`	`{`
`64`	`62`	`using (FileStream fileStream = File.Create(jobPlanFile.FilePath))`
`65`	`63`	`{`
`66`		`- await headerStream.CopyToAsync(fileStream, DefaultBufferSize, cancellationToken).ConfigureAwait(false);`
	`64`	`+ await headerStream.CopyToAsync(`
	`65`	`+ fileStream,`
	`66`	`+ DataMovementConstants.DefaultStreamCopyBufferSize,`
	`67`	`+ cancellationToken).ConfigureAwait(false);`
`67`	`68`	`}`
`68`	`69`	`}`
`69`	`70`	`catch (Exception)`
Original file line number	Diff line number	Diff line change
`@@ -142,11 +142,9 @@ protected internal override async Task CopyFromStreamAsync(`
`142`	`142`	`{`
`143`	`143`	`fileStream.Seek(position, SeekOrigin.Begin);`
`144`	`144`	`}`
`145`		`- await stream.CopyToAsync(`
`146`		`- fileStream,`
`147`		`- (int)streamLength,`
`148`		`- cancellationToken)`
`149`		`- .ConfigureAwait(false);`
	`145`	`+`
	`146`	`+ int bufferSize = Math.Min((int)streamLength, DataMovementConstants.DefaultStreamCopyBufferSize);`
	`147`	`+ await stream.CopyToAsync(fileStream, bufferSize, cancellationToken).ConfigureAwait(false);`
`150`	`148`	`}`
`151`	`149`	`}`
`152`	`150`	`}`
Original file line number	Diff line number	Diff line change
`@@ -152,8 +152,8 @@ public override async Task<Stream> ReadJobPlanFileAsync(`
`152`	`152`	`int length,`
`153`	`153`	`CancellationToken cancellationToken = default)`
`154`	`154`	`{`
`155`		`- int maxArraySize = length > 0 ? length : DataMovementConstants.DefaultArrayPoolArraySize;`
`156`		`- Stream copiedStream = new PooledMemoryStream(ArrayPool<byte>.Shared, maxArraySize);`
	`155`	`+ int bufferSize = length > 0 ? length : DataMovementConstants.DefaultStreamCopyBufferSize;`
	`156`	`+ Stream copiedStream = new PooledMemoryStream(ArrayPool<byte>.Shared, bufferSize);`
`157`	`157`
`158`	`158`	`CancellationHelper.ThrowIfCancellationRequested(cancellationToken);`
`159`	`159`	`if (_transferStates.TryGetValue(transferId, out JobPlanFile jobPlanFile))`
`@@ -164,7 +164,7 @@ public override async Task<Stream> ReadJobPlanFileAsync(`
`164`	`164`	`using (MemoryMappedFile mmf = MemoryMappedFile.CreateFromFile(jobPlanFile.FilePath))`
`165`	`165`	`using (MemoryMappedViewStream mmfStream = mmf.CreateViewStream(offset, length, MemoryMappedFileAccess.Read))`
`166`	`166`	`{`
`167`		`- await mmfStream.CopyToAsync(copiedStream).ConfigureAwait(false);`
	`167`	`+ await mmfStream.CopyToAsync(copiedStream, bufferSize, cancellationToken).ConfigureAwait(false);`
`168`	`168`	`}`
`169`	`169`
`170`	`170`	`copiedStream.Position = 0;`
`@@ -193,16 +193,16 @@ public override async Task<Stream> ReadJobPartPlanFileAsync(`
`193`	`193`	`{`
`194`	`194`	`if (jobPlanFile.JobParts.TryGetValue(partNumber, out JobPartPlanFile jobPartPlanFile))`
`195`	`195`	`{`
`196`		`- int maxArraySize = length > 0 ? length : DataMovementConstants.DefaultArrayPoolArraySize;`
`197`		`- Stream copiedStream = new PooledMemoryStream(ArrayPool<byte>.Shared, maxArraySize);`
	`196`	`+ int bufferSize = length > 0 ? length : DataMovementConstants.DefaultStreamCopyBufferSize;`
	`197`	`+ Stream copiedStream = new PooledMemoryStream(ArrayPool<byte>.Shared, bufferSize);`
`198`	`198`
`199`	`199`	`await jobPartPlanFile.WriteLock.WaitAsync(cancellationToken).ConfigureAwait(false);`
`200`	`200`	`try`
`201`	`201`	`{`
`202`	`202`	`using (MemoryMappedFile mmf = MemoryMappedFile.CreateFromFile(jobPartPlanFile.FilePath))`
`203`	`203`	`using (MemoryMappedViewStream mmfStream = mmf.CreateViewStream(offset, length, MemoryMappedFileAccess.Read))`
`204`	`204`	`{`
`205`		`- await mmfStream.CopyToAsync(copiedStream).ConfigureAwait(false);`
	`205`	`+ await mmfStream.CopyToAsync(copiedStream, bufferSize, cancellationToken).ConfigureAwait(false);`
`206`	`206`	`}`
`207`	`207`
`208`	`208`	`copiedStream.Position = 0;`
Original file line number	Diff line number	Diff line change
`@@ -472,7 +472,7 @@ public async Task WriteChunkToTempFile(string chunkFilePath, Stream source)`
`472`	`472`	`{`
`473`	`473`	`await source.CopyToAsync(`
`474`	`474`	`fileStream,`
`475`		`- Constants.DefaultDownloadCopyBufferSize,`
	`475`	`+ DataMovementConstants.DefaultStreamCopyBufferSize,`
`476`	`476`	`_cancellationToken)`
`477`	`477`	`.ConfigureAwait(false);`
`478`	`478`	`}`