Subscription mode (#1263)

* update subscribe with middleware

* fix lint & test

* update typing, fix test and lint

* refactor callback, rename, use refs for callbacks

* Delay unmount

* no swr destruction and add getters, revert pkg.json

* callback -> next

* rename fix types and update test

* change exports to unstable_subscription

* manage subs with useESE

* use serialize

* fix lint

* use cache helper to update error

* update exports path

* fix script

* pub check

* add exp jsdoc

* use cache-scoped storage

* use prefixed key; add more tests

* fix type

* remove subscriber ref

* rename

* fix lint

---------

Co-authored-by: Shu Ding <g@shud.in>

_internal/utils/mutate.ts

@@ -76,8 +76,8 @@ export async function internalMutate<Data>(
     for (let keyIt = it.next(); !keyIt.done; keyIt = it.next()) {
       const key = keyIt.value
       if (
-        // Skip the special useSWRInfinite keys.
-        !key.startsWith('$inf$') &&
+        // Skip the special useSWRInfinite and useSWRSubscription keys.
+        !/^\$(inf|sub)\$/.test(key) &&
         keyFilter((cache.get(key) as { _k: Arguments })._k)
       ) {
         matchedKeys.push(key)

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,136 @@
+import type { Key, SWRHook, Middleware, SWRConfiguration, SWRConfig } from 'swr'
+
+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 | null, 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>
+
+// [subscription count, disposer]
+type SubscriptionStates = [Map<string, number>, Map<string, () => void>]
+const subscriptionStorage = new WeakMap<object, SubscriptionStates>()
+
+const SUBSCRIPTION_PREFIX = '$sub$'
+
+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)
+
+    // Prefix the key to avoid conflicts with other SWR resources.
+    const subscriptionKey = key ? SUBSCRIPTION_PREFIX + key : undefined
+    const swr = useSWRNext(subscriptionKey, null, config)
+
+    const { cache } = config
+
+    // Ensure that the subscription state is scoped by the cache boundary, so
+    // you can have multiple SWR zones with subscriptions having the same key.
+    if (!subscriptionStorage.has(cache)) {
+      subscriptionStorage.set(cache, [
+        new Map<string, number>(),
+        new Map<string, () => void>()
+      ])
+    }
+
+    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+    const [subscriptions, disposers] = subscriptionStorage.get(cache)!
+
+    useIsomorphicLayoutEffect(() => {
+      if (!subscriptionKey) return
+
+      const [, set] = createCacheHelper<Data>(cache, subscriptionKey)
+      const refCount = subscriptions.get(subscriptionKey) || 0
+
+      const next = (error?: Error | null, data?: Data) => {
+        if (error !== null && typeof error !== 'undefined') {
+          set({ error })
+        } else {
+          swr.mutate(data, false)
+        }
+      }
+
+      // Increment the ref count.
+      subscriptions.set(subscriptionKey, refCount + 1)
+
+      if (!refCount) {
+        const dispose = subscribe(key, { next })
+        if (typeof dispose !== 'function') {
+          throw new Error(
+            'The `subscribe` function must return a function to unsubscribe.'
+          )
+        }
+        disposers.set(subscriptionKey, dispose)
+      }
+
+      return () => {
+        // Prevent frequent unsubscribe caused by unmount
+        setTimeout(() => {
+          // TODO: Throw error during development if count is undefined.
+          const count = subscriptions.get(subscriptionKey)
+          if (count == null) return
+
+          subscriptions.set(subscriptionKey, count - 1)
+
+          // Dispose if it's the last one.
+          if (count === 1) {
+            const dispose = disposers.get(subscriptionKey)
+            dispose?.()
+          }
+        })
+      }
+
+      // eslint-disable-next-line react-hooks/exhaustive-deps
+    }, [subscriptionKey])
+
+    return {
+      get data() {
+        return swr.data
+      },
+      get error() {
+        return swr.error
+      }
+    }
+  }) as unknown as Middleware
+
+/**
+ * A hook to subscribe a SWR resource to an external data source for continuous updates.
+ * @experimental This API is experimental and might change in the future.
+ * @example
+ * ```jsx
+ * import useSWRSubscription from 'swr/subscription'
+ *
+ * const { data, error } = useSWRSubscription(key, (key, { next }) => {
+ *   const unsubscribe = dataSource.subscribe(key, (err, data) => {
+ *     next(err, data)
+ *   })
+ *   return unsubscribe
+ * })
+ * ```
+ */
+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,139 @@
+import React from 'react'
+import { act, screen } from '@testing-library/react'
+import { sleep, renderWithConfig, createKey } from './utils'
+
+import useSWRSubscription from 'swr/subscription'
+import useSWR from 'swr'
+
+describe('useSWRSubscription', () => {
+  it('should update the state', async () => {
+    const swrKey = createKey()
+
+    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`)
+  })
+
+  it('should deduplicate subscriptions', async () => {
+    const swrKey = createKey()
+
+    let subscriptionCount = 0
+
+    function subscribe(key, { next }) {
+      ++subscriptionCount
+      let res = 0
+      const intervalId = setInterval(() => {
+        if (res === 3) {
+          const err = new Error(key + 'error')
+          next(err)
+        } else {
+          next(undefined, key + res)
+        }
+        res++
+      }, 100)
+
+      return () => {
+        clearInterval(intervalId)
+      }
+    }
+
+    function Page() {
+      const { data, error } = useSWRSubscription(swrKey, subscribe, {
+        fallbackData: 'fallback'
+      })
+      useSWRSubscription(swrKey, subscribe)
+      useSWRSubscription(swrKey, subscribe)
+
+      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`)
+
+    expect(subscriptionCount).toBe(1)
+  })
+
+  it('should not conflict with useSWR state', async () => {
+    const swrKey = createKey()
+
+    function subscribe(key, { next }) {
+      let res = 0
+      const intervalId = setInterval(() => {
+        if (res === 3) {
+          const err = new Error(key + 'error')
+          next(err)
+        } else {
+          next(undefined, key + res)
+        }
+        res++
+      }, 100)
+
+      return () => {
+        clearInterval(intervalId)
+      }
+    }
+
+    function Page() {
+      const { data, error } = useSWRSubscription(swrKey, subscribe, {
+        fallbackData: 'fallback'
+      })
+      const { data: swrData } = useSWR(swrKey, () => 'swr')
+      return (
+        <div>
+          {swrData}:{error ? error.message : data}
+        </div>
+      )
+    }
+
+    renderWithConfig(<Page />)
+    await act(() => sleep(10))
+    screen.getByText(`swr:fallback`)
+    await act(() => sleep(100))
+    screen.getByText(`swr:${swrKey}0`)
+    await act(() => sleep(100))
+    screen.getByText(`swr:${swrKey}1`)
+    await act(() => sleep(100))
+    screen.getByText(`swr:${swrKey}2`)
+  })
+})

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
   },