Fix distributed tracing in ASP.NET Core MVC (DataDog#303)

* add logging

* use DictionaryHeadersCollection to extract distributed tracing headers

* emit IL instead of using `dynamic`

Author: lucaspimentel

src/Datadog.Trace.ClrProfiler.Managed/Integrations/AspNetCoreMvc2Integration.cs

@@ -1,7 +1,10 @@
 using System;
+using System.Collections;
 using System.Collections.Generic;
 using System.Reflection;
 using System.Text;
+using Datadog.Trace.Headers;
+using Datadog.Trace.Logging;
 
 namespace Datadog.Trace.ClrProfiler.Integrations
 {
@@ -13,6 +16,8 @@ public sealed class AspNetCoreMvc2Integration : IDisposable
         internal const string OperationName = "aspnet-coremvc.request";
         private const string HttpContextKey = "__Datadog.Trace.ClrProfiler.Integrations." + nameof(AspNetCoreMvc2Integration);
 
+        private static readonly ILog Log = LogProvider.GetLogger(typeof(AspNetCoreMvc2Integration));
+
         private static Action<object, object, object, object> _beforeAction;
         private static Action<object, object, object, object> _afterAction;
 
@@ -36,7 +41,33 @@ public AspNetCoreMvc2Integration(object actionDescriptorObj, object httpContextO
                 string httpMethod = _httpContext.Request.Method.ToUpperInvariant();
                 string url = GetDisplayUrl(_httpContext.Request).ToLowerInvariant();
 
-                _scope = Tracer.Instance.StartActive(OperationName);
+                SpanContext propagatedContext = null;
+
+                if (Tracer.Instance.ActiveScope == null)
+                {
+                    try
+                    {
+                        // extract propagated http headers
+                        IEnumerable requestHeaders = _httpContext.Request.Headers;
+                        int headerCount = _httpContext.Request.Headers.Count;
+                        var headersCollection = new DictionaryHeadersCollection(headerCount);
+
+                        foreach (dynamic header in requestHeaders)
+                        {
+                            string key = header.Key;
+                            string[] values = header.Value.ToArray();
+                            headersCollection.Add(key, values);
+                        }
+
+                        propagatedContext = SpanContextPropagator.Instance.Extract(headersCollection);
+                    }
+                    catch (Exception ex)
+                    {
+                        Log.ErrorException("Error extracting propagated HTTP headers.", ex);
+                    }
+                }
+
+                _scope = Tracer.Instance.StartActive(OperationName, propagatedContext);
                 var span = _scope.Span;
                 span.Type = SpanTypes.Web;
                 span.ResourceName = $"{httpMethod} {controllerName}.{actionName}";
@@ -45,9 +76,10 @@ public AspNetCoreMvc2Integration(object actionDescriptorObj, object httpContextO
                 span.SetTag(Tags.AspNetController, controllerName);
                 span.SetTag(Tags.AspNetAction, actionName);
             }
-            catch
+            catch (Exception) when (DisposeObject(_scope))
             {
-                // TODO: logging
+                // unreachable code
+                throw;
             }
         }
 
@@ -63,22 +95,25 @@ public AspNetCoreMvc2Integration(object actionDescriptorObj, object httpContextO
             TargetAssembly = "Microsoft.AspNetCore.Mvc.Core",
             TargetType = "Microsoft.AspNetCore.Mvc.Internal.MvcCoreDiagnosticSourceExtensions")]
         public static void BeforeAction(
-                                        object diagnosticSource,
-                                        object actionDescriptor,
-                                        dynamic httpContext,
-                                        object routeData)
+            object diagnosticSource,
+            object actionDescriptor,
+            object httpContext,
+            object routeData)
         {
             AspNetCoreMvc2Integration integration = null;
 
             try
             {
                 integration = new AspNetCoreMvc2Integration(actionDescriptor, httpContext);
-                IDictionary<object, object> contextItems = httpContext.Items;
-                contextItems[HttpContextKey] = integration;
+
+                if (httpContext.TryGetPropertyValue("Items", out IDictionary<object, object> contextItems))
+                {
+                    contextItems[HttpContextKey] = integration;
+                }
             }
-            catch
+            catch (Exception ex)
             {
-                // TODO: log this as an instrumentation error, but continue calling instrumented method
+                Log.ErrorExceptionForFilter($"Error creating {nameof(AspNetCoreMvc2Integration)}.", ex);
             }
 
             try
@@ -89,25 +124,24 @@ public static void BeforeAction(
                     var type = assembly.GetType("Microsoft.AspNetCore.Mvc.Internal.MvcCoreDiagnosticSourceExtensions");
 
                     _beforeAction = DynamicMethodBuilder<Action<object, object, object, object>>.CreateMethodCallDelegate(
-                                                                                                                          type,
-                                                                                                                          "BeforeAction");
+                        type,
+                        "BeforeAction");
                 }
             }
-            catch
+            catch (Exception ex)
             {
-                // TODO: log this as an instrumentation error, we cannot call instrumented method,
                 // profiled app will continue working without DiagnosticSource
+                Log.ErrorException("Error calling Microsoft.AspNetCore.Mvc.Internal.MvcCoreDiagnosticSourceExtensions.BeforeAction()", ex);
             }
 
             try
             {
                 // call the original method, catching and rethrowing any unhandled exceptions
                 _beforeAction?.Invoke(diagnosticSource, actionDescriptor, httpContext, routeData);
             }
-            catch (Exception ex)
+            catch (Exception ex) when (integration?.SetException(ex) ?? false)
             {
-                integration?.SetException(ex);
-
+                // unreachable code
                 throw;
             }
         }
@@ -124,21 +158,23 @@ public static void BeforeAction(
             TargetAssembly = "Microsoft.AspNetCore.Mvc.Core",
             TargetType = "Microsoft.AspNetCore.Mvc.Internal.MvcCoreDiagnosticSourceExtensions")]
         public static void AfterAction(
-                                       object diagnosticSource,
-                                       object actionDescriptor,
-                                       dynamic httpContext,
-                                       object routeData)
+            object diagnosticSource,
+            object actionDescriptor,
+            object httpContext,
+            object routeData)
         {
             AspNetCoreMvc2Integration integration = null;
 
             try
             {
-                IDictionary<object, object> contextItems = httpContext?.Items;
-                integration = contextItems?[HttpContextKey] as AspNetCoreMvc2Integration;
+                if (httpContext.TryGetPropertyValue("Items", out IDictionary<object, object> contextItems))
+                {
+                    integration = contextItems?[HttpContextKey] as AspNetCoreMvc2Integration;
+                }
             }
-            catch
+            catch (Exception ex)
             {
-                // TODO: log this as an instrumentation error, but continue calling instrumented method
+                Log.ErrorExceptionForFilter($"Error accessing {nameof(AspNetCoreMvc2Integration)}.", ex);
             }
 
             try
@@ -148,8 +184,8 @@ public static void AfterAction(
                     var type = actionDescriptor.GetType().Assembly.GetType("Microsoft.AspNetCore.Mvc.Internal.MvcCoreDiagnosticSourceExtensions");
 
                     _afterAction = DynamicMethodBuilder<Action<object, object, object, object>>.CreateMethodCallDelegate(
-                                                                                                                         type,
-                                                                                                                         "AfterAction");
+                        type,
+                        "AfterAction");
                 }
             }
             catch
@@ -179,9 +215,11 @@ public static void AfterAction(
         /// Tags the current span as an error. Called when an unhandled exception is thrown in the instrumented method.
         /// </summary>
         /// <param name="ex">The exception that was thrown and not handled in the instrumented method.</param>
-        public void SetException(Exception ex)
+        /// <returns>Always <c>false</c>.</returns>
+        public bool SetException(Exception ex)
         {
             _scope?.Span?.SetException(ex);
+            return false;
         }
 
         /// <summary>
@@ -217,5 +255,11 @@ private static string GetDisplayUrl(dynamic request)
                                                                                                                                              .Append(queryString)
                                                                                                                                              .ToString();
         }
+
+        private bool DisposeObject(IDisposable disposable)
+        {
+            disposable?.Dispose();
+            return false;
+        }
     }
 }

src/Datadog.Trace/Headers/DictionaryHeadersCollection.cs

@@ -0,0 +1,62 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+
+namespace Datadog.Trace.Headers
+{
+    internal class DictionaryHeadersCollection : IHeadersCollection
+    {
+        private readonly IDictionary<string, IList<string>> _headers;
+
+        public DictionaryHeadersCollection()
+        {
+            _headers = new Dictionary<string, IList<string>>(StringComparer.OrdinalIgnoreCase);
+        }
+
+        public DictionaryHeadersCollection(int capacity)
+        {
+            _headers = new Dictionary<string, IList<string>>(capacity, StringComparer.OrdinalIgnoreCase);
+        }
+
+        public DictionaryHeadersCollection(IDictionary<string, IList<string>> dictionary)
+        {
+            _headers = dictionary ?? throw new ArgumentNullException(nameof(dictionary));
+        }
+
+        public IEnumerable<string> GetValues(string name)
+        {
+            return _headers.TryGetValue(name, out var values)
+                       ? values
+                       : Enumerable.Empty<string>();
+        }
+
+        public void Set(string name, string value)
+        {
+            _headers.Remove(name);
+            _headers.Add(name, new List<string> { value });
+        }
+
+        public void Add(string name, string value)
+        {
+            Add(name, new[] { value });
+        }
+
+        public void Add(string name, IEnumerable<string> values)
+        {
+            if (!_headers.TryGetValue(name, out var list))
+            {
+                list = new List<string>();
+            }
+
+            foreach (var value in values)
+            {
+                list.Add(value);
+            }
+        }
+
+        public void Remove(string name)
+        {
+            _headers.Remove(name);
+        }
+    }
+}