Fix RefCountCache.Ref panic race between concurrent snapshot builds

tsc/internal/project/project.go

@@ -411,24 +411,73 @@ func (p *Project) CreateProgram() CreateProgramResult {
 		newProgram, dirtyFile, programCloned = p.Program.UpdateProgram(p.dirtyFilePath, p.host, createCheckerPool)
 		if programCloned {
 			updateKind = ProgramUpdateKindCloned
+			// canonicalByKey/canonicalContentMappedByKey let the duplicate loop below fall
+			// back to recreating an entry it can't find: every duplicate's key matches some
+			// canonical file's key by construction (that's what makes it a duplicate), and
+			// canonical files are exactly what this loop visits.
+			var canonicalByKey map[ParseCacheKey]*ast.SourceFile
+			var canonicalContentMappedByKey map[ContentMappedParseCacheKey]*ast.SourceFile
 			for _, file := range newProgram.SourceFiles() {
+				if file.IsContentMapperFailureStub() || file.IsContentMapperSupplemental() {
+					continue
+				}
+				if file.ContentMapper() != "" {
+					if canonicalContentMappedByKey == nil {
+						canonicalContentMappedByKey = make(map[ContentMappedParseCacheKey]*ast.SourceFile)
+					}
+					canonicalContentMappedByKey[contentMappedParseCacheKeyForFile(file)] = file
+				} else {
+					if canonicalByKey == nil {
+						canonicalByKey = make(map[ParseCacheKey]*ast.SourceFile)
+					}
+					canonicalByKey[parseCacheKeyForFile(file)] = file
+				}
 				// Use pointer identity: dirtyFile is the exact instance UpdateProgram acquired,
 				// and it is the only file whose refcount is already accounted for.
-				if file != dirtyFile && !file.IsContentMapperFailureStub() && !file.IsContentMapperSupplemental() {
-					// UpdateProgram acquired the changed file only, so we need to ref everything else
+				if file != dirtyFile {
+					// UpdateProgram acquired the changed file only, so we need to ref everything else.
+					// We already hold file itself, so RefOrAcquire (rather than Ref) tolerates losing
+					// a benign race against a concurrent, independent snapshot build that drops the
+					// last other claim on this cache entry between our lookup and our lock (e.g. a
+					// normal edit racing a speculative auto-import clone sharing this file's old
+					// Program, see GetLanguageServiceWithAutoImports): recreating the entry from a
+					// value we already possess is always correct.
 					if file.ContentMapper() != "" {
-						p.host.builder.contentMappedParseCache.Ref(contentMappedParseCacheKeyForFile(file))
+						p.host.builder.contentMappedParseCache.RefOrAcquire(
+							contentMappedParseCacheKeyForFile(file),
+							contentmapper.SourceFiles{Canonical: file, Supplemental: file.SupplementalSourceFiles()},
+						)
 					} else {
-						p.host.builder.parseCache.Ref(parseCacheKeyForFile(file))
+						p.host.builder.parseCache.RefOrAcquire(parseCacheKeyForFile(file), file)
 					}
 				}
 			}
 			for _, file := range newProgram.DuplicateSourceFiles() {
 				if !file.IsContentMapperFailureStub {
+					// Duplicates are pure bookkeeping refs on an entry acquired elsewhere: we
+					// usually have no value to recreate it with, so RefIfPresent is tried first;
+					// it only fails to ref an entry if no trace of it survived even its own
+					// internal recovery, which can happen if this duplicate's canonical file was
+					// itself lost to the same benign race described above. In that case fall back
+					// to the canonical file visited by the loop above, which this duplicate's key
+					// is guaranteed to match, and recreate the entry from it directly.
 					if file.ContentMapper != "" {
-						p.host.builder.contentMappedParseCache.Ref(contentMappedParseCacheKeyForDuplicate(file))
+						key := contentMappedParseCacheKeyForDuplicate(file)
+						if !p.host.builder.contentMappedParseCache.RefIfPresent(key) {
+							if canonical, ok := canonicalContentMappedByKey[key]; ok {
+								p.host.builder.contentMappedParseCache.RefOrAcquire(
+									key,
+									contentmapper.SourceFiles{Canonical: canonical, Supplemental: canonical.SupplementalSourceFiles()},
+								)
+							}
+						}
 					} else {
-						p.host.builder.parseCache.Ref(parseCacheKeyForDuplicate(file))
+						key := parseCacheKeyForDuplicate(file)
+						if !p.host.builder.parseCache.RefIfPresent(key) {
+							if canonical, ok := canonicalByKey[key]; ok {
+								p.host.builder.parseCache.RefOrAcquire(key, canonical)
+							}
+						}
 					}
 				}
 			}

tsc/internal/project/refcountcache.go

@@ -80,23 +80,64 @@ func (c *RefCountCache[K, V, AcquireArgs]) AcquireOrError(identity K, produce fu
 	return value, nil
 }
 
-// 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) {
+// RefOrAcquire increments the reference count for an existing entry, or
+// installs value as a fresh entry with refCount 1 if none exists.
+//
+// It never panics on a missing entry. It exists for callers that already
+// hold value from elsewhere (e.g. a *ast.SourceFile reused from
+// an old Program while cloning a new one) and are re-establishing their own
+// claim on it. Such callers can legitimately race with a concurrent Deref of
+// the last other claim on the same identity: two independent snapshot builds
+// (for example a normal edit and a speculative auto-import clone, see
+// GetLanguageServiceWithAutoImports) can each be cloning from the same
+// shared Program concurrently, and the moment the file's last other owner
+// releases it can fall between this call's initial lookup and its lock
+// acquisition. Since the caller already possesses a valid value for
+// identity, recreating the entry is always safe: it never returns a value
+// the caller didn't already have.
+func (c *RefCountCache[K, V, AcquireArgs]) RefOrAcquire(identity K, value V) {
+	entry, loaded := c.loadOrStoreNewLockedEntry(identity)
+	if !loaded {
+		entry.value = value
+	}
+	entry.mu.Unlock()
+}
+
+// RefIfPresent increments the reference count for an existing entry and
+// reports true, or reports false if no trace of the entry could be found.
+//
+// It exists for callers that are recording an additional owner of an entry
+// they do not themselves have a value for (e.g. a duplicate source file,
+// which is only ever a bookkeeping reference to a canonical entry acquired
+// elsewhere). If the entry is concurrently deleted between this call's
+// lookup and its lock acquisition, it is resurrected using the value it
+// already held (the same recovery loadOrStoreNewLockedEntry performs for
+// Acquire/RefOrAcquire), so the ref this call records — and the Deref its
+// caller will issue later to release it — stay balanced. A plain no-op here
+// would leave that later Deref unmatched, and it could land on an unrelated
+// entry that happens to reuse the same key by the time it runs.
+//
+// Only if the entry was never observed at all (not even a stale, about to be
+// deleted one) does this return false; there is no value to recover in that
+// case, so the caller must skip the corresponding Deref to stay balanced.
+func (c *RefCountCache[K, V, AcquireArgs]) RefIfPresent(identity K) bool {
 	entry, ok := c.entries.Load(identity)
 	if !ok {
-		panic("cache entry not found")
+		return false
 	}
 	entry.mu.Lock()
-	defer entry.mu.Unlock()
 	if entry.refCount <= 0 && !c.Options.DisableDeletion {
-		// Entry was deleted while we were acquiring the lock
+		// Entry was deleted while we were acquiring the lock; resurrect it
+		// from the value it already held so this ref stays balanced.
+		entry.mu.Unlock()
 		newEntry, _ := c.loadOrStoreNewLockedEntry(identity)
-		defer newEntry.mu.Unlock()
 		newEntry.value = entry.value
-		return
+		newEntry.mu.Unlock()
+		return true
 	}
 	entry.refCount++
+	entry.mu.Unlock()
+	return true
 }
 
 // Deref decrements the reference count for an entry.

tsc/internal/project/refcountcache_test.go

@@ -89,6 +89,70 @@ func TestParseCacheBindsBeforePublishing(t *testing.T) {
 	assert.Assert(t, file.CommonJSModuleIndicator != nil)
 }
 
+func TestRefOrAcquireRecreatesConcurrentlyDeletedEntry(t *testing.T) {
+	t.Parallel()
+
+	cache := NewParseCache(RefCountCacheOptions{})
+	key := NewParseCacheKey(ast.SourceFileParseOptions{FileName: "/a.ts", Path: "/a.ts"}, xxh3.Hash128([]byte("a")), core.ScriptKindTS)
+	file := &ast.SourceFile{}
+
+	// A caller holding file (e.g. reused, by pointer, from an old Program while
+	// cloning a new one) can lose a benign race: some other, independent owner
+	// derefs the entry to zero and it's deleted from the map entirely before
+	// this caller gets a chance to record its own claim. Plain Ref would panic
+	// in that situation (see refcountcache.go); RefOrAcquire must instead
+	// recreate the entry from the value the caller already has.
+	assert.Assert(t, !cache.Has(key))
+	cache.RefOrAcquire(key, file)
+	assert.Assert(t, cache.Has(key))
+	entry, ok := cache.entries.Load(key)
+	assert.Assert(t, ok)
+	assert.Equal(t, entry.refCount, 1)
+	assert.Assert(t, entry.value == file)
+
+	// A second RefOrAcquire for a live entry behaves like Ref: it bumps the
+	// existing entry rather than replacing its value.
+	other := &ast.SourceFile{}
+	cache.RefOrAcquire(key, other)
+	entry, ok = cache.entries.Load(key)
+	assert.Assert(t, ok)
+	assert.Equal(t, entry.refCount, 2)
+	assert.Assert(t, entry.value == file)
+
+	cache.Deref(key)
+	cache.Deref(key)
+	assert.Assert(t, !cache.Has(key))
+}
+
+func TestRefIfPresentSkipsMissingEntry(t *testing.T) {
+	t.Parallel()
+
+	cache := NewParseCache(RefCountCacheOptions{})
+	key := NewParseCacheKey(ast.SourceFileParseOptions{FileName: "/a.ts", Path: "/a.ts"}, xxh3.Hash128([]byte("a")), core.ScriptKindTS)
+
+	// Duplicates are bookkeeping-only refs on an entry owned elsewhere: there's
+	// no value on hand to recreate it with, so a missing entry must be a no-op
+	// (never a panic) rather than fabricating a zero-value entry.
+	assert.Equal(t, cache.RefIfPresent(key), false)
+	assert.Assert(t, !cache.Has(key))
+
+	file := &ast.SourceFile{}
+	cache.RefOrAcquire(key, file)
+	assert.Equal(t, cache.RefIfPresent(key), true)
+	entry, ok := cache.entries.Load(key)
+	assert.Assert(t, ok)
+	assert.Equal(t, entry.refCount, 2)
+
+	cache.Deref(key)
+	cache.Deref(key)
+	assert.Assert(t, !cache.Has(key))
+
+	// The corresponding Deref for a duplicate whose RefIfPresent no-op'd must
+	// also no-op rather than panicking or corrupting an unrelated entry.
+	cache.Deref(key)
+	assert.Assert(t, !cache.Has(key))
+}
+
 func TestRefCountingCaches(t *testing.T) {
 	t.Parallel()
 

tsc/internal/project/snapshot_stress_test.go

@@ -0,0 +1,147 @@
+package project
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"sync/atomic"
+	"testing"
+
+	"github.com/microsoft/TypeScript/tsc/internal/bundled"
+	"github.com/microsoft/TypeScript/tsc/internal/lsp/lsproto"
+	"github.com/microsoft/TypeScript/tsc/internal/vfs/vfstest"
+	"gotest.tools/v3/assert"
+)
+
+// TestSnapshotConcurrentAutoImportCloneDoesNotPanic reproduces
+// https://github.com/microsoft/TypeScript/issues/63844: a "cache entry not
+// found" panic in RefCountCache.Ref hit by real monorepo users of the
+// language server.
+//
+// Two entry points can build a new Snapshot from a shared base: the normal,
+// serialized edit path (getSnapshot/updateSnapshot, under snapshotUpdateMu)
+// and the speculative auto-import clone used by completions needing
+// auto-imports (CloneSnapshotWithAutoImports, used by
+// GetLanguageServiceWithAutoImports and warmAutoImportCache), which does NOT
+// go through snapshotUpdateMu. Both read and mutate the same host-level,
+// ref-counted parseCache/contentMappedParseCache. When a project's Program is
+// unchanged across several edits, it (and its files) stay shared across many
+// snapshot generations; a concurrent auto-import clone that reuses one of
+// those files via Project.CreateProgram's clone path can lose a benign race
+// against a concurrent edit's disposal of an older generation, such that the
+// file's cache entry is gone by the time the clone tries to Ref it.
+//
+// Neither concurrent edits alone nor concurrent auto-import clones alone
+// (against an otherwise idle session) are sufficient to reproduce this; it
+// takes both running at once, which is what this test drives.
+func TestSnapshotConcurrentAutoImportCloneDoesNotPanic(t *testing.T) {
+	if !bundled.Embedded {
+		t.Skip("bundled files are not embedded")
+	}
+
+	const numProjects = 6
+
+	files := map[string]any{}
+	for i := range numProjects {
+		files[fmt.Sprintf("/home/projects/TS/p%d/tsconfig.json", i)] = "{}"
+		files[fmt.Sprintf("/home/projects/TS/p%d/index.ts", i)] = "import { foo } from './foo'; export const value = foo;"
+		files[fmt.Sprintf("/home/projects/TS/p%d/foo.ts", i)] = "export const foo = 1;"
+	}
+
+	fs := bundled.WrapFS(vfstest.FromMap(files, false /*useCaseSensitiveFileNames*/))
+	session := NewSession(&SessionInit{
+		BackgroundCtx: context.Background(),
+		Options: &SessionOptions{
+			CurrentDirectory:   "/",
+			DefaultLibraryPath: bundled.LibPath(),
+			TypingsLocation:    "/home/src/Library/Caches/typescript",
+			PositionEncoding:   lsproto.PositionEncodingKindUTF8,
+			WatchEnabled:       false,
+			LoggingEnabled:     false,
+		},
+		FS: fs,
+	})
+	defer session.Close()
+
+	ctx := context.Background()
+	uris := make([]lsproto.DocumentUri, numProjects)
+	for i := range numProjects {
+		uri := lsproto.DocumentUri(fmt.Sprintf("file:///home/projects/TS/p%d/index.ts", i))
+		uris[i] = uri
+		session.DidOpenFile(ctx, uri, 1, files[fmt.Sprintf("/home/projects/TS/p%d/index.ts", i)].(string), lsproto.LanguageKindTypeScript)
+		_, err := session.GetLanguageService(ctx, uri)
+		assert.NilError(t, err)
+	}
+
+	var version int32 = 1
+	var wg sync.WaitGroup
+	stop := make(chan struct{})
+
+	// Goroutines that keep editing files, forcing a steady stream of new
+	// snapshot generations (and disposal of old ones) via the normal,
+	// snapshotUpdateMu-serialized path.
+	for i := range numProjects {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			uri := uris[i]
+			for {
+				select {
+				case <-stop:
+					return
+				default:
+				}
+				v := atomic.AddInt32(&version, 1)
+				session.DidChangeFile(ctx, uri, v, []lsproto.TextDocumentContentChangePartialOrWholeDocument{
+					{
+						WholeDocument: &lsproto.TextDocumentContentChangeWholeDocument{
+							Text: fmt.Sprintf("import { foo } from './foo'; export const value = foo; export const v = %d;", v),
+						},
+					},
+				})
+				_, _ = session.GetLanguageService(ctx, uri)
+			}
+		}(i)
+	}
+
+	// Goroutines that repeatedly take a speculative auto-import clone off of
+	// whatever the current snapshot happens to be, mimicking the
+	// ErrNeedsAutoImports path used by completions (ctrl+space), which does
+	// NOT go through snapshotUpdateMu.
+	for i := range numProjects {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			uri := uris[i]
+			for {
+				select {
+				case <-stop:
+					return
+				default:
+				}
+				// session.Snapshot() only reads the current pointer; a concurrent edit can
+				// adopt a replacement and Deref this one immediately afterward. tryRef holds
+				// it alive for the clone (mirroring how GetLanguageServiceWithAutoImports's
+				// own callerRef protects its base snapshot in production), and we simply skip
+				// this iteration if we lost that race, rather than cloning from a snapshot
+				// that may already be disposed.
+				baseSnapshot := session.Snapshot()
+				if !baseSnapshot.tryRef() {
+					continue
+				}
+				preparedSnapshot := session.SnapshotHost.CloneSnapshotWithAutoImports(ctx, baseSnapshot, uri, nil)
+				session.TryAdoptSnapshotInBackground(baseSnapshot, preparedSnapshot)
+				preparedSnapshot.Deref()
+				baseSnapshot.Deref()
+			}
+		}(i)
+	}
+
+	// Let the race run for a bounded number of edit cycles rather than wall time.
+	for n := 0; n < 400; n++ {
+		_, _ = session.GetLanguageService(ctx, uris[n%numProjects])
+	}
+	close(stop)
+	wg.Wait()
+	session.WaitForBackgroundTasks()
+}