fix(rosetta): breaks when given a lot of snippets (#3075)

Rosetta started breaking recently (the process just exiting witout any error),
presumably because we increased the number of compiling code snippets and
hitting some undocumented limit in the worker message API.

Upon closer inspection the `ts.Diagnostic` objects we were sending over the
wire from the worker to the main process were holding on to a `ts.SourceFile`
object, and by extension to the entire parse tree and internal TypeScript
objects.

Instead, introduce a new type of object (`RosettaDiagnostic`) which is
a simple data holder with a couple of fields, and render the original
TypeScript diagnostics down to these and send these over the wire.

This should improve memory usage, and appears to make the rosetta
process finish again.



---

By submitting this pull request, I confirm that my contribution is made under the terms of the [Apache 2.0 license].

[Apache 2.0 license]: https://www.apache.org/licenses/LICENSE-2.0

packages/jsii-rosetta/bin/jsii-rosetta.ts

@@ -4,7 +4,7 @@ import * as fs from 'fs-extra';
 import * as path from 'path';
 import * as yargs from 'yargs';
 
-import { TranslateResult, DEFAULT_TABLET_NAME, translateTypeScript } from '../lib';
+import { TranslateResult, DEFAULT_TABLET_NAME, translateTypeScript, RosettaDiagnostic } from '../lib';
 import { translateMarkdown } from '../lib/commands/convert';
 import { extractSnippets } from '../lib/commands/extract';
 import { readTablet } from '../lib/commands/read';
@@ -13,7 +13,7 @@ import { TargetLanguage } from '../lib/languages';
 import { PythonVisitor } from '../lib/languages/python';
 import { VisualizeAstVisitor } from '../lib/languages/visualize';
 import * as logging from '../lib/logging';
-import { File, printDiagnostics, isErrorDiagnostic } from '../lib/util';
+import { File, printDiagnostics } from '../lib/util';
 
 function main() {
   const argv = yargs
@@ -41,7 +41,7 @@ function main() {
           }),
       wrapHandler(async (args) => {
         const result = translateTypeScript(await makeFileSource(args.FILE ?? '-', 'stdin.ts'), makeVisitor(args));
-        renderResult(result);
+        handleSingleResult(result);
       }),
     )
     .command(
@@ -60,7 +60,7 @@ function main() {
           }),
       wrapHandler(async (args) => {
         const result = translateMarkdown(await makeFileSource(args.FILE ?? '-', 'stdin.md'), makeVisitor(args));
-        renderResult(result);
+        handleSingleResult(result);
       }),
     )
     .command(
@@ -140,15 +140,7 @@ function main() {
           only: args.include,
         });
 
-        printDiagnostics(result.diagnostics, process.stderr);
-
-        if (result.diagnostics.length > 0) {
-          logging.warn(`${result.diagnostics.length} diagnostics encountered in ${result.tablet.count} snippets`);
-        }
-
-        if (result.diagnostics.some((diag) => isErrorDiagnostic(diag, { onlyStrict: !args.fail }))) {
-          process.exitCode = 1;
-        }
+        handleDiagnostics(result.diagnostics, args.fail, result.tablet.count);
       }),
     )
     .command(
@@ -285,7 +277,9 @@ function wrapHandler<A extends { verbose?: number }, R>(handler: (x: A) => Promi
   return (argv: A) => {
     logging.configure({ level: argv.verbose !== undefined ? argv.verbose : 0 });
     handler(argv).catch((e) => {
-      throw e;
+      logging.error(e.message);
+      logging.error(e.stack);
+      process.exitCode = 1;
     });
   };
 }
@@ -329,15 +323,58 @@ async function readStdin(): Promise<string> {
   });
 }
 
-function renderResult(result: TranslateResult) {
+function handleSingleResult(result: TranslateResult) {
   process.stdout.write(`${result.translation}\n`);
 
-  if (result.diagnostics.length > 0) {
-    printDiagnostics(result.diagnostics, process.stderr);
+  // For a single result, we always request implicit failure.
+  handleDiagnostics(result.diagnostics, 'implicit');
+}
 
-    if (result.diagnostics.some((diag) => isErrorDiagnostic(diag, { onlyStrict: false }))) {
-      process.exit(1);
+/**
+ * Print diagnostics and set exit code
+ *
+ * 'fail' is whether or not the user passed '--fail' for commands that accept
+ * it, or 'implicit' for commands that should always fail. 'implicit' will be
+ * treated as 'fail=true, but will not print to the user that the '--fail' is
+ * set (because for this particular command that switch does not exist and so it
+ * would be confusing).
+ */
+function handleDiagnostics(diagnostics: readonly RosettaDiagnostic[], fail: boolean | 'implicit', snippetCount = 1) {
+  if (fail !== false) {
+    // Fail on any diagnostic
+    if (diagnostics.length > 0) {
+      printDiagnostics(diagnostics, process.stderr);
+      logging.error(
+        [
+          `${diagnostics.length} diagnostics encountered in ${snippetCount} snippets`,
+          ...(fail === true ? ["(running with '--fail')"] : []),
+        ].join(' '),
+      );
+      process.exitCode = 1;
     }
+
+    return;
+  }
+
+  // Otherwise fail only on strict diagnostics. If we have strict diagnostics, print only those
+  // (so it's very clear what is failing the build), otherwise print everything.
+  const strictDiagnostics = diagnostics.filter((diag) => diag.isFromStrictAssembly);
+  if (strictDiagnostics.length > 0) {
+    printDiagnostics(strictDiagnostics, process.stderr);
+    const remaining = diagnostics.length - strictDiagnostics.length;
+    logging.warn(
+      [
+        `${strictDiagnostics.length} diagnostics from assemblies with 'strict' mode on`,
+        ...(remaining > 0 ? [`(and ${remaining} more non-strict diagnostics)`] : []),
+      ].join(' '),
+    );
+    process.exitCode = 1;
+    return;
+  }
+
+  if (diagnostics.length > 0) {
+    printDiagnostics(diagnostics, process.stderr);
+    logging.warn(`${diagnostics.length} diagnostics encountered in ${snippetCount} snippets`);
   }
 }
 

packages/jsii-rosetta/lib/commands/convert.ts

@@ -2,7 +2,7 @@ import { transformMarkdown } from '../markdown/markdown';
 import { MarkdownRenderer } from '../markdown/markdown-renderer';
 import { ReplaceTypeScriptTransform } from '../markdown/replace-typescript-transform';
 import { AstHandler, AstRendererOptions } from '../renderer';
-import { TranslateResult, Translator } from '../translate';
+import { TranslateResult, Translator, rosettaDiagFromTypescript } from '../translate';
 import { File } from '../util';
 
 export interface TranslateMarkdownOptions extends AstRendererOptions {
@@ -38,6 +38,6 @@ export function translateMarkdown(
 
   return {
     translation: translatedMarkdown,
-    diagnostics: translator.diagnostics,
+    diagnostics: translator.diagnostics.map(rosettaDiagFromTypescript),
   };
 }

packages/jsii-rosetta/lib/commands/extract.ts

@@ -8,11 +8,11 @@ import * as logging from '../logging';
 import { TypeScriptSnippet } from '../snippet';
 import { snippetKey } from '../tablets/key';
 import { LanguageTablet, TranslatedSnippet } from '../tablets/tablets';
-import { Translator } from '../translate';
+import { RosettaDiagnostic, Translator, rosettaDiagFromTypescript } from '../translate';
 import { divideEvenly } from '../util';
 
 export interface ExtractResult {
-  diagnostics: ts.Diagnostic[];
+  diagnostics: RosettaDiagnostic[];
   tablet: LanguageTablet;
 }
 
@@ -64,7 +64,7 @@ export async function extractSnippets(
 
 interface TranslateAllResult {
   translatedSnippets: TranslatedSnippet[];
-  diagnostics: ts.Diagnostic[];
+  diagnostics: RosettaDiagnostic[];
 }
 
 /**
@@ -133,7 +133,7 @@ export function singleThreadedTranslateAll(
 
   return {
     translatedSnippets,
-    diagnostics: [...translator.diagnostics, ...failures],
+    diagnostics: [...translator.diagnostics, ...failures].map(rosettaDiagFromTypescript),
   };
 }
 

packages/jsii-rosetta/lib/commands/extract_worker.ts

@@ -1,11 +1,12 @@
 /**
  * Pool worker for extract.ts
  */
-import * as ts from 'typescript';
 import * as worker from 'worker_threads';
 
+import * as logging from '../logging';
 import { TypeScriptSnippet } from '../snippet';
 import { TranslatedSnippetSchema } from '../tablets/schema';
+import { RosettaDiagnostic } from '../translate';
 import { singleThreadedTranslateAll } from './extract';
 
 export interface TranslateRequest {
@@ -14,7 +15,7 @@ export interface TranslateRequest {
 }
 
 export interface TranslateResponse {
-  diagnostics: ts.Diagnostic[];
+  diagnostics: RosettaDiagnostic[];
   // Cannot be 'TranslatedSnippet' because needs to be serializable
   translatedSnippetSchemas: TranslatedSnippetSchema[];
 }
@@ -36,6 +37,10 @@ if (worker.isMainThread) {
   throw new Error('This script should be run as a worker, not included directly.');
 }
 
-const request = worker.workerData;
+const request: TranslateRequest = worker.workerData;
+const startTime = Date.now();
 const response = translateSnippet(request);
+const delta = (Date.now() - startTime) / 1000;
+// eslint-disable-next-line prettier/prettier
+logging.info(`Finished translation of ${request.snippets.length} in ${delta.toFixed(0)}s (${response.translatedSnippetSchemas.length} responses)`);
 worker.parentPort!.postMessage(response);

packages/jsii-rosetta/lib/logging.ts

@@ -1,6 +1,7 @@
 import * as util from 'util';
 
 export enum Level {
+  ERROR = -2,
   WARN = -1,
   QUIET = 0,
   INFO = 1,
@@ -21,6 +22,10 @@ export function warn(fmt: string, ...args: any[]) {
   log(Level.WARN, fmt, ...args);
 }
 
+export function error(fmt: string, ...args: any[]) {
+  log(Level.ERROR, fmt, ...args);
+}
+
 export function info(fmt: string, ...args: any[]) {
   log(Level.INFO, fmt, ...args);
 }
@@ -32,6 +37,8 @@ export function debug(fmt: string, ...args: any[]) {
 function log(messageLevel: Level, fmt: string, ...args: any[]) {
   if (level >= messageLevel) {
     const levelName = Level[messageLevel];
-    process.stderr.write(`[jsii-rosetta] [${levelName}] ${util.format(fmt, ...args)}\n`);
+    // `console.error` will automatically be transported from worker child to worker parent,
+    // process.stderr.write() won't.
+    console.error(`[jsii-rosetta] [${levelName}] ${util.format(fmt, ...args)}`);
   }
 }

packages/jsii-rosetta/lib/rosetta.ts

@@ -11,7 +11,7 @@ import { ReplaceTypeScriptTransform } from './markdown/replace-typescript-transf
 import { CodeBlock } from './markdown/types';
 import { SnippetParameters, TypeScriptSnippet, updateParameters } from './snippet';
 import { DEFAULT_TABLET_NAME, LanguageTablet, Translation } from './tablets/tablets';
-import { Translator } from './translate';
+import { Translator, rosettaDiagFromTypescript } from './translate';
 import { printDiagnostics } from './util';
 
 export interface RosettaOptions {
@@ -71,7 +71,7 @@ export class Rosetta {
    * Diagnostics encountered while doing live translation
    */
   public get diagnostics() {
-    return this.translator.diagnostics;
+    return this.translator.diagnostics.map(rosettaDiagFromTypescript);
   }
 
   /**

packages/jsii-rosetta/lib/translate.ts

@@ -10,7 +10,7 @@ import { snippetKey } from './tablets/key';
 import { TranslatedSnippet } from './tablets/tablets';
 import { calculateVisibleSpans } from './typescript/ast-utils';
 import { TypeScriptCompiler, CompilationResult } from './typescript/ts-compiler';
-import { annotateStrictDiagnostic, File } from './util';
+import { annotateStrictDiagnostic, File, hasStrictBranding } from './util';
 
 export function translateTypeScript(
   source: File,
@@ -22,7 +22,7 @@ export function translateTypeScript(
 
   return {
     translation: translated,
-    diagnostics: translator.diagnostics,
+    diagnostics: translator.diagnostics.map(rosettaDiagFromTypescript),
   };
 }
 
@@ -95,7 +95,35 @@ export interface SnippetTranslatorOptions extends AstRendererOptions {
 
 export interface TranslateResult {
   translation: string;
-  diagnostics: readonly ts.Diagnostic[];
+  diagnostics: readonly RosettaDiagnostic[];
+}
+
+/**
+ * A translation of a TypeScript diagnostic into a data-only representation for Rosetta
+ *
+ * We cannot use the original `ts.Diagnostic` since it holds on to way too much
+ * state (the source file and by extension the entire parse tree), which grows
+ * too big to be properly serialized by a worker and also takes too much memory.
+ *
+ * Reduce it down to only the information we need.
+ */
+export interface RosettaDiagnostic {
+  /**
+   * If this is an error diagnostic or not
+   */
+  readonly isError: boolean;
+
+  /**
+   * If the diagnostic was emitted from an assembly that has its 'strict' flag set
+   */
+  readonly isFromStrictAssembly: boolean;
+
+  /**
+   * The formatted message, ready to be printed (will have colors and newlines in it)
+   *
+   * Ends in a newline.
+   */
+  readonly formattedMessage: string;
 }
 
 /**
@@ -178,3 +206,26 @@ function filterVisibleDiagnostics(diags: readonly ts.Diagnostic[], visibleSpans:
     (d) => d.source !== 'rosetta' || d.start === undefined || visibleSpans.some((s) => spanContains(s, d.start!)),
   );
 }
+
+/**
+ * Turn TypeScript diagnostics into Rosetta diagnostics
+ */
+export function rosettaDiagFromTypescript(diag: ts.Diagnostic): RosettaDiagnostic {
+  return {
+    isError: diag.category === ts.DiagnosticCategory.Error,
+    isFromStrictAssembly: hasStrictBranding(diag),
+    formattedMessage: ts.formatDiagnosticsWithColorAndContext([diag], DIAG_HOST),
+  };
+}
+
+const DIAG_HOST = {
+  getCurrentDirectory() {
+    return '.';
+  },
+  getCanonicalFileName(fileName: string) {
+    return fileName;
+  },
+  getNewLine() {
+    return '\n';
+  },
+};