From 71e170ed65cd4ff5f5863c9d2b6027eb0a5b9a51 Mon Sep 17 00:00:00 2001
From: Miguel Young <mcyoungsota@gmail.com>
Date: Fri, 11 Oct 2024 11:29:44 -0400
Subject: [PATCH] Temporarily serialize requests in the LSP, and fix the bugs
 that doing so revealed (#3370)

---
 private/buf/bufcli/cache.go                 |   4 +-
 private/buf/bufcli/controller.go            |   2 +-
 private/buf/buflsp/buflsp.go                |  39 +--
 private/buf/buflsp/file.go                  | 249 ++++++++++++--------
 private/buf/buflsp/file_manager.go          |   2 -
 private/buf/buflsp/mutex.go                 | 182 --------------
 private/buf/buflsp/report.go                |   7 +
 private/buf/buflsp/server.go                |  25 +-
 private/buf/buflsp/symbol.go                | 154 ++++++++----
 private/buf/cmd/buf/command/beta/lsp/lsp.go |  11 +-
 10 files changed, 295 insertions(+), 380 deletions(-)
 delete mode 100644 private/buf/buflsp/mutex.go

diff --git a/private/buf/bufcli/cache.go b/private/buf/bufcli/cache.go
index 568e1104f0..6b42f628fa 100644
--- a/private/buf/bufcli/cache.go
+++ b/private/buf/bufcli/cache.go
@@ -153,8 +153,8 @@ func CreateWasmRuntimeCacheDir(container appext.Container) (string, error) {
 	return fullCacheDirPath, nil
 }
 
-// newWKTStore returns a new bufwktstore.Store while creating the required cache directories.
-func newWKTStore(container appext.Container) (bufwktstore.Store, error) {
+// NewWKTStore returns a new bufwktstore.Store while creating the required cache directories.
+func NewWKTStore(container appext.Container) (bufwktstore.Store, error) {
 	if err := createCacheDir(container.CacheDirPath(), v3CacheWKTRelDirPath); err != nil {
 		return nil, err
 	}
diff --git a/private/buf/bufcli/controller.go b/private/buf/bufcli/controller.go
index 05ae2048f2..d4939a4e37 100644
--- a/private/buf/bufcli/controller.go
+++ b/private/buf/bufcli/controller.go
@@ -45,7 +45,7 @@ func NewController(
 	if err != nil {
 		return nil, err
 	}
-	wktStore, err := newWKTStore(container)
+	wktStore, err := NewWKTStore(container)
 	if err != nil {
 		return nil, err
 	}
diff --git a/private/buf/buflsp/buflsp.go b/private/buf/buflsp/buflsp.go
index 9ddb3121c0..9d04d840e3 100644
--- a/private/buf/buflsp/buflsp.go
+++ b/private/buf/buflsp/buflsp.go
@@ -21,11 +21,11 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"sync"
 	"sync/atomic"
 
 	"github.com/bufbuild/buf/private/buf/bufctl"
 	"github.com/bufbuild/buf/private/bufpkg/bufcheck"
-	"github.com/bufbuild/buf/private/bufpkg/bufimage"
 	"github.com/bufbuild/buf/private/pkg/app/appext"
 	"github.com/bufbuild/buf/private/pkg/slogext"
 	"github.com/bufbuild/buf/private/pkg/storage"
@@ -40,6 +40,7 @@ import (
 // Returns a context for managing the server.
 func Serve(
 	ctx context.Context,
+	wktBucket storage.ReadBucket,
 	container appext.Container,
 	controller bufctl.Controller,
 	checkClient bufcheck.Client,
@@ -68,6 +69,7 @@ func Serve(
 		controller:  controller,
 		checkClient: checkClient,
 		rootBucket:  bucket,
+		wktBucket:   wktBucket,
 	}
 	lsp.fileManager = newFileManager(lsp)
 	off := protocol.TraceOff
@@ -96,6 +98,10 @@ type lsp struct {
 	rootBucket  storage.ReadBucket
 	fileManager *fileManager
 
+	wktBucket storage.ReadBucket
+
+	lock sync.Mutex
+
 	// These are atomics, because they are read often and written to
 	// almost never, but potentially concurrently. Having them side-by-side
 	// is fine; they are almost never written to so false sharing is not a
@@ -107,7 +113,7 @@ type lsp struct {
 // init performs *actual* initialization of the server. This is called by Initialize().
 //
 // It may only be called once for a given server.
-func (l *lsp) init(params *protocol.InitializeParams) error {
+func (l *lsp) init(_ context.Context, params *protocol.InitializeParams) error {
 	if l.initParams.Load() != nil {
 		return fmt.Errorf("called the %q method more than once", protocol.MethodInitialize)
 	}
@@ -120,31 +126,6 @@ func (l *lsp) init(params *protocol.InitializeParams) error {
 	return nil
 }
 
-// findImportable finds all files that can potentially be imported by the proto file at
-// uri. This returns a map from potential Protobuf import path to the URI of the file it would import.
-//
-// Note that this performs no validation on these files, because those files might be open in the
-// editor and might contain invalid syntax at the moment. We only want to get their paths and nothing
-// more.
-func (l *lsp) findImportable(
-	ctx context.Context,
-	uri protocol.URI,
-) (map[string]bufimage.ImageFileInfo, error) {
-	fileInfos, err := l.controller.GetImportableImageFileInfos(ctx, uri.Filename())
-	if err != nil {
-		return nil, err
-	}
-
-	imports := make(map[string]bufimage.ImageFileInfo)
-	for _, fileInfo := range fileInfos {
-		imports[fileInfo.Path()] = fileInfo
-	}
-
-	l.logger.DebugContext(ctx, fmt.Sprintf("found imports for %q: %#v", uri, imports))
-
-	return imports, nil
-}
-
 // newHandler constructs an RPC handler that wraps the default one from jsonrpc2. This allows us
 // to inject debug logging, tracing, and timeouts to requests.
 func (l *lsp) newHandler() jsonrpc2.Handler {
@@ -152,8 +133,6 @@ func (l *lsp) newHandler() jsonrpc2.Handler {
 	return func(ctx context.Context, reply jsonrpc2.Replier, req jsonrpc2.Request) (retErr error) {
 		defer slogext.DebugProfile(l.logger, slog.String("method", req.Method()), slog.Any("params", req.Params()))()
 
-		ctx = withRequestID(ctx)
-
 		replier := l.wrapReplier(reply, req)
 
 		// Verify that the server has been initialized if this isn't the initialization
@@ -162,6 +141,8 @@ func (l *lsp) newHandler() jsonrpc2.Handler {
 			return replier(ctx, nil, fmt.Errorf("the first call to the server must be the %q method", protocol.MethodInitialize))
 		}
 
+		l.lock.Lock()
+		defer l.lock.Unlock()
 		return actual(ctx, replier, req)
 	}
 }
diff --git a/private/buf/buflsp/file.go b/private/buf/buflsp/file.go
index f1adae9d6e..c835f82a23 100644
--- a/private/buf/buflsp/file.go
+++ b/private/buf/buflsp/file.go
@@ -33,6 +33,7 @@ import (
 	"github.com/bufbuild/buf/private/bufpkg/bufmodule"
 	"github.com/bufbuild/buf/private/pkg/ioext"
 	"github.com/bufbuild/buf/private/pkg/slogext"
+	"github.com/bufbuild/buf/private/pkg/storage"
 	"github.com/bufbuild/protocompile"
 	"github.com/bufbuild/protocompile/ast"
 	"github.com/bufbuild/protocompile/linker"
@@ -50,32 +51,9 @@ const descriptorPath = "google/protobuf/descriptor.proto"
 //
 // Mutating a file is thread-safe.
 type file struct {
-	// lsp and uri are not protected by file.lock; they are immutable after
-	// file creation!
 	lsp *lsp
 	uri protocol.URI
 
-	// All variables after this lock variables are protected by file.lock.
-	//
-	// NOTE: this package must NEVER attempt to acquire a lock on a file while
-	// holding a lock on another file. This guarantees that any concurrent operations
-	// on distinct files can always make forward progress, even if the information they
-	// have is incomplete. This trades off up-to-date accuracy for responsiveness.
-	//
-	// For example, suppose g1 locks a.proto, and then attempts to lock b.proto
-	// because it followed a pointer in importMap. However, in the meantime, g2
-	// has acquired b.proto's lock already, and attempts to acquire a lock to a.proto,
-	// again because of a pointer in importMap. This will deadlock, and it will
-	// deadlock in such a way that will be undetectable to the Go scheduler, so the
-	// LSP will hang forever.
-	//
-	// This seems like a contrived scenario, but it can happen if a user creates two
-	// mutually-recursive Protobuf files. Although this is not permitted by Protobuf,
-	// the LSP must handle this invalid state gracefully.
-	//
-	// This is enforced by mutex.go.
-	lock mutex
-
 	text string
 	// Version is an opaque version identifier given to us by the LSP client. This
 	// is used in the protocol to disambiguate which version of a file e.g. publishing
@@ -84,19 +62,34 @@ type file struct {
 	hasText bool // Whether this file has ever had text read into it.
 	// Always set false->true. Once true, never becomes false again.
 
-	workspace     bufworkspace.Workspace
-	module        bufmodule.Module
-	imageFileInfo bufimage.ImageFileInfo
+	workspace bufworkspace.Workspace
+	module    bufmodule.Module
+
+	objectInfo             storage.ObjectInfo
+	importablePathToObject map[string]storage.ObjectInfo
 
-	isWKT bool
+	fileNode     *ast.FileNode
+	packageNode  *ast.PackageNode
+	diagnostics  []protocol.Diagnostic
+	importToFile map[string]*file
+	symbols      []*symbol
+	image        bufimage.Image
+}
 
-	fileNode          *ast.FileNode
-	packageNode       *ast.PackageNode
-	diagnostics       []protocol.Diagnostic
-	importableToImage map[string]bufimage.ImageFileInfo
-	importToFile      map[string]*file
-	symbols           []*symbol
-	image             bufimage.Image
+// IsWKT returns whether this file corresponds to a well-known type.
+func (f *file) IsWKT() bool {
+	_, ok := f.objectInfo.(wktObjectInfo)
+	return ok
+}
+
+// IsLocal returns whether this is a local file, i.e. a file that the editor
+// is editing and not something from e.g. the BSR.
+func (f *file) IsLocal() bool {
+	if f.objectInfo == nil {
+		return false
+	}
+
+	return f.objectInfo.LocalPath() == f.objectInfo.ExternalPath()
 }
 
 // Manager returns the file manager that owns this file.
@@ -117,26 +110,15 @@ func (f *file) Package() []string {
 func (f *file) Reset(ctx context.Context) {
 	f.lsp.logger.Debug(fmt.Sprintf("resetting file %v", f.uri))
 
-	// Lock and unlock to acquire the import map, then nil everything out
-	// This map is never mutated after being created, so we only
-	// need to read the pointer.
-	//
-	// We need to lock and unlock because Close() will call Reset() on other
-	// files, and this will deadlock if cyclic imports exist.
-	f.lock.Lock(ctx)
-	imports := f.importToFile
-
 	f.fileNode = nil
 	f.packageNode = nil
 	f.diagnostics = nil
-	f.importableToImage = nil
+	f.importablePathToObject = nil
 	f.importToFile = nil
 	f.symbols = nil
 	f.image = nil
-	f.lock.Unlock(ctx)
 
-	// Close all imported files while file.mu is not held.
-	for _, imported := range imports {
+	for _, imported := range f.importToFile {
 		imported.Close(ctx)
 	}
 }
@@ -154,8 +136,6 @@ func (f *file) Close(ctx context.Context) {
 // If it has been read from disk before, or has received updates from the LSP client, this
 // function returns nil.
 func (f *file) ReadFromDisk(ctx context.Context) (err error) {
-	f.lock.Lock(ctx)
-	defer f.lock.Unlock(ctx)
 	if f.hasText {
 		return nil
 	}
@@ -175,9 +155,6 @@ func (f *file) ReadFromDisk(ctx context.Context) (err error) {
 func (f *file) Update(ctx context.Context, version int32, text string) {
 	f.Reset(ctx)
 
-	f.lock.Lock(ctx)
-	defer f.lock.Unlock(ctx)
-
 	f.lsp.logger.Info(fmt.Sprintf("new file version: %v, %v -> %v", f.uri, f.version, version))
 	f.version = version
 	f.text = text
@@ -217,8 +194,6 @@ func (f *file) Refresh(ctx context.Context) {
 //
 // Returns whether a reparse was necessary.
 func (f *file) RefreshAST(ctx context.Context) bool {
-	f.lock.Lock(ctx)
-	defer f.lock.Unlock(ctx)
 	if f.fileNode != nil {
 		return false
 	}
@@ -257,9 +232,6 @@ func (f *file) RefreshAST(ctx context.Context) bool {
 func (f *file) PublishDiagnostics(ctx context.Context) {
 	defer slogext.DebugProfile(f.lsp.logger, slog.String("uri", string(f.uri)))()
 
-	f.lock.Lock(ctx)
-	defer f.lock.Unlock(ctx)
-
 	if f.diagnostics == nil {
 		return
 	}
@@ -307,39 +279,37 @@ func (f *file) FindModule(ctx context.Context) {
 		defer file.Close()
 	}
 
-	f.lock.Lock(ctx)
 	f.workspace = workspace
 	f.module = module
-	f.lock.Unlock(ctx)
 }
 
 // IndexImports finds URIs for all of the files imported by this file.
 func (f *file) IndexImports(ctx context.Context) {
 	defer slogext.DebugProfile(f.lsp.logger, slog.String("uri", string(f.uri)))()
 
-	unlock := f.lock.Lock(ctx)
-	defer unlock()
-
 	if f.fileNode == nil || f.importToFile != nil {
 		return
 	}
 
-	importable, err := f.lsp.findImportable(ctx, f.uri)
+	importable, err := findImportable(ctx, f.uri, f.lsp)
 	if err != nil {
 		f.lsp.logger.Warn(fmt.Sprintf("could not compute importable files for %s: %s", f.uri, err))
-		return
+		if f.importablePathToObject == nil {
+			return
+		}
+	} else if f.importablePathToObject == nil {
+		f.importablePathToObject = importable
 	}
-	f.importableToImage = importable
 
 	// Find the FileInfo for this path. The crazy thing is that it may appear in importable
 	// multiple times, with different path lengths! We want to pick the one with the longest path
 	// length.
 	for _, fileInfo := range importable {
 		if fileInfo.LocalPath() == f.uri.Filename() {
-			if f.imageFileInfo != nil && len(f.imageFileInfo.Path()) > len(fileInfo.Path()) {
+			if f.objectInfo != nil && len(f.objectInfo.Path()) > len(fileInfo.Path()) {
 				continue
 			}
-			f.imageFileInfo = fileInfo
+			f.objectInfo = fileInfo
 		}
 	}
 
@@ -350,12 +320,50 @@ func (f *file) IndexImports(ctx context.Context) {
 			continue
 		}
 
+		// If this is an external file, it will be in the cache and therefore
+		// finding imports via lsp.findImportable() will not work correctly:
+		// the bucket for the workspace found for a dependency will have
+		// truncated paths, and those workspace files will appear to be
+		// local rather than external.
+		//
+		// Thus, we search for name and all of its path suffixes. This is not
+		// ideal but is our only option in this case.
+		var fileInfo storage.ObjectInfo
+		var pathWasTruncated bool
 		name := node.Name.AsString()
-		fileInfo, ok := importable[name]
-		if !ok {
-			f.lsp.logger.Warn(fmt.Sprintf("could not find URI for import %q", name))
+		for {
+			fileInfo, ok = importable[name]
+			if ok {
+				break
+			}
+
+			idx := strings.Index(name, "/")
+			if idx == -1 {
+				break
+			}
+
+			name = name[idx+1:]
+			pathWasTruncated = true
+		}
+		if fileInfo == nil {
+			f.lsp.logger.Warn(fmt.Sprintf("could not find URI for import %q", node.Name.AsString()))
 			continue
 		}
+		if pathWasTruncated && !strings.HasSuffix(fileInfo.LocalPath(), node.Name.AsString()) {
+			// Verify that the file we found, with a potentially too-short path, does in fact have
+			// the "correct" full path as a prefix. E.g., suppose we import a/b/c.proto. We find
+			// c.proto in importable. Now, we look at the full local path, which we expect to be of
+			// the form /home/blah/.cache/blah/a/b/c.proto or similar. If it does not contain
+			// a/b/c.proto as a suffix, we didn't find our file.
+			f.lsp.logger.Warn(fmt.Sprintf("could not find URI for import %q, but found same-suffix path %q", node.Name.AsString(), fileInfo.LocalPath()))
+			continue
+		}
+
+		f.lsp.logger.Debug(
+			"mapped import -> path",
+			slog.String("import", name),
+			slog.String("path", fileInfo.LocalPath()),
+		)
 
 		var imported *file
 		if fileInfo.LocalPath() == f.uri.Filename() {
@@ -364,8 +372,7 @@ func (f *file) IndexImports(ctx context.Context) {
 			imported = f.Manager().Open(ctx, protocol.URI("file://"+fileInfo.LocalPath()))
 		}
 
-		imported.imageFileInfo = fileInfo
-		f.isWKT = strings.HasPrefix("google/protobuf/", fileInfo.Path())
+		imported.objectInfo = fileInfo
 		f.importToFile[node.Name.AsString()] = imported
 	}
 
@@ -377,19 +384,13 @@ func (f *file) IndexImports(ctx context.Context) {
 			f.importToFile[descriptorPath] = f
 		} else {
 			imported := f.Manager().Open(ctx, descriptorURI)
-			imported.imageFileInfo = descriptorFile
+			imported.objectInfo = descriptorFile
 			f.importToFile[descriptorPath] = imported
 		}
-		f.isWKT = true
 	}
 
 	// FIXME: This algorithm is not correct: it does not account for `import public`.
-
-	// Drop the lock after copying the pointer to the imports map. This
-	// particular map will not be mutated further, and since we're going to grab the lock of
-	// other files, we need to drop the currently held lock.
 	fileImports := f.importToFile
-	unlock()
 
 	for _, file := range fileImports {
 		if err := file.ReadFromDisk(ctx); err != nil {
@@ -412,10 +413,8 @@ func (f *file) IndexImports(ctx context.Context) {
 //
 // This operation requires IndexImports().
 func (f *file) BuildImage(ctx context.Context) {
-	f.lock.Lock(ctx)
-	importable := f.importableToImage
-	fileInfo := f.imageFileInfo
-	f.lock.Unlock(ctx)
+	importable := f.importablePathToObject
+	fileInfo := f.objectInfo
 
 	if importable == nil || fileInfo == nil {
 		return
@@ -450,9 +449,7 @@ func (f *file) BuildImage(ctx context.Context) {
 
 	compiled, err := compiler.Compile(ctx, fileInfo.Path())
 	if err != nil {
-		f.lock.Lock(ctx)
 		f.diagnostics = report.diagnostics
-		f.lock.Unlock(ctx)
 	}
 	if compiled[0] == nil {
 		return
@@ -530,25 +527,21 @@ func (f *file) BuildImage(ctx context.Context) {
 		return
 	}
 
-	f.lock.Lock(ctx)
 	f.image = image
-	f.lock.Unlock(ctx)
 }
 
 // RunLints runs linting on this file. Returns whether any lints failed.
 //
 // This operation requires BuildImage().
 func (f *file) RunLints(ctx context.Context) bool {
-	if f.isWKT {
+	if f.IsWKT() {
 		// Well-known types are not linted.
 		return false
 	}
 
-	f.lock.Lock(ctx)
 	workspace := f.workspace
 	module := f.module
 	image := f.image
-	f.lock.Unlock(ctx)
 
 	if module == nil || image == nil {
 		f.lsp.logger.Warn(fmt.Sprintf("could not find image for %q", f.uri))
@@ -578,8 +571,6 @@ func (f *file) RunLints(ctx context.Context) bool {
 
 	f.lsp.logger.Warn(fmt.Sprintf("lint generated %d error(s) for %s", len(annotations.FileAnnotations()), f.uri))
 
-	f.lock.Lock(ctx)
-	f.lock.Unlock(ctx)
 	for _, annotation := range annotations.FileAnnotations() {
 		f.lsp.logger.Info(annotation.FileInfo().Path(), " ", annotation.FileInfo().ExternalPath())
 
@@ -608,15 +599,14 @@ func (f *file) RunLints(ctx context.Context) bool {
 func (f *file) IndexSymbols(ctx context.Context) {
 	defer slogext.DebugProfile(f.lsp.logger, slog.String("uri", string(f.uri)))()
 
-	unlock := f.lock.Lock(ctx)
-	defer unlock()
-
 	// Throw away all the old symbols. Unlike other indexing functions, we rebuild
 	// symbols unconditionally.
 	f.symbols = nil
 
 	// Generate new symbols.
-	newWalker(f).Walk(f.fileNode, f.fileNode)
+	walker := newWalker(f)
+	walker.Walk(f.fileNode, f.fileNode)
+	f.symbols = walker.symbols
 
 	// Finally, sort the symbols in position order, with shorter symbols sorting smaller.
 	slices.SortFunc(f.symbols, func(s1, s2 *symbol) int {
@@ -629,7 +619,6 @@ func (f *file) IndexSymbols(ctx context.Context) {
 
 	// Now we can drop the lock and search for cross-file references.
 	symbols := f.symbols
-	unlock()
 	for _, symbol := range symbols {
 		symbol.ResolveCrossFile(ctx)
 	}
@@ -641,9 +630,6 @@ func (f *file) IndexSymbols(ctx context.Context) {
 //
 // Returns nil if no symbol is found.
 func (f *file) SymbolAt(ctx context.Context, cursor protocol.Position) *symbol {
-	f.lock.Lock(ctx)
-	defer f.lock.Unlock(ctx)
-
 	// Binary search for the symbol whose start is before or equal to cursor.
 	idx, found := slices.BinarySearchFunc(f.symbols, cursor, func(sym *symbol, cursor protocol.Position) int {
 		return comparePositions(sym.Range().Start, cursor)
@@ -665,3 +651,66 @@ func (f *file) SymbolAt(ctx context.Context, cursor protocol.Position) *symbol {
 
 	return symbol
 }
+
+// findImportable finds all files that can potentially be imported by the proto file at
+// uri. This returns a map from potential Protobuf import path to the URI of the file it would import.
+//
+// Note that this performs no validation on these files, because those files might be open in the
+// editor and might contain invalid syntax at the moment. We only want to get their paths and nothing
+// more.
+func findImportable(
+	ctx context.Context,
+	uri protocol.URI,
+	lsp *lsp,
+) (map[string]storage.ObjectInfo, error) {
+	// This does not use Controller.GetImportableImageFileInfos because:
+	//
+	// 1. That function throws away Module/ModuleSet information, because it
+	//    converts the module contents into ImageFileInfos.
+	//
+	// 2. That function does not specify which of the files it returns are
+	//    well-known imports. Previously, we were making an educated guess about
+	//    which files were well-known, but this resulted in subtle classification
+	//    bugs.
+	//
+	// Doing the file walk here manually helps us retain some control over what
+	// data is discarded.
+	workspace, err := lsp.controller.GetWorkspace(ctx, uri.Filename())
+	if err != nil {
+		return nil, err
+	}
+
+	imports := make(map[string]storage.ObjectInfo)
+	for _, module := range workspace.Modules() {
+		err = module.WalkFileInfos(ctx, func(fileInfo bufmodule.FileInfo) error {
+			if fileInfo.FileType() != bufmodule.FileTypeProto {
+				return nil
+			}
+
+			imports[fileInfo.Path()] = fileInfo
+
+			return nil
+		})
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	err = lsp.wktBucket.Walk(ctx, "", func(object storage.ObjectInfo) error {
+		imports[object.Path()] = wktObjectInfo{object}
+		return nil
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	lsp.logger.Debug(fmt.Sprintf("found imports for %q: %#v", uri, imports))
+
+	return imports, nil
+}
+
+// wktObjectInfo is a concrete type to help us identify WKTs among the
+// importable files.
+type wktObjectInfo struct {
+	storage.ObjectInfo
+}
diff --git a/private/buf/buflsp/file_manager.go b/private/buf/buflsp/file_manager.go
index 8d526478ae..dd604c9ff9 100644
--- a/private/buf/buflsp/file_manager.go
+++ b/private/buf/buflsp/file_manager.go
@@ -28,7 +28,6 @@ import (
 type fileManager struct {
 	lsp       *lsp
 	uriToFile refcount.Map[protocol.URI, file]
-	mutexPool mutexPool
 }
 
 // newFiles creates a new file manager.
@@ -44,7 +43,6 @@ func (fm *fileManager) Open(ctx context.Context, uri protocol.URI) *file {
 	if !found {
 		file.lsp = fm.lsp
 		file.uri = uri
-		file.lock = fm.mutexPool.NewMutex()
 	}
 
 	return file
diff --git a/private/buf/buflsp/mutex.go b/private/buf/buflsp/mutex.go
deleted file mode 100644
index c377b126b7..0000000000
--- a/private/buf/buflsp/mutex.go
+++ /dev/null
@@ -1,182 +0,0 @@
-// Copyright 2020-2024 Buf Technologies, Inc.
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-// This file defines various concurrency helpers.
-
-package buflsp
-
-import (
-	"context"
-	"fmt"
-	"sync"
-	"sync/atomic"
-)
-
-const poison = ^uint64(0)
-
-var nextRequestID atomic.Uint64
-
-// mutexPool represents a group of reentrant muteces that cannot be acquired simultaneously.
-//
-// A zero mutexPool is ready to use.
-type mutexPool struct {
-	lock sync.Mutex
-	held map[uint64]*mutex
-}
-
-// NewMutex creates a new mutex in this pool.
-func (mp *mutexPool) NewMutex() mutex {
-	return mutex{pool: mp}
-}
-
-// check checks what id is either not holding a lock, or is holding the given
-// map, depending on whether isUnlock is set.
-func (mp *mutexPool) check(id uint64, mu *mutex, isUnlock bool) {
-	if mp == nil {
-		return
-	}
-
-	mp.lock.Lock()
-	defer mp.lock.Unlock()
-
-	if mp.held == nil {
-		mp.held = make(map[uint64]*mutex)
-	}
-
-	if isUnlock {
-		if held := mp.held[id]; held != mu {
-			panic(fmt.Sprintf("buflsp/mutex.go: attempted to unlock incorrect non-reentrant lock: %p -> %p", held, mu))
-		}
-
-		delete(mp.held, id)
-	} else {
-		if held := mp.held[id]; held != nil {
-			panic(fmt.Sprintf("buflsp/mutex.go: attempted to acquire two non-reentrant locks at once: %p -> %p", mu, held))
-		}
-
-		mp.held[id] = mu
-	}
-}
-
-// mutex is a sync.Mutex with some extra features.
-//
-// The main feature is reentrancy-checking. Within the LSP, we need to lock-protect many structures,
-// and it is very easy to deadlock if the same request tries to lock something multiple times.
-// To achieve this, Lock() takes a context, which must be modified by withRequestID().
-type mutex struct {
-	lock sync.Mutex
-	// This is the id of the context currently holding the lock.
-	who  atomic.Uint64
-	pool *mutexPool
-}
-
-// Lock attempts to acquire this mutex or blocks.
-//
-// Unlike [sync.Mutex.Lock], this takes a Context. If that context was updated with withRequestID,
-// this function will panic when attempting to lock the mutex while it is already held by a
-// goroutine using this same context.
-//
-// NOTE: to Lock() and Unlock() with the same context DO NOT synchronize with each other. For example,
-// attempting to lock this mutex from two different goroutines with the same context will
-// result in undefined behavior.
-//
-// Also unlike [sync.Mutex.Lock], it returns an idempotent unlocker function. This can be used like
-// defer mu.Lock()(). Note that only the outer function call is deferred: this is part of the
-// definition of defer. See https://go.dev/play/p/RJNKRcoQRo1. This unlocker can also be used to
-// defer unlocking but also unlock before the function returns.
-//
-// The returned unlocker is not thread-safe.
-func (mu *mutex) Lock(ctx context.Context) (unlocker func()) {
-	var unlocked bool
-	unlocker = func() {
-		if unlocked {
-			return
-		}
-		mu.Unlock(ctx)
-		unlocked = true
-	}
-
-	id := getRequestID(ctx)
-
-	if mu.who.Load() == id && id > 0 {
-		// We seem to have tried to lock this lock twice. Panic, and poison the lock.
-		mu.who.Store(poison)
-		panic("buflsp/mutex.go: non-reentrant lock locked twice by the same request")
-	}
-
-	mu.pool.check(id, mu, false)
-
-	// Ok, we're definitely not holding a lock, so we can block until we acquire the lock.
-	mu.lock.Lock()
-	mu.storeWho(id)
-
-	return unlocker
-}
-
-// Unlock releases this mutex.
-//
-// Unlock must be called with the same context that locked it, otherwise this function panics.
-func (mu *mutex) Unlock(ctx context.Context) {
-	id := getRequestID(ctx)
-	if mu.who.Load() != id {
-		panic("buflsp/mutex.go: lock was locked by one request and unlocked by another")
-	}
-
-	mu.storeWho(0)
-
-	mu.pool.check(id, mu, true)
-	mu.lock.Unlock()
-}
-
-func (mu *mutex) storeWho(id uint64) {
-	for {
-		// This has to be a CAS loop to avoid races with a poisoning p.
-		old := mu.who.Load()
-		if old == poison {
-			panic("buflsp/mutex.go: non-reentrant lock locked twice by the same request")
-		}
-		if mu.who.CompareAndSwap(old, id) {
-			break
-		}
-	}
-}
-
-// withRequestID assigns a unique request ID to the given context, which can be retrieved
-// with with getRequestID.
-func withRequestID(ctx context.Context) context.Context {
-	// This will always be unique. It is impossible to increment a uint64 and wrap around before
-	// the heat death of the universe.
-	id := nextRequestID.Add(1)
-	// We need to give the context package a unique identifier for the request; it can be
-	// any value. The address of the global we mint new IDs from is actually great for this,
-	// because users can't access it outside of this package, nor can they extract it out
-	// of the context itself.
-	return context.WithValue(ctx, &nextRequestID, id)
-}
-
-// getRequestID returns the request ID for this context, or 0 if ctx is nil or has no
-// such ID.
-func getRequestID(ctx context.Context) uint64 {
-	if ctx == nil {
-		return 0
-	}
-	id, ok := ctx.Value(&nextRequestID).(uint64)
-	if !ok {
-		return 0
-	}
-
-	// Make sure we don't return 0. This is the only place where the id is actually
-	// witnessed so doing +1 won't affect anything.
-	return id + 1
-}
diff --git a/private/buf/buflsp/report.go b/private/buf/buflsp/report.go
index 6e71363546..6114146119 100644
--- a/private/buf/buflsp/report.go
+++ b/private/buf/buflsp/report.go
@@ -44,8 +44,15 @@ func (r *report) Warning(err reporter.ErrorWithPos) {
 	r.diagnostics = append(r.diagnostics, newDiagnostic(err, true))
 
 	if err.Unwrap() == parser.ErrNoSyntax {
+		if r.syntaxMissing == nil {
+			r.syntaxMissing = make(map[string]bool)
+		}
 		r.syntaxMissing[err.GetPosition().Filename] = true
 	} else if unusedImport, ok := err.Unwrap().(linker.ErrorUnusedImport); ok {
+		if r.pathToUnusedImports == nil {
+			r.pathToUnusedImports = make(map[string]map[string]bool)
+		}
+
 		path := err.GetPosition().Filename
 		unused, ok := r.pathToUnusedImports[path]
 		if !ok {
diff --git a/private/buf/buflsp/server.go b/private/buf/buflsp/server.go
index ada3cea0a7..eaf1f3c724 100644
--- a/private/buf/buflsp/server.go
+++ b/private/buf/buflsp/server.go
@@ -22,7 +22,6 @@ import (
 	"fmt"
 	"runtime/debug"
 	"strings"
-	"time"
 
 	"github.com/bufbuild/buf/private/buf/bufformat"
 	"github.com/bufbuild/protocompile/ast"
@@ -79,7 +78,7 @@ func (s *server) Initialize(
 	ctx context.Context,
 	params *protocol.InitializeParams,
 ) (*protocol.InitializeResult, error) {
-	if err := s.init(params); err != nil {
+	if err := s.init(ctx, params); err != nil {
 		return nil, err
 	}
 
@@ -173,7 +172,7 @@ func (s *server) DidOpen(
 ) error {
 	file := s.fileManager.Open(ctx, params.TextDocument.URI)
 	file.Update(ctx, params.TextDocument.Version, params.TextDocument.Text)
-	go file.Refresh(context.WithoutCancel(ctx))
+	file.Refresh(context.WithoutCancel(ctx))
 	return nil
 }
 
@@ -190,7 +189,7 @@ func (s *server) DidChange(
 	}
 
 	file.Update(ctx, params.TextDocument.Version, params.ContentChanges[0].Text)
-	go file.Refresh(context.WithoutCancel(ctx))
+	file.Refresh(context.WithoutCancel(ctx))
 	return nil
 }
 
@@ -207,9 +206,6 @@ func (s *server) Formatting(
 
 	// Currently we have no way to honor any of the parameters.
 	_ = params
-
-	file.lock.Lock(ctx)
-	defer file.lock.Unlock(ctx)
 	if file.fileNode == nil {
 		return nil, nil
 	}
@@ -322,23 +318,12 @@ func (s *server) SemanticTokensFull(
 	progress.Begin(ctx, "Processing Tokens")
 	defer progress.Done(ctx)
 
-	var symbols []*symbol
-	for {
-		file.lock.Lock(ctx)
-		symbols = file.symbols
-		file.lock.Unlock(ctx)
-		if symbols != nil {
-			break
-		}
-		time.Sleep(1 * time.Millisecond)
-	}
-
 	var (
 		encoded           []uint32
 		prevLine, prevCol uint32
 	)
-	for i, symbol := range symbols {
-		progress.Report(ctx, fmt.Sprintf("%d/%d", i+1, len(symbols)), float64(i)/float64(len(symbols)))
+	for i, symbol := range file.symbols {
+		progress.Report(ctx, fmt.Sprintf("%d/%d", i+1, len(file.symbols)), float64(i)/float64(len(file.symbols)))
 
 		var semanticType uint32
 
diff --git a/private/buf/buflsp/symbol.go b/private/buf/buflsp/symbol.go
index 1bb7e3a84a..8ea0da97f1 100644
--- a/private/buf/buflsp/symbol.go
+++ b/private/buf/buflsp/symbol.go
@@ -27,6 +27,7 @@ import (
 	"slices"
 	"strings"
 
+	"github.com/bufbuild/buf/private/bufpkg/bufmodule"
 	"github.com/bufbuild/buf/private/pkg/slicesext"
 	"github.com/bufbuild/protocompile/ast"
 	"go.lsp.dev/protocol"
@@ -115,8 +116,6 @@ func (s *symbol) Definition(ctx context.Context) (*symbol, ast.Node) {
 			return nil, nil
 		}
 
-		kind.file.lock.Lock(ctx)
-		defer kind.file.lock.Unlock(ctx)
 		for _, symbol := range kind.file.symbols {
 			def, ok := symbol.kind.(*definition)
 			if ok && slices.Equal(kind.path, def.path) {
@@ -272,22 +271,11 @@ func (s *symbol) ResolveCrossFile(ctx context.Context) {
 				return
 			}
 
-			// Make a copy of the import table pointer and then drop the lock,
-			// since searching inside of the imports will need to acquire other
-			// fileManager' locks.
-			s.file.lock.Lock(ctx)
+			// Look for a symbol with this exact path in descriptor proto.
 			descriptorProto := s.file.importToFile[descriptorPath]
-			s.file.lock.Unlock(ctx)
-
 			if descriptorProto == nil {
 				return
 			}
-
-			// Look for a symbol with this exact path in descriptor proto.
-
-			descriptorProto.lock.Lock(ctx)
-			defer descriptorProto.lock.Unlock(ctx)
-
 			var fieldSymbol *symbol
 			for _, symbol := range descriptorProto.symbols {
 				if def, ok := symbol.kind.(*definition); ok && slices.Equal(def.path, fieldPath) {
@@ -304,23 +292,35 @@ func (s *symbol) ResolveCrossFile(ctx context.Context) {
 			return
 		}
 
-		// Make a copy of the import table pointer and then drop the lock,
-		// since searching inside of the imports will need to acquire other
-		// fileManager' locks.
-		s.file.lock.Lock(ctx)
-		imports := s.file.importToFile
-		s.file.lock.Unlock(ctx)
-
-		if imports == nil {
+		if s.file.importToFile == nil {
 			// Hopeless. We'll have to try again once we have imports!
 			return
 		}
 
-		for _, imported := range imports {
-			// Remove a leading pkg from components.
+		for _, imported := range s.file.importToFile {
+			// If necessary, refresh the file. Note that this cannot hit
+			// cycles, because fileNode will become non-nil after calling
+			// Refresh but before calling IndexSymbols.
+			if imported.fileNode == nil {
+				imported.Refresh(ctx)
+			}
+
+			// Need to check two paths; components on its own, and components
+			// with s.file's package prepended.
+
+			// First, try removing the imported file's package from components.
 			path, ok := slicesext.TrimPrefix(components, imported.Package())
 			if !ok {
-				continue
+				// If that doesn't work, try appending the importee's package.
+				// This is necessary because protobuf allows for partial package
+				// names to appear in references.
+				path, ok = slicesext.TrimPrefix(
+					slices.Concat(s.file.Package(), components),
+					imported.Package(),
+				)
+				if !ok {
+					continue
+				}
 			}
 
 			if findDeclByPath(imported.fileNode.Decls, path) != nil {
@@ -332,7 +332,48 @@ func (s *symbol) ResolveCrossFile(ctx context.Context) {
 	}
 }
 
-// FormatDocs finds appropriate  documgntation for the given s and constructs a Markdown
+func (s *symbol) LogValue() slog.Value {
+	attrs := []slog.Attr{slog.String("file", s.file.uri.Filename())}
+
+	// pos converts an ast.SourcePos into a slog.Value.
+	pos := func(pos ast.SourcePos) slog.Value {
+		return slog.GroupValue(
+			slog.Int("line", pos.Line),
+			slog.Int("col", pos.Col),
+		)
+	}
+
+	attrs = append(attrs, slog.Any("start", pos(s.info.Start())))
+	attrs = append(attrs, slog.Any("end", pos(s.info.End())))
+
+	switch kind := s.kind.(type) {
+	case *builtin:
+		attrs = append(attrs, slog.String("builtin", kind.name))
+
+	case *import_:
+		if kind.file != nil {
+			attrs = append(attrs, slog.String("imports", kind.file.uri.Filename()))
+		}
+
+	case *definition:
+		attrs = append(attrs, slog.String("defines", strings.Join(kind.path, ".")))
+
+	case *reference:
+		if kind.file != nil {
+			attrs = append(attrs, slog.String("imports", kind.file.uri.Filename()))
+		}
+		if kind.path != nil {
+			attrs = append(attrs, slog.String("references", strings.Join(kind.path, ".")))
+		}
+		if kind.seeTypeOf != nil {
+			attrs = append(attrs, slog.Any("see_type_of", kind.seeTypeOf))
+		}
+	}
+
+	return slog.GroupValue(attrs...)
+}
+
+// FormatDocs finds appropriate documentation for the given s and constructs a Markdown
 // string for showing to the client.
 //
 // Returns the empty string if no docs are available.
@@ -371,6 +412,10 @@ func (s *symbol) FormatDocs(ctx context.Context) string {
 		return ""
 	}
 
+	if def == nil {
+		return ""
+	}
+
 	pkg := "<empty>"
 	if pkgNode := def.file.packageNode; pkgNode != nil {
 		pkg = string(pkgNode.Name.AsIdentifier())
@@ -413,19 +458,37 @@ func (s *symbol) FormatDocs(ctx context.Context) string {
 		return tooltip.String()
 	}
 
-	if def.file.imageFileInfo != nil {
-		path := strings.Join(path, ".")
+	// Do not show BSR links for local files, because those may contain symbols
+	// that dop not exist remotely.
+	if def.file != nil && !def.file.IsLocal() {
+		// Classify this node by whether or not the BSR has an anchor for it.
+		// Extensions are special: they are grouped under an #extensions anchor.
+		var hasAnchor, isExtension bool
+		switch node := node.(type) {
+		case *ast.MessageNode, *ast.EnumNode, *ast.ServiceNode:
+			hasAnchor = true
+		case *ast.FieldNode:
+			isExtension = node.Extendee != nil
+			hasAnchor = isExtension
+		}
 
-		fmt.Fprintf(
-			&tooltip,
-			"[`%s.%s` on the Buf Schema Registry](https://%s/docs/main:%s#%s.%s)\n\n",
-			pkg,
-			path,
-			def.file.imageFileInfo.ModuleFullName(),
-			pkg,
-			pkg,
-			path,
-		)
+		var bsrHost, bsrAnchor, bsrTooltip string
+		if def.file.IsWKT() {
+			bsrHost = "buf.build/protocolbuffers/wellknowntypes"
+		} else if fileInfo, ok := def.file.objectInfo.(bufmodule.FileInfo); ok {
+			bsrHost = fileInfo.Module().ModuleFullName().String()
+		}
+		if hasAnchor {
+			bsrTooltip = pkg + "." + strings.Join(path, ".")
+		} else {
+			bsrTooltip = pkg + "." + strings.Join(path[:len(path)-1], ".")
+		}
+		bsrAnchor = bsrTooltip
+		if isExtension {
+			bsrAnchor = "extensions"
+		}
+
+		fmt.Fprintf(&tooltip, "[`%s` on the Buf Schema Registry](https://%s/docs/main:%s#%s)\n\n", bsrTooltip, bsrHost, pkg, bsrAnchor)
 	}
 
 	// Dump all of the comments into the tooltip. These will be rendered as Markdown automatically
@@ -440,9 +503,13 @@ func (s *symbol) FormatDocs(ctx context.Context) string {
 			// The compiler does not currently provide comments without their
 			// delimited removed, so we have to do this ourselves.
 			if strings.HasPrefix(comment, "//") {
-				comment = strings.TrimSpace(strings.TrimPrefix(comment, "//"))
+				// NOTE: We do not trim the space here, because indentation is
+				// significant for Markdown code fences, and if every line
+				// starts with a space, Markdown will trim it for us, even off
+				// of code blocks.
+				comment = strings.TrimPrefix(comment, "//")
 			} else {
-				comment = strings.TrimSpace(strings.TrimSuffix(strings.TrimPrefix(comment, "/*"), "*/"))
+				comment = strings.TrimSuffix(strings.TrimPrefix(comment, "/*"), "*/")
 			}
 
 			if comment != "" {
@@ -463,7 +530,8 @@ func (s *symbol) FormatDocs(ctx context.Context) string {
 
 // symbolWalker is an AST walker that generates the symbol table for a file in IndexSymbols().
 type symbolWalker struct {
-	file *file
+	file    *file
+	symbols []*symbol
 
 	// This is the set of *ast.MessageNode, *ast.EnumNode, and *ast.ServiceNode that
 	// we have traversed. They are used for same-file symbol resolution, and for constructing
@@ -613,7 +681,7 @@ func (w *symbolWalker) Walk(node, parent ast.Node) {
 				}
 			} else {
 				// This depends on the type of the previous symbol.
-				prev := w.file.symbols[len(w.file.symbols)-1]
+				prev := w.symbols[len(w.symbols)-1]
 				next = w.newSymbol(part.Name)
 				next.kind = &reference{seeTypeOf: prev}
 			}
@@ -634,7 +702,7 @@ func (w *symbolWalker) newSymbol(name ast.Node) *symbol {
 		info: w.file.fileNode.NodeInfo(name),
 	}
 
-	w.file.symbols = append(w.file.symbols, symbol)
+	w.symbols = append(w.symbols, symbol)
 	return symbol
 }
 
diff --git a/private/buf/cmd/buf/command/beta/lsp/lsp.go b/private/buf/cmd/buf/command/beta/lsp/lsp.go
index fabb0c91be..adb32c87e2 100644
--- a/private/buf/cmd/buf/command/beta/lsp/lsp.go
+++ b/private/buf/cmd/buf/command/beta/lsp/lsp.go
@@ -89,6 +89,15 @@ func run(
 		return err
 	}
 
+	wktStore, err := bufcli.NewWKTStore(container)
+	if err != nil {
+		return err
+	}
+	wktBucket, err := wktStore.GetBucket(ctx)
+	if err != nil {
+		return err
+	}
+
 	controller, err := bufcli.NewController(container)
 	if err != nil {
 		return err
@@ -114,7 +123,7 @@ func run(
 		return err
 	}
 
-	conn, err := buflsp.Serve(ctx, container, controller, checkClient, jsonrpc2.NewStream(transport))
+	conn, err := buflsp.Serve(ctx, wktBucket, container, controller, checkClient, jsonrpc2.NewStream(transport))
 	if err != nil {
 		return err
 	}