From 41a3371e2ab6ab0cc1a68cbc3e5e189bc03d9f5c Mon Sep 17 00:00:00 2001
From: Utkarsh Umesan Pillai <66651184+utpilla@users.noreply.github.com>
Date: Sat, 13 Mar 2021 10:23:13 -0800
Subject: [PATCH] Optimize gRPC client instrumentation for suppress
 instrumentation (#1904)

* Optimize gRPC client instrumentation for SuppressInstrumentation

* Udated unit tests
---
 .../GrpcClientDiagnosticListener.cs           | 28 +++++++++++++++----
 .../GrpcTests.client.cs                       |  6 ++++
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/src/OpenTelemetry.Instrumentation.GrpcNetClient/Implementation/GrpcClientDiagnosticListener.cs b/src/OpenTelemetry.Instrumentation.GrpcNetClient/Implementation/GrpcClientDiagnosticListener.cs
index d741cf05100..1faaff01537 100644
--- a/src/OpenTelemetry.Instrumentation.GrpcNetClient/Implementation/GrpcClientDiagnosticListener.cs
+++ b/src/OpenTelemetry.Instrumentation.GrpcNetClient/Implementation/GrpcClientDiagnosticListener.cs
@@ -43,6 +43,22 @@ public GrpcClientDiagnosticListener(GrpcClientInstrumentationOptions options)
 
         public override void OnStartActivity(Activity activity, object payload)
         {
+            // The overall flow of what GrpcClient library does is as below:
+            // Activity.Start()
+            // DiagnosticSource.WriteEvent("Start", payload)
+            // DiagnosticSource.WriteEvent("Stop", payload)
+            // Activity.Stop()
+
+            // This method is in the WriteEvent("Start", payload) path.
+            // By this time, samplers have already run and
+            // activity.IsAllDataRequested populated accordingly.
+
+            if (Sdk.SuppressInstrumentation)
+            {
+                return;
+            }
+
+            // Ensure context propagation irrespective of sampling decision
             if (!this.startRequestFetcher.TryFetch(payload, out HttpRequestMessage request) || request == null)
             {
                 GrpcInstrumentationEventSource.Log.NullPayload(nameof(GrpcClientDiagnosticListener), nameof(this.OnStartActivity));
@@ -78,15 +94,15 @@ public override void OnStartActivity(Activity activity, object payload)
                     HttpRequestMessageContextPropagation.HeaderValueSetter);
             }
 
-            var grpcMethod = GrpcTagHelper.GetGrpcMethodFromActivity(activity);
+            if (activity.IsAllDataRequested)
+            {
+                ActivityInstrumentationHelper.SetActivitySourceProperty(activity, ActivitySource);
+                ActivityInstrumentationHelper.SetKindProperty(activity, ActivityKind.Client);
 
-            activity.DisplayName = grpcMethod?.Trim('/');
+                var grpcMethod = GrpcTagHelper.GetGrpcMethodFromActivity(activity);
 
-            ActivityInstrumentationHelper.SetActivitySourceProperty(activity, ActivitySource);
-            ActivityInstrumentationHelper.SetKindProperty(activity, ActivityKind.Client);
+                activity.DisplayName = grpcMethod?.Trim('/');
 
-            if (activity.IsAllDataRequested)
-            {
                 activity.SetTag(SemanticConventions.AttributeRpcSystem, GrpcTagHelper.RpcSystemGrpc);
 
                 if (GrpcTagHelper.TryParseRpcServiceAndRpcMethod(grpcMethod, out var rpcService, out var rpcMethod))
diff --git a/test/OpenTelemetry.Instrumentation.Grpc.Tests/GrpcTests.client.cs b/test/OpenTelemetry.Instrumentation.Grpc.Tests/GrpcTests.client.cs
index 6d617b98098..a735df41352 100644
--- a/test/OpenTelemetry.Instrumentation.Grpc.Tests/GrpcTests.client.cs
+++ b/test/OpenTelemetry.Instrumentation.Grpc.Tests/GrpcTests.client.cs
@@ -23,6 +23,7 @@
 using Microsoft.AspNetCore.Http;
 using Moq;
 using OpenTelemetry.Instrumentation.GrpcNetClient;
+using OpenTelemetry.Instrumentation.GrpcNetClient.Implementation;
 using OpenTelemetry.Trace;
 using Xunit;
 
@@ -143,6 +144,7 @@ public void GrpcAndHttpClientInstrumentationIsInvoked(bool shouldEnrich)
 
             ValidateGrpcActivity(grpcSpan);
             Assert.Equal($"greet.Greeter/SayHello", grpcSpan.DisplayName);
+            Assert.Equal(0, grpcSpan.GetTagValue(SemanticConventions.AttributeRpcGrpcStatusCode));
             Assert.Equal($"HTTP POST", httpSpan.DisplayName);
             Assert.Equal(grpcSpan.SpanId, httpSpan.ParentSpanId);
         }
@@ -274,6 +276,7 @@ public void GrpcPropagatesContextWithSuppressInstrumentation()
             Assert.Equal($"greet.Greeter/SayHello", serverActivity.DisplayName);
             Assert.Equal(clientActivity.TraceId, serverActivity.TraceId);
             Assert.Equal(clientActivity.SpanId, serverActivity.ParentSpanId);
+            Assert.Equal(0, clientActivity.GetTagValue(SemanticConventions.AttributeRpcGrpcStatusCode));
             Assert.Contains("item1=value1", serverActivity.GetCustomProperty("BaggageString") as string);
         }
 
@@ -286,11 +289,14 @@ public void Grpc_BadArgs()
 
         private static void ValidateGrpcActivity(Activity activityToValidate)
         {
+            Assert.Equal(GrpcClientDiagnosticListener.ActivitySourceName, activityToValidate.Source.Name);
+            Assert.Equal(GrpcClientDiagnosticListener.Version.ToString(), activityToValidate.Source.Version);
             Assert.Equal(ActivityKind.Client, activityToValidate.Kind);
         }
 
         private static void ActivityEnrichment(Activity activity, string method, object obj)
         {
+            Assert.True(activity.IsAllDataRequested);
             switch (method)
             {
                 case "OnStartActivity":