[test-proxy] Only remove session once body has been written

tools/test-proxy/Azure.Sdk.Tools.TestProxy/Common/RecordSession.cs

@@ -1,4 +1,4 @@
-// Copyright (c) Microsoft Corporation. All rights reserved.
+// Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
 using System;
@@ -84,7 +84,7 @@ public void Record(RecordEntry entry)
 
         public RecordEntry Lookup(RecordEntry requestEntry, RecordMatcher matcher, IEnumerable<RecordedTestSanitizer> sanitizers, bool remove = true)
         {
-            foreach(RecordedTestSanitizer sanitizer in sanitizers)
+            foreach (RecordedTestSanitizer sanitizer in sanitizers)
             {
                 sanitizer.Sanitize(requestEntry);
             }
@@ -103,6 +103,14 @@ public RecordEntry Lookup(RecordEntry requestEntry, RecordMatcher matcher, IEnum
             }
         }
 
+        public void Remove(RecordEntry entry)
+        {
+            lock (Entries)
+            {
+                Entries.Remove(entry);
+            }
+        }
+
         public void Sanitize(RecordedTestSanitizer sanitizer)
         {
             lock (Entries)

tools/test-proxy/Azure.Sdk.Tools.TestProxy/RecordingHandler.cs

@@ -457,16 +457,19 @@ public async Task HandlePlaybackRequest(string recordingId, HttpRequest incoming
 
             var sanitizers = SanitizerRegistry.GetSanitizers(session);
 
-            // we don't need to re-sanitize with recording-applicable sanitizers every time. just the very first one
-            lock (session)
+            if (!session.IsSanitized)
             {
-                if (!session.IsSanitized)
+                // we don't need to re-sanitize with recording-applicable sanitizers every time. just the very first one
+                lock (session)
                 {
-                    session.IsSanitized = true;
-
-                    foreach (RecordedTestSanitizer sanitizer in sanitizers)
+                    if (!session.IsSanitized)
                     {
-                        session.Session.Sanitize(sanitizer);
+                        foreach (RecordedTestSanitizer sanitizer in sanitizers)
+                        {
+                            session.Session.Sanitize(sanitizer);
+                        }
+
+                        session.IsSanitized = true;
                     }
                 }
             }
@@ -475,23 +478,14 @@ public async Task HandlePlaybackRequest(string recordingId, HttpRequest incoming
 
             var entry = (await CreateEntryAsync(incomingRequest).ConfigureAwait(false)).Item1;
 
-            // If request contains "x-recording-remove: false", then request is not removed from session after playback.
-            // Used by perf tests to play back the same request multiple times.
-            var remove = true;
-            if (incomingRequest.Headers.TryGetValue("x-recording-remove", out var removeHeader))
-            {
-                remove = bool.Parse(removeHeader);
-            }
-
-            var match = session.Session.Lookup(entry, session.CustomMatcher ?? Matcher, sanitizers, remove);
+            // Session may be removed later, but only after response has been fully written
+            var match = session.Session.Lookup(entry, session.CustomMatcher ?? Matcher, sanitizers, remove: false);
 
             foreach (ResponseTransform transform in Transforms.Concat(session.AdditionalTransforms))
             {
                 transform.Transform(incomingRequest, match);
             }
 
-            Interlocked.Increment(ref Startup.RequestsPlayedBack);
-
             outgoingResponse.StatusCode = match.StatusCode;
 
             foreach (var header in match.Response.Headers)
@@ -512,6 +506,23 @@ public async Task HandlePlaybackRequest(string recordingId, HttpRequest incoming
 
                 await WriteBodyBytes(bodyData, session.PlaybackResponseTime, outgoingResponse);
             }
+
+            Interlocked.Increment(ref Startup.RequestsPlayedBack);
+
+            // Only remove session once body has been written, to minimize probability client retries but test-proxy has already removed the session
+            var remove = true;
+
+            // If request contains "x-recording-remove: false", then request is not removed from session after playback.
+            // Used by perf tests to play back the same request multiple times.
+            if (incomingRequest.Headers.TryGetValue("x-recording-remove", out var removeHeader))
+            {
+                remove = bool.Parse(removeHeader);
+            }
+
+            if (remove)
+            {
+                session.Session.Remove(match);
+            }
         }
 
         public byte[][] GetBatches(byte[] bodyData, int batchCount)