浅谈 NextJS RSC/SSR 中数据水合和持久化数据（二）

上回说道在 SSR 中如何用 React Query 实现数据水合。这期来谈谈如何用 Jotai 实现。

使用 Jotai 对数据的管理后期对数据的修改和对 UI 的反应也会更加方便。

改造

注入数据

在 layout.tsx 我们依然使用 queryClient 去获取数据，然后我们需要把获取到的数据直接传入 Jotai 的 atom 中。后期 client 端的组件的消费数据全部从 atom 中获取。大概是这样的：

// layout.tsx
export default async (
  props: NextPageParams<{
    id: string
  }>,
) => {
  const id = props.params.id
  const query = queries.note.byNid(id)
  const data = await getQueryClient().fetchQuery(query)

  return (
    <>
      {/* 确保第一位，需要在组件渲染之前完成注入 */}
      <CurrentNoteDataProvider data={data} />
      {props.children}
    </>
  )
}

上面的代码简略，我这里实现了一个 <CurrentNoteDataProvider /> 的组件，主要是把当前页面的数据源直接灌入一个 atom 中。稍后我会再介绍这个 Provider 的实现。

注意看，和之前不同的是，我们这里不再使用 React Query 提供的 Hydrate 去 Query 的水合，Client 侧的组件之后也不会再使用 useQuery 拿数据了。由于数据都在 Jotai 中，对数据的掌管变得非常简单，你想过滤什么数据就很方便了。

现在服务端的 QueryClient 可以是一个单例，而不是为每次 React Tree 的构建而创建一个新的实例，创建新的实例虽然能避免跨请求数据的污染，但是却不能享受到 Query 缓存带来的福利。所以现在，我们改造 getQueryClient 方法。

// query-client.server.ts
const sharedClient = new QueryClient({
  defaultOptions: {
    queries: {
      staleTime: 1000 * 3,
      cacheTime: 1000 * 3,
    },
  },
})
export const getQueryClient = () => sharedClient

如遇到需要鉴权的数据，如何解决数据污染呢。可以在 layout 直接判断和处理，需不需要将数据注入到 CurrentDataProvider。或者可以选择性的注入数据，仅仅保留公开的内容。

// layout.tsx
export default async (
  props: NextPageParams<{
    id: string
  }>,
) => {
  const id = props.params.id
  const query = queries.note.byNid(id)
  const data = await getQueryClient().fetchQuery(query)

  const filteredData = omit(data, ['some-private-fieled']) // <---- 不会影响缓存的过滤
  return (
    <>
      <CurrentNoteDataProvider data={filteredData} />

      {props.children}
    </>
  )
}

消费数据

我们来到一个 Client 侧的组件，需要消费到当前页面的数据。下面是一个简单组件，用于显示文章的标题。

export const NoteTitle = () => {
  const title = useCurrentNoteDataSelector((data) => data?.data.title)
  if (!title) return null
  return (
    <h1 className="mt-8 text-left font-bold text-base-content/95">{title}</h1>
  )
}

我们同样实现了一个 hook 名为 useCurrentNoteDataSelector ，直接从刚刚注入的数据中提取 title 字段，而不消费其他任何字段，后期动态修改页面数据源，这样会非常方便实现细粒度的更新。

其他组件需要使用任何数据也是如此，都可以通过 Selector 实现细粒度渲染。

实现 Provider 和 Hooks

这是一套通用方法，如果页面上不只存在一个数据源，则可以创建多个类似的钩子。

所以我们创建一个工厂函数 createDataProvider，批量生产这些。

'use client'

import { memo, useCallback, useEffect } from 'react'
import { produce } from 'immer'
import { atom, useAtomValue } from 'jotai'
import { selectAtom } from 'jotai/utils'
import type { FC, PropsWithChildren } from 'react'

import { useBeforeMounted } from '~/hooks/common/use-before-mounted'
import { noopArr } from '~/lib/noop'
import { jotaiStore } from '~/lib/store'

export const createDataProvider = <Model>() => {
  const currentDataAtom = atom<null | Model>(null)
  const CurrentDataProvider: FC<
    {
      data: Model
    } & PropsWithChildren
  > = memo(({ data, children }) => {
    useBeforeMounted(() => {
      jotaiStore.set(currentDataAtom, data)
    })

    useEffect(() => {
      jotaiStore.set(currentDataAtom, data)
    }, [data])

    useEffect(() => {
      return () => {
        jotaiStore.set(currentDataAtom, null)
      }
    }, [])

    return children
  })

  CurrentDataProvider.displayName = 'CurrentDataProvider'
  const useCurrentDataSelector = <T>(
    selector: (data: Model | null) => T,
    deps?: any[],
  ) => {
    const nextSelector = useCallback((data: Model | null) => {
      return data ? selector(data) : null
    }, deps || noopArr)

    return useAtomValue(selectAtom(currentDataAtom, nextSelector))
  }

  const setCurrentData = (recipe: (draft: Model) => void) => {
    jotaiStore.set(
      currentDataAtom,
      produce(jotaiStore.get(currentDataAtom), recipe),
    )
  }

  const getCurrentData = () => {
    return jotaiStore.get(currentDataAtom)
  }

  return {
    CurrentDataProvider,
    useCurrentDataSelector,
    setCurrentData,
    getCurrentData,
  }
}

在工厂函数中我们首先创建一个最基本的 currentDataAtom 。这个 atom 用于托管页面数据源。然后 CurrentDataProvider 传入一个 props 名为 data，然后在页面渲染前就把数据扔到 currentDataAtom 这步很重要。我们需要确保在渲染页面组件之前，currentDataAtom 内数据已经准备就绪。所以我们需要 useBeforeMounted 来实现同步注入数据。实现如下：

// use-before-mounted.ts
import { useRef } from 'react'

export const useBeforeMounted = (fn: () => any) => {
  const effectOnce = useRef(false)

  if (!effectOnce.current) {
    effectOnce.current = true
    fn?.()
  }
}

上面的代码在开发环境中使用，可能会得到一个 Warning，告诉你不应该在 render 函数中直接使用 setState。

useBeforeMounted(() => {
  // React 会发出警告，但这是合理的，你可以忽略。
  // 不要排斥这种方式，因为在新版 React 文档中会告诉你善用这种方式去优化性能。
  jotaiStore.set(currentDataAtom, data)
})

useCurrentDataSelector 就比较简单了，就是 Jotai 提供的 selectAtom 套了一个壳。

以上就是最基本的创建方法了。

现在我们创建一个 CurrentNoteProvider。

const {
  CurrentDataProvider,
  getCurrentData,
  setCurrentData,
  useCurrentDataSelector,
} = createDataProvider<NoteWrappedPayload>()

非常简单。

响应式修改数据

Jotai 的好处是状态和 UI 分离，有了上面的方法之后，现在我们不在需要过于关注数据的变化带来的 UI 更新的问题。在任何地方我们都能通过修改数据去驱动 UI 的更新。

现在我们有一个 Socket 连接，当收到 NOTE_UPDATE 事件之后，立即更新数据，驱动 UI 更新。我们可以这样写。

// event-handler.ts

import {
  getCurrentNoteData,
  setCurrentNoteData,
} from '~/providers/note/CurrentNoteDataProvider'

import { EventTypes } from '~/types/events'

export const eventHandler = (
  type: EventTypes,
  data: any,
  router: AppRouterInstance,
) => {
  switch (type) {
    case 'NOTE_UPDATE': {
      const note = data as NoteModel
      if (getCurrentNoteData()?.data.id === note.id) {
        setCurrentNoteData((draft) => {
          // <----- 直接修改数据
          Object.assign(draft.data, note)
        })
        toast('手记已更新')
      }
      break
    }
    default: {
      if (isDev) {
        console.log(type, data)
      }
    }
  }
}

脱离 UI 之后，修改数据变得十分简单，我们直接获取到 Jotai atom 内部的数据源，判断和事件中的 id 是否一致，然后直接更新数据。我们无需关心 UI 如何，只要使用 setCurrentNoteData 更新数据，UI 会立刻更新。而且这一举动是非常细粒度的。未涉及到的组件永远不会发生更新。可以通过一下文件查看更多。

https://github.com/Innei/sprightly/blob/14234594f44956e6f56f1f92952ce82db37ef4df/src/socket/handler.ts

数据隔离

使用了上面的方法保证了页面上数据源和 UI 分离。现在又有一个新的问题。页面组件过于依赖 CurrentDataAtom 的数据，但是 CurrentDataAtom 却只有一个。

现在我们的页面主要存在多个组件，假设 NoteTitle 存在两个，第一个需要展示 NoteId 为 15 的标题，而第二个需要展示 NoteId 为 16 的标题。

按照上面的结构，这个需求基本是不可能实现的，因为在 NoteTitle 中都使用了 useCurrentDataSelector 获取数据，而 Atom 却只有一个。

为了解决这个问题，我们需要知道 React Context 的三个特征

当组件不在 Provider 内部时，useContext 会返回一个 Context 的默认值，而这个默认值我们是可以定义的。
如果在 Provider 内部时，则会消费传入 Provider 的值。
如果在多层相同的 Provider 内部时，则会消费离组件最近的 Provider 的值。

综上所述，我们可以对 CurrentDataProvider 进行改造。

// createDataProvider.tsx

export const createDataProvider = <Model,>() => {
  const CurrentDataAtomContext = createContext(
    null! as PrimitiveAtom<null | Model>,
  )
  const globalCurrentDataAtom = atom<null | Model>(null)
  const CurrentDataAtomProvider: FC<
    PropsWithChildren<{
      overrideAtom?: PrimitiveAtom<null | Model>
    }>
  > = ({ children, overrideAtom }) => {
    return (
      <CurrentDataAtomContext.Provider
        value={overrideAtom ?? globalCurrentDataAtom}
      >
        {children}
      </CurrentDataAtomContext.Provider>
    )
  }
  const CurrentDataProvider: FC<
    {
      data: Model
    } & PropsWithChildren
  > = memo(({ data, children }) => {
    const currentDataAtom =
      useContext(CurrentDataAtomContext) ?? globalCurrentDataAtom

    useBeforeMounted(() => {
      jotaiStore.set(currentDataAtom, data)
    })

    useEffect(() => {
      jotaiStore.set(currentDataAtom, data)
    }, [data])

    useEffect(() => {
      return () => {
        jotaiStore.set(currentDataAtom, null)
      }
    }, [])

    return children
  })

  CurrentDataProvider.displayName = 'CurrentDataProvider'

  const useCurrentDataSelector = <T,>(
    selector: (data: Model | null) => T,
    deps?: any[],
  ) => {
    const currentDataAtom =
      useContext(CurrentDataAtomContext) ?? globalCurrentDataAtom
    const nextSelector = useCallback((data: Model | null) => {
      return data ? selector(data) : null
    }, deps || noopArr)

    return useAtomValue(selectAtom(currentDataAtom, nextSelector))
  }

  const useSetCurrentData = () =>
    useSetAtom(useContext(CurrentDataAtomContext) ?? globalCurrentDataAtom)

  const setGlobalCurrentData = (recipe: (draft: Model) => void) => {
    jotaiStore.set(
      globalCurrentDataAtom,
      produce(jotaiStore.get(globalCurrentDataAtom), recipe),
    )
  }

  const getGlobalCurrentData = () => {
    return jotaiStore.get(globalCurrentDataAtom)
  }

  return {
    CurrentDataAtomProvider,
    CurrentDataProvider,
    useCurrentDataSelector,
    useSetCurrentData,
    setGlobalCurrentData,
    getGlobalCurrentData,
  }
}

我们新增加了 React Context，把原本的 currentDataAtom 改成了 globalCurrentDataAtom 用于顶层页面数据，顶层页面数据吃到默认值，也就是我们不需要修改原来的任何代码。增加了 Scope 的 CurrentProvider可以在组件内部实现数据隔离。

现在我们对第二个需要展示和当前页面数据的不同的 NoteTitle，只要包上 CurrentDataProvider 就行了。

// layout.tsx

export default async (
  props: NextPageParams<{
    id: string
  }>,
) => {
  const id = props.params.id
  const query = queries.note.byNid(id)
  const data = await getQueryClient().fetchQuery(query)

  return (
    <>
      // 确保第一位，需要在组件渲染之前完成注入 // 这里我们不需要包括
      props.children
      <CurrentNoteDataProvider data={data} />
      {props.children}
    </>
  )
}

// page.tsx

export default function Page() {
  return (
    <>
      <NoteTitle />
      // 这里需要包括，并传入不同的 data
      <CurrentNoteDataProvider data={otherData}>
        <NoteTitle />
      </CurrentNoteDataProvider>
    </>
  )
}

基于上面的方法，就可以做到数据隔离，页面存在多个由 Jotai 管理的数据源互不干扰。

基于这个特性，我熬夜实现了 Shiro 的 Peek 功能。

https://github.com/Innei/Shiro/commit/e1b0b57aaea0eec1b695c4f1961297b42b935044

OK，今天就先说这么多。