Fix parse cache ref race during cloned program updates

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

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.

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()