Fix race condition in EnC implementation that might cause VS crash (#19910)

Fix race condition in EnC implementation

src/VisualStudio/Core/Def/Implementation/EditAndContinue/Interop/NativeMethods.cs

@@ -20,9 +20,10 @@ public static int GetStreamForObject(object pUnk, out IntPtr stream)
             return CoMarshalInterThreadInterfaceInStream(ref s_IID_IUnknown, pUnk, out stream);
         }
 
-        public static int GetObjectForStream(IntPtr stream, out object pUnk)
+        public static object GetObjectAndRelease(IntPtr stream)
         {
-            return CoGetInterfaceAndReleaseStream(stream, ref s_IID_IUnknown, out pUnk);
+            Marshal.ThrowExceptionForHR(CoGetInterfaceAndReleaseStream(stream, ref s_IID_IUnknown, out object pUnk));
+            return pUnk;
         }
     }
 }

src/VisualStudio/Core/Def/Implementation/EditAndContinue/VsENCRebuildableProjectImpl.cs

@@ -90,9 +90,7 @@ internal sealed class VsENCRebuildableProjectImpl
         /// </summary>
         private ModuleMetadata _metadata;
 
-        private ISymUnmanagedReader3 _pdbReader;
-
-        private IntPtr _pdbReaderObjAsStream;
+        private Lazy<ISymUnmanagedReader3> _pdbReader;
 
         #endregion
 
@@ -348,16 +346,14 @@ public int StopDebuggingPE()
                 _activeStatementIds = null;
                 _projectBeingEmitted = null;
 
-                Debug.Assert(_pdbReaderObjAsStream == IntPtr.Zero || _pdbReader == null);
-
-                if (_pdbReader != null)
+                var pdbReader = Interlocked.Exchange(ref _pdbReader, null);
+                if (pdbReader?.IsValueCreated == true)
                 {
-                    if (Marshal.IsComObject(_pdbReader))
+                    var symReader = pdbReader.Value;
+                    if (Marshal.IsComObject(symReader))
                     {
-                        Marshal.ReleaseComObject(_pdbReader);
+                        Marshal.ReleaseComObject(symReader);
                     }
-
-                    _pdbReader = null;
                 }
 
                 // The HResult is ignored by the debugger.
@@ -990,11 +986,10 @@ public unsafe int BuildForEnc(object pUpdatePE)
                 var updater = (IDebugUpdateInMemoryPE2)pUpdatePE;
                 if (_committedBaseline == null)
                 {
-                    var hr = MarshalPdbReader(updater, out _pdbReaderObjAsStream, out _pdbReader);
-                    if (hr != VSConstants.S_OK)
-                    {
-                        return hr;
-                    }
+                    var previousPdbReader = Interlocked.Exchange(ref _pdbReader, MarshalPdbReader(updater));
+
+                    // PDB reader should have been nulled out when debugging stopped:
+                    Contract.ThrowIfFalse(previousPdbReader == null);
 
                     _committedBaseline = EmitBaseline.CreateInitialBaseline(_metadata, GetBaselineEncDebugInfo);
                 }
@@ -1114,31 +1109,20 @@ private Deltas EmitProjectDelta()
         private EditAndContinueMethodDebugInformation GetBaselineEncDebugInfo(MethodDefinitionHandle methodHandle)
         {
             Debug.Assert(Thread.CurrentThread.GetApartmentState() == ApartmentState.MTA);
-
-            InitializePdbReader();
-            return GetEditAndContinueMethodDebugInfo(_pdbReader, methodHandle);
+            return GetEditAndContinueMethodDebugInfo(_pdbReader.Value, methodHandle);
         }
 
-        private void InitializePdbReader()
+        // Unmarshal the symbol reader (being marshalled cross thread from STA -> MTA).
+        private static ISymUnmanagedReader3 UnmarshalSymReader(IntPtr stream)
         {
             Debug.Assert(Thread.CurrentThread.GetApartmentState() == ApartmentState.MTA);
-
-            if (_pdbReader == null)
+            try
             {
-                // Unmarshal the symbol reader (being marshalled cross thread from STA -> MTA).
-
-                Debug.Assert(_pdbReaderObjAsStream != IntPtr.Zero);
-                var exception = Marshal.GetExceptionForHR(NativeMethods.GetObjectForStream(_pdbReaderObjAsStream, out object pdbReaderObjMta));
-                if (exception != null)
-                {
-                    // likely a bug in the compiler/debugger
-                    FatalError.ReportWithoutCrash(exception);
-
-                    throw new InvalidDataException(exception.Message, exception);
-                }
-
-                _pdbReaderObjAsStream = IntPtr.Zero;
-                _pdbReader = (ISymUnmanagedReader3)pdbReaderObjMta;
+                return (ISymUnmanagedReader3)NativeMethods.GetObjectAndRelease(stream);
+            }
+            catch (Exception exception) when (FatalError.ReportWithoutCrash(exception))
+            {
+                throw new InvalidDataException(exception.Message, exception);
             }
         }
 
@@ -1309,7 +1293,7 @@ public int GetCurrentExceptionSpanPosition(uint exceptionRegionId, VsTextSpan[]
             }
         }
 
-        private static int MarshalPdbReader(IDebugUpdateInMemoryPE2 updater, out IntPtr pdbReaderPointer, out ISymUnmanagedReader3 managedSymReader)
+        private static Lazy<ISymUnmanagedReader3> MarshalPdbReader(IDebugUpdateInMemoryPE2 updater)
         {
             // ISymUnmanagedReader can only be accessed from an MTA thread, however, we need
             // fetch the IUnknown instance (call IENCSymbolReaderProvider.GetSymbolReader) here
@@ -1332,16 +1316,16 @@ private static int MarshalPdbReader(IDebugUpdateInMemoryPE2 updater, out IntPtr
             symbolReaderProvider.GetSymbolReader(out object pdbReaderObjSta);
             if (Marshal.IsComObject(pdbReaderObjSta))
             {
-                int hr = NativeMethods.GetStreamForObject(pdbReaderObjSta, out pdbReaderPointer);
+                int hr = NativeMethods.GetStreamForObject(pdbReaderObjSta, out IntPtr stream);
                 Marshal.ReleaseComObject(pdbReaderObjSta);
-                managedSymReader = null;
-                return hr;
+                Marshal.ThrowExceptionForHR(hr);
+
+                return new Lazy<ISymUnmanagedReader3>(() => UnmarshalSymReader(stream));
             }
             else
             {
-                pdbReaderPointer = IntPtr.Zero;
-                managedSymReader = (ISymUnmanagedReader3)pdbReaderObjSta;
-                return 0;
+                var managedSymReader = (ISymUnmanagedReader3)pdbReaderObjSta;
+                return new Lazy<ISymUnmanagedReader3>(() => managedSymReader);
             }
         }