Reuse source.Opener and ir.Session to fix executor memory leaks in LSP (#4230)

Author: doriable

private/buf/buflsp/buflsp.go

@@ -30,6 +30,8 @@ import (
 	"github.com/bufbuild/buf/private/pkg/storage"
 	"github.com/bufbuild/buf/private/pkg/wasm"
 	"github.com/bufbuild/protocompile/experimental/incremental"
+	"github.com/bufbuild/protocompile/experimental/ir"
+	"github.com/bufbuild/protocompile/experimental/source"
 	"go.lsp.dev/jsonrpc2"
 	"go.lsp.dev/protocol"
 	"go.uber.org/zap"
@@ -60,6 +62,8 @@ func Serve(
 		wasmRuntime:   wasmRuntime,
 		wktBucket:     wktBucket,
 		queryExecutor: queryExecutor,
+		opener:        source.NewMap(nil),
+		irSession:     new(ir.Session),
 	}
 	lsp.fileManager = newFileManager(lsp)
 	lsp.workspaceManager = newWorkspaceManager(lsp)
@@ -90,6 +94,8 @@ type lsp struct {
 	fileManager      *fileManager
 	workspaceManager *workspaceManager
 	queryExecutor    *incremental.Executor
+	opener           source.Map
+	irSession        *ir.Session
 	wktBucket        storage.ReadBucket
 	shutdown         bool
 

private/buf/buflsp/file.go

@@ -62,6 +62,8 @@ type file struct {
 	objectInfo storage.ObjectInfo // Info in the context of the workspace.
 
 	ir                   *ir.File
+	irQuery              queries.IR
+	fileQuery            queries.File
 	referenceableSymbols map[ir.FullName]*symbol
 	referenceSymbols     []*symbol
 	symbols              []*symbol
@@ -97,6 +99,11 @@ func (f *file) Reset(ctx context.Context) {
 		f.workspace.Release()
 		f.workspace = nil
 	}
+	// Evict the query key if there is a query cached on the file. We cache the [queries.File]
+	// query since this allows the executor to evict all dependent queries, e.g. AST and IR.
+	if f.fileQuery.Path != "" {
+		f.lsp.queryExecutor.Evict(f.fileQuery.Key())
+	}
 	// Clear the file as nothing should use it after reset. This asserts that.
 	*f = file{}
 }
@@ -229,21 +236,46 @@ func (f *file) RefreshIR(ctx context.Context) {
 	// Opener creates a cached view of all files in the workspace.
 	pathToFiles := f.workspace.PathToFile()
 	files := make([]*file, 0, len(pathToFiles))
-	openerMap := make(map[string]*source.File, len(pathToFiles))
+	var evictQueryKeys []any
+
+	openerMap := f.lsp.opener.Get()
 	for path, file := range pathToFiles {
-		openerMap[path] = file.file
+		current := openerMap[path]
+		// If there is no entry for the current path or if the file content has changed, we
+		// update the opener and set a new query.
+		if current == nil || current.Text() != file.file.Text() {
+			openerMap[path] = file.file
+			// Add the query key for eviction if there is a query cached on the file. We cache
+			// the [queries.File] query since this allows the executor to evict all dependent
+			// queries, e.g. AST and IR.
+			if file.fileQuery.Path != "" {
+				evictQueryKeys = append(evictQueryKeys, file.fileQuery.Key())
+			}
+			file.fileQuery = queries.File{
+				Opener:      file.lsp.opener,
+				Path:        path,
+				ReportError: true, // [queries.AST] sets this to be true.
+			}
+			file.irQuery = queries.IR{
+				Opener:  file.lsp.opener,
+				Path:    file.objectInfo.Path(),
+				Session: file.lsp.irSession,
+			}
+		}
 		files = append(files, file)
 	}
-	opener := source.NewMap(openerMap)
+	// Remove paths that are no longer in the current workspace and evict stale query keys.
+	for path := range openerMap {
+		if _, ok := pathToFiles[path]; !ok {
+			delete(openerMap, path)
+		}
+	}
+	f.lsp.queryExecutor.Evict(evictQueryKeys...)
 
-	session := new(ir.Session)
 	queries := xslices.Map(files, func(file *file) incremental.Query[*ir.File] {
-		return queries.IR{
-			Opener:  opener,
-			Path:    file.objectInfo.Path(),
-			Session: session,
-		}
+		return file.irQuery
 	})
+
 	results, diagnosticReport, err := incremental.Run(
 		ctx,
 		f.lsp.queryExecutor,

private/buf/buflsp/workspace.go

@@ -254,7 +254,7 @@ func (w *workspace) indexFiles(ctx context.Context) {
 		}
 
 		// Currently we only associate a file with one workspace. This assumption isn't accurate
-		// for shared dependencies. Here we update to the lastest, most recently used, workspace.
+		// for shared dependencies. Here we update to the latest, most recently used, workspace.
 		// This will make goto definition and find references only work in that workspace.
 		if oldWorkspace := file.workspace; oldWorkspace != nil && oldWorkspace != w {
 			oldWorkspace.Release()