Skip to content

Commit

Permalink
fix ref merging for callback refs that return a cleanup function (#68123
Browse files Browse the repository at this point in the history
)

### What

Fix handling of user-provided refs in `Image` and `Link` to support
[refs with cleanup
functions](https://react.dev/blog/2024/04/25/react-19#cleanup-functions-for-refs)
(new in React 19)

### Why

React 19 allows [a new form of callback
ref](https://react.dev/blog/2024/04/25/react-19#cleanup-functions-for-refs),
which returns a cleanup function:
```ts
(element: HTMLElement) => {
  console.log('hi ref!')
  return () => {
    console.log('bye ref!')
  }
}
```
Unfortunately, this can be a breaking change for code that attempts to
combine ("merge") multiple callback refs into one. With old-style
callback refs, we never had to care about the return values. But with
the introduction of cleanup functions, ignoring the return value can
lead to logic errors (because the cleanup will never be called)!

This was the case for `Image` and `Link`, both of which attempt to
combine user-provided refs with their own callback refs, but were
ignoring the return values, so they'd break if the user's ref returned a
cleanup function.

### How

This PR introduces a `useMergedRef(refA, refB)` hook that combines all
forms of refs correctly (thanks @eps1lon for sketching it out!) and
refactors Image and Link to use it.
  • Loading branch information
lubieowoce authored Jul 25, 2024
1 parent 16a55ae commit ccbcfff
Show file tree
Hide file tree
Showing 7 changed files with 218 additions and 61 deletions.
103 changes: 50 additions & 53 deletions packages/next/src/client/image-component.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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}
Expand All @@ -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(
Expand Down
13 changes: 5 additions & 8 deletions packages/next/src/client/link.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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> = {
Expand Down Expand Up @@ -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) {
Expand All @@ -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.
Expand Down
46 changes: 46 additions & 0 deletions packages/next/src/client/use-merged-ref.ts
Original file line number Diff line number Diff line change
@@ -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
}
}
}
52 changes: 52 additions & 0 deletions test/integration/link-ref/pages/child-ref-func-cleanup.js
Original file line number Diff line number Diff line change
@@ -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>
)
}
8 changes: 8 additions & 0 deletions test/integration/link-ref/test/index.test.js
Original file line number Diff line number Diff line change
Expand Up @@ -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)(
Expand All @@ -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')
})
}
)
})
43 changes: 43 additions & 0 deletions test/integration/next-image-new/app-dir/app/ref-cleanup/page.js
Original file line number Diff line number Diff line change
@@ -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>
)
}
14 changes: 14 additions & 0 deletions test/integration/next-image-new/app-dir/test/index.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import {
nextBuild,
nextStart,
renderViaHTTP,
retry,
waitFor,
} from 'next-test-utils'
import webdriver from 'next-webdriver'
Expand Down Expand Up @@ -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', () => {
Expand Down

0 comments on commit ccbcfff

Please sign in to comment.