useCancellationToken: Avoid memory leaks in react

Inspired by Sophia Brant's article on Memory Leaks With React SetState I set about creating a reusable hook that can be used to mitigate not being able to cancel promises. I recommend reading that article first if you're unsure about what might cause memory leaks and the different approaches to avoiding them.

I have gone with an approach that creates a cancellation token which can either be cancelled manually or gets cancelled automatically if the component unmounts. FYI: I'm using typescript.

The scenario

We have a component which performs an asynchronous task, most likely a fetch, and then updates the component state afterwards, but it's possible the component has been unmounted before that request completes. If the state gets updated at this point we have a memory leak.

const [movies, setMovies] = useState([] as Movies[])
useEffect(() => {
    const action = async () => {
        const result = await fetch('http://example.com/movies.json')
        setMovies(result)
    }
    action()
}, [setMovies])

React doesn't support async lambdas in useEffect, so creating an async lambda within the lambda and calling it, as we do here, is a common workaround.
We're going to refactor this to use a cancellation token approach.

The token

First off we need a token that we can check for cancellation on, and which we can cancel.

interface CancellationToken {
    isCancelled: boolean
    cancel(): void
}

export function useCancellationToken(): CancellationToken {
    return useMemo(() => {
        const token = {
            isCancelled: false,
            cancel: () => {}
        }

        token.cancel = () => token.isCancelled = true

        return token as CancellationToken
    }, [])
}

This hook can be used to create a cancellation token when the component is mounted. The use of useMemo ensures it only gets created once so that when we cancel it, it stays cancelled.

I'm going to change the original use of useEffect to check if the token has been cancelled, and to call the cancel method on the token if the component is unmounted.

const [movies, setMovies] = useState([] as Movies[])
const cancellationToken = useCancellationToken()
useEffect(() => {
    const action = async () => {
        const result = await fetch('http://example.com/movies.json')
        if (cancellationToken.isCancelled) {
            return
        }
        setMovies(result)
    }
    action()
}, [setMovies, cancellationToken])
// If a function is returned from useEffect it is called when the component unmounts.
useEffect(() => () => cancellationToken.cancel(), [])

At this point we're avoiding the memory leaks by checking if the cancellation token has been cancelled. By returning a lambda to useEffect which calls cancellationToken.cancel() we're cancelling the token when the component is unmounted.

I went one step further and wrapped this bit of functionality in another hook, which I call useCancellableEffect. This also allows me to write the async lambda directly in to my hook without needing to use the workaround above.

The hook itself is:

export default function useCancellableEffect(action: () => void, dependencies: any[], cancellationToken: CancellationToken) {
    useEffect(() => {
        action()
        // eslint-disable-next-line
    }, [...dependencies, cancellationToken])
    useEffect(() => () => cancellationToken.cancel()
        // eslint-disable-next-line
    , [])
}

and the usage becomes

const [movies, setMovies] = useState([] as Movies[])
const cancellationToken = useCancellationToken()
useCancellableEffect(async () => {
    const result = await fetch('http://example.com/movies.json')
    if (cancellationToken.isCancelled) {
        return
    }
    setMovies(result)
}, [setMovies], cancellationToken)

which keeps all the boiler plate locked away in the hook, and only keeps what is relevant on the page. Ofcourse it is still up to the developer to check for cancellation and avoid memory leaks, but at least this helps make that easier. I also don't like the need to ... spread the dependencies and ignore the action dependency in the use of useEffect. If anyone comes up with a nice way of doing that without the need to disable the linter please let me know. The only approach I could think of for now was wrapping the action in useCallback, but that's more boilerplate again.

Note: An earlier version of this article didn't take in to account that useEffect calls the cleanup on every re render!!! The code snippets have been edited to account for this and handle it only when the component is unmounted.