feat: add support for profiling through SANITY_DEBUG_PROFILING

Author: judofyr

packages/sanity/src/_internal/cli/debug.ts

@@ -1,3 +1,38 @@
+import {lstatSync} from 'node:fs'
+import {join} from 'node:path'
+
 import debugIt from 'debug'
 
 export const debug = debugIt('sanity:core')
+
+function isDir(path: string): boolean {
+  try {
+    return lstatSync(path).isDirectory()
+  } catch {
+    return false
+  }
+}
+
+/**
+ * Runs a function such that it will be profiled when the environment variable
+ * SANITY_DEBUG_PROFILING is set to a directory. A file (starting with `key`) will
+ * be placed in said directory. The generated file can be inspected by using the
+ * Speedscpe NPM package: `speedscope ${filename}` opens a UI in the browser.
+ */
+export async function withTracingProfiling<T>(key: string, fn: () => Promise<T>): Promise<T> {
+  const dir = process.env.SANITY_DEBUG_PROFILING
+  if (!dir) return await fn()
+
+  if (!isDir(dir))
+    throw new Error(`SANITY_DEBUG_PROFILING (${JSON.stringify(dir)}) must be set to a directory`)
+
+  let profiling
+  try {
+    profiling = await import('./util/profiling')
+  } catch (err) {
+    throw new Error(`Failed to load SANITY_DEBUG_PROFILING: ${err}`)
+  }
+
+  const filenamePrefix = join(dir, key)
+  return profiling.withTracing(filenamePrefix, fn)
+}

packages/sanity/src/_internal/cli/threads/extractManifest.ts

@@ -1,6 +1,7 @@
 import {isMainThread, parentPort, workerData as _workerData} from 'node:worker_threads'
 
 import {extractCreateWorkspaceManifest} from '../../manifest/extractWorkspaceManifest'
+import {withTracingProfiling} from '../debug'
 import {getStudioWorkspaces} from '../util/getStudioWorkspaces'
 import {mockBrowserEnvironment} from '../util/mockBrowserEnvironment'
 
@@ -30,4 +31,4 @@ async function main() {
   }
 }
 
-main()
+withTracingProfiling('extractManifest', main)

packages/sanity/src/_internal/cli/threads/extractSchema.ts

@@ -3,6 +3,7 @@ import {isMainThread, parentPort, workerData as _workerData} from 'node:worker_t
 import {extractSchema} from '@sanity/schema/_internal'
 import {type Workspace} from 'sanity'
 
+import {withTracingProfiling} from '../debug'
 import {getStudioWorkspaces} from '../util/getStudioWorkspaces'
 import {mockBrowserEnvironment} from '../util/mockBrowserEnvironment'
 
@@ -48,7 +49,7 @@ async function main() {
   }
 }
 
-main()
+withTracingProfiling('extractSchema', main)
 
 function getWorkspace({
   workspaces,

packages/sanity/src/_internal/cli/threads/getGraphQLAPIs.ts

@@ -7,13 +7,16 @@ import oneline from 'oneline'
 import {type Workspace} from 'sanity'
 
 import {type SchemaDefinitionish, type TypeResolvedGraphQLAPI} from '../actions/graphql/types'
+import {withTracingProfiling} from '../debug'
 import {getStudioWorkspaces} from '../util/getStudioWorkspaces'
 
-if (isMainThread || !parentPort) {
+const port = parentPort
+
+if (isMainThread || !port) {
   throw new Error('This module must be run as a worker thread')
 }
 
-getGraphQLAPIsForked(parentPort)
+withTracingProfiling('getGraphQLAPIs', async () => getGraphQLAPIsForked(port))
 
 async function getGraphQLAPIsForked(parent: MessagePort): Promise<void> {
   const {cliConfig, cliConfigPath, workDir} = workerData

packages/sanity/src/_internal/cli/threads/validateDocuments.ts

@@ -14,6 +14,7 @@ import {
 import {isReference, type ValidationContext, type ValidationMarker} from '@sanity/types'
 import {isRecord, validateDocument} from 'sanity'
 
+import {withTracingProfiling} from '../debug'
 import {extractDocumentsFromNdjsonOrTarball} from '../util/extractDocumentsFromNdjsonOrTarball'
 import {getStudioWorkspaces} from '../util/getStudioWorkspaces'
 import {mockBrowserEnvironment} from '../util/mockBrowserEnvironment'
@@ -129,7 +130,7 @@ async function* readerToGenerator(reader: ReadableStreamDefaultReader<Uint8Array
   }
 }
 
-validateDocuments()
+withTracingProfiling('validateDocuments', validateDocuments)
 
 async function loadWorkspace() {
   const workspaces = await getStudioWorkspaces({basePath: workDir, configPath})

packages/sanity/src/_internal/cli/threads/validateSchema.ts

@@ -4,6 +4,7 @@ import {groupProblems, validateSchema} from '@sanity/schema/_internal'
 import {type SchemaValidationProblem, type SchemaValidationProblemGroup} from '@sanity/types'
 import {resolveSchemaTypes} from 'sanity'
 
+import {withTracingProfiling} from '../debug'
 import {getStudioConfig} from '../util/getStudioWorkspaces'
 import {mockBrowserEnvironment} from '../util/mockBrowserEnvironment'
 
@@ -29,49 +30,53 @@ if (isMainThread || !parentPort) {
   throw new Error('This module must be run as a worker thread')
 }
 
-const cleanup = mockBrowserEnvironment(workDir)
+async function main() {
+  const cleanup = mockBrowserEnvironment(workDir)
 
-try {
-  const workspaces = getStudioConfig({basePath: workDir})
+  try {
+    const workspaces = getStudioConfig({basePath: workDir})
 
-  if (!workspaces.length) {
-    throw new Error(`Configuration did not return any workspaces.`)
-  }
-
-  let workspace
-  if (workspaceName) {
-    workspace = workspaces.find((w) => w.name === workspaceName)
-    if (!workspace) {
-      throw new Error(`Could not find any workspaces with name \`${workspaceName}\``)
+    if (!workspaces.length) {
+      throw new Error(`Configuration did not return any workspaces.`)
     }
-  } else {
-    if (workspaces.length !== 1) {
-      throw new Error(
-        "Multiple workspaces found. Please specify which workspace to use with '--workspace'.",
-      )
+
+    let workspace
+    if (workspaceName) {
+      workspace = workspaces.find((w) => w.name === workspaceName)
+      if (!workspace) {
+        throw new Error(`Could not find any workspaces with name \`${workspaceName}\``)
+      }
+    } else {
+      if (workspaces.length !== 1) {
+        throw new Error(
+          "Multiple workspaces found. Please specify which workspace to use with '--workspace'.",
+        )
+      }
+      workspace = workspaces[0]
     }
-    workspace = workspaces[0]
-  }
 
-  const schemaTypes = resolveSchemaTypes({
-    config: workspace,
-    context: {dataset: workspace.dataset, projectId: workspace.projectId},
-  })
+    const schemaTypes = resolveSchemaTypes({
+      config: workspace,
+      context: {dataset: workspace.dataset, projectId: workspace.projectId},
+    })
 
-  const validation = groupProblems(validateSchema(schemaTypes).getTypes())
+    const validation = groupProblems(validateSchema(schemaTypes).getTypes())
 
-  const result: ValidateSchemaWorkerResult = {
-    validation: validation
-      .map((group) => ({
-        ...group,
-        problems: group.problems.filter((problem) =>
-          level === 'error' ? problem.severity === 'error' : true,
-        ),
-      }))
-      .filter((group) => group.problems.length),
-  }
+    const result: ValidateSchemaWorkerResult = {
+      validation: validation
+        .map((group) => ({
+          ...group,
+          problems: group.problems.filter((problem) =>
+            level === 'error' ? problem.severity === 'error' : true,
+          ),
+        }))
+        .filter((group) => group.problems.length),
+    }
 
-  parentPort?.postMessage(result)
-} finally {
-  cleanup()
+    parentPort?.postMessage(result)
+  } finally {
+    cleanup()
+  }
 }
+
+withTracingProfiling('validateSchema', main)

packages/sanity/src/_internal/cli/util/profiling.ts

@@ -0,0 +1,30 @@
+import {writeFileSync} from 'node:fs'
+import {close as closeInspector, open as openInspector, Session} from 'node:inspector/promises'
+
+// This file should not be imported directly since it depends on
+// `inspector/promises` which is only available since Node v19 (and we want to
+// support earlier Node versions as well.
+
+/**
+ * Runs a function with a tracing profiler and writes the result into a file.
+ *
+ * @param filenamePrefix - The filename where the report will be written. The full name
+ * will be `{filenamePrefix}-{random}.cpuprofile`.
+ */
+export async function withTracing<T>(filenamePrefix: string, fn: () => Promise<T>): Promise<T> {
+  // Make it available in the Chrome DevTools as well
+
+  openInspector()
+  const session = new Session()
+  session.connect()
+  await session.post('Profiler.enable')
+  await session.post('Profiler.start')
+  try {
+    return await fn()
+  } finally {
+    closeInspector()
+    const fullname = `${filenamePrefix}-${Date.now()}-${Math.floor(Math.random() * 10000)}.cpuprofile`
+    const {profile} = await session.post('Profiler.stop')
+    writeFileSync(fullname, JSON.stringify(profile))
+  }
+}