perf: collect performance improvements

examples/vue3/src/components/HugeGrid.story.vue

@@ -1,5 +1,6 @@
 <script lang="ts" setup>
 import BaseButton from './BaseButton.vue'
+import { count } from './huge-grid'
 </script>
 
 <template>
@@ -11,7 +12,7 @@ import BaseButton from './BaseButton.vue'
     }"
   >
     <Variant
-      v-for="n in 1000"
+      v-for="n in count"
       :key="n"
       :title="`Variant ${n}`"
     >

examples/vue3/src/components/huge-grid.js

@@ -0,0 +1 @@
+export const count = 1000

packages/histoire/src/node/collect/index.ts

@@ -11,14 +11,16 @@ import type { ServerStoryFile } from '@histoire/shared'
 import { createPath } from '../tree.js'
 import type { Context } from '../context.js'
 import type { Payload, ReturnData } from './worker.js'
+import { slash } from '../util/fs.js'
 
 export interface UseCollectStoriesOptions {
   server: ViteDevServer
+  mainServer?: ViteDevServer
   throws?: boolean
 }
 
 export function useCollectStories (options: UseCollectStoriesOptions, ctx: Context) {
-  const { server } = options
+  const { server, mainServer } = options
 
   const node = new ViteNodeServer(server, {
     deps: {
@@ -35,8 +37,9 @@ export function useCollectStories (options: UseCollectStoriesOptions, ctx: Conte
   })
 
   const threadsCount = ctx.mode === 'dev'
-    ? Math.max(cpus().length / 2, 1)
+    ? Math.max(Math.min(cpus().length / 2, 4), 1)
     : Math.max(cpus().length - 1, 1)
+  console.log(pc.blue(`Using ${threadsCount} threads for story collection`))
 
   const threadPool = new Tinypool({
     filename: new URL('./worker.js', import.meta.url).href,
@@ -73,15 +76,31 @@ export function useCollectStories (options: UseCollectStoriesOptions, ctx: Conte
     }
   }
 
+  const invalidates = new Set<string>()
+
+  if (mainServer) {
+    mainServer.watcher.on('change', (file) => {
+      file = slash(file)
+      if (invalidates.has(file)) return
+      invalidates.add(file)
+    })
+  }
+
+  function clearInvalidates () {
+    invalidates.clear()
+  }
+
   async function executeStoryFile (storyFile: ServerStoryFile) {
     try {
       const { workerPort } = createChannel()
-      const { storyData } = await threadPool.run({
+      const payload: Payload = {
         root: server.config.root,
         base: server.config.base,
         storyFile,
         port: workerPort,
-      } as Payload, {
+        invalidates: Array.from(invalidates),
+      }
+      const { storyData } = await threadPool.run(payload, {
         transferList: [
           workerPort,
         ],
@@ -132,5 +151,6 @@ export function useCollectStories (options: UseCollectStoriesOptions, ctx: Conte
     clearCache,
     executeStoryFile,
     destroy,
+    clearInvalidates,
   }
 }

packages/histoire/src/node/collect/worker.ts

@@ -1,9 +1,11 @@
 import { fileURLToPath } from 'node:url'
 import type { MessagePort } from 'node:worker_threads'
-import { ViteNodeRunner } from 'vite-node/client'
+import { performance } from 'node:perf_hooks'
+import { ModuleCacheMap, ViteNodeRunner } from 'vite-node/client'
 import { createBirpc } from 'birpc'
 import type { FetchFunction, ResolveIdFunction } from 'vite-node'
 import { dirname, resolve } from 'pathe'
+import pc from 'picocolors'
 import type { ServerStoryFile, ServerStory, ServerRunPayload } from '@histoire/shared'
 import { createDomEnv } from '../dom/env.js'
 
@@ -14,33 +16,48 @@ export interface Payload {
   base: string
   port: MessagePort
   storyFile: ServerStoryFile
+  invalidates: string[]
 }
 
 export interface ReturnData {
   storyData: ServerStory[]
 }
 
+const _moduleCache = new ModuleCacheMap()
+let _runner: ViteNodeRunner
+let _rpc: ReturnType<typeof createBirpc<{
+  fetchModule: FetchFunction
+  resolveId: ResolveIdFunction
+}>>
+
 export default async (payload: Payload): Promise<ReturnData> => {
+  const startTime = performance.now()
   process.env.HST_COLLECT = 'true'
 
-  const rpc = createBirpc<{
+  _rpc = createBirpc<{
     fetchModule: FetchFunction
     resolveId: ResolveIdFunction
   }>({}, {
     post: data => payload.port.postMessage(data),
     on: data => payload.port.on('message', data),
   })
 
-  const runner = new ViteNodeRunner({
+  const runner = _runner ?? (_runner = new ViteNodeRunner({
     root: payload.root,
     base: payload.base,
+    moduleCache: _moduleCache,
     fetchModule (id) {
-      return rpc.fetchModule(id)
+      return _rpc.fetchModule(id)
     },
     resolveId (id, importer) {
-      return rpc.resolveId(id, importer)
+      return _rpc.resolveId(id, importer)
     },
-  })
+  }))
+
+  // Cleanup cache
+  for (const file of payload.invalidates) {
+    _moduleCache.delete(file)
+  }
 
   const { destroy: destroyDomEnv } = createDomEnv()
   if (!global.CSS.supports) {
@@ -49,14 +66,17 @@ export default async (payload: Payload): Promise<ReturnData> => {
 
   const el = window.document.createElement('div')
 
+  const beforeExectureTime = performance.now()
   // Mount app to collect stories/variants
   const { run } = (await runner.executeFile(resolve(__dirname, './run.js'))) as { run: (payload: ServerRunPayload) => Promise<any> }
+  const afterExectureTime = performance.now()
   const storyData: ServerStory[] = []
   await run({
     file: payload.storyFile,
     storyData,
     el,
   })
+  const afterRunTime = performance.now()
 
   if (payload.storyFile.markdownFile) {
     const el = document.createElement('div')
@@ -69,6 +89,9 @@ export default async (payload: Payload): Promise<ReturnData> => {
 
   destroyDomEnv()
 
+  const endTime = performance.now()
+  console.log(pc.dim(`${payload.storyFile.relativePath} ${Math.round(endTime - startTime)}ms (setup:${Math.round(beforeExectureTime - startTime)}ms, execute:${Math.round(afterExectureTime - beforeExectureTime)}ms, run:${Math.round(afterRunTime - afterExectureTime)}ms)`))
+
   return {
     storyData,
   }

packages/histoire/src/node/server.ts

@@ -1,4 +1,6 @@
+import { performance } from 'node:perf_hooks'
 import { createServer as createViteServer } from 'vite'
+import pc from 'picocolors'
 import type { ServerStoryFile } from '@histoire/shared'
 import { Context } from './context.js'
 import { getViteConfigWithPlugins, RESOLVED_MARKDOWN_FILES, RESOLVED_SEARCH_TITLE_DATA_ID, RESOLVED_STORIES_ID } from './vite.js'
@@ -46,16 +48,17 @@ export async function createServer (ctx: Context, options: CreateServerOptions =
     clearCache,
     executeStoryFile,
     destroy: destroyCollectStories,
+    clearInvalidates,
   } = useCollectStories({
     server: nodeServer,
+    mainServer: server,
   }, ctx)
 
   // onStoryChange debouncing
-  let queue: Promise<void>
-  let queueResolve: () => void
   let queued = false
   let queuedFiles: ServerStoryFile[] = []
-  let queuedAll = false
+  let queueTimer
+  let collecting = false
   let didAllStoriesYet = false
 
   // Invalidate modules
@@ -86,36 +89,38 @@ export async function createServer (ctx: Context, options: CreateServerOptions =
     if (changedFile && !didAllStoriesYet) {
       return
     }
-    if (queue) {
-      if (queued) {
-        if (changedFile) {
-          if (!queuedFiles.includes(changedFile)) {
-            queuedFiles.push(changedFile)
-          }
-        } else {
-          queuedAll = true
-        }
-        return
-      } else if (queuedFiles.includes(changedFile)) {
-        return
-      } else {
-        queued = true
-        // Next call
-        await queue
+
+    if (changedFile) {
+      if (!queuedFiles.includes(changedFile)) {
+        queuedFiles.push(changedFile)
       }
+    } else {
+      queuedFiles = []
     }
-    queuedFiles.unshift(changedFile)
-    queued = false
-    queue = new Promise(resolve => {
-      queueResolve = resolve
-    })
+
+    if (!queued) {
+      queued = true
+      if (!collecting) {
+        clearTimeout(queueTimer)
+        // Debounce
+        queueTimer = setTimeout(collect, 100)
+      }
+    }
+  })
+
+  async function collect () {
+    collecting = true
 
     clearCache()
 
-    console.log('Collect stories start', changedFile?.path ?? 'all')
-    const time = Date.now()
-    if (changedFile && !queuedAll) {
-      await Promise.all(queuedFiles.map(async storyFile => {
+    const currentFiles = queuedFiles.slice()
+    queuedFiles = []
+    queued = false
+
+    console.log('Collect stories start', currentFiles.length ? currentFiles.map(f => f.fileName).join(', ') : 'all')
+    const time = performance.now()
+    if (currentFiles.length) {
+      await Promise.all(currentFiles.map(async storyFile => {
         await executeStoryFile(storyFile)
         if (storyFile.story) {
           await invalidateModule(`/__resolved__virtual:story-source:${storyFile.story.id}`)
@@ -145,15 +150,19 @@ export async function createServer (ctx: Context, options: CreateServerOptions =
       didAllStoriesYet = true
       server.ws.send('histoire:all-stories-loaded', {})
     }
-    console.log('Collect stories end', Date.now() - time, 'ms')
+    console.log(`Collect stories end ${pc.bold(pc.blue(Math.round(performance.now() - time)))}ms`)
 
-    queuedFiles = []
-    queuedAll = false
-    queueResolve()
+    clearInvalidates()
 
     invalidateModule(RESOLVED_STORIES_ID)
     invalidateModule(RESOLVED_SEARCH_TITLE_DATA_ID)
-  })
+
+    collecting = false
+
+    if (queued) {
+      await collect()
+    }
+  }
 
   onStoryListChange(() => {
     invalidateModule(RESOLVED_STORIES_ID)

packages/histoire/src/node/util/fs.ts

@@ -0,0 +1 @@
+export const slash = (str: string) => str.replace(/\\/g, '/')