Nested slice

[armpogart]

I'm trying to implement some sort of nested slice. Below is minimal reproduction of what I'm doing. For simplicity all other slices are removed.

slices/chat.ts

import type { StateCreator } from 'zustand/vanilla';

import type { AllSlices } from '../index';

export interface Chat {
  id: number;
  name: string;
}

interface ChatState {
  resource: Chat | null;
  subscribed: boolean;
}

const chatStateInitial: ChatState = {
  resource: null,
  subscribed: false,
};

interface ChatActions {
  setSubscribed: (subscribed: boolean) => void;
}

export type ChatSlice = ChatState & ChatActions;

export const initChatSlice: (id: number) => (...params: Parameters<StateCreator<AllSlices>>) => ChatSlice =
  (id) => (set) => ({
    ...chatStateInitial,
    setSubscribed: (subscribed) =>
      set((state) => ({
        chats: {
          ...state.chats,
          byId: {
            ...state.chats.byId,
            [id]: {
              ...state.chats.byId[id],
              subscribed: subscribed,
            },
          },
        },
      })),
  });

slices/chats.ts

import type { StateCreator } from 'zustand/vanilla';

import type { AllSlices } from '../index';
import type { Chat, ChatSlice } from './chat';
import { initChatSlice } from './chat';

interface ChatsState {
  byId: Record<number, ChatSlice>;
  allIds: number[];
}

const chatsStateInitialValues: ChatsState = {
  byId: {},
  allIds: [],
};

interface ChatsActions {
  initChats: (chats: Chat[]) => void;
}

export interface ChatsSlice {
  chats: ChatsState & ChatsActions;
}

export const createChatsSlice: StateCreator<AllSlices, [], [], ChatsSlice> = (set, get, store) => ({
  chats: {
    ...chatsStateInitialValues,
    initChats: (chats) =>
      set((state) => ({
        chats: {
          ...state.chats,
          byId: Object.fromEntries(
            chats.map((chat) => [
              chat.id,
              {
                ...initChatSlice(chat.id)(set, get, store),
                resource: chat,
              },
            ]),
          ),
        },
      })),
  },
});

index.ts

import { create } from 'zustand/react';

import type { ChatsSlice } from './slices//chats';
import { createChatsSlice } from './slices/chats';

export type AllSlices = ChatsSlice;

export const useBoundStore = create<AllSlices>()((...state) => ({
  ...createChatsSlice(...state),
}));

initChats action is used in loader to fetch and initialized the conversations, it's doing it's job as intended.

Now after that I have the folllowing selector (hook):
selectors/chat.ts

import useSocket from '@hooks/use-socket';
import { useEffect } from 'react';
import { useShallow } from 'zustand/react/shallow';

import { useBoundStore } from '../index';

export function useChat(id: number, subscribe = true) {
  const socket = useSocket();
  const slice = useBoundStore(useShallow((state) => state.chats.byId[id]));
  const chat = useBoundStore(useShallow((state) => state.chats.byId[id].resource));

  useEffect(() => {
    const handler = () => console.log('ws message received');

    if (subscribe) {
      socket.on(`conversation/${id}`, handler);
      slice.setSubscribed(true);
    }

    return () => {
      if (subscribe && slice.subscribed) {
        socket.off(`conversation/${id}`, handler);
        slice.setSubscribed(false);
      }
    };
  }, [id, slice, socket, subscribe]);

  return {
    chat,
    subscribed: slice.subscribed,
  };
}

useSocket is from Context and holds socket connection and is working correct outside of this code.

1. First question is, whether this approach is ok, or I'm doing something incorrect completely? I'm new to zustand.
2. When I use useChat selector/hook, as far as I see it's working correct initially, but as soon as I do some local change on files and HMR triggers (I'm using vite.js), I'm getting the following error: Maximum update depth exceeded. This can happen when a component repeatedly calls setState inside componentWillUpdate or componentDidUpdate.. I guess that happens on unmount, when the useEffect return is called in my hook and slice.setSubscribed(false) is done. Do I have some unstable selector output, why does that happen?

Any help is appreciated!

Update 1:

selectors/chat.ts

import { useEffect } from 'react';
import { useShallow } from 'zustand/react/shallow';

import { useBoundStore } from '../index';

export function useChat(id: number, subscribe = true) {
  const { conversation, subscribed, setSubscribed } = useBoundStore(
    useShallow((state) => ({
      chat: state.chats.byId[id].resource,
      subscribed: state.chats.byId[id].subscribed,
      setSubscribed: state.chats.byId[id].setSubscribed,
    })),
  );

  useEffect(() => {
    if (subscribe && !subscribed) {
      console.log(`Subscribing to ${id}`);
      setSubscribed(true);
    }

    return () => {
      if (subscribe && subscribed) {
        console.log(`Unsubscribing from ${id}`);
        setSubscribed(false);
      }
    };
  }, [id, setSubscribed, subscribe, subscribed]);

  return {
    chat,
    subscribed,,
  };
}

Now with this modified version, where I've tried to make the selectors stable. I'm seeing double subscription (the console.log message) and I never see unsubscription, the useEffect return is called, but subscribed is never true there. As soon as I modify the return condition if (subscribe) .... removing subscribed, I'm immediately getting infinite loop (the mentioned error).

[dbritto-dev]

@armpogart I saw some issues here, that code could cause infinite loop

 useEffect(() => {
    const handler = () => console.log('ws message received');

    if (subscribe) {
      socket.on(`conversation/${id}`, handler);
      slice.setSubscribed(true);
    }

    return () => {
      if (subscribe && slice.subscribed) {
        socket.off(`conversation/${id}`, handler);
        slice.setSubscribed(false);
      }
    };
  }, [id, slice, socket, subscribe]);

[armpogart]

@dbritto-dev I guess due to setting setSubscribed on which the effect hook depends? Do you have any suggestion how to prevent the infinite loop in that case? I had something similar with pure react state (useState) in the component itself before integrating zustand and it was working. As soon as I remove slice.setSubscribed(false) in useEffect return, the error is gone.

Also is the approach of creating nested slice I use ok, or there is some more recommended (correct) approach with zustand also not to have unnecessary re-renders later the road.

[armpogart]

@dbritto-dev I've updated the question/post with modified selector/hook, trying to fix selectors stability and debug what's going on there and added more info.

[dbritto-dev]

@armpogart you should consider reentrancy which means that that can be called multiple times, might updating this part fix your issue and only subscribe when is not slice.subscribed

 if (subscribe && !slice.subscribed) {
      socket.on(`conversation/${id}`, handler);
      slice.setSubscribed(true);
    }

Caution

without checking if is not slice.subscribed you would ended up setting slice.subscribed to true which would
trigger a new render and that would rerun the useEffect again causing the infinite loop

[armpogart]

import { useEffect } from 'react';
import { useShallow } from 'zustand/react/shallow';

import { useBoundStore } from '../index';

export function useChat(id: number, subscribe = true) {
  const { conversation, subscribed, setSubscribed } = useBoundStore(
    useShallow((state) => ({
      chat: state.chats.byId[id].resource,
      subscribed: state.chats.byId[id].subscribed,
      setSubscribed: state.chats.byId[id].setSubscribed,
    })),
  );

  useEffect(() => {
    console.log(`Mount called for ${id}`);
    if (subscribe && !subscribed) {
      console.log(`Subscribing to ${id}`);
      setSubscribed(true);
    }

    return () => {
      console.log(`Unmount called for ${id}`);
      if (subscribe && subscribed) {
        console.log(`Unsubscribing from ${id}`);
        setSubscribed(false);
      }
    };
  }, [id, setSubscribed, subscribe, subscribed]);

  return {
    chat,
    subscribed,,
  };
}

@dbritto-dev With the above modified version I get the following logs

Mount called for #
Subscribing to #
Unmount called for #
Mount called for #
Subscribing to #
Unmount called for #
Mount called for #

where # is specific ID for my chats. So I get double subscription there, unmount is called, but subscribed is never true there for it to unsubscribe. After any change on the component using the hook, HMR triggers, and I get some proper unmounts with unsubsribe, but eventually I get 1 more subscription, so each HMR adds +1 subscription. If I'm adding the real socket subscription there and sending messages from backend, I see 2 handler calls initially, and +1 handler with each HMR change.

[dbritto-dev]

@armpogart double run on useEffect is expected when you use StrictMode

Ref: https://react.dev/reference/react/StrictMode