Refactor DiagnosticsHandler

dotnet#54437 but without the ActivitySource support

src/libraries/System.Net.Http/src/ILLink/ILLink.Substitutions.xml

@@ -1,7 +1,7 @@
 <linker>
   <assembly fullname="System.Net.Http">
     <type fullname="System.Net.Http.DiagnosticsHandler">
-      <method signature="System.Boolean IsGloballyEnabled()" body="stub" value="false" feature="System.Net.Http.EnableActivityPropagation" featurevalue="false" />
+      <method signature="System.Boolean get_IsGloballyEnabled()" body="stub" value="false" feature="System.Net.Http.EnableActivityPropagation" featurevalue="false" />
     </type>
     <type fullname="System.Net.Http.HttpClient">
       <method signature="System.Boolean IsNativeHandlerEnabled()" body="stub" value="true" feature="System.Net.Http.UseNativeHttpHandler" featurevalue="true" />

src/libraries/System.Net.Http/src/System.Net.Http.csproj

@@ -496,7 +496,6 @@
   <!-- Common -->
   <ItemGroup>
     <Compile Include="System\Net\Http\DiagnosticsHandler.cs" />
-    <Compile Include="System\Net\Http\DiagnosticsHandlerLoggingStrings.cs" />
     <Compile Include="System\Net\Http\GlobalHttpSettings.cs" />
     <Compile Include="System\Net\Http\SocketsHttpHandler\SocketsHttpPlaintextStreamFilterContext.cs" />
     <Compile Include="$(CommonPath)System\Threading\Tasks\TaskCompletionSourceWithCancellation.cs"

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

@@ -15,103 +15,101 @@ namespace System.Net.Http
     /// </summary>
     internal sealed class DiagnosticsHandler : DelegatingHandler
     {
-        private static readonly DiagnosticListener s_diagnosticListener =
-                new DiagnosticListener(DiagnosticsHandlerLoggingStrings.DiagnosticListenerName);
-
-        /// <summary>
-        /// DiagnosticHandler constructor
-        /// </summary>
-        /// <param name="innerHandler">Inner handler: Windows or Unix implementation of HttpMessageHandler.
-        /// Note that DiagnosticHandler is the latest in the pipeline </param>
-        public DiagnosticsHandler(HttpMessageHandler innerHandler) : base(innerHandler)
-        {
-        }
-
-        internal static bool IsEnabled()
-        {
-            // check if there is a parent Activity (and propagation is not suppressed)
-            // or if someone listens to HttpHandlerDiagnosticListener
-            return IsGloballyEnabled && (Activity.Current != null || s_diagnosticListener.IsEnabled());
-        }
-
-        internal static bool IsGloballyEnabled => GlobalHttpSettings.DiagnosticsHandler.EnableActivityPropagation;
+        private const string Namespace = "System.Net.Http";
+        private const string RequestWriteNameDeprecated = Namespace + ".Request";
+        private const string ResponseWriteNameDeprecated = Namespace + ".Response";
+        private const string ExceptionEventName = Namespace + ".Exception";
+        private const string ActivityName = Namespace + ".HttpRequestOut";
+        private const string ActivityStartName = ActivityName + ".Start";
+        private const string ActivityStopName = ActivityName + ".Stop";
 
-        // SendAsyncCore returns already completed ValueTask for when async: false is passed.
-        // Internally, it calls the synchronous Send method of the base class.
-        protected internal override HttpResponseMessage Send(HttpRequestMessage request, CancellationToken cancellationToken) =>
-            SendAsyncCore(request, async: false, cancellationToken).AsTask().GetAwaiter().GetResult();
+        private static readonly DiagnosticListener s_diagnosticListener = new("HttpHandlerDiagnosticListener");
+        //private static readonly ActivitySource s_activitySource = new(Namespace);
 
-        protected internal override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) =>
-            SendAsyncCore(request, async: true, cancellationToken).AsTask();
+        public static bool IsGloballyEnabled => GlobalHttpSettings.DiagnosticsHandler.EnableActivityPropagation;
 
-        private async ValueTask<HttpResponseMessage> SendAsyncCore(HttpRequestMessage request, bool async,
-            CancellationToken cancellationToken)
+        public DiagnosticsHandler(HttpMessageHandler innerHandler) : base(innerHandler)
         {
-            // HttpClientHandler is responsible to call static DiagnosticsHandler.IsEnabled() before forwarding request here.
-            // It will check if propagation is on (because parent Activity exists or there is a listener) or off (forcibly disabled)
-            // This code won't be called unless consumer unsubscribes from DiagnosticListener right after the check.
-            // So some requests happening right after subscription starts might not be instrumented. Similarly,
-            // when consumer unsubscribes, extra requests might be instrumented
+            Debug.Assert(IsGloballyEnabled);
+        }
 
-            if (request == null)
+        private static bool ShouldLogDiagnostics(HttpRequestMessage request, out Activity? activity)
+        {
+            if (request is null)
             {
                 throw new ArgumentNullException(nameof(request), SR.net_http_handler_norequest);
             }
 
-            Activity? activity = null;
-            DiagnosticListener diagnosticListener = s_diagnosticListener;
+            activity = null;
+
+            //if (s_activitySource.HasListeners())
+            //{
+            //    activity = s_activitySource.CreateActivity(ActivityName, ActivityKind.Client);
+            //}
 
-            // if there is no listener, but propagation is enabled (with previous IsEnabled() check)
-            // do not write any events just start/stop Activity and propagate Ids
-            if (!diagnosticListener.IsEnabled())
+            if (activity is null)
             {
-                activity = new Activity(DiagnosticsHandlerLoggingStrings.ActivityName);
-                activity.Start();
-                InjectHeaders(activity, request);
+                bool diagnosticListenerEnabled = s_diagnosticListener.IsEnabled();
 
-                try
+                if (Activity.Current is not null || (diagnosticListenerEnabled && s_diagnosticListener.IsEnabled(ActivityName, request)))
                 {
-                    return async ?
-                        await base.SendAsync(request, cancellationToken).ConfigureAwait(false) :
-                        base.Send(request, cancellationToken);
+                    // If a diagnostics listener is enabled for the Activity, always create one
+                    activity = new Activity(ActivityName);
                 }
-                finally
+                else
                 {
-                    activity.Stop();
+                    // There is no Activity, but we may still want to use the instrumented SendAsyncCore if diagnostic listeners are interested in other events
+                    return diagnosticListenerEnabled;
                 }
             }
 
-            Guid loggingRequestId = Guid.Empty;
+            activity.Start();
 
-            // There is a listener. Check if listener wants to be notified about HttpClient Activities
-            if (diagnosticListener.IsEnabled(DiagnosticsHandlerLoggingStrings.ActivityName, request))
+            if (s_diagnosticListener.IsEnabled(ActivityStartName))
             {
-                activity = new Activity(DiagnosticsHandlerLoggingStrings.ActivityName);
+                Write(ActivityStartName, new ActivityStartData(request));
+            }
 
-                // Only send start event to users who subscribed for it, but start activity anyway
-                if (diagnosticListener.IsEnabled(DiagnosticsHandlerLoggingStrings.ActivityStartName))
-                {
-                    StartActivity(diagnosticListener, activity, new ActivityStartData(request));
-                }
-                else
-                {
-                    activity.Start();
-                }
+            InjectHeaders(activity, request);
+
+            return true;
+        }
+
+        protected internal override HttpResponseMessage Send(HttpRequestMessage request, CancellationToken cancellationToken)
+        {
+            if (ShouldLogDiagnostics(request, out Activity? activity))
+            {
+                ValueTask<HttpResponseMessage> sendTask = SendAsyncCore(request, activity, async: false, cancellationToken);
+                return sendTask.IsCompleted ?
+                    sendTask.Result :
+                    sendTask.AsTask().GetAwaiter().GetResult();
             }
-            // try to write System.Net.Http.Request event (deprecated)
-            if (diagnosticListener.IsEnabled(DiagnosticsHandlerLoggingStrings.RequestWriteNameDeprecated))
+            else
             {
-                long timestamp = Stopwatch.GetTimestamp();
-                loggingRequestId = Guid.NewGuid();
-                Write(diagnosticListener, DiagnosticsHandlerLoggingStrings.RequestWriteNameDeprecated,
-                    new RequestData(request, loggingRequestId, timestamp));
+                return base.Send(request, cancellationToken);
+            }
+        }
+
+        protected internal override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
+        {
+            if (ShouldLogDiagnostics(request, out Activity? activity))
+            {
+                return SendAsyncCore(request, activity, async: true, cancellationToken).AsTask();
             }
+            else
+            {
+                return base.SendAsync(request, cancellationToken);
+            }
+        }
 
-            // If we are on at all, we propagate current activity information
-            Activity? currentActivity = Activity.Current;
-            if (currentActivity != null)
+        private async ValueTask<HttpResponseMessage> SendAsyncCore(HttpRequestMessage request, Activity? activity, bool async, CancellationToken cancellationToken)
+        {
+            Guid loggingRequestId = default;
+
+            if (s_diagnosticListener.IsEnabled(RequestWriteNameDeprecated))
             {
-                InjectHeaders(currentActivity, request);
+                loggingRequestId = Guid.NewGuid();
+                Write(RequestWriteNameDeprecated, new RequestData(request, loggingRequestId, Stopwatch.GetTimestamp()));
             }
 
             HttpResponseMessage? response = null;
@@ -126,52 +124,39 @@ await base.SendAsync(request, cancellationToken).ConfigureAwait(false) :
             catch (OperationCanceledException)
             {
                 taskStatus = TaskStatus.Canceled;
-
-                // we'll report task status in HttpRequestOut.Stop
                 throw;
             }
             catch (Exception ex)
             {
                 taskStatus = TaskStatus.Faulted;
 
-                if (diagnosticListener.IsEnabled(DiagnosticsHandlerLoggingStrings.ExceptionEventName))
+                if (s_diagnosticListener.IsEnabled(ExceptionEventName))
                 {
-                    // If request was initially instrumented, Activity.Current has all necessary context for logging
-                    // Request is passed to provide some context if instrumentation was disabled and to avoid
-                    // extensive Activity.Tags usage to tunnel request properties
-                    Write(diagnosticListener, DiagnosticsHandlerLoggingStrings.ExceptionEventName, new ExceptionData(ex, request));
+                    Write(ExceptionEventName, new ExceptionData(ex, request));
                 }
                 throw;
             }
             finally
             {
-                // always stop activity if it was started
-                if (activity != null)
+                if (activity is not null)
                 {
-                    StopActivity(diagnosticListener, activity, new ActivityStopData(
-                        response,
-                        // If request is failed or cancelled, there is no response, therefore no information about request;
-                        // pass the request in the payload, so consumers can have it in Stop for failed/canceled requests
-                        // and not retain all requests in Start
-                        request,
-                        taskStatus));
+                    activity.SetEndTime(DateTime.UtcNow);
+
+                    if (s_diagnosticListener.IsEnabled(ActivityStopName))
+                    {
+                        Write(ActivityStopName, new ActivityStopData(response, request, taskStatus));
+                    }
+
+                    activity.Stop();
                 }
-                // Try to write System.Net.Http.Response event (deprecated)
-                if (diagnosticListener.IsEnabled(DiagnosticsHandlerLoggingStrings.ResponseWriteNameDeprecated))
+
+                if (s_diagnosticListener.IsEnabled(ResponseWriteNameDeprecated))
                 {
-                    long timestamp = Stopwatch.GetTimestamp();
-                    Write(diagnosticListener, DiagnosticsHandlerLoggingStrings.ResponseWriteNameDeprecated,
-                        new ResponseData(
-                            response,
-                            loggingRequestId,
-                            timestamp,
-                            taskStatus));
+                    Write(ResponseWriteNameDeprecated, new ResponseData(response, loggingRequestId, Stopwatch.GetTimestamp(), taskStatus));
                 }
             }
         }
 
-        #region private
-
         private sealed class ActivityStartData
         {
             // matches the properties selected in https://github.com/dotnet/diagnostics/blob/ffd0254da3bcc47847b1183fa5453c0877020abd/src/Microsoft.Diagnostics.Monitoring.EventPipe/Configuration/HttpRequestSourceConfiguration.cs#L36-L40
@@ -271,22 +256,27 @@ internal ResponseData(HttpResponseMessage? response, Guid loggingRequestId, long
 
         private static void InjectHeaders(Activity currentActivity, HttpRequestMessage request)
         {
+            const string TraceParentHeaderName = "traceparent";
+            const string TraceStateHeaderName = "tracestate";
+            const string RequestIdHeaderName = "Request-Id";
+            const string CorrelationContextHeaderName = "Correlation-Context";
+
             if (currentActivity.IdFormat == ActivityIdFormat.W3C)
             {
-                if (!request.Headers.Contains(DiagnosticsHandlerLoggingStrings.TraceParentHeaderName))
+                if (!request.Headers.Contains(TraceParentHeaderName))
                 {
-                    request.Headers.TryAddWithoutValidation(DiagnosticsHandlerLoggingStrings.TraceParentHeaderName, currentActivity.Id);
+                    request.Headers.TryAddWithoutValidation(TraceParentHeaderName, currentActivity.Id);
                     if (currentActivity.TraceStateString != null)
                     {
-                        request.Headers.TryAddWithoutValidation(DiagnosticsHandlerLoggingStrings.TraceStateHeaderName, currentActivity.TraceStateString);
+                        request.Headers.TryAddWithoutValidation(TraceStateHeaderName, currentActivity.TraceStateString);
                     }
                 }
             }
             else
             {
-                if (!request.Headers.Contains(DiagnosticsHandlerLoggingStrings.RequestIdHeaderName))
+                if (!request.Headers.Contains(RequestIdHeaderName))
                 {
-                    request.Headers.TryAddWithoutValidation(DiagnosticsHandlerLoggingStrings.RequestIdHeaderName, currentActivity.Id);
+                    request.Headers.TryAddWithoutValidation(RequestIdHeaderName, currentActivity.Id);
                 }
             }
 
@@ -302,41 +292,16 @@ private static void InjectHeaders(Activity currentActivity, HttpRequestMessage r
                         baggage.Add(new NameValueHeaderValue(WebUtility.UrlEncode(item.Key), WebUtility.UrlEncode(item.Value)).ToString());
                     }
                     while (e.MoveNext());
-                    request.Headers.TryAddWithoutValidation(DiagnosticsHandlerLoggingStrings.CorrelationContextHeaderName, baggage);
+                    request.Headers.TryAddWithoutValidation(CorrelationContextHeaderName, baggage);
                 }
             }
         }
 
         [UnconditionalSuppressMessage("ReflectionAnalysis", "IL2026:UnrecognizedReflectionPattern",
             Justification = "The values being passed into Write have the commonly used properties being preserved with DynamicDependency.")]
-        private static void Write<[DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicProperties)] T>(
-            DiagnosticSource diagnosticSource,
-            string name,
-            T value)
-        {
-            diagnosticSource.Write(name, value);
-        }
-
-        [UnconditionalSuppressMessage("ReflectionAnalysis", "IL2026:UnrecognizedReflectionPattern",
-            Justification = "The args being passed into StartActivity have the commonly used properties being preserved with DynamicDependency.")]
-        private static Activity StartActivity<[DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicProperties)] T>(
-            DiagnosticSource diagnosticSource,
-            Activity activity,
-            T? args)
+        private static void Write<[DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicProperties)] T>(string name, T value)
         {
-            return diagnosticSource.StartActivity(activity, args);
+            s_diagnosticListener.Write(name, value);
         }
-
-        [UnconditionalSuppressMessage("ReflectionAnalysis", "IL2026:UnrecognizedReflectionPattern",
-            Justification = "The args being passed into StopActivity have the commonly used properties being preserved with DynamicDependency.")]
-        private static void StopActivity<[DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicProperties)] T>(
-            DiagnosticSource diagnosticSource,
-            Activity activity,
-            T? args)
-        {
-            diagnosticSource.StopActivity(activity, args);
-        }
-
-        #endregion
     }
 }