Subscription mode

_internal/utils/mutate.ts

@@ -54,6 +54,8 @@ export async function internalMutate<Data>(
     typeof _opts === 'boolean' ? { revalidate: _opts } : _opts || {}
   )
 
+  console.log('opts', options)
+
   let populateCache = options.populateCache
 
   const rollbackOnErrorOption = options.rollbackOnError

jest.config.js

@@ -8,6 +8,7 @@ module.exports = {
     '^swr$': '<rootDir>/core/index.ts',
     '^swr/infinite$': '<rootDir>/infinite/index.ts',
     '^swr/immutable$': '<rootDir>/immutable/index.ts',
+    '^swr/subscription$': '<rootDir>/subscription/index.ts',
     '^swr/mutation$': '<rootDir>/mutation/index.ts',
     '^swr/_internal$': '<rootDir>/_internal/index.ts'
   },

package.json

@@ -34,6 +34,12 @@
       "module": "./immutable/dist/index.esm.js",
       "require": "./immutable/dist/index.js"
     },
+    "./subscription": {
+      "types": "./subscription/dist/index.d.ts",
+      "import": "./subscription/dist/index.mjs",
+      "module": "./subscription/dist/index.esm.js",
+      "require": "./subscription/dist/index.js"
+    },
     "./mutation": {
       "types": "./mutation/dist/mutation/index.d.ts",
       "import": "./mutation/dist/index.mjs",
@@ -53,11 +59,13 @@
     "immutable/dist/**/*.{js,d.ts,mjs}",
     "mutation/dist/**/*.{js,d.ts,mjs}",
     "_internal/dist/**/*.{js,d.ts,mjs}",
+    "subscription/dist/*.{js,d.ts,mjs}",
     "core/package.json",
     "infinite/package.json",
     "immutable/package.json",
     "mutation/package.json",
-    "_internal/package.json"
+    "_internal/package.json",
+    "subscription/package.json"
   ],
   "repository": "github:vercel/swr",
   "homepage": "https://swr.vercel.app",
@@ -68,7 +76,7 @@
     "csb:build": "pnpm build",
     "clean": "pnpm -r run clean && rimraf playwright-report test-result",
     "watch": "pnpm -r run watch",
-    "build": "pnpm build-package _internal && pnpm build-package core && pnpm build-package infinite && pnpm build-package immutable && pnpm build-package mutation",
+    "build": "pnpm build-package _internal && pnpm build-package core && pnpm build-package infinite && pnpm build-package immutable && pnpm build-package mutation && pnpm build-package subscription",
     "build:e2e": "pnpm next build e2e/site",
     "build-package": "bunchee index.ts --cwd",
     "types:check": "pnpm -r run types:check",

subscription/index.ts

@@ -0,0 +1,92 @@
+import type { Key, SWRHook, Middleware, SWRConfiguration, SWRConfig } from 'swr'
+import { useRef } from 'react'
+import useSWR from 'swr'
+import {
+  withMiddleware,
+  serialize,
+  useIsomorphicLayoutEffect,
+  createCacheHelper
+} from 'swr/_internal'
+
+export type SWRSubscription<Data = any, Error = any> = (
+  key: Key,
+  { next }: { next: (err?: Error, data?: Data) => void }
+) => void
+
+export type SWRSubscriptionResponse<Data = any, Error = any> = {
+  data?: Data
+  error?: Error
+}
+
+export type SWRSubscriptionHook<Data = any, Error = any> = (
+  key: Key,
+  subscribe: SWRSubscription<Data, Error>,
+  config?: SWRConfiguration
+) => SWRSubscriptionResponse<Data, Error>
+
+const subscriptions = new Map<Key, number>()
+const disposers = new Map()
+
+export const subscription = (<Data, Error>(useSWRNext: SWRHook) =>
+  (
+    _key: Key,
+    subscribe: SWRSubscription<Data, Error>,
+    config: SWRConfiguration & typeof SWRConfig.defaultValue
+  ): SWRSubscriptionResponse<Data, Error> => {
+    const [key] = serialize(_key)
+    const swr = useSWRNext(key, null, config)
+    const subscribeRef = useRef(subscribe)
+
+    useIsomorphicLayoutEffect(() => {
+      subscribeRef.current = subscribe
+    })
+
+    const { cache } = config
+    const [, set] = createCacheHelper<Data>(cache, key)
+
+    useIsomorphicLayoutEffect(() => {
+      subscriptions.set(key, (subscriptions.get(key) || 0) + 1)
+
+      const onData = (val?: Data) => swr.mutate(val, false)
+      const onError = async (err: any) => set({ error: err })
+
+      const next = (err_?: any, data_?: Data) => {
+        console.log('next:_err', err_)
+        if (err_) onError(err_)
+        else onData(data_)
+      }
+
+      if (subscriptions.get(key) === 1) {
+        const dispose = subscribeRef.current(key, { next })
+        disposers.set(key, dispose)
+      }
+      return () => {
+        // Prevent frequent unsubscribe caused by unmount
+        setTimeout(() => {
+          const count = subscriptions.get(key) || 1
+          subscriptions.set(key, count - 1)
+          // Dispose if it's last one
+          if (count === 1) {
+            disposers.get(key)()
+          }
+        })
+      }
+      // eslint-disable-next-line react-hooks/exhaustive-deps
+    }, [key])
+
+    return {
+      get data() {
+        return swr.data
+      },
+      get error() {
+        return swr.error
+      }
+    }
+  }) as unknown as Middleware
+
+const useSWRSubscription = withMiddleware(
+  useSWR,
+  subscription
+) as SWRSubscriptionHook
+
+export default useSWRSubscription

subscription/package.json

@@ -0,0 +1,17 @@
+{
+  "main": "./dist/index.js",
+  "module": "./dist/index.esm.js",
+  "types": "./dist/index.d.ts",
+  "exports": "./dist/index.mjs",
+  "private": true,
+  "scripts": {
+    "watch": "bunchee index.ts -w",
+    "build": "bunchee index.ts",
+    "types:check": "tsc --noEmit",
+    "clean": "rimraf dist"
+  },
+  "peerDependencies": {
+    "swr": "*",
+    "react": "*"
+  }
+}

subscription/tsconfig.json

@@ -0,0 +1,9 @@
+{
+  "extends": "../tsconfig.json",
+  "compilerOptions": {
+    "outDir": "./dist",
+    "rootDir": "..",
+  },
+  "include": [".", "../src"],
+  "exclude": ["./dist"]
+}

test/use-swr-subscription.test.tsx

@@ -0,0 +1,47 @@
+import React from 'react'
+import { act, screen } from '@testing-library/react'
+import { sleep, renderWithConfig } from './utils'
+import useSWRSubscription from 'swr/subscription'
+
+describe('useSWRSubscription', () => {
+  it('should update state when fetcher is a subscription', async () => {
+    const swrKey = 'sub-0'
+    let intervalId
+    let res = 0
+    function subscribe(key, { next }) {
+      intervalId = setInterval(() => {
+        if (res === 3) {
+          const err = new Error(key + 'error')
+          next(err)
+        } else {
+          next(undefined, key + res)
+        }
+        res++
+      }, 100)
+
+      return () => {}
+    }
+
+    function Page() {
+      const { data, error } = useSWRSubscription(swrKey, subscribe, {
+        fallbackData: 'fallback'
+      })
+      return <div>{error ? error.message : data}</div>
+    }
+
+    renderWithConfig(<Page />)
+    await act(() => sleep(10))
+    screen.getByText(`fallback`)
+    await act(() => sleep(100))
+    screen.getByText(`${swrKey}0`)
+    await act(() => sleep(100))
+    screen.getByText(`${swrKey}1`)
+    await act(() => sleep(100))
+    screen.getByText(`${swrKey}2`)
+    await act(() => sleep(100))
+    screen.getByText(`${swrKey}error`)
+    clearInterval(intervalId)
+    await sleep(100)
+    screen.getByText(`${swrKey}error`)
+  })
+})

tsconfig.json

@@ -22,7 +22,8 @@
       "swr/infinite": ["./infinite/index.ts"],
       "swr/immutable": ["./immutable/index.ts"],
       "swr/mutation": ["./mutation/index.ts"],
-      "swr/_internal": ["./_internal/index.ts"]
+      "swr/_internal": ["./_internal/index.ts"],
+      "swr/subscription": ["subscription/index.ts"],
     },
     "incremental": true
   },