Support for explicit resource management "using statements" on RTK Query subscriptions

packages/toolkit/src/query/core/buildInitiate.ts

@@ -16,6 +16,8 @@ import type { BaseQueryError, QueryReturnValue } from '../baseQueryTypes'
 import type { QueryResultSelectorResult } from './buildSelectors'
 import type { Dispatch } from 'redux'
 import { isNotNullish } from '../utils/isNotNullish'
+import type { DisposableIfAvailable } from '../utils/disposableIfAvailable'
+import { disposableIfAvailable } from '../utils/disposableIfAvailable'
 
 declare module './module' {
   export interface ApiEndpointQuery<
@@ -65,7 +67,7 @@ export type QueryActionCreatorResult<
   refetch(): QueryActionCreatorResult<D>
   updateSubscriptionOptions(options: SubscriptionOptions): void
   queryCacheKey: string
-}
+} & DisposableIfAvailable
 
 type StartMutationActionCreator<
   D extends MutationDefinition<any, any, any, any>
@@ -362,6 +364,16 @@ You must add the middleware for RTK-Query to function correctly!`
 
         const runningQuery = runningQueries.get(dispatch)?.[queryCacheKey]
         const selectFromState = () => selector(getState())
+        const unsubscribe = () => {
+          if (subscribe) {
+            dispatch(
+              unsubscribeQueryResult({
+                queryCacheKey,
+                requestId,
+              })
+            )
+          }
+        }
 
         const statePromise: QueryActionCreatorResult<any> = Object.assign(
           forceQueryFn
@@ -394,15 +406,7 @@ You must add the middleware for RTK-Query to function correctly!`
               dispatch(
                 queryAction(arg, { subscribe: false, forceRefetch: true })
               ),
-            unsubscribe() {
-              if (subscribe)
-                dispatch(
-                  unsubscribeQueryResult({
-                    queryCacheKey,
-                    requestId,
-                  })
-                )
-            },
+            unsubscribe,
             updateSubscriptionOptions(options: SubscriptionOptions) {
               statePromise.subscriptionOptions = options
               dispatch(
@@ -414,6 +418,7 @@ You must add the middleware for RTK-Query to function correctly!`
                 })
               )
             },
+            ...disposableIfAvailable(unsubscribe),
           }
         )
 

packages/toolkit/src/query/tests/buildInitiate.test.tsx

@@ -52,3 +52,30 @@ test('multiple synchonrous initiate calls with pre-existing cache entry', async
     requestId: thirdValue.requestId,
   })
 })
+
+test('subscriptions can be unsubscribed with the using keyword', async () => {
+  const { store, api } = storeRef
+
+  // Polyfill Symbol.dispose for this test
+  // Explicit any can be removed once this library is upgraded to TypeScript 5.2
+  const anySymbol = Symbol as any;
+  anySymbol.dispose = anySymbol.dispose ?? Symbol.for('Symbol.dispose');
+
+  // Helper to count the number of subscriptions for the increment endpoint
+  const getSubscriptionCount = () =>
+    Object.keys(
+      store.getState().api.subscriptions['increment(undefined)'] ?? {}
+    ).length
+
+  const initialSubscriptionCount = getSubscriptionCount()
+  const subscription = store.dispatch(api.endpoints.increment.initiate())
+  await subscription;
+
+  // Simulate the dispose call made by the using keyword
+  (subscription as any)[anySymbol.dispose]();
+
+  // Wait for the unsubscribe to be processed as it's done asynchronously
+  await new Promise((resolve) => setTimeout(resolve, 0))
+
+  expect(getSubscriptionCount()).toBe(initialSubscriptionCount)
+})

packages/toolkit/src/query/utils/disposableIfAvailable.ts

@@ -0,0 +1,46 @@
+/**
+ * DisposableIfAvailable
+ * 
+ * The explicit resource management TC39 proposal introduces a new symbol
+ * `symbol.dispose` that can be used to mark objects as disposable.
+ * 
+ * It's important that this type can be compiled in older TypeScript versions
+ * that don't include the symbol so consumers of the library don't face type
+ * errors.
+ * 
+ * At compile time DisposableIfAvailable will be either:
+ * 
+ * - `{ [Symbol.dispose]: () => void }` if the symbol is defined
+ * - `{}` if the symbol is not defined
+ */
+type DisposeSymbolType = SymbolConstructor extends { dispose: symbol }
+  ? typeof Symbol['dispose']
+  : 'unused-literal'
+
+// Explicit any can be removed once this library is upgraded to TypeScript 5.2
+const disposeSymbol: DisposeSymbolType = (Symbol as any)['dispose']
+export type DisposableIfAvailable = SymbolConstructor extends {
+  dispose: symbol
+}
+  ? { [disposeSymbol]: () => void }
+  : {}
+
+/**
+ * At runtime check for the existence of Symbol.dispose and if available return
+ * an object with a dispose method.
+ * 
+ * This needs to be a runtime check as the symbol is not available in every environment.
+ */
+export const disposableIfAvailable = (
+  disposeFn: () => void
+): DisposableIfAvailable => {
+
+  // Explicit any can be removed once we upgrade to TypeScript 5.2
+  const dispose: DisposeSymbolType = (Symbol as any)['dispose']
+  if (dispose) {
+    return {
+      [dispose]: () => disposeFn(),
+    }
+  }
+  return {} as DisposableIfAvailable
+}