waitForStore

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";
 

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<T>(
     clearTimeout(timeout);
   });
 }
+
+export async function waitForStore<T extends Store>(
+  filter: string,
+  options?: { timeout?: number },
+): Promise<T>;
+
+/**
+ * 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<T extends Store>(
+  name: string,
+  options: { timeout?: number } = {},
+): Promise<T> {
+  const { flux } = await import("@common");
+  const existingStores = flux.Store.getAll() as Store[] & {
+    listeners?: Set<Array<string | ((store: Store) => 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<T>((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<typeof setTimeout>;
+
+  // Promise that rejects if the module takes too long to appear
+  const timeoutPromise = new Promise<never>((_, 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);
+  });
+}