From ee8d7cbb8ccf06ef8a3e58266b9c95ddf3a71902 Mon Sep 17 00:00:00 2001
From: AJ Lauer Barinov <abarinov@mux.com>
Date: Tue, 5 Mar 2024 17:32:35 -0800
Subject: [PATCH] avoid keeping a reference to the file chunk in ChunkWorker

When the Worker actor loops over a file, it does not always deallocate
file chunks until the loop completes.

Because each ChunkWorker is intended to be only used once, it is not
necessary to keep a reference to its file chunk beyond the upload task
call.
---
 Sources/MuxUploadSDK/Upload/ChunkWorker.swift | 36 +++++++++++--------
 .../Upload/ChunkedFileUploader.swift          |  6 ++--
 2 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/Sources/MuxUploadSDK/Upload/ChunkWorker.swift b/Sources/MuxUploadSDK/Upload/ChunkWorker.swift
index 4f6b4e0e..eab4be9c 100644
--- a/Sources/MuxUploadSDK/Upload/ChunkWorker.swift
+++ b/Sources/MuxUploadSDK/Upload/ChunkWorker.swift
@@ -7,12 +7,13 @@
 
 import Foundation
 
-/// Uploads a single chunk. Starts an internal Task on creation, which you can get with ``getTask``
-///  This class takes care of retries and backoff on a per-chunk basis
-///  This class provides no thread safety to the outside world
+/// Uploads a single chunk. Starts an internal Task on creation, 
+/// which can be accessed with ``makeUploadTaskIfNeeded``.
+///
+/// This class takes care of retries and backoff on a per-chunk basis
+/// This class provides no thread safety to the outside world
 class ChunkWorker {
     let uploadURL: URL
-    let chunk: FileChunk
     let maxRetries: Int
     let chunkProgress: Progress
     
@@ -33,11 +34,12 @@ class ChunkWorker {
         self.progressDelegate = delegatePair
     }
     
-    func getTask() -> Task<Success, Error> {
-
+    func makeUploadTaskIfNeeded(
+        chunk: FileChunk
+    ) -> Task<Success, Error> {
         guard let uploadTask else {
             chunkStartTime = Date().timeIntervalSince1970
-            let uploadTask = makeUploadTask()
+            let uploadTask = makeUploadTask(chunk: chunk)
             self.uploadTask = uploadTask
             return uploadTask
         }
@@ -51,12 +53,14 @@ class ChunkWorker {
         }
     }
     
-    private func makeUploadTask() -> Task<Success, Error> {
+    private func makeUploadTask(
+        chunk: FileChunk
+    ) -> Task<Success, Error> {
         return Task { [self] in
             var retries = 0
             var requestError: Error?
             let repsonseValidator = ChunkResponseValidator()
-    
+
             repeat {
                 do {
                     let chunkActor = ChunkActor(
@@ -145,21 +149,25 @@ class ChunkWorker {
     struct Success : Sendable {
         let finalState: Update
         let tries: Int
-        // TODO: Also AF Response
     }
     
-    convenience init(uploadInfo: UploadInfo, fileChunk: FileChunk, chunkProgress: Progress) {
+    convenience init(
+        uploadInfo: UploadInfo,
+        chunkProgress: Progress
+    ) {
         self.init(
             uploadURL: uploadInfo.uploadURL,
-            fileChunk: fileChunk,
             chunkProgress: chunkProgress,
             maxRetries: uploadInfo.options.transport.retryLimitPerChunk
         )
     }
     
-    init(uploadURL: URL, fileChunk: FileChunk, chunkProgress: Progress, maxRetries: Int) {
+    init(
+        uploadURL: URL,
+        chunkProgress: Progress,
+        maxRetries: Int
+    ) {
         self.uploadURL = uploadURL
-        self.chunk = fileChunk
         self.maxRetries = maxRetries
         self.chunkProgress = chunkProgress
     }
diff --git a/Sources/MuxUploadSDK/Upload/ChunkedFileUploader.swift b/Sources/MuxUploadSDK/Upload/ChunkedFileUploader.swift
index 139d4606..5bf50864 100644
--- a/Sources/MuxUploadSDK/Upload/ChunkedFileUploader.swift
+++ b/Sources/MuxUploadSDK/Upload/ChunkedFileUploader.swift
@@ -449,11 +449,9 @@ fileprivate actor Worker {
             
             let wideChunkSize = Int64(chunk.size())
             let chunkProgress = Progress(totalUnitCount: wideChunkSize)
-            //overallProgress.addChild(chunkProgress, withPendingUnitCount: wideChunkSize)
             
             let chunkWorker = ChunkWorker(
                 uploadURL: uploadInfo.uploadURL,
-                fileChunk: chunk,
                 chunkProgress: chunkProgress,
                 maxRetries: uploadInfo.options.transport.retryLimitPerChunk
             )
@@ -470,7 +468,9 @@ fileprivate actor Worker {
             // Problem is in line bellow, task will retain the reference to a chunk read from file and will
             // not release it until the for loop is exited, we need to find a way to implicitly release task memory
             // withouth breaking the for loop.
-            let chunkResult = try await chunkWorker.getTask().value
+            let chunkResult = try await chunkWorker.makeUploadTaskIfNeeded(
+                chunk: chunk
+            ).value
             SDKLogger.logger?.info("Completed Chunk:\n \(String(describing: chunkResult))")
         } while (readBytes == uploadInfo.options.transport.chunkSizeInBytes)