Subscription mode

jest.config.js

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

package.json

@@ -33,6 +33,12 @@
       "require": "./immutable/dist/index.js",
       "types": "./immutable/dist/immutable/index.d.ts"
     },
+    "./unstable_subscription": {
+      "import": "./subscription/dist/index.mjs",
+      "module": "./subscription/dist/index.esm.js",
+      "require": "./subscription/dist/index.js",
+      "types": "./subscription/dist/subscription/index.d.ts"
+    },
     "./mutation": {
       "import": "./mutation/dist/index.mjs",
       "module": "./mutation/dist/index.esm.js",
@@ -51,6 +57,8 @@
     "infinite/dist/**",
     "immutable/dist/**",
     "mutation/dist/**",
+    "subscription/dist/**",
+    "subscription/package.json",
     "_internal/dist/**",
     "core/dist/package.json",
     "infinite/package.json",
@@ -64,16 +72,18 @@
   "license": "MIT",
   "scripts": {
     "clean": "rimraf core/dist infinite/dist immutable/dist mutation/dist",
-    "build": "yarn build:internal && yarn build:core && yarn build:infinite && yarn build:immutable && yarn build:mutation",
-    "watch": "npm-run-all -p watch:core watch:infinite watch:immutable watch:mutation",
+    "build": "yarn build:internal && yarn build:core && yarn build:infinite && yarn build:immutable && yarn build:mutation && yarn build:subscription",
+    "watch": "npm-run-all -p watch:core watch:infinite watch:immutable watch:mutation watch:subscription",
     "watch:core": "yarn build:core -w",
     "watch:infinite": "yarn build:infinite -w",
     "watch:immutable": "yarn build:immutable -w",
     "watch:mutation": "yarn build:mutation -w",
+    "watch:subscription": "yarn build:subscription -w",
     "build:core": "bunchee index.ts --cwd core --no-sourcemap",
     "build:infinite": "bunchee index.ts --cwd infinite --no-sourcemap",
     "build:immutable": "bunchee index.ts --cwd immutable --no-sourcemap",
     "build:mutation": "bunchee index.ts --cwd mutation --no-sourcemap",
+    "build:subscription": "bunchee index.ts --cwd subscription --no-sourcemap",
     "build:internal": "bunchee index.ts --cwd _internal --no-sourcemap",
     "prepublishOnly": "yarn clean && yarn build",
     "publish-beta": "yarn publish --tag beta",

subscription/index.ts

@@ -0,0 +1,100 @@
+import { useRef } from 'react'
+import useSWR from 'swr'
+import {
+  withMiddleware,
+  Key,
+  SWRHook,
+  Middleware,
+  serialize,
+  SWRConfiguration,
+  useIsomorphicLayoutEffect
+} 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
+  ): SWRSubscriptionResponse<Data, Error> => {
+    const [key] = serialize(_key)
+    const swr = useSWRNext(key, null, config)
+    const subscribeRef = useRef(subscribe)
+
+    useIsomorphicLayoutEffect(() => {
+      subscribeRef.current = subscribe
+    })
+
+    useIsomorphicLayoutEffect(() => {
+      subscriptions.set(key, (subscriptions.get(key) || 0) + 1)
+
+      const onData = (val?: Data) => swr.mutate(val, false)
+      const onError = async (err: any) => {
+        // Avoid thrown errors from `mutate`
+        // eslint-disable-next-line no-empty
+        try {
+          await swr.mutate(() => {
+            throw err
+          }, false)
+        } catch (_) {
+          /* eslint-disable-line no-empty */
+        }
+      }
+
+      const next = (_err?: any, _data?: Data) => {
+        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,12 @@
+{
+  "name": "swr-subscription",
+  "version": "0.0.1",
+  "main": "./dist/index.js",
+  "module": "./dist/index.esm.js",
+  "types": "./dist/subscription",
+  "exports": "./dist/index.mjs",
+  "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/unstable_subscription'
+
+describe('useSWRSubscription', () => {
+  it('should update state when fetcher is a subscription', async () => {
+    const key = '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(key, 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(`${key}0`)
+    await act(() => sleep(100))
+    screen.getByText(`${key}1`)
+    await act(() => sleep(100))
+    screen.getByText(`${key}2`)
+    await act(() => sleep(100))
+    screen.getByText(`${key}error`)
+    clearInterval(intervalId)
+    await sleep(100)
+    screen.getByText(`${key}error`)
+  })
+})

tsconfig.check.json

@@ -3,5 +3,5 @@
   "compilerOptions": {
     "rootDir": "."
   },
-  "include": ["src", "immutable", "infinite", "mutation"]
+  "include": ["src", "immutable", "infinite", "mutation", "subscription"]
 }

tsconfig.json

@@ -21,7 +21,8 @@
       "swr": ["./core/index.ts"],
       "swr/infinite": ["./infinite/index.ts"],
       "swr/immutable": ["./immutable/index.ts"],
-      "swr/mutation": ["./mutation/index.ts"]
+      "swr/mutation": ["./mutation/index.ts"],
+      "swr/unstable_subscription": ["subscription/index.ts"]
     },
     "typeRoots": ["./node_modules/@types"],
     "incremental": true