[9.3] Speed up telemetry check by sharing a single TS program across roots !! (#267651)

src/platform/packages/private/kbn-telemetry-tools/src/tools/extract_collectors.ts

@@ -10,6 +10,7 @@
 import { readFileSync } from 'fs';
 import globby from 'globby';
 import * as path from 'path';
+import type ts from 'typescript';
 import { parseUsageCollection } from './ts_parser';
 import type { TelemetryRC } from './config';
 import { createKibanaProgram, getAllSourceFiles } from './ts_program';
@@ -40,8 +41,7 @@ export async function getProgramPaths({
   );
 
   if (filePaths.length === 0) {
-    return []; // Temporarily accept empty directories while https://github.com/elastic/kibana-team/issues/1066 is completed
-    // throw Error(`No files found in ${root}`);
+    return [];
   }
 
   const fullPaths = filePaths
@@ -55,11 +55,12 @@ export async function getProgramPaths({
   return fullPaths;
 }
 
+export function filterCollectorPaths(fullPaths: string[]): string[] {
+  return fullPaths.filter((p) => COLLECTOR_RE.test(readFileSync(p, 'utf-8')));
+}
+
 export function* extractCollectors(fullPaths: string[], tsConfig: any) {
-  // Pre-filter to only files that reference collector APIs so TS doesn't
-  // parse thousands of unrelated source files (36K → ~70 root files).
-  // TS still resolves transitive imports needed for type-checking.
-  const collectorPaths = fullPaths.filter((p) => COLLECTOR_RE.test(readFileSync(p, 'utf-8')));
+  const collectorPaths = filterCollectorPaths(fullPaths);
 
   if (collectorPaths.length === 0) {
     return;
@@ -72,3 +73,13 @@ export function* extractCollectors(fullPaths: string[], tsConfig: any) {
     yield* parseUsageCollection(sourceFile, program);
   }
 }
+
+export function* extractCollectorsWithProgram(collectorPaths: string[], program: ts.Program) {
+  if (collectorPaths.length === 0) {
+    return;
+  }
+  const sourceFiles = getAllSourceFiles(collectorPaths, program);
+  for (const sourceFile of sourceFiles) {
+    yield* parseUsageCollection(sourceFile, program);
+  }
+}

src/platform/packages/private/kbn-telemetry-tools/src/tools/tasks/extract_collectors_task.ts

@@ -10,39 +10,65 @@
 import ts from 'typescript';
 import * as path from 'path';
 import type { TaskContext } from './task_context';
-import { extractCollectors, getProgramPaths } from '../extract_collectors';
+import {
+  extractCollectorsWithProgram,
+  filterCollectorPaths,
+  getProgramPaths,
+} from '../extract_collectors';
+import { createKibanaProgram } from '../ts_program';
 
 export function extractCollectorsTask(
   { roots }: TaskContext,
   restrictProgramToPath?: string | string[]
 ) {
-  return roots.map((root) => ({
-    task: async () => {
-      const tsConfig = ts.findConfigFile('./', ts.sys.fileExists, 'tsconfig.json');
-      if (!tsConfig) {
-        throw new Error('Could not find a valid tsconfig.json.');
-      }
-      const programPaths = await getProgramPaths(root.config);
-
-      if (typeof restrictProgramToPath !== 'undefined') {
-        const restrictProgramToPaths = Array.isArray(restrictProgramToPath)
-          ? restrictProgramToPath
-          : [restrictProgramToPath];
-
-        const fullRestrictedPaths = restrictProgramToPaths.map((collectorPath) =>
-          path.resolve(process.cwd(), collectorPath)
-        );
-        const restrictedProgramPaths = programPaths.filter((programPath) =>
-          fullRestrictedPaths.includes(programPath)
+  return [
+    {
+      task: async () => {
+        const tsConfig = ts.findConfigFile('./', ts.sys.fileExists, 'tsconfig.json');
+        if (!tsConfig) {
+          throw new Error('Could not find a valid tsconfig.json.');
+        }
+
+        const rootPathsMap = new Map<number, string[]>();
+        await Promise.all(
+          roots.map(async (root, idx) => {
+            const programPaths = await getProgramPaths(root.config);
+            rootPathsMap.set(idx, programPaths);
+          })
         );
-        if (restrictedProgramPaths.length) {
-          root.parsedCollections = [...extractCollectors(restrictedProgramPaths, tsConfig)];
+
+        const rootCollectorMap = new Map<number, string[]>();
+        let allCollectorPaths: string[] = [];
+
+        for (const [idx, programPaths] of rootPathsMap) {
+          let paths = programPaths;
+
+          if (typeof restrictProgramToPath !== 'undefined') {
+            const restrictProgramToPaths = Array.isArray(restrictProgramToPath)
+              ? restrictProgramToPath
+              : [restrictProgramToPath];
+            const fullRestrictedPaths = restrictProgramToPaths.map((collectorPath) =>
+              path.resolve(process.cwd(), collectorPath)
+            );
+            paths = paths.filter((p) => fullRestrictedPaths.includes(p));
+          }
+
+          const collectorPaths = filterCollectorPaths(paths);
+          rootCollectorMap.set(idx, collectorPaths);
+          allCollectorPaths = allCollectorPaths.concat(collectorPaths);
+        }
+
+        if (allCollectorPaths.length === 0) {
+          return;
         }
-        return;
-      }
 
-      root.parsedCollections = [...extractCollectors(programPaths, tsConfig)];
+        const program = createKibanaProgram(allCollectorPaths, tsConfig);
+
+        for (const [idx, collectorPaths] of rootCollectorMap) {
+          roots[idx].parsedCollections = [...extractCollectorsWithProgram(collectorPaths, program)];
+        }
+      },
+      title: 'Extracting collectors across all roots',
     },
-    title: `Extracting collectors in ${root.config.root}`,
-  }));
+  ];
 }