From 8127491b8beee39cfd6758dcf6e44ecfe49872fd Mon Sep 17 00:00:00 2001
From: Patrik Oldsberg <poldsberg@gmail.com>
Date: Sun, 21 Jun 2020 14:21:51 +0200
Subject: [PATCH] backend-common: change useHotCleanup to clean up on parents
 being disposed as well

---
 packages/backend-common/src/hot.ts | 48 ++++++++++++++++++++++++++----
 1 file changed, 43 insertions(+), 5 deletions(-)

diff --git a/packages/backend-common/src/hot.ts b/packages/backend-common/src/hot.ts
index 3d4252b378..bd6454c537 100644
--- a/packages/backend-common/src/hot.ts
+++ b/packages/backend-common/src/hot.ts
@@ -14,10 +14,38 @@
  * limitations under the License.
  */
 
+// Find all active hot module APIs of all ancestors of a module, including the module itself
+function findAllAncestors(_module: NodeModule): NodeModule[] {
+  const ancestors = new Array<NodeModule>();
+  const parentIds = new Set<string | number>();
+
+  function add(id: string | number, m: NodeModule) {
+    if (parentIds.has(id)) {
+      return;
+    }
+    parentIds.add(id);
+    ancestors.push(m);
+
+    for (const parentId of (m as any).parents) {
+      const parent = require.cache[parentId];
+      if (parent) {
+        add(parentId, parent);
+      }
+    }
+  }
+
+  add(_module.id, _module);
+
+  return ancestors;
+}
+
 /**
- * This function allows devs to cleanup
- * ongoing effects when module gets hot-reloaded
+ * useHotCleanup allows cleanup of ongoing effects when a module is
+ * hot-reloaded during development. The cleanup function will be called
+ * whenever the module itself or any of its parent modules is hot-reloaded.
+ *
  * Useful for cleaning intervals, timers, requests etc
+ *
  * @example
  * ```ts
  * const intervalId = setInterval(doStuff, 1000);
@@ -28,9 +56,19 @@
  */
 export function useHotCleanup(_module: NodeModule, cancelEffect: () => void) {
   if (_module.hot) {
-    _module.hot.addDisposeHandler(() => {
-      cancelEffect();
-    });
+    const ancestors = findAllAncestors(_module);
+    let cancelled = false;
+
+    const handler = () => {
+      if (!cancelled) {
+        cancelled = true;
+        cancelEffect();
+      }
+    };
+
+    for (const m of ancestors) {
+      m.hot?.addDisposeHandler(handler);
+    }
   }
 }