add metadata and cancellation support

thebrowsercompany · May 15, 2024 · dd13a79 · dd13a79
1 parent b028a06
commit dd13a79
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 9 deletions.
diff --git a/Sources/FirebaseStorage/StorageReference+Swift.swift b/Sources/FirebaseStorage/StorageReference+Swift.swift
@@ -52,13 +52,24 @@ public class StorageReference {
     metadata: StorageMetadata? = nil,
     onProgress: ((Progress?) -> Void)? = nil
   ) async throws -> StorageMetadata {
-    try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<StorageMetadata, any Error>) in
-      uploadData.withUnsafeBytes { ptr in
-        let future = swift_firebase.swift_cxx_shims.firebase.storage.storage_reference_put_bytes(
-          self.impl, ptr.baseAddress!.assumingMemoryBound(to: UInt8.self), uploadData.count
-          // XXX support `metadata`
-          // XXX support `onProgress`
-        )
+    // TODO(PRENG-63978): Add support for `onProgress` callback.
+    assert(onProgress == nil, "Missing support for non-nil onProgress")
+    let controller = ControllerRef()
+    return try await withTaskCancellationHandler {
+      try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<StorageMetadata, any Error>) in
+        let future = uploadData.withUnsafeBytes { ptr in
+          let bytes = ptr.baseAddress!.assumingMemoryBound(to: UInt8.self)
+          let numBytes = uploadData.count
+          if let metadata {
+            return swift_firebase.swift_cxx_shims.firebase.storage.storage_reference_put_bytes(
+              self.impl, bytes, numBytes, metadata.impl, &controller.impl
+            )
+          } else {
+            return swift_firebase.swift_cxx_shims.firebase.storage.storage_reference_put_bytes(
+              self.impl, bytes, numBytes, &controller.impl
+            )
+          }
+        }
         future.setCompletion({
           let (result, error) = future.resultAndError { StorageErrorCode($0) }
           if let error {
@@ -68,6 +79,13 @@ public class StorageReference {
           }
         })
       }
+    } onCancel: {
+      controller.impl.Cancel()
     }
   }
 }
+
+// The underlying `firebase.storage.Controller` type is thread-safe.
+private class ControllerRef: @unchecked Sendable {
+  var impl: firebase.storage.Controller = .init()
+}
diff --git a/Sources/firebase/include/FirebaseStorage.hh b/Sources/firebase/include/FirebaseStorage.hh
@@ -36,8 +36,17 @@ storage_reference_get_download_url(::firebase::storage::StorageReference ref) {
 
 inline ::swift_firebase::swift_cxx_shims::firebase::Future<::firebase::storage::Metadata>
 storage_reference_put_bytes(::firebase::storage::StorageReference ref,
-                            const void* buffer, size_t buffer_size) {
-  return ref.PutBytes(buffer, buffer_size);
+                            const void* buffer, size_t buffer_size,
+                            ::firebase::storage::Controller* controller) {
+  return ref.PutBytes(buffer, buffer_size, nullptr, controller);
+}
+
+inline ::swift_firebase::swift_cxx_shims::firebase::Future<::firebase::storage::Metadata>
+storage_reference_put_bytes(::firebase::storage::StorageReference ref,
+                            const void* buffer, size_t buffer_size,
+                            const ::firebase::storage::Metadata& metadata,
+                            ::firebase::storage::Controller* controller) {
+  return ref.PutBytes(buffer, buffer_size, metadata, nullptr, controller);
 }
 
 typedef std::map<std::string, std::string> CustomMetadata;