Add AsyncLazy<T>.Dispose() method

This addresses a common need for a class with an `AsyncLazy<T>` field where the value is disposable to ensure on the class disposal that it can dispose of the value if and when it is constructed.

src/Microsoft.VisualStudio.Threading/AsyncLazy`1.cs

@@ -12,6 +12,11 @@ namespace Microsoft.VisualStudio.Threading;
 /// A thread-safe, lazily and asynchronously evaluated value factory.
 /// </summary>
 /// <typeparam name="T">The type of value generated by the value factory.</typeparam>
+/// <remarks>
+/// This class does not itself carry any resources needful of disposing.
+/// But the value factory may produce a value that needs to be disposed of,
+/// which is why this class carries a <see cref="Dispose"/> method but does not implement <see cref="IDisposable"/>.
+/// </remarks>
 public class AsyncLazy<T>
 {
     /// <summary>
@@ -20,6 +25,11 @@ public class AsyncLazy<T>
     /// </summary>
     private static readonly object RecursiveCheckSentinel = new object();
 
+    /// <summary>
+    /// A value set on the <see cref="value"/> field when this object is disposed.
+    /// </summary>
+    private static readonly Task<T> DisposedSentinel = Task.FromException<T>(new ObjectDisposedException(nameof(AsyncLazy<T>)));
+
     /// <summary>
     /// The object to lock to provide thread-safety.
     /// </summary>
@@ -65,24 +75,30 @@ public AsyncLazy(Func<Task<T>> valueFactory, JoinableTaskFactory? joinableTaskFa
     /// <summary>
     /// Gets a value indicating whether the value factory has been invoked.
     /// </summary>
+    /// <remarks>
+    /// This returns <see langword="false" /> after a call to <see cref="Dispose"/>.
+    /// </remarks>
     public bool IsValueCreated
     {
         get
         {
             Interlocked.MemoryBarrier();
-            return this.valueFactory is null;
+            return this.valueFactory is null && this.value != DisposedSentinel;
         }
     }
 
     /// <summary>
     /// Gets a value indicating whether the value factory has been invoked and has run to completion.
     /// </summary>
+    /// <remarks>
+    /// This returns <see langword="false" /> after a call to <see cref="Dispose"/>.
+    /// </remarks>
     public bool IsValueFactoryCompleted
     {
         get
         {
             Interlocked.MemoryBarrier();
-            return this.value is object && this.value.IsCompleted;
+            return this.value is object && this.value.IsCompleted && this.value != DisposedSentinel;
         }
     }
 
@@ -93,6 +109,7 @@ public bool IsValueFactoryCompleted
     /// <exception cref="InvalidOperationException">
     /// Thrown when the value factory calls <see cref="GetValueAsync()"/> on this instance.
     /// </exception>
+    /// <exception cref="ObjectDisposedException">Thrown after <see cref="Dispose"/> is called.</exception>
     public Task<T> GetValueAsync() => this.GetValueAsync(CancellationToken.None);
 
     /// <summary>
@@ -108,6 +125,7 @@ public bool IsValueFactoryCompleted
     /// <exception cref="InvalidOperationException">
     /// Thrown when the value factory calls <see cref="GetValueAsync()"/> on this instance.
     /// </exception>
+    /// <exception cref="ObjectDisposedException">Thrown after <see cref="Dispose"/> is called.</exception>
     public Task<T> GetValueAsync(CancellationToken cancellationToken)
     {
         if (!((this.value is object && this.value.IsCompleted) || this.recursiveFactoryCheck.Value is null))
@@ -234,6 +252,69 @@ public T GetValue(CancellationToken cancellationToken)
         }
     }
 
+    /// <summary>
+    /// Disposes of the lazily-initialized value if disposable, and causes all subsequent attempts to obtain the value to fail.
+    /// </summary>
+    /// <remarks>
+    /// <para>Calling this method will put this object into a disposed state where future calls to obtain the value will throw <see cref="ObjectDisposedException"/>.</para>
+    /// <para>If the value has already been produced and implements <see cref="IDisposable"/>, it will be disposed of.
+    /// If the value factory has already started but has not yet completed, its value will be disposed of when the value factory completes.</para>
+    /// <para>If prior calls to obtain the value are in flight when this method is called, those calls <em>may</em> complete and their callers may obtain the value, although <see cref="IDisposable.Dispose"/>
+    /// may have been or will soon be called on the value, leading those users to experience a <see cref="ObjectDisposedException"/>.</para>
+    /// <para>Note all conditions based on the value implementing <see cref="IDisposable"/> is based on the actual value, rather than the <typeparamref name="T"/> type argument.
+    /// This means that although <typeparamref name="T"/> may be <c>IFoo</c> (which does not implement <see cref="IDisposable"/>), the concrete type that implements <c>IFoo</c> may implement <see cref="IDisposable"/>
+    /// and thus be treated as a disposable object as described above.</para>
+    /// </remarks>
+    public void Dispose()
+    {
+        Task<T>? localValueTask = null;
+        IDisposable? localValue = default;
+        lock (this.syncObject)
+        {
+            if (this.value == DisposedSentinel)
+            {
+                return;
+            }
+
+            switch (this.value?.Status)
+            {
+                case TaskStatus.RanToCompletion:
+                    // We'll dispose of the value inline, outside the lock.
+                    localValue = this.value.Result as IDisposable;
+                    break;
+                case TaskStatus.Faulted:
+                case TaskStatus.Canceled:
+                    // Nothing left to do.
+                    break;
+                default:
+                    // We'll schedule the value for disposal outside the lock so it can be synchronous with the value factory,
+                    // but will not execute within our lock.
+                    localValueTask = this.value;
+                    break;
+            }
+
+            // Shut out all future callers from obtaining the value.
+            this.value = DisposedSentinel;
+
+            // Release associated memory.
+            this.joinableTask = null;
+            this.valueFactory = null;
+        }
+
+        if (localValue is not null)
+        {
+            localValue.Dispose();
+        }
+        else if (localValueTask is not null)
+        {
+            localValueTask.ContinueWith(
+                static v => (v.Result as IDisposable)?.Dispose(),
+                CancellationToken.None,
+                TaskContinuationOptions.OnlyOnRanToCompletion | TaskContinuationOptions.ExecuteSynchronously,
+                TaskScheduler.Default);
+        }
+    }
+
     /// <summary>
     /// Renders a string describing an uncreated value, or the string representation of the created value.
     /// </summary>

src/Microsoft.VisualStudio.Threading/net472/PublicAPI.Unshipped.txt

@@ -0,0 +1 @@
+Microsoft.VisualStudio.Threading.AsyncLazy<T>.Dispose() -> void

src/Microsoft.VisualStudio.Threading/net6.0-windows/PublicAPI.Unshipped.txt

@@ -0,0 +1 @@
+Microsoft.VisualStudio.Threading.AsyncLazy<T>.Dispose() -> void

src/Microsoft.VisualStudio.Threading/net6.0/PublicAPI.Unshipped.txt

@@ -0,0 +1 @@
+Microsoft.VisualStudio.Threading.AsyncLazy<T>.Dispose() -> void

src/Microsoft.VisualStudio.Threading/netstandard2.0/PublicAPI.Unshipped.txt

@@ -0,0 +1 @@
+Microsoft.VisualStudio.Threading.AsyncLazy<T>.Dispose() -> void

test/Microsoft.VisualStudio.Threading.Tests/AsyncLazyTests.cs

@@ -6,6 +6,7 @@
 using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft;
 using Microsoft.VisualStudio.Threading;
 using Xunit;
 using Xunit.Abstractions;
@@ -644,6 +645,102 @@ public async Task ExecutionContextFlowsFromFirstCaller_JTF()
         await asyncLazy.GetValueAsync();
     }
 
+    [Fact]
+    public async Task Dispose_ValueType_Completed()
+    {
+        AsyncLazy<int> lazy = new(() => Task.FromResult(3));
+        lazy.GetValue();
+        lazy.Dispose();
+        await this.AssertDisposedLazyAsync(lazy);
+    }
+
+    [Fact]
+    public async Task Dispose_Disposable_Completed()
+    {
+        AsyncLazy<object> lazy = new(() => Task.FromResult<object>(new Disposable()));
+        Disposable value = (Disposable)lazy.GetValue();
+        lazy.Dispose();
+        Assert.True(value.IsDisposed);
+        await this.AssertDisposedLazyAsync(lazy);
+    }
+
+    [Fact]
+    public async Task Dispose_NonDisposable_Completed()
+    {
+        AsyncLazy<object> lazy = new(() => Task.FromResult(new object()));
+        lazy.GetValue();
+        lazy.Dispose();
+        await this.AssertDisposedLazyAsync(lazy);
+    }
+
+    [Fact]
+    public async Task Dispose_Disposable_Incomplete()
+    {
+        AsyncManualResetEvent unblock = new();
+        AsyncLazy<object> lazy = new(async delegate
+        {
+            await unblock;
+            return new Disposable();
+        });
+        Task<object> lazyTask = lazy.GetValueAsync(this.TimeoutToken);
+        lazy.Dispose();
+        unblock.Set();
+        Disposable value = (Disposable)await lazyTask;
+        await this.AssertDisposedLazyAsync(lazy);
+        await value.Disposed.WithCancellation(this.TimeoutToken);
+    }
+
+    [Fact]
+    public async Task Dispose_NonDisposable_Incomplete()
+    {
+        AsyncManualResetEvent unblock = new();
+        AsyncLazy<object> lazy = new(async delegate
+        {
+            await unblock;
+            return new object();
+        });
+        Task<object> lazyTask = lazy.GetValueAsync(this.TimeoutToken);
+        lazy.Dispose();
+        unblock.Set();
+        await lazyTask;
+        await this.AssertDisposedLazyAsync(lazy);
+    }
+
+    [Fact]
+    public async Task Dispose_CalledTwice_NotStarted()
+    {
+        bool valueFactoryExecuted = false;
+        AsyncLazy<object> lazy = new(() =>
+        {
+            valueFactoryExecuted = true;
+            return Task.FromResult(new object());
+        });
+        lazy.Dispose();
+        lazy.Dispose();
+        await this.AssertDisposedLazyAsync(lazy);
+        Assert.False(valueFactoryExecuted);
+    }
+
+    [Fact]
+    public async Task Dispose_CalledTwice_NonDisposable_Completed()
+    {
+        AsyncLazy<object> lazy = new(() => Task.FromResult(new object()));
+        lazy.GetValue();
+        lazy.Dispose();
+        lazy.Dispose();
+        await this.AssertDisposedLazyAsync(lazy);
+    }
+
+    [Fact]
+    public async Task Dispose_CalledTwice_Disposable_Completed()
+    {
+        AsyncLazy<Disposable> lazy = new(() => Task.FromResult(new Disposable()));
+        lazy.GetValue();
+        lazy.Dispose();
+        lazy.Dispose();
+        await this.AssertDisposedLazyAsync(lazy);
+    }
+
     [Fact(Skip = "Hangs. This test documents a deadlock scenario that is not fixed (by design, IIRC).")]
     public async Task ValueFactoryRequiresReadLockHeldByOther()
     {
@@ -727,4 +824,23 @@ private async Task<WeakReference> AsyncPumpReleasedAfterExecution_Helper(bool th
         await lazy.GetValueAsync().NoThrowAwaitable();
         return collectible;
     }
+
+    private async Task AssertDisposedLazyAsync<T>(AsyncLazy<T> lazy)
+    {
+        Assert.False(lazy.IsValueCreated);
+        Assert.False(lazy.IsValueFactoryCompleted);
+        Assert.Throws<ObjectDisposedException>(() => lazy.GetValue());
+        await Assert.ThrowsAsync<ObjectDisposedException>(lazy.GetValueAsync);
+    }
+
+    private class Disposable : IDisposableObservable
+    {
+        private readonly AsyncManualResetEvent disposalEvent = new();
+
+        public Task Disposed => this.disposalEvent.WaitAsync();
+
+        public bool IsDisposed => this.disposalEvent.IsSet;
+
+        public void Dispose() => this.disposalEvent.Set();
+    }
 }