Ensure Http Telemetry correctness

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

@@ -6,16 +6,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 MessageAlreadySent_StopNotYetCalled = 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;
@@ -183,7 +185,31 @@ 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()
+        {
+            Debug.Assert(_sendStatus != MessageAlreadySent_StopNotYetCalled);
+            _sendStatus = MessageAlreadySent_StopNotYetCalled;
+        }
+
+        internal void OnAborted() => OnStopped(aborted: true);
+
+        internal void OnStopped(bool aborted = false)
+        {
+            if (HttpTelemetry.Log.IsEnabled())
+            {
+                if (Interlocked.Exchange(ref _sendStatus, MessageAlreadySent) == MessageAlreadySent_StopNotYetCalled)
+                {
+                    if (aborted)
+                    {
+                        HttpTelemetry.Log.RequestAborted();
+                    }
+
+                    HttpTelemetry.Log.RequestStop();
+                }
+            }
         }
 
         #region IDisposable Members
@@ -200,6 +226,8 @@ protected virtual void Dispose(bool disposing)
                     _content.Dispose();
                 }
             }
+
+            OnStopped();
         }
 
         public void Dispose()

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

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

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

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

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

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

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

@@ -351,7 +351,7 @@ private void Complete()
                     _creditWaiter = null;
                 }
 
-                if (HttpTelemetry.Log.IsEnabled()) HttpTelemetry.Log.RequestStop();
+                if (HttpTelemetry.Log.IsEnabled()) _request.OnStopped();
             }
 
             private void Cancel()
@@ -387,7 +387,7 @@ private void Cancel()
                     _waitSource.SetResult(true);
                 }
 
-                if (HttpTelemetry.Log.IsEnabled()) HttpTelemetry.Log.RequestAborted();
+                if (HttpTelemetry.Log.IsEnabled()) _request.OnAborted();
             }
 
             // Returns whether the waiter should be signalled or not.
@@ -1147,6 +1147,10 @@ private void CloseResponseBody()
                 {
                     Cancel();
                 }
+                else
+                {
+                    _request.OnStopped();
+                }
 
                 _responseBuffer.Dispose();
             }

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

@@ -90,6 +90,8 @@ public HttpConnection(
             if (NetEventSource.Log.IsEnabled()) TraceConnection(_stream);
         }
 
+        ~HttpConnection() => Dispose(disposing: false);
+
         public void Dispose() => Dispose(disposing: true);
 
         protected void Dispose(bool disposing)
@@ -100,6 +102,9 @@ protected void Dispose(bool disposing)
             {
                 if (NetEventSource.Log.IsEnabled()) Trace("Connection closing.");
                 _pool.DecrementConnectionCount();
+
+                if (HttpTelemetry.Log.IsEnabled()) _currentRequest?.OnAborted();
+
                 if (disposing)
                 {
                     GC.SuppressFinalize(this);
@@ -630,7 +635,6 @@ public async Task<HttpResponseMessage> SendAsyncCore(HttpRequestMessage request,
                 Stream responseStream;
                 if (ReferenceEquals(normalizedMethod, HttpMethod.Head) || response.StatusCode == HttpStatusCode.NoContent || response.StatusCode == HttpStatusCode.NotModified)
                 {
-                    if (HttpTelemetry.Log.IsEnabled()) HttpTelemetry.Log.RequestStop();
                     responseStream = EmptyReadStream.Instance;
                     CompleteResponse();
                 }
@@ -653,7 +657,6 @@ public async Task<HttpResponseMessage> SendAsyncCore(HttpRequestMessage request,
                     long contentLength = response.Content.Headers.ContentLength.GetValueOrDefault();
                     if (contentLength <= 0)
                     {
-                        if (HttpTelemetry.Log.IsEnabled()) HttpTelemetry.Log.RequestStop();
                         responseStream = EmptyReadStream.Instance;
                         CompleteResponse();
                     }
@@ -1850,6 +1853,8 @@ private void CompleteResponse()
             Debug.Assert(_currentRequest != null, "Expected the connection to be associated with a request.");
             Debug.Assert(_writeOffset == 0, "Everything in write buffer should have been flushed.");
 
+            if (HttpTelemetry.Log.IsEnabled()) _currentRequest.OnStopped();
+
             // Disassociate the connection from a request.
             _currentRequest = null;
 
@@ -1963,13 +1968,4 @@ public sealed override void Trace(string message, [CallerMemberName] string? mem
                 memberName,                          // method name
                 message);                            // message
     }
-
-    internal sealed class HttpConnectionWithFinalizer : HttpConnection
-    {
-        public HttpConnectionWithFinalizer(HttpConnectionPool pool, Connection connection, TransportContext? transportContext) : base(pool, connection, transportContext) { }
-
-        // This class is separated from HttpConnection so we only pay the price of having a finalizer
-        // when it's actually needed, e.g. when MaxConnectionsPerServer is enabled.
-        ~HttpConnectionWithFinalizer() => Dispose(disposing: false);
-    }
 }

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

@@ -1238,9 +1238,7 @@ private ValueTask<Connection> ConnectToTcpHostAsync(string host, int port, HttpR
 
         private HttpConnection ConstructHttp11Connection(Connection connection, TransportContext? transportContext)
         {
-            return _maxConnections == int.MaxValue ?
-                new HttpConnection(this, connection, transportContext) :
-                new HttpConnectionWithFinalizer(this, connection, transportContext); // finalizer needed to signal the pool when a connection is dropped
+            return new HttpConnection(this, connection, transportContext);
         }
 
         // Returns the established stream or an HttpResponseMessage from the proxy indicating failure.

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

@@ -354,20 +354,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

@@ -1,17 +1,12 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
-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;
@@ -45,14 +40,6 @@ protected HttpConnection GetConnectionOrThrow()
                 ThrowObjectDisposedException();
         }
 
-        protected void LogRequestStop()
-        {
-            if (Interlocked.Exchange(ref _requestStopCalled, 1) == 0)
-            {
-                HttpTelemetry.Log.RequestStop();
-            }
-        }
-
         private HttpConnection ThrowObjectDisposedException() => throw new ObjectDisposedException(GetType().Name);
     }
 }

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

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

src/libraries/System.Net.Http/tests/FunctionalTests/System.Net.Http.Functional.Tests.csproj

@@ -224,6 +224,7 @@
     <Compile Include="$(CommonTestPath)System\Net\Http\SchSendAuxRecordHttpTest.cs"
              Link="Common\System\Net\Http\SchSendAuxRecordHttpTest.cs" />
     <Compile Include="SyncHttpHandlerTest.cs" />
+    <Compile Include="TelemetryTest.cs" />
     <Compile Include="StreamContentTest.cs" />
     <Compile Include="StringContentTest.cs" />
     <Compile Include="$(CommonTestPath)System\Net\Http\SyncBlockingContent.cs"