internal/lsp/cache: use persistent map for storing gofiles in the sna…

…pshot

Use treap (https://en.wikipedia.org/wiki/Treap) as a persistent map to avoid copying s.goFiles across generations.
Maintain an additional s.parseKeysByURIMap to avoid scanning s.goFiles on individual file's content invalidation.

This on average reduces didChange latency on internal codebase from 160ms to 150ms.

In a followup the same approach can be generically extended to avoid copying s.files and s.packages.

Updates golang/go#45686

Change-Id: Ic4a9b3c8fb2b66256f224adf9896ddcaaa6865b1
GitHub-Last-Rev: b211d6c
GitHub-Pull-Request: golang#382

internal/lsp/cache/maps.go

@@ -0,0 +1,112 @@
+// Copyright 2022 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package cache
+
+import (
+	"golang.org/x/tools/internal/persistent"
+	"golang.org/x/tools/internal/span"
+)
+
+// TODO(euroelessar): Use generics once support for go1.17 is dropped.
+
+type goFilesMap struct {
+	impl *persistent.Map
+}
+
+func newGoFilesMap() *goFilesMap {
+	return &goFilesMap{
+		impl: persistent.NewMap(func(a, b interface{}) bool {
+			return parseKeyLess(a.(parseKey), b.(parseKey))
+		}),
+	}
+}
+
+func parseKeyLess(a, b parseKey) bool {
+	if a.mode != b.mode {
+		return a.mode < b.mode
+	}
+	if a.file.Hash != b.file.Hash {
+		return a.file.Hash.Less(b.file.Hash)
+	}
+	return a.file.URI < b.file.URI
+}
+
+func (m *goFilesMap) Clone() *goFilesMap {
+	return &goFilesMap{
+		impl: m.impl.Clone(),
+	}
+}
+
+func (m *goFilesMap) Destroy() {
+	m.impl.Destroy()
+}
+
+func (m *goFilesMap) Load(key parseKey) (*parseGoHandle, bool) {
+	value, ok := m.impl.Load(key)
+	if !ok {
+		return nil, false
+	}
+	return value.(*parseGoHandle), true
+}
+
+func (m *goFilesMap) Range(do func(key parseKey, value *parseGoHandle)) {
+	m.impl.Range(func(key, value interface{}) {
+		do(key.(parseKey), value.(*parseGoHandle))
+	})
+}
+
+func (m *goFilesMap) Store(key parseKey, value *parseGoHandle, release func()) {
+	m.impl.Store(key, value, func(key, value interface{}) {
+		release()
+	})
+}
+
+func (m *goFilesMap) Delete(key parseKey) {
+	m.impl.Delete(key)
+}
+
+type parseKeysByURIMap struct {
+	impl *persistent.Map
+}
+
+func newParseKeysByURIMap() *parseKeysByURIMap {
+	return &parseKeysByURIMap{
+		impl: persistent.NewMap(func(a, b interface{}) bool {
+			return a.(span.URI) < b.(span.URI)
+		}),
+	}
+}
+
+func (m *parseKeysByURIMap) Clone() *parseKeysByURIMap {
+	return &parseKeysByURIMap{
+		impl: m.impl.Clone(),
+	}
+}
+
+func (m *parseKeysByURIMap) Destroy() {
+	m.impl.Destroy()
+}
+
+func (m *parseKeysByURIMap) Load(key span.URI) ([]parseKey, bool) {
+	value, ok := m.impl.Load(key)
+	if !ok {
+		return nil, false
+	}
+	return value.([]parseKey), true
+}
+
+func (m *parseKeysByURIMap) Range(do func(key span.URI, value []parseKey)) {
+	m.impl.Range(func(key, value interface{}) {
+		do(key.(span.URI), value.([]parseKey))
+	})
+}
+
+func (m *parseKeysByURIMap) Store(key span.URI, value []parseKey) {
+	m.impl.Store(key, value, nil)
+}
+
+func (m *parseKeysByURIMap) Delete(key span.URI) {
+	m.impl.Delete(key)
+}

internal/lsp/cache/parse.go

@@ -58,7 +58,7 @@ func (s *snapshot) parseGoHandle(ctx context.Context, fh source.FileHandle, mode
 	if pgh := s.getGoFile(key); pgh != nil {
 		return pgh
 	}
-	parseHandle := s.generation.Bind(key, func(ctx context.Context, arg memoize.Arg) interface{} {
+	parseHandle, release := s.generation.NewHandle(key, func(ctx context.Context, arg memoize.Arg) interface{} {
 		snapshot := arg.(*snapshot)
 		return parseGo(ctx, snapshot.FileSet(), fh, mode)
 	}, nil)
@@ -68,7 +68,7 @@ func (s *snapshot) parseGoHandle(ctx context.Context, fh source.FileHandle, mode
 		file:   fh,
 		mode:   mode,
 	}
-	return s.addGoFile(key, pgh)
+	return s.addGoFile(key, pgh, release)
 }
 
 func (pgh *parseGoHandle) String() string {

internal/lsp/cache/session.go

@@ -234,7 +234,8 @@ func (s *Session) createView(ctx context.Context, name string, folder span.URI,
 		packages:          make(map[packageKey]*packageHandle),
 		meta:              &metadataGraph{},
 		files:             make(map[span.URI]source.VersionedFileHandle),
-		goFiles:           newGoFileMap(),
+		goFiles:           newGoFilesMap(),
+		parseKeysByURI:    newParseKeysByURIMap(),
 		symbols:           make(map[span.URI]*symbolHandle),
 		actions:           make(map[actionKey]*actionHandle),
 		workspacePackages: make(map[PackageID]PackagePath),

internal/lsp/cache/snapshot.go

@@ -77,7 +77,8 @@ type snapshot struct {
 	files map[span.URI]source.VersionedFileHandle
 
 	// goFiles maps a parseKey to its parseGoHandle.
-	goFiles *goFileMap
+	goFiles        *goFilesMap
+	parseKeysByURI *parseKeysByURIMap
 
 	// TODO(rfindley): consider merging this with files to reduce burden on clone.
 	symbols map[span.URI]*symbolHandle
@@ -133,6 +134,12 @@ type actionKey struct {
 	analyzer *analysis.Analyzer
 }
 
+func (s *snapshot) Destroy(destroyedBy string) {
+	s.generation.Destroy(destroyedBy)
+	s.goFiles.Destroy()
+	s.parseKeysByURI.Destroy()
+}
+
 func (s *snapshot) ID() uint64 {
 	return s.id
 }
@@ -665,17 +672,23 @@ func (s *snapshot) transitiveReverseDependencies(id PackageID, ids map[PackageID
 func (s *snapshot) getGoFile(key parseKey) *parseGoHandle {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-	return s.goFiles.get(key)
+	if result, ok := s.goFiles.Load(key); ok {
+		return result
+	}
+	return nil
 }
 
-func (s *snapshot) addGoFile(key parseKey, pgh *parseGoHandle) *parseGoHandle {
+func (s *snapshot) addGoFile(key parseKey, pgh *parseGoHandle, release func()) *parseGoHandle {
 	s.mu.Lock()
 	defer s.mu.Unlock()
-
-	if prev := s.goFiles.get(key); prev != nil {
-		return prev
-	}
-	s.goFiles.set(key, pgh)
+	if result, ok := s.goFiles.Load(key); ok {
+		release()
+		return result
+	}
+	s.goFiles.Store(key, pgh, release)
+	keys, _ := s.parseKeysByURI.Load(key.file.URI)
+	keys = append([]parseKey{key}, keys...)
+	s.parseKeysByURI.Store(key.file.URI, keys)
 	return pgh
 }
 
@@ -1661,6 +1674,9 @@ func (ac *unappliedChanges) GetFile(ctx context.Context, uri span.URI) (source.F
 }
 
 func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]*fileChange, forceReloadMetadata bool) *snapshot {
+	ctx, done := event.Start(ctx, "snapshot.clone")
+	defer done()
+
 	var vendorChanged bool
 	newWorkspace, workspaceChanged, workspaceReload := s.workspace.invalidate(ctx, changes, &unappliedChanges{
 		originalSnapshot: s,
@@ -1684,7 +1700,8 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]*fileC
 		packages:          make(map[packageKey]*packageHandle, len(s.packages)),
 		actions:           make(map[actionKey]*actionHandle, len(s.actions)),
 		files:             make(map[span.URI]source.VersionedFileHandle, len(s.files)),
-		goFiles:           s.goFiles.clone(),
+		goFiles:           s.goFiles.Clone(),
+		parseKeysByURI:    s.parseKeysByURI.Clone(),
 		symbols:           make(map[span.URI]*symbolHandle, len(s.symbols)),
 		workspacePackages: make(map[PackageID]PackagePath, len(s.workspacePackages)),
 		unloadableFiles:   make(map[span.URI]struct{}, len(s.unloadableFiles)),
@@ -1729,27 +1746,14 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]*fileC
 		result.parseWorkHandles[k] = v
 	}
 
-	// Copy the handles of all Go source files.
-	// There may be tens of thousands of files,
-	// but changes are typically few, so we
-	// use a striped map optimized for this case
-	// and visit its stripes in parallel.
-	var (
-		toDeleteMu sync.Mutex
-		toDelete   []parseKey
-	)
-	s.goFiles.forEachConcurrent(func(k parseKey, v *parseGoHandle) {
-		if changes[k.file.URI] == nil {
-			// no change (common case)
-			newGen.Inherit(v.handle)
-		} else {
-			toDeleteMu.Lock()
-			toDelete = append(toDelete, k)
-			toDeleteMu.Unlock()
+	for uri := range changes {
+		keys, ok := result.parseKeysByURI.Load(uri)
+		if ok {
+			for _, key := range keys {
+				result.goFiles.Delete(key)
+			}
+			result.parseKeysByURI.Delete(uri)
 		}
-	})
-	for _, k := range toDelete {
-		result.goFiles.delete(k)
 	}
 
 	// Copy all of the go.mod-related handles. They may be invalidated later,
@@ -2206,7 +2210,7 @@ func metadataChanges(ctx context.Context, lockedSnapshot *snapshot, oldFH, newFH
 // lockedSnapshot must be locked.
 func peekOrParse(ctx context.Context, lockedSnapshot *snapshot, fh source.FileHandle, mode source.ParseMode) (*source.ParsedGoFile, error) {
 	key := parseKey{file: fh.FileIdentity(), mode: mode}
-	if pgh := lockedSnapshot.goFiles.get(key); pgh != nil {
+	if pgh, ok := lockedSnapshot.goFiles.Load(key); ok {
 		cached := pgh.handle.Cached(lockedSnapshot.generation)
 		if cached != nil {
 			cached := cached.(*parseGoData)
@@ -2494,89 +2498,3 @@ func readGoSum(dst map[module.Version][]string, file string, data []byte) error
 	}
 	return nil
 }
-
-// -- goFileMap --
-
-// A goFileMap is conceptually a map[parseKey]*parseGoHandle,
-// optimized for cloning all or nearly all entries.
-type goFileMap struct {
-	// The map is represented as a map of 256 stripes, one per
-	// distinct value of the top 8 bits of key.file.Hash.
-	// Each stripe has an associated boolean indicating whether it
-	// is shared, and thus immutable, and thus must be copied before any update.
-	// (The bits could be packed but it hasn't been worth it yet.)
-	stripes   [256]map[parseKey]*parseGoHandle
-	exclusive [256]bool // exclusive[i] means stripe[i] is not shared and may be safely mutated
-}
-
-// newGoFileMap returns a new empty goFileMap.
-func newGoFileMap() *goFileMap {
-	return new(goFileMap) // all stripes are shared (non-exclusive) nil maps
-}
-
-// clone returns a copy of m.
-// For concurrency, it counts as an update to m.
-func (m *goFileMap) clone() *goFileMap {
-	m.exclusive = [256]bool{} // original and copy are now nonexclusive
-	copy := *m
-	return &copy
-}
-
-// get returns the value for key k.
-func (m *goFileMap) get(k parseKey) *parseGoHandle {
-	return m.stripes[m.hash(k)][k]
-}
-
-// set updates the value for key k to v.
-func (m *goFileMap) set(k parseKey, v *parseGoHandle) {
-	m.unshare(k)[k] = v
-}
-
-// delete deletes the value for key k, if any.
-func (m *goFileMap) delete(k parseKey) {
-	// TODO(adonovan): opt?: skip unshare if k isn't present.
-	delete(m.unshare(k), k)
-}
-
-// forEachConcurrent calls f for each entry in the map.
-// Calls may be concurrent.
-// f must not modify m.
-func (m *goFileMap) forEachConcurrent(f func(parseKey, *parseGoHandle)) {
-	// Visit stripes in parallel chunks.
-	const p = 16 // concurrency level
-	var wg sync.WaitGroup
-	wg.Add(p)
-	for i := 0; i < p; i++ {
-		chunk := m.stripes[i*p : (i+1)*p]
-		go func() {
-			for _, stripe := range chunk {
-				for k, v := range stripe {
-					f(k, v)
-				}
-			}
-			wg.Done()
-		}()
-	}
-	wg.Wait()
-}
-
-// -- internal--
-
-// hash returns 8 bits from the key's file digest.
-func (*goFileMap) hash(k parseKey) byte { return k.file.Hash[0] }
-
-// unshare makes k's stripe exclusive, allocating a copy if needed, and returns it.
-func (m *goFileMap) unshare(k parseKey) map[parseKey]*parseGoHandle {
-	i := m.hash(k)
-	if !m.exclusive[i] {
-		m.exclusive[i] = true
-
-		// Copy the map.
-		copy := make(map[parseKey]*parseGoHandle, len(m.stripes[i]))
-		for k, v := range m.stripes[i] {
-			copy[k] = v
-		}
-		m.stripes[i] = copy
-	}
-	return m.stripes[i]
-}

internal/lsp/cache/view.go

@@ -524,7 +524,7 @@ func (v *View) shutdown(ctx context.Context) {
 	v.mu.Unlock()
 	v.snapshotMu.Lock()
 	if v.snapshot != nil {
-		go v.snapshot.generation.Destroy("View.shutdown")
+		go v.snapshot.Destroy("View.shutdown")
 		v.snapshot = nil
 	}
 	v.snapshotMu.Unlock()
@@ -721,7 +721,7 @@ func (v *View) invalidateContent(ctx context.Context, changes map[span.URI]*file
 	oldSnapshot := v.snapshot
 
 	v.snapshot = oldSnapshot.clone(ctx, v.baseCtx, changes, forceReloadMetadata)
-	go oldSnapshot.generation.Destroy("View.invalidateContent")
+	go oldSnapshot.Destroy("View.invalidateContent")
 
 	return v.snapshot, v.snapshot.generation.Acquire()
 }

internal/lsp/source/view.go

@@ -551,6 +551,11 @@ func (h Hash) String() string {
 	return fmt.Sprintf("%64x", [sha256.Size]byte(h))
 }
 
+// Less returns true if the given hash is less than the other.
+func (h Hash) Less(other Hash) bool {
+	return bytes.Compare(h[:], other[:]) < 0
+}
+
 // FileIdentity uniquely identifies a file at a version from a FileSystem.
 type FileIdentity struct {
 	URI  span.URI