use async/await when instrumenting asynchronous methods (redux)

src/Datadog.Trace.ClrProfiler.Managed/Datadog.Trace.ClrProfiler.Managed.csproj

@@ -26,6 +26,7 @@
 
   <ItemGroup Condition="'$(TargetFramework)' == 'net45'">
     <Reference Include="System.Configuration" />
+    <Reference Include="System.Net.Http" />
     <Reference Include="System.Web" />
   </ItemGroup>
 

src/Datadog.Trace.ClrProfiler.Managed/DynamicMethodBuilder.cs

@@ -58,18 +58,15 @@ public static TDelegate CreateMethodCallDelegate(
             Type delegateType = typeof(TDelegate);
             Type[] genericTypeArguments = delegateType.GenericTypeArguments;
 
-            Type returnType;
             Type[] parameterTypes;
 
             if (delegateType.Name.StartsWith("Func`"))
             {
                 // last generic type argument is the return type
-                returnType = genericTypeArguments.Last();
                 parameterTypes = genericTypeArguments.Take(genericTypeArguments.Length - 1).ToArray();
             }
             else if (delegateType.Name.StartsWith("Action`"))
             {
-                returnType = typeof(void);
                 parameterTypes = genericTypeArguments;
             }
             else

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

@@ -2,7 +2,9 @@
 
 using System;
 using System.Collections.Generic;
+using System.Net.Http;
 using System.Threading;
+using System.Threading.Tasks;
 using Datadog.Trace.ExtensionMethods;
 
 namespace Datadog.Trace.ClrProfiler.Integrations
@@ -14,52 +16,60 @@ public static class AspNetWebApi2Integration
     {
         internal const string OperationName = "aspnet-web-api.request";
 
-        private static readonly Type HttpControllerContextType = Type.GetType("System.Web.Http.Controllers.HttpControllerContext, System.Web.Http", throwOnError: false);
-
         /// <summary>
-        /// ExecuteAsync calls the underlying ExecuteAsync and traces the request.
+        /// Calls the underlying ExecuteAsync and traces the request.
         /// </summary>
-        /// <param name="this">The Api Controller</param>
+        /// <param name="apiController">The Api Controller</param>
         /// <param name="controllerContext">The controller context for the call</param>
         /// <param name="cancellationTokenSource">The cancellation token source</param>
         /// <returns>A task with the result</returns>
-        public static object ExecuteAsync(object @this, object controllerContext, object cancellationTokenSource)
+        public static object ExecuteAsync(object apiController, object controllerContext, object cancellationTokenSource)
         {
-            Type controllerType = @this.GetType();
+            var cancellationToken = ((CancellationTokenSource)cancellationTokenSource).Token;
+            return ExecuteAsyncInternal(apiController, controllerContext, cancellationToken);
+        }
+
+        /// <summary>
+        /// Calls the underlying ExecuteAsync and traces the request.
+        /// </summary>
+        /// <param name="apiController">The Api Controller</param>
+        /// <param name="controllerContext">The controller context for the call</param>
+        /// <param name="cancellationToken">The cancellation token</param>
+        /// <returns>A task with the result</returns>
+        private static async Task<HttpResponseMessage> ExecuteAsyncInternal(object apiController, object controllerContext, CancellationToken cancellationToken)
+        {
+            Type controllerType = apiController.GetType();
 
             // in some cases, ExecuteAsync() is an explicit interface implementation,
             // which is not public and has a different name, so try both
             var executeAsyncFunc =
-                DynamicMethodBuilder<Func<object, object, CancellationToken, object>>
+                DynamicMethodBuilder<Func<object, object, CancellationToken, Task<HttpResponseMessage>>>
                    .GetOrCreateMethodCallDelegate(controllerType, "ExecuteAsync") ??
-                DynamicMethodBuilder<Func<object, object, CancellationToken, object>>
+                DynamicMethodBuilder<Func<object, object, CancellationToken, Task<HttpResponseMessage>>>
                    .GetOrCreateMethodCallDelegate(controllerType, "System.Web.Http.Controllers.IHttpController.ExecuteAsync");
 
             using (Scope scope = CreateScope(controllerContext))
             {
-                return scope.Span.Trace(
-                                        () =>
-                                        {
-                                            CancellationToken cancellationToken = ((CancellationTokenSource)cancellationTokenSource).Token;
-                                            return executeAsyncFunc(@this, controllerContext, cancellationToken);
-                                        },
-                                        onComplete: e =>
-                                        {
-                                            if (e != null)
-                                            {
-                                                scope.Span.SetException(e);
-                                            }
+                try
+                {
+                    var responseMessage = await executeAsyncFunc(apiController, controllerContext, cancellationToken).ConfigureAwait(false);
+
+                    // some fields aren't set till after execution, so populate anything missing
+                    UpdateSpan(controllerContext, scope.Span);
 
-                                            // some fields aren't set till after execution, so repopulate anything missing
-                                            UpdateSpan(controllerContext, scope.Span);
-                                            scope.Span.Finish();
-                                        });
+                    return responseMessage;
+                }
+                catch (Exception ex)
+                {
+                    scope.Span.SetException(ex);
+                    throw;
+                }
             }
         }
 
-        private static Scope CreateScope(dynamic controllerContext)
+        private static Scope CreateScope(object controllerContext)
         {
-            var scope = Tracer.Instance.StartActive(OperationName, finishOnClose: false);
+            var scope = Tracer.Instance.StartActive(OperationName);
             UpdateSpan(controllerContext, scope.Span);
             return scope;
         }

src/Datadog.Trace.ClrProfiler.Managed/Integrations/Elasticsearch.Net/Pipeline.cs

@@ -21,68 +21,100 @@ public static class Pipeline
         private const string ElasticsearchUrlKey = "elasticsearch.url";
         private const string ElasticsearchParamsKey = "elasticsearch.params";
 
+        private static readonly Type CancellationTokenType = typeof(CancellationToken);
+
         private static Type _requestDataType;
-        private static Type _cancellationTokenType = typeof(CancellationToken);
 
         /// <summary>
-        /// CallElasticsearch traces a call to Elasticsearch
+        /// Traces a synchronous call to Elasticsearch.
         /// </summary>
         /// <typeparam name="TResponse">The type of the response</typeparam>
         /// <param name="pipeline">The pipeline for the original method</param>
         /// <param name="requestData">The request data</param>
         /// <returns>The original result</returns>
         public static object CallElasticsearch<TResponse>(object pipeline, object requestData)
         {
-            var originalMethod = DynamicMethodBuilder<Func<object, object, TResponse>>.GetOrCreateMethodCallDelegate(
-                pipeline.GetType(),
-                "CallElasticsearch",
-                methodGenericArguments: new Type[] { typeof(TResponse) });
-            return CreateScope(pipeline, requestData).Span.Trace(() =>
+            var originalMethod = DynamicMethodBuilder<Func<object, object, TResponse>>
+               .GetOrCreateMethodCallDelegate(
+                    pipeline.GetType(),
+                    "CallElasticsearch",
+                    methodGenericArguments: new Type[] { typeof(TResponse) });
+
+            using (var scope = CreateScope(pipeline, requestData))
             {
-                return originalMethod(pipeline, requestData);
-            });
+                try
+                {
+                    return originalMethod(pipeline, requestData);
+                }
+                catch (Exception ex)
+                {
+                    scope.Span.SetException(ex);
+                    throw;
+                }
+            }
         }
 
         /// <summary>
-        /// CallElasticsearchAsync traces an asynchronous call to Elasticsearch
+        /// Traces an asynchronous call to Elasticsearch.
         /// </summary>
         /// <typeparam name="TResponse">Type type of the response</typeparam>
         /// <param name="pipeline">The pipeline for the original method</param>
         /// <param name="requestData">The request data</param>
         /// <param name="cancellationTokenSource">A cancellation token</param>
         /// <returns>The original result</returns>
         public static object CallElasticsearchAsync<TResponse>(object pipeline, object requestData, object cancellationTokenSource)
+        {
+            var cancellationToken = ((CancellationTokenSource)cancellationTokenSource)?.Token ?? CancellationToken.None;
+            return CallElasticsearchAsyncInternal<TResponse>(pipeline, requestData, cancellationToken);
+        }
+
+        /// <summary>
+        /// Traces an asynchronous call to Elasticsearch.
+        /// </summary>
+        /// <typeparam name="TResponse">Type type of the response</typeparam>
+        /// <param name="pipeline">The pipeline for the original method</param>
+        /// <param name="requestData">The request data</param>
+        /// <param name="cancellationToken">A cancellation token</param>
+        /// <returns>The original result</returns>
+        private static async Task<TResponse> CallElasticsearchAsyncInternal<TResponse>(object pipeline, object requestData, CancellationToken cancellationToken)
         {
             if (_requestDataType == null)
             {
-                _requestDataType = pipeline.GetType().Assembly.GetType("Elasticsearch.Net.RequestData");
+                _requestDataType = requestData.GetType();
             }
 
-            var cancellationToken = (cancellationTokenSource as CancellationTokenSource)?.Token ?? CancellationToken.None;
-
-            var originalMethod = DynamicMethodBuilder<Func<object, object, CancellationToken, Task<TResponse>>>.
-                GetOrCreateMethodCallDelegate(
+            var originalMethod = DynamicMethodBuilder<Func<object, object, CancellationToken, Task<TResponse>>>
+               .GetOrCreateMethodCallDelegate(
                     pipeline.GetType(),
                     "CallElasticsearchAsync",
-                    methodParameterTypes: new Type[] { _requestDataType, _cancellationTokenType },
-                    methodGenericArguments: new Type[] { typeof(TResponse) });
+                    new[] { _requestDataType, CancellationTokenType },
+                    new[] { typeof(TResponse) });
 
-            return CreateScope(pipeline, requestData).Span.Trace(() =>
+            using (var scope = CreateScope(pipeline, requestData))
             {
-                return originalMethod(pipeline, requestData, cancellationToken);
-            });
+                try
+                {
+                    return await originalMethod(pipeline, requestData, cancellationToken).ConfigureAwait(false);
+                }
+                catch (Exception ex)
+                {
+                    scope.Span.SetException(ex);
+                    throw;
+                }
+            }
         }
 
-        private static Scope CreateScope(dynamic pipeline, dynamic requestData)
+        private static Scope CreateScope(object pipeline, dynamic requestData)
         {
             string requestName = null;
             try
             {
-                var requestParameters = DynamicMethodBuilder<Func<object, dynamic>>.GetOrCreateMethodCallDelegate(
-                    pipeline.GetType(),
-                    "get_RequestParameters")(pipeline);
-                requestName = requestParameters?.GetType()?.Name;
-                requestName = requestName?.Replace("RequestParameters", string.Empty);
+                var requestParameters = DynamicMethodBuilder<Func<object, object>>
+                   .GetOrCreateMethodCallDelegate(
+                        pipeline.GetType(),
+                        "get_RequestParameters")(pipeline);
+
+                requestName = requestParameters?.GetType().Name.Replace("RequestParameters", string.Empty);
             }
             catch
             {
@@ -117,7 +149,7 @@ private static Scope CreateScope(dynamic pipeline, dynamic requestData)
 
             var serviceName = string.Join("-", Tracer.Instance.DefaultServiceName, ServiceName);
 
-            var scope = Tracer.Instance.StartActive(OperationName, serviceName: serviceName, finishOnClose: false);
+            var scope = Tracer.Instance.StartActive(OperationName, serviceName: serviceName);
             scope.Span.ResourceName = requestName ?? pathAndQuery ?? string.Empty;
             scope.Span.Type = SpanType;
             scope.Span.SetTag(ComponentKey, ComponentValue);