Cancel auto-import cache warming if changes come in while building (#3454)

Author: andrewbranch

internal/project/session.go

@@ -135,6 +135,13 @@ type Session struct {
 	diagnosticsRefreshCancel context.CancelFunc
 	diagnosticsRefreshMu     sync.Mutex
 
+	// warmAutoImportCancel is the cancelation function for a running
+	// auto-import cache warming task. It is cancelled on file opens,
+	// closes, changes, watched-file changes, new auto-import warming
+	// requests, and when the session closes.
+	warmAutoImportCancel context.CancelFunc
+	warmAutoImportMu     sync.Mutex
+
 	// idleCacheCleanTimer is a resettable timer for scheduling idle disk
 	// cache cleans. The timer resets on any file event (open, close,
 	// change, save, watch) and fires after 30 seconds of inactivity.
@@ -270,6 +277,7 @@ func (s *Session) InitializeWithUserConfig(config lsutil.UserPreferences) {
 
 func (s *Session) DidOpenFile(ctx context.Context, uri lsproto.DocumentUri, version int32, content string, languageKind lsproto.LanguageKind) {
 	s.cancelDiagnosticsRefresh()
+	s.cancelWarmAutoImportCache()
 	s.scheduleIdleCacheClean()
 	s.snapshotUpdateMu.Lock()
 	defer s.snapshotUpdateMu.Unlock()
@@ -294,6 +302,7 @@ func (s *Session) DidOpenFile(ctx context.Context, uri lsproto.DocumentUri, vers
 
 func (s *Session) DidCloseFile(ctx context.Context, uri lsproto.DocumentUri) {
 	s.cancelDiagnosticsRefresh()
+	s.cancelWarmAutoImportCache()
 	s.scheduleIdleCacheClean()
 	s.pendingFileChangesMu.Lock()
 	defer s.pendingFileChangesMu.Unlock()
@@ -305,6 +314,7 @@ func (s *Session) DidCloseFile(ctx context.Context, uri lsproto.DocumentUri) {
 
 func (s *Session) DidChangeFile(ctx context.Context, uri lsproto.DocumentUri, version int32, changes []lsproto.TextDocumentContentChangePartialOrWholeDocument) {
 	s.cancelDiagnosticsRefresh()
+	s.cancelWarmAutoImportCache()
 	s.scheduleIdleCacheClean()
 	s.pendingFileChangesMu.Lock()
 	defer s.pendingFileChangesMu.Unlock()
@@ -353,6 +363,7 @@ func (s *Session) DidChangeWatchedFiles(ctx context.Context, changes []*lsproto.
 
 	// Schedule a debounced diagnostics refresh
 	s.ScheduleDiagnosticsRefresh()
+	s.cancelWarmAutoImportCache()
 	s.scheduleIdleCacheClean()
 }
 
@@ -416,6 +427,15 @@ func (s *Session) cancelDiagnosticsRefresh() {
 	}
 }
 
+func (s *Session) cancelWarmAutoImportCache() {
+	s.warmAutoImportMu.Lock()
+	defer s.warmAutoImportMu.Unlock()
+	if s.warmAutoImportCancel != nil {
+		s.warmAutoImportCancel()
+		s.warmAutoImportCancel = nil
+	}
+}
+
 const idleCacheCleanDelay = 30 * time.Second
 
 func (s *Session) scheduleIdleCacheClean() {
@@ -1245,6 +1265,8 @@ func (s *Session) updateWatches(oldSnapshot *Snapshot, newSnapshot *Snapshot) er
 func (s *Session) Close() {
 	// Cancel any pending diagnostics refresh
 	s.cancelDiagnosticsRefresh()
+	// Cancel any pending auto-import cache warming
+	s.cancelWarmAutoImportCache()
 	// Cancel any pending idle cache clean
 	s.cancelIdleCacheClean()
 	// Cancel periodic performance telemetry
@@ -1590,6 +1612,56 @@ func (s *Session) warmAutoImportCache(ctx context.Context, change SnapshotChange
 		) {
 			return
 		}
-		_, _ = s.GetCurrentLanguageServiceWithAutoImports(ctx, changedFile)
+
+		// Cancel any previous auto-import warming and create a new cancellable context.
+		// Only publish the new cancel func if the derived context is still active,
+		// and make the stored cancel func a no-op once that warming task is done.
+		s.warmAutoImportMu.Lock()
+		if s.warmAutoImportCancel != nil {
+			s.warmAutoImportCancel()
+		}
+		warmCtx, cancel := context.WithCancel(ctx)
+		if warmCtx.Err() == nil {
+			s.warmAutoImportCancel = func() {
+				if warmCtx.Err() != nil {
+					return
+				}
+				s.logger.Logf("Cancelling auto-import warming for file %s", changedFile.FileName())
+				cancel()
+			}
+		}
+		s.warmAutoImportMu.Unlock()
+
+		if warmCtx.Err() != nil {
+			cancel()
+			return
+		}
+		defer cancel()
+
+		// Clone the snapshot with auto-imports using warmCtx so the expensive
+		// extraction work is cancelled if a file change arrives.
+		if !newSnapshot.tryRef() {
+			return
+		}
+		defer newSnapshot.Deref(s)
+
+		warmChange := SnapshotChange{
+			reason: UpdateReasonRequestedLanguageServiceWithAutoImports,
+			ResourceRequest: ResourceRequest{
+				Documents:   []lsproto.DocumentUri{changedFile},
+				AutoImports: changedFile,
+			},
+		}
+		clonedSnapshot := newSnapshot.Clone(warmCtx, warmChange, newSnapshot.fs.overlays, s)
+
+		// If cancelled during clone, discard the incomplete result.
+		if warmCtx.Err() != nil {
+			clonedSnapshot.Deref(s)
+			return
+		}
+
+		// Conditionally adopt: if the session hasn't moved past newSnapshot,
+		// promote the clone so future requests benefit from the warmed cache.
+		s.adoptSnapshotChange(newSnapshot, clonedSnapshot)
 	}
 }

internal/project/snapshot.go

@@ -503,6 +503,21 @@ func (s *Snapshot) ref() {
 	}
 }
 
+// tryRef attempts to increment the snapshot's reference count. If the
+// snapshot is already disposed (refCount == 0), it returns false without
+// modifying the count. On success the caller must eventually call Deref.
+func (s *Snapshot) tryRef() bool {
+	for {
+		rc := s.refCount.Load()
+		if rc <= 0 {
+			return false
+		}
+		if s.refCount.CompareAndSwap(rc, rc+1) {
+			return true
+		}
+	}
+}
+
 // Deref decrements the snapshot's reference count. When the count reaches
 // zero, the snapshot is disposed and its resources are released.
 func (s *Snapshot) Deref(session *Session) {