Make sure to cancel tasks with the associated token

When we were cancelling tasks returned from GetValueAsync, we were
not passing along the CancellationToken because TaskCompletionSource
doesn't allow us to pass that in .NET 4.5. The appropriate method was
added in .NET 4.6 but we haven't moved to it yet. In the mean time
we'll just use AsyncTaskMethodBuilder which does provide the
necessary functionality.

Fixes GitHub issue #447.

src/Workspaces/Core/Portable/Utilities/AsyncLazy`1.cs

@@ -2,6 +2,7 @@
 
 using System;
 using System.Collections.Generic;
+using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.ErrorReporting;
@@ -396,11 +397,13 @@ private void StartAsynchronousComputation(AsynchronousComputationToStart computa
             }
             catch (OperationCanceledException oce) when (CrashIfCanceledWithDifferentToken(oce, cancellationToken))
             {
-                // The underlying computation cancelled with the correct token, but we must ourselves insure that the caller
+                // The underlying computation cancelled with the correct token, but we must ourselves ensure that the caller
                 // on our stack gets an OperationCanceledException thrown with the right token
                 requestToCompleteSynchronouslyCancellationToken.ThrowIfCancellationRequested();
 
-                // The cancellation must have been requested, so this is unreachable
+                // We can only be here if the computation was cancelled, which means all requests for the value
+                // must have been cancelled. Therefore, the ThrowIfCancellationRequested above must have thrown
+                // because that token from the requestor was cancelled.
                 throw ExceptionUtilities.Unreachable;
             }
         }
@@ -504,13 +507,38 @@ private void OnAsynchronousRequestCancelled(object state)
             }
         }
 
-        // Using inheritance instead of wrapping a TaskCompletionSource to avoid a second allocation
-        private class Request : TaskCompletionSource<T>
+        private sealed class Request
         {
+            /// <summary>
+            /// The <see cref="CancellationToken"/> associated with this request. This field will be initialized before
+            /// any cancellation is observed from the token.
+            /// </summary>
+            private CancellationToken _cancellationToken;
             private CancellationTokenRegistration _cancellationTokenRegistration;
 
+            // We use a AsyncTaskMethodBuilder so we have the ability to cancel the task with a given cancellation token
+            // TODO: remove this once we're on .NET 4.6 and can move back to using TaskCompletionSource.
+            // WARNING: this is a mutable struct, and thus cannot be made readonly
+            private AsyncTaskMethodBuilder<T> _taskBuilder;
+
+            public Request()
+            {
+                // .Task on AsyncTaskMethodBuilder is lazily created in a non-synchronized way, so we must request it
+                // once before we start doing fancy stuff
+                var ignored = _taskBuilder.Task;
+            }
+
+            public Task<T> Task
+            {
+                get
+                {
+                    return _taskBuilder.Task;
+                }
+            }
+
             public void RegisterForCancellation(Action<object> callback, CancellationToken cancellationToken)
             {
+                _cancellationToken = cancellationToken;
                 _cancellationTokenRegistration = cancellationToken.Register(callback, this);
             }
 
@@ -526,24 +554,36 @@ private void CompleteFromTaskSynchronouslyStub(object task)
 
             public void CompleteFromTaskSynchronously(Task<T> task)
             {
-                if (task.Status == TaskStatus.RanToCompletion)
+                // AsyncTaskMethodBuilder doesn't give us Try* methods, and the Set methods may throw if the task
+                // is already completed. The belief is that the race is somewhere between rare to impossible, and
+                // so we'll do a quick check to see if the task is already completed or otherwise just give it a shot
+                // and catch it if it fails
+                if (_taskBuilder.Task.IsCompleted)
                 {
-                    if (TrySetResult(task.Result))
-                    {
-                        _cancellationTokenRegistration.Dispose();
-                    }
+                    return;
                 }
-                else if (task.Status == TaskStatus.Faulted)
+
+                try
                 {
-                    if (TrySetException(task.Exception))
+                    if (task.Status == TaskStatus.RanToCompletion)
                     {
-                        _cancellationTokenRegistration.Dispose();
+                        _taskBuilder.SetResult(task.Result);
+                    }
+                    else if (task.Status == TaskStatus.Faulted)
+                    {
+                        _taskBuilder.SetException(task.Exception);
+                    }
+                    else
+                    {
+                        CancelSynchronously();
                     }
                 }
-                else
+                catch (InvalidOperationException)
                 {
-                    CancelSynchronously();
+                    // Something else beat us to setting the state, so bail
                 }
+
+                _cancellationTokenRegistration.Dispose();
             }
 
             public void CancelAsynchronously()
@@ -555,11 +595,22 @@ public void CancelAsynchronously()
 
             private void CancelSynchronously()
             {
-                if (TrySetCanceled())
+                // AsyncTaskMethodBuilder doesn't give us Try* methods, and the Set methods may throw if the task
+                // is already completed. The belief is that the race is somewhere between rare to impossible, and
+                // so we'll do a quick check to see if the task is already completed or otherwise just give it a shot
+                // and catch it if it fails
+                if (_taskBuilder.Task.IsCompleted)
+                {
+                    return;
+                }
+
+                try
+                {
+                    _taskBuilder.SetException(new OperationCanceledException(_cancellationToken));
+                }
+                catch (InvalidOperationException)
                 {
-                    // Paranoia: the only reason we should ever get here is if the CancellationToken that
-                    // we registered against was cancelled, but just in case, dispose the registration
-                    _cancellationTokenRegistration.Dispose();
+                    // Something else beat us to setting the state, so bail
                 }
             }
         }

src/Workspaces/CoreTest/UtilityTest/AsyncLazyTests.cs

@@ -296,5 +296,35 @@ private static void GetValueOrGetValueAsyncThrowsCorrectExceptionDuringCancellat
                 Assert.Equal(cancellationTokenSource.Token, oce.CancellationToken);
             }
         }
+
+        [Fact]
+        [Trait(Traits.Feature, Traits.Features.AsyncLazy)]
+        public void GetValueAsyncThatIsCancelledReturnsTaskCancelledWithCorrectToken()
+        {
+            var cancellationTokenSource = new CancellationTokenSource();
+
+            var lazy = new AsyncLazy<object>(c => Task.Run((Func<object>)(() =>
+            {
+                cancellationTokenSource.Cancel();
+                while (true)
+                {
+                    c.ThrowIfCancellationRequested();
+                }
+            }), c), cacheResult: true);
+
+            var task = lazy.GetValueAsync(cancellationTokenSource.Token);
+
+            // Now wait until the task completes
+            try
+            {
+                task.Wait();
+                AssertEx.Fail(nameof(AsyncLazy<object>.GetValueAsync) + " did not throw an exception.");
+            }
+            catch (AggregateException ex)
+            {
+                var operationCancelledException = (OperationCanceledException)ex.Flatten().InnerException;
+                Assert.Equal(cancellationTokenSource.Token, operationCancelledException.CancellationToken);
+            }
+        }
     }
 }