Merge pull request #4313 from dpoeschl/RenameTrackingBufferLeak

Fix ITextBuffer leak after RenameTracking commit

src/EditorFeatures/Core/Implementation/RenameTracking/RenameTrackingTaggerProvider.RenameTrackingCommitter.cs

@@ -137,7 +137,9 @@ private bool ApplyChangesToWorkspace(CancellationToken cancellationToken)
 
                 // When this action is undone (the user has undone twice), restore the state
                 // machine to so that they can continue their original rename tracking session.
-                UpdateWorkspaceForResetOfTypedIdentifier(workspace, renameTrackingSolutionSet.OriginalSolution);
+
+                var trackingSessionId = _stateMachine.StoreCurrentTrackingSessionAndGenerateId();
+                UpdateWorkspaceForResetOfTypedIdentifier(workspace, renameTrackingSolutionSet.OriginalSolution, trackingSessionId);
 
                 // Now that the solution is back in its original state, notify third parties about
                 // the coming rename operation.
@@ -164,7 +166,15 @@ private bool ApplyChangesToWorkspace(CancellationToken cancellationToken)
                 }
 
                 // Undo/redo on this action must always clear the state machine
-                UpdateWorkspaceForGlobalIdentifierRename(workspace, finalSolution, workspace.CurrentSolution, _displayText, changedDocuments, renameTrackingSolutionSet.Symbol, newName);
+                UpdateWorkspaceForGlobalIdentifierRename(
+                    workspace, 
+                    finalSolution, 
+                    workspace.CurrentSolution, 
+                    _displayText, 
+                    changedDocuments, 
+                    renameTrackingSolutionSet.Symbol, 
+                    newName,
+                    trackingSessionId);
 
                 RenameTrackingDismisser.DismissRenameTracking(workspace, changedDocuments);
                 return true;
@@ -209,7 +219,7 @@ private async Task<ISymbol> TryGetSymbolAsync(Solution solutionWithOriginalName,
                 return tokenRenameInfo.HasSymbols ? tokenRenameInfo.Symbols.First() : null;
             }
 
-            private void UpdateWorkspaceForResetOfTypedIdentifier(Workspace workspace, Solution newSolution)
+            private void UpdateWorkspaceForResetOfTypedIdentifier(Workspace workspace, Solution newSolution, int trackingSessionId)
             {
                 AssertIsForeground();
 
@@ -219,7 +229,7 @@ private void UpdateWorkspaceForResetOfTypedIdentifier(Workspace workspace, Solut
                 var undoHistory = _undoHistoryRegistry.RegisterHistory(_stateMachine.Buffer);
                 using (var localUndoTransaction = undoHistory.CreateTransaction(EditorFeaturesResources.TextBufferChange))
                 {
-                    var undoPrimitiveBefore = new UndoPrimitive(_stateMachine, shouldRestoreStateOnUndo: true);
+                    var undoPrimitiveBefore = new UndoPrimitive(_stateMachine.Buffer, trackingSessionId, shouldRestoreStateOnUndo: true);
                     localUndoTransaction.AddUndo(undoPrimitiveBefore);
 
                     if (!workspace.TryApplyChanges(newSolution))
@@ -228,7 +238,7 @@ private void UpdateWorkspaceForResetOfTypedIdentifier(Workspace workspace, Solut
                     }
 
                     // Never resume tracking session on redo
-                    var undoPrimitiveAfter = new UndoPrimitive(_stateMachine, shouldRestoreStateOnUndo: false);
+                    var undoPrimitiveAfter = new UndoPrimitive(_stateMachine.Buffer, trackingSessionId, shouldRestoreStateOnUndo: false);
                     localUndoTransaction.AddUndo(undoPrimitiveAfter);
 
                     localUndoTransaction.Complete();
@@ -242,7 +252,8 @@ private void UpdateWorkspaceForResetOfTypedIdentifier(Workspace workspace, Solut
                 string undoName,
                 IEnumerable<DocumentId> changedDocuments,
                 ISymbol symbol,
-                string newName)
+                string newName,
+                int trackingSessionId)
             {
                 AssertIsForeground();
 
@@ -254,7 +265,7 @@ private void UpdateWorkspaceForResetOfTypedIdentifier(Workspace workspace, Solut
                 using (var workspaceUndoTransaction = workspace.OpenGlobalUndoTransaction(undoName))
                 using (var localUndoTransaction = undoHistory.CreateTransaction(undoName))
                 {
-                    var undoPrimitiveBefore = new UndoPrimitive(_stateMachine, shouldRestoreStateOnUndo: false);
+                    var undoPrimitiveBefore = new UndoPrimitive(_stateMachine.Buffer, trackingSessionId, shouldRestoreStateOnUndo: false);
                     localUndoTransaction.AddUndo(undoPrimitiveBefore);
 
                     if (!workspace.TryApplyChanges(newSolution))
@@ -272,7 +283,7 @@ private void UpdateWorkspaceForResetOfTypedIdentifier(Workspace workspace, Solut
                     }
 
                     // Never resume tracking session on redo
-                    var undoPrimitiveAfter = new UndoPrimitive(_stateMachine, shouldRestoreStateOnUndo: false);
+                    var undoPrimitiveAfter = new UndoPrimitive(_stateMachine.Buffer, trackingSessionId, shouldRestoreStateOnUndo: false);
                     localUndoTransaction.AddUndo(undoPrimitiveAfter);
 
                     localUndoTransaction.Complete();

src/EditorFeatures/Core/Implementation/RenameTracking/RenameTrackingTaggerProvider.StateMachine.cs

@@ -33,6 +33,11 @@ private class StateMachine : ForegroundThreadAffinitizedObject
             private readonly ITextBuffer _buffer;
             private readonly IDiagnosticAnalyzerService _diagnosticAnalyzerService;
 
+            // Store committed sessions so they can be restored on undo/redo. The undo transactions
+            // may live beyond the lifetime of the buffer tracked by this StateMachine, so storing
+            // them here allows them to be correctly cleaned up when the buffer goes away.
+            private readonly IList<TrackingSession> _committedSessions = new List<TrackingSession>();
+
             private int _refCount;
 
             public TrackingSession TrackingSession { get; private set; }
@@ -242,6 +247,21 @@ public bool ClearVisibleTrackingSession()
                 return false;
             }
 
+            internal int StoreCurrentTrackingSessionAndGenerateId()
+            {
+                AssertIsForeground();
+
+                var existingIndex = _committedSessions.IndexOf(TrackingSession);
+                if (existingIndex >= 0)
+                {
+                    return existingIndex;
+                }
+
+                var index = _committedSessions.Count;
+                _committedSessions.Insert(index, TrackingSession);
+                return index;
+            }
+
             public bool CanInvokeRename(out TrackingSession trackingSession, bool isSmartTagCheck = false, bool waitForResult = false, CancellationToken cancellationToken = default(CancellationToken))
             {
                 // This needs to be able to run on a background thread for the diagnostic.
@@ -300,12 +320,12 @@ internal async Task<IEnumerable<Diagnostic>> GetDiagnostic(SyntaxTree tree, Diag
                 }
             }
 
-            public void RestoreTrackingSession(TrackingSession trackingSession)
+            public void RestoreTrackingSession(int trackingSessionId)
             {
                 AssertIsForeground();
                 ClearTrackingSession();
 
-                this.TrackingSession = trackingSession;
+                this.TrackingSession = _committedSessions[trackingSessionId];
                 TrackingSessionUpdated();
             }
 

src/EditorFeatures/Core/Implementation/RenameTracking/RenameTrackingTaggerProvider.UndoPrimitive.cs

 // Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
 
 using System;
+using Microsoft.VisualStudio.Text;
 using Microsoft.VisualStudio.Text.Operations;
 
 namespace Microsoft.CodeAnalysis.Editor.Implementation.RenameTracking
@@ -8,14 +9,19 @@ namespace Microsoft.CodeAnalysis.Editor.Implementation.RenameTracking
     internal sealed partial class RenameTrackingTaggerProvider
     {
         /// <summary>
-        /// Clears the state machine on relevant undo/redo actions.
+        /// Clears or restores the state machine on relevant undo/redo actions.
+        /// 
+        /// These may stay alive on the global undo stack well beyond the lifetime of the
+        /// <see cref="ITextBuffer"/> on which they were created, so we must avoid strong
+        /// references to anything that may hold that <see cref="ITextBuffer"/> alive.
         /// </summary>
         private class UndoPrimitive : ITextUndoPrimitive
         {
-            private readonly StateMachine _stateMachine;
-            private readonly TrackingSession _trackingSession;
-            private ITextUndoTransaction _parent;
+            private readonly WeakReference<ITextBuffer> _weakTextBuffer;
+            private readonly int _trackingSessionId;
+            private readonly bool _shouldRestoreStateOnUndo;
 
+            private ITextUndoTransaction _parent;
             public ITextUndoTransaction Parent
             {
                 get { return _parent; }
@@ -32,29 +38,46 @@ public bool CanUndo
                 get { return true; }
             }
 
-            public UndoPrimitive(StateMachine stateMachine, bool shouldRestoreStateOnUndo)
+            public UndoPrimitive(ITextBuffer textBuffer, int trackingSessionId, bool shouldRestoreStateOnUndo)
             {
-                _stateMachine = stateMachine;
-                _trackingSession = shouldRestoreStateOnUndo ? stateMachine.TrackingSession : null;
+                _weakTextBuffer = new WeakReference<ITextBuffer>(textBuffer);
+                _trackingSessionId = trackingSessionId;
+                _shouldRestoreStateOnUndo = shouldRestoreStateOnUndo;
             }
 
             public void Do()
             {
-                _stateMachine.ClearTrackingSession();
+                StateMachine stateMachine;
+                if (TryGetStateMachine(out stateMachine))
+                {
+                    stateMachine.ClearTrackingSession();
+                }
             }
 
             public void Undo()
             {
-                if (_trackingSession != null)
+                StateMachine stateMachine;
+                if (TryGetStateMachine(out stateMachine))
                 {
-                    _stateMachine.RestoreTrackingSession(_trackingSession);
-                }
-                else
-                {
-                    _stateMachine.ClearTrackingSession();
+                    if (_shouldRestoreStateOnUndo)
+                    {
+                        stateMachine.RestoreTrackingSession(_trackingSessionId);
+                    }
+                    else
+                    {
+                        stateMachine.ClearTrackingSession();
+                    }
                 }
             }
 
+            private bool TryGetStateMachine(out StateMachine stateMachine)
+            {
+                stateMachine = null;
+                ITextBuffer textBuffer;
+                return _weakTextBuffer.TryGetTarget(out textBuffer) &&
+                    textBuffer.Properties.TryGetProperty(typeof(StateMachine), out stateMachine);
+            }
+
             public bool CanMerge(ITextUndoPrimitive older)
             {
                 return false;