atomWithDebounce - Update without debounce when needed

[VQuesada]

Hello,

I'm new to Jotai and I'm loving this library so far. It's so easy to create global state and custom hooks to modify/access it.

Recently, I've needed to have a state with debounce, for which I decided to use the atomWithDebounce code provided in Recipes https://jotai.org/docs/recipes/atom-with-debounce.

It works as expected, but I needed an extra feature: the possiblity to update both the currentValueAtom and the debouncedValueAtom immediately when needed, without waiting for the timeout.

A couple examples of when such use case is useful are:
1- When the user removes the text from the search bar: you don't need to wait to cleanup the searched elements
2- When the user selects a suggested search text: you don't need to debounce to avoid unneeded fetches, the user wants that text and you can fetch it right away

What I've done to add this extra option has been the following:

export interface UpdateParams<T> {
  update: SetStateAction<T>
  setWithoutDebounce?: boolean
}

export default function atomWithDebounce<T>(
  initialValue: T,
  delayMilliseconds = 500,
  shouldDebounceOnReset = false
) {
  const prevTimeoutAtom = atom<ReturnType<typeof setTimeout> | undefined>(
    undefined
  )

  // DO NOT EXPORT currentValueAtom as using this atom to set state can cause
  // inconsistent state between currentValueAtom and debouncedValueAtom
  const _currentValueAtom = atom(initialValue)
  const isDebouncingAtom = atom(false)

  const debouncedValueAtom = atom(
    initialValue,
    (get, set, updateParams: UpdateParams<T>) => {
      const { update, } = updateParams
      clearTimeout(get(prevTimeoutAtom))

      const prevValue = get(_currentValueAtom)
      const nextValue =
        typeof update === 'function'
          ? (update as (prev: T) => T)(prevValue)
          : update

      const onDebounceStart = () => {
        set(_currentValueAtom, nextValue)
        set(isDebouncingAtom, true)
      }

      const onDebounceEnd = () => {
        set(debouncedValueAtom, nextValue)
        set(isDebouncingAtom, false)
      }

      onDebounceStart()

      if ((!shouldDebounceOnReset && nextValue === initialValue) || updateParams.setWithoutDebounce === true) {
        onDebounceEnd()
        return
      }

      const nextTimeoutId = setTimeout(() => {
        onDebounceEnd()
      }, delayMilliseconds)

      // set previous timeout atom in case it needs to get cleared
      set(prevTimeoutAtom, nextTimeoutId)
    }
  )

  // exported atom setter to clear timeout if needed
  const clearTimeoutAtom = atom(null, (get, set, _arg) => {
    clearTimeout(get(prevTimeoutAtom))
    set(isDebouncingAtom, false)
  })

  return {
    currentValueAtom: atom((get) => get(_currentValueAtom)),
    isDebouncingAtom,
    clearTimeoutAtom,
    debouncedValueAtom,
  }
}

Basically, I made the update param from the atom setter to be an object:

export interface UpdateParams<T> {
  update: SetStateAction<T>
  setWithoutDebounce?: boolean
}

And I then take the update value from it to update both currentValueAtom and debouncedValueAtom, proceeding with it as in the default version of this code.

//  ...previous code
const debouncedValueAtom = atom(
    initialValue,
    (get, set, updateParams: UpdateParams<T>) => {
      const { update, } = updateParams
// ...following code

If I tell it to update the atoms without debouncing it just does so

if ((!shouldDebounceOnReset && nextValue === initialValue) || updateParams.setWithoutDebounce === true) {
  onDebounceEnd()
  return
}

The thing is that, evendough the code works as expected, typescript shouts at me on this line of code:

const onDebounceEnd = () => {
  set(debouncedValueAtom, nextValue) // Argument of type 'T' is not assignable to parameter of type 'UpdateParams<T>'
  set(isDebouncingAtom, false)
}

As I told it that updateParams is an object of type UpdateParams<T> it expects me to pass a compatible object to this setter. Instead, I am passing it just the new value since I don't need it to store the setWithoutDebounce value.

A workaround would be to set the debouncedValueAtom with the whole object and take the value within the custom hook like this:

export const useSearchTextDebouncedValue = () => {
  const searchTextDebounced = useAtomValue(searchTextDebouncedAtom)
  
  return {
    searchTextDebounced :  searchTextDebounced.update,
  }
}

But then I am just moving the error from the previous file to this one, as typescript expects searchText to be of type string, not UpdateParams<string>. I also prefer to set the atom with the actual value, it doesn't need to store if it was updated with or without debounce.

The code works, so it's obviously just a typing-related issue. Is there anything I'm missing that could make typescript chill out?

Thank you in anticipation.

[DavidJFelix]

Hey -- I forked the atomWithDebounce to try to demonstrate how I might introduce an action that bypasses the debouncer. Take a look here: https://codesandbox.io/s/elegant-shape-63lpyx

[DavidJFelix]

I made a few modifications -- first, I created a new "action" type modeled after reducer behavior and created a simple action creator that defaults to debounce but if you pass a second argument you can make it immediate. There might be a cleaner way or more preferable way to do this kind of switching.

I also added a new internal atom for the debouncedValue -- having a setter that takes in a complex type and unwraps it cannot use the simplified write only path that was being done before, so what we do is turn our debouncedValue atom into a derived writer that unwraps that action type.

import { atom, SetStateAction } from "jotai";

type DebounceAction<T> = {
  type: "debounce" | "immediate";
  value: T;
};

export function debounceAction<T>(
  value: T,
  immediate: boolean | undefined = false
): DebounceAction<T> {
  return {
    type: immediate ? "immediate" : "debounce",
    value
  };
}

export function atomWithDebounce<T>(
  initialValue: T,
  delayMilliseconds = 500,
  shouldDebounceOnReset = false
) {
  const prevTimeoutAtom = atom<ReturnType<typeof setTimeout> | undefined>(
    undefined
  );

  // DO NOT EXPORT currentValueAtom as using this atom to set state can cause
  // inconsistent state between currentValueAtom and debouncedValueAtom
  const _currentValueAtom = atom(initialValue);
  const baseDebouncedValueAtom = atom(initialValue);
  const isDebouncingAtom = atom(false);

  const debouncedValueAtom = atom(
    (get) => get(baseDebouncedValueAtom),
    (get, set, update: DebounceAction<SetStateAction<T>>) => {
      clearTimeout(get(prevTimeoutAtom));

      const prevValue = get(_currentValueAtom);
      const nextValue =
        typeof update.value === "function"
          ? (update.value as (prev: T) => T)(prevValue)
          : update.value;

      const onDebounceStart = () => {
        set(_currentValueAtom, nextValue);
        set(isDebouncingAtom, true);
      };

      const onDebounceEnd = () => {
        set(baseDebouncedValueAtom, nextValue);
        set(isDebouncingAtom, false);
      };

      onDebounceStart();

      if (
        (!shouldDebounceOnReset && nextValue === initialValue) ||
        update.type === "immediate"
      ) {
        onDebounceEnd();
        return;
      }

      const nextTimeoutId = setTimeout(() => {
        onDebounceEnd();
      }, delayMilliseconds);

      // set previous timeout atom in case it needs to get cleared
      set(prevTimeoutAtom, nextTimeoutId);
    }
  );

  // exported atom setter to clear timeout if needed
  const clearTimeoutAtom = atom(null, (get, set, _arg) => {
    clearTimeout(get(prevTimeoutAtom));
    set(isDebouncingAtom, false);
  });

  return {
    currentValueAtom: atom((get) => get(_currentValueAtom)),
    isDebouncingAtom,
    clearTimeoutAtom,
    debouncedValueAtom
  };
}

[VQuesada]

That's brilliant. Also, I see this makes the debouncedValueAtom get the baseDebouncedValue instead of the opposite, which may seem counterintuitive at first glance but it actually allows you to have the value and the setter within the same atom.

Thank you for your code, I think this should be displayed in the Recipes page.

[zknicker]

This is fantastic, thank you @DavidJFelix