Proposal: Support cleanup function in subscribe()

[ycmjason]

Hello,

In my use cases of store.subscribe, I often need to clean up resources when the store state updates. However, subscribe() does not currently support returning a cleanup function, making this process a bit cumbersome.

Currently, I have to manually track and call cleanup functions like this:

const useStore = create(() => ({ ... }));

let cleanup;
useStore.subscribe((state) => {
  cleanup?.(); // Manually call the previous cleanup function
  cleanup = someSideEffectThatReturnsCleanupFn(state);
});

To make this more ergonomic, I created a utility called withCleanup, which mimics useEffect's cleanup behavior:

const withCleanup = <T>(fn: (t: T) => (() => void) | undefined): ((t: T) => void) => {
  let cleanup: undefined | (() => void);
  return (t: T) => {
    cleanup?.();
    cleanup = fn(t);
  };
};

const useStore = create(() => ({ ... }));

useStore.subscribe(
  withCleanup((state) => {
    return someSideEffectThatReturnsCleanupFn(state);
  })
);

This works well, but I think it would be even better if Zustand's subscribe() natively supported cleanup functions, allowing this:

const useStore = create(() => ({ ... }));

useStore.subscribe((state) => {
  return someSideEffectThatReturnsCleanupFn(state); // Cleanup function is handled automatically
});

Proposal

Would it make sense to either:

1. Expose a built-in withCleanup utility in Zustand?
2. Modify subscribe() to allow returning a cleanup function?

I'm keen to hear your thoughts on whether this could be integrated into Zustand or if there are alternative recommendations for handling cleanup in subscribe().

Thanks for considering this! 🚀

[dbritto-dev]

@ycmjason what you mean? the .subscribe() method returns an unsubscribe function

Ref: https://zustand.docs.pmnd.rs/apis/create#subscribing-to-state-updates

[ycmjason]

agree this can be done in the userland as i have done already. just proposing as i thought it'd be a handy feature to have!

thank you for considering!

[dbritto-dev]

let userPostDatabaseUnsubscribe;
useUserStore.subscribe(({ currentUser }) => {
  if (!currentUser) return;
  if (userPostDatabaseUnsubscribe) userPostDatabaseUnsubscribe()
  userPostDatabaseUnsubscribe = userPostDatabase.onData(currentUser.id, (data) => {
    useUserPostStore.setState(data);
  });
});

[ycmjason]

Thank you @dbritto-dev !

I understand that this can be implemented in userland, as I already demonstrated in my the first example of the original post:

const useStore = create(() => ({ ... }));

let cleanup;
useStore.subscribe((state) => {
  cleanup?.(); // Manually call the previous cleanup function
  cleanup = someSideEffectThatReturnsCleanupFn(state);
});

My proposal is simply that Zustand could consider supporting this use case more ergonomically. Instead of manually tracking the cleanup variable, Zustand could provide built-in support for returning a cleanup function, similar to useEffect:

const useStore = create(() => ({ ... }));

useStore.subscribe((state) => {
  return someSideEffectThatReturnsCleanupFn(state); // Cleanup function runs automatically before the next update
});

This would streamline cases where subscriptions require per-update cleanup, making it easier to manage resources and prevent memory leaks.

It sounds like you don't feel like this is something Zustand would consider adding support for, which is totally fine! I just hadn’t seen any previous discussions on this, so I created this post to explore the idea. 😄

[dbritto-dev]

I'd say this good for a third-party lib middleware, btw would be inefficient using .subscribe() you might need subscribeWithSelector to avoid subscribing/unsubscribing same user

[ycmjason]

Thank you for pointing out.

I am just providing some really simple example to illustrate my proposal. I have been using subscribeWithSelector in my real implementation!