Subscription Middleware

Purpose

This middleware addresses the challenge of efficiently managing multiple concurrent subscriptions to external data sources within the SWR ecosystem. When many components subscribe to the same data stream, naive implementations risk redundant subscriptions and potential memory leaks. The Subscription Middleware introduces a reference-counting mechanism coupled with proper cleanup (disposers), ensuring that:

• Multiple subscribers to the same key share a single underlying subscription.

• Subscriptions are established only when needed.

• Subscriptions are properly disposed of when no longer in use.

This prevents resource leaks and excessive network or event listener usage, optimizing real-time data handling in React applications using SWR.

Functionality

At its core, the middleware wraps a SWR hook to intercept subscription logic:

1. Key Serialization and Prefixing:
   Input keys are serialized and prefixed to isolate subscription cache entries from regular SWR cache entries, preventing collisions.

2. Subscription State Storage:
   Using a WeakMap keyed by the SWR cache object, it maintains two maps per cache instance:

   • A map of subscription keys to their reference counts.

   • A map of subscription keys to their disposer functions.

3. Reference Counting:
   When a component invokes the subscription hook with a given key:

   • The reference count for that subscription key is incremented.

   • If this is the first subscriber (count was zero), the middleware calls the provided subscribe function to start the subscription and stores the disposer.

4. Data and Error Updates:
   The middleware provides a next callback to the subscribe function, which is used to update SWR's internal cache with new data or errors as they arrive asynchronously.

5. Cleanup on Unmount:
   When a subscriber unmounts or stops using the subscription, the reference count is decremented. If the count drops to zero, the stored disposer is called to clean up the subscription resource.

6. Exposing Data and Error:
   The middleware returns an object exposing the current data and error from SWR's cache, allowing components to reactively consume subscription updates.

Critical Interaction Snippet

useIsomorphicLayoutEffect(() => {
  if (!subscriptionKey) return

  const [, set] = createCacheHelper<Data>(cache, subscriptionKey)
  const refCount = subscriptions.get(subscriptionKey) || 0

  const next = (error, data) => {
    if (error != null) {
      set({ error })
    } else {
      set({ error: undefined })
      swr.mutate(data, false)
    }
  }

  subscriptions.set(subscriptionKey, refCount + 1)

  if (!refCount) {
    const dispose = subscribe(args, { next })
    if (typeof dispose !== 'function') {
      throw new Error('The `subscribe` function must return a function to unsubscribe.')
    }
    disposers.set(subscriptionKey, dispose)
  }

  return () => {
    const count = subscriptions.get(subscriptionKey)! - 1
    subscriptions.set(subscriptionKey, count)
    if (!count) {
      const dispose = disposers.get(subscriptionKey)
      dispose?.()
    }
  }
}, [subscriptionKey])

This code manages the lifecycle of subscriptions tied to the component lifecycle, ensuring subscriptions are shared and disposed correctly.

Integration

The Subscription Middleware integrates into the broader subscription system as a composable middleware layer that wraps the core SWR hook. It complements other subscription subtopics by:

• Providing the underlying mechanism to manage subscription lifecycles and resource cleanup.

• Enabling the useSWRSubscription hook to expose reactive data streams to React components efficiently.

• Working within the SWR cache boundaries to isolate subscription state per cache instance, allowing multiple SWR zones to operate independently.

• Collaborating with subscription types that define the contract for subscriber callbacks and options.

By abstracting subscription reference counting and cleanup, it allows higher-level subscription logic and components to focus on data consumption and UI without worrying about redundant subscriptions or leaks.

Diagram

flowchart TD
  A[Component Mounts with SubscriptionKey] --> B[Increment Ref Count]
  B --> C{Ref Count == 1?}
  C -->|Yes| D[Call subscribe() to Start Subscription]
  C -->|No| E[Reuse Existing Subscription]
  D --> F[Store disposer function]
  E --> F
  F --> G[Listen for next(error, data)]
  G --> H[Update SWR Cache with new data or error]
  A --> I[Component Unmounts]
  I --> J[Decrement Ref Count]
  J --> K{Ref Count == 0?}
  K -->|Yes| L[Call disposer() to clean up]
  K -->|No| M[Keep subscription alive]

This flowchart illustrates how subscription reference counting governs the creation, sharing, update, and disposal of subscriptions, ensuring efficient real-time data management within SWR.