internal/lsp/cache: persist known subdirs

This on average reduces latency from 12ms to 4ms on internal codebase.

Updates golang/go#45686

Change-Id: Id376fcd97ce375210f2ad8b88e42f6ca283d29d3
Reviewed-on: https://go-review.googlesource.com/c/tools/+/413657
Reviewed-by: Robert Findley <[email protected]>
Reviewed-by: Alan Donovan <[email protected]>

internal/lsp/cache/maps.go

@@ -206,3 +206,48 @@ func (m packagesMap) Set(key packageKey, value *packageHandle, release func()) {
 func (m packagesMap) Delete(key packageKey) {
 	m.impl.Delete(key)
 }
+
+type knownDirsSet struct {
+	impl *persistent.Map
+}
+
+func newKnownDirsSet() knownDirsSet {
+	return knownDirsSet{
+		impl: persistent.NewMap(func(a, b interface{}) bool {
+			return a.(span.URI) < b.(span.URI)
+		}),
+	}
+}
+
+func (s knownDirsSet) Clone() knownDirsSet {
+	return knownDirsSet{
+		impl: s.impl.Clone(),
+	}
+}
+
+func (s knownDirsSet) Destroy() {
+	s.impl.Destroy()
+}
+
+func (s knownDirsSet) Contains(key span.URI) bool {
+	_, ok := s.impl.Get(key)
+	return ok
+}
+
+func (s knownDirsSet) Range(do func(key span.URI)) {
+	s.impl.Range(func(key, value interface{}) {
+		do(key.(span.URI))
+	})
+}
+
+func (s knownDirsSet) SetAll(other knownDirsSet) {
+	s.impl.SetAll(other.impl)
+}
+
+func (s knownDirsSet) Insert(key span.URI) {
+	s.impl.Set(key, nil, nil)
+}
+
+func (s knownDirsSet) Remove(key span.URI) {
+	s.impl.Delete(key)
+}

internal/lsp/cache/session.go

@@ -244,6 +244,7 @@ func (s *Session) createView(ctx context.Context, name string, folder span.URI,
 		parseWorkHandles:  make(map[span.URI]*parseWorkHandle),
 		modTidyHandles:    make(map[span.URI]*modTidyHandle),
 		modWhyHandles:     make(map[span.URI]*modWhyHandle),
+		knownSubdirs:      newKnownDirsSet(),
 		workspace:         workspace,
 	}
 
@@ -537,9 +538,11 @@ func (s *Session) ExpandModificationsToDirectories(ctx context.Context, changes
 		snapshots = append(snapshots, snapshot)
 	}
 	knownDirs := knownDirectories(ctx, snapshots)
+	defer knownDirs.Destroy()
+
 	var result []source.FileModification
 	for _, c := range changes {
-		if _, ok := knownDirs[c.URI]; !ok {
+		if !knownDirs.Contains(c.URI) {
 			result = append(result, c)
 			continue
 		}
@@ -561,16 +564,17 @@ func (s *Session) ExpandModificationsToDirectories(ctx context.Context, changes
 
 // knownDirectories returns all of the directories known to the given
 // snapshots, including workspace directories and their subdirectories.
-func knownDirectories(ctx context.Context, snapshots []*snapshot) map[span.URI]struct{} {
-	result := map[span.URI]struct{}{}
+// It is responsibility of the caller to destroy the returned set.
+func knownDirectories(ctx context.Context, snapshots []*snapshot) knownDirsSet {
+	result := newKnownDirsSet()
 	for _, snapshot := range snapshots {
 		dirs := snapshot.workspace.dirs(ctx, snapshot)
 		for _, dir := range dirs {
-			result[dir] = struct{}{}
-		}
-		for _, dir := range snapshot.getKnownSubdirs(dirs) {
-			result[dir] = struct{}{}
+			result.Insert(dir)
 		}
+		knownSubdirs := snapshot.getKnownSubdirs(dirs)
+		result.SetAll(knownSubdirs)
+		knownSubdirs.Destroy()
 	}
 	return result
 }

internal/lsp/cache/snapshot.go

@@ -123,7 +123,7 @@ type snapshot struct {
 
 	// knownSubdirs is the set of subdirectories in the workspace, used to
 	// create glob patterns for file watching.
-	knownSubdirs             map[span.URI]struct{}
+	knownSubdirs             knownDirsSet
 	knownSubdirsPatternCache string
 	// unprocessedSubdirChanges are any changes that might affect the set of
 	// subdirectories in the workspace. They are not reflected to knownSubdirs
@@ -147,6 +147,7 @@ func (s *snapshot) Destroy(destroyedBy string) {
 	s.files.Destroy()
 	s.goFiles.Destroy()
 	s.parseKeysByURI.Destroy()
+	s.knownSubdirs.Destroy()
 
 	if s.workspaceDir != "" {
 		if err := os.RemoveAll(s.workspaceDir); err != nil {
@@ -842,17 +843,20 @@ func (s *snapshot) getKnownSubdirsPattern(wsDirs []span.URI) string {
 	// It may change list of known subdirs and therefore invalidate the cache.
 	s.applyKnownSubdirsChangesLocked(wsDirs)
 
-	if len(s.knownSubdirs) == 0 {
-		return ""
-	}
-
 	if s.knownSubdirsPatternCache == "" {
-		dirNames := make([]string, 0, len(s.knownSubdirs))
-		for uri := range s.knownSubdirs {
-			dirNames = append(dirNames, uri.Filename())
+		var builder strings.Builder
+		s.knownSubdirs.Range(func(uri span.URI) {
+			if builder.Len() == 0 {
+				builder.WriteString("{")
+			} else {
+				builder.WriteString(",")
+			}
+			builder.WriteString(uri.Filename())
+		})
+		if builder.Len() > 0 {
+			builder.WriteString("}")
+			s.knownSubdirsPatternCache = builder.String()
 		}
-		sort.Strings(dirNames)
-		s.knownSubdirsPatternCache = fmt.Sprintf("{%s}", strings.Join(dirNames, ","))
 	}
 
 	return s.knownSubdirsPatternCache
@@ -867,26 +871,23 @@ func (s *snapshot) collectAllKnownSubdirs(ctx context.Context) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
-	s.knownSubdirs = map[span.URI]struct{}{}
+	s.knownSubdirs.Destroy()
+	s.knownSubdirs = newKnownDirsSet()
 	s.knownSubdirsPatternCache = ""
 	s.files.Range(func(uri span.URI, fh source.VersionedFileHandle) {
 		s.addKnownSubdirLocked(uri, dirs)
 	})
 }
 
-func (s *snapshot) getKnownSubdirs(wsDirs []span.URI) []span.URI {
+func (s *snapshot) getKnownSubdirs(wsDirs []span.URI) knownDirsSet {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
 	// First, process any pending changes and update the set of known
 	// subdirectories.
 	s.applyKnownSubdirsChangesLocked(wsDirs)
 
-	result := make([]span.URI, 0, len(s.knownSubdirs))
-	for uri := range s.knownSubdirs {
-		result = append(result, uri)
-	}
-	return result
+	return s.knownSubdirs.Clone()
 }
 
 func (s *snapshot) applyKnownSubdirsChangesLocked(wsDirs []span.URI) {
@@ -907,7 +908,7 @@ func (s *snapshot) addKnownSubdirLocked(uri span.URI, dirs []span.URI) {
 	dir := filepath.Dir(uri.Filename())
 	// First check if the directory is already known, because then we can
 	// return early.
-	if _, ok := s.knownSubdirs[span.URIFromPath(dir)]; ok {
+	if s.knownSubdirs.Contains(span.URIFromPath(dir)) {
 		return
 	}
 	var matched span.URI
@@ -926,10 +927,10 @@ func (s *snapshot) addKnownSubdirLocked(uri span.URI, dirs []span.URI) {
 			break
 		}
 		uri := span.URIFromPath(dir)
-		if _, ok := s.knownSubdirs[uri]; ok {
+		if s.knownSubdirs.Contains(uri) {
 			break
 		}
-		s.knownSubdirs[uri] = struct{}{}
+		s.knownSubdirs.Insert(uri)
 		dir = filepath.Dir(dir)
 		s.knownSubdirsPatternCache = ""
 	}
@@ -939,11 +940,11 @@ func (s *snapshot) removeKnownSubdirLocked(uri span.URI) {
 	dir := filepath.Dir(uri.Filename())
 	for dir != "" {
 		uri := span.URIFromPath(dir)
-		if _, ok := s.knownSubdirs[uri]; !ok {
+		if !s.knownSubdirs.Contains(uri) {
 			break
 		}
 		if info, _ := os.Stat(dir); info == nil {
-			delete(s.knownSubdirs, uri)
+			s.knownSubdirs.Remove(uri)
 			s.knownSubdirsPatternCache = ""
 		}
 		dir = filepath.Dir(dir)
@@ -1714,7 +1715,7 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]*fileC
 		parseWorkHandles:  make(map[span.URI]*parseWorkHandle, len(s.parseWorkHandles)),
 		modTidyHandles:    make(map[span.URI]*modTidyHandle, len(s.modTidyHandles)),
 		modWhyHandles:     make(map[span.URI]*modWhyHandle, len(s.modWhyHandles)),
-		knownSubdirs:      make(map[span.URI]struct{}, len(s.knownSubdirs)),
+		knownSubdirs:      s.knownSubdirs.Clone(),
 		workspace:         newWorkspace,
 	}
 
@@ -1771,9 +1772,7 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]*fileC
 	// Add all of the known subdirectories, but don't update them for the
 	// changed files. We need to rebuild the workspace module to know the
 	// true set of known subdirectories, but we don't want to do that in clone.
-	for k, v := range s.knownSubdirs {
-		result.knownSubdirs[k] = v
-	}
+	result.knownSubdirs = s.knownSubdirs.Clone()
 	result.knownSubdirsPatternCache = s.knownSubdirsPatternCache
 	for _, c := range changes {
 		result.unprocessedSubdirChanges = append(result.unprocessedSubdirChanges, c)

internal/persistent/map.go

@@ -28,6 +28,8 @@ import (
 // client-provided function that implements a strict weak order.
 //
 // Maps can be Cloned in constant time.
+// Get, Store, and Delete operations are done on average in logarithmic time.
+// Maps can be Updated in O(m log(n/m)) time for maps of size n and m, where m < n.
 //
 // Values are reference counted, and a client-supplied release function
 // is called when a value is no longer referenced by a map or any clone.
@@ -156,6 +158,15 @@ func (pm *Map) Get(key interface{}) (interface{}, bool) {
 	return nil, false
 }
 
+// SetAll updates the map with key/value pairs from the other map, overwriting existing keys.
+// It is equivalent to calling Set for each entry in the other map but is more efficient.
+// Both maps must have the same comparison function, otherwise behavior is undefined.
+func (pm *Map) SetAll(other *Map) {
+	root := pm.root
+	pm.root = union(root, other.root, pm.less, true)
+	root.decref()
+}
+
 // Set updates the value associated with the specified key.
 // If release is non-nil, it will be called with entry's key and value once the
 // key is no longer contained in the map or any clone.

internal/persistent/map_test.go

@@ -151,6 +151,31 @@ func TestRandomMap(t *testing.T) {
 	assertSameMap(t, seenEntries, deletedEntries)
 }
 
+func TestUpdate(t *testing.T) {
+	deletedEntries := make(map[mapEntry]struct{})
+	seenEntries := make(map[mapEntry]struct{})
+
+	m1 := &validatedMap{
+		impl: NewMap(func(a, b interface{}) bool {
+			return a.(int) < b.(int)
+		}),
+		expected: make(map[int]int),
+		deleted:  deletedEntries,
+		seen:     seenEntries,
+	}
+	m2 := m1.clone()
+
+	m1.set(t, 1, 1)
+	m1.set(t, 2, 2)
+	m2.set(t, 2, 20)
+	m2.set(t, 3, 3)
+	m1.setAll(t, m2)
+
+	m1.destroy()
+	m2.destroy()
+	assertSameMap(t, seenEntries, deletedEntries)
+}
+
 func (vm *validatedMap) onDelete(t *testing.T, key, value int) {
 	entry := mapEntry{key: key, value: value}
 	if _, ok := vm.deleted[entry]; ok {
@@ -254,6 +279,14 @@ func validateNode(t *testing.T, node *mapNode, less func(a, b interface{}) bool)
 	validateNode(t, node.right, less)
 }
 
+func (vm *validatedMap) setAll(t *testing.T, other *validatedMap) {
+	vm.impl.SetAll(other.impl)
+	for key, value := range other.expected {
+		vm.expected[key] = value
+	}
+	vm.validate(t)
+}
+
 func (vm *validatedMap) set(t *testing.T, key, value int) {
 	vm.seen[mapEntry{key: key, value: value}] = struct{}{}
 	vm.impl.Set(key, value, func(deletedKey, deletedValue interface{}) {