Auto batch enhancer

packages/toolkit/src/autoBatchEnhancer.ts

@@ -0,0 +1,68 @@
+import type { StoreEnhancer } from 'redux'
+
+export const shouldAutoBatch = 'RTK_autoBatch'
+
+// Copied from https://github.com/feross/queue-microtask
+let promise: Promise<any>
+const queueMicrotaskShim =
+  typeof queueMicrotask === 'function'
+    ? queueMicrotask.bind(typeof window !== 'undefined' ? window : global)
+    : // reuse resolved promise, and allocate it lazily
+      (cb: () => void) =>
+        (promise || (promise = Promise.resolve())).then(cb).catch((err: any) =>
+          setTimeout(() => {
+            throw err
+          }, 0)
+        )
+
+export const autoBatchEnhancer =
+  (): StoreEnhancer =>
+  (next) =>
+  (...args) => {
+    const store = next(...args)
+
+    let notifying = true
+    let notificationQueued = false
+    // let nextNotification: NodeJS.Timeout | undefined = undefined
+    const listeners = new Set<() => void>()
+    const notifyListeners = () => {
+      //nextNotification = void
+      notificationQueued = false
+      listeners.forEach((l) => l())
+    }
+
+    return Object.assign({}, store, {
+      subscribe(listener: () => void) {
+        const wrappedListener: typeof listener = () => notifying && listener()
+        const unsubscribe = store.subscribe(wrappedListener)
+        listeners.add(listener)
+        return () => {
+          unsubscribe()
+          listeners.delete(listener)
+        }
+      },
+      dispatch(action: any) {
+        try {
+          notifying = !action?.meta?.[shouldAutoBatch]
+          if (notifying) {
+            notificationQueued = false
+            // if (nextNotification) {
+            //   nextNotification = void clearTimeout(nextNotification)
+            // }
+          } else {
+            if (!notificationQueued) {
+              notificationQueued = true
+              queueMicrotaskShim(notifyListeners)
+            }
+            // nextNotification ||= setTimeout(
+            //   notifyListeners,
+            //   batchTimeout
+            // ) as any
+          }
+          return store.dispatch(action)
+        } finally {
+          notifying = true
+        }
+      },
+    })
+  }

packages/toolkit/src/configureStore.ts

@@ -34,7 +34,7 @@ const IS_PRODUCTION = process.env.NODE_ENV === 'production'
  * @public
  */
 export type ConfigureEnhancersCallback<E extends Enhancers = Enhancers> = (
-    defaultEnhancers: readonly StoreEnhancer[]
+  defaultEnhancers: readonly StoreEnhancer[]
 ) => [...E]
 
 /**
@@ -104,10 +104,10 @@ type Middlewares<S> = ReadonlyArray<Middleware<{}, S>>
 
 type Enhancers = ReadonlyArray<StoreEnhancer>
 
-interface ToolkitStore<
+export interface ToolkitStore<
   S = any,
   A extends Action = AnyAction,
-  M extends Middlewares<S> = Middlewares<S>,
+  M extends Middlewares<S> = Middlewares<S>
 > extends Store<S, A> {
   /**
    * The `dispatch` method of your store, enhanced by all its middlewares.

packages/toolkit/src/createAsyncThunk.ts

@@ -7,6 +7,7 @@ import { createAction } from './createAction'
 import type { ThunkDispatch } from 'redux-thunk'
 import type { FallbackIfUnknown, Id, IsAny, IsUnknown } from './tsHelpers'
 import { nanoid } from './nanoid'
+import { shouldAutoBatch } from './autoBatchEnhancer'
 
 // @ts-ignore we need the import of these types due to a bundling issue.
 type _Keep = PayloadAction | ActionCreatorWithPreparedPayload<any, unknown>
@@ -547,6 +548,10 @@ export const createAsyncThunk = (() => {
             requestStatus: 'rejected' as const,
             aborted: error?.name === 'AbortError',
             condition: error?.name === 'ConditionError',
+            // TODO: this is a hack to showcase the autobatching behaviour
+            // currently there is no way to add `meta` to a "condition rejected"
+            // action - that would need to be added
+            [shouldAutoBatch]: true,
           },
         })
       )

packages/toolkit/src/index.ts

@@ -17,6 +17,7 @@ export type {
 } from 'reselect'
 export { createDraftSafeSelector } from './createDraftSafeSelector'
 export type { ThunkAction, ThunkDispatch, ThunkMiddleware } from 'redux-thunk'
+export { shouldAutoBatch, autoBatchEnhancer } from './autoBatchEnhancer'
 
 // We deliberately enable Immer's ES5 support, on the grounds that
 // we assume RTK will be used with React Native and other Proxy-less

packages/toolkit/src/query/core/buildMiddleware/index.ts

@@ -79,9 +79,9 @@ export function buildMiddleware<
 
         const stateBefore = mwApi.getState()
 
-        if (!batchedActionsHandler(action, mwApi, stateBefore)) {
-          return
-        }
+        // if (!batchedActionsHandler(action, mwApi, stateBefore)) {
+        //   return
+        // }
 
         const res = next(action)
 

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

@@ -1,5 +1,6 @@
 import type { AnyAction, PayloadAction } from '@reduxjs/toolkit'
 import {
+  shouldAutoBatch,
   combineReducers,
   createAction,
   createSlice,
@@ -85,6 +86,12 @@ function updateMutationSubstateIfExists(
 }
 
 const initialState = {} as any
+const prepareAutoBatched =
+  <T>() =>
+  (payload: T): { payload: T; meta: unknown } => ({
+    payload,
+    meta: { [shouldAutoBatch]: true },
+  })
 
 export function buildSlice({
   reducerPath,
@@ -114,11 +121,14 @@ export function buildSlice({
     name: `${reducerPath}/queries`,
     initialState: initialState as QueryState<any>,
     reducers: {
-      removeQueryResult(
-        draft,
-        { payload: { queryCacheKey } }: PayloadAction<QuerySubstateIdentifier>
-      ) {
-        delete draft[queryCacheKey]
+      removeQueryResult: {
+        reducer(
+          draft,
+          { payload: { queryCacheKey } }: PayloadAction<QuerySubstateIdentifier>
+        ) {
+          delete draft[queryCacheKey]
+        },
+        prepare: prepareAutoBatched<QuerySubstateIdentifier>(),
       },
       queryResultPatched(
         draft,
@@ -243,14 +253,14 @@ export function buildSlice({
     name: `${reducerPath}/mutations`,
     initialState: initialState as MutationState<any>,
     reducers: {
-      removeMutationResult(
-        draft,
-        { payload }: PayloadAction<MutationSubstateIdentifier>
-      ) {
-        const cacheKey = getMutationCacheKey(payload)
-        if (cacheKey in draft) {
-          delete draft[cacheKey]
-        }
+      removeMutationResult: {
+        reducer(draft, { payload }: PayloadAction<MutationSubstateIdentifier>) {
+          const cacheKey = getMutationCacheKey(payload)
+          if (cacheKey in draft) {
+            delete draft[cacheKey]
+          }
+        },
+        prepare: prepareAutoBatched<MutationSubstateIdentifier>(),
       },
     },
     extraReducers(builder) {
@@ -369,35 +379,42 @@ export function buildSlice({
     },
   })
 
+  type SubscriptionUpdate = {
+    endpointName: string
+    requestId: string
+    options: Subscribers[number]
+  } & QuerySubstateIdentifier
   const subscriptionSlice = createSlice({
     name: `${reducerPath}/subscriptions`,
     initialState: initialState as SubscriptionState,
     reducers: {
-      updateSubscriptionOptions(
-        draft,
-        {
-          payload: { queryCacheKey, requestId, options },
-        }: PayloadAction<
+      updateSubscriptionOptions: {
+        reducer(
+          draft,
           {
-            endpointName: string
-            requestId: string
-            options: Subscribers[number]
-          } & QuerySubstateIdentifier
-        >
-      ) {
-        if (draft?.[queryCacheKey]?.[requestId]) {
-          draft[queryCacheKey]![requestId] = options
-        }
+            payload: { queryCacheKey, requestId, options },
+          }: PayloadAction<SubscriptionUpdate>
+        ) {
+          if (draft?.[queryCacheKey]?.[requestId]) {
+            draft[queryCacheKey]![requestId] = options
+          }
+        },
+        prepare: prepareAutoBatched<SubscriptionUpdate>(),
       },
-      unsubscribeQueryResult(
-        draft,
-        {
-          payload: { queryCacheKey, requestId },
-        }: PayloadAction<{ requestId: string } & QuerySubstateIdentifier>
-      ) {
-        if (draft[queryCacheKey]) {
-          delete draft[queryCacheKey]![requestId]
-        }
+      unsubscribeQueryResult: {
+        reducer(
+          draft,
+          {
+            payload: { queryCacheKey, requestId },
+          }: PayloadAction<{ requestId: string } & QuerySubstateIdentifier>
+        ) {
+          if (draft[queryCacheKey]) {
+            delete draft[queryCacheKey]![requestId]
+          }
+        },
+        prepare: prepareAutoBatched<
+          { requestId: string } & QuerySubstateIdentifier
+        >(),
       },
       subscriptionRequestsRejected(
         draft,
@@ -435,6 +452,18 @@ export function buildSlice({
               arg.subscriptionOptions ?? substate[requestId] ?? {}
           }
         })
+        // original case added back in as otherwise nothing would happen without the batching middleware
+        .addCase(
+          queryThunk.rejected,
+          (draft, { meta: { condition, arg, requestId }, error, payload }) => {
+            // request was aborted due to condition (another query already running)
+            if (condition && arg.subscribe) {
+              const substate = (draft[arg.queryCacheKey] ??= {})
+              substate[requestId] =
+                arg.subscriptionOptions ?? substate[requestId] ?? {}
+            }
+          }
+        )
         // update the state to be a new object to be picked up as a "state change"
         // by redux-persist's `autoMergeLevel2`
         .addMatcher(hasRehydrationInfo, (draft) => ({ ...draft }))

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

@@ -38,7 +38,7 @@ import type {
   ThunkDispatch,
   AsyncThunk,
 } from '@reduxjs/toolkit'
-import { createAsyncThunk } from '@reduxjs/toolkit'
+import { createAsyncThunk, shouldAutoBatch } from '@reduxjs/toolkit'
 
 import { HandledError } from '../HandledError'
 
@@ -122,13 +122,18 @@ export interface MutationThunkArg {
 export type ThunkResult = unknown
 
 export type ThunkApiMetaConfig = {
-  pendingMeta: { startedTimeStamp: number }
+  pendingMeta: {
+    startedTimeStamp: number
+    [shouldAutoBatch]: true
+  }
   fulfilledMeta: {
     fulfilledTimeStamp: number
     baseQueryMeta: unknown
+    [shouldAutoBatch]: true
   }
   rejectedMeta: {
     baseQueryMeta: unknown
+    [shouldAutoBatch]: true
   }
 }
 export type QueryThunk = AsyncThunk<
@@ -398,6 +403,7 @@ export function buildThunks<
         {
           fulfilledTimeStamp: Date.now(),
           baseQueryMeta: result.meta,
+          [shouldAutoBatch]: true,
         }
       )
     } catch (error) {
@@ -422,7 +428,7 @@ export function buildThunks<
               catchedError.meta,
               arg.originalArgs
             ),
-            { baseQueryMeta: catchedError.meta }
+            { baseQueryMeta: catchedError.meta, [shouldAutoBatch]: true }
           )
         } catch (e) {
           catchedError = e
@@ -472,7 +478,7 @@ In the case of an unhandled error, no tags will be "provided" or "invalidated".`
     ThunkApiMetaConfig & { state: RootState<any, string, ReducerPath> }
   >(`${reducerPath}/executeQuery`, executeEndpoint, {
     getPendingMeta() {
-      return { startedTimeStamp: Date.now() }
+      return { startedTimeStamp: Date.now(), [shouldAutoBatch]: true }
     },
     condition(arg, { getState }) {
       const state = getState()
@@ -506,7 +512,7 @@ In the case of an unhandled error, no tags will be "provided" or "invalidated".`
     ThunkApiMetaConfig & { state: RootState<any, string, ReducerPath> }
   >(`${reducerPath}/executeMutation`, executeEndpoint, {
     getPendingMeta() {
-      return { startedTimeStamp: Date.now() }
+      return { startedTimeStamp: Date.now(), [shouldAutoBatch]: true }
     },
   })
 

packages/toolkit/src/query/react/module.ts

@@ -150,6 +150,7 @@ export const reactHooksModule = ({
       context,
     })
     safeAssign(anyApi, { usePrefetch })
+    // even with React batching completely out of the picture, we should get similar results now
     safeAssign(context, { batch })
 
     return {