Merge pull request #2307 from Clockwork-Muse/PLINQ_ArrayCancelImmediate

PLINQ - immediate cancellation of ToArray

Author: stephentoub

src/System.Linq.Parallel/src/System/Linq/Parallel/QueryOperators/QueryOperator.cs

@@ -209,6 +209,18 @@ internal TOutput[] ExecuteAndGetResultsAsArray()
 
                 QueryResults<TOutput> results = GetQueryResults(querySettings);
 
+                // Top-level preemptive cancellation test.
+                // This handles situations where cancellation has occured before execution commences
+                // The handling for in-execution occurs in QueryTaskGroupState.QueryEnd()
+
+                if (querySettings.CancellationState.MergedCancellationToken.IsCancellationRequested)
+                {
+                    if (querySettings.CancellationState.ExternalCancellationToken.IsCancellationRequested)
+                        throw new OperationCanceledException(querySettings.CancellationState.ExternalCancellationToken);
+                    else
+                        throw new OperationCanceledException();
+                }
+
                 if (results.IsIndexible && OutputOrdered)
                 {
                     // The special array-based merge performs better if the output is ordered, because

src/System.Linq.Parallel/tests/Helpers/UnorderedSources.cs

@@ -190,6 +190,13 @@ public static IEnumerable<object[]> Ranges<T>(IEnumerable<int> counts, params Fu
             }
         }
 
+        // Return an enumerable which throws on first MoveNext.
+        // Useful for testing promptness of cancellation.
+        public static IEnumerable<object[]> ThrowOnFirstEnumeration()
+        {
+            yield return new object[] { Labeled.Label("ThrowOnFirstEnumeration", Enumerables<int>.ThrowOnEnumeration().AsParallel()), 8 };
+        }
+
         private static IEnumerable<Labeled<ParallelQuery<int>>> LabeledRanges(int start, int count)
         {
             yield return Labeled.Label("ParallelEnumerable.Range", ParallelEnumerable.Range(start, count));

src/System.Linq.Parallel/tests/QueryOperators/ToArrayTests.cs

@@ -48,6 +48,7 @@ public static void ToArray_Longrunning(Labeled<ParallelQuery<int>> labeled, int
 
         [Theory]
         [MemberData("Ranges", (object)(new int[] { 1 }), MemberType = typeof(Sources))]
+        [MemberData("ThrowOnFirstEnumeration", MemberType = typeof(UnorderedSources))]
         public static void ToArray_OperationCanceledException_PreCanceled(Labeled<ParallelQuery<int>> labeled, int count)
         {
             CancellationTokenSource cs = new CancellationTokenSource();