Ensure Http Telemetry correctness

src/libraries/System.Net.Http/src/System/Net/Http/HttpRequestMessage.cs

@@ -7,16 +7,18 @@
 using System.Text;
 using System.Threading;
 using System.Collections.Generic;
+using System.Diagnostics;
 
 namespace System.Net.Http
 {
     public class HttpRequestMessage : IDisposable
     {
         private const int MessageNotYetSent = 0;
         private const int MessageAlreadySent = 1;
+        private const int MessageShouldEmitTelemetry = 2;
 
         // Track whether the message has been sent.
-        // The message shouldn't be sent again if this field is equal to MessageAlreadySent.
+        // The message should only be sent if this field is equal to MessageNotYetSent.
         private int _sendStatus = MessageNotYetSent;
 
         private HttpMethod _method;
@@ -26,6 +28,7 @@ public class HttpRequestMessage : IDisposable
         private HttpContent? _content;
         private bool _disposed;
         private IDictionary<string, object?>? _properties;
+        private HttpRequestMessageFinalizer? _finalizer;
 
         public Version Version
         {
@@ -197,7 +200,40 @@ private void InitializeValues(HttpMethod method, Uri? requestUri)
 
         internal bool MarkAsSent()
         {
-            return Interlocked.Exchange(ref _sendStatus, MessageAlreadySent) == MessageNotYetSent;
+            return Interlocked.CompareExchange(ref _sendStatus, MessageAlreadySent, MessageNotYetSent) == MessageNotYetSent;
+        }
+
+        internal void MarkAsTrackedByTelemetry()
+        {
+            if (_finalizer is null)
+            {
+                _finalizer = new HttpRequestMessageFinalizer();
+            }
+            else
+            {
+                GC.ReRegisterForFinalize(_finalizer);
+            }
+
+            Debug.Assert(_sendStatus != MessageShouldEmitTelemetry);
+            _sendStatus = MessageShouldEmitTelemetry;
+        }
+
+        internal void OnAborted() => OnStopped(aborted: true);
+
+        internal void OnStopped(bool aborted = false)
+        {
+            if (_sendStatus == MessageShouldEmitTelemetry && Interlocked.Exchange(ref _sendStatus, MessageAlreadySent) == MessageShouldEmitTelemetry)
+            {
+                if (aborted)
+                {
+                    HttpTelemetry.Log.RequestAborted();
+                }
+
+                HttpTelemetry.Log.RequestStop();
+
+                Debug.Assert(_finalizer != null);
+                GC.SuppressFinalize(_finalizer);
+            }
         }
 
         #region IDisposable Members
@@ -214,6 +250,8 @@ protected virtual void Dispose(bool disposing)
                     _content.Dispose();
                 }
             }
+
+            OnStopped();
         }
 
         public void Dispose()
@@ -231,5 +269,16 @@ private void CheckDisposed()
                 throw new ObjectDisposedException(this.GetType().ToString());
             }
         }
+
+        /// <summary>
+        /// This class will only be allocated if Telemetry is enabled.
+        /// We can't use HttpRequestMessage's own finalizer because it is not sealed.
+        /// The finalizer will run iff OnStopped/OnAborted/Dispose were never called.
+        /// This way we ensure that RequestStop is always called if we call RequestStart.
+        /// </summary>
+        private sealed class HttpRequestMessageFinalizer
+        {
+            ~HttpRequestMessageFinalizer() => HttpTelemetry.Log.RequestStop();
+        }
     }
 }

src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/ChunkedEncodingReadStream.cs

@@ -57,7 +57,7 @@ public override int Read(Span<byte> buffer)
                     if (_connection == null)
                     {
                         // Fully consumed the response in ReadChunksFromConnectionBuffer.
-                        if (HttpTelemetry.IsEnabled) LogRequestStop();
+                        LogRequestStop();
                         return 0;
                     }
 
@@ -363,7 +363,7 @@ private ReadOnlyMemory<byte> ReadChunkFromConnectionBuffer(int maxBytesToRead, C
                                     cancellationRegistration.Dispose();
                                     CancellationHelper.ThrowIfCancellationRequested(cancellationRegistration.Token);
 
-                                    if (HttpTelemetry.IsEnabled) LogRequestStop();
+                                    LogRequestStop();
                                     _state = ParsingState.Done;
                                     _connection.CompleteResponse();
                                     _connection = null;

src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/ConnectionCloseReadStream.cs

@@ -29,7 +29,7 @@ public override int Read(Span<byte> buffer)
                 if (bytesRead == 0)
                 {
                     // We cannot reuse this connection, so close it.
-                    if (HttpTelemetry.IsEnabled) LogRequestStop();
+                    LogRequestStop();
                     _connection = null;
                     connection.Dispose();
                 }
@@ -82,7 +82,7 @@ public override async ValueTask<int> ReadAsync(Memory<byte> buffer, Cancellation
                     CancellationHelper.ThrowIfCancellationRequested(cancellationToken);
 
                     // We cannot reuse this connection, so close it.
-                    if (HttpTelemetry.IsEnabled) LogRequestStop();
+                    LogRequestStop();
                     _connection = null;
                     connection.Dispose();
                 }
@@ -144,7 +144,7 @@ private async Task CompleteCopyToAsync(Task copyTask, HttpConnection connection,
             private void Finish(HttpConnection connection)
             {
                 // We cannot reuse this connection, so close it.
-                if (HttpTelemetry.IsEnabled) LogRequestStop();
+                LogRequestStop();
                 _connection = null;
                 connection.Dispose();
             }

src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/ContentLengthReadStream.cs

@@ -48,7 +48,7 @@ public override int Read(Span<byte> buffer)
                 if (_contentBytesRemaining == 0)
                 {
                     // End of response body
-                    if (HttpTelemetry.IsEnabled) LogRequestStop();
+                    LogRequestStop();
                     _connection.CompleteResponse();
                     _connection = null;
                 }
@@ -111,7 +111,7 @@ public override async ValueTask<int> ReadAsync(Memory<byte> buffer, Cancellation
                 if (_contentBytesRemaining == 0)
                 {
                     // End of response body
-                    if (HttpTelemetry.IsEnabled) LogRequestStop();
+                    LogRequestStop();
                     _connection.CompleteResponse();
                     _connection = null;
                 }
@@ -166,7 +166,7 @@ private async Task CompleteCopyToAsync(Task copyTask, CancellationToken cancella
 
             private void Finish()
             {
-                if (HttpTelemetry.IsEnabled) LogRequestStop();
+                LogRequestStop();
                 _contentBytesRemaining = 0;
                 _connection!.CompleteResponse();
                 _connection = null;

src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/Http2Stream.cs

@@ -350,7 +350,7 @@ private void Complete()
                     _creditWaiter = null;
                 }
 
-                if (HttpTelemetry.IsEnabled) HttpTelemetry.Log.RequestStop();
+                _request.OnStopped();
             }
 
             private void Cancel()
@@ -386,7 +386,7 @@ private void Cancel()
                     _waitSource.SetResult(true);
                 }
 
-                if (HttpTelemetry.IsEnabled) HttpTelemetry.Log.RequestAborted();
+                _request.OnAborted();
             }
 
             // Returns whether the waiter should be signalled or not.

src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/HttpConnection.cs

@@ -49,7 +49,7 @@ internal partial class HttpConnection : HttpConnectionBase, IDisposable
         private readonly TransportContext? _transportContext;
         private readonly WeakReference<HttpConnection> _weakThisRef;
 
-        private HttpRequestMessage? _currentRequest;
+        internal HttpRequestMessage? _currentRequest;
         private readonly byte[] _writeBuffer;
         private int _writeOffset;
         private int _allowedReadLineBytes;
@@ -622,7 +622,7 @@ public async Task<HttpResponseMessage> SendAsyncCore(HttpRequestMessage request,
                 Stream responseStream;
                 if (ReferenceEquals(normalizedMethod, HttpMethod.Head) || response.StatusCode == HttpStatusCode.NoContent || response.StatusCode == HttpStatusCode.NotModified)
                 {
-                    if (HttpTelemetry.IsEnabled) HttpTelemetry.Log.RequestStop();
+                    _currentRequest.OnStopped();
                     responseStream = EmptyReadStream.Instance;
                     CompleteResponse();
                 }
@@ -645,7 +645,7 @@ public async Task<HttpResponseMessage> SendAsyncCore(HttpRequestMessage request,
                     long contentLength = response.Content.Headers.ContentLength.GetValueOrDefault();
                     if (contentLength <= 0)
                     {
-                        if (HttpTelemetry.IsEnabled) HttpTelemetry.Log.RequestStop();
+                        _currentRequest.OnStopped();
                         responseStream = EmptyReadStream.Instance;
                         CompleteResponse();
                     }

src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/HttpConnectionPoolManager.cs

@@ -355,20 +355,22 @@ private async ValueTask<HttpResponseMessage> SendAsyncWithLogging(HttpRequestMes
                 request.RequestUri.PathAndQuery,
                 request.Version.Major,
                 request.Version.Minor);
+
+            request.MarkAsTrackedByTelemetry();
+
             try
             {
                 return await SendAsyncHelper(request, async, doRequestAuth, cancellationToken).ConfigureAwait(false);
             }
-            catch (Exception e) when (LogException(e))
+            catch when (LogException(request))
             {
                 // This code should never run.
                 throw;
             }
 
-            static bool LogException(Exception e)
+            static bool LogException(HttpRequestMessage request)
             {
-                HttpTelemetry.Log.RequestAborted();
-                HttpTelemetry.Log.RequestStop();
+                request.OnAborted();
 
                 // Returning false means the catch handler isn't run.
                 // So the exception isn't considered to be caught so it will now propagate up the stack.

src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/HttpContentStream.cs

@@ -2,17 +2,12 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-using System.Threading;
-
 namespace System.Net.Http
 {
     internal abstract class HttpContentStream : HttpBaseStream
     {
         protected HttpConnection? _connection;
 
-        // Makes sure we don't call HttpTelemetry events more than once.
-        private int _requestStopCalled; // 0==no, 1==yes
-
         public HttpContentStream(HttpConnection connection)
         {
             _connection = connection;
@@ -48,10 +43,7 @@ protected HttpConnection GetConnectionOrThrow()
 
         protected void LogRequestStop()
         {
-            if (Interlocked.Exchange(ref _requestStopCalled, 1) == 0)
-            {
-                HttpTelemetry.Log.RequestStop();
-            }
+            _connection?._currentRequest?.OnStopped();
         }
 
         private HttpConnection ThrowObjectDisposedException() => throw new ObjectDisposedException(GetType().Name);

src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/RawConnectionStream.cs

@@ -34,7 +34,7 @@ public override int Read(Span<byte> buffer)
                 if (bytesRead == 0)
                 {
                     // We cannot reuse this connection, so close it.
-                    if (HttpTelemetry.IsEnabled) LogRequestStop();
+                    LogRequestStop();
                     _connection = null;
                     connection.Dispose();
                 }
@@ -82,7 +82,7 @@ public override async ValueTask<int> ReadAsync(Memory<byte> buffer, Cancellation
                     CancellationHelper.ThrowIfCancellationRequested(cancellationToken);
 
                     // We cannot reuse this connection, so close it.
-                    if (HttpTelemetry.IsEnabled) LogRequestStop();
+                    LogRequestStop();
                     _connection = null;
                     connection.Dispose();
                 }
@@ -144,7 +144,7 @@ private async Task CompleteCopyToAsync(Task copyTask, HttpConnection connection,
             private void Finish(HttpConnection connection)
             {
                 // We cannot reuse this connection, so close it.
-                if (HttpTelemetry.IsEnabled) LogRequestStop();
+                LogRequestStop();
                 connection.Dispose();
                 _connection = null;
             }

src/libraries/System.Net.Http/tests/UnitTests/Fakes/HttpTelemetry.cs

@@ -0,0 +1,15 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+namespace System.Net.Http
+{
+    internal class HttpTelemetry
+    {
+        public static HttpTelemetry Log => new HttpTelemetry();
+
+        public void RequestStop() { }
+
+        public void RequestAborted() { }
+    }
+}

src/libraries/System.Net.Http/tests/UnitTests/System.Net.Http.Unit.Tests.csproj

@@ -1,4 +1,4 @@
-<Project Sdk="Microsoft.NET.Sdk">
+<Project Sdk="Microsoft.NET.Sdk">
   <PropertyGroup>
     <StringResourcesPath>../../src/Resources/Strings.resx</StringResourcesPath>
     <AllowUnsafeBlocks>true</AllowUnsafeBlocks>
@@ -242,6 +242,7 @@
              Link="ProductionCode\System\Net\Http\HttpHandlerDefaults.cs" />
     <Compile Include="DigestAuthenticationTests.cs" />
     <Compile Include="Fakes\HttpClientHandler.cs" />
+    <Compile Include="Fakes\HttpTelemetry.cs" />
     <Compile Include="Fakes\MacProxy.cs" Condition=" ('$(TargetsOSX)' == 'true' or '$(TargetsiOS)' == 'true' or '$(TargetstvOS)' == 'true') and '$(TargetFramework)' == '$(NetCoreAppCurrent)'" />
     <Compile Include="Headers\AltSvcHeaderParserTest.cs" />
     <Compile Include="Headers\AuthenticationHeaderValueTest.cs" />