Perf optimization for protobuf instrumentation

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Protobuf/BufferMessageInternalMergeFromIntegration.cs

@@ -8,6 +8,7 @@
 using System;
 using System.ComponentModel;
 using Datadog.Trace.ClrProfiler.CallTarget;
+using Datadog.Trace.ClrProfiler.CallTarget.Handlers;
 
 namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Protobuf;
 
@@ -28,14 +29,25 @@ namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Protobuf;
 [EditorBrowsable(EditorBrowsableState.Never)]
 public class BufferMessageInternalMergeFromIntegration
 {
+    // For performance reasons, we want to do the actual instrumentation work with a Duck constraint,
+    // but to be able to disable the instrumentation we need the raw type
+    // so we use 2 different methods to have access to both when we need it.
+    // Note: Disabling OnMethodBegin means the OnMethodEnd will not be called afterward.
+
     internal static CallTargetState OnMethodBegin<TTarget, TOutput>(TTarget instance, ref TOutput? output)
+    {
+        Helper.DisableInstrumentationIfInternalProtobufType<TTarget>();
+        return CallTargetState.GetDefault();
+    }
+
+    internal static CallTargetReturn OnMethodEnd<TTarget>(TTarget instance, Exception exception, in CallTargetState state)
         where TTarget : IMessageProxy
     {
-        if (Helper.TryGetDescriptor(instance, out var descriptor))
+        if (instance.Instance != null)
         {
-            SchemaExtractor.EnrichActiveSpanWith(descriptor, "deserialization");
+            SchemaExtractor.EnrichActiveSpanWith(instance.Descriptor, "deserialization");
         }
 
-        return CallTargetState.GetDefault();
+        return CallTargetReturn.GetDefault();
     }
 }

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Protobuf/BufferMessageInternalWriteToIntegration.cs

@@ -5,8 +5,10 @@
 
 #nullable enable
 
+using System;
 using System.ComponentModel;
 using Datadog.Trace.ClrProfiler.CallTarget;
+using Datadog.Trace.ClrProfiler.CallTarget.Handlers;
 
 namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Protobuf;
 
@@ -27,14 +29,25 @@ namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Protobuf;
 [EditorBrowsable(EditorBrowsableState.Never)]
 public class BufferMessageInternalWriteToIntegration
 {
-    internal static CallTargetState OnMethodBegin<TTarget, TCtx>(TTarget instance, ref TCtx ctx)
+    // For performance reasons, we want to do the actual instrumentation work with a Duck constraint,
+    // but to be able to disable the instrumentation we need the raw type
+    // so we use 2 different methods to have access to both when we need it.
+    // Note: Disabling OnMethodBegin means the OnMethodEnd will not be called afterward.
+
+    internal static CallTargetState OnMethodBegin<TTarget, TOutput>(TTarget instance, ref TOutput? output)
+    {
+        Helper.DisableInstrumentationIfInternalProtobufType<TTarget>();
+        return CallTargetState.GetDefault();
+    }
+
+    internal static CallTargetReturn OnMethodEnd<TTarget>(TTarget instance, Exception exception, in CallTargetState state)
         where TTarget : IMessageProxy
     {
-        if (instance.Instance is not null)
+        if (instance.Instance != null)
         {
-            SchemaExtractor.EnrichActiveSpanWith(instance.Descriptor, "deserialization");
+            SchemaExtractor.EnrichActiveSpanWith(instance.Descriptor, "serialization");
         }
 
-        return CallTargetState.GetDefault();
+        return CallTargetReturn.GetDefault();
     }
 }

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Protobuf/Helper.cs

@@ -4,63 +4,25 @@
 // </copyright>
 
 #nullable enable
-
-using System;
-using System.Threading;
-using Datadog.Trace.DuckTyping;
-using Datadog.Trace.Logging;
+using Datadog.Trace.ClrProfiler.CallTarget;
+using Datadog.Trace.ClrProfiler.CallTarget.Handlers;
 
 namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Protobuf;
 
-internal class Helper
+internal static class Helper
 {
-    private static readonly IDatadogLogger Log = DatadogLogging.GetLoggerFor<Helper>();
-
-    private static readonly Lazy<IDescriptorReflectionProxy?> DescriptorReflectionProxy = new(
-        () =>
-        {
-            // prepare incantations to access the static property DescriptorReflection.Descriptor
-            var staticType = Type.GetType("Google.Protobuf.Reflection.DescriptorReflection,Google.Protobuf");
-            if (staticType == null)
-            {
-                return null;
-            }
-
-            var proxyType = typeof(IDescriptorReflectionProxy);
-
-            var proxyResult = DuckType.GetOrCreateProxyType(proxyType, staticType);
-            if (!proxyResult.Success)
-            {
-                Log.Warning("Cannot create proxy for type Google.Protobuf.Reflection.DescriptorReflection, protobuf instrumentation may malfunction.");
-                return null;
-            }
-
-            return (IDescriptorReflectionProxy)proxyResult.CreateInstance(null!);
-        });
-
-    public interface IDescriptorReflectionProxy
+    /// <typeparam name="TMessage">needs to be the raw type (not a DuckType)</typeparam>
+    public static void DisableInstrumentationIfInternalProtobufType<TMessage>()
     {
-        object? Descriptor { get; } // this is actually a static property
-    }
-
-    public static bool TryGetDescriptor(IMessageProxy messageProxy, out MessageDescriptorProxy? descriptor)
-    {
-        descriptor = null;
-        if (messageProxy.Instance is null)
+        var typeName = typeof(TMessage).FullName;
+        if (typeName != null && typeName.StartsWith("Google.Protobuf."))
         {
-            return false;
+            // Google uses protobuf internally in the protobuf library, we don't want to capture those.
+            // We disable the integrations once and for all here.
+            IntegrationOptions<MessageWriteToIntegration, TMessage>.DisableIntegration();
+            IntegrationOptions<MessageMergeFromIntegration, TMessage>.DisableIntegration();
+            IntegrationOptions<BufferMessageInternalWriteToIntegration, TMessage>.DisableIntegration();
+            IntegrationOptions<BufferMessageInternalMergeFromIntegration, TMessage>.DisableIntegration();
         }
-
-        // some public functions that we are instrumenting are also called internally by protobuf,
-        // and there is one case where trying to access the descriptor at that point results in a nullref
-        // because it relies on this property. We check it here to make sure we're not going to generate an exception by accessing it.
-        if (DescriptorReflectionProxy.Value?.Descriptor == null)
-        {
-            return false;
-        }
-
-        // now we know it's safe to access the Descriptor property on the message
-        descriptor = messageProxy.Descriptor;
-        return true;
     }
 }

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Protobuf/IMessageProxy.cs

@@ -17,8 +17,6 @@ internal interface IMessageProxy : IDuckType
 {
     /// <summary>
     /// Gets the descriptor.
-    /// Accessing this property can generate a nullref in some cases,
-    /// use <see cref="Helper.TryGetDescriptor"/> to avoid that.
     /// </summary>
     [Duck(ExplicitInterfaceTypeName = "Google.Protobuf.IMessage")]
     MessageDescriptorProxy Descriptor { get; }

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Protobuf/MessageMergeFromIntegration.cs

@@ -8,6 +8,7 @@
 using System;
 using System.ComponentModel;
 using Datadog.Trace.ClrProfiler.CallTarget;
+using Datadog.Trace.ClrProfiler.CallTarget.Handlers;
 
 namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Protobuf;
 
@@ -28,14 +29,25 @@ namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Protobuf;
 [EditorBrowsable(EditorBrowsableState.Never)]
 public class MessageMergeFromIntegration
 {
+    // For performance reasons, we want to do the actual instrumentation work with a Duck constraint,
+    // but to be able to disable the instrumentation we need the raw type
+    // so we use 2 different methods to have access to both when we need it.
+    // Note: Disabling OnMethodBegin means the OnMethodEnd will not be called afterward.
+
     internal static CallTargetState OnMethodBegin<TTarget, TOutput>(TTarget instance, ref TOutput? output)
+    {
+        Helper.DisableInstrumentationIfInternalProtobufType<TTarget>();
+        return CallTargetState.GetDefault();
+    }
+
+    internal static CallTargetReturn OnMethodEnd<TTarget>(TTarget instance, Exception exception, in CallTargetState state)
         where TTarget : IMessageProxy
     {
-        if (Helper.TryGetDescriptor(instance, out var descriptor))
+        if (instance.Instance != null)
         {
-            SchemaExtractor.EnrichActiveSpanWith(descriptor, "deserialization");
+            SchemaExtractor.EnrichActiveSpanWith(instance.Descriptor, "deserialization");
         }
 
-        return CallTargetState.GetDefault();
+        return CallTargetReturn.GetDefault();
     }
 }

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Protobuf/MessageWriteToIntegration.cs

@@ -5,8 +5,10 @@
 
 #nullable enable
 
+using System;
 using System.ComponentModel;
 using Datadog.Trace.ClrProfiler.CallTarget;
+using Datadog.Trace.ClrProfiler.CallTarget.Handlers;
 
 namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Protobuf;
 
@@ -27,14 +29,25 @@ namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Protobuf;
 [EditorBrowsable(EditorBrowsableState.Never)]
 public class MessageWriteToIntegration
 {
+    // For performance reasons, we want to do the actual instrumentation work with a Duck constraint,
+    // but to be able to disable the instrumentation we need the raw type
+    // so we use 2 different methods to have access to both when we need it.
+    // Note: Disabling OnMethodBegin means the OnMethodEnd will not be called afterward.
+
     internal static CallTargetState OnMethodBegin<TTarget, TOutput>(TTarget instance, ref TOutput? output)
+    {
+        Helper.DisableInstrumentationIfInternalProtobufType<TTarget>();
+        return CallTargetState.GetDefault();
+    }
+
+    internal static CallTargetReturn OnMethodEnd<TTarget>(TTarget instance, Exception exception, in CallTargetState state)
         where TTarget : IMessageProxy
     {
-        if (instance.Instance is not null)
+        if (instance.Instance != null)
         {
             SchemaExtractor.EnrichActiveSpanWith(instance.Descriptor, "serialization");
         }
 
-        return CallTargetState.GetDefault();
+        return CallTargetReturn.GetDefault();
     }
 }

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Protobuf/SchemaExtractor.cs

@@ -53,13 +53,6 @@ internal static void EnrichActiveSpanWith(MessageDescriptorProxy? descriptor, st
             return;
         }
 
-        if (descriptor.Value.File.Name.StartsWith("google/protobuf/", StringComparison.OrdinalIgnoreCase))
-        {
-            // it's a protobuf operation internal to the protobuf library, not the one we want
-            Log.Debug("Skipping instrumentation for internal protobuf schema {Schema}", descriptor.Value.File.Name);
-            return;
-        }
-
         activeSpan.SetTag(Tags.SchemaType, "protobuf");
         activeSpan.SetTag(Tags.SchemaName, descriptor.Value.Name);
         activeSpan.SetTag(Tags.SchemaOperation, operationName);

tracer/test/Datadog.Trace.ClrProfiler.IntegrationTests/GoogleProtobufTests.cs

@@ -46,7 +46,7 @@ public async Task TagTraces(string metadataSchemaVersion)
         SetEnvironmentVariable(ConfigurationKeys.DataStreamsMonitoring.Enabled, "1");
         using var telemetry = this.ConfigureTelemetry();
         using var agent = EnvironmentHelper.GetMockAgent();
-        using (await RunSampleAndWaitForExit(agent, $"{TestPrefix}"))
+        using (await RunSampleAndWaitForExit(agent, "AddressBook"))
         {
             using var assertionScope = new AssertionScope();
             var spans = agent.WaitForSpans(2);
@@ -72,14 +72,31 @@ await VerifyHelper.VerifySpans(
         }
     }
 
+    [Fact]
+    [Trait("Category", "EndToEnd")]
+    public async Task NoInstrumentationForGoogleTypes()
+    {
+        SetEnvironmentVariable(ConfigurationKeys.DataStreamsMonitoring.Enabled, "1");
+        using var telemetry = this.ConfigureTelemetry();
+        using var agent = EnvironmentHelper.GetMockAgent();
+        using (await RunSampleAndWaitForExit(agent, "TimeStamp"))
+        {
+            var spans = agent.WaitForSpans(2);
+            foreach (var span in spans)
+            {
+                span.Tags.Should().NotContain(t => t.Key.StartsWith("schema."));
+            }
+        }
+    }
+
     [Fact]
     [Trait("Category", "EndToEnd")]
     public async Task OnlyEnabledWithDsm()
     {
         SetEnvironmentVariable(ConfigurationKeys.DataStreamsMonitoring.Enabled, "0");
         using var telemetry = this.ConfigureTelemetry();
         using var agent = EnvironmentHelper.GetMockAgent();
-        using (await RunSampleAndWaitForExit(agent, $"{TestPrefix}"))
+        using (await RunSampleAndWaitForExit(agent, "AddressBook"))
         {
             var spans = agent.WaitForSpans(2);
             foreach (var span in spans)

tracer/test/test-applications/integrations/Samples.GoogleProtobuf/Program.cs

@@ -1,34 +1,56 @@
 using System;
 using Google.Protobuf;
+using Google.Protobuf.WellKnownTypes;
 using Sample;
 using static Sample.Person.Types;
 
 namespace Samples.GoogleProtobuf;
 
 internal class Program
 {
-    private static void Main()
+    private const string Usage = "Usage: GoogleProtobuf.exe [AddressBook|TimeStamp]";
+
+    private static void Main(string[] args)
     {
-        // sample data
-        var john = new Person { Name = "John Doe", Email = "[email protected]", Phones = { new PhoneNumber { Number = "12345", Type = PhoneType.Home } } };
-        var jane = new Person { Name = "Jane Doe", Email = "[email protected]", Phones = { new PhoneNumber { Number = "67890", Type = PhoneType.Work }, new PhoneNumber { Number = "54321", Type = PhoneType.Mobile } } };
-        var addressBook = new AddressBook { People = { { 12, john }, { 21, jane } } };
+        if (args.Length == 0)
+        {
+            throw new ArgumentException(Usage);
+        }
 
-        // serialization
+        IMessage outMessage;
+        IMessage inMessage;
+        if (args[0].Equals("TimeStamp", StringComparison.OrdinalIgnoreCase))
+        {
+            // use a raw google protobuf type (shouldn't be instrumented)
+            outMessage = Timestamp.FromDateTime(DateTime.UtcNow);
+            inMessage = new Timestamp();
+        }
+        else if (args[0].Equals("AddressBook", StringComparison.OrdinalIgnoreCase))
+        {
+            // our own sample data
+            var john = new Person { Name = "John Doe", Email = "[email protected]", Phones = { new PhoneNumber { Number = "12345", Type = PhoneType.Home } } };
+            var jane = new Person { Name = "Jane Doe", Email = "[email protected]", Phones = { new PhoneNumber { Number = "67890", Type = PhoneType.Work }, new PhoneNumber { Number = "54321", Type = PhoneType.Mobile } } };
+            outMessage = new AddressBook { People = { { 12, john }, { 21, jane } } };
+            inMessage = new AddressBook();
+        }
+        else
+        {
+            throw new ArgumentException(Usage);
+        }
 
+        // serialization
         byte[] serialized;
         using (SampleHelpers.CreateScope("Ser"))
         {
-            serialized = addressBook.ToByteArray();
+            serialized = outMessage.ToByteArray();
         }
 
         // deserialization
-        var deserialized = new AddressBook();
         using (SampleHelpers.CreateScope("Deser"))
         {
-            deserialized.MergeFrom(serialized);
+            inMessage.MergeFrom(serialized);
         }
 
-        Console.WriteLine(deserialized.People.Count);
+        Console.WriteLine(inMessage.ToString());
     }
 }