From 0afff581266c9db39699369d82b354a5ae11eb2f Mon Sep 17 00:00:00 2001
From: Mike Harder <mharder@microsoft.com>
Date: Thu, 20 May 2021 17:01:18 -0700
Subject: [PATCH] Only ignore OperationCanceledException if cancellationToken
 is cancelled

---
 .../OperationCanceledTest.cs                  | 32 +++++++++++++++++++
 common/Perf/Azure.Test.Perf/PerfProgram.cs    | 17 ++++++++--
 2 files changed, 46 insertions(+), 3 deletions(-)
 create mode 100644 common/Perf/Azure.Sample.Perf/OperationCanceledTest.cs

diff --git a/common/Perf/Azure.Sample.Perf/OperationCanceledTest.cs b/common/Perf/Azure.Sample.Perf/OperationCanceledTest.cs
new file mode 100644
index 0000000000000..eeb493e256ffa
--- /dev/null
+++ b/common/Perf/Azure.Sample.Perf/OperationCanceledTest.cs
@@ -0,0 +1,32 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using Azure.Test.Perf;
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Azure.Sample.Perf
+{
+    // The perf framework should ignore OperationCanceledException thrown by Run/RunAsync(), but only
+    // if the test is being cancelled.  This verifies that an OperationCanceledException thrown earlier
+    // will cause the perf framework to fail fast.
+    // https://github.com/Azure/azure-sdk-for-net/issues/21241
+
+    public class OperationCanceledTest : PerfTest<PerfOptions>
+    {
+        public OperationCanceledTest(PerfOptions options) : base(options)
+        {
+        }
+
+        public override void Run(CancellationToken cancellationToken)
+        {
+            throw new OperationCanceledException();
+        }
+
+        public override Task RunAsync(CancellationToken cancellationToken)
+        {
+            throw new OperationCanceledException();
+        }
+    }
+}
diff --git a/common/Perf/Azure.Test.Perf/PerfProgram.cs b/common/Perf/Azure.Test.Perf/PerfProgram.cs
index 2ef3d49943205..c4684f317cf98 100644
--- a/common/Perf/Azure.Test.Perf/PerfProgram.cs
+++ b/common/Perf/Azure.Test.Perf/PerfProgram.cs
@@ -441,8 +441,16 @@ private static void RunLoop(IPerfTest test, int index, bool latency, Cancellatio
                     _lastCompletionTimes[index] = sw.Elapsed;
                 }
             }
-            catch (OperationCanceledException)
+            catch (Exception e)
             {
+                if (cancellationToken.IsCancellationRequested && PerfStressUtilities.ContainsOperationCanceledException(e))
+                {
+                    // If the test has been canceled, ignore if any part of the exception chain is OperationCanceledException.
+                }
+                else
+                {
+                    throw;
+                }
             }
         }
 
@@ -484,8 +492,11 @@ private static async Task RunLoopAsync(IPerfTest test, int index, bool latency,
             }
             catch (Exception e)
             {
-                // Ignore if any part of the exception chain is type OperationCanceledException
-                if (!PerfStressUtilities.ContainsOperationCanceledException(e))
+                if (cancellationToken.IsCancellationRequested && PerfStressUtilities.ContainsOperationCanceledException(e))
+                {
+                    // If the test has been canceled, ignore if any part of the exception chain is OperationCanceledException.
+                }
+                else
                 {
                     throw;
                 }