Plan memory leak (Azure#47495)

* initial implementation

* test

* feedback

Author: jaschrep-msft

sdk/storage/Azure.Storage.DataMovement/src/JobPlan/JobPlanFile.cs

@@ -43,6 +43,12 @@ private JobPlanFile(string id, string filePath)
             WriteLock = new SemaphoreSlim(1);
         }
 
+        private static string ToFullPath(string checkpointerPath, string transferId)
+        {
+            string fileName = $"{transferId}{DataMovementConstants.JobPlanFile.FileExtension}";
+            return Path.Combine(checkpointerPath, fileName);
+        }
+
         public static async Task<JobPlanFile> CreateJobPlanFileAsync(
             string checkpointerPath,
             string id,
@@ -53,8 +59,7 @@ public static async Task<JobPlanFile> CreateJobPlanFileAsync(
             Argument.AssertNotNullOrEmpty(id, nameof(id));
             Argument.AssertNotNull(headerStream, nameof(headerStream));
 
-            string fileName = $"{id}{DataMovementConstants.JobPlanFile.FileExtension}";
-            string filePath = Path.Combine(checkpointerPath, fileName);
+            string filePath = ToFullPath(checkpointerPath, id);
 
             JobPlanFile jobPlanFile = new(id, filePath);
             try
@@ -92,6 +97,11 @@ public static JobPlanFile LoadExistingJobPlanFile(string fullPath)
             return new JobPlanFile(transferId, fullPath);
         }
 
+        public static JobPlanFile LoadExistingJobPlanFile(string checkpointerPath, string transferId)
+        {
+            return LoadExistingJobPlanFile(ToFullPath(checkpointerPath, transferId));
+        }
+
         public void Dispose()
         {
             WriteLock.Dispose();

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

@@ -27,7 +27,7 @@ internal class LocalTransferCheckpointer : SerializerTransferCheckpointer
         /// <summary>
         /// Stores references to the memory mapped files stored by IDs.
         /// </summary>
-        private Dictionary<string, JobPlanFile> _transferStates;
+        internal readonly Dictionary<string, JobPlanFile> _transferStates;
 
         /// <summary>
         /// Initializes a new instance of <see cref="LocalTransferCheckpointer"/> class.
@@ -52,10 +52,22 @@ public LocalTransferCheckpointer(string folderPath)
             else
             {
                 _pathToCheckpointer = folderPath;
-                InitializeExistingCheckpointer();
             }
         }
 
+        private bool TryGetJobPlanFile(string transferId, out JobPlanFile result)
+        {
+            if (!_transferStates.TryGetValue(transferId, out result))
+            {
+                RefreshCache(transferId);
+                if (!_transferStates.TryGetValue(transferId, out result))
+                {
+                    return false;
+                }
+            }
+            return true;
+        }
+
         public override async Task AddNewJobAsync(
             string transferId,
             StorageResource source,
@@ -138,11 +150,11 @@ public override Task<int> CurrentJobPartCountAsync(
             CancellationToken cancellationToken = default)
         {
             CancellationHelper.ThrowIfCancellationRequested(cancellationToken);
-            if (_transferStates.TryGetValue(transferId, out JobPlanFile result))
+            if (!TryGetJobPlanFile(transferId, out JobPlanFile result))
             {
-                return Task.FromResult(result.JobParts.Count);
+                throw Errors.MissingTransferIdCheckpointer(transferId);
             }
-            throw Errors.MissingTransferIdCheckpointer(transferId);
+            return Task.FromResult(result.JobParts.Count);
         }
 
         public override async Task<Stream> ReadJobPlanFileAsync(
@@ -155,28 +167,26 @@ public override async Task<Stream> ReadJobPlanFileAsync(
             Stream copiedStream = new PooledMemoryStream(ArrayPool<byte>.Shared, bufferSize);
 
             CancellationHelper.ThrowIfCancellationRequested(cancellationToken);
-            if (_transferStates.TryGetValue(transferId, out JobPlanFile jobPlanFile))
+            if (!TryGetJobPlanFile(transferId, out JobPlanFile jobPlanFile))
             {
-                await jobPlanFile.WriteLock.WaitAsync(cancellationToken).ConfigureAwait(false);
-                try
-                {
-                    using (MemoryMappedFile mmf = MemoryMappedFile.CreateFromFile(jobPlanFile.FilePath))
-                    using (MemoryMappedViewStream mmfStream = mmf.CreateViewStream(offset, length, MemoryMappedFileAccess.Read))
-                    {
-                        await mmfStream.CopyToAsync(copiedStream, bufferSize, cancellationToken).ConfigureAwait(false);
-                    }
+                throw Errors.MissingTransferIdCheckpointer(transferId);
+            }
 
-                    copiedStream.Position = 0;
-                    return copiedStream;
-                }
-                finally
+            await jobPlanFile.WriteLock.WaitAsync(cancellationToken).ConfigureAwait(false);
+            try
+            {
+                using (MemoryMappedFile mmf = MemoryMappedFile.CreateFromFile(jobPlanFile.FilePath))
+                using (MemoryMappedViewStream mmfStream = mmf.CreateViewStream(offset, length, MemoryMappedFileAccess.Read))
                 {
-                    jobPlanFile.WriteLock.Release();
+                    await mmfStream.CopyToAsync(copiedStream, bufferSize, cancellationToken).ConfigureAwait(false);
                 }
+
+                copiedStream.Position = 0;
+                return copiedStream;
             }
-            else
+            finally
             {
-                throw Errors.MissingTransferIdCheckpointer(transferId);
+                jobPlanFile.WriteLock.Release();
             }
         }
 
@@ -188,38 +198,33 @@ public override async Task<Stream> ReadJobPartPlanFileAsync(
             CancellationToken cancellationToken = default)
         {
             CancellationHelper.ThrowIfCancellationRequested(cancellationToken);
-            if (_transferStates.TryGetValue(transferId, out JobPlanFile jobPlanFile))
+            if (!TryGetJobPlanFile(transferId, out JobPlanFile jobPlanFile))
             {
-                if (jobPlanFile.JobParts.TryGetValue(partNumber, out JobPartPlanFile jobPartPlanFile))
-                {
-                    int bufferSize = length > 0 ? length : DataMovementConstants.DefaultStreamCopyBufferSize;
-                    Stream copiedStream = new PooledMemoryStream(ArrayPool<byte>.Shared, bufferSize);
+                throw Errors.MissingTransferIdCheckpointer(transferId);
+            }
+            if (!jobPlanFile.JobParts.TryGetValue(partNumber, out JobPartPlanFile jobPartPlanFile))
+            {
+                throw Errors.MissingPartNumberCheckpointer(transferId, partNumber);
+            }
 
-                    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, bufferSize, cancellationToken).ConfigureAwait(false);
-                        }
-
-                        copiedStream.Position = 0;
-                        return copiedStream;
-                    }
-                    finally
-                    {
-                        jobPartPlanFile.WriteLock.Release();
-                    }
-                }
-                else
+            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))
                 {
-                    throw Errors.MissingPartNumberCheckpointer(transferId, partNumber);
+                    await mmfStream.CopyToAsync(copiedStream, bufferSize, cancellationToken).ConfigureAwait(false);
                 }
+
+                copiedStream.Position = 0;
+                return copiedStream;
             }
-            else
+            finally
             {
-                throw Errors.MissingTransferIdCheckpointer(transferId);
+                jobPartPlanFile.WriteLock.Release();
             }
         }
 
@@ -261,7 +266,7 @@ public override Task<bool> TryRemoveStoredTransferAsync(string transferId, Cance
 
             List<string> filesToDelete = new List<string>();
 
-            if (_transferStates.TryGetValue(transferId, out JobPlanFile jobPlanFile))
+            if (TryGetJobPlanFile(transferId, out JobPlanFile jobPlanFile))
             {
                 filesToDelete.Add(jobPlanFile.FilePath);
             }
@@ -303,6 +308,7 @@ public override Task<bool> TryRemoveStoredTransferAsync(string transferId, Cance
 
         public override Task<List<string>> GetStoredTransfersAsync(CancellationToken cancellationToken = default)
         {
+            RefreshCache();
             return Task.FromResult(_transferStates.Keys.ToList());
         }
 
@@ -316,26 +322,37 @@ public override async Task SetJobTransferStatusAsync(
 
             CancellationHelper.ThrowIfCancellationRequested(cancellationToken);
 
-            if (_transferStates.TryGetValue(transferId, out JobPlanFile jobPlanFile))
+            if (!TryGetJobPlanFile(transferId, out JobPlanFile jobPlanFile))
             {
-                await jobPlanFile.WriteLock.WaitAsync(cancellationToken).ConfigureAwait(false);
-                try
-                {
-                    using (MemoryMappedFile mmf = MemoryMappedFile.CreateFromFile(jobPlanFile.FilePath, FileMode.Open))
-                    using (MemoryMappedViewAccessor accessor = mmf.CreateViewAccessor(offset, length))
-                    {
-                        accessor.Write(0, (int)status.ToJobPlanStatus());
-                        accessor.Flush();
-                    }
-                }
-                finally
+                throw Errors.MissingTransferIdCheckpointer(transferId);
+            }
+
+            // if completed successfully, get rid of all checkpointing info
+            if (status.HasCompletedSuccessfully)
+            {
+                await TryRemoveStoredTransferAsync(transferId, cancellationToken).ConfigureAwait(false);
+                return;
+            }
+
+            // if paused or other completion state, remove the memory cache but still write state to the plan file for later resume
+            if (status.State == DataTransferState.Completed || status.State == DataTransferState.Paused)
+            {
+                _transferStates.Remove(transferId);
+            }
+
+            await jobPlanFile.WriteLock.WaitAsync(cancellationToken).ConfigureAwait(false);
+            try
+            {
+                using (MemoryMappedFile mmf = MemoryMappedFile.CreateFromFile(jobPlanFile.FilePath, FileMode.Open))
+                using (MemoryMappedViewAccessor accessor = mmf.CreateViewAccessor(offset, length))
                 {
-                    jobPlanFile.WriteLock.Release();
+                    accessor.Write(0, (int)status.ToJobPlanStatus());
+                    accessor.Flush();
                 }
             }
-            else
+            finally
             {
-                throw Errors.MissingTransferIdCheckpointer(transferId);
+                jobPlanFile.WriteLock.Release();
             }
         }
 
@@ -350,47 +367,39 @@ public override async Task SetJobPartTransferStatusAsync(
 
             CancellationHelper.ThrowIfCancellationRequested(cancellationToken);
 
-            if (_transferStates.TryGetValue(transferId, out JobPlanFile jobPlanFile))
+            if (!TryGetJobPlanFile(transferId, out JobPlanFile jobPlanFile))
             {
-                if (jobPlanFile.JobParts.TryGetValue(partNumber, out JobPartPlanFile file))
-                {
-                    await file.WriteLock.WaitAsync(cancellationToken).ConfigureAwait(false);
-                    try
-                    {
-                        using (MemoryMappedFile mmf = MemoryMappedFile.CreateFromFile(file.FilePath, FileMode.Open))
-                        using (MemoryMappedViewAccessor accessor = mmf.CreateViewAccessor(offset, length))
-                        {
-                            accessor.Write(0, (int)status.ToJobPlanStatus());
-                            accessor.Flush();
-                        }
-                    }
-                    finally
-                    {
-                        file.WriteLock.Release();
-                    }
-                }
-                else
+                throw Errors.MissingTransferIdCheckpointer(transferId);
+            }
+            if (!jobPlanFile.JobParts.TryGetValue(partNumber, out JobPartPlanFile file))
+            {
+                throw Errors.MissingPartNumberCheckpointer(transferId, partNumber);
+            }
+            await file.WriteLock.WaitAsync(cancellationToken).ConfigureAwait(false);
+            try
+            {
+                using (MemoryMappedFile mmf = MemoryMappedFile.CreateFromFile(file.FilePath, FileMode.Open))
+                using (MemoryMappedViewAccessor accessor = mmf.CreateViewAccessor(offset, length))
                 {
-                    throw Errors.MissingPartNumberCheckpointer(transferId, partNumber);
+                    accessor.Write(0, (int)status.ToJobPlanStatus());
+                    accessor.Flush();
                 }
             }
-            else
+            finally
             {
-                throw Errors.MissingTransferIdCheckpointer(transferId);
+                file.WriteLock.Release();
             }
         }
 
         /// <summary>
-        /// Takes the path of the checkpointer reads all the files in the top directory level
-        /// and populates the _transferStates.
+        /// Clears cached transfer states and repopulates by enumerating directory.
         /// </summary>
-        private void InitializeExistingCheckpointer()
+        private void RefreshCache()
         {
-            // Enumerate the filesystem
-            IEnumerable<string> checkpointFiles = Directory.EnumerateFiles(_pathToCheckpointer);
+            _transferStates.Clear();
 
             // First, retrieve all valid job plan files
-            foreach (string path in checkpointFiles
+            foreach (string path in Directory.EnumerateFiles(_pathToCheckpointer)
                 .Where(p => Path.GetExtension(p) == DataMovementConstants.JobPlanFile.FileExtension))
             {
                 // TODO: Should we check for valid schema version inside file now?
@@ -406,7 +415,7 @@ private void InitializeExistingCheckpointer()
             }
 
             // Retrieve all valid job part plan files stored in the checkpointer path.
-            foreach (string path in checkpointFiles
+            foreach (string path in Directory.EnumerateFiles(_pathToCheckpointer)
                 .Where(p => Path.GetExtension(p) == DataMovementConstants.JobPartPlanFile.FileExtension))
             {
                 // Ensure each file has the correct format
@@ -423,6 +432,32 @@ private void InitializeExistingCheckpointer()
             }
         }
 
+        /// <summary>
+        /// Clears cache for a given trandfer ID and repopulates from disk if any.
+        /// </summary>
+        private void RefreshCache(string transferId)
+        {
+            _transferStates.Remove(transferId);
+            JobPlanFile jobPlanFile = JobPlanFile.LoadExistingJobPlanFile(_pathToCheckpointer, transferId);
+            if (!File.Exists(jobPlanFile.FilePath))
+            {
+                return;
+            }
+            _transferStates.Add(transferId, jobPlanFile);
+            foreach (string path in Directory.EnumerateFiles(_pathToCheckpointer)
+                .Where(p => Path.GetExtension(p) == DataMovementConstants.JobPartPlanFile.FileExtension))
+            {
+                // Ensure each file has the correct format
+                if (JobPartPlanFileName.TryParseJobPartPlanFileName(path, out JobPartPlanFileName partPlanFileName) &&
+                    partPlanFileName.Id == transferId)
+                {
+                    jobPlanFile.JobParts.TryAdd(
+                        partPlanFileName.JobPartNumber,
+                        JobPartPlanFile.CreateExistingPartPlanFile(partPlanFileName));
+                }
+            }
+        }
+
         private static JobPlanOperation GetOperationType(StorageResource source, StorageResource destination)
         {
             if (source.IsLocalResource() && !destination.IsLocalResource())