From c3b34de3960f907608f96f165ddcccb0d4e3766d Mon Sep 17 00:00:00 2001
From: Bnaya Peretz <me@bnaya.net>
Date: Thu, 5 Nov 2020 21:12:59 +0200
Subject: [PATCH] Iteration

---
 src/useObserver.ts                                       | 2 +-
 ...teReactionCleanupTrackingUsingFinalizationRegister.ts | 9 +++++----
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/src/useObserver.ts b/src/useObserver.ts
index ad3609b..2ff48b8 100644
--- a/src/useObserver.ts
+++ b/src/useObserver.ts
@@ -15,7 +15,7 @@ function observerComponentNameFor(baseComponentName: string) {
 }
 
 /**
- * We use class to make it easier to detect in heap explorer
+ * We use class to make it easier to detect in heap snapshots by name
  */
 class ObjectToBeRetainedByReact {}
 
diff --git a/src/utils/createReactionCleanupTrackingUsingFinalizationRegister.ts b/src/utils/createReactionCleanupTrackingUsingFinalizationRegister.ts
index f5fc352..44f2038 100644
--- a/src/utils/createReactionCleanupTrackingUsingFinalizationRegister.ts
+++ b/src/utils/createReactionCleanupTrackingUsingFinalizationRegister.ts
@@ -35,15 +35,16 @@ export function createReactionCleanupTrackingUsingFinalizationRegister(
             reactionTrackingRef.current = createTrackingData(reaction)
             reactionTrackingRef.current.finalizationRegistryCleanupToken = token
             cleanupTokenToReactionTrackingMap.set(token, reactionTrackingRef.current)
+
             return reactionTrackingRef.current
         },
         recordReactionAsCommitted(reactionRef: React.MutableRefObject<IReactionTracking | null>) {
             registry.unregister(reactionRef)
 
-            for (const [key, value] of cleanupTokenToReactionTrackingMap.entries()) {
-                if (value === reactionRef.current) {
-                    cleanupTokenToReactionTrackingMap.delete(key)
-                }
+            if (reactionRef.current && reactionRef.current.finalizationRegistryCleanupToken) {
+                cleanupTokenToReactionTrackingMap.delete(
+                    reactionRef.current.finalizationRegistryCleanupToken
+                )
             }
         },
         forceCleanupTimerToRunNowForTests() {