diff --git a/src/renderer/modules/webpack/index.ts b/src/renderer/modules/webpack/index.ts index 719e010bb..e000bd7c1 100644 --- a/src/renderer/modules/webpack/index.ts +++ b/src/renderer/modules/webpack/index.ts @@ -1,4 +1,4 @@ -export { waitForModule } from "./lazy"; +export { waitForModule, waitForStore } from "./lazy"; export { getFunctionBySource, getFunctionKeyBySource } from "./inner-search"; diff --git a/src/renderer/modules/webpack/lazy.ts b/src/renderer/modules/webpack/lazy.ts index c2427846e..59031a8a8 100644 --- a/src/renderer/modules/webpack/lazy.ts +++ b/src/renderer/modules/webpack/lazy.ts @@ -2,6 +2,8 @@ import type { Filter, LazyCallback, LazyListener, RawModule, WaitForOptions } fr import { getExports, getModule } from "./get-modules"; +import type { Store } from "@common/flux"; + /** * Listeners that will be checked when each module is initialized */ @@ -116,3 +118,83 @@ export async function waitForModule( clearTimeout(timeout); }); } + +export async function waitForStore( + filter: string, + options?: { timeout?: number }, +): Promise; + +/** + * Wait for a module that matches the given filter + * @param filter Store Name + * @param options Options + * @param options.timeout Timeout in milliseconds + * + * @see {@link filters} + * + * @remarks + * Some modules may not be available immediately when Discord starts and will take up to a few seconds. + * This is useful to ensure that the module is available before using it. + */ +export async function waitForStore( + name: string, + options: { timeout?: number } = {}, +): Promise { + const { flux } = await import("@common"); + const existingStores = flux.Store.getAll() as Store[] & { + listeners?: Set void)>>; + onPush?: (name: string, callback: (store: Store) => void) => () => void; + }; + if (!existingStores.listeners) { + existingStores.listeners = new Set<[string, (store: Store) => void]>(); + existingStores.onPush = (name, callback) => { + const store = existingStores.find((store) => store.getName() === name); + if (store) { + callback(store); + } + const listener = [name, callback]; + existingStores.listeners!.add(listener); + return () => existingStores.listeners!.delete(listener); + }; + existingStores.push = (...stores: Store[]) => { + for (const [name, callback] of existingStores.listeners!) { + const store = stores.find((store) => store.getName() === name); + if (store) { + (callback as (store: Store) => void)(store); + } + } + return Array.prototype.push.call(existingStores, ...stores); + }; + } + const existingStore = existingStores.find((store) => store.getName() === name) as T | undefined; + + if (existingStore) { + return existingStore; + } + + // Promise that resolves with the module + const promise = new Promise((resolve) => { + const unregister = existingStores.onPush!(name, (store: Store) => { + unregister(); + resolve(store as T); + }); + }); + + // If no timeout, then wait for as long as it takes + if (!options.timeout) return promise; + + // Different in Node and browser environments--number in browser, NodeJS.Timeout in Node + let timeout: ReturnType; + + // Promise that rejects if the module takes too long to appear + const timeoutPromise = new Promise((_, reject) => { + timeout = setTimeout(() => { + reject(new Error(`waitForStore timed out after ${options.timeout}ms`)); + }, options.timeout); + }); + + // Go with whichever happens first + return Promise.race([promise, timeoutPromise]).finally(() => { + clearTimeout(timeout); + }); +}