perf(lsp): resolve code lenses eagerly to fix ~25s document-open stall

internal/lsp/code_lens.go

@@ -3,6 +3,7 @@ package lsp
 import (
 	"context"
 	"fmt"
+	"path/filepath"
 
 	"go.lsp.dev/protocol"
 
@@ -22,50 +23,102 @@ func (s *Server) CodeLens(ctx context.Context, params *protocol.CodeLensParams)
 	if !ok {
 		return []protocol.CodeLens{}, nil
 	}
-
-	semanticDoc := newSemanticDocument(doc.Path, doc.Content, s.parsePathForNavigation(doc.Path))
+	// Check cancellation before the (potentially expensive) initial parse so a
+	// cancelled request never surfaces as a misleading empty success.
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	parse := s.newRequestParse(ctx)
+	rootResult := parse(doc.Path)
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+	semanticDoc := newSemanticDocument(doc.Path, doc.Content, rootResult)
 	if !semanticDoc.valid() {
 		return []protocol.CodeLens{}, nil
 	}
 
-	return semanticDoc.codeLenses(), nil
-}
-
-func (s *Server) CodeLensResolve(ctx context.Context, params *protocol.CodeLens) (*protocol.CodeLens, error) {
-	if params == nil || params.Command != nil {
-		return params, nil
+	// Resolve every lens up front in one workspace pass: the client then never
+	// calls CodeLensResolve, so there is no per-version command memo to go stale
+	// when a cross-file reference changes.
+	candidatePaths := s.referenceCandidatePaths()
+	resolveType := func(path string, result *ridl.ParseResult, name string) *definitionMatch {
+		return resolveTypeDefinitionWith(parse, path, result, name)
 	}
-
-	data, ok := decodeCodeLensData(params.Data)
-	if !ok {
-		return params, nil
+	resolveError := func(path string, result *ridl.ParseResult, name string) *definitionMatch {
+		return resolveErrorDefinitionWith(parse, path, result, name)
 	}
 
-	doc, ok := s.docs.Get(data.URI)
-	if !ok {
-		return params, nil
+	base := semanticDoc.codeLenses()
+	lenses := make([]protocol.CodeLens, 0, len(base))
+	for i := range base {
+		data, ok := decodeCodeLensData(base[i].Data)
+		if !ok {
+			continue
+		}
+		target := semanticDoc.referenceTargetAt(
+			protocol.Position{Line: data.Line, Character: data.Character},
+			resolveType, resolveError,
+		)
+		if target == nil || target.definition == nil {
+			continue
+		}
+		locations := collectReferenceLocationsWith(parse, candidatePaths, s.contentForPath, target, false, resolveType, resolveError)
+		// Check after the (potentially expensive) cross-file collection so a
+		// cancelled request — including on the last/only symbol — never returns
+		// a partial under-counted result to the client.
+		if err := ctx.Err(); err != nil {
+			return nil, err
+		}
+		lens := base[i]
+		lens.Command = &protocol.Command{
+			Title:     referenceCountTitle(len(locations)),
+			Command:   showReferencesCommand,
+			Arguments: []any{protocol.DocumentURI(data.URI), base[i].Range.Start, locations},
+		}
+		lens.Data = nil
+		lenses = append(lenses, lens)
 	}
 
-	semanticDoc := newSemanticDocument(doc.Path, doc.Content, s.parsePathForNavigation(doc.Path))
-	if !semanticDoc.valid() {
-		return params, nil
-	}
+	return lenses, nil
+}
 
-	target := semanticDoc.referenceTargetAt(protocol.Position{
-		Line:      data.Line,
-		Character: data.Character,
-	}, s.resolveTypeDefinition, s.resolveErrorDefinition)
-	if target == nil || target.definition == nil {
-		return params, nil
-	}
+// CodeLensResolve is a passthrough: CodeLens returns fully-resolved lenses, so
+// resolution never needs to recompute.
+func (s *Server) CodeLensResolve(ctx context.Context, params *protocol.CodeLens) (*protocol.CodeLens, error) {
+	return params, nil
+}
 
-	locations := s.collectReferenceLocations(target, false)
-	params.Command = &protocol.Command{
-		Title:     referenceCountTitle(len(locations)),
-		Command:   showReferencesCommand,
-		Arguments: []any{protocol.DocumentURI(data.URI), params.Range.Start, locations},
+// newRequestParse builds a request-scoped parser that AST-parses each file at
+// most once: an open buffer's already-built result is reused; everything else is
+// parsed AST-only (no schema build, no import recursion). Keyed by cleaned path
+// so each file is parsed once per request.
+//
+// The overlay map is snapshotted once here; s.docs.FindByPath and s.contentForPath
+// read live document state on each miss. Handlers run concurrently, so a request
+// is eventually-consistent-by-design: a mid-request DidChange self-heals on the
+// client's next codeLens request.
+func (s *Server) newRequestParse(ctx context.Context) parseFn {
+	overlays := s.overlayContents()
+	memo := map[string]*ridl.ParseResult{}
+	return func(path string) *ridl.ParseResult {
+		// Honour request cancellation before each (possibly expensive) AST parse.
+		if ctx.Err() != nil {
+			return nil
+		}
+		key := filepath.Clean(path)
+		if result, ok := memo[key]; ok {
+			return result
+		}
+		var result *ridl.ParseResult
+		if doc, ok := s.docs.FindByPath(path); ok && doc.Result != nil && doc.Result.Root != nil {
+			result = doc.Result
+		} else {
+			result, _ = s.parser.ParseAST(ctx, s.workspace.Root(), path, overlays)
+		}
+		memo[key] = result
+		return result
 	}
-	return params, nil
 }
 
 func (d *semanticDocument) codeLenses() []protocol.CodeLens {

internal/lsp/code_lens_perf_test.go

@@ -0,0 +1,149 @@
+package lsp
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"strconv"
+	"testing"
+
+	"go.lsp.dev/protocol"
+	"go.uber.org/zap"
+)
+
+const baseRIDL = `webrpc = v1
+
+name = test
+version = v0.0.1
+
+struct User
+  - id: uint64
+`
+
+// TestNewRequestParseMemoizesPerPath asserts that calling the request-scoped
+// parse function twice with the same path returns the identical *ParseResult
+// pointer, proving the file is parsed at most once per request.
+func TestNewRequestParseMemoizesPerPath(t *testing.T) {
+	srv, _, dir := setupServer(t)
+
+	basePath := filepath.Join(dir, "base.ridl")
+	if err := os.WriteFile(basePath, []byte(baseRIDL), 0o644); err != nil {
+		t.Fatal(err)
+	}
+
+	parse := srv.newRequestParse(context.Background())
+
+	r1 := parse(basePath)
+	if r1 == nil {
+		t.Fatal("expected non-nil ParseResult, got nil — test would be vacuous")
+	}
+	if r1.Root == nil {
+		t.Fatal("expected ParseResult.Root != nil, got nil — test would be vacuous")
+	}
+
+	r2 := parse(basePath)
+	if r1 != r2 {
+		t.Errorf("parse returned different pointers for the same path — file was parsed more than once")
+	}
+}
+
+// TestNewRequestParseKeyedByCleanPath asserts that path variants that resolve
+// to the same cleaned path (e.g. "dir/./file.ridl" vs "dir/file.ridl") share
+// the same memo slot and produce the identical *ParseResult pointer.
+func TestNewRequestParseKeyedByCleanPath(t *testing.T) {
+	srv, _, dir := setupServer(t)
+
+	basePath := filepath.Join(dir, "base.ridl")
+	if err := os.WriteFile(basePath, []byte(baseRIDL), 0o644); err != nil {
+		t.Fatal(err)
+	}
+
+	parse := srv.newRequestParse(context.Background())
+
+	// Two string paths that differ only by a redundant "." segment.
+	p1 := dir + "/base.ridl"
+	p2 := dir + "/./base.ridl" // deliberately un-cleaned via string concat
+
+	r1 := parse(p1)
+	if r1 == nil {
+		t.Fatal("expected non-nil ParseResult for p1, got nil — test would be vacuous")
+	}
+
+	r2 := parse(p2)
+	if r2 == nil {
+		t.Fatal("expected non-nil ParseResult for p2, got nil")
+	}
+	if r1 != r2 {
+		t.Errorf("different un-cleaned paths that resolve to the same file returned different ParseResult pointers — clean-path keying is broken")
+	}
+}
+
+// BenchmarkCodeLensEager measures the cost of a single CodeLens request
+// against a workspace of 40 importer files all referencing two symbols in
+// a shared base file. Run with -bench BenchmarkCodeLensEager -run x.
+func BenchmarkCodeLensEager(b *testing.B) {
+	// Replicate setupServer body for *testing.B (setupServer only accepts *testing.T).
+	dir := b.TempDir()
+	client := newMockClient()
+	srv := NewServer(zap.NewNop())
+	srv.SetClient(client)
+	srv.workspace.SetRoot(dir)
+
+	baseContent := `webrpc = v1
+
+name = base
+version = v0.0.1
+
+struct User
+  - id: uint64
+
+struct Account
+  - userId: uint64
+`
+	basePath := filepath.Join(dir, "base.ridl")
+	if err := os.WriteFile(basePath, []byte(baseContent), 0o644); err != nil {
+		b.Fatal(err)
+	}
+
+	// Write 40 importer files each referencing User and Account from base.
+	for i := 0; i < 40; i++ {
+		content := `webrpc = v1
+
+name = f` + strconv.Itoa(i) + `
+version = v0.0.1
+
+import
+  - base.ridl
+
+struct S` + strconv.Itoa(i) + `
+  - u: User
+  - a: Account
+`
+		p := filepath.Join(dir, "f"+strconv.Itoa(i)+".ridl")
+		if err := os.WriteFile(p, []byte(content), 0o644); err != nil {
+			b.Fatal(err)
+		}
+	}
+
+	// DidOpen base.ridl so it is the target document.
+	baseURI := string(PathToURI(basePath))
+	ctx := context.Background()
+	if err := srv.DidOpen(ctx, &protocol.DidOpenTextDocumentParams{
+		TextDocument: protocol.TextDocumentItem{
+			URI:     protocol.DocumentURI(baseURI),
+			Text:    baseContent,
+			Version: 1,
+		},
+	}); err != nil {
+		b.Fatal(err)
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		if _, err := srv.CodeLens(ctx, &protocol.CodeLensParams{
+			TextDocument: protocol.TextDocumentIdentifier{URI: protocol.DocumentURI(baseURI)},
+		}); err != nil {
+			b.Fatal(err)
+		}
+	}
+}