Support async callback in linter (#8942)

- Add 'exitRoot' final event for semantic walker and linter rules
- Support 'async' in linter definition and async function as callback
for 'exitRoot' event.

---------

Co-authored-by: Timothee Guerin <[email protected]>

Author: RodgeFu

.chronus/changes/linter-async-callback-2025-10-6-12-33-19.md

@@ -0,0 +1,8 @@
+---
+changeKind: feature
+packages:
+  - "@typespec/compiler"
+---
+
+- Add 'exit' final event for linter rules
+- Support 'async' in linter definition and async function as callback for 'exit' event.

packages/compiler/src/core/linter.ts

@@ -1,3 +1,4 @@
+import { isPromise } from "../utils/misc.js";
 import { DiagnosticCollector, compilerAssert, createDiagnosticCollector } from "./diagnostics.js";
 import { getLocationContext } from "./helpers/location-context.js";
 import { defineLinter } from "./library.js";
@@ -6,7 +7,7 @@ import { createUnusedUsingLinterRule } from "./linter-rules/unused-using.rule.js
 import { createDiagnostic } from "./messages.js";
 import type { Program } from "./program.js";
 import { EventEmitter, mapEventEmitterToNodeListener, navigateProgram } from "./semantic-walker.js";
-import { startTimer, time } from "./stats.js";
+import { startTimer } from "./stats.js";
 import {
   Diagnostic,
   DiagnosticMessages,
@@ -26,7 +27,7 @@ type LinterLibraryInstance = { linter: LinterResolvedDefinition };
 export interface Linter {
   extendRuleSet(ruleSet: LinterRuleSet): Promise<readonly Diagnostic[]>;
   registerLinterLibrary(name: string, lib?: LinterLibraryInstance): void;
-  lint(): LinterResult;
+  lint(): Promise<LinterResult>;
 }
 
 export interface LinterStats {
@@ -157,7 +158,25 @@ export function createLinter(
     return diagnostics.diagnostics;
   }
 
-  function lint(): LinterResult {
+  async function lint(): Promise<LinterResult> {
+    const syncLintResult = await lintInternal(false /* asyncRules */);
+    const asyncLintResult = await lintInternal(true /* asyncRules */);
+
+    return {
+      diagnostics: [...syncLintResult.diagnostics, ...asyncLintResult.diagnostics],
+      stats: {
+        runtime: {
+          total: syncLintResult.stats.runtime.total + asyncLintResult.stats.runtime.total,
+          rules: {
+            ...syncLintResult.stats.runtime.rules,
+            ...asyncLintResult.stats.runtime.rules,
+          },
+        },
+      },
+    };
+  }
+
+  async function lintInternal(asyncRules: boolean): Promise<LinterResult> {
     const diagnostics = createDiagnosticCollector();
     const eventEmitter = new EventEmitter<SemanticNodeListener>();
     const stats: LinterStats = {
@@ -166,26 +185,61 @@ export function createLinter(
         rules: {},
       },
     };
+    const filteredRules = new Map<string, LinterRule<string, any>>();
+    for (const [ruleId, rule] of enabledRules) {
+      if ((rule.async ?? false) === asyncRules) {
+        filteredRules.set(ruleId, rule);
+      }
+    }
     tracer.trace(
       "lint",
-      `Running linter with following rules:\n` +
-        [...enabledRules.keys()].map((x) => ` - ${x}`).join("\n"),
+      `Running ${asyncRules ? "async" : "sync"} linter with following rules:\n` +
+        [...filteredRules.keys()].map((x) => ` - ${x}`).join("\n"),
     );
 
     const timer = startTimer();
-    for (const rule of enabledRules.values()) {
+    const exitCallbacks = [];
+    const EXIT_EVENT_NAME = "exit";
+    const allPromises: Promise<any>[] = [];
+    for (const rule of filteredRules.values()) {
       const createTiming = startTimer();
       const listener = rule.create(createLinterRuleContext(program, rule, diagnostics));
       stats.runtime.rules[rule.id] = createTiming.end();
       for (const [name, cb] of Object.entries(listener)) {
         const timedCb = (...args: any[]) => {
-          const duration = time(() => (cb as any)(...args));
-          stats.runtime.rules[rule.id] += duration;
+          const timer = startTimer();
+          const result = (cb as any)(...args);
+          if (name === EXIT_EVENT_NAME && isPromise(result)) {
+            compilerAssert(
+              rule.async,
+              `Linter rule "${rule.id}" is not marked as async but returned a promise from the "${name}" callback.`,
+            );
+            const rr = result.finally(() => {
+              const duration = timer.end();
+              stats.runtime.rules[rule.id] += duration;
+            });
+            allPromises.push(rr);
+          } else {
+            const duration = timer.end();
+            stats.runtime.rules[rule.id] += duration;
+          }
         };
-        eventEmitter.on(name as any, timedCb);
+        if (name === EXIT_EVENT_NAME) {
+          // we need to trigger 'exit' callbacks explicitly after semantic walker is done
+          exitCallbacks.push(timedCb);
+        } else {
+          eventEmitter.on(name as any, timedCb);
+        }
       }
     }
     navigateProgram(program, mapEventEmitterToNodeListener(eventEmitter));
+    for (const cb of exitCallbacks) {
+      cb(program);
+    }
+    if (allPromises.length > 0) {
+      await Promise.all(allPromises);
+    }
+
     stats.runtime.total = timer.end();
     return { diagnostics: diagnostics.diagnostics, stats };
   }

packages/compiler/src/core/program.ts

@@ -325,7 +325,7 @@ async function createProgram(
   }
 
   // Linter stage
-  const lintResult = linter.lint();
+  const lintResult = await linter.lint();
   runtimeStats.linter = lintResult.stats.runtime;
   program.reportDiagnostics(lintResult.diagnostics);
 

packages/compiler/src/core/types.ts

@@ -2323,7 +2323,7 @@ export interface LinterResolvedDefinition {
   };
 }
 
-export interface LinterRuleDefinition<N extends string, DM extends DiagnosticMessages> {
+interface LinterRuleDefinitionBase<N extends string, DM extends DiagnosticMessages> {
   /** Rule name (without the library name) */
   name: N;
   /** Rule default severity. */
@@ -2334,16 +2334,40 @@ export interface LinterRuleDefinition<N extends string, DM extends DiagnosticMes
   url?: string;
   /** Messages that can be reported with the diagnostic. */
   messages: DM;
+}
+
+interface LinterRuleDefinitionSync<N extends string, DM extends DiagnosticMessages>
+  extends LinterRuleDefinitionBase<N, DM> {
+  /** Whether this is an async rule. Default is false */
+  async?: false;
+  /** Creator */
+  create(
+    context: LinterRuleContext<DM>,
+  ): SemanticNodeListener & { exit?: (context: Program) => void | undefined };
+}
+
+interface LinterRuleDefinitionAsync<N extends string, DM extends DiagnosticMessages>
+  extends LinterRuleDefinitionBase<N, DM> {
+  /** Whether this is an async rule. Default is false */
+  async: true;
   /** Creator */
-  create(context: LinterRuleContext<DM>): SemanticNodeListener;
+  create(
+    context: LinterRuleContext<DM>,
+  ): SemanticNodeListener & { exit?: (context: Program) => Promise<void | undefined> };
 }
 
+export type LinterRuleDefinition<N extends string, DM extends DiagnosticMessages> =
+  | LinterRuleDefinitionSync<N, DM>
+  | LinterRuleDefinitionAsync<N, DM>;
+
 /** Resolved instance of a linter rule that will run. */
-export interface LinterRule<N extends string, DM extends DiagnosticMessages>
-  extends LinterRuleDefinition<N, DM> {
+export type LinterRule<N extends string, DM extends DiagnosticMessages> = LinterRuleDefinition<
+  N,
+  DM
+> & {
   /** Expanded rule id in format `<library-name>:<rule-name>` */
   id: string;
-}
+};
 
 /** Reference to a rule. In this format `<library name>:<rule/ruleset name>` */
 export type RuleRef = `${string}/${string}`;

packages/compiler/src/testing/rule-tester.ts

@@ -133,6 +133,9 @@ export function createLinterRuleTester(
     const context = createLinterRuleContext(runner.program, rule, diagnostics);
     const listener = ruleDef.create(context);
     navigateProgram(runner.program, listener);
+    if (listener.exit) {
+      await listener.exit(runner.program);
+    }
     // No diagnostics should have been reported to the program. If it happened the rule is calling reportDiagnostic directly and should NOT be doing that.
     expectDiagnosticEmpty(runner.program.diagnostics);
     return [res, diagnostics.diagnostics];

packages/compiler/src/utils/misc.ts

@@ -484,3 +484,7 @@ class RekeyableMapImpl<K, V> implements RekeyableMap<K, V> {
     return true;
   }
 }
+
+export function isPromise(value: unknown): value is Promise<unknown> {
+  return !!value && typeof (value as any).then === "function";
+}