Remove HttpMetricsEnrichmentContext caching

src/libraries/System.Net.Http/src/System/Net/Http/Metrics/HttpMetricsEnrichmentContext.cs

@@ -1,12 +1,9 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Diagnostics;
 using System.Diagnostics.Metrics;
-using System.Runtime.InteropServices;
-using System.Threading;
 
 namespace System.Net.Http.Metrics
 {
@@ -19,22 +16,18 @@ namespace System.Net.Http.Metrics
     /// information exposed on the <see cref="HttpMetricsEnrichmentContext"/> instance.
     ///
     /// > [!IMPORTANT]
-    /// > The <see cref="HttpMetricsEnrichmentContext"/> intance must not be used from outside of the enrichment callbacks.
+    /// > The <see cref="HttpMetricsEnrichmentContext"/> instance must not be used from outside of the enrichment callbacks.
     /// </remarks>
     public sealed class HttpMetricsEnrichmentContext
     {
-        private static readonly HttpRequestOptionsKey<HttpMetricsEnrichmentContext> s_optionsKeyForContext = new(nameof(HttpMetricsEnrichmentContext));
-        private static readonly ConcurrentQueue<HttpMetricsEnrichmentContext> s_pool = new();
-        private static int s_poolItemCount;
-        private const int PoolCapacity = 1024;
+        private static readonly HttpRequestOptionsKey<List<Action<HttpMetricsEnrichmentContext>>> s_optionsKeyForCallbacks = new(nameof(HttpMetricsEnrichmentContext));
 
-        private readonly List<Action<HttpMetricsEnrichmentContext>> _callbacks = new();
         private HttpRequestMessage? _request;
         private HttpResponseMessage? _response;
         private Exception? _exception;
-        private List<KeyValuePair<string, object?>> _tags = new(capacity: 16);
+        private TagList _tags;
 
-        internal HttpMetricsEnrichmentContext() { } // Hide the default parameterless constructor.
+        private HttpMetricsEnrichmentContext() { } // Hide the default parameterless constructor.
 
         /// <summary>
         /// Gets the <see cref="HttpRequestMessage"/> that has been sent.
@@ -68,7 +61,7 @@ public sealed class HttpMetricsEnrichmentContext
         /// <remarks>
         /// This method must not be used from outside of the enrichment callbacks.
         /// </remarks>
-        public void AddCustomTag(string name, object? value) => _tags.Add(new KeyValuePair<string, object?>(name, value));
+        public void AddCustomTag(string name, object? value) => _tags.Add(name, value);
 
         /// <summary>
         /// Adds a callback to register custom tags for request metrics `http-client-request-duration` and `http-client-failed-requests`.
@@ -77,85 +70,55 @@ public sealed class HttpMetricsEnrichmentContext
         /// <param name="callback">The callback responsible to add custom tags via <see cref="AddCustomTag(string, object?)"/>.</param>
         public static void AddCallback(HttpRequestMessage request, Action<HttpMetricsEnrichmentContext> callback)
         {
+            ArgumentNullException.ThrowIfNull(request);
+            ArgumentNullException.ThrowIfNull(callback);
+
             HttpRequestOptions options = request.Options;
 
-            // We associate an HttpMetricsEnrichmentContext with the request on the first call to AddCallback(),
-            // and store the callbacks in the context. This allows us to cache all the enrichment objects together.
-            if (!options.TryGetValue(s_optionsKeyForContext, out HttpMetricsEnrichmentContext? context))
+            if (options.TryGetValue(s_optionsKeyForCallbacks, out List<Action<HttpMetricsEnrichmentContext>>? callbacks))
+            {
+                callbacks.Add(callback);
+            }
+            else
             {
-                if (s_pool.TryDequeue(out context))
-                {
-                    Debug.Assert(context._callbacks.Count == 0);
-                    Interlocked.Decrement(ref s_poolItemCount);
-                }
-                else
-                {
-                    context = new HttpMetricsEnrichmentContext();
-                }
-
-                options.Set(s_optionsKeyForContext, context);
+                options.Set(s_optionsKeyForCallbacks, [callback]);
             }
-            context._callbacks.Add(callback);
         }
 
-        internal static HttpMetricsEnrichmentContext? GetEnrichmentContextForRequest(HttpRequestMessage request)
+        internal static List<Action<HttpMetricsEnrichmentContext>>? GetEnrichmentCallbacksForRequest(HttpRequestMessage request)
         {
-            if (request._options is null)
+            if (request._options is HttpRequestOptions options &&
+                options.Remove(s_optionsKeyForCallbacks.Key, out object? callbacks))
             {
-                return null;
+                return (List<Action<HttpMetricsEnrichmentContext>>)callbacks!;
             }
-            request._options.TryGetValue(s_optionsKeyForContext, out HttpMetricsEnrichmentContext? context);
-            return context;
+
+            return null;
         }
 
-        internal void RecordDurationWithEnrichment(HttpRequestMessage request,
+        internal static void RecordDurationWithEnrichment(
+            List<Action<HttpMetricsEnrichmentContext>> callbacks,
+            HttpRequestMessage request,
             HttpResponseMessage? response,
             Exception? exception,
             TimeSpan durationTime,
             in TagList commonTags,
             Histogram<double> requestDuration)
         {
-            _request = request;
-            _response = response;
-            _exception = exception;
-
-            Debug.Assert(_tags.Count == 0);
-
-            // Adding the enrichment tags to the TagList would likely exceed its' on-stack capacity, leading to an allocation.
-            // To avoid this, we add all the tags to a List<T> which is cached together with HttpMetricsEnrichmentContext.
-            // Use a for loop to iterate over the TagList, since TagList.GetEnumerator() allocates, see
-            // https://github.com/dotnet/runtime/issues/87022.
-            for (int i = 0; i < commonTags.Count; i++)
+            var context = new HttpMetricsEnrichmentContext
             {
-                _tags.Add(commonTags[i]);
-            }
+                _request = request,
+                _response = response,
+                _exception = exception,
+                _tags = commonTags,
+            };
 
-            try
+            foreach (Action<HttpMetricsEnrichmentContext> callback in callbacks)
             {
-                foreach (Action<HttpMetricsEnrichmentContext> callback in _callbacks)
-                {
-                    callback(this);
-                }
-
-                requestDuration.Record(durationTime.TotalSeconds, CollectionsMarshal.AsSpan(_tags));
-            }
-            finally
-            {
-                _request = null;
-                _response = null;
-                _exception = null;
-                _callbacks.Clear();
-                _tags.Clear();
-
-                if (Interlocked.Increment(ref s_poolItemCount) <= PoolCapacity)
-                {
-                    s_pool.Enqueue(this);
-                }
-                else
-                {
-                    Interlocked.Decrement(ref s_poolItemCount);
-                }
+                callback(context);
             }
+
+            requestDuration.Record(durationTime.TotalSeconds, context._tags);
         }
     }
 }

src/libraries/System.Net.Http/src/System/Net/Http/Metrics/MetricsHandler.cs

@@ -120,14 +120,14 @@ private void RequestStop(HttpRequestMessage request, HttpResponseMessage? respon
 
             TimeSpan durationTime = Stopwatch.GetElapsedTime(startTimestamp, Stopwatch.GetTimestamp());
 
-            HttpMetricsEnrichmentContext? enrichmentContext = HttpMetricsEnrichmentContext.GetEnrichmentContextForRequest(request);
-            if (enrichmentContext is null)
+            List<Action<HttpMetricsEnrichmentContext>>? callbacks = HttpMetricsEnrichmentContext.GetEnrichmentCallbacksForRequest(request);
+            if (callbacks is null)
             {
                 _requestsDuration.Record(durationTime.TotalSeconds, tags);
             }
             else
             {
-                enrichmentContext.RecordDurationWithEnrichment(request, response, exception, durationTime, tags, _requestsDuration);
+                HttpMetricsEnrichmentContext.RecordDurationWithEnrichment(callbacks, request, response, exception, durationTime, tags, _requestsDuration);
             }
         }
 

src/libraries/System.Net.Http/tests/FunctionalTests/MetricsTest.cs

@@ -354,6 +354,49 @@ public Task RequestDuration_CustomTags_Recorded()
                 VerifyRequestDuration(m, uri, UseVersion, 200);
                 Assert.Equal("/test", m.Tags.ToArray().Single(t => t.Key == "route").Value);
 
+            }, async server =>
+            {
+                await server.HandleRequestAsync();
+            });
+        }
+
+        [Fact]
+        public Task RequestDuration_MultipleCallbacksPerRequest_AllCalledInOrder()
+        {
+            return LoopbackServerFactory.CreateClientAndServerAsync(async uri =>
+            {
+                using HttpMessageInvoker client = CreateHttpMessageInvoker();
+                using InstrumentRecorder<double> recorder = SetupInstrumentRecorder<double>(InstrumentNames.RequestDuration);
+                using HttpRequestMessage request = new(HttpMethod.Get, uri) { Version = UseVersion };
+
+                int lastCallback = -1;
+
+                HttpMetricsEnrichmentContext.AddCallback(request, ctx =>
+                {
+                    Assert.Equal(-1, lastCallback);
+                    lastCallback = 1;
+                    ctx.AddCustomTag("custom1", "foo");
+                });
+                HttpMetricsEnrichmentContext.AddCallback(request, ctx =>
+                {
+                    Assert.Equal(1, lastCallback);
+                    lastCallback = 2;
+                    ctx.AddCustomTag("custom2", "bar");
+                });
+                HttpMetricsEnrichmentContext.AddCallback(request, ctx =>
+                {
+                    Assert.Equal(2, lastCallback);
+                    ctx.AddCustomTag("custom3", "baz");
+                });
+
+                using HttpResponseMessage response = await SendAsync(client, request);
+
+                Measurement<double> m = Assert.Single(recorder.GetMeasurements());
+                VerifyRequestDuration(m, uri, UseVersion, 200);
+                Assert.Equal("foo", Assert.Single(m.Tags.ToArray(), t => t.Key == "custom1").Value);
+                Assert.Equal("bar", Assert.Single(m.Tags.ToArray(), t => t.Key == "custom2").Value);
+                Assert.Equal("baz", Assert.Single(m.Tags.ToArray(), t => t.Key == "custom3").Value);
+
             }, async server =>
             {
                 await server.AcceptConnectionSendResponseAndCloseAsync();
@@ -928,6 +971,96 @@ public class HttpMetricsTest_Http11_Async_HttpMessageInvoker : HttpMetricsTest_H
         public HttpMetricsTest_Http11_Async_HttpMessageInvoker(ITestOutputHelper output) : base(output)
         {
         }
+
+        [Fact]
+        public async Task RequestDuration_RequestReused_EnrichmentCallbacksAreCleared()
+        {
+            await LoopbackServerFactory.CreateClientAndServerAsync(async uri =>
+            {
+                using HttpMessageInvoker client = CreateHttpMessageInvoker();
+                using InstrumentRecorder<double> recorder = SetupInstrumentRecorder<double>(InstrumentNames.RequestDuration);
+
+                using HttpRequestMessage request = new(HttpMethod.Get, uri);
+
+                int firstCallbackCalls = 0;
+
+                HttpMetricsEnrichmentContext.AddCallback(request, ctx =>
+                {
+                    firstCallbackCalls++;
+                    ctx.AddCustomTag("key1", "foo");
+                });
+
+                (await SendAsync(client, request)).Dispose();
+                Assert.Equal(1, firstCallbackCalls);
+
+                Measurement<double> m = Assert.Single(recorder.GetMeasurements());
+                Assert.Equal("key1", Assert.Single(m.Tags.ToArray(), t => t.Value as string == "foo").Key);
+
+                HttpMetricsEnrichmentContext.AddCallback(request, static ctx =>
+                {
+                    ctx.AddCustomTag("key2", "foo");
+                });
+
+                (await SendAsync(client, request)).Dispose();
+                Assert.Equal(1, firstCallbackCalls);
+
+                Assert.Equal(2, recorder.GetMeasurements().Count);
+                m = recorder.GetMeasurements()[1];
+                Assert.Equal("key2", Assert.Single(m.Tags.ToArray(), t => t.Value as string == "foo").Key);
+            }, async server =>
+            {
+                await server.HandleRequestAsync();
+                await server.HandleRequestAsync();
+            });
+        }
+
+        [ConditionalFact(typeof(PlatformDetection), nameof(PlatformDetection.IsThreadingSupported))]
+        public async Task RequestDuration_ConcurrentRequestsSeeDifferentContexts()
+        {
+            await LoopbackServerFactory.CreateClientAndServerAsync(async uri =>
+            {
+                using HttpMessageInvoker client = CreateHttpMessageInvoker();
+                using var _ = SetupInstrumentRecorder<double>(InstrumentNames.RequestDuration);
+
+                using HttpRequestMessage request1 = new(HttpMethod.Get, uri);
+                using HttpRequestMessage request2 = new(HttpMethod.Get, uri);
+
+                HttpMetricsEnrichmentContext.AddCallback(request1, _ => { });
+                (await client.SendAsync(request1, CancellationToken.None)).Dispose();
+
+                HttpMetricsEnrichmentContext context1 = null;
+                HttpMetricsEnrichmentContext context2 = null;
+                CountdownEvent countdownEvent = new(2);
+
+                HttpMetricsEnrichmentContext.AddCallback(request1, ctx =>
+                {
+                    context1 = ctx;
+                    countdownEvent.Signal();
+                    Assert.True(countdownEvent.Wait(TestHelper.PassingTestTimeout));
+                });
+                HttpMetricsEnrichmentContext.AddCallback(request2, ctx =>
+                {
+                    context2 = ctx;
+                    countdownEvent.Signal();
+                    Assert.True(countdownEvent.Wait(TestHelper.PassingTestTimeout));
+                });
+
+                Task<HttpResponseMessage> task1 = Task.Run(() => client.SendAsync(request1, CancellationToken.None));
+                Task<HttpResponseMessage> task2 = Task.Run(() => client.SendAsync(request2, CancellationToken.None));
+
+                (await task1).Dispose();
+                (await task2).Dispose();
+
+                Assert.NotSame(context1, context2);
+            }, async server =>
+            {
+                await server.HandleRequestAsync();
+
+                await Task.WhenAll(
+                    server.HandleRequestAsync(),
+                    server.HandleRequestAsync());
+            }, options: new GenericLoopbackOptions { ListenBacklog = 2 });
+        }
     }
 
     [ConditionalClass(typeof(PlatformDetection), nameof(PlatformDetection.IsNotMobile))]