feat: restart profiler

vitejs · Dec 7, 2022 · 3647c78 · 3647c78
1 parent 023f47d
commit 3647c78
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 26 deletions.
diff --git a/packages/vite/src/node/cli.ts b/packages/vite/src/node/cli.ts
@@ -9,6 +9,7 @@ import type { ServerOptions } from './server'
 import type { LogLevel } from './logger'
 import { createLogger } from './logger'
 import { VERSION } from './constants'
+import { bindShortcuts } from './shortcuts'
 import { resolveConfig } from '.'
 
 const cli = cac('vite')
@@ -32,25 +33,33 @@ interface GlobalCLIOptions {
 }
 
 // @ts-ignore
-const profileSession: Session | undefined = global.__vite_profile_session
+let profileSession: Session | undefined = global.__vite_profile_session
+let profileCount = 0
 
-export const stopProfiler = (log: (message: string) => void): void => {
-  if (profileSession) {
-    profileSession.post('Profiler.stop', (err: any, { profile }: any) => {
+export const stopProfiler = (
+  log: (message: string) => void,
+): void | Promise<void> => {
+  if (!profileSession) return
+  return new Promise((res, rej) => {
+    profileSession!.post('Profiler.stop', (err: any, { profile }: any) => {
       // Write profile to disk, upload, etc.
       if (!err) {
-        const outPath = path.resolve('./vite-profile.cpuprofile')
+        const outPath = path.resolve(
+          `./vite-profile-${profileCount++}.cpuprofile`,
+        )
         fs.writeFileSync(outPath, JSON.stringify(profile))
         log(
           colors.yellow(
             `CPU profile written to ${colors.white(colors.dim(outPath))}`,
           ),
         )
+        profileSession = undefined
+        res()
       } else {
-        throw err
+        rej(err)
       }
     })
-  }
+  })
 }
 
 const filterDuplicateOptions = <T extends object>(options: T) => {
@@ -150,14 +159,30 @@ cli
       )
 
       server.printUrls()
-      server.bindShortcuts({
+      bindShortcuts(server, {
         print: true,
         customShortcuts: [
           profileSession && {
-            key: 's',
-            description: 'stop the profiler',
-            action(server) {
-              stopProfiler(server.config.logger.info)
+            key: 'p',
+            description: 'start/stop the profiler',
+            async action(server) {
+              if (profileSession) {
+                await stopProfiler(server.config.logger.info)
+              } else {
+                const inspector = await import('node:inspector').then(
+                  (r) => r.default,
+                )
+                await new Promise<void>((res) => {
+                  profileSession = new inspector.Session()
+                  profileSession.connect()
+                  profileSession.post('Profiler.enable', () => {
+                    profileSession!.post('Profiler.start', () => {
+                      server.config.logger.info('Profiler started')
+                      res()
+                    })
+                  })
+                })
+              }
             },
           },
         ],

diff --git a/packages/vite/src/node/index.ts b/packages/vite/src/node/index.ts
@@ -61,7 +61,6 @@ export type {
   LogType,
   LoggerOptions,
 } from './logger'
-export type { BindShortcutsOptions, CLIShortcut } from './shortcuts'
 export type {
   IndexHtmlTransform,
   IndexHtmlTransformHook,

diff --git a/packages/vite/src/node/server/index.ts b/packages/vite/src/node/server/index.ts
@@ -264,10 +264,6 @@ export interface ViteDevServer {
    * Print server urls
    */
   printUrls(): void
-  /**
-   * Listen to `process.stdin` for keyboard shortcuts.
-   */
-  bindShortcuts(opts?: BindShortcutsOptions): void
   /**
    * Restart the server.
    *
@@ -308,8 +304,11 @@ export interface ViteDevServer {
   _fsDenyGlob: Matcher
   /**
    * @internal
+   * Actually BindShortcutsOptions | undefined but api-extractor checks for
+   * export before trimming internal types :(
+   * And I don't want to complexity prePatchTypes for that
    */
-  _shortcutsBound: boolean
+  _shortcutsOptions: any | undefined
 }
 
 export interface ResolvedServerUrls {
@@ -445,10 +444,6 @@ export async function createServer(
         )
       }
     },
-    bindShortcuts(opts = {}) {
-      bindShortcuts(server, opts)
-      server._shortcutsBound = true
-    },
     async restart(forceOptimize?: boolean) {
       if (!server._restartPromise) {
         server._forceOptimizeOnRestart = !!forceOptimize
@@ -466,7 +461,7 @@ export async function createServer(
     _forceOptimizeOnRestart: false,
     _pendingRequests: new Map(),
     _fsDenyGlob: picomatch(config.server.fs.deny, { matchBase: true }),
-    _shortcutsBound: false,
+    _shortcutsOptions: undefined,
   }
 
   server.transformIndexHtml = createDevHtmlTransformFn(server)
@@ -786,7 +781,7 @@ async function restartServer(server: ViteDevServer) {
   // @ts-ignore
   global.__vite_start_time = performance.now()
   const { port: prevPort, host: prevHost } = server.config.server
-  const bindShortcuts = server._shortcutsBound
+  const shortcutsOptions: BindShortcutsOptions = server._shortcutsOptions
 
   await server.close()
 
@@ -835,8 +830,9 @@ async function restartServer(server: ViteDevServer) {
     logger.info('server restarted.', { timestamp: true })
   }
 
-  if (bindShortcuts) {
-    newServer.bindShortcuts()
+  if (shortcutsOptions) {
+    shortcutsOptions.print = false
+    bindShortcuts(newServer, shortcutsOptions)
   }
 
   // new server (the current server) can restart now

diff --git a/packages/vite/src/node/shortcuts.ts b/packages/vite/src/node/shortcuts.ts
@@ -22,6 +22,7 @@ export function bindShortcuts(
   opts: BindShortcutsOptions,
 ): void {
   if (!server.httpServer) return
+  server._shortcutsOptions = opts
 
   if (opts.print) {
     server.config.logger.info(