Add support for ValueTask instrumentations for < NET Core 3.1

tracer/src/Datadog.Trace/ClrProfiler/CallTarget/Handlers/Continuations/IValueTaskDuckType.cs

@@ -0,0 +1,18 @@
+// <copyright file="IValueTaskDuckType.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+
+#if !NETCOREAPP3_1_OR_GREATER
+#nullable enable
+using System.Threading.Tasks;
+
+namespace Datadog.Trace.ClrProfiler.CallTarget.Handlers.Continuations;
+
+internal interface IValueTaskDuckType
+{
+    bool IsCompletedSuccessfully { get; }
+
+    Task AsTask();
+}
+#endif

tracer/src/Datadog.Trace/ClrProfiler/CallTarget/Handlers/Continuations/ValueTaskActivator.cs

@@ -0,0 +1,65 @@
+// <copyright file="ValueTaskActivator.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+
+#if !NETCOREAPP3_1_OR_GREATER
+#nullable enable
+
+using System;
+using System.Reflection.Emit;
+using System.Threading.Tasks;
+using Datadog.Trace.DuckTyping;
+using Datadog.Trace.Logging;
+using Datadog.Trace.Util;
+
+namespace Datadog.Trace.ClrProfiler.CallTarget.Handlers.Continuations;
+
+internal static class ValueTaskActivator<TValueTask>
+{
+    private static readonly Func<Task, TValueTask> Activator;
+
+    static ValueTaskActivator()
+    {
+        try
+        {
+            Activator = CreateActivator();
+        }
+        catch (Exception ex)
+        {
+            DatadogLogging.GetLoggerFor<ActivatorHelper>()
+                          .Error(ex, "Error creating the custom activator for: {Type}", typeof(TValueTask).FullName);
+
+            // Unfortunately this will box the ValueTask, but I think it's still the best we can do in this scenario
+            Activator = FallbackActivator;
+        }
+    }
+
+    // Internal for testing
+    internal static Func<Task, TValueTask> CreateActivator()
+    {
+        var valueTaskType = typeof(TValueTask);
+        var ctor = valueTaskType.GetConstructor([typeof(Task)])!;
+
+        var createValueTaskMethod = new DynamicMethod(
+            $"TypeActivator" + valueTaskType.Name,
+            returnType: valueTaskType,
+            parameterTypes: [typeof(Task)],
+            typeof(DuckType).Module,
+            true);
+
+        var il = createValueTaskMethod.GetILGenerator();
+        il.Emit(OpCodes.Ldarg_0);
+        il.Emit(OpCodes.Newobj, ctor);
+        il.Emit(OpCodes.Ret);
+
+        return (Func<Task, TValueTask>)createValueTaskMethod.CreateDelegate(typeof(Func<Task, TValueTask>));
+    }
+
+    // Internal for testing
+    internal static TValueTask FallbackActivator(Task task)
+        => (TValueTask)System.Activator.CreateInstance(typeof(TValueTask), task)!;
+
+    public static TValueTask CreateInstance(Task task) => Activator(task);
+}
+#endif

tracer/src/Datadog.Trace/ClrProfiler/CallTarget/Handlers/Continuations/ValueTaskActivator`1.cs

@@ -0,0 +1,94 @@
+// <copyright file="ValueTaskActivator`1.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+
+#if !NETCOREAPP3_1_OR_GREATER
+#nullable enable
+
+using System;
+using System.Reflection.Emit;
+using System.Threading.Tasks;
+using Datadog.Trace.DuckTyping;
+using Datadog.Trace.Logging;
+using Datadog.Trace.Util;
+
+#pragma warning disable SA1649 // File name must match first type name
+
+namespace Datadog.Trace.ClrProfiler.CallTarget.Handlers.Continuations;
+
+internal static class ValueTaskActivator<TValueTask, TResult>
+{
+    private static readonly Func<Task<TResult>, TValueTask> TaskActivator;
+    private static readonly Func<TResult, TValueTask> ResultActivator;
+
+    static ValueTaskActivator()
+    {
+        try
+        {
+            TaskActivator = CreateTaskActivator();
+            ResultActivator = CreateResultActivator();
+        }
+        catch (Exception ex)
+        {
+            DatadogLogging.GetLoggerFor<ActivatorHelper>()
+                          .Error(ex, "Error creating the custom activator for: {Type}", typeof(TValueTask).FullName);
+
+            // Unfortunately this will box the ValueTask, but I think it's still the best we can do in this scenario
+            TaskActivator = FallbackTaskActivator;
+            ResultActivator = FallbackResultActivator;
+        }
+    }
+
+    // Internal for testing
+    internal static Func<Task<TResult>, TValueTask> CreateTaskActivator()
+    {
+        var valueTaskType = typeof(TValueTask);
+        var ctor = valueTaskType.GetConstructor([typeof(Task<TResult>)])!;
+
+        var createValueTaskMethod = new DynamicMethod(
+            $"TypeActivatorTask" + valueTaskType.Name,
+            returnType: valueTaskType,
+            parameterTypes: [typeof(Task<TResult>)],
+            typeof(DuckType).Module,
+            true);
+
+        var il = createValueTaskMethod.GetILGenerator();
+        il.Emit(OpCodes.Ldarg_0);
+        il.Emit(OpCodes.Newobj, ctor);
+        il.Emit(OpCodes.Ret);
+
+        return (Func<Task<TResult>, TValueTask>)createValueTaskMethod.CreateDelegate(typeof(Func<Task<TResult>, TValueTask>));
+    }
+
+    // Internal for testing
+    internal static Func<TResult, TValueTask> CreateResultActivator()
+    {
+        var valueTaskType = typeof(TValueTask);
+        var ctor = valueTaskType.GetConstructor([typeof(TResult)])!;
+
+        var createValueTaskMethod = new DynamicMethod(
+            $"TypeActivatorResult" + valueTaskType.Name,
+            returnType: valueTaskType,
+            parameterTypes: [typeof(TResult)],
+            typeof(DuckType).Module,
+            true);
+
+        var il = createValueTaskMethod.GetILGenerator();
+        il.Emit(OpCodes.Ldarg_0);
+        il.Emit(OpCodes.Newobj, ctor);
+        il.Emit(OpCodes.Ret);
+
+        return (Func<TResult, TValueTask>)createValueTaskMethod.CreateDelegate(typeof(Func<TResult, TValueTask>));
+    }
+
+    // Internal for testing
+    internal static TValueTask FallbackTaskActivator(Task<TResult> task)
+        => (TValueTask)Activator.CreateInstance(typeof(TValueTask), task)!;
+
+    internal static TValueTask FallbackResultActivator(TResult task)
+        => (TValueTask)Activator.CreateInstance(typeof(TValueTask), task)!;
+
+    public static TValueTask CreateInstance(Task<TResult> task) => TaskActivator(task);
+}
+#endif

tracer/src/Datadog.Trace/ClrProfiler/CallTarget/Handlers/Continuations/ValueTaskContinuationGenerator.NetFx.cs

@@ -0,0 +1,205 @@
+// <copyright file="ValueTaskContinuationGenerator.NetFx.cs" company="Datadog">
+// Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
+// This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
+// </copyright>
+#nullable enable
+
+#if !NETCOREAPP3_1_OR_GREATER
+using System;
+using System.Threading.Tasks;
+using Datadog.Trace.DuckTyping;
+using Datadog.Trace.Vendors.Serilog.Events;
+
+namespace Datadog.Trace.ClrProfiler.CallTarget.Handlers.Continuations;
+
+internal class ValueTaskContinuationGenerator<TIntegration, TTarget, TReturn> : ContinuationGenerator<TTarget, TReturn>
+{
+    private static readonly CallbackHandler Resolver;
+
+    static ValueTaskContinuationGenerator()
+    {
+        var result = IntegrationMapper.CreateAsyncEndMethodDelegate(typeof(TIntegration), typeof(TTarget), typeof(object));
+        if (result.Method is not null)
+        {
+            if (result.Method.ReturnType == typeof(Task) ||
+                (result.Method.ReturnType.IsGenericType && typeof(Task).IsAssignableFrom(result.Method.ReturnType)))
+            {
+                var asyncContinuation = (AsyncObjectContinuationMethodDelegate)result.Method.CreateDelegate(typeof(AsyncObjectContinuationMethodDelegate));
+                Resolver = new AsyncCallbackHandler(asyncContinuation, result.PreserveContext);
+            }
+            else
+            {
+                var continuation = (ObjectContinuationMethodDelegate)result.Method.CreateDelegate(typeof(ObjectContinuationMethodDelegate));
+                Resolver = new SyncCallbackHandler(continuation, result.PreserveContext);
+            }
+        }
+        else
+        {
+            Resolver = new NoOpCallbackHandler();
+        }
+
+        if (Log.IsEnabled(LogEventLevel.Debug))
+        {
+            Log.Debug(
+                "== {TaskContinuationGenerator} using Resolver: {Resolver}",
+                $"TaskContinuationGenerator<{typeof(TIntegration).FullName}, {typeof(TTarget).FullName}, {typeof(TReturn).FullName}>",
+                Resolver.GetType().FullName);
+        }
+    }
+
+    public override TReturn? SetContinuation(TTarget? instance, TReturn? returnValue, Exception? exception, in CallTargetState state)
+    {
+        return Resolver.ExecuteCallback(instance, returnValue, exception, in state);
+    }
+
+    private class SyncCallbackHandler : CallbackHandler
+    {
+        private readonly ObjectContinuationMethodDelegate _continuation;
+        private readonly bool _preserveContext;
+
+        public SyncCallbackHandler(ObjectContinuationMethodDelegate continuation, bool preserveContext)
+        {
+            _continuation = continuation;
+            _preserveContext = preserveContext;
+        }
+
+        public override TReturn? ExecuteCallback(TTarget? instance, TReturn? returnValue, Exception? exception, in CallTargetState state)
+        {
+            if (exception != null || returnValue == null)
+            {
+                _continuation(instance, default, exception, in state);
+                return returnValue;
+            }
+
+            var previousValueTask = returnValue.DuckCast<IValueTaskDuckType>();
+            if (previousValueTask.IsCompletedSuccessfully)
+            {
+                // ok all good, just run synchronously
+                _continuation(instance, returnValue, exception, in state);
+                return returnValue;
+            }
+
+            // uh oh, need to extract the task, await it, run the continuation
+            var task = previousValueTask.AsTask();
+            var secondTask = ContinuationAction(task, instance, state);
+            // need to wrap the secondTask in a ValueTask to return it
+            return ValueTaskActivator<TReturn>.CreateInstance(secondTask);
+        }
+
+        private async Task ContinuationAction(Task previousValueTask, TTarget? target, CallTargetState state)
+        {
+            try
+            {
+                await previousValueTask.ConfigureAwait(_preserveContext);
+            }
+            catch (Exception ex)
+            {
+                try
+                {
+                    // *
+                    // Calls the CallTarget integration continuation, exceptions here should never bubble up to the application
+                    // *
+                    _continuation(target, default, ex, in state);
+                }
+                catch (Exception contEx)
+                {
+                    IntegrationOptions<TIntegration, TTarget>.LogException(contEx);
+                }
+
+                throw;
+            }
+
+            try
+            {
+                // *
+                // Calls the CallTarget integration continuation, exceptions here should never bubble up to the application
+                // *
+                _continuation(target, default, default, in state);
+            }
+            catch (Exception contEx)
+            {
+                IntegrationOptions<TIntegration, TTarget>.LogException(contEx);
+            }
+        }
+    }
+
+    private class AsyncCallbackHandler : CallbackHandler
+    {
+        private readonly AsyncObjectContinuationMethodDelegate _asyncContinuation;
+        private readonly bool _preserveContext;
+
+        public AsyncCallbackHandler(AsyncObjectContinuationMethodDelegate asyncContinuation, bool preserveContext)
+        {
+            _asyncContinuation = asyncContinuation;
+            _preserveContext = preserveContext;
+        }
+
+        public override TReturn ExecuteCallback(TTarget? instance, TReturn? returnValue, Exception? exception, in CallTargetState state)
+        {
+            Task task;
+            if (returnValue == null)
+            {
+                task = Task.CompletedTask;
+            }
+            else
+            {
+                var previousValueTask = returnValue.DuckCast<IValueTaskDuckType>();
+                if (previousValueTask.IsCompletedSuccessfully)
+                {
+                    // ok all good, just run synchronously
+                    task = Task.CompletedTask;
+                }
+                else
+                {
+                    task = previousValueTask.AsTask();
+                }
+            }
+
+            var secondTask = ContinuationAction(task, instance, state, exception);
+            // need to wrap the secondTask in a ValueTask to return it
+            return ValueTaskActivator<TReturn>.CreateInstance(secondTask);
+        }
+
+        private async Task ContinuationAction(Task previousValueTask, TTarget? target, CallTargetState state, Exception? exception)
+        {
+            if (exception != null)
+            {
+                await _asyncContinuation(target, default, exception, in state).ConfigureAwait(_preserveContext);
+            }
+
+            try
+            {
+                await previousValueTask.ConfigureAwait(_preserveContext);
+            }
+            catch (Exception ex)
+            {
+                try
+                {
+                    // *
+                    // Calls the CallTarget integration continuation, exceptions here should never bubble up to the application
+                    // *
+                    await _asyncContinuation(target, default, ex, in state).ConfigureAwait(_preserveContext);
+                }
+                catch (Exception contEx)
+                {
+                    IntegrationOptions<TIntegration, TTarget>.LogException(contEx);
+                }
+
+                throw;
+            }
+
+            try
+            {
+                // *
+                // Calls the CallTarget integration continuation, exceptions here should never bubble up to the application
+                // *
+                await _asyncContinuation(target, default, default, in state).ConfigureAwait(_preserveContext);
+            }
+            catch (Exception contEx)
+            {
+                IntegrationOptions<TIntegration, TTarget>.LogException(contEx);
+            }
+        }
+    }
+}
+#endif