Understanding Jotai: A Fresh Take on React State Management

React state management has evolved significantly over the years, from the simple useState hook to complex state management libraries. Today, let's explore Jotai, a primitive and flexible state management library that brings an atom-based approach to React applications.

What is Jotai?

Jotai (序態, which means "state" in Japanese) is a state management library that uses the concept of atomic state pieces. Think of atoms as the smallest possible units of state that can't be broken down further – just like atoms in chemistry!

The Mental Model: Thinking in Atoms

1. Atoms as Building Blocks

Imagine you're building with LEGO blocks. Each LEGO piece is like an atom in Jotai:

• It's a standalone unit

• It can connect with other pieces

• It can be reused across different constructions

import { atom } from 'jotai'

// Creating a basic atom
const counterAtom = atom(0)

2. Derived States as Computed Values

Just like how you can combine LEGO pieces to create more complex structures, Jotai allows you to derive new atoms from existing ones:

const counterAtom = atom(0)
const doubledCounterAtom = atom(
  (get) => get(counterAtom) * 2
)

3. Write-Only Atoms as Actions

Think of write-only atoms as remote controls that can modify other atoms:

const incrementAtom = atom(
  null, // read is null
  (get, set) => set(counterAtom, get(counterAtom) + 1)
)

Key Concepts Through Examples

1. Basic Usage

function Counter() {
  const [count, setCount] = useAtom(counterAtom)

  return (
    <button onClick={() => setCount(c => c + 1)}>
      Count: {count}
    </button>
  )
}

2. Loadable Utility for Async State

import { atom, useAtom } from 'jotai'
import { loadable } from 'jotai/utils'

// Create an async atom
const asyncDataAtom = atom(async () => {
  const response = await fetch('https://api.example.com/data')
  return response.json()
})

// Wrap it with loadable
const loadableDataAtom = loadable(asyncDataAtom)

function DataComponent() {
  const [data] = useAtom(loadableDataAtom)

  switch (data.state) {
    case 'loading':
      return <div>Loading...</div>
    case 'hasError':
      return <div>Error: {data.error}</div>
    case 'hasData':
      return <div>Data: {JSON.stringify(data.data)}</div>
  }
}

3. Complex Derived Atoms

import { atom } from 'jotai'

// Base atoms
const usersAtom = atom([
  { id: 1, name: 'John', role: 'admin' },
  { id: 2, name: 'Jane', role: 'user' }
])
const searchQueryAtom = atom('')
const filterRoleAtom = atom('all')

// Derived atom with multiple dependencies
const filteredUsersAtom = atom((get) => {
  const users = get(usersAtom)
  const query = get(searchQueryAtom).toLowerCase()
  const roleFilter = get(filterRoleAtom)

  return users
    .filter(user => 
      user.name.toLowerCase().includes(query) &&
      (roleFilter === 'all' || user.role === roleFilter)
    )
})

// Usage Example
function UserList() {
  const [filteredUsers] = useAtom(filteredUsersAtom)
  const [query, setQuery] = useAtom(searchQueryAtom)
  const [roleFilter, setRoleFilter] = useAtom(filterRoleAtom)

  return (
    <div>
      <input
        value={query}
        onChange={(e) => setQuery(e.target.value)}
        placeholder="Search users..."
      />
      <select
        value={roleFilter}
        onChange={(e) => setRoleFilter(e.target.value)}
      >
        <option value="all">All Roles</option>
        <option value="admin">Admin</option>
        <option value="user">User</option>
      </select>
      <ul>
        {filteredUsers.map(user => (
          <li key={user.id}>{user.name} ({user.role})</li>
        ))}
      </ul>
    </div>
  )
}

4. Advanced Async Patterns with Error Handling

const userAtom = atom(null)

const fetchUserAtom = atom(
  null,
  async (get, set, userId) => {
    try {
      set(userAtom, { status: 'loading' })
      const response = await fetch(`/api/users/${userId}`)
      if (!response.ok) throw new Error('Failed to fetch user')
      const data = await response.json()
      set(userAtom, { status: 'success', data })
    } catch (error) {
      set(userAtom, { status: 'error', error: error.message })
    }
  }
)

function UserProfile({ userId }) {
  const [user] = useAtom(userAtom)
  const [, fetchUser] = useAtom(fetchUserAtom)

  useEffect(() => {
    fetchUser(userId)
  }, [userId])

  if (user?.status === 'loading') return <div>Loading...</div>
  if (user?.status === 'error') return <div>Error: {user.error}</div>
  if (user?.status === 'success') {
    return <div>Welcome, {user.data.name}!</div>
  }
  return null
}

5. Using Scope Provider for Isolated State

import { Provider, atom, useAtom } from 'jotai'

const counterAtom = atom(0)

function Counter() {
  const [count, setCount] = useAtom(counterAtom)
  return (
    <button onClick={() => setCount(c => c + 1)}>
      Count: {count}
    </button>
  )
}

function App() {
  return (
    <div>
      <h2>Counter A</h2>
      <Provider>
        <Counter />
      </Provider>

      <h2>Counter B (separate state)</h2>
      <Provider>
        <Counter />
      </Provider>
    </div>
  )
}

6. Real-time Updates with WebSocket

const websocketAtom = atom((get) => {
  const ws = new WebSocket('wss://api.example.com')
  return ws
})

const messagesAtom = atom([])

const messageListenerAtom = atom(null, (get, set) => {
  const ws = get(websocketAtom)

  ws.addEventListener('message', (event) => {
    const message = JSON.parse(event.data)
    set(messagesAtom, (prev) => [...prev, message])
  })

  return () => {
    ws.close()
  }
})

function ChatRoom() {
  const [messages] = useAtom(messagesAtom)
  const [, initializeWebSocket] = useAtom(messageListenerAtom)

  useEffect(() => {
    initializeWebSocket()
  }, [])

  return (
    <div>
      {messages.map((msg, index) => (
        <div key={index}>{msg.content}</div>
      ))}
    </div>
  )
}

Best Practices for Scaling

1. Organize Atoms by Feature

// atoms/auth.js
export const userAtom = atom(null)
export const isAuthenticatedAtom = atom(
  (get) => get(userAtom) !== null
)

// atoms/todos.js
export const todosAtom = atom([])
export const activeTodosAtom = atom(
  (get) => get(todosAtom).filter(todo => !todo.completed)
)

2. Create Custom Hooks for Complex Logic

function useAuthActions() {
  const [, setUser] = useAtom(userAtom)

  const login = useCallback(async (credentials) => {
    const response = await fetch('/api/login', {
      method: 'POST',
      body: JSON.stringify(credentials)
    })
    const user = await response.json()
    setUser(user)
  }, [])

  return { login }
}

Why Choose Jotai?

1. Primitive First: Jotai is built on primitive atoms, making it easy to compose and reason about state.

2. No Boilerplate: Unlike Redux or MobX, there's minimal setup required.

3. TypeScript-First: Built with TypeScript from the ground up.

4. Tree-Shakeable: Only bundle what you use.

Conclusion

Jotai's atomic approach to state management brings a fresh perspective to React applications. Its mental model of small, composable pieces of state makes it intuitive to use and reason about. The examples above demonstrate how Jotai can handle everything from simple counters to complex real-time applications.

Key takeaways:

• Think in atoms (small, composable units of state)

• Use derived atoms for computed values

• Leverage async atoms and loadable utilities for data fetching

• Organize atoms by feature for better maintainability

• Use providers for state isolation when needed

Whether you're building a small application or a complex system, Jotai's flexibility and simplicity make it a compelling choice for React state management.

Happy coding! 🚀

---

Don't forget to follow me for more React and JavaScript content! If you found this helpful, feel free to share and discuss in the comments below.