Retain a notification observer in StatsForegroundObservable

wordpress-mobile · Jun 30, 2023 · d204b03 · d204b03
1 parent b319b72
commit d204b03
Showing 1 changed file with 16 additions and 4 deletions.
diff --git a/WordPress/Classes/ViewRelated/Stats/StatsForegroundObservable.swift b/WordPress/Classes/ViewRelated/Stats/StatsForegroundObservable.swift
@@ -4,18 +4,30 @@ protocol StatsForegroundObservable: AnyObject {
     func reloadStatsData()
 }
 
+private var observerKey = 0
+
 extension StatsForegroundObservable where Self: UIViewController {
     func addWillEnterForegroundObserver() {
-        NotificationCenter.default.addObserver(forName: UIApplication.willEnterForegroundNotification,
+        enterForegroundObserver = NotificationCenter.default.addObserver(forName: UIApplication.willEnterForegroundNotification,
                                                object: nil,
                                                queue: nil) { [weak self] _ in
             self?.reloadStatsData()
         }
     }
 
     func removeWillEnterForegroundObserver() {
-        NotificationCenter.default.removeObserver(self,
-                                                  name: UIApplication.willEnterForegroundNotification,
-                                                  object: nil)
+        if let enterForegroundObserver {
+            NotificationCenter.default.removeObserver(enterForegroundObserver)
+        }
+        enterForegroundObserver = nil
+    }
+
+    private var enterForegroundObserver: NSObjectProtocol? {
+        get {
+            objc_getAssociatedObject(self, &observerKey) as? NSObjectProtocol
+        }
+        set {
+            objc_setAssociatedObject(self, &observerKey, newValue, .OBJC_ASSOCIATION_RETAIN)
+        }
     }
 }