fix(pacmak): greatly reduce go code-gen memory footprint (#3362)

The go code generator was previously holding multiple copies of the
`RootModule` objects for dependencies (one for the generated package's
own `RootModule`, and one for each submodule that uses a type from that
dependency), including their entire type hierarchy, methods and
properties.

On large packages, especially with many submodules, this results in an
extremely large resident set size. Pathological cases have required
heaps as big as 6 GiB.

De-duplicating the `RootModule` instances, and not retaining `GoTypeRef`
instances (which are stateless, and cheap to construct), allows reducing
those pathological modules to only occupy around 300 to 400 MiB worth of
heap, which is a very significant improvement.

The reduction of duplication also has significant performance benefits
on affected code-bases. The go bindings generate significantly faster on
top of needing significantly less memory.



---

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-pacmak/lib/targets/go/package.ts

@@ -41,7 +41,7 @@ export abstract class Package {
   public readonly submodules: InternalPackage[];
   public readonly types: GoType[];
 
-  private readonly embeddedTypes: { [fqn: string]: EmbeddedType } = {};
+  private readonly embeddedTypes = new Map<string, EmbeddedType>();
 
   public constructor(
     private readonly typeSpec: readonly Type[],
@@ -158,7 +158,7 @@ export abstract class Package {
       };
     }
 
-    const exists = this.embeddedTypes[type.fqn];
+    const exists = this.embeddedTypes.get(type.fqn);
     if (exists) {
       return exists;
     }
@@ -169,14 +169,15 @@ export abstract class Package {
     const slug = original.replace(/[^A-Za-z0-9]/g, '');
     const aliasName = `Type__${slug}`;
 
-    this.embeddedTypes[type.fqn] = {
+    const embeddedType: EmbeddedType = {
       foriegnTypeName: original,
       foriegnType: typeref,
       fieldName: aliasName,
       embed: `${INTERNAL_PACKAGE_NAME}.${aliasName}`,
     };
+    this.embeddedTypes.set(type.fqn, embeddedType);
 
-    return this.resolveEmbeddedType(type);
+    return embeddedType;
   }
 
   protected emitHeader(code: CodeMaker) {
@@ -272,9 +273,7 @@ export abstract class Package {
   }
 
   private emitInternal(context: EmitContext) {
-    const aliases = Object.values(this.embeddedTypes);
-
-    if (aliases.length === 0) {
+    if (this.embeddedTypes.size === 0) {
       return;
     }
 
@@ -287,7 +286,7 @@ export abstract class Package {
 
     const imports = new Set<string>();
 
-    for (const alias of aliases) {
+    for (const alias of this.embeddedTypes.values()) {
       if (!alias.foriegnType) {
         continue;
       }
@@ -303,7 +302,7 @@ export abstract class Package {
     }
     code.close(')');
 
-    for (const alias of aliases) {
+    for (const alias of this.embeddedTypes.values()) {
       code.line(`type ${alias.fieldName} = ${alias.foriegnTypeName}`);
     }
 
@@ -322,7 +321,13 @@ export class RootPackage extends Package {
   private readonly readme?: ReadmeFile;
   private readonly versionFile: VersionFile;
 
-  public constructor(assembly: Assembly) {
+  // This cache of root packages is shared across all root packages derived created by this one (via dependencies).
+  private readonly rootPackageCache: Map<string, RootPackage>;
+
+  public constructor(
+    assembly: Assembly,
+    rootPackageCache = new Map<string, RootPackage>(),
+  ) {
     const goConfig = assembly.targets?.go ?? {};
     const packageName = goPackageNameForAssembly(assembly);
     const filePath = '';
@@ -338,6 +343,9 @@ export class RootPackage extends Package {
       version,
     );
 
+    this.rootPackageCache = rootPackageCache;
+    this.rootPackageCache.set(assembly.name, this);
+
     this.assembly = assembly;
     this.version = version;
     this.versionFile = new VersionFile(this.version);
@@ -423,7 +431,9 @@ export class RootPackage extends Package {
    */
   public get packageDependencies(): RootPackage[] {
     return this.assembly.dependencies.map(
-      (dep) => new RootPackage(dep.assembly),
+      (dep) =>
+        this.rootPackageCache.get(dep.assembly.name) ??
+        new RootPackage(dep.assembly, this.rootPackageCache),
     );
   }
 

packages/jsii-pacmak/lib/targets/go/runtime/emit-arguments.ts

@@ -19,7 +19,7 @@ export function emitArguments(
   if (argsList.length === 0) {
     return undefined;
   }
-  if (parameters[parameters.length - 1].parameter.variadic) {
+  if (parameters[parameters.length - 1].isVariadic) {
     // For variadic methods, we must build up the []interface{} slice by hand,
     // as there would not be any implicit conversion happening when passing
     // the variadic argument as a splat to the append function...

packages/jsii-pacmak/lib/targets/go/types/class.ts

@@ -23,7 +23,7 @@ import { GoMethod, GoProperty, GoTypeMember } from './type-member';
 /*
  * GoClass wraps a Typescript class as a Go custom struct type
  */
-export class GoClass extends GoType {
+export class GoClass extends GoType<ClassType> {
   public readonly methods: ClassMethod[];
   public readonly staticMethods: StaticMethod[];
   public readonly properties: GoProperty[];
@@ -34,7 +34,7 @@ export class GoClass extends GoType {
 
   private readonly initializer?: GoClassConstructor;
 
-  public constructor(pkg: Package, public type: ClassType) {
+  public constructor(pkg: Package, type: ClassType) {
     super(pkg, type);
 
     const methods = new Array<ClassMethod>();

packages/jsii-pacmak/lib/targets/go/types/enum.ts

@@ -7,10 +7,10 @@ import { Package } from '../package';
 import { JSII_RT_ALIAS } from '../runtime';
 import { GoType } from './go-type';
 
-export class Enum extends GoType {
+export class Enum extends GoType<EnumType> {
   private readonly members: readonly GoEnumMember[];
 
-  public constructor(pkg: Package, public type: EnumType) {
+  public constructor(pkg: Package, type: EnumType) {
     super(pkg, type);
 
     this.members = type.members.map((mem) => new GoEnumMember(this, mem));

packages/jsii-pacmak/lib/targets/go/types/go-type-reference.ts

@@ -47,6 +47,7 @@ type TypeMap =
  */
 export class GoTypeRef {
   private _typeMap?: TypeMap;
+
   public constructor(
     public readonly root: Package,
     public readonly reference: TypeReference,
@@ -126,18 +127,18 @@ export class GoTypeRef {
     switch (this.typeMap.type) {
       case 'interface':
         if (this.type?.pkg) {
-          ret.push(this.type?.pkg);
+          ret.push(this.type.pkg);
         }
         break;
 
       case 'array':
       case 'map':
-        ret.push(...(this.typeMap.value.dependencies ?? []));
+        ret.push(...this.typeMap.value.dependencies);
         break;
 
       case 'union':
         for (const t of this.typeMap.value) {
-          ret.push(...(t.dependencies ?? []));
+          ret.push(...t.dependencies);
         }
         break;
 

packages/jsii-pacmak/lib/targets/go/types/go-type.ts

@@ -8,12 +8,12 @@ import { JSII_RT_ALIAS } from '../runtime';
 import { GoClass } from './class';
 import { GoInterface } from './interface';
 
-export abstract class GoType {
+export abstract class GoType<T extends Type = Type> {
   public readonly name: string;
   public readonly fqn: string;
   public readonly proxyName: string;
 
-  public constructor(public pkg: Package, public type: Type) {
+  public constructor(public readonly pkg: Package, public readonly type: T) {
     this.name = type.name;
 
     // Prefix with the nesting parent name(s), using an _ delimiter.

packages/jsii-pacmak/lib/targets/go/types/interface.ts

@@ -11,13 +11,13 @@ import { GoType } from './go-type';
 import { GoTypeRef } from './go-type-reference';
 import { GoMethod, GoProperty } from './type-member';
 
-export class GoInterface extends GoType {
+export class GoInterface extends GoType<InterfaceType> {
   public readonly methods: InterfaceMethod[];
   public readonly reimplementedMethods?: readonly InterfaceMethod[];
   public readonly properties: InterfaceProperty[];
   public readonly reimplementedProperties?: readonly InterfaceProperty[];
 
-  public constructor(pkg: Package, public type: InterfaceType) {
+  public constructor(pkg: Package, type: InterfaceType) {
     super(pkg, type);
 
     this.methods = type.ownMethods
@@ -191,8 +191,6 @@ export class GoInterface extends GoType {
 }
 
 class InterfaceProperty extends GoProperty {
-  public readonly reference?: GoTypeRef;
-
   public constructor(
     public readonly parent: GoInterface,
     public readonly property: Property,
@@ -201,10 +199,7 @@ class InterfaceProperty extends GoProperty {
   }
 
   public get returnType(): string {
-    return (
-      this.reference?.scopedReference(this.parent.pkg) ??
-      this.property.type.toString()
-    );
+    return this.reference.scopedReference(this.parent.pkg);
   }
 
   public emit({ code, documenter }: EmitContext) {

packages/jsii-pacmak/lib/targets/go/types/struct.ts

@@ -13,10 +13,10 @@ import { GoProperty } from './type-member';
 /*
  * Struct wraps a JSII datatype interface aka, structs
  */
-export class Struct extends GoType {
+export class Struct extends GoType<InterfaceType> {
   private readonly properties: readonly GoProperty[];
 
-  public constructor(parent: Package, public readonly type: InterfaceType) {
+  public constructor(parent: Package, type: InterfaceType) {
     super(parent, type);
 
     assert(

packages/jsii-pacmak/lib/targets/go/types/type-member.ts

@@ -1,8 +1,15 @@
-import { Callable, Method, Parameter, Property } from 'jsii-reflect';
+import {
+  Callable,
+  Method,
+  Parameter,
+  Property,
+  TypeReference,
+} from 'jsii-reflect';
 
 import { jsiiToPascalCase } from '../../../naming-util';
 import { SpecialDependencies } from '../dependencies';
 import { EmitContext } from '../emit-context';
+import { Package } from '../package';
 import { GetProperty, JSII_RT_ALIAS, SetProperty } from '../runtime';
 import { substituteReservedWords } from '../util';
 
@@ -26,7 +33,6 @@ export interface GoTypeMember {
 */
 export class GoProperty implements GoTypeMember {
   public readonly name: string;
-  public readonly reference?: GoTypeRef;
   public readonly immutable: boolean;
 
   public constructor(
@@ -35,10 +41,10 @@ export class GoProperty implements GoTypeMember {
   ) {
     this.name = jsiiToPascalCase(this.property.name);
     this.immutable = property.immutable;
+  }
 
-    if (property.type) {
-      this.reference = new GoTypeRef(parent.pkg.root, property.type);
-    }
+  public get reference(): GoTypeRef {
+    return new GoTypeRef(this.parent.pkg.root, this.property.type);
   }
 
   public get specialDependencies(): SpecialDependencies {
@@ -143,17 +149,13 @@ export class GoProperty implements GoTypeMember {
 
 export abstract class GoMethod implements GoTypeMember {
   public readonly name: string;
-  public readonly reference?: GoTypeRef;
   public readonly parameters: GoParameter[];
 
   public constructor(
     public readonly parent: GoClass | GoInterface,
     public readonly method: Callable,
   ) {
     this.name = jsiiToPascalCase(method.name);
-    if (Method.isMethod(method) && method.returns.type) {
-      this.reference = new GoTypeRef(parent.pkg.root, method.returns.type);
-    }
     this.parameters = this.method.parameters.map(
       (param) => new GoParameter(parent, param),
     );
@@ -163,6 +165,13 @@ export abstract class GoMethod implements GoTypeMember {
 
   public abstract get specialDependencies(): SpecialDependencies;
 
+  public get reference(): GoTypeRef | undefined {
+    if (Method.isMethod(this.method) && this.method.returns.type) {
+      return new GoTypeRef(this.parent.pkg.root, this.method.returns.type);
+    }
+    return undefined;
+  }
+
   public get returnsRef(): boolean {
     if (
       this.reference?.type?.type.isClassType() ||
@@ -197,18 +206,23 @@ export abstract class GoMethod implements GoTypeMember {
 
 export class GoParameter {
   public readonly name: string;
-  public readonly reference: GoTypeRef;
+  public readonly isVariadic: boolean;
+  private readonly type: TypeReference;
+  private readonly pkg: Package;
 
-  public constructor(
-    public parent: GoClass | GoInterface,
-    public readonly parameter: Parameter,
-  ) {
+  public constructor(parent: GoClass | GoInterface, parameter: Parameter) {
     this.name = substituteReservedWords(parameter.name);
-    this.reference = new GoTypeRef(parent.pkg.root, parameter.type);
+    this.isVariadic = parameter.variadic;
+    this.type = parameter.type;
+    this.pkg = parent.pkg;
+  }
+
+  public get reference(): GoTypeRef {
+    return new GoTypeRef(this.pkg.root, this.type);
   }
 
   public toString(): string {
-    const paramType = this.reference.scopedReference(this.parent.pkg);
-    return `${this.name} ${this.parameter.variadic ? '...' : ''}${paramType}`;
+    const paramType = this.reference.scopedReference(this.pkg);
+    return `${this.name} ${this.isVariadic ? '...' : ''}${paramType}`;
   }
 }