From abe7409c53657191481fa39d52d19bd7094aab1b Mon Sep 17 00:00:00 2001
From: Dorian Maliszewski <dmaliszewski@scaleway.com>
Date: Mon, 31 May 2021 11:59:14 +0200
Subject: [PATCH] fix: enabled effect and drop unecessary refs (#213)

---
 .../src/__tests__/useDataLoader.test.js       |  2 +-
 packages/use-dataloader/src/useDataLoader.js  | 84 +++++++++++--------
 2 files changed, 51 insertions(+), 35 deletions(-)

diff --git a/packages/use-dataloader/src/__tests__/useDataLoader.test.js b/packages/use-dataloader/src/__tests__/useDataLoader.test.js
index 56093bd82..3fc7db3f5 100644
--- a/packages/use-dataloader/src/__tests__/useDataLoader.test.js
+++ b/packages/use-dataloader/src/__tests__/useDataLoader.test.js
@@ -256,8 +256,8 @@ describe('useDataLoader', () => {
     expect(result.current.isPolling).toBe(true)
     expect(result.current.isLoading).toBe(true)
     expect(result.current.isSuccess).toBe(false)
-    expect(method2).toBeCalledTimes(1)
     await waitForNextUpdate()
+    expect(method2).toBeCalledTimes(1)
     expect(result.current.isPolling).toBe(true)
     expect(result.current.isLoading).toBe(false)
     expect(result.current.isSuccess).toBe(true)
diff --git a/packages/use-dataloader/src/useDataLoader.js b/packages/use-dataloader/src/useDataLoader.js
index 3da360b7c..b3d7264c6 100644
--- a/packages/use-dataloader/src/useDataLoader.js
+++ b/packages/use-dataloader/src/useDataLoader.js
@@ -1,4 +1,11 @@
-import { useEffect, useLayoutEffect, useMemo, useReducer, useRef } from 'react'
+import {
+  useCallback,
+  useEffect,
+  useLayoutEffect,
+  useMemo,
+  useReducer,
+  useRef,
+} from 'react'
 import { useDataLoaderContext } from './DataLoaderProvider'
 import { ActionEnum, StatusEnum } from './constants'
 import reducer from './reducer'
@@ -75,10 +82,6 @@ const useDataLoader = (
   }, [cacheKeyPrefix, fetchKey])
 
   const previousDataRef = useRef()
-  const isMountedRef = useRef(enabled)
-  const methodRef = useRef(method)
-  const onSuccessRef = useRef(onSuccess)
-  const onErrorRef = useRef(onError)
 
   const isLoading = useMemo(() => status === StatusEnum.LOADING, [status])
   const isIdle = useMemo(() => status === StatusEnum.IDLE, [status])
@@ -90,53 +93,62 @@ const useDataLoader = (
     [isSuccess, isLoading, enabled, pollingInterval],
   )
 
-  const handleRequest = useRef(async cacheKey => {
-    try {
-      dispatch(Actions.createOnLoading())
-      const result = await methodRef.current?.()
+  const handleRequest = useCallback(
+    async cacheKey => {
+      try {
+        dispatch(Actions.createOnLoading())
+        const result = await method()
 
-      if (keepPreviousData) {
-        previousDataRef.current = getCachedData(cacheKey)
-      }
-      if (result !== undefined && result !== null && cacheKey)
-        addCachedData(cacheKey, result)
+        if (keepPreviousData) {
+          previousDataRef.current = getCachedData(cacheKey)
+        }
+        if (result !== undefined && result !== null && cacheKey)
+          addCachedData(cacheKey, result)
 
-      dispatch(Actions.createOnSuccess())
+        dispatch(Actions.createOnSuccess())
 
-      await onSuccessRef.current?.(result)
-    } catch (err) {
-      dispatch(Actions.createOnError(err))
-      await onErrorRef.current?.(err)
-    }
-  })
+        await onSuccess?.(result)
+      } catch (err) {
+        dispatch(Actions.createOnError(err))
+        await onError?.(err)
+      }
+    },
+    [
+      addCachedData,
+      getCachedData,
+      keepPreviousData,
+      method,
+      onError,
+      onSuccess,
+    ],
+  )
+
+  const handleRequestRef = useRef(handleRequest)
 
   useEffect(() => {
     let handler
-    if (isMountedRef.current) {
+    if (enabled) {
       if (isIdle) {
-        handleRequest.current(key)
+        handleRequestRef.current(key)
       }
-      if (pollingInterval && isSuccess) {
-        handler = setTimeout(() => handleRequest.current(key), pollingInterval)
+      if (pollingInterval && (isSuccess || isError)) {
+        handler = setTimeout(
+          () => handleRequestRef.current(key),
+          pollingInterval,
+        )
       }
     }
 
     return () => {
       if (handler) clearTimeout(handler)
     }
-    // Why can't put empty array for componentDidMount, componentWillUnmount ??? No array act like componentDidMount and componentDidUpdate
-  }, [key, pollingInterval, isIdle, isSuccess])
-
-  useLayoutEffect(() => {
-    methodRef.current = method
-  }, [method])
+  }, [key, pollingInterval, isIdle, isSuccess, isError, enabled])
 
   useLayoutEffect(() => {
-    isMountedRef.current = enabled
     dispatch(Actions.createReset())
     if (key && typeof key === 'string') {
       addReloadRef.current?.(key, reloadArgs =>
-        handleRequest.current(key, reloadArgs),
+        handleRequestRef.current(key, reloadArgs),
       )
     }
 
@@ -152,6 +164,10 @@ const useDataLoader = (
     addReloadRef.current = addReload
   }, [clearReload, addReload])
 
+  useLayoutEffect(() => {
+    handleRequestRef.current = handleRequest
+  }, [handleRequest])
+
   return {
     data: getCachedData(key) || initialData,
     error,
@@ -161,7 +177,7 @@ const useDataLoader = (
     isPolling,
     isSuccess,
     previousData: previousDataRef.current,
-    reload: args => handleRequest.current(key, args),
+    reload: args => handleRequestRef.current(key, args),
   }
 }