Refactor link actions

src/components/LinkWithStatus.tsx

@@ -1,6 +1,6 @@
 'use client';
 
-import { ComponentProps, ReactNode, useEffect, useRef, useState } from 'react';
+import { ComponentProps, ReactNode, useState } from 'react';
 import Link from 'next/link';
 import LinkWithStatusChild from './primitives/LinkWithStatusChild';
 import clsx from 'clsx/lite';
@@ -14,7 +14,7 @@ export default function LinkWithStatus({
   onLoad,
   flickerThreshold,
   ...props
-}: Omit<ComponentProps<typeof Link>, 'children'> & {
+}: Omit<ComponentProps<typeof Link>, 'children' | 'onLoad'> & {
   children: ReactNode | ((props: { isLoading: boolean }) => ReactNode)
   loadingClassName?: string
   // For hoisting state to a parent component, e.g., <EntityLink />
@@ -29,24 +29,6 @@ export default function LinkWithStatus({
 
   const isControlled = typeof children === 'function';
 
-  const didStartLoading = useRef(false);
-  useEffect(() => {
-    console.log('link: 01', {isLoading, ref: didStartLoading.current});
-    if (isLoading) {
-      console.log('link: 02', {isLoading, ref: didStartLoading.current});
-      didStartLoading.current = true;
-    } else if (didStartLoading.current) {
-      console.log('link: 04', {isLoading, ref: didStartLoading.current});
-      onLoad?.();
-      didStartLoading.current = false;
-    }
-    return () => {
-      // Call onload when component unmounts while loading
-      console.log('link: 03', {isLoading, ref: didStartLoading.current});
-      if (isLoading) { onLoad?.(); }
-    };
-  }, [isLoading, onLoad]);
-
   return <Link
     {...props}
     className={clsx(
@@ -58,7 +40,7 @@ export default function LinkWithStatus({
       isLoading && loadingClassName,
     )}
   >
-    <LinkWithStatusChild {...{ setIsLoading, flickerThreshold }}>
+    <LinkWithStatusChild {...{ setIsLoading, flickerThreshold, onLoad }}>
       {typeof children === 'function'
         ? children({ isLoading })
         : children}

src/components/primitives/LinkWithStatusChild.tsx

@@ -8,10 +8,12 @@ const DEFAULT_FLICKER_THRESHOLD = 400;
 export default function LinkWithStatusChild({
   children,
   setIsLoading,
+  onLoad,
   flickerThreshold = DEFAULT_FLICKER_THRESHOLD,
 }: {
   children: ReactNode
   setIsLoading: (isLoading: boolean) => void
+  onLoad?: () => void
   flickerThreshold?: number
 }) {
   const { pending } = useLinkStatus();
@@ -38,12 +40,22 @@ export default function LinkWithStatusChild({
         isLoadingStartTime.current = undefined;
       }, Math.max(0, flickerThreshold - loadingDuration));
     }
+    return () => {
+      clearTimeout(startLoadingTimeout.current);
+      clearTimeout(stopLoadingTimeout.current);
+    };
   }, [pending, setIsLoading, flickerThreshold]);
 
-  useEffect(() => () => {
-    clearTimeout(startLoadingTimeout.current);
-    clearTimeout(stopLoadingTimeout.current);
-  }, []);
+  useEffect(() => {
+    if (!pending && startLoadingTimeout.current) {
+      onLoad?.();
+    }
+    return () => {
+      if (pending && startLoadingTimeout.current) {
+        onLoad?.();
+      }
+    };
+  }, [pending, onLoad]);
 
   return <>{children}</>;
 }