DataDog · pierotibou · Nov 12, 2021 · Nov 4, 2021 · Nov 2, 2021 · Nov 3, 2021
@@ -17,6 +17,8 @@ public class CallTargetDefinitionSource
 
         public string TargetMethod { get; init; }
 
+        public ushort[] TargetMethodArgumentsToLoad { get; init; }
+
         public string[] TargetSignatureTypes { get; init; }
 
         public ushort TargetMinimumMajor { get; init; }

@@ -130,6 +130,7 @@ from assemblyNames in GetPropertyValue<string[]>(attribute, "AssemblyNames")
                                              TargetAssembly = assemblyNames,
                                              TargetType = GetPropertyValue<string>(attribute, "TypeName"),
                                              TargetMethod = GetPropertyValue<string>(attribute, "MethodName"),
+                                             TargetMethodArgumentsToLoad = (GetPropertyValue<ushort[]>(attribute, "TargetMethodArgumentsToLoad") ?? Enumerable.Empty<ushort>()).ToArray(),
                                              TargetSignatureTypes = new string[] { GetPropertyValue<string>(attribute, "ReturnTypeName") }
                                                                    .Concat(GetPropertyValue<string[]>(attribute, "ParameterTypeNames") ?? Enumerable.Empty<string>())
                                                                    .ToArray(),
@@ -187,6 +188,28 @@ static void WriteCallTargetDefinitionFile(StreamWriter swriter, IEnumerable<Call
 
                     swriter.Write(" }, ");
 
+                    if (integration.TargetMethodArgumentsToLoad.Length > 0)
+                    {
+                        swriter.Write($" new ushort[] {{ ");
+                        for (var s = 0; s < integration.TargetMethodArgumentsToLoad.Length; s++)
+                        {
+                            if (s == integration.TargetMethodArgumentsToLoad.Length - 1)
+                            {
+                                swriter.Write($"{integration.TargetMethodArgumentsToLoad[s]}");
+                            }
+                            else
+                            {
+                                swriter.Write($"{integration.TargetMethodArgumentsToLoad[s]}, ");
+                            }
+                        }
+
+                        swriter.Write(" }, ");
+                    }
+                    else
+                    {
+                        swriter.Write("null, ");
+                    }
+
                     swriter.Write($"{integration.TargetMinimumMajor}, ");
                     swriter.Write($"{integration.TargetMinimumMinor}, ");
                     swriter.Write($"{integration.TargetMinimumPatch}, ");

@@ -873,7 +873,8 @@ HRESULT CallTargetTokens::ModifyLocalSigAndInitialize(void* rewriterWrapperPtr,
 
 HRESULT CallTargetTokens::WriteBeginMethod(void* rewriterWrapperPtr, mdTypeRef integrationTypeRef,
                                            const TypeInfo* currentType,
-                                           std::vector<FunctionMethodArgument>& methodArguments, ILInstr** instruction)
+                                           const std::vector<FunctionMethodArgument>& methodArguments,
+                                           const std::vector<USHORT>& argsToLoad, ILInstr** instruction)
 {
     auto hr = EnsureBaseCalltargetTokens();
     if (FAILED(hr))
@@ -884,7 +885,9 @@ HRESULT CallTargetTokens::WriteBeginMethod(void* rewriterWrapperPtr, mdTypeRef i
     ILRewriterWrapper* rewriterWrapper = (ILRewriterWrapper*) rewriterWrapperPtr;
     ModuleMetadata* module_metadata = GetMetadata();
 
-    auto numArguments = (int) methodArguments.size();
+    auto numArgsToLoad = argsToLoad.size();
+    bool useArgToLoad = numArgsToLoad > 0;
+    auto numArguments = useArgToLoad ? (int) numArgsToLoad : (int) methodArguments.size();
     if (numArguments >= FASTPATH_COUNT)
     {
         return WriteBeginMethodWithArgumentsArray(rewriterWrapperPtr, integrationTypeRef, currentType, instruction);
@@ -947,7 +950,17 @@ HRESULT CallTargetTokens::WriteBeginMethod(void* rewriterWrapperPtr, mdTypeRef i
     ULONG argumentsSignatureSize[FASTPATH_COUNT];
     for (auto i = 0; i < numArguments; i++)
     {
-        auto signatureSize = methodArguments[i].GetSignature(argumentsSignatureBuffer[i]);
+        FunctionMethodArgument methodArgument;
+        if (useArgToLoad)
+        {
+            methodArgument = methodArguments[argsToLoad[i]];
+        }
+        else
+        {
+            methodArgument = methodArguments[i];
+        }
+
+        auto signatureSize = methodArgument.GetSignature(argumentsSignatureBuffer[i]);
         argumentsSignatureSize[i] = signatureSize;
         signatureLength += signatureSize;
     }

@@ -86,7 +86,8 @@ class CallTargetTokens
                                         mdToken* callTargetReturnToken, ILInstr** firstInstruction);
 
     HRESULT WriteBeginMethod(void* rewriterWrapperPtr, mdTypeRef integrationTypeRef, const TypeInfo* currentType,
-                             std::vector<FunctionMethodArgument>& methodArguments, ILInstr** instruction);
+                             const std::vector<FunctionMethodArgument>& methodArguments, const std::vector<USHORT>& argsToLoad,
+                             ILInstr** instruction);
 
     HRESULT WriteEndVoidReturnMemberRef(void* rewriterWrapperPtr, mdTypeRef integrationTypeRef,
                                         const TypeInfo* currentType, ILInstr** instruction);

@@ -1251,6 +1251,14 @@ void CorProfiler::InitializeProfiler(WCHAR* id, CallTargetDefinition* items, int
                 wrapperType = WSTRING(current.wrapperType);
             }
 
+            bool useTargetMethodArgumentsToLoad = current.useTargetMethodArgumentsToLoad;
+            std::vector<USHORT> targetMethodArgumentsToLoad;
+            for (int sIdx = 0; sIdx < current.targetMethodArgumentsToLoadLength; sIdx++)
+            {
+                const auto currentArgumentIndex = current.targetMethodArgumentsToLoad[sIdx];
+                targetMethodArgumentsToLoad.push_back(currentArgumentIndex);
+            }
+
             std::vector<WSTRING> signatureTypes;
             for (int sIdx = 0; sIdx < current.signatureTypesLength; sIdx++)
             {
@@ -1271,9 +1279,9 @@ void CorProfiler::InitializeProfiler(WCHAR* id, CallTargetDefinition* items, int
                 MethodReplacement(
                     {},
                     MethodReference(targetAssembly, targetType, targetMethod, EmptyWStr, minVersion, maxVersion,
-                        {}, signatureTypes),
+                        {}, signatureTypes, useTargetMethodArgumentsToLoad, targetMethodArgumentsToLoad),
                     MethodReference(wrapperAssembly, wrapperType, EmptyWStr, calltarget_modification_action, {}, {}, {},
-                                    {})));
+                                    {}, false, {})));
 
             if (Logger::IsDebugEnabled())
             {
@@ -3432,7 +3440,11 @@ HRESULT CorProfiler::CallTarget_RewriterCallback(RejitHandlerModule* moduleHandl
     bool isVoid = (retTypeFlags & TypeFlagVoid) > 0;
     bool isStatic = !(caller->method_signature.CallingConvention() & IMAGE_CEE_CS_CALLCONV_HASTHIS);
     std::vector<FunctionMethodArgument> methodArguments = caller->method_signature.GetMethodArguments();
-    int numArgs = caller->method_signature.NumberOfArguments();
+    bool useCustomArgumentsTargetMethodArguments = method_replacement->target_method.use_target_method_arguments_to_load;
+    std::vector<USHORT> methodArgumentsToLoad = method_replacement->target_method.target_method_arguments_to_load;
+    int numArgs = useCustomArgumentsTargetMethodArguments
+                      ? (int) methodArgumentsToLoad.size()
+                      : caller->method_signature.NumberOfArguments();
     auto metaEmit = module_metadata->metadata_emit;
     auto metaImport = module_metadata->metadata_import;
 
@@ -3547,8 +3559,11 @@ HRESULT CorProfiler::CallTarget_RewriterCallback(RejitHandlerModule* moduleHandl
         // Load the arguments directly (FastPath)
         for (int i = 0; i < numArgs; i++)
         {
-            reWriterWrapper.LoadArgument(i + (isStatic ? 0 : 1));
-            auto argTypeFlags = methodArguments[i].GetTypeFlags(elementType);
+            int loadIndex = useCustomArgumentsTargetMethodArguments
+                                ? methodArgumentsToLoad[i]
+                                : i;
+            reWriterWrapper.LoadArgument(loadIndex + (isStatic ? 0 : 1));
+            auto argTypeFlags = methodArguments[loadIndex].GetTypeFlags(elementType);
             if (argTypeFlags & TypeFlagByRef)
             {
                 Logger::Warn("*** CallTarget_RewriterCallback(): Methods with ref parameters "
@@ -3564,8 +3579,12 @@ HRESULT CorProfiler::CallTarget_RewriterCallback(RejitHandlerModule* moduleHandl
         for (int i = 0; i < numArgs; i++)
         {
             reWriterWrapper.BeginLoadValueIntoArray(i);
-            reWriterWrapper.LoadArgument(i + (isStatic ? 0 : 1));
-            auto argTypeFlags = methodArguments[i].GetTypeFlags(elementType);
+
+            int loadIndex = useCustomArgumentsTargetMethodArguments
+                                ? methodArgumentsToLoad[i]
+                                : i;
+            reWriterWrapper.LoadArgument(loadIndex + (isStatic ? 0 : 1));
+            auto argTypeFlags = methodArguments[loadIndex].GetTypeFlags(elementType);
             if (argTypeFlags & TypeFlagByRef)
             {
                 Logger::Warn("*** CallTarget_RewriterCallback(): Methods with ref parameters "
@@ -3574,7 +3593,7 @@ HRESULT CorProfiler::CallTarget_RewriterCallback(RejitHandlerModule* moduleHandl
             }
             if (argTypeFlags & TypeFlagBoxedType)
             {
-                auto tok = methodArguments[i].GetTypeTok(metaEmit, callTargetTokens->GetCorLibAssemblyRef());
+                auto tok = methodArguments[loadIndex].GetTypeTok(metaEmit, callTargetTokens->GetCorLibAssemblyRef());
                 if (tok == mdTokenNil)
                 {
                     return S_FALSE;
@@ -3621,6 +3640,7 @@ HRESULT CorProfiler::CallTarget_RewriterCallback(RejitHandlerModule* moduleHandl
 
     ILInstr* beginCallInstruction;
     hr = callTargetTokens->WriteBeginMethod(&reWriterWrapper, wrapper_type_ref, &caller->type, methodArguments,
+                                            methodArgumentsToLoad,
                                             &beginCallInstruction);
     if (FAILED(hr))
     {

@@ -247,23 +247,27 @@ struct MethodReference
     const Version min_version;
     const Version max_version;
     const std::vector<WSTRING> signature_types;
+    const bool use_target_method_arguments_to_load;
+    const std::vector<USHORT> target_method_arguments_to_load;
 
     MethodReference() :
-        min_version(Version(0, 0, 0, 0)), max_version(Version(USHRT_MAX, USHRT_MAX, USHRT_MAX, USHRT_MAX))
+        min_version(Version(0, 0, 0, 0)), max_version(Version(USHRT_MAX, USHRT_MAX, USHRT_MAX, USHRT_MAX)), use_target_method_arguments_to_load(false)
     {
     }
 
     MethodReference(const WSTRING& assembly_name, WSTRING type_name, WSTRING method_name, WSTRING action,
                     Version min_version, Version max_version, const std::vector<BYTE>& method_signature,
-                    const std::vector<WSTRING>& signature_types) :
+                    const std::vector<WSTRING>& signature_types, const bool use_target_method_arguments_to_load, const std::vector<USHORT>& target_method_arguments_to_load) :
         assembly(*AssemblyReference::GetFromCache(assembly_name)),
         type_name(type_name),
         method_name(method_name),
         action(action),
         method_signature(method_signature),
         min_version(min_version),
         max_version(max_version),
-        signature_types(signature_types)
+        signature_types(signature_types),
+        use_target_method_arguments_to_load(use_target_method_arguments_to_load),
+        target_method_arguments_to_load(target_method_arguments_to_load)
     {
     }
 
@@ -336,6 +340,9 @@ typedef struct _CallTargetDefinition
     WCHAR* targetMethod;
     WCHAR** signatureTypes;
     USHORT signatureTypesLength;
+    bool useTargetMethodArgumentsToLoad;
+    USHORT* targetMethodArgumentsToLoad;
+    USHORT targetMethodArgumentsToLoadLength;
     USHORT targetMinimumMajor;
     USHORT targetMinimumMinor;
     USHORT targetMinimumPatch;

@@ -212,6 +212,7 @@ namespace
         USHORT max_minor = USHRT_MAX;
         USHORT max_patch = USHRT_MAX;
         std::vector<WSTRING> signature_type_array;
+        std::vector<USHORT> target_method_arguments_to_load_array; // Do not populate this because this only applies to CallTarget and we no longer use integrations.json
         WSTRING action = EmptyWStr;
 
         if (is_target_method)
@@ -306,7 +307,7 @@ namespace
             }
         }
         return MethodReference(assembly, type, method, action, Version(min_major, min_minor, min_patch, 0),
-                               Version(max_major, max_minor, max_patch, USHRT_MAX), signature, signature_type_array);
+                               Version(max_major, max_minor, max_patch, USHRT_MAX), signature, signature_type_array, false, target_method_arguments_to_load_array);
     }
 
 } // namespace

@@ -0,0 +1,62 @@
+// <copyright file="AsyncMethodInvoker_InvokeBegin_Integration.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 NETFRAMEWORK
+using System;
+using System.ComponentModel;
+using System.Reflection;
+using Datadog.Trace.ClrProfiler.CallTarget;
+using Datadog.Trace.ClrProfiler.Emit;
+using Datadog.Trace.ClrProfiler.Integrations;
+using Datadog.Trace.Configuration;
+
+namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Wcf
+{
+    /// <summary>
+    /// System.ServiceModel.Dispatcher.AsyncMethodInvoker calltarget instrumentation
+    /// </summary>
+    [InstrumentMethod(
+        AssemblyName = "System.ServiceModel",
+        TypeName = "System.ServiceModel.Dispatcher.AsyncMethodInvoker",
+        MethodName = "InvokeBegin",
+        ReturnTypeName = ClrNames.IAsyncResult,
+        ParameterTypeNames = new[] { ClrNames.Object, "System.Object[]", ClrNames.AsyncCallback, ClrNames.Object },
+        MinimumVersion = "4.0.0",
+        MaximumVersion = "4.*.*",
+        IntegrationName = WcfCommon.IntegrationName)]
+    [Browsable(false)]
+    [EditorBrowsable(EditorBrowsableState.Never)]
+    public class AsyncMethodInvoker_InvokeBegin_Integration
+    {
+        /// <summary>
+        /// OnMethodBegin callback
+        /// </summary>
+        /// <typeparam name="TTarget">Type of the target</typeparam>
+        /// <param name="instance">Instance value, aka `this` of the instrumented method.</param>
+        /// <param name="instanceArg">RequestContext instance</param>
+        /// <param name="inputs">Input arguments</param>
+        /// <param name="callback">Callback argument</param>
+        /// <param name="state">State argument</param>
+        /// <returns>Calltarget state value</returns>
+        public static CallTargetState OnMethodBegin<TTarget>(TTarget instance, object instanceArg, object[] inputs, AsyncCallback callback, object state)
+        {
+            // TODO Just use the OperationContext.Current object to get the span information
+            // context.IncomingMessageHeaders contains:
+            //  - Action
+            //  - To
+            //
+            // context.IncomingMessageProperties contains:
+            // - ["httpRequest"] key to find distributed tracing headers
+            if (!Tracer.Instance.Settings.IsIntegrationEnabled(WcfCommon.IntegrationId) || !Tracer.Instance.Settings.DelayWcfInstrumentationEnabled || WcfCommon.GetCurrentOperationContext is null)
+            {
+                return CallTargetState.GetDefault();
+            }
+
+            var requestContext = WcfCommon.GetCurrentOperationContext()?.GetProperty<object>("RequestContext").GetValueOrDefault();
+            return new CallTargetState(WcfIntegration.CreateScope(requestContext));
+        }
+    }
+}
+#endif
@@ -0,0 +1,50 @@
+// <copyright file="AsyncMethodInvoker_InvokeEnd_Integration.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 NETFRAMEWORK
+using System;
+using System.ComponentModel;
+using System.Reflection;
+using Datadog.Trace.ClrProfiler.CallTarget;
+using Datadog.Trace.Configuration;
+
+namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Wcf
+{
+    /// <summary>
+    /// System.ServiceModel.Dispatcher.AsyncMethodInvoker calltarget instrumentation
+    /// </summary>
+    [InstrumentMethod(
+        AssemblyName = "System.ServiceModel",
+        TypeName = "System.ServiceModel.Dispatcher.AsyncMethodInvoker",
+        MethodName = "InvokeEnd",
+        ReturnTypeName = ClrNames.Object,
+        ParameterTypeNames = new[] { ClrNames.Object, "System.Object[]&", ClrNames.IAsyncResult },
+        TargetMethodArgumentsToLoad = new ushort[] { 0, 2 }, // DO NOT pass the "out object[]" parameter into the instrumentation method
+        MinimumVersion = "4.0.0",
+        MaximumVersion = "4.*.*",
+        IntegrationName = WcfCommon.IntegrationName)]
+    [Browsable(false)]
+    [EditorBrowsable(EditorBrowsableState.Never)]
+    public class AsyncMethodInvoker_InvokeEnd_Integration
+    {
+        /// <summary>
+        /// OnMethodEnd callback
+        /// </summary>
+        /// <typeparam name="TTarget">Type of the target</typeparam>
+        /// <typeparam name="TReturn">Type of the response</typeparam>
+        /// <param name="instance">Instance value, aka `this` of the instrumented method.</param>
+        /// <param name="returnValue">Return value</param>
+        /// <param name="exception">Exception instance in case the original code threw an exception.</param>
+        /// <param name="state">Calltarget state value</param>
+        /// <returns>A response value, in an async scenario will be T of Task of T</returns>
+        public static CallTargetReturn<TReturn> OnMethodEnd<TTarget, TReturn>(TTarget instance, TReturn returnValue, Exception exception, CallTargetState state)
+        {
+            var scope = Tracer.Instance.ActiveScope;
+            scope.DisposeWithException(exception);
+            return new CallTargetReturn<TReturn>(returnValue);
+        }
+    }
+}
+#endif
@@ -42,6 +42,11 @@ public class ChannelHandlerIntegration
         /// <returns>Calltarget state value</returns>
         public static CallTargetState OnMethodBegin<TTarget, TRequestContext, TOperationContext>(TTarget instance, TRequestContext request, TOperationContext currentOperationContext)
         {
+            if (Tracer.Instance.Settings.DelayWcfInstrumentationEnabled)
+            {
+                return CallTargetState.GetDefault();
+            }
+
             return new CallTargetState(WcfIntegration.CreateScope(request));
         }