From 417a4d538e57f02ae7f8070bf923fdba9a2e3bbc Mon Sep 17 00:00:00 2001
From: Ivan Shmidt <bahoc47@gmail.com>
Date: Mon, 15 Jun 2020 22:36:36 +0200
Subject: [PATCH] feat: more clean hot hookz

---
 packages/backend-common/src/hot.ts            | 73 +++++++++++++++----
 .../src/service/ServiceBuilderImpl.ts         | 13 ++--
 packages/backend/src/plugins/catalog.ts       |  5 +-
 3 files changed, 66 insertions(+), 25 deletions(-)

diff --git a/packages/backend-common/src/hot.ts b/packages/backend-common/src/hot.ts
index 3043bcb681..f684b5cba1 100644
--- a/packages/backend-common/src/hot.ts
+++ b/packages/backend-common/src/hot.ts
@@ -13,11 +13,20 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-export function useHotEffect(
-  _module: NodeModule,
-  effectFactory: () => () => void,
-) {
-  const cancelEffect = effectFactory();
+
+/**
+ * This function allows devs to cleanup
+ * ongoing effects when module gets hot-reloaded
+ * Useful for cleaning intervals, timers, requests etc
+ * @example
+ * ```ts
+ * const intervalId = setInterval(doStuff, 1000);
+ * useHotCleanup(module, () => clearInterval(intervalId));
+ * ```
+ * @param _module Reference to the current module where you invoke the fn
+ * @param cancelEffect Fn that cleans up the ongoing effects
+ */
+export function useHotCleanup(_module: NodeModule, cancelEffect: () => void) {
   if (_module.hot) {
     _module.hot.addDisposeHandler(() => {
       cancelEffect();
@@ -25,27 +34,59 @@ export function useHotEffect(
   }
 }
 
+/**
+ * This function allows devs to preserve
+ * some value between hot-reloads.
+ * Useful for stateful parts of the backend
+ * @example
+ * ```ts
+ * const db = useHotMemoize(module, () => createDB(dbParams));
+ * ```
+ * @param _module Reference to the current module where you invoke the fn
+ * @param valueFactory Fn that returns the value you want to memoize
+ * @warning Don't use inside conditionals or loops,
+ * same rules as for hooks apply (https://reactjs.org/docs/hooks-rules.html)
+ */
 export function useHotMemoize<T>(
   _module: NodeModule,
   valueFactory: () => T,
 ): T {
+  const CURRENT_HOT_MEMOIZE_INDEX_KEY = 'backstage.io/hmr-memoize-key';
+
   if (!_module.hot) {
+    // Just return value straight away
     return valueFactory();
   }
-  const index = (useHotMemoize as any).index ?? 0;
-  (useHotMemoize as any).index += 1;
-  const prevValue = _module.hot?.data?.[index];
+
+  if (_module.hot && typeof _module.hot.data === 'undefined') {
+    // First run, init the module data
+    _module.hot.data = {
+      [CURRENT_HOT_MEMOIZE_INDEX_KEY]: 0,
+    };
+  }
+
+  // Let's store data per module based on the order of the code invocation
+  const index = _module.hot.data?.[CURRENT_HOT_MEMOIZE_INDEX_KEY];
+  // Increasing the counter after each call
+  _module.hot.data[CURRENT_HOT_MEMOIZE_INDEX_KEY] += 1;
+
+  const prevValue = _module.hot.data?.[index];
+  const createDisposeHandler = (value: any) => (data: {
+    [key: number]: any;
+    [indexKey: string]: number;
+  }) => {
+    // Preserving the value through the HMR process
+    data[index] = value;
+    // Decreasing the counter after each handler
+    data[CURRENT_HOT_MEMOIZE_INDEX_KEY] = index - 1;
+  };
+
   if (prevValue) {
-    _module.hot!.addDisposeHandler(data => {
-      data[index] = prevValue;
-    });
+    _module.hot!.addDisposeHandler(createDisposeHandler(prevValue));
     return prevValue;
   }
+
   const newValue = valueFactory();
-  if (_module.hot) {
-    _module.hot.addDisposeHandler(data => {
-      data[index] = newValue;
-    });
-  }
+  _module.hot.addDisposeHandler(createDisposeHandler(newValue));
   return newValue;
 }
diff --git a/packages/backend-common/src/service/ServiceBuilderImpl.ts b/packages/backend-common/src/service/ServiceBuilderImpl.ts
index 844990a39d..931d29b841 100644
--- a/packages/backend-common/src/service/ServiceBuilderImpl.ts
+++ b/packages/backend-common/src/service/ServiceBuilderImpl.ts
@@ -28,7 +28,7 @@ import {
   requestLoggingHandler,
 } from '../middleware';
 import { ServiceBuilder } from './types';
-import { useHotEffect } from '../hot';
+import { useHotCleanup } from '../hot';
 
 const DEFAULT_PORT = 7000;
 
@@ -96,12 +96,11 @@ export class ServiceBuilderImpl implements ServiceBuilder {
         0,
       );
 
-      useHotEffect(this.module, () => {
-        return () =>
-          server.stop((e: any) => {
-            if (e) console.error(e);
-          });
-      });
+      useHotCleanup(this.module, () =>
+        server.stop((e: any) => {
+          if (e) console.error(e);
+        }),
+      );
 
       resolve(server);
     });
diff --git a/packages/backend/src/plugins/catalog.ts b/packages/backend/src/plugins/catalog.ts
index 6f82836a79..41e4226001 100644
--- a/packages/backend/src/plugins/catalog.ts
+++ b/packages/backend/src/plugins/catalog.ts
@@ -24,7 +24,7 @@ import {
   runPeriodically,
 } from '@backstage/plugin-catalog-backend';
 import { PluginEnvironment } from '../types';
-import { useHotEffect } from '@backstage/backend-common';
+import { useHotCleanup } from '@backstage/backend-common';
 
 export default async function createPlugin({
   logger,
@@ -42,7 +42,8 @@ export default async function createPlugin({
     logger,
   );
 
-  useHotEffect(module, () =>
+  useHotCleanup(
+    module,
     runPeriodically(() => higherOrderOperation.refreshAllLocations(), 10000),
   );