[test-proxy] Only remove session once body has been written (#8709)

* Only remove session once body has been written, to minimize probability client retries but test-proxy has already removed the session
* remove the necessity of always locking to check if a session is sanitized
* add entry remove, call it
* optimize a few lock calls to happen once where possible

---------

Co-authored-by: Scott Beddall <scbedd@microsoft.com>

Author: mikeharder

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,12 +103,30 @@ 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)
             {
                 sanitizer.Sanitize(this);
             }
         }
+        public void Sanitize(IEnumerable<RecordedTestSanitizer> sanitizers)
+        {
+            lock (Entries)
+            {
+                foreach(var sanitizer in sanitizers)
+                {
+                    sanitizer.Sanitize(this);
+                }
+            }
+        }
     }
 }

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

@@ -114,11 +114,7 @@ public void StopRecording(string sessionId, IDictionary<string, string> variable
             }
 
             var sanitizers = SanitizerRegistry.GetSanitizers(recordingSession);
-
-            foreach (RecordedTestSanitizer sanitizer in sanitizers)
-            {
-                recordingSession.Session.Sanitize(sanitizer);
-            }
+            recordingSession.Session.Sanitize(sanitizers);
 
             if (variables != null)
             {
@@ -457,16 +453,15 @@ 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.SanitizerLock)
                 {
-                    session.IsSanitized = true;
-
-                    foreach (RecordedTestSanitizer sanitizer in sanitizers)
+                    if (!session.IsSanitized)
                     {
-                        session.Session.Sanitize(sanitizer);
+                        session.Session.Sanitize(sanitizers);
+                        session.IsSanitized = true;
                     }
                 }
             }
@@ -475,23 +470,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 +498,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)