From aa64b588fff41c13df1bec922b24a416d4928d53 Mon Sep 17 00:00:00 2001
From: Leonel Fernandez Mir <leonelfmir@meta.com>
Date: Thu, 26 Oct 2023 21:06:22 -0700
Subject: [PATCH] Unsubscribing from resolvers before removing the record while
 runing the garbage collector

Reviewed By: captbaritone

Differential Revision: D50540816

fbshipit-source-id: 2ddcb625b1e194bebdf3a1308b8c07944f41442e
---
 .../__tests__/LiveResolvers-test.js           | 21 +++++++++++--------
 .../resolvers/ExampleExternalStateStore.js    |  3 +++
 .../LiveResolverCache.js                      |  1 +
 .../LiveResolverStore.js                      | 17 ++++++++++++++-
 4 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/packages/react-relay/__tests__/LiveResolvers-test.js b/packages/react-relay/__tests__/LiveResolvers-test.js
index b3ed92c8133a8..7da3a15d6c4c4 100644
--- a/packages/react-relay/__tests__/LiveResolvers-test.js
+++ b/packages/react-relay/__tests__/LiveResolvers-test.js
@@ -43,7 +43,6 @@ const RelayRecordSource = require('relay-runtime/store/RelayRecordSource');
 const {
   disallowConsoleErrors,
   disallowWarnings,
-  expectToWarn,
 } = require('relay-test-utils-internal');
 
 disallowWarnings();
@@ -1539,10 +1538,19 @@ describe.each([
       ),
     );
 
+    const subscriptionsCountBeforeGCRun = GLOBAL_STORE.getSubscriptionsCount();
+
     // Go-go-go! Clean the store!
     store.scheduleGC();
     jest.runAllImmediates();
-    // This will clean the store, but won't unsubscribe from the external states
+    // This will clean the store, and unsubscribe from the external states
+
+    const subscriptionsCountAfterGCRun = GLOBAL_STORE.getSubscriptionsCount();
+
+    // this will verify that we unsubscribed from the external store
+    expect(subscriptionsCountAfterGCRun).toEqual(
+      subscriptionsCountBeforeGCRun - 1,
+    );
 
     // Re-reading resolvers will create new records for them (but) the
     // `live_counter_with_possible_missing_fragment_data` will have missing required data at this
@@ -1550,13 +1558,8 @@ describe.each([
     // from the external state.
     environment.lookup(operation.fragment);
 
-    // this will dispatch an action from the extenrnal store and the callback that was created before GC
-    expectToWarn(
-      'Unexpected callback for a incomplete live resolver record',
-      () => {
-        GLOBAL_STORE.dispatch({type: 'INCREMENT'});
-      },
-    );
+    // this will dispatch an action from the external store and the callback that was created before GC
+    GLOBAL_STORE.dispatch({type: 'INCREMENT'});
 
     // The data for the live resolver is missing (it has missing dependecies)
     snapshot = environment.lookup(operation.fragment);
diff --git a/packages/relay-runtime/store/__tests__/resolvers/ExampleExternalStateStore.js b/packages/relay-runtime/store/__tests__/resolvers/ExampleExternalStateStore.js
index 73819681630cc..81c7bdca672c4 100644
--- a/packages/relay-runtime/store/__tests__/resolvers/ExampleExternalStateStore.js
+++ b/packages/relay-runtime/store/__tests__/resolvers/ExampleExternalStateStore.js
@@ -44,6 +44,9 @@ class Store {
     this._state = 0;
     this._subscriptions = [];
   }
+  getSubscriptionsCount(): number {
+    return this._subscriptions.length;
+  }
 }
 
 const Selectors = {
diff --git a/packages/relay-runtime/store/experimental-live-resolvers/LiveResolverCache.js b/packages/relay-runtime/store/experimental-live-resolvers/LiveResolverCache.js
index 7b260315a720c..0831b7caadec3 100644
--- a/packages/relay-runtime/store/experimental-live-resolvers/LiveResolverCache.js
+++ b/packages/relay-runtime/store/experimental-live-resolvers/LiveResolverCache.js
@@ -895,4 +895,5 @@ function getConcreteTypename(
 module.exports = {
   LiveResolverCache,
   getUpdatedDataIDs,
+  RELAY_RESOLVER_LIVE_STATE_SUBSCRIPTION_KEY,
 };
diff --git a/packages/relay-runtime/store/experimental-live-resolvers/LiveResolverStore.js b/packages/relay-runtime/store/experimental-live-resolvers/LiveResolverStore.js
index a23308caae120..4de21cb2f7a58 100644
--- a/packages/relay-runtime/store/experimental-live-resolvers/LiveResolverStore.js
+++ b/packages/relay-runtime/store/experimental-live-resolvers/LiveResolverStore.js
@@ -48,7 +48,11 @@ const RelayReferenceMarker = require('../RelayReferenceMarker');
 const RelayStoreSubscriptions = require('../RelayStoreSubscriptions');
 const RelayStoreUtils = require('../RelayStoreUtils');
 const {ROOT_ID, ROOT_TYPE} = require('../RelayStoreUtils');
-const {LiveResolverCache, getUpdatedDataIDs} = require('./LiveResolverCache');
+const {
+  LiveResolverCache,
+  RELAY_RESOLVER_LIVE_STATE_SUBSCRIPTION_KEY,
+  getUpdatedDataIDs,
+} = require('./LiveResolverCache');
 const invariant = require('invariant');
 
 export type LiveState<+T> = {
@@ -706,6 +710,17 @@ class LiveResolverStore implements Store {
         for (let ii = 0; ii < storeIDs.length; ii++) {
           const dataID = storeIDs[ii];
           if (!references.has(dataID)) {
+            const record = this._recordSource.get(dataID);
+            if (record != null) {
+              const maybeResolverSubscription = RelayModernRecord.getValue(
+                record,
+                RELAY_RESOLVER_LIVE_STATE_SUBSCRIPTION_KEY,
+              );
+              if (maybeResolverSubscription != null) {
+                // $FlowFixMe - this value if it is not null, it is a function
+                maybeResolverSubscription();
+              }
+            }
             this._recordSource.remove(dataID);
           }
         }