From f1f07aa6b7d069ee08d4b323e069195a414ceff7 Mon Sep 17 00:00:00 2001
From: Tom Carey <794635+TomCarey@users.noreply.github.com>
Date: Tue, 17 Dec 2024 10:44:14 +1300
Subject: [PATCH] fix(firestore): Prevent iOS crash caused by concurrency issue
 (#772)

---
 .changeset/eleven-hornets-joke.md             |  5 ++
 .../ios/Plugin/FirebaseFirestore.swift        | 51 ++++++++++++++-----
 .../ios/Plugin/FirebaseFirestorePlugin.swift  | 13 +++--
 3 files changed, 50 insertions(+), 19 deletions(-)
 create mode 100644 .changeset/eleven-hornets-joke.md

diff --git a/.changeset/eleven-hornets-joke.md b/.changeset/eleven-hornets-joke.md
new file mode 100644
index 00000000..fa010cce
--- /dev/null
+++ b/.changeset/eleven-hornets-joke.md
@@ -0,0 +1,5 @@
+---
+'@capacitor-firebase/firestore': patch
+---
+
+fix(ios): prevent crash caused by concurrency issue
diff --git a/packages/firestore/ios/Plugin/FirebaseFirestore.swift b/packages/firestore/ios/Plugin/FirebaseFirestore.swift
index 023d509e..18d832f1 100644
--- a/packages/firestore/ios/Plugin/FirebaseFirestore.swift
+++ b/packages/firestore/ios/Plugin/FirebaseFirestore.swift
@@ -2,10 +2,32 @@ import Foundation
 import FirebaseCore
 import FirebaseFirestore
 
-@objc public class FirebaseFirestore: NSObject {
-    private let plugin: FirebaseFirestorePlugin
+private actor ListenerRegistrationMap {
     private var listenerRegistrationMap: [String: ListenerRegistration] = [:]
+    
+    func addRegistration(_ listenerRegistration: ListenerRegistration, listenerId: String) async {
+        listenerRegistrationMap[listenerId] = listenerRegistration
+    }
+    
+    func removeRegistration(listenerId: String) async {
+        if let listenerRegistration = listenerRegistrationMap[listenerId] {
+            listenerRegistration.remove()
+        }
+        listenerRegistrationMap.removeValue(forKey: listenerId)
+    }
+    
+    func removeAll() async {
+        listenerRegistrationMap.forEach { _, value in
+            value.remove()
+        }
+        listenerRegistrationMap.removeAll()
+    }
+}
 
+@objc public class FirebaseFirestore: NSObject {
+    private let plugin: FirebaseFirestorePlugin
+    private var listenerRegistrationMap = ListenerRegistrationMap()
+    
     init(plugin: FirebaseFirestorePlugin) {
         self.plugin = plugin
         super.init()
@@ -208,7 +230,9 @@ import FirebaseFirestore
                 completion(result, nil)
             }
         }
-        self.listenerRegistrationMap[callbackId] = listenerRegistration
+        Task {
+            await self.listenerRegistrationMap.addRegistration(listenerRegistration, listenerId: callbackId)
+        }
     }
 
     @objc public func addCollectionSnapshotListener(_ options: AddCollectionSnapshotListenerOptions, completion: @escaping (Result?, Error?) -> Void) {
@@ -241,7 +265,7 @@ import FirebaseFirestore
                         completion(result, nil)
                     }
                 }
-                self.listenerRegistrationMap[callbackId] = listenerRegistration
+                await listenerRegistrationMap.addRegistration(listenerRegistration, listenerId: callbackId)
             } catch {
                 completion(nil, error)
             }
@@ -278,26 +302,25 @@ import FirebaseFirestore
                         completion(result, nil)
                     }
                 }
-                self.listenerRegistrationMap[callbackId] = listenerRegistration
+                await listenerRegistrationMap.addRegistration(listenerRegistration, listenerId: callbackId)
             } catch {
                 completion(nil, error)
             }
         }
     }
 
-    @objc public func removeSnapshotListener(_ options: RemoveSnapshotListenerOptions) {
+    @objc public func removeSnapshotListener(_ options: RemoveSnapshotListenerOptions, completion: @escaping () -> Void) {
         let callbackId = options.getCallbackId()
-
-        if let listenerRegistration = self.listenerRegistrationMap[callbackId] {
-            listenerRegistration.remove()
+        Task {
+            await listenerRegistrationMap.removeRegistration(listenerId: callbackId)
+            completion()
         }
-        self.listenerRegistrationMap.removeValue(forKey: callbackId)
     }
 
-    @objc public func removeAllListeners() {
-        for listenerRegistration in self.listenerRegistrationMap.values {
-            listenerRegistration.remove()
+    @objc public func removeAllListeners(completion: @escaping () -> Void) {
+        Task {
+            await listenerRegistrationMap.removeAll()
+            completion()
         }
-        self.listenerRegistrationMap.removeAll()
     }
 }
diff --git a/packages/firestore/ios/Plugin/FirebaseFirestorePlugin.swift b/packages/firestore/ios/Plugin/FirebaseFirestorePlugin.swift
index 5d06ae95..79572479 100644
--- a/packages/firestore/ios/Plugin/FirebaseFirestorePlugin.swift
+++ b/packages/firestore/ios/Plugin/FirebaseFirestorePlugin.swift
@@ -348,8 +348,9 @@ public class FirebaseFirestorePlugin: CAPPlugin {
 
         let options = RemoveSnapshotListenerOptions(callbackId: callbackId)
 
-        implementation?.removeSnapshotListener(options)
-        call.resolve()
+        implementation?.removeSnapshotListener(options) {
+            call.resolve()
+        }
     }
 
     @objc override public func removeAllListeners(_ call: CAPPluginCall) {
@@ -357,8 +358,10 @@ public class FirebaseFirestorePlugin: CAPPlugin {
             bridge?.releaseCall(savedCall)
         }
         self.pluginCallMap.removeAll()
-
-        implementation?.removeAllListeners()
-        super.removeAllListeners(call)
+        self.eventListeners?.removeAllObjects()
+        
+        implementation?.removeAllListeners {
+            call.resolve()
+        }
     }
 }