[image, link] fix ref merging for callback refs that return a cleanup function

packages/next/src/client/image-component.tsx

@@ -31,6 +31,7 @@ import { RouterContext } from '../shared/lib/router-context.shared-runtime'
 
 // @ts-ignore - This is replaced by webpack alias
 import defaultLoader from 'next/dist/shared/lib/image-loader'
+import { useMergedRef } from './use-merged-ref'
 
 // This is replaced by webpack define plugin
 const configEnv = process.env.__NEXT_IMAGE_OPTS as any as ImageConfigComplete
@@ -206,6 +207,54 @@ const ImageElement = forwardRef<HTMLImageElement | null, ImageElementProps>(
     },
     forwardedRef
   ) => {
+    const ownRef = useCallback(
+      (img: ImgElementWithDataProp | null) => {
+        if (!img) {
+          return
+        }
+        if (onError) {
+          // If the image has an error before react hydrates, then the error is lost.
+          // The workaround is to wait until the image is mounted which is after hydration,
+          // then we set the src again to trigger the error handler (if there was an error).
+          // eslint-disable-next-line no-self-assign
+          img.src = img.src
+        }
+        if (process.env.NODE_ENV !== 'production') {
+          if (!src) {
+            console.error(`Image is missing required "src" property:`, img)
+          }
+          if (img.getAttribute('alt') === null) {
+            console.error(
+              `Image is missing required "alt" property. Please add Alternative Text to describe the image for screen readers and search engines.`
+            )
+          }
+        }
+        if (img.complete) {
+          handleLoading(
+            img,
+            placeholder,
+            onLoadRef,
+            onLoadingCompleteRef,
+            setBlurComplete,
+            unoptimized,
+            sizesInput
+          )
+        }
+      },
+      [
+        src,
+        placeholder,
+        onLoadRef,
+        onLoadingCompleteRef,
+        setBlurComplete,
+        onError,
+        unoptimized,
+        sizesInput,
+      ]
+    )
+
+    const ref = useMergedRef(forwardedRef, ownRef)
+
     return (
       <img
         {...rest}
@@ -229,59 +278,7 @@ const ImageElement = forwardRef<HTMLImageElement | null, ImageElementProps>(
         sizes={sizes}
         srcSet={srcSet}
         src={src}
-        ref={useCallback(
-          (img: ImgElementWithDataProp | null) => {
-            if (forwardedRef) {
-              if (typeof forwardedRef === 'function') forwardedRef(img)
-              else if (typeof forwardedRef === 'object') {
-                // @ts-ignore - .current is read only it's usually assigned by react internally
-                forwardedRef.current = img
-              }
-            }
-            if (!img) {
-              return
-            }
-            if (onError) {
-              // If the image has an error before react hydrates, then the error is lost.
-              // The workaround is to wait until the image is mounted which is after hydration,
-              // then we set the src again to trigger the error handler (if there was an error).
-              // eslint-disable-next-line no-self-assign
-              img.src = img.src
-            }
-            if (process.env.NODE_ENV !== 'production') {
-              if (!src) {
-                console.error(`Image is missing required "src" property:`, img)
-              }
-              if (img.getAttribute('alt') === null) {
-                console.error(
-                  `Image is missing required "alt" property. Please add Alternative Text to describe the image for screen readers and search engines.`
-                )
-              }
-            }
-            if (img.complete) {
-              handleLoading(
-                img,
-                placeholder,
-                onLoadRef,
-                onLoadingCompleteRef,
-                setBlurComplete,
-                unoptimized,
-                sizesInput
-              )
-            }
-          },
-          [
-            src,
-            placeholder,
-            onLoadRef,
-            onLoadingCompleteRef,
-            setBlurComplete,
-            onError,
-            unoptimized,
-            sizesInput,
-            forwardedRef,
-          ]
-        )}
+        ref={ref}
         onLoad={(event) => {
           const img = event.currentTarget as ImgElementWithDataProp
           handleLoading(

packages/next/src/client/link.tsx

@@ -22,6 +22,7 @@ import { useIntersection } from './use-intersection'
 import { getDomainLocale } from './get-domain-locale'
 import { addBasePath } from './add-base-path'
 import { PrefetchKind } from './components/router-reducer/router-reducer-types'
+import { useMergedRef } from './use-merged-ref'
 
 type Url = string | UrlObject
 type RequiredKeys<T> = {
@@ -546,7 +547,7 @@ const Link = React.forwardRef<HTMLAnchorElement, LinkPropsReal>(
       rootMargin: '200px',
     })
 
-    const setRef = React.useCallback(
+    const setIntersectionWithResetRef = React.useCallback(
       (el: Element) => {
         // Before the link getting observed, check if visible state need to be reset
         if (previousAs.current !== as || previousHref.current !== href) {
@@ -556,16 +557,12 @@ const Link = React.forwardRef<HTMLAnchorElement, LinkPropsReal>(
         }
 
         setIntersectionRef(el)
-        if (childRef) {
-          if (typeof childRef === 'function') childRef(el)
-          else if (typeof childRef === 'object') {
-            childRef.current = el
-          }
-        }
       },
-      [as, childRef, href, resetVisible, setIntersectionRef]
+      [as, href, resetVisible, setIntersectionRef]
     )
 
+    const setRef = useMergedRef(setIntersectionWithResetRef, childRef)
+
     // Prefetch the URL if we haven't already and it's visible.
     React.useEffect(() => {
       // in dev, we only prefetch on hover to avoid wasting resources as the prefetch will trigger compiling the page.

packages/next/src/client/use-merged-ref.ts

@@ -0,0 +1,46 @@
+import { useMemo, type Ref } from 'react'
+
+export function useMergedRef<TElement>(
+  refA: Ref<TElement>,
+  refB: Ref<TElement>
+): Ref<TElement> {
+  return useMemo(() => mergeRefs(refA, refB), [refA, refB])
+}
+
+export function mergeRefs<TElement>(
+  refA: Ref<TElement>,
+  refB: Ref<TElement>
+): Ref<TElement> {
+  if (!refA || !refB) {
+    return refA || refB
+  }
+
+  return (current: TElement) => {
+    const cleanupA = applyRef(refA, current)
+    const cleanupB = applyRef(refB, current)
+
+    return () => {
+      cleanupA()
+      cleanupB()
+    }
+  }
+}
+
+function applyRef<TElement>(
+  refA: NonNullable<Ref<TElement>>,
+  current: TElement
+) {
+  if (typeof refA === 'function') {
+    const cleanup = refA(current)
+    if (typeof cleanup === 'function') {
+      return cleanup
+    } else {
+      return () => refA(null)
+    }
+  } else {
+    refA.current = current
+    return () => {
+      refA.current = null
+    }
+  }
+}

test/integration/link-ref/pages/child-ref-func-cleanup.js

@@ -0,0 +1,52 @@
+import React from 'react'
+import Link from 'next/link'
+import { useCallback, useRef, useEffect, useState } from 'react'
+import { flushSync } from 'react-dom'
+
+export default function Page() {
+  const [isVisible, setIsVisible] = useState(true)
+
+  const statusRef = useRef({ wasInitialized: false, wasCleanedUp: false })
+
+  const refWithCleanup = useCallback((el) => {
+    if (!el) {
+      console.error(
+        'callback refs that returned a cleanup should never be called with null'
+      )
+      return
+    }
+
+    statusRef.current.wasInitialized = true
+    return () => {
+      statusRef.current.wasCleanedUp = true
+    }
+  }, [])
+
+  useEffect(() => {
+    const timeout = setTimeout(
+      () => {
+        flushSync(() => {
+          setIsVisible(false)
+        })
+        if (!statusRef.current.wasInitialized) {
+          console.error('callback ref was not initialized')
+        }
+        if (!statusRef.current.wasCleanedUp) {
+          console.error('callback ref was not cleaned up')
+        }
+      },
+      100 // if we hide the Link too quickly, the prefetch won't fire, failing a test
+    )
+    return () => clearTimeout(timeout)
+  }, [])
+
+  if (!isVisible) {
+    return null
+  }
+
+  return (
+    <Link href="/" ref={refWithCleanup}>
+      Click me
+    </Link>
+  )
+}

test/integration/link-ref/test/index.test.js

@@ -84,6 +84,10 @@ describe('Invalid hrefs', () => {
       it('should handle child ref that is a function', async () => {
         await noError('/child-ref-func')
       })
+
+      it('should handle child ref that is a function that returns a cleanup function', async () => {
+        await noError('/child-ref-func-cleanup')
+      })
     }
   )
   ;(process.env.TURBOPACK_DEV ? describe.skip : describe)(
@@ -109,6 +113,10 @@ describe('Invalid hrefs', () => {
       it('should preload with child ref with function', async () => {
         await didPrefetch('/child-ref-func')
       })
+
+      it('should preload with child ref with function that returns a cleanup function', async () => {
+        await didPrefetch('/child-ref-func-cleanup')
+      })
     }
   )
 })

test/integration/next-image-new/app-dir/app/ref-cleanup/page.js

@@ -0,0 +1,43 @@
+'use client'
+import Image from 'next/image'
+import { useCallback, useEffect, useState } from 'react'
+
+export default function Home() {
+  const [displayImage, setDisplayImage] = useState(true)
+
+  const refWithCleanup = useCallback((el) => {
+    if (!el) {
+      throw new Error(
+        'callback refs that returned a cleanup should never be called with null'
+      )
+    }
+
+    return () => {
+      console.log('callback ref was cleaned up')
+    }
+  }, [])
+
+  useEffect(() => {
+    setDisplayImage(false)
+  }, [])
+
+  return (
+    <main>
+      <h1>Should call ref cleanup on unmount</h1>
+      <section>
+        {displayImage ? (
+          <div style={{ position: 'relative', width: 10, height: 10 }}>
+            <Image
+              ref={refWithCleanup}
+              priority
+              fill
+              src="/test.jpg"
+              alt="alt"
+              sizes="10px"
+            />
+          </div>
+        ) : null}
+      </section>
+    </main>
+  )
+}

test/integration/next-image-new/app-dir/test/index.test.ts

@@ -14,6 +14,7 @@ import {
   nextBuild,
   nextStart,
   renderViaHTTP,
+  retry,
   waitFor,
 } from 'next-test-utils'
 import webdriver from 'next-webdriver'
@@ -1574,6 +1575,19 @@ function runTests(mode) {
       }
     }
   })
+
+  it('should call callback ref cleanups when unmounting', async () => {
+    const browser = await webdriver(appPort, '/ref-cleanup')
+    const getLogs = async () => (await browser.log()).map((log) => log.message)
+
+    await retry(async () => {
+      expect(await getLogs()).toContain('callback ref was cleaned up')
+    })
+
+    expect(await getLogs()).not.toContain(
+      'callback refs that returned a cleanup should never be called with null'
+    )
+  })
 }
 
 describe('Image Component Default Tests', () => {