Core: Batch the loading of CSF files for extract() etc

code/lib/preview-api/src/modules/store/StoryStore.test.ts

@@ -22,6 +22,14 @@ jest.mock('global', () => ({
   },
 }));
 
+const createGate = (): [Promise<any | undefined>, (_?: any) => void] => {
+  let openGate = (_?: any) => {};
+  const gate = new Promise<any | undefined>((resolve) => {
+    openGate = resolve;
+  });
+  return [gate, openGate];
+};
+
 const componentOneExports = {
   default: { title: 'Component One' },
   a: { args: { foo: 'a' } },
@@ -465,6 +473,24 @@ describe('StoryStore', () => {
         './src/ComponentTwo.stories.js',
       ]);
     });
+
+    it('imports in batches', async () => {
+      const [gate, openGate] = createGate();
+      const blockedImportFn = jest.fn(async (file) => {
+        await gate;
+        return importFn(file);
+      });
+      const store = new StoryStore();
+      store.setProjectAnnotations(projectAnnotations);
+      store.initialize({ storyIndex, importFn: blockedImportFn, cache: false });
+
+      const promise = store.loadAllCSFFiles({ batchSize: 1 });
+      expect(blockedImportFn).toHaveBeenCalledTimes(1);
+
+      openGate();
+      await promise;
+      expect(blockedImportFn).toHaveBeenCalledTimes(3);
+    });
   });
 
   describe('extract', () => {

code/lib/preview-api/src/modules/store/StoryStore.ts

@@ -31,9 +31,9 @@ import { ArgsStore } from './ArgsStore';
 import { GlobalsStore } from './GlobalsStore';
 import { processCSFFile, prepareStory, normalizeProjectAnnotations } from './csf';
 
-// TODO -- what are reasonable values for these?
 const CSF_CACHE_SIZE = 1000;
 const STORY_CACHE_SIZE = 10000;
+const EXTRACT_BATCH_SIZE = 20;
 
 export class StoryStore<TRenderer extends Renderer> {
   storyIndex?: StoryIndexStore;
@@ -143,22 +143,38 @@ export class StoryStore<TRenderer extends Renderer> {
     );
   }
 
-  loadAllCSFFiles(): Promise<StoryStore<TRenderer>['cachedCSFFiles']> {
+  loadAllCSFFiles({ batchSize = EXTRACT_BATCH_SIZE } = {}): Promise<
+    StoryStore<TRenderer>['cachedCSFFiles']
+  > {
     if (!this.storyIndex) throw new Error(`loadAllCSFFiles called before initialization`);
 
-    const importPaths: Record<Path, StoryId> = {};
-    Object.entries(this.storyIndex.entries).forEach(([storyId, { importPath }]) => {
-      importPaths[importPath] = storyId;
-    });
+    const importPaths = Object.entries(this.storyIndex.entries).map(([storyId, { importPath }]) => [
+      importPath,
+      storyId,
+    ]);
+
+    const loadInBatches = (
+      remainingImportPaths: typeof importPaths
+    ): Promise<{ importPath: Path; csfFile: CSFFile<TRenderer> }[]> => {
+      if (remainingImportPaths.length === 0) return SynchronousPromise.resolve([]);
+
+      const csfFilePromiseList = remainingImportPaths
+        .slice(0, batchSize)
+        .map(([importPath, storyId]) =>
+          this.loadCSFFileByStoryId(storyId).then((csfFile) => ({
+            importPath,
+            csfFile,
+          }))
+        );
 
-    const csfFilePromiseList = Object.entries(importPaths).map(([importPath, storyId]) =>
-      this.loadCSFFileByStoryId(storyId).then((csfFile) => ({
-        importPath,
-        csfFile,
-      }))
-    );
+      return SynchronousPromise.all(csfFilePromiseList).then((firstResults) =>
+        loadInBatches(remainingImportPaths.slice(batchSize)).then((restResults) =>
+          firstResults.concat(restResults)
+        )
+      );
+    };
 
-    return SynchronousPromise.all(csfFilePromiseList).then((list) =>
+    return loadInBatches(importPaths).then((list) =>
       list.reduce((acc, { importPath, csfFile }) => {
         acc[importPath] = csfFile;
         return acc;