From d22217277a2e7099dc5f67b4a73f2822e111d6bf Mon Sep 17 00:00:00 2001 From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com> Date: Sun, 19 Jan 2025 20:10:44 +0000 Subject: [PATCH 1/5] Add `SubscriptionManager`, dynamic `enabled` support, optimize functions --- src/state/internal/createQueryStore.ts | 352 +++++++++++------- .../queryStore/classes/SubscriptionManager.ts | 109 ++++++ 2 files changed, 327 insertions(+), 134 deletions(-) create mode 100644 src/state/internal/queryStore/classes/SubscriptionManager.ts diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts index 19dbb811f65..2c524f5d1a5 100644 --- a/src/state/internal/createQueryStore.ts +++ b/src/state/internal/createQueryStore.ts @@ -4,6 +4,7 @@ import { subscribeWithSelector } from 'zustand/middleware'; import { IS_DEV } from '@/env'; import { RainbowError, logger } from '@/logger'; import { time } from '@/utils'; +import { SubscriptionManager } from './queryStore/classes/SubscriptionManager'; import { RainbowPersistConfig, createRainbowStore, omitStoreMethods } from './createRainbowStore'; import { $, AttachValue, SignalFunction, Unsubscribe, attachValueSubscriptionMap } from './signal'; @@ -129,7 +130,7 @@ export interface QueryStore< * status can be obtained by directly reading the `status` property. * @example * ```ts - * const isInitialLoad = useQueryStore(state => state.getStatus().isInitialLoad); + * const isInitialLoading = useMyQueryStore(state => state.getStatus().isInitialLoading); * ``` * @returns An object containing boolean flags for each status. */ @@ -155,10 +156,9 @@ export interface QueryStore< /** * The private state managed by the query store, omitted from the store's public interface. + * This is currently a placeholder type. */ -type PrivateStoreState = { - subscriptionCount: number; -}; +type PrivateStoreState = Record; /** * The full state structure managed by the query store. This type is generally internal, @@ -179,17 +179,23 @@ type StoreState> = Pick< */ export type QueryStoreConfig, TData, S extends StoreState> = { /** - * A function responsible for fetching data from a remote source. - * Receives parameters of type `TParams` and returns either a promise or a raw data value of type `TQueryFnData`. + * **A function responsible for fetching data from a remote source.** + * Receives parameters of type TParams and optional abort/cancel controls. + * Returns either a promise or a raw data value of type TQueryFnData. + * + * --- + * `abortController` and `cancel` are by default available, unless either: + * - `abortInterruptedFetches` is set to `false` in the store's config + * - The fetch was manually triggered with `skipStoreUpdates: true` */ - fetcher: (params: TParams, abortController: AbortController | null) => TQueryFnData | Promise; + fetcher: (params: TParams, abortController: AbortController | null, cancel: (() => void) | null) => TQueryFnData | Promise; /** - * A callback invoked whenever a fetch operation fails. + * **A callback invoked whenever a fetch operation fails.** * Receives the error and the current retry count. */ onError?: (error: Error, retryCount: number) => void; /** - * A callback invoked whenever fresh data is successfully fetched. + * **A callback invoked whenever fresh data is successfully fetched.** * Receives the transformed data and the store's set function, which can optionally be used to update store state. */ onFetched?: (info: { @@ -199,7 +205,7 @@ export type QueryStoreConfig | ((state: S) => S | Partial)) => void; }) => void; /** - * A function that overrides the default behavior of setting the fetched data in the store's query cache. + * **A function that overrides the default behavior of setting the fetched data in the store's query cache.** * Receives an object containing the transformed data, the query parameters, the query key, and the store's set function. * * When using `setData`, it’s important to note that you are taking full responsibility for managing query data. If your @@ -218,7 +224,7 @@ export type QueryStoreConfig | ((state: S) => S | Partial)) => void; }) => void; /** - * A function to transform the raw fetched data (`TQueryFnData`) into another form (`TData`). + * **A function to transform the raw fetched data** (`TQueryFnData`) into another form (`TData`). * If not provided, the raw data returned by `fetcher` is used. */ transform?: (data: TQueryFnData, params: TParams) => TData; @@ -258,9 +264,9 @@ export type QueryStoreConfig; /** - * If `true`, the store's `getData` method will always return existing data from the cache if it exists, + * When `true`, the store's `getData` method will always return existing data from the cache if it exists, * regardless of whether the cached data is expired, until the data is pruned following a successful fetch. * * Additionally, when params change while the store is enabled, `getData` will return the previous data until @@ -311,15 +317,30 @@ type ParamResolvable, S extends Store | T | (($: SignalFunction, store: QueryStore) => AttachValue); -interface ResolvedParamsResult { +interface ResolvedEnabledResult { /** - * Direct, non-reactive values resolved from the initial configuration. + * The reactive enabled state, if provided as a function returning an AttachValue. */ - directValues: Partial; + enabledAttachVal: AttachValue | null; + /** + * The static enabled state, if provided as a direct boolean value. + */ + enabledDirectValue: boolean | null; + /** + * The final enabled state, derived from either the reactive or static value. + */ + resolvedEnabled: boolean; +} + +interface ResolvedParamsResult { /** * Reactive parameter values wrapped in `AttachValue`, which trigger refetches when they change. */ - paramAttachVals: Partial>>; + attachVals: Partial>>; + /** + * Direct, non-reactive values resolved from the initial configuration. + */ + directValues: Partial; /** * Fully resolved parameters, merging both direct and reactive values. */ @@ -349,7 +370,6 @@ const SHOULD_PERSIST_INTERNAL_STATE_MAP: Record = { isDataExpired: discard, isStale: discard, reset: discard, - subscriptionCount: discard, } satisfies Record; const ABORT_ERROR = new Error('[createQueryStore: AbortError] Fetch interrupted'); @@ -436,8 +456,8 @@ export function createQueryStore< ); } - let directValues: Partial | null = null; - let paramAttachVals: Partial>> | null = null; + let attachVals: { enabled: AttachValue | null; params: Partial>> } | null = null; + let directValues: { enabled: boolean | null; params: Partial } | null = null; let paramUnsubscribes: Unsubscribe[] = []; let fetchAfterParamCreation = false; @@ -446,29 +466,28 @@ export function createQueryStore< let activeRefetchTimeout: NodeJS.Timeout | null = null; let lastFetchKey: string | null = null; - const enableLogs = IS_DEV && debugMode; const cacheTimeIsFunction = typeof cacheTime === 'function'; + const enableLogs = IS_DEV && debugMode; const initialData = { - enabled, + enabled: typeof enabled === 'function' ? false : enabled, error: null, lastFetchedAt: null, queryCache: {}, queryKey: '', status: QueryStatuses.Idle, - subscriptionCount: 0, }; - const getQueryKey = (params: TParams): string => JSON.stringify(Object.values(params)); + const subscriptionManager = new SubscriptionManager({ + disableAutoRefetching, + initialEnabled: initialData.enabled, + }); - const getCurrentResolvedParams = () => { - const currentParams: Partial = directValues ? { ...directValues } : {}; - for (const k in paramAttachVals) { - const attachVal = paramAttachVals[k as keyof TParams]; - if (!attachVal) continue; - currentParams[k as keyof TParams] = attachVal.value as TParams[keyof TParams]; + const abortActiveFetch = () => { + if (activeAbortController) { + activeAbortController.abort(); + activeAbortController = null; } - return currentParams as TParams; }; const fetchWithAbortControl = async (params: TParams): Promise => { @@ -479,7 +498,7 @@ export function createQueryStore< return await new Promise((resolve, reject) => { abortController.signal.addEventListener('abort', () => reject(ABORT_ERROR), { once: true }); - Promise.resolve(fetcher(params, abortController)).then(resolve, reject); + Promise.resolve(fetcher(params, abortController, abortController.abort)).then(resolve, reject); }); } finally { if (activeAbortController === abortController) { @@ -488,14 +507,38 @@ export function createQueryStore< } }; - const abortActiveFetch = () => { - if (activeAbortController) { - activeAbortController.abort(); - activeAbortController = null; - } - }; - const createState: StateCreator = (set, get, api) => { + subscriptionManager.init({ + onSubscribe: (enabled, isFirstSubscription, shouldThrottle) => { + if (!directValues && !attachVals && params !== undefined) return (fetchAfterParamCreation = true); + if (!enabled) return; + + if (isFirstSubscription) { + const { fetch, isStale, queryKey: storeQueryKey } = get(); + const currentParams = getCurrentResolvedParams(attachVals, directValues); + const currentQueryKey = getQueryKey(currentParams); + + if (storeQueryKey !== currentQueryKey) set(state => ({ ...state, queryKey: currentQueryKey })); + + if (isStale()) { + fetch(currentParams); + } else { + scheduleNextFetch(currentParams, undefined); + } + } else if (disableAutoRefetching && !shouldThrottle) { + get().fetch(); + } + }, + + onLastUnsubscribe: () => { + if (activeRefetchTimeout) { + clearTimeout(activeRefetchTimeout); + activeRefetchTimeout = null; + } + if (abortInterruptedFetches) abortActiveFetch(); + }, + }); + const scheduleNextFetch = (params: TParams, options: FetchOptions | undefined) => { if (disableAutoRefetching) return; const effectiveStaleTime = options?.staleTime ?? staleTime; @@ -512,20 +555,20 @@ export function createQueryStore< const timeUntilRefetch = lastFetchedAt ? effectiveStaleTime - (Date.now() - lastFetchedAt) : effectiveStaleTime; activeRefetchTimeout = setTimeout(() => { - const { enabled, fetch, subscriptionCount } = get(); + const { enabled, subscriptionCount } = subscriptionManager.get(); if (enabled && subscriptionCount > 0) { - fetch(params, { force: true }); + get().fetch(params, { force: true }); } }, timeUntilRefetch); }; const baseMethods = { async fetch(params: TParams | undefined, options: FetchOptions | undefined) { - const { enabled, error, status } = get(); + if (!options?.force && !subscriptionManager.get().enabled) return null; - if (!options?.force && !enabled) return null; + const { error, status } = get(); - const effectiveParams = params ?? getCurrentResolvedParams(); + const effectiveParams = params ?? getCurrentResolvedParams(attachVals, directValues); const currentQueryKey = getQueryKey(effectiveParams); const isLoading = status === QueryStatuses.Loading; const skipStoreUpdates = !!options?.skipStoreUpdates; @@ -542,7 +585,6 @@ export function createQueryStore< lastFetchedAt: storeLastFetchedAt, queryCache: { [currentQueryKey]: cacheEntry }, queryKey: storeQueryKey, - subscriptionCount, } = get(); const { errorInfo, lastFetchedAt: cachedLastFetchedAt } = cacheEntry ?? {}; @@ -551,7 +593,7 @@ export function createQueryStore< const isStale = !lastFetchedAt || Date.now() - lastFetchedAt >= (options?.staleTime ?? staleTime); if (!isStale && (!errorInfo || errorRetriesExhausted)) { - if (!activeRefetchTimeout && subscriptionCount > 0 && staleTime !== 0 && staleTime !== Infinity) { + if (!activeRefetchTimeout && subscriptionManager.get().subscriptionCount > 0 && staleTime !== 0 && staleTime !== Infinity) { scheduleNextFetch(effectiveParams, options); } if (enableLogs) console.log('[πŸ’Ύ Returning Cached Data πŸ’Ύ] for params:', JSON.stringify(effectiveParams)); @@ -570,7 +612,7 @@ export function createQueryStore< if (enableLogs) console.log('[πŸ”„ Fetching πŸ”„] for params:', JSON.stringify(effectiveParams)); const rawResult = await (abortInterruptedFetches && !skipStoreUpdates ? fetchWithAbortControl(effectiveParams) - : fetcher(effectiveParams, null)); + : fetcher(effectiveParams, null, null)); const lastFetchedAt = Date.now(); if (enableLogs) console.log('[βœ… Fetch Successful βœ…] for params:', JSON.stringify(effectiveParams)); @@ -616,7 +658,7 @@ export function createQueryStore< }, }; } else if (setData) { - if (enableLogs) console.log('[πŸ’Ύ Setting Data πŸ’Ύ] for params:\n', JSON.stringify(effectiveParams)); + if (enableLogs) console.log('[πŸ’Ύ Setting Data πŸ’Ύ] for params:', JSON.stringify(effectiveParams)); setData({ data: transformedData, params: effectiveParams, @@ -679,13 +721,13 @@ export function createQueryStore< onError?.(typedError, currentRetryCount); if (currentRetryCount < maxRetries) { - if (get().subscriptionCount > 0) { + if (subscriptionManager.get().subscriptionCount > 0) { const errorRetryDelay = typeof retryDelay === 'function' ? retryDelay(currentRetryCount, typedError) : retryDelay; if (errorRetryDelay !== Infinity) { activeRefetchTimeout = setTimeout(() => { - const { enabled, fetch, subscriptionCount } = get(); + const { enabled, subscriptionCount } = subscriptionManager.get(); if (enabled && subscriptionCount > 0) { - fetch(params, { force: true }); + get().fetch(params, { force: true }); } }, errorRetryDelay); } @@ -799,7 +841,7 @@ export function createQueryStore< if (abortInterruptedFetches) abortActiveFetch(); activeFetch = null; lastFetchKey = null; - set(state => ({ ...state, ...initialData, queryKey: getQueryKey(getCurrentResolvedParams()) })); + set(state => ({ ...state, ...initialData, queryKey: getQueryKey(getCurrentResolvedParams(attachVals, directValues)) })); }, }; @@ -807,8 +849,9 @@ export function createQueryStore< const handleSetEnabled = (state: S, prevState: S) => { if (state.enabled !== prevState.enabled && lastHandledEnabled !== state.enabled) { lastHandledEnabled = state.enabled; + subscriptionManager.setEnabled(state.enabled); if (state.enabled) { - const currentParams = getCurrentResolvedParams(); + const currentParams = getCurrentResolvedParams(attachVals, directValues); const currentKey = state.queryKey; if (currentKey !== lastFetchKey || state.isStale()) { state.fetch(currentParams); @@ -822,67 +865,18 @@ export function createQueryStore< } }; - const handleSubscribe = () => { - if (!directValues && !paramAttachVals && params !== undefined) { - fetchAfterParamCreation = true; - return; - } - const { enabled, fetch, isStale, queryKey: storeQueryKey, subscriptionCount } = get(); - - if (!enabled || (subscriptionCount !== 1 && !disableAutoRefetching)) return; - - if (subscriptionCount === 1) { - const currentParams = getCurrentResolvedParams(); - const currentQueryKey = getQueryKey(currentParams); - - if (storeQueryKey !== currentQueryKey) set(state => ({ ...state, queryKey: currentQueryKey })); - - if (isStale()) { - fetch(currentParams); - } else { - scheduleNextFetch(currentParams, undefined); - } - } else if (disableAutoRefetching) { - fetch(); - } - }; - - const handleUnsubscribe = (unsubscribe: () => void) => { - return () => { - unsubscribe(); - set(state => { - const newCount = Math.max(state.subscriptionCount - 1, 0); - if (newCount === 0) { - if (activeRefetchTimeout) { - clearTimeout(activeRefetchTimeout); - activeRefetchTimeout = null; - } - if (abortInterruptedFetches) abortActiveFetch(); - } - return { ...state, subscriptionCount: newCount }; - }); - }; - }; - - const incrementSubscriptionCount = () => { - set(state => ({ - ...state, - subscriptionCount: state.subscriptionCount + 1, - })); - }; - const subscribeWithSelector = api.subscribe; - api.subscribe = (listener: (state: S, prevState: S) => void) => { - incrementSubscriptionCount(); + api.subscribe = (listener: (state: S, prevState: S) => void) => { + const internalUnsubscribe = subscriptionManager.subscribe(); const unsubscribe = subscribeWithSelector((state: S, prevState: S) => { listener(state, prevState); handleSetEnabled(state, prevState); }); - - handleSubscribe(); - - return handleUnsubscribe(unsubscribe); + return () => { + internalUnsubscribe(); + unsubscribe(); + }; }; const userState = customStateCreator?.(set, get, api) ?? ({} as U); @@ -904,10 +898,10 @@ export function createQueryStore< const baseStore = persistConfig?.storageKey ? createRainbowStore(createState, combinedPersistConfig) - : create(subscribeWithSelector(createState)); + : create()(subscribeWithSelector(createState)); - const queryCapableStore: QueryStore = Object.assign(baseStore, { - enabled, + const queryCapableStore = Object.assign(baseStore, { + enabled: baseStore.getState().enabled, queryKey: baseStore.getState().queryKey, fetch: (params?: TParams, options?: FetchOptions) => baseStore.getState().fetch(params, options), getData: () => baseStore.getState().getData(), @@ -918,18 +912,30 @@ export function createQueryStore< for (const unsub of paramUnsubscribes) unsub(); paramUnsubscribes = []; queryCapableStore.getState().reset(); - queryCapableStore.setState(state => ({ ...state, ...initialData, queryKey: getQueryKey(getCurrentResolvedParams()) })); + queryCapableStore.setState(state => ({ + ...state, + ...initialData, + queryKey: getQueryKey(getCurrentResolvedParams(attachVals, directValues)), + })); }, }); + const { error, queryKey } = queryCapableStore.getState(); + if (queryKey && !error) lastFetchKey = queryKey; + if (params) { - const result = resolveParams(params, queryCapableStore); - paramAttachVals = result.paramAttachVals; - directValues = result.directValues; + const { + directValues: resolvedDirectValues, + enabledAttachVal: resolvedEnabledAttachVal, + enabledDirectValue: resolvedEnabledDirectValue, + attachVals: resolvedAttachVals, + } = resolveParams(enabled, params, queryCapableStore); + attachVals = { enabled: resolvedEnabledAttachVal, params: resolvedAttachVals }; + directValues = { enabled: resolvedEnabledDirectValue, params: resolvedDirectValues }; } const onParamChange = () => { - const newParams = getCurrentResolvedParams(); + const newParams = getCurrentResolvedParams(attachVals, directValues); if (!keepPreviousData) { const newQueryKey = getQueryKey(newParams); queryCapableStore.setState(state => ({ ...state, queryKey: newQueryKey })); @@ -937,8 +943,32 @@ export function createQueryStore< queryCapableStore.fetch(newParams); }; - for (const k in paramAttachVals) { - const attachVal = paramAttachVals[k]; + if (attachVals?.enabled) { + const attachVal = attachVals.enabled; + const subscribeFn = attachValueSubscriptionMap.get(attachVal); + + if (subscribeFn) { + let oldVal = attachVal.value; + if (oldVal) { + queryCapableStore.setState(state => ({ ...state, enabled: true })); + subscriptionManager.setEnabled(true); + } + if (enableLogs) console.log('[πŸŒ€ Enabled Subscription πŸŒ€] Initial value:', oldVal); + + const unsub = subscribeFn(() => { + const newVal = attachVal.value; + if (newVal !== oldVal) { + if (enableLogs) console.log('[πŸŒ€ Enabled Change πŸŒ€] - [Old]:', `${oldVal},`, '[New]:', newVal); + oldVal = newVal; + queryCapableStore.setState(state => ({ ...state, enabled: newVal })); + } + }); + paramUnsubscribes.push(unsub); + } + } + + for (const k in attachVals?.params) { + const attachVal = attachVals.params[k]; if (!attachVal) continue; const subscribeFn = attachValueSubscriptionMap.get(attachVal); @@ -963,46 +993,100 @@ export function createQueryStore< return queryCapableStore; } +function getQueryKey>(params: TParams): string { + const values = new Array(Object.keys(params).length); + let i = 0; + // eslint-disable-next-line no-plusplus + for (const key in params) values[i++] = params[key]; + return JSON.stringify(values); +} + +function getCurrentResolvedParams>( + attachVals: { enabled: AttachValue | null; params: Partial>> } | null, + directValues: { enabled: boolean | null; params: Partial } | null +): TParams { + const currentParams: Partial = directValues?.params ?? {}; + for (const k in attachVals?.params) { + const attachVal = attachVals.params[k]; + if (!attachVal) continue; + currentParams[k as keyof TParams] = attachVal.value as TParams[keyof TParams]; + } + return currentParams as TParams; +} + function pruneCache, TData, TParams extends Record>( keepPreviousData: boolean, keyToPreserve: string | null, state: S | Partial ): S | Partial { if (!state.queryCache) return state; + const clonedCache = { ...state.queryCache }; const effectiveKeyToPreserve = keyToPreserve ?? (keepPreviousData ? state.queryKey ?? null : null); - const newCache: Record> = {}; const pruneTime = Date.now(); + let hasChanges = false; - Object.entries(state.queryCache).forEach(([key, entry]) => { - if (entry && (pruneTime - entry.lastFetchedAt <= entry.cacheTime || key === effectiveKeyToPreserve)) { - newCache[key] = entry; + for (const key in clonedCache) { + const entry = clonedCache[key]; + if (entry && !(pruneTime - entry.lastFetchedAt <= entry.cacheTime || key === effectiveKeyToPreserve)) { + delete clonedCache[key]; + hasChanges = true; } - }); + } + return hasChanges ? { ...state, queryCache: clonedCache } : state; +} + +function isResolvableParam, S extends StoreState, TData>( + param: ParamResolvable +): param is ($: SignalFunction, store: QueryStore) => AttachValue { + return typeof param === 'function'; +} - return { ...state, queryCache: newCache }; +type StaticParamValue, S extends StoreState, TData> = Exclude< + ParamResolvable, + ($: SignalFunction, store: QueryStore) => AttachValue +>; + +function isStaticParam, S extends StoreState, TData>( + param: ParamResolvable +): param is StaticParamValue { + return !isResolvableParam(param); } function resolveParams, S extends StoreState & U, TData, U = unknown>( + enabled: boolean | ParamResolvable, params: { [K in keyof TParams]: ParamResolvable }, store: QueryStore -): ResolvedParamsResult { +): ResolvedParamsResult & ResolvedEnabledResult { + const attachVals: Partial>> = {}; const directValues: Partial = {}; - const paramAttachVals: Partial>> = {}; - const resolvedParams = {} as TParams; + const resolvedParams: TParams = {} as TParams; for (const key in params) { const param = params[key]; - if (typeof param === 'function') { + if (isResolvableParam(param)) { const attachVal = param($, store); - resolvedParams[key] = attachVal.value as TParams[typeof key]; - paramAttachVals[key] = attachVal; - } else { - resolvedParams[key] = param as TParams[typeof key]; - directValues[key] = param as TParams[typeof key]; + attachVals[key] = attachVal; + resolvedParams[key] = attachVal.value; + } else if (isStaticParam(param)) { + directValues[key] = param; + resolvedParams[key] = param; } } - return { directValues, paramAttachVals, resolvedParams }; + let enabledAttachVal: AttachValue | null = null; + let enabledDirectValue: boolean | null = null; + let resolvedEnabled: boolean; + + if (isResolvableParam(enabled)) { + const attachVal = enabled($, store); + resolvedEnabled = attachVal.value; + enabledAttachVal = attachVal; + } else { + resolvedEnabled = enabled; + enabledDirectValue = enabled; + } + + return { attachVals, directValues, enabledAttachVal, enabledDirectValue, resolvedEnabled, resolvedParams }; } function createBlendedPartialize, S extends StoreState & U, U = unknown>( diff --git a/src/state/internal/queryStore/classes/SubscriptionManager.ts b/src/state/internal/queryStore/classes/SubscriptionManager.ts new file mode 100644 index 00000000000..b75bb70878e --- /dev/null +++ b/src/state/internal/queryStore/classes/SubscriptionManager.ts @@ -0,0 +1,109 @@ +import { time } from '@/utils'; + +/** + * Initialization config for the SubscriptionManager. + */ +interface SubscriptionManagerConfig { + /** + * The store's `disableAutoRefetching` option, passed through from the query config. + */ + disableAutoRefetching: boolean; + /** + * The store's initial `enabled` state, as determined in `initialData`. + */ + initialEnabled: boolean; +} + +/** + * Lazy initialization config for subscription handlers. + */ +interface SubscriptionHandlerConfig { + /** + * Callback executed when a subscription is added. + * @param enabled - The current enabled state + * @param isFirstSubscription - Whether this is the first subscription + * @param shouldThrottle - Whether to throttle the fetch + */ + onSubscribe: (enabled: boolean, isFirstSubscription: boolean, shouldThrottle: boolean) => void; + /** Callback executed when the last remaining subscription is removed. */ + onLastUnsubscribe: () => void; +} + +/** + * Manages subscription state and lifecycle events for a `createQueryStore` instance. + */ +export class SubscriptionManager { + private count = 0; + private enabled: boolean; + private lastSubscriptionTime: number | null = null; + private readonly fetchThrottleMs: number | null = null; + + private onSubscribe: SubscriptionHandlerConfig['onSubscribe'] | null = null; + private onLastUnsubscribe: SubscriptionHandlerConfig['onLastUnsubscribe'] | null = null; + + /** + * Creates a new SubscriptionManager instance. + */ + constructor({ disableAutoRefetching, initialEnabled }: SubscriptionManagerConfig) { + this.enabled = initialEnabled; + if (disableAutoRefetching) { + this.fetchThrottleMs = time.seconds(5); + } + } + + /** + * Initializes subscription event handlers. + */ + init({ onSubscribe, onLastUnsubscribe }: SubscriptionHandlerConfig): void { + this.onSubscribe = onSubscribe; + this.onLastUnsubscribe = onLastUnsubscribe; + } + + /** + * Returns the current subscription state. + * @returns An object containing `enabled`, `lastSubscriptionTime`, and `subscriptionCount` + */ + get(): { enabled: boolean; lastSubscriptionTime: number | null; subscriptionCount: number } { + return { + enabled: this.enabled, + lastSubscriptionTime: this.lastSubscriptionTime, + subscriptionCount: this.count, + }; + } + + /** + * Updates the enabled state for queries. + * @param enabled - The new enabled state + */ + setEnabled(enabled: boolean): void { + this.enabled = enabled; + } + + /** + * Adds a new subscription and triggers relevant lifecycle callbacks. + * @returns A cleanup function that removes the subscription when called + */ + subscribe(): () => void { + const isFirstSubscription = this.count === 0; + const shouldThrottle = + this.fetchThrottleMs !== null && + this.lastSubscriptionTime !== null && + !isFirstSubscription && + Date.now() - this.lastSubscriptionTime <= this.fetchThrottleMs; + + this.onSubscribe?.(this.enabled, isFirstSubscription, shouldThrottle); + + this.count += 1; + this.lastSubscriptionTime = Date.now(); + + return () => { + const isLastSubscription = this.count === 1; + this.count = Math.max(this.count - 1, 0); + + if (isLastSubscription) { + this.onLastUnsubscribe?.(); + this.lastSubscriptionTime = null; + } + }; + } +} From 211e617033d9e6b9fe9f28be3383c5fba4aa9b82 Mon Sep 17 00:00:00 2001 From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com> Date: Mon, 20 Jan 2025 03:45:20 +0000 Subject: [PATCH 2/5] abc --- src/state/internal/createQueryStore.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts index 2c524f5d1a5..26971edc3d2 100644 --- a/src/state/internal/createQueryStore.ts +++ b/src/state/internal/createQueryStore.ts @@ -4,8 +4,8 @@ import { subscribeWithSelector } from 'zustand/middleware'; import { IS_DEV } from '@/env'; import { RainbowError, logger } from '@/logger'; import { time } from '@/utils'; -import { SubscriptionManager } from './queryStore/classes/SubscriptionManager'; import { RainbowPersistConfig, createRainbowStore, omitStoreMethods } from './createRainbowStore'; +import { SubscriptionManager } from './queryStore/classes/SubscriptionManager'; import { $, AttachValue, SignalFunction, Unsubscribe, attachValueSubscriptionMap } from './signal'; /** From f627e5516881f45e79e3d880d102ce85034d7468 Mon Sep 17 00:00:00 2001 From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com> Date: Tue, 21 Jan 2025 07:01:58 +0000 Subject: [PATCH 3/5] `getQueryKey`: Ensure consistent params ordering --- src/state/internal/createQueryStore.ts | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts index 26971edc3d2..163b9d4ef6e 100644 --- a/src/state/internal/createQueryStore.ts +++ b/src/state/internal/createQueryStore.ts @@ -994,11 +994,11 @@ export function createQueryStore< } function getQueryKey>(params: TParams): string { - const values = new Array(Object.keys(params).length); - let i = 0; - // eslint-disable-next-line no-plusplus - for (const key in params) values[i++] = params[key]; - return JSON.stringify(values); + return JSON.stringify( + Object.keys(params) + .sort() + .map(key => params[key]) + ); } function getCurrentResolvedParams>( From 0b6ab1cc4d373dc0e41992995f6ab32af78930c9 Mon Sep 17 00:00:00 2001 From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com> Date: Tue, 21 Jan 2025 09:54:37 +0000 Subject: [PATCH 4/5] Improve `CacheEntry` type safety, make `pruneCache` faster --- src/state/internal/createQueryStore.ts | 66 ++++++++++++++++---------- 1 file changed, 42 insertions(+), 24 deletions(-) diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts index 163b9d4ef6e..d93e666d37b 100644 --- a/src/state/internal/createQueryStore.ts +++ b/src/state/internal/createQueryStore.ts @@ -74,16 +74,27 @@ interface FetchOptions { * Represents an entry in the query cache, which stores fetched data along with metadata, and error information * in the event the most recent fetch failed. */ -interface CacheEntry { +type CacheEntry = { cacheTime: number; data: TData | null; - errorInfo: { - error: Error; - lastFailedAt: number; - retryCount: number; - } | null; - lastFetchedAt: number; -} +} & ( + | { + errorInfo: { + error: Error; + lastFailedAt: number; + retryCount: number; + }; + lastFetchedAt: null; + } + | { + errorInfo: { + error: Error; + lastFailedAt: number; + retryCount: number; + } | null; + lastFetchedAt: number; + } +); /** * A specialized store interface that combines Zustand's store capabilities with remote data fetching support. @@ -655,7 +666,7 @@ export function createQueryStore< data: transformedData, errorInfo: null, lastFetchedAt, - }, + } satisfies CacheEntry, }; } else if (setData) { if (enableLogs) console.log('[πŸ’Ύ Setting Data πŸ’Ύ] for params:', JSON.stringify(effectiveParams)); @@ -675,7 +686,7 @@ export function createQueryStore< data: null, errorInfo: null, lastFetchedAt, - }, + } satisfies CacheEntry, }; } } @@ -739,13 +750,15 @@ export function createQueryStore< queryCache: { ...state.queryCache, [currentQueryKey]: { - ...entry, + cacheTime: entry?.cacheTime ?? (cacheTimeIsFunction ? cacheTime(effectiveParams) : cacheTime), + data: entry?.data ?? null, + lastFetchedAt: entry?.lastFetchedAt ?? null, errorInfo: { error: typedError, lastFailedAt: Date.now(), retryCount: currentRetryCount + 1, }, - }, + } satisfies CacheEntry, }, queryKey: keepPreviousData ? currentQueryKey : state.queryKey, status: QueryStatuses.Error, @@ -758,13 +771,15 @@ export function createQueryStore< queryCache: { ...state.queryCache, [currentQueryKey]: { - ...entry, + cacheTime: entry?.cacheTime ?? (cacheTimeIsFunction ? cacheTime(effectiveParams) : cacheTime), + data: entry?.data ?? null, + lastFetchedAt: entry?.lastFetchedAt ?? null, errorInfo: { error: typedError, lastFailedAt: Date.now(), retryCount: maxRetries, }, - }, + } satisfies CacheEntry, }, queryKey: keepPreviousData ? currentQueryKey : state.queryKey, status: QueryStatuses.Error, @@ -1020,19 +1035,22 @@ function pruneCache, TData, TParams extends state: S | Partial ): S | Partial { if (!state.queryCache) return state; - const clonedCache = { ...state.queryCache }; - const effectiveKeyToPreserve = keyToPreserve ?? (keepPreviousData ? state.queryKey ?? null : null); const pruneTime = Date.now(); - let hasChanges = false; + const preserve = keyToPreserve ?? ((keepPreviousData && state.queryKey) || null); - for (const key in clonedCache) { - const entry = clonedCache[key]; - if (entry && !(pruneTime - entry.lastFetchedAt <= entry.cacheTime || key === effectiveKeyToPreserve)) { - delete clonedCache[key]; - hasChanges = true; - } + let prunedSomething = false; + const newCache: Record> = Object.create(null); + + for (const key in state.queryCache) { + const entry = state.queryCache[key]; + const isValid = !!entry && (pruneTime - (entry.lastFetchedAt ?? entry.errorInfo.lastFailedAt) <= entry.cacheTime || key === preserve); + if (isValid) newCache[key] = entry; + else prunedSomething = true; } - return hasChanges ? { ...state, queryCache: clonedCache } : state; + + if (!prunedSomething) return state; + + return { ...state, queryCache: newCache }; } function isResolvableParam, S extends StoreState, TData>( From 452262b0d81b3c6cca0b67ddc81b40c6f4f752be Mon Sep 17 00:00:00 2001 From: Christian Baroni <7061887+christianbaroni@users.noreply.github.com> Date: Tue, 21 Jan 2025 18:06:57 +0000 Subject: [PATCH 5/5] Catch all AbortErrors --- src/state/internal/createQueryStore.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/state/internal/createQueryStore.ts b/src/state/internal/createQueryStore.ts index d93e666d37b..131e2e3995c 100644 --- a/src/state/internal/createQueryStore.ts +++ b/src/state/internal/createQueryStore.ts @@ -713,7 +713,7 @@ export function createQueryStore< return transformedData ?? null; } catch (error) { - if (error === ABORT_ERROR) { + if (error === ABORT_ERROR || (error instanceof Error && error.name === 'AbortError')) { if (enableLogs) console.log('[❌ Fetch Aborted ❌] for params:', JSON.stringify(effectiveParams)); return null; }