Merge ad5b1e5 into 34b0216

FirebaseStorageSwift/Sources/AsyncAwait.swift

@@ -14,6 +14,10 @@
 
 import FirebaseStorage
 
+public enum StorageError: Error {
+  case cancelled
+}
+
 #if swift(>=5.5)
   @available(iOS 15, tvOS 15, macOS 12, *)
   public extension StorageReference {
@@ -66,12 +70,27 @@ import FirebaseStorage
     func putFileAsync(from url: URL,
                       metadata: StorageMetadata? = nil) async throws -> StorageMetadata {
       typealias MetadataContinuation = CheckedContinuation<StorageMetadata, Error>
-      return try await withCheckedThrowingContinuation { (continuation: MetadataContinuation) in
-        // TODO: Use task to handle progress and cancellation.
-        _ = self.putFile(from: url, metadata: metadata) { result in
-          continuation.resume(with: result)
+
+      var storageTask: StorageUploadTask?
+
+      return try await withTaskCancellationHandler(operation: {
+        try await withCheckedThrowingContinuation { (continuation: MetadataContinuation) in
+
+          if Task.isCancelled {
+            continuation.resume(throwing: CancellationError())
+            return
+          }
+          storageTask = self.putFile(from: url, metadata: metadata) { result in
+            continuation.resume(with: result)
+          }
         }
-      }
+      }, onCancel: { [storageTask] in
+        // A possible race condition here. `onCancel` and `operation` are performed concurrently,
+        // so it's possible for `storageTask` still be `nil` but `Task.isCancelled` still be `false`,
+        // so `storageTask` can be initialized and started after `onCancel` block has been
+        // performed.
+        storageTask?.cancel()
+      })
     }
 
     /// Asynchronously downloads the object at the current path to a specified system filepath.

FirebaseStorageSwift/Tests/Integration/StorageAsyncAwait.swift

@@ -155,17 +155,58 @@ import XCTest
       XCTAssertNotNil(result)
     }
 
-    func testSimplePutFileNoMetadata() async throws {
-      let fileName = "hello&+@_ .txt"
-      let ref = storage.reference(withPath: "ios/public/" + fileName)
+    private func setupFile(_ fileName: String) throws -> URL {
       let data = try XCTUnwrap("Hello Swift World".data(using: .utf8), "Data construction failed")
       let tmpDirURL = URL(fileURLWithPath: NSTemporaryDirectory())
       let fileURL = tmpDirURL.appendingPathComponent("hello.txt")
       try data.write(to: fileURL, options: .atomicWrite)
+      return fileURL
+    }
+
+    func testSimplePutFileNoMetadata() async throws {
+      let fileName = "hello&+@_ .txt"
+      let ref = storage.reference(withPath: "ios/public/" + fileName)
+      let fileURL = try setupFile(fileName)
       let result = try await ref.putFileAsync(from: fileURL)
       XCTAssertNotNil(result)
     }
 
+    func testSimplePutFileFromTask() async throws {
+      let fileName = "hello&+@_ .txt"
+      let ref = storage.reference(withPath: "ios/public/" + fileName)
+      let fileURL = try setupFile(fileName)
+
+      let task = Task { () -> StorageMetadata in
+        let metadata = try await ref.putFileAsync(from: fileURL)
+        return metadata
+      }
+
+      let result = try await task.value
+      XCTAssertNotNil(result)
+    }
+
+    func testSimplePutFileCancel() async throws {
+      let fileName = "hello&+@_ .txt"
+      let ref = storage.reference(withPath: "ios/public/" + fileName)
+      let fileURL = try setupFile(fileName)
+
+      let task = Task { () -> StorageMetadata in
+        try await ref.putFileAsync(from: fileURL)
+      }
+
+      // Cancel the task after a while.
+      DispatchQueue.main.asyncAfter(deadline: .now() + 0.1) {
+        task.cancel()
+      }
+
+      do {
+        _ = try await task.value
+        XCTFail("Unexpected success from putFileCancel")
+      } catch {
+        XCTAssert(error is CancellationError)
+      }
+    }
+
     func testSimpleGetData() async throws {
       let ref = storage.reference(withPath: "ios/public/1mb2")
       let result = try await ref.data(maxSize: 1024 * 1024)

FirebaseStorageSwift/Tests/Integration/StorageIntegration.swift

@@ -193,6 +193,26 @@ class StorageResultTests: StorageIntegrationCommon {
     waitForExpectations()
   }
 
+  func testSimplePutFileCancel() throws {
+    let expectation = self.expectation(description: #function)
+    let ref = storage.reference(withPath: "ios/public/testSimplePutFile")
+    let data = try XCTUnwrap("Hello Swift World".data(using: .utf8), "Data construction failed")
+    let tmpDirURL = URL(fileURLWithPath: NSTemporaryDirectory())
+    let fileURL = tmpDirURL.appendingPathComponent("hello.txt")
+    try data.write(to: fileURL, options: .atomicWrite)
+    let task = ref.putFile(from: fileURL) { result in
+      switch result {
+      case .success:
+        XCTFail("Unexpected success from putFileCancel")
+      case let .failure(error as NSError):
+        XCTAssertEqual(error.code, StorageErrorCode.cancelled.rawValue)
+        expectation.fulfill()
+      }
+    }
+    task.cancel()
+    waitForExpectations()
+  }
+
   func testAttemptToUploadDirectoryShouldFail() throws {
     // This `.numbers` file is actually a directory.
     let fileName = "HomeImprovement.numbers"