Concurrent LSP with cancelation (microsoft#869)

Author: andrewbranch

go.mod

@@ -7,14 +7,14 @@ require (
 	github.com/go-json-experiment/json v0.0.0-20250223041408-d3c622f1b874
 	github.com/google/go-cmp v0.7.0
 	github.com/pkg/diff v0.0.0-20241224192749-4e6772a4315c
+	golang.org/x/sync v0.11.0
 	golang.org/x/sys v0.31.0
 	gotest.tools/v3 v3.5.2
 )
 
 require (
 	github.com/matryer/moq v0.5.3 // indirect
 	golang.org/x/mod v0.23.0 // indirect
-	golang.org/x/sync v0.11.0 // indirect
 	golang.org/x/tools v0.30.0 // indirect
 )
 

internal/api/api.go

@@ -1,6 +1,7 @@
 package api
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -129,7 +130,7 @@ func (api *API) IsWatchEnabled() bool {
 	return false
 }
 
-func (api *API) HandleRequest(id int, method string, payload []byte) ([]byte, error) {
+func (api *API) HandleRequest(ctx context.Context, method string, payload []byte) ([]byte, error) {
 	params, err := unmarshalPayload(method, payload)
 	if err != nil {
 		return nil, err
@@ -155,27 +156,27 @@ func (api *API) HandleRequest(id int, method string, payload []byte) ([]byte, er
 		return encodeJSON(api.LoadProject(params.(*LoadProjectParams).ConfigFileName))
 	case MethodGetSymbolAtPosition:
 		params := params.(*GetSymbolAtPositionParams)
-		return encodeJSON(api.GetSymbolAtPosition(params.Project, params.FileName, int(params.Position)))
+		return encodeJSON(api.GetSymbolAtPosition(ctx, params.Project, params.FileName, int(params.Position)))
 	case MethodGetSymbolsAtPositions:
 		params := params.(*GetSymbolsAtPositionsParams)
 		return encodeJSON(core.TryMap(params.Positions, func(position uint32) (any, error) {
-			return api.GetSymbolAtPosition(params.Project, params.FileName, int(position))
+			return api.GetSymbolAtPosition(ctx, params.Project, params.FileName, int(position))
 		}))
 	case MethodGetSymbolAtLocation:
 		params := params.(*GetSymbolAtLocationParams)
-		return encodeJSON(api.GetSymbolAtLocation(params.Project, params.Location))
+		return encodeJSON(api.GetSymbolAtLocation(ctx, params.Project, params.Location))
 	case MethodGetSymbolsAtLocations:
 		params := params.(*GetSymbolsAtLocationsParams)
 		return encodeJSON(core.TryMap(params.Locations, func(location Handle[ast.Node]) (any, error) {
-			return api.GetSymbolAtLocation(params.Project, location)
+			return api.GetSymbolAtLocation(ctx, params.Project, location)
 		}))
 	case MethodGetTypeOfSymbol:
 		params := params.(*GetTypeOfSymbolParams)
-		return encodeJSON(api.GetTypeOfSymbol(params.Project, params.Symbol))
+		return encodeJSON(api.GetTypeOfSymbol(ctx, params.Project, params.Symbol))
 	case MethodGetTypesOfSymbols:
 		params := params.(*GetTypesOfSymbolsParams)
 		return encodeJSON(core.TryMap(params.Symbols, func(symbol Handle[ast.Symbol]) (any, error) {
-			return api.GetTypeOfSymbol(params.Project, symbol)
+			return api.GetTypeOfSymbol(ctx, params.Project, symbol)
 		}))
 	default:
 		return nil, fmt.Errorf("unhandled API method %q", method)
@@ -223,12 +224,14 @@ func (api *API) LoadProject(configFileName string) (*ProjectResponse, error) {
 	return data, nil
 }
 
-func (api *API) GetSymbolAtPosition(projectId Handle[project.Project], fileName string, position int) (*SymbolResponse, error) {
+func (api *API) GetSymbolAtPosition(ctx context.Context, projectId Handle[project.Project], fileName string, position int) (*SymbolResponse, error) {
 	project, ok := api.projects[projectId]
 	if !ok {
 		return nil, errors.New("project not found")
 	}
-	symbol, err := project.LanguageService().GetSymbolAtPosition(fileName, position)
+	languageService, done := project.GetLanguageServiceForRequest(ctx)
+	defer done()
+	symbol, err := languageService.GetSymbolAtPosition(ctx, fileName, position)
 	if err != nil || symbol == nil {
 		return nil, err
 	}
@@ -239,7 +242,7 @@ func (api *API) GetSymbolAtPosition(projectId Handle[project.Project], fileName
 	return data, nil
 }
 
-func (api *API) GetSymbolAtLocation(projectId Handle[project.Project], location Handle[ast.Node]) (*SymbolResponse, error) {
+func (api *API) GetSymbolAtLocation(ctx context.Context, projectId Handle[project.Project], location Handle[ast.Node]) (*SymbolResponse, error) {
 	project, ok := api.projects[projectId]
 	if !ok {
 		return nil, errors.New("project not found")
@@ -262,7 +265,9 @@ func (api *API) GetSymbolAtLocation(projectId Handle[project.Project], location
 	if node == nil {
 		return nil, fmt.Errorf("node of kind %s not found at position %d in file %q", kind.String(), pos, sourceFile.FileName())
 	}
-	symbol := project.LanguageService().GetSymbolAtLocation(node)
+	languageService, done := project.GetLanguageServiceForRequest(ctx)
+	defer done()
+	symbol := languageService.GetSymbolAtLocation(ctx, node)
 	if symbol == nil {
 		return nil, nil
 	}
@@ -273,7 +278,7 @@ func (api *API) GetSymbolAtLocation(projectId Handle[project.Project], location
 	return data, nil
 }
 
-func (api *API) GetTypeOfSymbol(projectId Handle[project.Project], symbolHandle Handle[ast.Symbol]) (*TypeResponse, error) {
+func (api *API) GetTypeOfSymbol(ctx context.Context, projectId Handle[project.Project], symbolHandle Handle[ast.Symbol]) (*TypeResponse, error) {
 	project, ok := api.projects[projectId]
 	if !ok {
 		return nil, errors.New("project not found")
@@ -284,7 +289,9 @@ func (api *API) GetTypeOfSymbol(projectId Handle[project.Project], symbolHandle
 	if !ok {
 		return nil, fmt.Errorf("symbol %q not found", symbolHandle)
 	}
-	t := project.LanguageService().GetTypeOfSymbol(symbol)
+	languageService, done := project.GetLanguageServiceForRequest(ctx)
+	defer done()
+	t := languageService.GetTypeOfSymbol(ctx, symbol)
 	if t == nil {
 		return nil, nil
 	}

internal/api/server.go

@@ -2,13 +2,16 @@ package api
 
 import (
 	"bufio"
+	"context"
 	"encoding/binary"
 	"encoding/json"
 	"fmt"
 	"io"
+	"strconv"
 	"sync"
 
 	"github.com/microsoft/typescript-go/internal/bundled"
+	"github.com/microsoft/typescript-go/internal/core"
 	"github.com/microsoft/typescript-go/internal/project"
 	"github.com/microsoft/typescript-go/internal/vfs"
 	"github.com/microsoft/typescript-go/internal/vfs/osvfs"
@@ -254,7 +257,7 @@ func (s *Server) handleRequest(method string, payload []byte) ([]byte, error) {
 	case "echo":
 		return payload, nil
 	default:
-		return s.api.HandleRequest(s.requestId, method, payload)
+		return s.api.HandleRequest(core.WithRequestID(context.Background(), strconv.Itoa(s.requestId)), method, payload)
 	}
 }
 

internal/checker/checker_test.go

@@ -39,7 +39,8 @@ foo.bar;`
 	}
 	p := compiler.NewProgram(opts)
 	p.BindSourceFiles()
-	c := p.GetTypeChecker()
+	c, done := p.GetTypeChecker(t.Context())
+	defer done()
 	file := p.GetSourceFile("/foo.ts")
 	interfaceId := file.Statements.Nodes[0].Name()
 	varId := file.Statements.Nodes[1].AsVariableStatement().DeclarationList.AsVariableDeclarationList().Declarations.Nodes[0].Name()

internal/checker/exports.go

@@ -48,3 +48,7 @@ func (c *Checker) GetTypeOfPropertyOfContextualType(t *Type, name string) *Type
 func GetDeclarationModifierFlagsFromSymbol(s *ast.Symbol) ast.ModifierFlags {
 	return getDeclarationModifierFlagsFromSymbol(s)
 }
+
+func (c *Checker) WasCanceled() bool {
+	return c.wasCanceled
+}

internal/compiler/checkerpool.go

@@ -0,0 +1,90 @@
+package compiler
+
+import (
+	"context"
+	"iter"
+	"slices"
+	"sync"
+
+	"github.com/microsoft/typescript-go/internal/ast"
+	"github.com/microsoft/typescript-go/internal/checker"
+	"github.com/microsoft/typescript-go/internal/core"
+)
+
+type CheckerPool interface {
+	GetChecker(ctx context.Context) (*checker.Checker, func())
+	GetCheckerForFile(ctx context.Context, file *ast.SourceFile) (*checker.Checker, func())
+	GetAllCheckers(ctx context.Context) ([]*checker.Checker, func())
+	Files(checker *checker.Checker) iter.Seq[*ast.SourceFile]
+}
+
+type checkerPool struct {
+	checkerCount int
+	program      *Program
+
+	createCheckersOnce sync.Once
+	checkers           []*checker.Checker
+	fileAssociations   map[*ast.SourceFile]*checker.Checker
+}
+
+var _ CheckerPool = (*checkerPool)(nil)
+
+func newCheckerPool(checkerCount int, program *Program) *checkerPool {
+	pool := &checkerPool{
+		program:      program,
+		checkerCount: checkerCount,
+		checkers:     make([]*checker.Checker, checkerCount),
+	}
+
+	return pool
+}
+
+func (p *checkerPool) GetCheckerForFile(ctx context.Context, file *ast.SourceFile) (*checker.Checker, func()) {
+	p.createCheckers()
+	checker := p.fileAssociations[file]
+	return checker, noop
+}
+
+func (p *checkerPool) GetChecker(ctx context.Context) (*checker.Checker, func()) {
+	p.createCheckers()
+	checker := p.checkers[0]
+	return checker, noop
+}
+
+func (p *checkerPool) createCheckers() {
+	p.createCheckersOnce.Do(func() {
+		wg := core.NewWorkGroup(p.program.singleThreaded())
+		for i := range p.checkerCount {
+			wg.Queue(func() {
+				p.checkers[i] = checker.NewChecker(p.program)
+			})
+		}
+
+		wg.RunAndWait()
+
+		p.fileAssociations = make(map[*ast.SourceFile]*checker.Checker, len(p.program.files))
+		for i, file := range p.program.files {
+			p.fileAssociations[file] = p.checkers[i%p.checkerCount]
+		}
+	})
+}
+
+func (p *checkerPool) GetAllCheckers(ctx context.Context) ([]*checker.Checker, func()) {
+	p.createCheckers()
+	return p.checkers, noop
+}
+
+func (p *checkerPool) Files(checker *checker.Checker) iter.Seq[*ast.SourceFile] {
+	checkerIndex := slices.Index(p.checkers, checker)
+	return func(yield func(*ast.SourceFile) bool) {
+		for i, file := range p.program.files {
+			if i%p.checkerCount == checkerIndex {
+				if !yield(file) {
+					return
+				}
+			}
+		}
+	}
+}
+
+func noop() {}

internal/compiler/emitHost.go

@@ -1,6 +1,8 @@
 package compiler
 
 import (
+	"context"
+
 	"github.com/microsoft/typescript-go/internal/ast"
 	"github.com/microsoft/typescript-go/internal/core"
 	"github.com/microsoft/typescript-go/internal/printer"
@@ -32,7 +34,11 @@ func (host *emitHost) WriteFile(fileName string, text string, writeByteOrderMark
 }
 
 func (host *emitHost) GetEmitResolver(file *ast.SourceFile, skipDiagnostics bool) printer.EmitResolver {
-	checker := host.program.GetTypeCheckerForFile(file)
+	// The context and done function don't matter in tsc, currently the only caller of this function.
+	// But if this ever gets used by LSP code, we'll need to thread the context properly and pass the
+	// done function to the caller to ensure resources are cleaned up at the end of the request.
+	checker, done := host.program.GetTypeCheckerForFile(context.TODO(), file)
+	defer done()
 	return checker.GetEmitResolver(file, skipDiagnostics)
 }