[Profiler] Grep for processes (#216770)

Grep for running Node.js processes if specified.

x-pack/platform/packages/shared/kbn-profiler-cli/README.md

@@ -1,6 +1,6 @@
 # @kbn/profiler-cli
 
-Profile Kibana while it's running, and open the CPU profile in Speedscope.
+Profile Kibana (or any other Node.js processes) while it's running, and open the CPU or Heap profile in Speedscope.
 
 ## Usage
 
@@ -18,6 +18,26 @@ Or with a timeout:
 
 `node scripts/profile.js --timeout=10000`
 
+### Heap profiling
+
+If you want to collect a heap profile, simply add `--heap`:
+
+`node scripts/profile.js --heap --timeout=10000`
+
+### Selecting a process
+
+By default, the profiler will look for a process running on 5603 or 5601 (in that order), where Kibana runs by default. But you can attach the profiler to any process. Add `--pid` to specify a specific process id:
+
+`node scripts/profile.js --pid 14782`
+
+Or, use `--grep` to list Node.js processes you can attach to:
+
+`node scripts/profile.js --grep`
+
+You can also already specify a filter:
+
+`node scripts/profile.js --grep myProcess`
+
 ## Examples
 
 ### Commands

x-pack/platform/packages/shared/kbn-profiler-cli/index.ts

@@ -6,19 +6,27 @@
  */
 import { run } from '@kbn/dev-cli-runner';
 import { compact, once, uniq } from 'lodash';
-import { getKibanaProcessId } from './src/get_kibana_process_id';
+import { getProcessId } from './src/get_process_id';
 import { runCommand } from './src/run_command';
 import { runUntilSigInt } from './src/run_until_sigint';
 import { getProfiler } from './src/get_profiler';
 import { untilStdinCompletes } from './src/until_stdin_completes';
 
+const NO_GREP = '__NO_GREP__';
+
 export function cli() {
   run(
     async ({ flags, log, addCleanupTask }) => {
+      // flags.grep can only be a string, and defaults to an empty string,
+      // so we override the default with a const and check for that
+      // to differentiate between ``, `--grep` and `--grep myString`
+      const grep = flags.grep === NO_GREP ? false : flags.grep === '' ? true : String(flags.grep);
+
       const pid = flags.pid
         ? Number(flags.pid)
-        : await getKibanaProcessId({
+        : await getProcessId({
             ports: uniq(compact([Number(flags.port), 5603, 5601])),
+            grep,
           });
 
       const controller = new AbortController();
@@ -28,9 +36,11 @@ export function cli() {
         }, Number(flags.timeout));
       }
 
+      log.debug(`Sending SIGUSR1 to ${pid}`);
+
       process.kill(pid, 'SIGUSR1');
 
-      const stop = once(await getProfiler({ log, type: flags.heap ? 'heap' : 'cpu' }));
+      const stop = once(await getProfiler({ pid, log, type: flags.heap ? 'heap' : 'cpu' }));
 
       addCleanupTask(() => {
         // exit-hook, which is used by addCleanupTask,
@@ -40,9 +50,18 @@ export function cli() {
         // process.exit a noop for a bit until the
         // profile has been collected and opened
         const exit = process.exit.bind(process);
+        const kill = process.kill.bind(process);
 
         // @ts-expect-error
         process.exit = () => {};
+        process.kill = (pidToKill, signal) => {
+          // inquirer sends a SIGINT kill signal to the process,
+          // that we need to handle here
+          if (pidToKill === process.pid && signal === 'SIGINT') {
+            return true;
+          }
+          return kill(pidToKill, signal);
+        };
 
         stop()
           .then(() => {
@@ -86,7 +105,7 @@ export function cli() {
     },
     {
       flags: {
-        string: ['port', 'pid', 't', 'timeout', 'c', 'connections', 'a', 'amount'],
+        string: ['port', 'pid', 't', 'timeout', 'c', 'connections', 'a', 'amount', 'grep'],
         boolean: ['heap'],
         help: `
           Usage: node scripts/profiler.js <args> <command>
@@ -97,7 +116,11 @@ export function cli() {
           --c, --connections  Number of commands that can be run in parallel.
           --a, --amount       Amount of times the command should be run
           --heap              Collect a heap snapshot
+          --grep              Grep through running Node.js processes
         `,
+        default: {
+          grep: NO_GREP,
+        },
         allowUnexpected: false,
       },
     }

x-pack/platform/packages/shared/kbn-profiler-cli/src/get_kibana_process_id.ts

@@ -1,29 +0,0 @@
-/*
- * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
- * or more contributor license agreements. Licensed under the Elastic License
- * 2.0; you may not use this file except in compliance with the Elastic License
- * 2.0.
- */
-import execa from 'execa';
-
-async function getProcessIdAtPort(port: number) {
-  return await execa
-    .command(`lsof -ti :${port}`)
-    .then(({ stdout }) => {
-      return parseInt(stdout.trim().split('\n')[0], 10);
-    })
-    .catch((error) => {
-      return undefined;
-    });
-}
-
-export async function getKibanaProcessId({ ports }: { ports: number[] }): Promise<number> {
-  for (const port of ports) {
-    const pid = await getProcessIdAtPort(port);
-    if (pid) {
-      return pid;
-    }
-  }
-
-  throw new Error(`Kibana process id not found at ports ${ports.join(', ')}`);
-}

x-pack/platform/packages/shared/kbn-profiler-cli/src/get_node_processes.ts

@@ -0,0 +1,78 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import execa from 'execa';
+
+export async function getNodeProcesses(
+  grep?: string
+): Promise<Array<{ pid: number; command: string; ports: number[] }>> {
+  const candidates = await execa
+    .command(
+      `ps -eo pid,command | grep -E '^[[:space:]0-9]+[[:space:]]*node[[:space:]]?' | grep -v grep ${
+        grep ? `| grep "${grep}" ` : ''
+      }`,
+      { shell: true, reject: false }
+    )
+    .then(({ stdout, exitCode }) => {
+      if (exitCode !== 0) {
+        return [];
+      }
+
+      // example
+      // 6915 /Users/dariogieselaar/.nvm/versions/node/v20.18.2/bin/node scripts/es snapshot
+      const lines = stdout.split('\n');
+      return lines.map((line) => {
+        const [pid, ...command] = line.trim().split(' ');
+        return {
+          pid: Number(pid.trim()),
+          command: command.join(' ').trim(),
+        };
+      });
+    });
+
+  if (!candidates.length) {
+    return [];
+  }
+
+  const pids = candidates.map((candidate) => candidate.pid);
+  const portsByPid: Record<string, number[]> = {};
+
+  await execa
+    .command(`lsof -Pan -i -iTCP -sTCP:LISTEN -p ${pids.join(',')}`, { shell: true, reject: false })
+    .then(({ stdout, exitCode }) => {
+      // exitCode 1 is returned when some of the ports don't match
+      if (exitCode !== 0 && exitCode !== 1) {
+        return;
+      }
+
+      const lines = stdout.split('\n').slice(1);
+
+      lines.forEach((line) => {
+        const values = line.trim().split(/\s+/);
+        const pid = values[1];
+        const name = values.slice(8).join(' ');
+
+        if (!name) {
+          return;
+        }
+
+        const portMatch = name.match(/:(\d+)(?:\s|\()/);
+        const port = portMatch ? Number(portMatch[1]) : undefined;
+        if (port) {
+          (portsByPid[pid] = portsByPid[pid] || []).push(port);
+        }
+      });
+    });
+
+  return candidates.map(({ pid, command }) => {
+    return {
+      pid,
+      command,
+      ports: portsByPid[pid] ?? [],
+    };
+  });
+}

x-pack/platform/packages/shared/kbn-profiler-cli/src/get_process_id.ts

@@ -0,0 +1,67 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+import execa from 'execa';
+import inquirer from 'inquirer';
+import { getNodeProcesses } from './get_node_processes';
+
+class ProcessNotFoundError extends Error {
+  constructor() {
+    super(`No Node.js processes found to attach to`);
+  }
+}
+
+async function getProcessIdAtPort(port: number) {
+  return await execa
+    .command(`lsof -ti :${port}`)
+    .then(({ stdout }) => {
+      return parseInt(stdout.trim().split('\n')[0], 10);
+    })
+    .catch((error) => {
+      return undefined;
+    });
+}
+
+export async function getProcessId({
+  ports,
+  grep,
+}: {
+  ports: number[];
+  grep: boolean | string;
+}): Promise<number> {
+  if (grep) {
+    const candidates = await getNodeProcesses(typeof grep === 'boolean' ? '' : grep);
+
+    if (!candidates.length) {
+      throw new ProcessNotFoundError();
+    }
+
+    const { pid } = await inquirer.prompt({
+      type: 'list',
+      name: 'pid',
+      message: `Select a Node.js process to attach to`,
+      choices: candidates.map((candidate) => {
+        return {
+          name: `${candidate.command}${
+            candidate.ports.length ? ` (Listening on ${candidate.ports.join(', ')})` : ``
+          }`,
+          value: candidate.pid,
+        };
+      }),
+    });
+
+    return pid as number;
+  }
+
+  for (const port of ports) {
+    const pid = await getProcessIdAtPort(port);
+    if (pid) {
+      return pid;
+    }
+  }
+
+  throw new Error(`Kibana process id not found at ports ${ports.join(', ')}`);
+}

x-pack/platform/packages/shared/kbn-profiler-cli/src/get_profiler.ts

@@ -10,6 +10,14 @@ import Os from 'os';
 import Path from 'path';
 import { ToolingLog } from '@kbn/tooling-log';
 import execa from 'execa';
+import getPort from 'get-port';
+import { getNodeProcesses } from './get_node_processes';
+
+class InspectorSessionConflictError extends Error {
+  constructor() {
+    super(`An inspector session is already running in another process. Close the process first.`);
+  }
+}
 
 async function getHeapProfiler({ client, log }: { client: CDP.Client; log: ToolingLog }) {
   await client.HeapProfiler.enable();
@@ -52,11 +60,28 @@ async function getCpuProfiler({ client, log }: { client: CDP.Client; log: Toolin
 export async function getProfiler({
   log,
   type,
+  pid,
 }: {
   log: ToolingLog;
   type: 'cpu' | 'heap';
+  pid: number;
 }): Promise<() => Promise<void>> {
-  log.debug(`Attaching to remote debugger at 9229`);
+  const port = await getPort({
+    host: '127.0.0.1',
+    port: 9229,
+  });
+
+  if (port !== 9229) {
+    // Inspector is already running, see if it's attached to the selected process
+    await getNodeProcesses()
+      .then((processes) => processes.find((process) => process.pid === pid))
+      .then((candidate) => {
+        if (!candidate?.ports.includes(9229)) {
+          throw new InspectorSessionConflictError();
+        }
+      });
+  }
+
   const client = await CDP({ port: 9229 });
 
   log.info(`Attached to remote debugger at 9229`);

x-pack/platform/packages/shared/kbn-profiler-cli/src/run_command.ts

@@ -58,8 +58,10 @@ export async function runCommand({
   const limiter = pLimit(connections);
 
   await Promise.allSettled(
-    range(0, amount).map(async () => {
-      await limiter(() => Promise.race([abortPromise, executeCommand()]));
+    range(0, amount).map(async (index) => {
+      await limiter(() => {
+        return Promise.race([abortPromise, executeCommand()]);
+      });
     })
   ).then((results) => {
     const errors = results.flatMap((result) =>