Micro-State management with Zustand in React.

Micro-State management in React with Zustand

State management is one of the most important parts of frontend development, As the web apps are getting more and more complex with time to provide the most optimal user experience, and to achieve this they are experimenting with different architectures to have a frictionless development experience.

The larger the application, the more complex the state management. As the team starts to grow it becomes harder to manage the state of the application, thus the states can be broken down into smaller chunks and can be limited to components or module level, or business logic.

The concept of the Micro-State management is derived from the concept of microservices, The idea is to slice a monolith frontend into smaller chunks that can be managed, developed, and deployed independently by smaller dev teams.

We will see how we can do the micro-state management with Zustand in React. Zustand is a tiny library created to provide a module state for React. It is based on the Immutable update model, in which states objects cannot be modified but always have to be newly created.

Using selectors we can optimize the render manually, it also comes with a straightforward, yet powerful state creator interface.

Understanding module state and immutable state.

Module state

In module state, stores are defined in modules and we export them, they work on the Singleton + Factory design pattern where all the boilerplate logic to create the store and update them will reside in a module (function) that can be exported and used.

The store will share the state wherever used, which means if the same state has been accessed in two different components and one component updates the state, the second component will also re-render.

We can create a store and use that as a hook within the functional component.

//store.js
import { create } from "zustand";
export const store = create(() => ({ count: 0 }));

Immutable state.

Zustand is based on the Immutable state model, in which you are not allowed to modify the state objects directly, if you want to update the state, you will have to create a new object, you can reuse the unmodified state objects. The benefit of the immutable state model is that you only need to check state object referential equality to know if there’s any update; you don’t have to check equality deeply.

Let’s import the store and see the working of it with the following example.

// App.js
import { store } from "./store";

function App(){
  console.log(store.getState()); // {count: 0};
  return null;
}

store function exposes few methods like getState(), setState(), and subscribe().

getState() is used to get the state in the store and setState() is used to update the state in the store.

console.log(store.getState()); // ---> { count: 0 }

store.setState({ count: 1 });

console.log(store.getState()); // ---> { count: 1 }

The state objects are immutable, updating the state properties directly won’t update the state and thus will not trigger the re-render.

const state1 = store.getState();

state1.count = 2; // invalid, won't work;

store.setState(state1);

The store must be always updated with the new object, store.setState({count: 2}), it also accepts a callback function, which has the prevState as the argument, which can be used to create the new state, known as functional update.

store.setState((prev) => ({ count: prev.count + 1 }));

Our state has only one property currently, but it can have more than that, Let’s add one more property.

export const store = create(() => ({ count: 0, message: 'Hello World!' }));

This will also be updated through immutability.

store.setState({count: 1, message: 'Hello World!'});

However, it is not necessary to update all the properties every time, setState merges the old state with the new state, thus we can specify the only property which we want to update.

store.setState({count: 1});
console.log(store.getState()); // {count: 1, message: 'Hello World!'}

It uses Object.assign internally to merge the object.

Object.assign({}, oldState, newState);

The last piece of the puzzle is the store.subscribe() method, it registers a callback function which is triggered every time the state is updated in store.

store.subscribe(() => {
  console.log("state is updated");
});

store.setState({ message: 'Bye World!' });
//"state is updated"

Optimizing re-rendering.

When using Global states, it becomes necessary to optimize the re-renders because not all the components will be using all the properties of the state, let us see how can we address this using Zustand.

The good thing about the Zustand is that we can create hooks with different stores and use them within the components as per the need.

import { create } from "zustand";

const useStore = create((set) => ({
  count: 0,
  message: "Hello World!",
  setCount: (value) => set((state) => ({ count: state.count + value })),
  setMessage: (message) => set(() => ({ message })),
}));

export default useStore;

The create invokes the callback function which has a parameter set that can be used to update the properties of the state.

With the above structure, we get separate methods to update each state independently providing us better control over that state.

Using the useStore hook.

import useStore from "./store";

function App() {
  const {count, message} = useStore();
  console.log(count, message); // 0, "Hello World!"
  return null;
}

export default App;

The useStore hook by default returns the state object, which has all the properties, the bad thing about this is that, if any of those properties updates, it will trigger re-render in each component where this useStore hook is used.

Alternatively, you can pull the particular properties from the state to optimize the re-render, this will only render the component again when that particular property changes.

import useStore from "./store";

function App() {
  const count = useStore((state) => state.count);
  console.log(count); // 0
  return null;
}

export default App;

It is referred to as Manual render optimization when we do selector-based extra re-render control.The way the selector works to avoid re-renders is to compare the results of what the selector function returns. You need to be careful when you’re defining a selector function to return stable results to avoid re-renders.

For example, the below code won’t trigger render if we click on the Update Message button because the count property is not being changed.

import useStore from "./store";

function App() {
  const count = useStore((state) => state.count);
  const setMessage = useStore((state) => state.setMessage);
  const increaseCount = useStore((state) => state.increaseCount);

  console.log(count);
  return (
    <>
      <p>Count: {count}</p>
      <button onClick={() => increaseCount(1)}>Increment</button>
      <button onClick={() => setMessage("Bye World!")}>Update Message</button>
    </>
  );
}

export default App;

Alternatively if you do the same thing this way, by accessing all the properties once. This will trigger re-render everytime as the count, setMessage, increaseCount all of them are being returned as new object.

import useStore from "./store";

function App() {
  const { count, setMessage, increaseCount } = useStore((state) => ({
    count: state.count,
    setMessage: state.setMessage,
    increaseCount: state.increaseCount,
  }));

  console.log(count);

  return (
    <>
      <p>Count: {count}</p>
      <button onClick={() => increaseCount(1)}>Increment</button>
      <button onClick={() => setMessage("Bye World!")}>Update Message</button>
    </>
  );
}

export default App;

Middlewares

Zustand comes with multiple middlewares which can be used to increase its power further.

devtools

devtools middleware can be used to integrate zustand with the Redux-Devtools for a better developer experience.

import { create } from "zustand";
import { devtools } from "zustand/middleware";

const store = (set) => ({
  count: 0,
  message: "Hello World!",
  increaseCount: (value) => set((state) => ({ count: state.count + value })),
  setMessage: (message) => set(() => ({ message })),
});

//add devtools (redux devtools)
const useStore = create(devtools(store));

export default useStore;

persist

As the name suggests persist can be used to persist the state in the local storage. Allowing cross-module sharing of the state, allowing different teams to access the states.

import { create } from "zustand";
import { devtools, persist } from "zustand/middleware";

let store = (set) => ({
  count: 0,
  message: "Hello World!",
  increaseCount: (value) => set((state) => ({ count: state.count + value })),
  setMessage: (message) => set(() => ({ message })),
});

//add devtools (redux devtools)
store = devtools(store);

//persist the state with key "randomKey"
store = persist(store, { name: "randomKey" });

//create the store
let useStore = create(store);

export default useStore;

You can create multiple states in a single module and export them and utilize them as per the requirement, it is very flexible yet extremely powerful, experiment with it.


Summary

To summarize, the key point is that React is based on object immutability for optimization, Zustand having the same model as React gives us a huge benefit in terms of library simplicity and its small bundle size.

  • Reading state: This utilizes selector functions to optimize re-renders.
  • Writing state: This is based on the immutable state model.