Add Storage API to limit upload chunk size (#10609)

Author: paulb777

FirebaseStorage/CHANGELOG.md

@@ -1,3 +1,6 @@
+# 10.5.0
+- [added] Added Storage API to limit upload chunk size. (#10137)
+
 # 10.3.0
 - [fixed] Use dedicated serial queue for Storage uploads and downloads instead of a (concurrent) global queue.
   Fixes regression introduced in 10.0.0. (#10487)

FirebaseStorage/Sources/Storage.swift

@@ -123,6 +123,12 @@ import FirebaseAuthInterop
     }
   }
 
+  /**
+   * Specify the maximum upload chunk size. Values less than 256K (262144) will be rounded up to 256K. Values
+   * above 256K will be rounded down to the nearest 256K multiple. The default is no maximum.
+   */
+  @objc public var uploadChunkSizeBytes: Int64 = .max
+
   /**
    * A `DispatchQueue` that all developer callbacks are fired on. Defaults to the main queue.
    */

FirebaseStorage/Sources/StorageUploadTask.swift

@@ -76,7 +76,7 @@ import Foundation
       let uploadFetcher = GTMSessionUploadFetcher(
         request: request,
         uploadMIMEType: contentType,
-        chunkSize: Int64.max,
+        chunkSize: self.reference.storage.uploadChunkSizeBytes,
         fetcherService: self.fetcherService
       )
       if let data = self.uploadData {

FirebaseStorage/Tests/Integration/StorageIntegration.swift

@@ -265,6 +265,106 @@ class StorageResultTests: StorageIntegrationCommon {
     waitForExpectations()
   }
 
+  func testPutFileLimitedChunk() throws {
+    defer {
+      // Reset since tests share storage instance.
+      storage.uploadChunkSizeBytes = Int64.max
+    }
+    let expectation = self.expectation(description: #function)
+    let putFileExpectation = self.expectation(description: "putFile")
+    let ref = storage.reference(withPath: "ios/public/testPutFilePauseResume")
+    let bundle = Bundle(for: StorageIntegrationCommon.self)
+    let filePath = try XCTUnwrap(bundle.path(forResource: "1mb", ofType: "dat"),
+                                 "Failed to get filePath")
+    let data = try XCTUnwrap(try Data(contentsOf: URL(fileURLWithPath: filePath)),
+                             "Failed to load file")
+    let tmpDirURL = URL(fileURLWithPath: NSTemporaryDirectory())
+    let fileURL = tmpDirURL.appendingPathComponent("LargePutFile.txt")
+    var progressCount = 0
+
+    try data.write(to: fileURL, options: .atomicWrite)
+
+    // Limit the upload chunk size
+    storage.uploadChunkSizeBytes = 256_000
+
+    let task = ref.putFile(from: fileURL) { result in
+      XCTAssertGreaterThanOrEqual(progressCount, 4)
+      self.assertResultSuccess(result)
+      putFileExpectation.fulfill()
+    }
+
+    task.observe(StorageTaskStatus.success) { snapshot in
+      XCTAssertEqual(snapshot.description, "<State: Success>")
+      expectation.fulfill()
+    }
+
+    var uploadedBytes: Int64 = -1
+
+    task.observe(StorageTaskStatus.progress) { snapshot in
+      XCTAssertTrue(snapshot.description.starts(with: "<State: Progress") ||
+        snapshot.description.starts(with: "<State: Resume"))
+      guard let progress = snapshot.progress else {
+        XCTFail("Failed to get snapshot.progress")
+        return
+      }
+      progressCount = progressCount + 1
+      XCTAssertGreaterThanOrEqual(progress.completedUnitCount, uploadedBytes)
+      uploadedBytes = progress.completedUnitCount
+    }
+    waitForExpectations()
+  }
+
+  func testPutFileTinyChunk() throws {
+    defer {
+      // Reset since tests share storage instance.
+      storage.uploadChunkSizeBytes = Int64.max
+    }
+    let expectation = self.expectation(description: #function)
+    let putFileExpectation = self.expectation(description: "putFile")
+    let ref = storage.reference(withPath: "ios/public/testPutFilePauseResume")
+    let bundle = Bundle(for: StorageIntegrationCommon.self)
+    let filePath = try XCTUnwrap(bundle.path(forResource: "1mb", ofType: "dat"),
+                                 "Failed to get filePath")
+    let data = try XCTUnwrap(try Data(contentsOf: URL(fileURLWithPath: filePath)),
+                             "Failed to load file")
+    let tmpDirURL = URL(fileURLWithPath: NSTemporaryDirectory())
+    let fileURL = tmpDirURL.appendingPathComponent("LargePutFile.txt")
+    var progressCount = 0
+
+    try data.write(to: fileURL, options: .atomicWrite)
+
+    // Limit the upload chunk size. This should behave exactly like the previous
+    // test since small chunk sizes are rounded up to 256K.
+    storage.uploadChunkSizeBytes = 1
+
+    let task = ref.putFile(from: fileURL) { result in
+      XCTAssertGreaterThanOrEqual(progressCount, 4)
+      XCTAssertLessThanOrEqual(progressCount, 6)
+      self.assertResultSuccess(result)
+      putFileExpectation.fulfill()
+    }
+
+    task.observe(StorageTaskStatus.success) { snapshot in
+      XCTAssertEqual(snapshot.description, "<State: Success>")
+      expectation.fulfill()
+    }
+
+    var uploadedBytes: Int64 = -1
+
+    task.observe(StorageTaskStatus.progress) { snapshot in
+      XCTAssertTrue(snapshot.description.starts(with: "<State: Progress") ||
+        snapshot.description.starts(with: "<State: Resume"))
+      guard let progress = snapshot.progress else {
+        XCTFail("Failed to get snapshot.progress")
+        return
+      }
+      progressCount = progressCount + 1
+      XCTAssertGreaterThanOrEqual(progress.completedUnitCount, uploadedBytes)
+      uploadedBytes = progress.completedUnitCount
+    }
+    waitForExpectations()
+  }
+
   func testAttemptToUploadDirectoryShouldFail() throws {
     // This `.numbers` file is actually a directory.
     let fileName = "HomeImprovement.numbers"