internal/lsp: watch directories in replace targets and update on changes

This change adds the notion of a "workspace directory", which is
basically the set of directories that contains workspace packages. These
are mainly used for replace targets right now. It's a little trickier
than expected because the set of workspace directories can technically
change on any go.mod change.

At first, I wanted DidModifyFiles to report whether there was a change,
but I don't think it's actually that expensive to check on each call
and it complicates the code a bit. I can change it back if you think
it's worth doing.

The parse mod handle changes are because I needed an unlocked way of
parsing the mod file, but I imagine they'll conflict with CL 244769
anyway.

The next CL will be to "promote" replace targets to the level of
workspace packages, meaning we will be able to find references in them.

Change-Id: I5dd58fe29415473496ca6634a94a3134923228dc
Reviewed-on: https://go-review.googlesource.com/c/tools/+/245327
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
Reviewed-by: Heschi Kreinick <heschi@google.com>
diff --git a/internal/lsp/cache/session.go b/internal/lsp/cache/session.go
index ef1104e..4374eee 100644
--- a/internal/lsp/cache/session.go
+++ b/internal/lsp/cache/session.go
@@ -204,6 +204,14 @@
 		Env:         v.goEnv,
 	}
 
+	// Set the first snapshot's workspace directories. The view's modURI was
+	// set by setBuildInformation.
+	var fh source.FileHandle
+	if v.modURI != "" {
+		fh, _ = s.GetFile(ctx, v.modURI)
+	}
+	v.snapshot.workspaceDirectories = v.snapshot.findWorkspaceDirectories(ctx, fh)
+
 	// Initialize the view without blocking.
 	initCtx, initCancel := context.WithCancel(xcontext.Detach(ctx))
 	v.initCancelFirstAttempt = initCancel
diff --git a/internal/lsp/cache/snapshot.go b/internal/lsp/cache/snapshot.go
index e940133..87c164e 100644
--- a/internal/lsp/cache/snapshot.go
+++ b/internal/lsp/cache/snapshot.go
@@ -75,6 +75,10 @@
 	// when the view is created.
 	workspacePackages map[packageID]packagePath
 
+	// workspaceDirectories are the directories containing workspace packages.
+	// They are the view's root, as well as any replace targets.
+	workspaceDirectories map[span.URI]struct{}
+
 	// unloadableFiles keeps track of files that we've failed to load.
 	unloadableFiles map[span.URI]struct{}
 
@@ -419,6 +423,17 @@
 	return ids
 }
 
+func (s *snapshot) WorkspaceDirectories(ctx context.Context) []span.URI {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	var dirs []span.URI
+	for d := range s.workspaceDirectories {
+		dirs = append(dirs, d)
+	}
+	return dirs
+}
+
 func (s *snapshot) WorkspacePackages(ctx context.Context) ([]source.Package, error) {
 	if err := s.awaitLoaded(ctx); err != nil {
 		return nil, err
@@ -778,22 +793,23 @@
 	defer s.mu.Unlock()
 
 	result := &snapshot{
-		id:                s.id + 1,
-		view:              s.view,
-		builtin:           s.builtin,
-		ids:               make(map[span.URI][]packageID),
-		importedBy:        make(map[packageID][]packageID),
-		metadata:          make(map[packageID]*metadata),
-		packages:          make(map[packageKey]*packageHandle),
-		actions:           make(map[actionKey]*actionHandle),
-		files:             make(map[span.URI]source.VersionedFileHandle),
-		goFiles:           make(map[parseKey]*parseGoHandle),
-		workspacePackages: make(map[packageID]packagePath),
-		unloadableFiles:   make(map[span.URI]struct{}),
-		parseModHandles:   make(map[span.URI]*parseModHandle),
-		modTidyHandle:     s.modTidyHandle,
-		modUpgradeHandle:  s.modUpgradeHandle,
-		modWhyHandle:      s.modWhyHandle,
+		id:                   s.id + 1,
+		view:                 s.view,
+		builtin:              s.builtin,
+		ids:                  make(map[span.URI][]packageID),
+		importedBy:           make(map[packageID][]packageID),
+		metadata:             make(map[packageID]*metadata),
+		packages:             make(map[packageKey]*packageHandle),
+		actions:              make(map[actionKey]*actionHandle),
+		files:                make(map[span.URI]source.VersionedFileHandle),
+		goFiles:              make(map[parseKey]*parseGoHandle),
+		workspaceDirectories: make(map[span.URI]struct{}),
+		workspacePackages:    make(map[packageID]packagePath),
+		unloadableFiles:      make(map[span.URI]struct{}),
+		parseModHandles:      make(map[span.URI]*parseModHandle),
+		modTidyHandle:        s.modTidyHandle,
+		modUpgradeHandle:     s.modUpgradeHandle,
+		modWhyHandle:         s.modWhyHandle,
 	}
 
 	// Copy all of the FileHandles.
@@ -808,6 +824,10 @@
 	for k, v := range s.parseModHandles {
 		result.parseModHandles[k] = v
 	}
+	// Copy all of the workspace directories. They may be reset later.
+	for k, v := range s.workspaceDirectories {
+		result.workspaceDirectories[k] = v
+	}
 
 	for k, v := range s.goFiles {
 		if _, ok := withoutURIs[k.file.URI]; ok {
@@ -850,7 +870,15 @@
 					directIDs[k.id] = struct{}{}
 				}
 			}
+
 			delete(result.parseModHandles, withoutURI)
+
+			if currentFH.URI() == s.view.modURI {
+				// The go.mod's replace directives may have changed. We may
+				// need to update our set of workspace directories. Use the new
+				// snapshot, as it can be locked without causing issues.
+				result.workspaceDirectories = result.findWorkspaceDirectories(ctx, currentFH)
+			}
 		}
 
 		// If this is a file we don't yet know about,
@@ -1033,6 +1061,45 @@
 	return false
 }
 
+// findWorkspaceDirectoriesLocked returns all of the directories that are
+// considered to be part of the view's workspace. For GOPATH workspaces, this
+// is just the view's root. For modules-based workspaces, this is the module
+// root and any replace targets. It also returns the parseModHandle for the
+// view's go.mod file if it has one.
+//
+// It assumes that the file handle is the view's go.mod file, if it has one.
+// The caller need not be holding the snapshot's mutex, but it might be.
+func (s *snapshot) findWorkspaceDirectories(ctx context.Context, modFH source.FileHandle) map[span.URI]struct{} {
+	m := map[span.URI]struct{}{
+		s.view.root: {},
+	}
+	// If the view does not have a go.mod file, only the root directory
+	// is known. In GOPATH mode, we should really watch the entire GOPATH,
+	// but that's too expensive.
+	modURI := s.view.modURI
+	if modURI == "" {
+		return m
+	}
+	if modFH == nil {
+		return m
+	}
+	// Ignore parse errors. An invalid go.mod is not fatal.
+	mod, err := s.ParseMod(ctx, modFH)
+	if err != nil {
+		return m
+	}
+	for _, r := range mod.File.Replace {
+		// We may be replacing a module with a different version. not a path
+		// on disk.
+		if r.New.Version != "" {
+			continue
+		}
+		uri := span.URIFromPath(r.New.Path)
+		m[uri] = struct{}{}
+	}
+	return m
+}
+
 func (s *snapshot) BuiltinPackage(ctx context.Context) (*source.BuiltinPackage, error) {
 	s.view.awaitInitialized(ctx)
 
diff --git a/internal/lsp/fake/client.go b/internal/lsp/fake/client.go
index e053a3d..2133621 100644
--- a/internal/lsp/fake/client.go
+++ b/internal/lsp/fake/client.go
@@ -19,6 +19,8 @@
 	OnProgress               func(context.Context, *protocol.ProgressParams) error
 	OnShowMessage            func(context.Context, *protocol.ShowMessageParams) error
 	OnShowMessageRequest     func(context.Context, *protocol.ShowMessageRequestParams) error
+	OnRegistration           func(context.Context, *protocol.RegistrationParams) error
+	OnUnregistration         func(context.Context, *protocol.UnregistrationParams) error
 }
 
 // Client is an adapter that converts an *Editor into an LSP Client. It mosly
@@ -80,11 +82,17 @@
 	return results, nil
 }
 
-func (c *Client) RegisterCapability(context.Context, *protocol.RegistrationParams) error {
+func (c *Client) RegisterCapability(ctx context.Context, params *protocol.RegistrationParams) error {
+	if c.hooks.OnRegistration != nil {
+		return c.hooks.OnRegistration(ctx, params)
+	}
 	return nil
 }
 
-func (c *Client) UnregisterCapability(context.Context, *protocol.UnregistrationParams) error {
+func (c *Client) UnregisterCapability(ctx context.Context, params *protocol.UnregistrationParams) error {
+	if c.hooks.OnUnregistration != nil {
+		return c.hooks.OnUnregistration(ctx, params)
+	}
 	return nil
 }
 
diff --git a/internal/lsp/fake/editor.go b/internal/lsp/fake/editor.go
index 32771e8..a0ae0ed 100644
--- a/internal/lsp/fake/editor.go
+++ b/internal/lsp/fake/editor.go
@@ -202,6 +202,12 @@
 	// TODO: set client capabilities
 	params.InitializationOptions = e.configuration()
 
+	// This is a bit of a hack, since the fake editor doesn't actually support
+	// watching changed files that match a specific glob pattern. However, the
+	// editor does send didChangeWatchedFiles notifications, so set this to
+	// true.
+	params.Capabilities.Workspace.DidChangeWatchedFiles.DynamicRegistration = true
+
 	params.Trace = "messages"
 	// TODO: support workspace folders.
 	if e.Server != nil {
diff --git a/internal/lsp/general.go b/internal/lsp/general.go
index ed9a74e..3aeab58 100644
--- a/internal/lsp/general.go
+++ b/internal/lsp/general.go
@@ -169,29 +169,11 @@
 	}
 	s.pendingFolders = nil
 
-	if options.DynamicWatchedFilesSupported {
-		for _, view := range s.session.Views() {
-			dirs, err := view.WorkspaceDirectories(ctx)
-			if err != nil {
-				return err
-			}
-			for _, dir := range dirs {
-				registrations = append(registrations, protocol.Registration{
-					ID:     "workspace/didChangeWatchedFiles",
-					Method: "workspace/didChangeWatchedFiles",
-					RegisterOptions: protocol.DidChangeWatchedFilesRegistrationOptions{
-						Watchers: []protocol.FileSystemWatcher{{
-							GlobPattern: fmt.Sprintf("%s/**/*.{go,mod,sum}", dir),
-							Kind:        float64(protocol.WatchChange + protocol.WatchDelete + protocol.WatchCreate),
-						}},
-					},
-				})
-			}
-		}
-		if len(registrations) > 0 {
-			s.client.RegisterCapability(ctx, &protocol.RegistrationParams{
-				Registrations: registrations,
-			})
+	if len(registrations) > 0 {
+		if err := s.client.RegisterCapability(ctx, &protocol.RegistrationParams{
+			Registrations: registrations,
+		}); err != nil {
+			return err
 		}
 	}
 	return nil
@@ -211,6 +193,7 @@
 			}()
 		}()
 	}
+	dirsToWatch := map[span.URI]struct{}{}
 	for _, folder := range folders {
 		uri := span.URIFromURI(folder.URI)
 		view, snapshot, release, err := s.addView(ctx, folder.Name, uri)
@@ -218,6 +201,10 @@
 			viewErrors[uri] = err
 			continue
 		}
+		for _, dir := range snapshot.WorkspaceDirectories(ctx) {
+			dirsToWatch[dir] = struct{}{}
+		}
+
 		// Print each view's environment.
 		buf := &bytes.Buffer{}
 		if err := view.WriteEnv(ctx, buf); err != nil {
@@ -234,6 +221,13 @@
 			wg.Done()
 		}()
 	}
+	// Register for file watching notifications, if they are supported.
+	s.watchedDirectoriesMu.Lock()
+	err := s.registerWatchedDirectoriesLocked(ctx, dirsToWatch)
+	s.watchedDirectoriesMu.Unlock()
+	if err != nil {
+		return err
+	}
 	if len(viewErrors) > 0 {
 		errMsg := fmt.Sprintf("Error loading workspace folders (expected %v, got %v)\n", len(folders), len(s.session.Views())-originalViews)
 		for uri, err := range viewErrors {
@@ -247,6 +241,113 @@
 	return nil
 }
 
+// updateWatchedDirectories compares the current set of directories to watch
+// with the previously registered set of directories. If the set of directories
+// has changed, we unregister and re-register for file watching notifications.
+// updatedSnapshots is the set of snapshots that have been updated.
+func (s *Server) updateWatchedDirectories(ctx context.Context, updatedSnapshots []source.Snapshot) error {
+	dirsToWatch := map[span.URI]struct{}{}
+	seenViews := map[source.View]struct{}{}
+
+	// Collect all of the workspace directories from the updated snapshots.
+	for _, snapshot := range updatedSnapshots {
+		seenViews[snapshot.View()] = struct{}{}
+		for _, dir := range snapshot.WorkspaceDirectories(ctx) {
+			dirsToWatch[dir] = struct{}{}
+		}
+	}
+	// Not all views were necessarily updated, so check the remaining views.
+	for _, view := range s.session.Views() {
+		if _, ok := seenViews[view]; ok {
+			continue
+		}
+		snapshot, release := view.Snapshot()
+		for _, dir := range snapshot.WorkspaceDirectories(ctx) {
+			dirsToWatch[dir] = struct{}{}
+		}
+		release()
+	}
+
+	s.watchedDirectoriesMu.Lock()
+	defer s.watchedDirectoriesMu.Unlock()
+
+	// Nothing to do if the set of workspace directories is unchanged.
+	if equalURISet(s.watchedDirectories, dirsToWatch) {
+		return nil
+	}
+
+	// If the set of directories to watch has changed, register the updates and
+	// unregister the previously watched directories. This ordering avoids a
+	// period where no files are being watched. Still, if a user makes on-disk
+	// changes before these updates are complete, we may miss them for the new
+	// directories.
+	if s.watchRegistrationCount > 0 {
+		prevID := s.watchRegistrationCount - 1
+		if err := s.registerWatchedDirectoriesLocked(ctx, dirsToWatch); err != nil {
+			return err
+		}
+		return s.client.UnregisterCapability(ctx, &protocol.UnregistrationParams{
+			Unregisterations: []protocol.Unregistration{{
+				ID:     watchedFilesCapabilityID(prevID),
+				Method: "workspace/didChangeWatchedFiles",
+			}},
+		})
+	}
+	return nil
+}
+
+func watchedFilesCapabilityID(id uint64) string {
+	return fmt.Sprintf("workspace/didChangeWatchedFiles-%d", id)
+}
+
+func equalURISet(m1, m2 map[span.URI]struct{}) bool {
+	if len(m1) != len(m2) {
+		return false
+	}
+	for k := range m1 {
+		_, ok := m2[k]
+		if !ok {
+			return false
+		}
+	}
+	return true
+}
+
+// registerWatchedDirectoriesLocked sends the workspace/didChangeWatchedFiles
+// registrations to the client and updates s.watchedDirectories.
+func (s *Server) registerWatchedDirectoriesLocked(ctx context.Context, dirs map[span.URI]struct{}) error {
+	if !s.session.Options().DynamicWatchedFilesSupported {
+		return nil
+	}
+	for k := range s.watchedDirectories {
+		delete(s.watchedDirectories, k)
+	}
+	var watchers []protocol.FileSystemWatcher
+	for dir := range dirs {
+		watchers = append(watchers, protocol.FileSystemWatcher{
+			GlobPattern: fmt.Sprintf("%s/**/*.{go,mod,sum}", dir),
+			Kind:        float64(protocol.WatchChange + protocol.WatchDelete + protocol.WatchCreate),
+		})
+	}
+	if err := s.client.RegisterCapability(ctx, &protocol.RegistrationParams{
+		Registrations: []protocol.Registration{{
+			ID:     watchedFilesCapabilityID(s.watchRegistrationCount),
+			Method: "workspace/didChangeWatchedFiles",
+			RegisterOptions: protocol.DidChangeWatchedFilesRegistrationOptions{
+				Watchers: watchers,
+			},
+		}},
+	}); err != nil {
+		return err
+	}
+	s.watchRegistrationCount++
+
+	for dir := range dirs {
+		s.watchedDirectories[dir] = struct{}{}
+	}
+	return nil
+}
+
 func (s *Server) fetchConfig(ctx context.Context, name string, folder span.URI, o *source.Options) error {
 	if !s.session.Options().ConfigurationSupported {
 		return nil
diff --git a/internal/lsp/regtest/env.go b/internal/lsp/regtest/env.go
index 0232558..6fbd877 100644
--- a/internal/lsp/regtest/env.go
+++ b/internal/lsp/regtest/env.go
@@ -47,6 +47,10 @@
 	logs               []*protocol.LogMessageParams
 	showMessage        []*protocol.ShowMessageParams
 	showMessageRequest []*protocol.ShowMessageRequestParams
+
+	registrations   []*protocol.RegistrationParams
+	unregistrations []*protocol.UnregistrationParams
+
 	// outstandingWork is a map of token->work summary. All tokens are assumed to
 	// be string, though the spec allows for numeric tokens as well.  When work
 	// completes, it is deleted from this map.
@@ -129,6 +133,8 @@
 			OnProgress:               env.onProgress,
 			OnShowMessage:            env.onShowMessage,
 			OnShowMessageRequest:     env.onShowMessageRequest,
+			OnRegistration:           env.onRegistration,
+			OnUnregistration:         env.onUnregistration,
 		}
 	}
 	editor, err := fake.NewEditor(sandbox, editorConfig).Connect(ctx, conn, hooks)
@@ -210,6 +216,24 @@
 	return nil
 }
 
+func (e *Env) onRegistration(_ context.Context, m *protocol.RegistrationParams) error {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+
+	e.state.registrations = append(e.state.registrations, m)
+	e.checkConditionsLocked()
+	return nil
+}
+
+func (e *Env) onUnregistration(_ context.Context, m *protocol.UnregistrationParams) error {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+
+	e.state.unregistrations = append(e.state.unregistrations, m)
+	e.checkConditionsLocked()
+	return nil
+}
+
 func (e *Env) checkConditionsLocked() {
 	for id, condition := range e.waiters {
 		if v, _, _ := checkExpectations(e.state, condition.expectations); v != Unmet {
@@ -496,6 +520,86 @@
 	}
 }
 
+// RegistrationExpectation is an expectation on the capability registrations
+// received by the editor from gopls.
+type RegistrationExpectation struct {
+	check       func([]*protocol.RegistrationParams) (Verdict, interface{})
+	description string
+}
+
+// Check implements the Expectation interface.
+func (e RegistrationExpectation) Check(s State) (Verdict, interface{}) {
+	return e.check(s.registrations)
+}
+
+// Description implements the Expectation interface.
+func (e RegistrationExpectation) Description() string {
+	return e.description
+}
+
+// RegistrationMatching asserts that the client has received a capability
+// registration matching the given regexp.
+func RegistrationMatching(re string) RegistrationExpectation {
+	rec, err := regexp.Compile(re)
+	if err != nil {
+		panic(err)
+	}
+	check := func(params []*protocol.RegistrationParams) (Verdict, interface{}) {
+		for _, p := range params {
+			for _, r := range p.Registrations {
+				if rec.Match([]byte(r.Method)) {
+					return Met, r
+				}
+			}
+		}
+		return Unmet, nil
+	}
+	return RegistrationExpectation{
+		check:       check,
+		description: fmt.Sprintf("registration matching %q", re),
+	}
+}
+
+// UnregistrationExpectation is an expectation on the capability
+// unregistrations received by the editor from gopls.
+type UnregistrationExpectation struct {
+	check       func([]*protocol.UnregistrationParams) (Verdict, interface{})
+	description string
+}
+
+// Check implements the Expectation interface.
+func (e UnregistrationExpectation) Check(s State) (Verdict, interface{}) {
+	return e.check(s.unregistrations)
+}
+
+// Description implements the Expectation interface.
+func (e UnregistrationExpectation) Description() string {
+	return e.description
+}
+
+// UnregistrationMatching asserts that the client has received an
+// unregistration whose ID matches the given regexp.
+func UnregistrationMatching(re string) UnregistrationExpectation {
+	rec, err := regexp.Compile(re)
+	if err != nil {
+		panic(err)
+	}
+	check := func(params []*protocol.UnregistrationParams) (Verdict, interface{}) {
+		for _, p := range params {
+			for _, r := range p.Unregisterations {
+				if rec.Match([]byte(r.Method)) {
+					return Met, r
+				}
+			}
+		}
+		return Unmet, nil
+	}
+	return UnregistrationExpectation{
+		check:       check,
+		description: fmt.Sprintf("unregistration matching %q", re),
+	}
+}
+
 // A DiagnosticExpectation is a condition that must be met by the current set
 // of diagnostics for a file.
 type DiagnosticExpectation struct {
diff --git a/internal/lsp/regtest/workspace_test.go b/internal/lsp/regtest/workspace_test.go
index a3c54f3..987a4c3 100644
--- a/internal/lsp/regtest/workspace_test.go
+++ b/internal/lsp/regtest/workspace_test.go
@@ -5,7 +5,10 @@
 package regtest
 
 import (
+	"fmt"
 	"testing"
+
+	"golang.org/x/tools/internal/lsp"
 )
 
 const workspaceProxy = `
@@ -19,29 +22,44 @@
 func SaySomething() {
 	fmt.Println("something")
 }
+-- random.org@v1.2.3/go.mod --
+module random.org
+
+go 1.12
+-- random.org@v1.2.3/bye/bye.go --
+package bye
+
+func Goodbye() {
+	println("Bye")
+}
 `
 
 // TODO: Add a replace directive.
 const workspaceModule = `
--- go.mod --
+-- pkg/go.mod --
 module mod.com
 
 go 1.14
 
-require example.com v1.2.3
--- main.go --
+require (
+	example.com v1.2.3
+	random.org v1.2.3
+)
+-- pkg/main.go --
 package main
 
 import (
 	"example.com/blah"
 	"mod.com/inner"
+	"random.org/bye"
 )
 
 func main() {
 	blah.SaySomething()
 	inner.Hi()
+	bye.Goodbye()
 }
--- main2.go --
+-- pkg/main2.go --
 package main
 
 import "fmt"
@@ -49,7 +67,7 @@
 func _() {
 	fmt.Print("%s")
 }
--- inner/inner.go --
+-- pkg/inner/inner.go --
 package inner
 
 import "example.com/blah"
@@ -57,6 +75,14 @@
 func Hi() {
 	blah.SaySomething()
 }
+-- goodbye/bye/bye.go --
+package bye
+
+func Bye() {}
+-- goodbye/go.mod --
+module random.org
+
+go 1.12
 `
 
 // Confirm that find references returns all of the references in the module,
@@ -66,11 +92,12 @@
 		name, rootPath string
 	}{
 		{
-			name: "module root",
+			name:     "module root",
+			rootPath: "pkg",
 		},
 		{
 			name:     "subdirectory",
-			rootPath: "inner",
+			rootPath: "pkg/inner",
 		},
 	} {
 		t.Run(tt.name, func(t *testing.T) {
@@ -79,8 +106,8 @@
 				opts = append(opts, WithRootPath(tt.rootPath))
 			}
 			withOptions(opts...).run(t, workspaceModule, func(t *testing.T, env *Env) {
-				env.OpenFile("inner/inner.go")
-				locations := env.ReferencesAtRegexp("inner/inner.go", "SaySomething")
+				env.OpenFile("pkg/inner/inner.go")
+				locations := env.ReferencesAtRegexp("pkg/inner/inner.go", "SaySomething")
 				want := 3
 				if got := len(locations); got != want {
 					t.Fatalf("expected %v locations, got %v", want, got)
@@ -95,14 +122,36 @@
 // VS Code, where clicking on a reference result triggers a
 // textDocument/didOpen without a corresponding textDocument/didClose.
 func TestClearAnalysisDiagnostics(t *testing.T) {
-	withOptions(WithProxyFiles(workspaceProxy), WithRootPath("inner")).run(t, workspaceModule, func(t *testing.T, env *Env) {
-		env.OpenFile("main.go")
+	withOptions(WithProxyFiles(workspaceProxy), WithRootPath("pkg/inner")).run(t, workspaceModule, func(t *testing.T, env *Env) {
+		env.OpenFile("pkg/main.go")
 		env.Await(
-			env.DiagnosticAtRegexp("main2.go", "fmt.Print"),
+			env.DiagnosticAtRegexp("pkg/main2.go", "fmt.Print"),
 		)
-		env.CloseBuffer("main.go")
+		env.CloseBuffer("pkg/main.go")
 		env.Await(
-			EmptyDiagnostics("main2.go"),
+			EmptyDiagnostics("pkg/main2.go"),
+		)
+	})
+}
+
+// This test checks that gopls updates the set of files it watches when a
+// replace target is added to the go.mod.
+func TestWatchReplaceTargets(t *testing.T) {
+	withOptions(WithProxyFiles(workspaceProxy), WithRootPath("pkg")).run(t, workspaceModule, func(t *testing.T, env *Env) {
+		env.Await(
+			CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromInitialWorkspaceLoad), 1),
+		)
+		// Add a replace directive and expect the files that gopls is watching
+		// to change.
+		dir := env.Sandbox.Workdir.URI("goodbye").SpanURI().Filename()
+		goModWithReplace := fmt.Sprintf(`%s
+replace random.org => %s
+`, env.ReadWorkspaceFile("pkg/go.mod"), dir)
+		env.WriteWorkspaceFile("pkg/go.mod", goModWithReplace)
+		env.Await(
+			CompletedWork(lsp.DiagnosticWorkTitle(lsp.FromDidChangeWatchedFiles), 1),
+			UnregistrationMatching("didChangeWatchedFiles"),
+			RegistrationMatching("didChangeWatchedFiles"),
 		)
 	})
 }
diff --git a/internal/lsp/server.go b/internal/lsp/server.go
index 88f44ef..dda2100 100644
--- a/internal/lsp/server.go
+++ b/internal/lsp/server.go
@@ -25,6 +25,7 @@
 	return &Server{
 		delivered:            make(map[span.URI]sentDiagnostics),
 		gcOptimizatonDetails: make(map[span.URI]struct{}),
+		watchedDirectories:   make(map[span.URI]struct{}),
 		session:              session,
 		client:               client,
 		diagnosticsSema:      make(chan struct{}, concurrentAnalyses),
@@ -71,6 +72,13 @@
 	// set of folders to build views for when we are ready
 	pendingFolders []protocol.WorkspaceFolder
 
+	// watchedDirectories is the set of directories that we have requested that
+	// the client watch on disk. It will be updated as the set of directories
+	// that the server should watch changes.
+	watchedDirectoriesMu   sync.Mutex
+	watchedDirectories     map[span.URI]struct{}
+	watchRegistrationCount uint64
+
 	// delivered is a cache of the diagnostics that the server has sent.
 	deliveredMu sync.Mutex
 	delivered   map[span.URI]sentDiagnostics
diff --git a/internal/lsp/source/view.go b/internal/lsp/source/view.go
index 37c72c6..9494f94 100644
--- a/internal/lsp/source/view.go
+++ b/internal/lsp/source/view.go
@@ -110,6 +110,10 @@
 
 	// WorkspacePackages returns the snapshot's top-level packages.
 	WorkspacePackages(ctx context.Context) ([]Package, error)
+
+	// WorkspaceDirectories returns any directory known by the view. For views
+	// within a module, this is the module root and any replace targets.
+	WorkspaceDirectories(ctx context.Context) []span.URI
 }
 
 // View represents a single workspace.
@@ -169,10 +173,6 @@
 	// IgnoredFile reports if a file would be ignored by a `go list` of the whole
 	// workspace.
 	IgnoredFile(uri span.URI) bool
-
-	// WorkspaceDirectories returns any directory known by the view. For views
-	// within a module, this is the module root and any replace targets.
-	WorkspaceDirectories(ctx context.Context) ([]string, error)
 }
 
 type BuiltinPackage struct {
@@ -242,8 +242,8 @@
 	// GetFile returns a handle for the specified file.
 	GetFile(ctx context.Context, uri span.URI) (FileHandle, error)
 
-	// DidModifyFile reports a file modification to the session.
-	// It returns the resulting snapshots, a guaranteed one per view.
+	// DidModifyFile reports a file modification to the session. It returns the
+	// resulting snapshots, a guaranteed one per view.
 	DidModifyFiles(ctx context.Context, changes []FileModification) ([]Snapshot, []func(), []span.URI, error)
 
 	// Overlays returns a slice of file overlays for the session.
diff --git a/internal/lsp/text_synchronization.go b/internal/lsp/text_synchronization.go
index 6f9a166..f5db5b6 100644
--- a/internal/lsp/text_synchronization.go
+++ b/internal/lsp/text_synchronization.go
@@ -308,7 +308,12 @@
 			release()
 		}
 	}()
-
+	// After any file modifications, we need to update our watched files,
+	// in case something changed. Compute the new set of directories to watch,
+	// and if it differs from the current set, send updated registrations.
+	if err := s.updateWatchedDirectories(ctx, snapshots); err != nil {
+		return err
+	}
 	return nil
 }