go/token: clear cache after grabbing the mutex in RemoveFile

This fixes a race, that was possible to hit in RemoveFile.
The file cache could have been populated concurrently just before
grabbing of the mutex.

Change-Id: I6a6a696452d8bd79ac4ac6574297390978353444
Reviewed-on: https://go-review.googlesource.com/c/go/+/705756
Reviewed-by: Junyang Shao <shaojunyang@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Alan Donovan <adonovan@google.com>
Auto-Submit: Alan Donovan <adonovan@google.com>

src/go/token/position.go

@@ -531,11 +531,11 @@ func (s *FileSet) AddExistingFiles(files ...*File) {
 //
 // Removing a file that does not belong to the set has no effect.
 func (s *FileSet) RemoveFile(file *File) {
-	s.last.CompareAndSwap(file, nil) // clear last file cache
-
 	s.mutex.Lock()
 	defer s.mutex.Unlock()
 
+	s.last.CompareAndSwap(file, nil) // clear last file cache
+
 	pn, _ := s.tree.locate(file.key())
 	if *pn != nil && (*pn).file == file {
 		s.tree.delete(pn)

src/go/token/position_test.go

@@ -579,3 +579,45 @@ func fsetString(fset *FileSet) string {
 	buf.WriteRune('}')
 	return buf.String()
 }
+
+// Test that File() does not return the already removed file, while used concurrently.
+func TestRemoveFileRace(t *testing.T) {
+	fset := NewFileSet()
+
+	// Create bunch of files.
+	var files []*File
+	for i := range 20000 {
+		f := fset.AddFile("f", -1, (i+1)*10)
+		files = append(files, f)
+	}
+
+	// governor goroutine
+	race1, race2 := make(chan *File), make(chan *File)
+	start := make(chan struct{})
+	go func() {
+		for _, f := range files {
+			<-start
+			race1 <- f
+			race2 <- f
+		}
+		<-start // unlock main test goroutine
+		close(race1)
+		close(race2)
+	}()
+
+	go func() {
+		for f := range race1 {
+			fset.File(Pos(f.Base()) + 5) // populates s.last with f
+		}
+	}()
+
+	start <- struct{}{}
+	for f := range race2 {
+		fset.RemoveFile(f)
+		got := fset.File(Pos(f.Base()) + 5)
+		if got != nil {
+			t.Fatalf("file was not removed correctly")
+		}
+		start <- struct{}{}
+	}
+}