internal/lsp/cache: use persistent map for storing packages in the snapshot

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

Updates golang/go#45686

Change-Id: I49c8f09f8e54b7b486d7ff7eb8f4ba9f0d90b278
Reviewed-on: https://go-review.googlesource.com/c/tools/+/413655
gopls-CI: kokoro <noreply+kokoro@google.com>
Run-TryBot: Robert Findley <rfindley@google.com>
TryBot-Result: Gopher Robot <gobot@golang.org>
Reviewed-by: Robert Findley <rfindley@google.com>
Reviewed-by: Alan Donovan <adonovan@google.com>

internal/lsp/cache/check.go

@@ -108,7 +108,7 @@ func (s *snapshot) buildPackageHandle(ctx context.Context, id PackageID, mode so
 	m := ph.m
 	key := ph.key
 
-	h := s.generation.Bind(key, func(ctx context.Context, arg memoize.Arg) interface{} {
+	h, release := s.generation.GetHandle(key, func(ctx context.Context, arg memoize.Arg) interface{} {
 		snapshot := arg.(*snapshot)
 
 		// Begin loading the direct dependencies, in parallel.
@@ -128,14 +128,14 @@ func (s *snapshot) buildPackageHandle(ctx context.Context, id PackageID, mode so
 		wg.Wait()
 
 		return data
-	}, nil)
+	})
 	ph.handle = h
 
 	// Cache the handle in the snapshot. If a package handle has already
 	// been cached, addPackage will return the cached value. This is fine,
 	// since the original package handle above will have no references and be
 	// garbage collected.
-	ph = s.addPackageHandle(ph)
+	ph = s.addPackageHandle(ph, release)
 
 	return ph, nil
 }

internal/lsp/cache/load.go

@@ -228,9 +228,9 @@ func (s *snapshot) load(ctx context.Context, allowNetwork bool, scopes ...interf
 	// TODO(rfindley): this should not be necessary, as we should have already
 	// invalidated in snapshot.clone.
 	for id := range invalidatedPackages {
-		for _, mode := range []source.ParseMode{source.ParseHeader, source.ParseExported, source.ParseFull} {
+		for _, mode := range source.AllParseModes {
 			key := packageKey{mode, id}
-			delete(s.packages, key)
+			s.packages.Delete(key)
 		}
 	}
 

internal/lsp/cache/maps.go

@@ -155,3 +155,54 @@ func (m parseKeysByURIMap) Set(key span.URI, value []parseKey) {
 func (m parseKeysByURIMap) Delete(key span.URI) {
 	m.impl.Delete(key)
 }
+
+type packagesMap struct {
+	impl *persistent.Map
+}
+
+func newPackagesMap() packagesMap {
+	return packagesMap{
+		impl: persistent.NewMap(func(a, b interface{}) bool {
+			left := a.(packageKey)
+			right := b.(packageKey)
+			if left.mode != right.mode {
+				return left.mode < right.mode
+			}
+			return left.id < right.id
+		}),
+	}
+}
+
+func (m packagesMap) Clone() packagesMap {
+	return packagesMap{
+		impl: m.impl.Clone(),
+	}
+}
+
+func (m packagesMap) Destroy() {
+	m.impl.Destroy()
+}
+
+func (m packagesMap) Get(key packageKey) (*packageHandle, bool) {
+	value, ok := m.impl.Get(key)
+	if !ok {
+		return nil, false
+	}
+	return value.(*packageHandle), true
+}
+
+func (m packagesMap) Range(do func(key packageKey, value *packageHandle)) {
+	m.impl.Range(func(key, value interface{}) {
+		do(key.(packageKey), value.(*packageHandle))
+	})
+}
+
+func (m packagesMap) Set(key packageKey, value *packageHandle, release func()) {
+	m.impl.Set(key, value, func(key, value interface{}) {
+		release()
+	})
+}
+
+func (m packagesMap) Delete(key packageKey) {
+	m.impl.Delete(key)
+}

internal/lsp/cache/session.go

@@ -231,7 +231,7 @@ func (s *Session) createView(ctx context.Context, name string, folder span.URI,
 		cancel:            cancel,
 		initializeOnce:    &sync.Once{},
 		generation:        s.cache.store.Generation(generationName(v, 0)),
-		packages:          make(map[packageKey]*packageHandle),
+		packages:          newPackagesMap(),
 		meta:              &metadataGraph{},
 		files:             newFilesMap(),
 		goFiles:           newGoFilesMap(),

internal/lsp/cache/snapshot.go

@@ -89,7 +89,7 @@ type snapshot struct {
 
 	// packages maps a packageKey to a set of packageHandles to which that file belongs.
 	// It may be invalidated when a file's content changes.
-	packages map[packageKey]*packageHandle
+	packages packagesMap
 
 	// actions maps an actionkey to its actionHandle.
 	actions map[actionKey]*actionHandle
@@ -140,6 +140,7 @@ type actionKey struct {
 
 func (s *snapshot) Destroy(destroyedBy string) {
 	s.generation.Destroy(destroyedBy)
+	s.packages.Destroy()
 	s.files.Destroy()
 	s.goFiles.Destroy()
 	s.parseKeysByURI.Destroy()
@@ -711,16 +712,17 @@ func (s *snapshot) getImportedBy(id PackageID) []PackageID {
 	return s.meta.importedBy[id]
 }
 
-func (s *snapshot) addPackageHandle(ph *packageHandle) *packageHandle {
+func (s *snapshot) addPackageHandle(ph *packageHandle, release func()) *packageHandle {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
 	// If the package handle has already been cached,
 	// return the cached handle instead of overriding it.
-	if ph, ok := s.packages[ph.packageKey()]; ok {
-		return ph
+	if result, ok := s.packages.Get(ph.packageKey()); ok {
+		release()
+		return result
 	}
-	s.packages[ph.packageKey()] = ph
+	s.packages.Set(ph.packageKey(), ph, release)
 	return ph
 }
 
@@ -1090,10 +1092,10 @@ func (s *snapshot) CachedImportPaths(ctx context.Context) (map[string]source.Pac
 	defer s.mu.Unlock()
 
 	results := map[string]source.Package{}
-	for _, ph := range s.packages {
+	s.packages.Range(func(key packageKey, ph *packageHandle) {
 		cachedPkg, err := ph.cached(s.generation)
 		if err != nil {
-			continue
+			return
 		}
 		for importPath, newPkg := range cachedPkg.imports {
 			if oldPkg, ok := results[string(importPath)]; ok {
@@ -1105,7 +1107,7 @@ func (s *snapshot) CachedImportPaths(ctx context.Context) (map[string]source.Pac
 				results[string(importPath)] = newPkg
 			}
 		}
-	}
+	})
 	return results, nil
 }
 
@@ -1134,7 +1136,8 @@ func (s *snapshot) getPackage(id PackageID, mode source.ParseMode) *packageHandl
 		id:   id,
 		mode: mode,
 	}
-	return s.packages[key]
+	ph, _ := s.packages.Get(key)
+	return ph
 }
 
 func (s *snapshot) getSymbolHandle(uri span.URI) *symbolHandle {
@@ -1690,7 +1693,7 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]*fileC
 		builtin:           s.builtin,
 		initializeOnce:    s.initializeOnce,
 		initializedErr:    s.initializedErr,
-		packages:          make(map[packageKey]*packageHandle, len(s.packages)),
+		packages:          s.packages.Clone(),
 		actions:           make(map[actionKey]*actionHandle, len(s.actions)),
 		files:             s.files.Clone(),
 		goFiles:           s.goFiles.Clone(),
@@ -1894,13 +1897,12 @@ func (s *snapshot) clone(ctx, bgCtx context.Context, changes map[span.URI]*fileC
 		addRevDeps(id, invalidateMetadata)
 	}
 
-	// Copy the package type information.
-	for k, v := range s.packages {
-		if _, ok := idsToInvalidate[k.id]; ok {
-			continue
+	// Delete invalidated package type information.
+	for id := range idsToInvalidate {
+		for _, mode := range source.AllParseModes {
+			key := packageKey{mode, id}
+			result.packages.Delete(key)
 		}
-		newGen.Inherit(v.handle)
-		result.packages[k] = v
 	}
 
 	// Copy the package analysis information.

internal/lsp/source/view.go

@@ -478,6 +478,10 @@ const (
 	ParseFull
 )
 
+// AllParseModes contains all possible values of ParseMode.
+// It is used for cache invalidation on a file content change.
+var AllParseModes = []ParseMode{ParseHeader, ParseExported, ParseFull}
+
 // TypecheckMode controls what kind of parsing should be done (see ParseMode)
 // while type checking a package.
 type TypecheckMode int