Recycle type checkers under a configurable memory budget

packages/typescript/src/api/proto.generated.ts

@@ -1927,4 +1927,5 @@ export interface RawCompilerOptions {
     singleThreaded?: boolean | undefined;
     quiet?: boolean | undefined;
     checkers?: number | undefined;
+    maxMemory?: number | undefined;
 }

tools/scripts/tsc/options.ts

@@ -1626,6 +1626,17 @@ export const options: OptionsModel = {
                 minValue: 1,
             },
         },
+        {
+            name: "maxMemory",
+            type: "*int",
+            internal: true,
+            declaration: {
+                group: "commonOptionsWithBuild",
+                category: diagnostic("Command-line Options"),
+                description: diagnostic("Set the memory limit in MiB for recycling checkers during no-emit type checking."),
+                minValue: 1,
+            },
+        },
     ],
     // Preserve help and config output order independently of compiler fields.
     declarationOrder: {
@@ -1657,6 +1668,7 @@ export const options: OptionsModel = {
             "singleThreaded",
             "pprofDir",
             "checkers",
+            "maxMemory",
             "runExternalCode",
         ],
         optionsForCompiler: [

tsc/internal/compiler/checkerpool.go

@@ -3,9 +3,11 @@ package compiler
 import (
 	"context"
 	"math"
+	"runtime"
 	"slices"
 	"sort"
 	"sync"
+	"sync/atomic"
 
 	"github.com/microsoft/TypeScript/tsc/internal/ast"
 	"github.com/microsoft/TypeScript/tsc/internal/checker"
@@ -32,7 +34,16 @@ type checkerPool struct {
 	createCheckersOnce sync.Once
 	checkers           []*checker.Checker
 	locks              []*sync.Mutex
-	fileAssociations   map[*ast.SourceFile]*checker.Checker
+	fileAssociations   map[*ast.SourceFile]int
+
+	maxMemoryBytes              uint64
+	recycleMu                   sync.Mutex
+	retiredGlobalDiagnostics    []*ast.Diagnostic
+	retiredSymbolCount          atomic.Uint64
+	retiredTypeCount            atomic.Uint64
+	retiredInstantiationCount   atomic.Uint64
+	retiredCheckerCount         atomic.Uint64
+	readMemoryAllocationInBytes func() uint64
 }
 
 var _ CheckerPool = (*checkerPool)(nil)
@@ -321,6 +332,21 @@ func newCheckerPoolWithTracing(program *Program, tr *tracing.Tracing) *checkerPo
 		checkers: make([]*checker.Checker, checkerCount),
 		locks:    make([]*sync.Mutex, checkerCount),
 		tracing:  tr,
+		readMemoryAllocationInBytes: func() uint64 {
+			var stats runtime.MemStats
+			runtime.ReadMemStats(&stats)
+			return stats.Alloc
+		},
+	}
+	options := program.Options()
+	if maxMemory := options.MaxMemory; maxMemory != nil &&
+		*maxMemory > 0 &&
+		options.NoEmit.IsTrue() &&
+		!options.GetEmitDeclarations() &&
+		!options.IsIncremental() &&
+		options.GenerateTrace == "" &&
+		options.Watch.IsFalseOrUnknown() {
+		pool.maxMemoryBytes = uint64(*maxMemory) * 1024 * 1024
 	}
 
 	return pool
@@ -345,14 +371,14 @@ func (p *checkerPool) GetChecker(ctx context.Context, file *ast.SourceFile) (*ch
 // e.g. for read-only operations like obtaining an emit resolver.
 func (p *checkerPool) getCheckerForFileNonExclusive(file *ast.SourceFile) (*checker.Checker, func()) {
 	p.createCheckers()
-	return p.fileAssociations[file], noop
+	return p.checkers[p.fileAssociations[file]], noop
 }
 
 func (p *checkerPool) getCheckerForFileExclusive(ctx context.Context, file *ast.SourceFile) (*checker.Checker, func()) {
 	p.createCheckers()
-	c := p.fileAssociations[file]
-	idx := slices.Index(p.checkers, c)
+	idx := p.fileAssociations[file]
 	p.locks[idx].Lock()
+	c := p.checkers[idx]
 	return c, sync.OnceFunc(func() {
 		p.locks[idx].Unlock()
 	})
@@ -412,9 +438,9 @@ func (p *checkerPool) createCheckers() {
 			fileOrder := getCheckerAssociationOrder(fileWeights, isDeclarationFile, policy.prioritizeSourceFiles)
 			associations = getCheckerAssociationsInOrder(fileWeights, adjacentFiles, fileOrder, checkerCount, policy.balancePenaltyMultiplier)
 		}
-		p.fileAssociations = make(map[*ast.SourceFile]*checker.Checker, len(p.program.files))
+		p.fileAssociations = make(map[*ast.SourceFile]int, len(p.program.files))
 		for i, file := range p.program.files {
-			p.fileAssociations[file] = p.checkers[associations[i]]
+			p.fileAssociations[file] = associations[i]
 		}
 	})
 }
@@ -451,11 +477,11 @@ func (p *checkerPool) getImportAdjacency() [][]int {
 func (p *checkerPool) forEachCheckerParallel(cb func(idx int, c *checker.Checker)) {
 	p.createCheckers()
 	wg := core.NewWorkGroup(p.program.SingleThreaded())
-	for idx, checker := range p.checkers {
+	for idx := range p.checkers {
 		wg.Queue(func() {
 			p.locks[idx].Lock()
 			defer p.locks[idx].Unlock()
-			cb(idx, checker)
+			cb(idx, p.checkers[idx])
 		})
 	}
 	wg.RunAndWait()
@@ -467,7 +493,10 @@ func (p *checkerPool) GetGlobalDiagnostics() []*ast.Diagnostic {
 	p.forEachCheckerParallel(func(idx int, checker *checker.Checker) {
 		globalDiagnostics[idx] = checker.GetGlobalDiagnostics()
 	})
-	return SortAndDeduplicateDiagnostics(slices.Concat(globalDiagnostics...))
+	p.recycleMu.Lock()
+	retired := slices.Clone(p.retiredGlobalDiagnostics)
+	p.recycleMu.Unlock()
+	return SortAndDeduplicateDiagnostics(append(slices.Concat(globalDiagnostics...), retired...))
 }
 
 // forEachCheckerGroupDo runs one task per checker in parallel. Each task iterates
@@ -477,19 +506,57 @@ func (p *checkerPool) forEachCheckerGroupDo(ctx context.Context, files []*ast.So
 	p.createCheckers()
 
 	checkerCount := len(p.checkers)
+	remainingCheckable := make([]int, checkerCount)
+	for _, file := range files {
+		if !p.program.SkipTypeChecking(file, false) {
+			remainingCheckable[p.fileAssociations[file]]++
+		}
+	}
 	wg := core.NewWorkGroup(singleThreaded)
 	for checkerIdx := range checkerCount {
 		wg.Queue(func() {
 			p.locks[checkerIdx].Lock()
 			defer p.locks[checkerIdx].Unlock()
+			current := p.checkers[checkerIdx]
 			for i, file := range files {
-				if checker := p.checkers[checkerIdx]; checker == p.fileAssociations[file] {
-					cb(checker, i, file)
+				if checkerIdx == p.fileAssociations[file] {
+					cb(current, i, file)
+					if p.program.SkipTypeChecking(file, false) {
+						continue
+					}
+					remainingCheckable[checkerIdx]--
+					if remainingCheckable[checkerIdx] != 0 {
+						current = p.recycleCheckerIfOverBudget(checkerIdx, current)
+					}
 				}
 			}
 		})
 	}
 	wg.RunAndWait()
 }
 
+func (p *checkerPool) recycleCheckerIfOverBudget(checkerIdx int, current *checker.Checker) *checker.Checker {
+	if p.maxMemoryBytes == 0 || p.readMemoryAllocationInBytes() <= p.maxMemoryBytes {
+		return current
+	}
+
+	p.recycleMu.Lock()
+	defer p.recycleMu.Unlock()
+	if p.readMemoryAllocationInBytes() <= p.maxMemoryBytes {
+		return current
+	}
+
+	p.retiredGlobalDiagnostics = append(p.retiredGlobalDiagnostics, current.GetGlobalDiagnostics()...)
+	p.retiredSymbolCount.Add(uint64(current.SymbolCount))
+	p.retiredTypeCount.Add(uint64(current.TypeCount))
+	p.retiredInstantiationCount.Add(uint64(current.TotalInstantiationCount))
+	p.retiredCheckerCount.Add(1)
+
+	replacement, _ := checker.NewChecker(p.program, nil)
+	p.checkers[checkerIdx] = replacement
+	current = nil
+	runtime.GC()
+	return replacement
+}
+
 func noop() {}

tsc/internal/compiler/checkerpool_test.go

@@ -1,12 +1,16 @@
 package compiler
 
 import (
+	"context"
 	"slices"
 	"testing"
 
 	"github.com/microsoft/TypeScript/tsc/internal/ast"
+	"github.com/microsoft/TypeScript/tsc/internal/core"
 	"github.com/microsoft/TypeScript/tsc/internal/module"
+	"github.com/microsoft/TypeScript/tsc/internal/tsoptions"
 	"github.com/microsoft/TypeScript/tsc/internal/tspath"
+	"github.com/microsoft/TypeScript/tsc/internal/vfs/vfstest"
 )
 
 func TestGetSourceFileForResolvedModuleUsesResolvedPath(t *testing.T) {
@@ -120,6 +124,48 @@ func TestGetCheckerAssociationOrder(t *testing.T) {
 	}
 }
 
+func TestCheckerRecyclingPreservesDiagnosticsAndStatistics(t *testing.T) {
+	t.Parallel()
+	fs := vfstest.FromMap(map[string]any{
+		"/tsconfig.json": `{
+			"compilerOptions": { "noLib": true, "skipLibCheck": true },
+			"files": ["/a.ts", "/skip.d.ts", "/b.ts", "/c.ts"]
+		}`,
+		"/a.ts":      `const a: string = 1;`,
+		"/skip.d.ts": `declare const skipped: Missing;`,
+		"/b.ts":      `const b: number = "b";`,
+		"/c.ts":      `const c: boolean = 0;`,
+	}, tspath.CaseSensitive)
+	config, configDiagnostics := tsoptions.GetParsedCommandLineOfConfigFile("/tsconfig.json", nil, nil, fs, nil)
+	if len(configDiagnostics) != 0 {
+		t.Fatalf("unexpected config diagnostics: %v", configDiagnostics)
+	}
+	config.SetCompilerOptions(config.CompilerOptions().Clone())
+	config.CompilerOptions().NoEmit = core.TSTrue
+	config.CompilerOptions().Checkers = new(1)
+	config.CompilerOptions().MaxMemory = new(1)
+
+	program := NewProgram(ProgramOptions{
+		Config: config,
+		Host:   NewCompilerHost(fs, "/", nil, nil, nil),
+	})
+	pool := program.compilerCheckerPool
+	pool.readMemoryAllocationInBytes = func() uint64 {
+		return pool.maxMemoryBytes + 1
+	}
+
+	diagnostics := program.GetSemanticDiagnostics(context.Background(), nil)
+	if len(diagnostics) != 3 {
+		t.Fatalf("got %d diagnostics, want 3", len(diagnostics))
+	}
+	if got := program.RetiredCheckerCount(); got != 2 {
+		t.Fatalf("RetiredCheckerCount() = %d, want 2", got)
+	}
+	if program.TypeCount() == 0 || program.SymbolCount() == 0 {
+		t.Fatal("retired checker statistics were not preserved")
+	}
+}
+
 func TestGetCheckerAssociationWeights(t *testing.T) {
 	t.Parallel()
 

tsc/internal/compiler/program.go

@@ -1747,6 +1747,9 @@ func (p *Program) SymbolCount() int {
 	p.ForEachCheckerParallel(func(_ int, c *checker.Checker) {
 		val.Add(c.SymbolCount)
 	})
+	if p.compilerCheckerPool != nil {
+		val.Add(uint32(p.compilerCheckerPool.retiredSymbolCount.Load()))
+	}
 	return int(val.Load())
 }
 
@@ -1755,6 +1758,9 @@ func (p *Program) TypeCount() int {
 	p.ForEachCheckerParallel(func(_ int, c *checker.Checker) {
 		val.Add(c.TypeCount)
 	})
+	if p.compilerCheckerPool != nil {
+		val.Add(uint32(p.compilerCheckerPool.retiredTypeCount.Load()))
+	}
 	return int(val.Load())
 }
 
@@ -1763,9 +1769,19 @@ func (p *Program) InstantiationCount() int {
 	p.ForEachCheckerParallel(func(_ int, c *checker.Checker) {
 		val.Add(c.TotalInstantiationCount)
 	})
+	if p.compilerCheckerPool != nil {
+		val.Add(uint32(p.compilerCheckerPool.retiredInstantiationCount.Load()))
+	}
 	return int(val.Load())
 }
 
+func (p *Program) RetiredCheckerCount() int {
+	if p.compilerCheckerPool == nil {
+		return 0
+	}
+	return int(p.compilerCheckerPool.retiredCheckerCount.Load())
+}
+
 func (p *Program) Program() *Program {
 	return p
 }

tsc/internal/diagnostics/diagnosticMessages.generated.json

@@ -1695,6 +1695,7 @@
     "Set_the_number_of_projects_to_build_concurrently_6938": "Set the number of projects to build concurrently.",
     "Deduplicate_packages_with_the_same_name_and_version_6939": "Deduplicate packages with the same name and version.",
     "Allow_loading_external_content_mapper_plugins_that_execute_code_during_compilation_6940": "Allow loading external content mapper plugins that execute code during compilation.",
+    "Set_the_memory_limit_in_MiB_for_recycling_checkers_during_no_emit_type_checking_6941": "Set the memory limit in MiB for recycling checkers during no-emit type checking.",
     "Variable_0_implicitly_has_an_1_type_7005": "Variable '{0}' implicitly has an '{1}' type.",
     "Parameter_0_implicitly_has_an_1_type_7006": "Parameter '{0}' implicitly has an '{1}' type.",
     "Member_0_implicitly_has_an_1_type_7008": "Member '{0}' implicitly has an '{1}' type.",

tsc/internal/diagnostics/diagnosticMessages.json

@@ -6815,6 +6815,10 @@
         "category": "Message",
         "code": 6940
     },
+    "Set the memory limit in MiB for recycling checkers during no-emit type checking.": {
+        "category": "Message",
+        "code": 6941
+    },
 
     "Variable '{0}' implicitly has an '{1}' type.": {
         "category": "Error",