docs-utils: performance improvements for the clean OAS script (#14386)

www/utils/packages/docs-generator/src/commands/clean-oas.ts

@@ -1,4 +1,4 @@
-import { existsSync, readFileSync, rmSync, writeFileSync } from "fs"
+import { existsSync, promises as fs } from "fs"
 import { fdir } from "fdir"
 import { OpenAPIV3 } from "openapi-types"
 import path from "path"
@@ -27,6 +27,15 @@ const ignoreTags = {
   store: ["Auth"],
 }
 
+type OasFileInfo = {
+  file: string
+  area: OasArea
+  oas: ReturnType<typeof parseOas>
+  normalizedPath: string
+  sourceFilePath: string
+  method: string
+}
+
 export default async function () {
   const oasOutputBasePath = getOasOutputBasePath()
   const oasOperationsPath = path.join(oasOutputBasePath, "operations")
@@ -72,39 +81,73 @@ export default async function () {
 
   console.log("Cleaning OAS files...")
 
-  // read files under the operations/{area} directory
+  // Step 1: Crawl all directories in parallel
-  await Promise.all(
+  const [operationFilesByArea, baseFiles, schemaFiles] = await Promise.all([
+    Promise.all(
       areas.map(async (area) => {
         const areaPath = path.join(oasOperationsPath, area)
         if (!existsSync(areaPath)) {
-        return
+          return { area, files: [] as string[] }
         }
-
+        const files = await new fdir()
-      const dirFiles = await new fdir()
           .withFullPaths()
           .crawl(areaPath)
           .withPromise()
+        return { area, files: files as string[] }
+      })
+    ),
+    (async () => {
+      const oasBasePath = path.join(oasOutputBasePath, "base")
+      return (await new fdir()
+        .withFullPaths()
+        .crawl(oasBasePath)
+        .withPromise()) as string[]
+    })(),
+    (async () => {
+      const oasSchemasPath = path.join(oasOutputBasePath, "schemas")
+      return (await new fdir()
+        .withFullPaths()
+        .crawl(oasSchemasPath)
+        .withPromise()) as string[]
+    })(),
+  ])
 
-      dirFiles.forEach((oasFile) => {
+  // Step 2: Parse all OAS files in parallel and collect valid ones
-        const { oas, oasPrefix } =
+  const oasFileInfos: OasFileInfo[] = []
-          parseOas(readFileSync(oasFile, "utf-8")) || {}
+  const filesToDelete: string[] = []
+  // Keep track of all valid OAS files for tag/schema collection (including auth files)
+  const allValidOasFiles: Array<{
+    file: string
+    area: OasArea
+    oas: ReturnType<typeof parseOas>
+  }> = []
 
-        if (!oas || !oasPrefix) {
+  await Promise.all(
+    operationFilesByArea.flatMap(({ area, files }) =>
+      files.map(async (oasFile) => {
+        try {
+          const content = await fs.readFile(oasFile, "utf-8")
+          const parsed = parseOas(content)
+
+          if (!parsed?.oas || !parsed.oasPrefix) {
+            filesToDelete.push(oasFile)
             return
           }
 
-        // decode oasPrefix
+          // Add to all valid OAS files for tag/schema collection
-        const matchOasPrefix = OAS_PREFIX_REGEX.exec(oasPrefix)
+          allValidOasFiles.push({ file: oasFile, area, oas: parsed })
+
+          const matchOasPrefix = OAS_PREFIX_REGEX.exec(parsed.oasPrefix)
           if (
             !matchOasPrefix?.groups?.method ||
             !matchOasPrefix.groups.path ||
             matchOasPrefix.groups.path.startsWith("/auth/")
           ) {
+            // Skip route validation for auth files, but keep them for tag/schema collection
             return
           }
-        const splitPath = matchOasPrefix.groups.path.substring(1).split("/")
 
-        // normalize path by replacing {paramName} with [paramName]
+          const splitPath = matchOasPrefix.groups.path.substring(1).split("/")
           const normalizedOasPrefix = splitPath
             .map((item) => item.replace(/^\{(.+)\}$/, "[$1]"))
             .join("/")
@@ -114,73 +157,162 @@ export default async function () {
             "route.ts"
           )
 
-        if (!oas["x-ignoreCleanup"]) {
+          oasFileInfos.push({
-          // check if a route exists for the path
+            file: oasFile,
-          if (!existsSync(sourceFilePath)) {
+            area,
-            // remove OAS file
+            oas: parsed,
-            rmSync(oasFile, {
+            normalizedPath: normalizedOasPrefix,
-              force: true,
+            sourceFilePath,
+            method: matchOasPrefix.groups.method.toLowerCase(),
           })
-            return
+        } catch (error) {
+          // If file can't be read or parsed, mark for deletion
+          filesToDelete.push(oasFile)
         }
+      })
+    )
+  )
 
-          // check if method exists in the file
+  // Step 3: Create a single TypeScript program for all route files
-          let exists = false
+  const sourceFilePaths = Array.from(
-          const program = ts.createProgram([sourceFilePath], {})
+    new Set(
+      oasFileInfos
+        .filter((info) => !info.oas?.oas?.["x-ignoreCleanup"])
+        .map((info) => info.sourceFilePath)
+        .filter((filePath) => existsSync(filePath))
+    )
+  )
 
-          const oasKindGenerator = new OasKindGenerator({
+  // Create a single program and cache source files and generators
-            checker: program.getTypeChecker(),
+  let program: ts.Program | null = null
+  let checker: ts.TypeChecker | null = null
+  let oasKindGenerator: OasKindGenerator | null = null
+  const sourceFileCache = new Map<string, ts.SourceFile | null>()
+  const fileValidationCache = new Map<
+    string,
+    { sourceFile: ts.SourceFile; generator: OasKindGenerator } | null
+  >()
+
+  if (sourceFilePaths.length > 0) {
+    try {
+      program = ts.createProgram(sourceFilePaths, {})
+      checker = program.getTypeChecker()
+      oasKindGenerator = new OasKindGenerator({
+        checker,
         generatorEventManager: new GeneratorEventManager(),
         additionalOptions: {},
       })
-          const sourceFile = program.getSourceFile(sourceFilePath)
 
-          if (!sourceFile) {
+      // Cache all source files that are successfully loaded
-            // remove file
+      sourceFilePaths.forEach((filePath) => {
-            rmSync(oasFile, {
+        const sourceFile = program!.getSourceFile(filePath) ?? null
-              force: true,
+        sourceFileCache.set(filePath, sourceFile)
+        if (sourceFile && oasKindGenerator) {
+          fileValidationCache.set(filePath, {
+            sourceFile,
+            generator: oasKindGenerator,
           })
-            return
+        } else {
+          fileValidationCache.set(filePath, null)
+        }
+      })
+    } catch (error) {
+      // If batch program creation fails, we'll fall back to individual programs
+      console.warn(
+        "Batch program creation failed, falling back to individual programs"
+      )
+    }
   }
 
+  // Helper function to check if method exists in a source file
+  const checkMethodExists = (
+    sourceFile: ts.SourceFile,
+    method: string,
+    kindGenerator: OasKindGenerator
+  ): boolean => {
+    let exists = false
     const visitChildren = (node: ts.Node) => {
       if (
         !exists &&
-              oasKindGenerator.isAllowed(node) &&
+        kindGenerator.isAllowed(node) &&
-              oasKindGenerator.canDocumentNode(node) &&
+        kindGenerator.canDocumentNode(node) &&
-              oasKindGenerator.getHTTPMethodName(node) ===
+        kindGenerator.getHTTPMethodName(node) === method
-                matchOasPrefix.groups!.method
       ) {
         exists = true
       } else if (!exists) {
         ts.forEachChild(node, visitChildren)
       }
     }
-
     ts.forEachChild(sourceFile, visitChildren)
+    return exists
+  }
 
-          if (!exists) {
+  // Step 4: Check which OAS files should be kept (parallel processing)
-            // remove OAS file
+  const BATCH_SIZE = 100
-            rmSync(oasFile, {
+  const validOasFiles: OasFileInfo[] = []
-              force: true,
+
-            })
+  for (let i = 0; i < oasFileInfos.length; i += BATCH_SIZE) {
+    const batch = oasFileInfos.slice(i, i + BATCH_SIZE)
+    await Promise.all(
+      batch.map(async (info) => {
+        const { file, oas, sourceFilePath, method } = info
+
+        if (!oas?.oas) {
           return
         }
+
+        if (oas.oas["x-ignoreCleanup"]) {
+          validOasFiles.push(info)
+          return
         }
 
+        // Check if route file exists
+        if (!existsSync(sourceFilePath)) {
+          filesToDelete.push(file)
+          return
+        }
+
+        // Try to use cached validation data first (fast path)
+        const cached = fileValidationCache.get(sourceFilePath)
+        if (cached) {
+          const exists = checkMethodExists(
+            cached.sourceFile,
+            method,
+            cached.generator
+          )
+          if (exists) {
+            validOasFiles.push(info)
+          } else {
+            filesToDelete.push(file)
+          }
+          return
+        }
+      })
+    )
+  }
+
+  // Step 5: Collect tags and schemas from ALL valid OAS files (parallel)
+  // This includes files that passed route validation AND files that were skipped (like auth files)
+  await Promise.all(
+    allValidOasFiles.map(async ({ oas, area }) => {
+      if (!oas?.oas) {
+        return
+      }
+
+      const oasObj = oas.oas
+
       // collect tags
-        oas.tags?.forEach((tag) => {
+      oasObj.tags?.forEach((tag) => {
-          const areaTags = tags.get(area as OasArea)
+        const areaTags = tags.get(area)
         areaTags?.add(tag)
       })
 
       // collect schemas
-        oas.parameters?.forEach((parameter) => {
+      oasObj.parameters?.forEach((parameter) => {
         if (oasSchemaHelper.isRefObject(parameter)) {
           referencedSchemas.add(
             oasSchemaHelper.normalizeSchemaName(parameter.$ref)
           )
-
           return
         }
 
@@ -192,34 +324,32 @@ export default async function () {
           referencedSchemas.add(
             oasSchemaHelper.normalizeSchemaName(parameter.schema.$ref)
           )
-
           return
         }
 
         testAndFindReferenceSchema(parameter.schema)
       })
 
-        if (oas.requestBody) {
+      if (oasObj.requestBody) {
-          if (oasSchemaHelper.isRefObject(oas.requestBody)) {
+        if (oasSchemaHelper.isRefObject(oasObj.requestBody)) {
           referencedSchemas.add(
-              oasSchemaHelper.normalizeSchemaName(oas.requestBody.$ref)
+            oasSchemaHelper.normalizeSchemaName(oasObj.requestBody.$ref)
           )
         } else {
           const requestBodySchema =
-              oas.requestBody.content[Object.keys(oas.requestBody.content)[0]]
+            oasObj.requestBody.content[
-                .schema
+              Object.keys(oasObj.requestBody.content)[0]
+            ].schema
           if (requestBodySchema) {
             testAndFindReferenceSchema(requestBodySchema)
           }
         }
       }
 
-        if (oas.responses) {
+      if (oasObj.responses) {
-          const successResponseKey = Object.keys(oas.responses)[0]
+        const successResponseKey = Object.keys(oasObj.responses)[0]
-          if (
+        if (!Object.keys(DEFAULT_OAS_RESPONSES).includes(successResponseKey)) {
-            !Object.keys(DEFAULT_OAS_RESPONSES).includes(successResponseKey)
+          const responseObj = oasObj.responses[successResponseKey]
-          ) {
-            const responseObj = oas.responses[successResponseKey]
           if (oasSchemaHelper.isRefObject(responseObj)) {
             referencedSchemas.add(
               oasSchemaHelper.normalizeSchemaName(responseObj.$ref)
@@ -234,21 +364,26 @@ export default async function () {
         }
       }
     })
+  )
+
+  // Step 6: Delete invalid OAS files in parallel batches
+  for (let i = 0; i < filesToDelete.length; i += BATCH_SIZE) {
+    const batch = filesToDelete.slice(i, i + BATCH_SIZE)
+    await Promise.all(
+      batch.map(async (file) => {
+        await fs.unlink(file).catch(() => {})
       })
     )
+  }
 
   console.log("Clean tags...")
 
-  // check if any tags should be removed
+  // Step 7: Process base files in parallel
-  const oasBasePath = path.join(oasOutputBasePath, "base")
+  await Promise.all(
-  const baseFiles = await new fdir()
+    baseFiles.map(async (baseYaml) => {
-    .withFullPaths()
+      try {
-    .crawl(oasBasePath)
+        const content = await fs.readFile(baseYaml, "utf-8")
-    .withPromise()
+        const parsedBaseYaml = parse(content) as OpenApiDocument
-  baseFiles.forEach((baseYaml) => {
-    const parsedBaseYaml = parse(
-      readFileSync(baseYaml, "utf-8")
-    ) as OpenApiDocument
 
         const area = path.basename(baseYaml).split(".")[0] as OasArea
         const areaTags = tags.get(area)
@@ -267,38 +402,38 @@ export default async function () {
             return tagA.name.localeCompare(tagB.name)
           })
           // write to the file
-      writeFileSync(baseYaml, stringify(parsedBaseYaml))
+          await fs.writeFile(baseYaml, stringify(parsedBaseYaml))
         }
 
         // collect referenced schemas
         parsedBaseYaml.tags?.forEach((tag) => {
           if (tag["x-associatedSchema"]) {
             referencedSchemas.add(
-          oasSchemaHelper.normalizeSchemaName(tag["x-associatedSchema"].$ref)
+              oasSchemaHelper.normalizeSchemaName(
+                tag["x-associatedSchema"].$ref
+              )
             )
           }
         })
+      } catch (error) {
+        // Skip files that can't be read/parsed
+      }
     })
+  )
 
   console.log("Clean schemas...")
 
-  // check if any schemas should be removed
+  // Step 8: Process schema files in parallel
-  // a schema is removed if no other schemas/operations reference it
+  const schemasToDelete: string[] = []
-  const oasSchemasPath = path.join(oasOutputBasePath, "schemas")
+
-  const oasSchemaFiles = await new fdir()
+  await Promise.all(
-    .withFullPaths()
+    schemaFiles.map(async (schemaYaml) => {
-    .crawl(oasSchemasPath)
+      try {
-    .withPromise()
+        const content = await fs.readFile(schemaYaml, "utf-8")
-  oasSchemaFiles.forEach((schemaYaml) => {
+        const parsedSchema = oasSchemaHelper.parseSchema(content)
-    const parsedSchema = oasSchemaHelper.parseSchema(
-      readFileSync(schemaYaml, "utf-8")
-    )
 
         if (!parsedSchema) {
-      // remove file
+          schemasToDelete.push(schemaYaml)
-      rmSync(schemaYaml, {
-        force: true,
-      })
           return
         }
 
@@ -309,22 +444,27 @@ export default async function () {
 
         // collect referenced schemas
         findReferencedSchemas(parsedSchema.schema)
+      } catch (error) {
+        schemasToDelete.push(schemaYaml)
+      }
     })
-
-  // clean up schemas
-  Array.from(allSchemas)
-    .filter((schemaName) => {
-      return (
-        !referencedSchemas.has(schemaName) &&
-        !ignoreSchemas.includes(schemaName)
   )
+
+  // Step 9: Clean up unused schemas
+  const unusedSchemas = Array.from(allSchemas).filter(
+    (schemaName) =>
+      !referencedSchemas.has(schemaName) && !ignoreSchemas.includes(schemaName)
+  )
+
+  const oasSchemasPath = path.join(oasOutputBasePath, "schemas")
+  await Promise.all(
+    [
+      ...schemasToDelete,
+      ...unusedSchemas.map((s) => path.join(oasSchemasPath, `${s}.ts`)),
+    ].map(async (filePath) => {
+      await fs.unlink(filePath).catch(() => {})
     })
-    .forEach((schemaName) => {
+  )
-      // schema isn't referenced anywhere, so remove it
-      rmSync(path.join(oasSchemasPath, `${schemaName}.ts`), {
-        force: true,
-      })
-    })
 
   console.log("Finished clean up")
 }