Wire up DiagnosticsHandler (#456)

* Wire up DiagnosticsHandler * Remove deprecated events * Fix broken tests * Add a toggle to the HttpClientOptions * Add short-circuit if not enabled * Remove existing header if we already have a parent * Fixes after rebasing * Greatly simplify DiagnosticsHandler * Further simplification * Yet even simpler * Rename to ActivityPropagationHandler Fix comments * Update src/ReverseProxy/Telemetry/ActivityPropagationHandler.cs Co-authored-by: Chris Ross <[email protected]> * Make PropagateActivityContext nullable * Fix TODO in test * Fix PR comments * Add ActivityPropagationHandler test * Add missing using * Made handler public so that it can be re-used * Moved test to appropriate folder * Fix namespace Co-authored-by: Chris Ross <[email protected]> Co-authored-by: MihaZupan <[email protected]>
microsoft · Nov 12, 2020 · 31ca104 · 31ca104
1 parent ac86b14
commit 31ca104
Show file tree

Hide file tree

Showing 11 changed files with 262 additions and 51 deletions.
diff --git a/src/ReverseProxy/Abstractions/ClusterDiscovery/Contract/ProxyHttpClientOptions.cs b/src/ReverseProxy/Abstractions/ClusterDiscovery/Contract/ProxyHttpClientOptions.cs
@@ -17,6 +17,8 @@ public sealed class ProxyHttpClientOptions
 
         public int? MaxConnectionsPerServer { get; set; }
 
+        public bool? PropagateActivityContext { get; set; }
+
         // TODO: Add this property once we have migrated to SDK version that supports it.
         //public bool? EnableMultipleHttp2Connections { get; set; }
 
@@ -28,7 +30,8 @@ internal ProxyHttpClientOptions DeepClone()
                 DangerousAcceptAnyServerCertificate = DangerousAcceptAnyServerCertificate,
                 // TODO: Clone certificate?
                 ClientCertificate = ClientCertificate,
-                MaxConnectionsPerServer = MaxConnectionsPerServer
+                MaxConnectionsPerServer = MaxConnectionsPerServer,
+                PropagateActivityContext = PropagateActivityContext,
             };
         }
 
@@ -47,7 +50,8 @@ internal static bool Equals(ProxyHttpClientOptions options1, ProxyHttpClientOpti
             return options1.SslProtocols == options2.SslProtocols
                 && Equals(options1.ClientCertificate, options2.ClientCertificate)
                 && options1.DangerousAcceptAnyServerCertificate == options2.DangerousAcceptAnyServerCertificate
-                && options1.MaxConnectionsPerServer == options2.MaxConnectionsPerServer;
+                && options1.MaxConnectionsPerServer == options2.MaxConnectionsPerServer
+                && options1.PropagateActivityContext == options2.PropagateActivityContext;
         }
 
         private static bool Equals(X509Certificate2 certificate1, X509Certificate2 certificate2)

diff --git a/src/ReverseProxy/Configuration/ConfigurationConfigProvider.cs b/src/ReverseProxy/Configuration/ConfigurationConfigProvider.cs
@@ -380,7 +380,8 @@ private ProxyHttpClientOptions CreateProxyHttpClientOptions(IConfigurationSectio
                 SslProtocols = sslProtocols,
                 DangerousAcceptAnyServerCertificate = section.ReadBool(nameof(ProxyHttpClientOptions.DangerousAcceptAnyServerCertificate)) ?? true,
                 ClientCertificate = clientCertificate,
-                MaxConnectionsPerServer = section.ReadInt32(nameof(ProxyHttpClientOptions.MaxConnectionsPerServer))
+                MaxConnectionsPerServer = section.ReadInt32(nameof(ProxyHttpClientOptions.MaxConnectionsPerServer)),
+                PropagateActivityContext = section.ReadBool(nameof(ProxyHttpClientOptions.PropagateActivityContext))
             };
         }
 

diff --git a/src/ReverseProxy/Service/Management/ProxyConfigManager.cs b/src/ReverseProxy/Service/Management/ProxyConfigManager.cs
@@ -512,7 +512,8 @@ private ClusterProxyHttpClientOptions ConvertProxyHttpClientOptions(ProxyHttpCli
                 httpClientOptions.SslProtocols,
                 httpClientOptions.DangerousAcceptAnyServerCertificate,
                 httpClientOptions.ClientCertificate,
-                httpClientOptions.MaxConnectionsPerServer);
+                httpClientOptions.MaxConnectionsPerServer,
+                httpClientOptions.PropagateActivityContext);
         }
 
         public void Dispose()

diff --git a/src/ReverseProxy/Service/Proxy/Infrastructure/ProxyHttpClientFactory.cs b/src/ReverseProxy/Service/Proxy/Infrastructure/ProxyHttpClientFactory.cs
@@ -6,6 +6,7 @@
 using System.Net.Http;
 using System.Security.Cryptography.X509Certificates;
 using Microsoft.Extensions.Logging;
+using Microsoft.ReverseProxy.Telemetry;
 
 namespace Microsoft.ReverseProxy.Service.Proxy.Infrastructure
 {
@@ -65,6 +66,12 @@ public HttpMessageInvoker CreateClient(ProxyHttpClientContext context)
             }
 
             Log.ProxyClientCreated(_logger, context.ClusterId);
+
+            if (newClientOptions.PropagateActivityContext.GetValueOrDefault(true))
+            {
+                return new HttpMessageInvoker(new ActivityPropagationHandler(handler), disposeHandler: true);
+            }
+
             return new HttpMessageInvoker(handler, disposeHandler: true);
         }
 

diff --git a/src/ReverseProxy/Service/RuntimeModel/ClusterProxyHttpClientOptions.cs b/src/ReverseProxy/Service/RuntimeModel/ClusterProxyHttpClientOptions.cs
@@ -14,12 +14,14 @@ public ClusterProxyHttpClientOptions(
             SslProtocols? sslProtocols,
             bool acceptAnyServerCertificate,
             X509Certificate2 clientCertificate,
-            int? maxConnectionsPerServer)
+            int? maxConnectionsPerServer,
+            bool? propagateActivityContext)
         {
             SslProtocols = sslProtocols;
             DangerousAcceptAnyServerCertificate = acceptAnyServerCertificate;
             ClientCertificate = clientCertificate;
             MaxConnectionsPerServer = maxConnectionsPerServer;
+            PropagateActivityContext = propagateActivityContext;
         }
 
         public SslProtocols? SslProtocols { get; }
@@ -30,6 +32,8 @@ public ClusterProxyHttpClientOptions(
 
         public int? MaxConnectionsPerServer { get; }
 
+        public bool? PropagateActivityContext { get; }
+
         // TODO: Add this property once we have migrated to SDK version that supports it.
         //public bool? EnableMultipleHttp2Connections { get; }
 
@@ -43,7 +47,8 @@ public bool Equals(ClusterProxyHttpClientOptions other)
             return SslProtocols == other.SslProtocols &&
                    DangerousAcceptAnyServerCertificate == other.DangerousAcceptAnyServerCertificate &&
                    EqualityComparer<X509Certificate2>.Default.Equals(ClientCertificate, other.ClientCertificate) &&
-                   MaxConnectionsPerServer == other.MaxConnectionsPerServer;
+                   MaxConnectionsPerServer == other.MaxConnectionsPerServer &&
+                   PropagateActivityContext == other.PropagateActivityContext;
         }
 
         public override int GetHashCode()

diff --git a/src/ReverseProxy/Telemetry/ActivityPropagationHandler.cs b/src/ReverseProxy/Telemetry/ActivityPropagationHandler.cs
@@ -0,0 +1,87 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using System;
+using System.Collections.Generic;
+using System.Diagnostics;
+using System.Net.Http;
+using System.Net.Http.Headers;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.ReverseProxy.Telemetry
+{
+    /// <summary>
+    /// ActivityPropagationHandler propagates the current Activity to the downstream service
+    /// </summary>
+    public sealed class ActivityPropagationHandler : DelegatingHandler
+    {
+        private const string RequestIdHeaderName = "Request-Id";
+        private const string CorrelationContextHeaderName = "Correlation-Context";
+
+        private const string TraceParentHeaderName = "traceparent";
+        private const string TraceStateHeaderName = "tracestate";
+
+        public ActivityPropagationHandler(HttpMessageHandler innerHandler) : base(innerHandler)
+        {
+        }
+
+        protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request,
+            CancellationToken cancellationToken)
+        {
+            // This handler is conditionally inserted by the ProxyHttpClientFactory based on the configuration
+            // If inserted it will insert the necessary headers to propagate the current activity context to
+            // the downstream service, if there is a current activity
+
+            if (request == null)
+            {
+                throw new ArgumentNullException(nameof(request));
+            }
+
+            // If we are on at all, we propagate current activity information
+            var currentActivity = Activity.Current;
+            if (currentActivity != null)
+            {
+                InjectHeaders(currentActivity, request);
+            }
+
+            return base.SendAsync(request, cancellationToken);
+        }
+
+        private void InjectHeaders(Activity currentActivity, HttpRequestMessage request)
+        {
+            if (currentActivity.IdFormat == ActivityIdFormat.W3C)
+            {
+                request.Headers.Remove(TraceParentHeaderName);
+                request.Headers.Remove(TraceStateHeaderName);
+
+                request.Headers.TryAddWithoutValidation(TraceParentHeaderName, currentActivity.Id);
+                if (currentActivity.TraceStateString != null)
+                {
+                    request.Headers.TryAddWithoutValidation(TraceStateHeaderName, currentActivity.TraceStateString);
+                }
+            }
+            else
+            {
+                request.Headers.Remove(RequestIdHeaderName);
+                request.Headers.TryAddWithoutValidation(RequestIdHeaderName, currentActivity.Id);
+            }
+
+            // we expect baggage to be empty or contain a few items
+            using (var e = currentActivity.Baggage.GetEnumerator())
+            {
+                if (e.MoveNext())
+                {
+                    var baggage = new List<string>();
+                    do
+                    {
+                        var item = e.Current;
+                        baggage.Add(new NameValueHeaderValue(Uri.EscapeDataString(item.Key), Uri.EscapeDataString(item.Value)).ToString());
+                    }
+                    while (e.MoveNext());
+                    request.Headers.TryAddWithoutValidation(CorrelationContextHeaderName, baggage);
+                }
+            }
+        }
+    }
+}
diff --git a/test/ReverseProxy.Tests/Common/MockHttpHandler.cs b/test/ReverseProxy.Tests/Common/MockHttpHandler.cs
@@ -0,0 +1,31 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Net.Http;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.ReverseProxy.Common.Tests
+{
+    internal class MockHttpHandler : HttpMessageHandler
+    {
+        private readonly Func<HttpRequestMessage, CancellationToken, Task<HttpResponseMessage>> _func;
+
+        public MockHttpHandler(Func<HttpRequestMessage, CancellationToken, Task<HttpResponseMessage>> func)
+        {
+            _func = func ?? throw new ArgumentNullException(nameof(func));
+        }
+
+        public static HttpMessageInvoker CreateClient(Func<HttpRequestMessage, CancellationToken, Task<HttpResponseMessage>> func)
+        {
+            var handler = new MockHttpHandler(func);
+            return new HttpMessageInvoker(handler);
+        }
+
+        protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
+        {
+            return _func(request, cancellationToken);
+        }
+    }
+}
diff --git a/test/ReverseProxy.Tests/Configuration/ConfigurationConfigProviderTests.cs b/test/ReverseProxy.Tests/Configuration/ConfigurationConfigProviderTests.cs
@@ -88,6 +88,7 @@ public class ConfigurationConfigProviderTests
                             SslProtocols = SslProtocols.Tls11 | SslProtocols.Tls12,
                             MaxConnectionsPerServer = 10,
                             DangerousAcceptAnyServerCertificate = true,
+                            PropagateActivityContext = true,
                         },
                         HttpRequest = new ProxyHttpRequestOptions()
                         {
@@ -227,7 +228,8 @@ public class ConfigurationConfigProviderTests
                     ""Location"": null,
                     ""AllowInvalid"": null
                 },
-                ""MaxConnectionsPerServer"": 10
+                ""MaxConnectionsPerServer"": 10,
+                ""PropagateActivityContext"": true,
             },
             ""HttpRequest"": {
                 ""RequestTimeout"": ""00:01:00"",
@@ -618,6 +620,7 @@ private void VerifyValidAbstractConfig(IProxyConfig validConfig, X509Certificate
             Assert.Equal(cluster1.SessionAffinity.Settings, abstractCluster1.SessionAffinity.Settings);
             Assert.Same(certificate, abstractCluster1.HttpClient.ClientCertificate);
             Assert.Equal(cluster1.HttpClient.MaxConnectionsPerServer, abstractCluster1.HttpClient.MaxConnectionsPerServer);
+            Assert.Equal(cluster1.HttpClient.PropagateActivityContext, abstractCluster1.HttpClient.PropagateActivityContext);
             Assert.Equal(SslProtocols.Tls11 | SslProtocols.Tls12, abstractCluster1.HttpClient.SslProtocols);
             Assert.Equal(cluster1.HttpRequest.RequestTimeout, abstractCluster1.HttpRequest.RequestTimeout);
             Assert.Equal(HttpVersion.Version10, abstractCluster1.HttpRequest.Version);

diff --git a/test/ReverseProxy.Tests/Service/Proxy/HttpProxyTests.cs b/test/ReverseProxy.Tests/Service/Proxy/HttpProxyTests.cs
@@ -1606,27 +1606,6 @@ private static string StreamToString(Stream stream)
             return reader.ReadToEnd();
         }
 
-        private class MockHttpHandler : HttpMessageHandler
-        {
-            private readonly Func<HttpRequestMessage, CancellationToken, Task<HttpResponseMessage>> func;
-
-            private MockHttpHandler(Func<HttpRequestMessage, CancellationToken, Task<HttpResponseMessage>> func)
-            {
-                this.func = func ?? throw new ArgumentNullException(nameof(func));
-            }
-
-            public static HttpMessageInvoker CreateClient(Func<HttpRequestMessage, CancellationToken, Task<HttpResponseMessage>> func)
-            {
-                var handler = new MockHttpHandler(func);
-                return new HttpMessageInvoker(handler);
-            }
-
-            protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
-            {
-                return func(request, cancellationToken);
-            }
-        }
-
         private class DuplexStream : Stream
         {
             public DuplexStream(Stream readStream, Stream writeStream)