Unmap memory mapped view accessors when they are no longer in use

Further improves #19604, #19493

src/Workspaces/Core/Desktop/Workspace/Host/TemporaryStorage/TemporaryStorageServiceFactory.CopiedMemoryMappedViewHandle.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 System.Runtime.CompilerServices;
-using System.Runtime.InteropServices;
-using Microsoft.Win32.SafeHandles;
-
-namespace Microsoft.CodeAnalysis.Host
-{
-    internal partial class TemporaryStorageServiceFactory
-    {
-        /// <summary>
-        /// This critical handle increments the reference count of a <see cref="SafeMemoryMappedViewHandle"/>. The view
-        /// will be released when it is disposed <em>and</em> all pointers acquired through
-        /// <see cref="SafeBuffer.AcquirePointer"/> (which this class uses) are released.
-        /// </summary>
-        /// <remarks>
-        /// <para><see cref="CriticalHandle"/> types are not reference counted, and are thus somewhat limited in their
-        /// usefulness. However, this handle class has tightly restricted accessibility and is only used by managed
-        /// code which does not rely on it counting references.</para>
-        ///
-        /// <para>This is a supporting class for <see cref="MemoryMappedInfo"/>. See additional comments on that
-        /// class.</para>
-        /// </remarks>
-        private unsafe sealed class CopiedMemoryMappedViewHandle : CriticalHandleZeroOrMinusOneIsInvalid
-        {
-            private readonly SafeMemoryMappedViewHandle _viewHandle;
-
-            public CopiedMemoryMappedViewHandle(SafeMemoryMappedViewHandle viewHandle)
-            {
-                _viewHandle = viewHandle;
-
-                byte* pointer = null;
-
-                // The following code uses a constrained execution region (CER) to ensure that the code ends up in one
-                // of the following states, even in the presence of asynchronous exceptions like ThreadAbortException:
-                //
-                // 1. The pointer is not acquired, and the current handle is invalid (thus will not be released)
-                // 2. The pointer is acquired, and the current handle is fully initialized for later cleanup
-                RuntimeHelpers.PrepareConstrainedRegions();
-                try
-                {
-                    viewHandle.AcquirePointer(ref pointer);
-                }
-                finally
-                {
-                    SetHandle((IntPtr)pointer);
-                }
-            }
-
-            public byte* Pointer => (byte*)handle;
-
-            protected override bool ReleaseHandle()
-            {
-                // The operating system will release these handles when the process terminates, so do not spend time
-                // releasing them manually in that case. Doing so causes unnecessary delays during application shutdown.
-                if (!Environment.HasShutdownStarted)
-                {
-                    _viewHandle.ReleasePointer();
-                }
-
-                return true;
-            }
-        }
-    }
-}

src/Workspaces/Core/Desktop/Workspace/Host/TemporaryStorage/TemporaryStorageServiceFactory.MemoryMappedInfo.cs

@@ -50,13 +50,13 @@ internal sealed class MemoryMappedInfo : IDisposable
             /// actual memory accessor that owns the VM
             /// </summary>
             /// <remarks>
-            /// <para>It is possible for this accessor to be disposed prior to the streams which use it. However, the
-            /// streams interact directly with the underlying memory buffer, and keep a
-            /// <see cref="CopiedMemoryMappedViewHandle"/> to prevent that buffer from being released while still in
-            /// use. The <see cref="SafeHandle"/> used by this accessor is reference counted, and is not finally
-            /// released until the reference count reaches zero.</para>
+            /// <para>This holds a weak reference to current <see cref="MemoryMappedViewAccessor"/>, which allows
+            /// additional accessors for the same address space to be obtained up until the point when no external code
+            /// is using it. When the memory is no longer being used by any <see cref="SharedReadableStream"/> objects,
+            /// it is unmapped automatically and a new accessor is created for the memory mapped file the next time it's
+            /// needed.</para>
             /// </remarks>
-            private MemoryMappedViewAccessor _accessor;
+            private ReferenceCountedDisposable<MemoryMappedViewAccessor>.WeakReference _accessor;
 
             public MemoryMappedInfo(long size)
             {
@@ -99,13 +99,16 @@ public Stream CreateReadableStream()
                 // CreateViewAccessor is not guaranteed to be thread-safe
                 lock (_memoryMappedFile)
                 {
-                    if (_accessor == null)
+                    var streamAccessor = _accessor.TryAddReference();
+                    if (streamAccessor == null)
                     {
-                        _accessor = RunWithCompactingGCFallback(info => info._memoryMappedFile.CreateViewAccessor(0, info._size, MemoryMappedFileAccess.Read), this);
+                        var rawAccessor = RunWithCompactingGCFallback(info => info._memoryMappedFile.CreateViewAccessor(0, info._size, MemoryMappedFileAccess.Read), this);
+                        streamAccessor = new ReferenceCountedDisposable<MemoryMappedViewAccessor>(rawAccessor);
+                        _accessor = new ReferenceCountedDisposable<MemoryMappedViewAccessor>.WeakReference(streamAccessor);
                     }
 
-                    Contract.Assert(_accessor.CanRead);
-                    return new SharedReadableStream(this, new CopiedMemoryMappedViewHandle(_accessor.SafeMemoryMappedViewHandle), _accessor.PointerOffset, _size);
+                    Contract.Assert(streamAccessor.Target.CanRead);
+                    return new SharedReadableStream(this, streamAccessor, _size);
                 }
             }
 
@@ -162,16 +165,6 @@ private void Dispose(bool disposing)
             {
                 if (disposing)
                 {
-                    lock (_memoryMappedFile)
-                    {
-                        if (_accessor != null)
-                        {
-                            // (see remarks on accessor for relation between _accessor and the streams)
-                            _accessor.Dispose();
-                            _accessor = null;
-                        }
-                    }
-
                     // (see remarks on accessor for relation between _memoryMappedFile and the views/streams)
                     _memoryMappedFile.Dispose();
                 }
@@ -184,16 +177,16 @@ public static string CreateUniqueName(long size)
 
             private unsafe sealed class SharedReadableStream : Stream, ISupportDirectMemoryAccess
             {
-                private readonly CopiedMemoryMappedViewHandle _handle;
+                private readonly ReferenceCountedDisposable<MemoryMappedViewAccessor> _accessor;
 
                 private byte* _start;
                 private byte* _current;
                 private readonly byte* _end;
 
-                public SharedReadableStream(MemoryMappedInfo owner, CopiedMemoryMappedViewHandle handle, long offset, long length)
+                public SharedReadableStream(MemoryMappedInfo owner, ReferenceCountedDisposable<MemoryMappedViewAccessor> accessor, long length)
                 {
-                    _handle = handle;
-                    _current = _start = handle.Pointer + offset;
+                    _accessor = accessor;
+                    _current = _start = (byte*)_accessor.Target.SafeMemoryMappedViewHandle.DangerousGetHandle() + _accessor.Target.PointerOffset;
                     _end = checked(_start + length);
                 }
 
@@ -331,7 +324,7 @@ protected override void Dispose(bool disposing)
 
                     if (disposing)
                     {
-                        _handle.Dispose();
+                        _accessor.Dispose();
                     }
 
                     _start = null;

src/Workspaces/Core/Desktop/Workspace/Host/TemporaryStorage/TemporaryStorageServiceFactory.ReferenceCountedDisposable.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 System.Runtime.CompilerServices;
+using System.Threading;
+
+#if DEBUG
+using System.Diagnostics;
+#endif
+
+namespace Microsoft.CodeAnalysis.Host
+{
+    internal partial class TemporaryStorageServiceFactory
+    {
+        private sealed class ReferenceCountedDisposable<T> : IDisposable
+            where T : IDisposable
+        {
+            private T _instance;
+            private StrongBox<int> _referenceCount;
+
+            /// <summary>
+            /// Initializes a new reference counting wrapper around an <see cref="IDisposable"/> object.
+            /// </summary>
+            /// <remarks>
+            /// <para>The reference count is initialized to 1.</para>
+            /// </remarks>
+            /// <param name="instance">The object owned by this wrapper.</param>
+            public ReferenceCountedDisposable(T instance)
+            {
+                _instance = instance;
+                _referenceCount = new StrongBox<int>(1);
+            }
+
+            private ReferenceCountedDisposable(T instance, StrongBox<int> referenceCount)
+            {
+                _instance = instance;
+
+                // The reference count has already been incremented for this instance
+                _referenceCount = referenceCount;
+            }
+
+            /// <summary>
+            /// Gets the target object.
+            /// </summary>
+            /// <remarks>
+            /// <para>This call is not valid after <see cref="Dispose"/> is called.</para>
+            /// </remarks>
+            public T Target => _instance;
+
+            /// <summary>
+            /// Increments the reference count for the disposable object, and returns a new disposable reference to it.
+            /// </summary>
+            /// <remarks>
+            /// <para>The returned object is an independent reference to the same underlying object. Disposing of the
+            /// returned value multiple times will only cause the reference count to be decreased once.</para>
+            /// </remarks>
+            /// <returns>A new <see cref="ReferenceCountedDisposable{T}"/> pointing to the same underlying object, if it
+            /// has not yet been disposed; otherwise, <see langword="null"/> if this reference to the underlying object
+            /// has already been disposed.</returns>
+            public ReferenceCountedDisposable<T> TryAddReference()
+            {
+                // This value can be set in Dispose. However, even if we have a race condition with Dispose where it's
+                // set to null in the middle of the call, the reference count will be incremented by this method if and
+                // only if a valid new reference was created. This is because the validity of this instance for creating
+                // new references is determined by _referenceCount.
+                var target = _instance;
+                Interlocked.MemoryBarrier();
+                var referenceCount = Volatile.Read(ref _referenceCount);
+
+                // Cannot use Interlocked.Increment because we need to latch the reference count when it reaches 0.
+                while (true)
+                {
+                    var currentValue = Volatile.Read(ref referenceCount.Value);
+                    if (currentValue == 0)
+                    {
+                        // The target is already disposed, and cannot be reused
+                        return null;
+                    }
+
+                    if (Interlocked.CompareExchange(ref referenceCount.Value, currentValue + 1, currentValue) == currentValue)
+                    {
+                        return new ReferenceCountedDisposable<T>(target, referenceCount);
+                    }
+                }
+            }
+
+            public void Dispose()
+            {
+                var referenceCount = Interlocked.Exchange(ref _referenceCount, null);
+                if (referenceCount == null)
+                {
+                    // Already disposed; allow multiple without error.
+                    return;
+                }
+
+                // This is the thread which will write to _instance, so this will act as an atomic read
+                var instance = _instance;
+
+                // Set the instance back to its default value, so in case someone forgets to dispose of one of the
+                // counted references the finalizer of the underlying disposable object will have a chance to clean up
+                // the unmanaged resources as soon as possible.
+                Interlocked.MemoryBarrier();
+                _instance = default(T);
+
+#if !DEBUG
+                var decrementedValue = Interlocked.Decrement(ref referenceCount.Value);
+                if (decrementedValue == 0)
+                {
+                    instance.Dispose();
+                }
+#else
+                while (true)
+                {
+                    var currentValue = Volatile.Read(ref referenceCount.Value);
+                    Debug.Assert(currentValue > 0, "Dispose should have protected itself against races.");
+
+                    if (Interlocked.CompareExchange(ref referenceCount.Value, currentValue - 1, currentValue) == currentValue)
+                    {
+                        if (currentValue == 1)
+                        {
+                            // Reference count hit 0 for this call
+                            instance.Dispose();
+                        }
+
+                        break;
+                    }
+                }
+#endif
+            }
+
+            /// <summary>
+            /// Represents a weak reference to a <see cref="ReferenceCountedDisposable{T}"/> which is capable of
+            /// obtaining a new counted reference up until the point when the object is no longer accessible.
+            /// </summary>
+            public struct WeakReference
+            {
+                /// <summary>
+                /// DO NOT DISPOSE OF THE TARGET.
+                /// </summary>
+                private readonly WeakReference<ReferenceCountedDisposable<T>> _instance;
+
+                public WeakReference(ReferenceCountedDisposable<T> reference)
+                    : this()
+                {
+                    var instance = reference._instance;
+                    Interlocked.MemoryBarrier();
+                    var referenceCount = Volatile.Read(ref reference._referenceCount);
+                    if (referenceCount == null)
+                    {
+                        // The specified reference is already not valid.
+                        return;
+                    }
+
+                    if (referenceCount.Value == 0)
+                    {
+                        // We were able to read the reference, but it's already disposed.
+                        return;
+                    }
+
+                    var innerReference = new ReferenceCountedDisposable<T>(instance, referenceCount);
+                    _instance = new WeakReference<ReferenceCountedDisposable<T>>(innerReference);
+                }
+
+                /// <summary>
+                /// Increments the reference count for the disposable object, and returns a new disposable reference to
+                /// it.
+                /// </summary>
+                /// <remarks>
+                /// <para>Unlike <see cref="ReferenceCountedDisposable{T}.TryAddReference"/>, this method is capable of
+                /// adding a reference to the underlying instance all the way up to the point where it is finally
+                /// disposed.</para>
+                ///
+                /// <para>The returned object is an independent reference to the same underlying object. Disposing of
+                /// the returned value multiple times will only cause the reference count to be decreased once.</para>
+                /// </remarks>
+                /// <returns>A new <see cref="ReferenceCountedDisposable{T}"/> pointing to the same underlying object,
+                /// if it has not yet been disposed; otherwise, <see langword="null"/> if the underlying object has
+                /// already been disposed.</returns>
+                public ReferenceCountedDisposable<T> TryAddReference()
+                {
+                    var instance = _instance;
+                    if (instance == null)
+                    {
+                        return null;
+                    }
+
+                    if (!_instance.TryGetTarget(out var target))
+                    {
+                        return null;
+                    }
+
+                    return target.TryAddReference();
+                }
+            }
+        }
+    }
+}

src/Workspaces/Core/Desktop/Workspaces.Desktop.csproj

@@ -95,9 +95,9 @@
     <Compile Include="Workspace\Host\Mef\DesktopMefHostServices.cs" />
     <Compile Include="Workspace\Host\Mef\MefV1HostServices.cs" />
     <Compile Include="Workspace\Host\SimpleAnalyzerAssemblyLoaderService.cs" />
-    <Compile Include="Workspace\Host\TemporaryStorage\TemporaryStorageServiceFactory.CopiedMemoryMappedViewHandle.cs" />
     <Compile Include="Workspace\Host\TemporaryStorage\TemporaryStorageServiceFactory.cs" />
     <Compile Include="Workspace\Host\TemporaryStorage\TemporaryStorageServiceFactory.MemoryMappedInfo.cs" />
+    <Compile Include="Workspace\Host\TemporaryStorage\TemporaryStorageServiceFactory.ReferenceCountedDisposable.cs" />
     <Compile Include="Workspace\Host\TextFactory\DesktopTextFactoryService.cs" />
     <Compile Include="Workspace\MSBuild\CSharp\CSharpProjectFileLoader.cs" />
     <Compile Include="Workspace\MSBuild\CSharp\CSharpProjectFileLoader.CSharpProjectFile.cs" />