diff --git a/internal/project/project.go b/internal/project/project.go index 39c77aa78d..32b570cafa 100644 --- a/internal/project/project.go +++ b/internal/project/project.go @@ -374,11 +374,14 @@ func (p *Project) CreateProgram() CreateProgramResult { // and it is the only file whose refcount is already accounted for. if file != dirtyFile { // UpdateProgram acquired the changed file only, so we need to ref everything else - p.host.builder.parseCache.Ref(NewParseCacheKey(file.ParseOptions(), file.Hash, file.ScriptKind)) + p.host.builder.parseCache.RefValue(NewParseCacheKey(file.ParseOptions(), file.Hash, file.ScriptKind), file) } } for _, file := range newProgram.DuplicateSourceFiles() { - p.host.builder.parseCache.Ref(NewParseCacheKey(file.ParseOptions, file.Hash, file.ScriptKind)) + // Duplicate entries come from the old program's bookkeeping and only + // balance entries that were already acquired; unlike reused source files, + // they don't carry a SourceFile value that could restore a released entry. + p.host.builder.parseCache.RefIfPresent(NewParseCacheKey(file.ParseOptions, file.Hash, file.ScriptKind)) } } else if dirtyFile != nil { // UpdateProgram always acquires the dirty file before deciding whether it can diff --git a/internal/project/refcountcache.go b/internal/project/refcountcache.go index 909b5aa4b1..b00c1710b0 100644 --- a/internal/project/refcountcache.go +++ b/internal/project/refcountcache.go @@ -60,20 +60,62 @@ func (c *RefCountCache[K, V, AcquireArgs]) Has(identity K) bool { // Ref increments the reference count for an existing entry. // Panics if the entry does not exist. func (c *RefCountCache[K, V, AcquireArgs]) Ref(identity K) { + if c.TryRef(identity) { + return + } + panic("cache entry not found") +} + +// TryRef increments the reference count for an existing entry. +// Returns false if the entry does not exist. +func (c *RefCountCache[K, V, AcquireArgs]) TryRef(identity K) bool { entry, ok := c.entries.Load(identity) if !ok { - panic("cache entry not found") + return false + } + return c.tryRefEntry(entry) +} + +// RefIfPresent increments the reference count for an existing entry and does +// nothing if the entry does not exist. +func (c *RefCountCache[K, V, AcquireArgs]) RefIfPresent(identity K) { + c.TryRef(identity) +} + +// RefValue increments the reference count for an entry, restoring it with value +// if it was deleted before the ref could be taken. If another goroutine restores +// the entry first, this refs that entry and ignores value. +func (c *RefCountCache[K, V, AcquireArgs]) RefValue(identity K, value V) { + for { + entry, ok := c.entries.Load(identity) + if !ok { + entry = &refCountCacheEntry[V]{ + value: value, + refCount: 1, + } + existing, loaded := c.entries.LoadOrStore(identity, entry) + if !loaded { + return + } + entry = existing + } + if c.tryRefEntry(entry) { + return + } + // The deleted entry was already removed from entries before we could ref it, + // so continue until this goroutine stores a replacement or another goroutine + // does. If another goroutine wins, we'll ref its entry on the next iteration. } +} + +func (c *RefCountCache[K, V, AcquireArgs]) tryRefEntry(entry *refCountCacheEntry[V]) bool { entry.mu.Lock() defer entry.mu.Unlock() if entry.refCount <= 0 && !c.Options.DisableDeletion { - // Entry was deleted while we were acquiring the lock - newEntry, _ := c.loadOrStoreNewLockedEntry(identity) - defer newEntry.mu.Unlock() - newEntry.value = entry.value - return + return false } entry.refCount++ + return true } // Deref decrements the reference count for an entry. diff --git a/internal/project/refcountcache_test.go b/internal/project/refcountcache_test.go index 439588cd3e..819151fd17 100644 --- a/internal/project/refcountcache_test.go +++ b/internal/project/refcountcache_test.go @@ -13,6 +13,35 @@ import ( "gotest.tools/v3/assert" ) +func TestRefCountCacheRefValueRestoresDeletedEntry(t *testing.T) { + t.Parallel() + + cache := NewRefCountCache(RefCountCacheOptions{}, func(_ string, value string) string { + return "parsed:" + value + }) + + assert.Equal(t, cache.Acquire("key", "original"), "parsed:original") + cache.Deref("key") + assert.Assert(t, !cache.TryRef("key")) + + cache.RefValue("key", "restored") + assert.Assert(t, cache.TryRef("key")) + cache.Deref("key") + assert.Equal(t, cache.Acquire("key", "ignored"), "restored") + cache.Deref("key") + cache.Deref("key") + assert.Assert(t, !cache.Has("key")) + + cache.RefValue("key", "restored") + cache.RefValue("key", "ignored") + assert.Equal(t, cache.Acquire("key", "ignored"), "restored") + + cache.Deref("key") + cache.Deref("key") + cache.Deref("key") + assert.Assert(t, !cache.Has("key")) +} + func TestRefCountingCaches(t *testing.T) { t.Parallel()