From 192fc215c11395bff389dee5e391189e25d75e72 Mon Sep 17 00:00:00 2001
From: Utkarsh Umesan Pillai <utpilla@microsoft.com>
Date: Fri, 12 Mar 2021 12:40:53 -0800
Subject: [PATCH] Optimize AspNet instrumentation for SuppressInstrumentation

---
 .../Implementation/HttpInListener.cs          | 60 ++++++++++++-------
 .../HttpInListenerTests.cs                    |  2 +
 2 files changed, 40 insertions(+), 22 deletions(-)

diff --git a/src/OpenTelemetry.Instrumentation.AspNet/Implementation/HttpInListener.cs b/src/OpenTelemetry.Instrumentation.AspNet/Implementation/HttpInListener.cs
index 61f6538b5fa..458f814178b 100644
--- a/src/OpenTelemetry.Instrumentation.AspNet/Implementation/HttpInListener.cs
+++ b/src/OpenTelemetry.Instrumentation.AspNet/Implementation/HttpInListener.cs
@@ -47,26 +47,26 @@ public HttpInListener(string name, AspNetInstrumentationOptions options)
         [System.Diagnostics.CodeAnalysis.SuppressMessage("Reliability", "CA2000:Dispose objects before losing scope", Justification = "Activity is retrieved from Activity.Current later and disposed.")]
         public override void OnStartActivity(Activity activity, object payload)
         {
-            var context = HttpContext.Current;
-            if (context == null)
+            // The overall flow of what AspNet 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)
             {
-                AspNetInstrumentationEventSource.Log.NullPayload(nameof(HttpInListener), nameof(this.OnStartActivity));
                 return;
             }
 
-            try
-            {
-                if (this.options.Filter?.Invoke(context) == false)
-                {
-                    AspNetInstrumentationEventSource.Log.RequestIsFilteredOut(activity.OperationName);
-                    activity.IsAllDataRequested = false;
-                    return;
-                }
-            }
-            catch (Exception ex)
+            // Ensure context extraction irrespective of sampling decision
+            var context = HttpContext.Current;
+            if (context == null)
             {
-                AspNetInstrumentationEventSource.Log.RequestFilterException(ex);
-                activity.IsAllDataRequested = false;
+                AspNetInstrumentationEventSource.Log.NullPayload(nameof(HttpInListener), nameof(this.OnStartActivity));
                 return;
             }
 
@@ -108,15 +108,31 @@ public override void OnStartActivity(Activity activity, object payload)
                 }
             }
 
-            // see the spec https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/data-semantic-conventions.md
-            var path = requestValues.Path;
-            activity.DisplayName = path;
-
-            ActivityInstrumentationHelper.SetActivitySourceProperty(activity, ActivitySource);
-            ActivityInstrumentationHelper.SetKindProperty(activity, ActivityKind.Server);
-
             if (activity.IsAllDataRequested)
             {
+                try
+                {
+                    if (this.options.Filter?.Invoke(context) == false)
+                    {
+                        AspNetInstrumentationEventSource.Log.RequestIsFilteredOut(activity.OperationName);
+                        activity.IsAllDataRequested = false;
+                        return;
+                    }
+                }
+                catch (Exception ex)
+                {
+                    AspNetInstrumentationEventSource.Log.RequestFilterException(ex);
+                    activity.IsAllDataRequested = false;
+                    return;
+                }
+
+                ActivityInstrumentationHelper.SetActivitySourceProperty(activity, ActivitySource);
+                ActivityInstrumentationHelper.SetKindProperty(activity, ActivityKind.Server);
+
+                // see the spec https://github.com/open-telemetry/opentelemetry-specification/blob/main/specification/data-semantic-conventions.md
+                var path = requestValues.Path;
+                activity.DisplayName = path;
+
                 if (request.Url.Port == 80 || request.Url.Port == 443)
                 {
                     activity.SetTag(SemanticConventions.AttributeHttpHost, request.Url.Host);
diff --git a/test/OpenTelemetry.Instrumentation.AspNet.Tests/HttpInListenerTests.cs b/test/OpenTelemetry.Instrumentation.AspNet.Tests/HttpInListenerTests.cs
index 1cb850d4ce2..68ee629e1b0 100644
--- a/test/OpenTelemetry.Instrumentation.AspNet.Tests/HttpInListenerTests.cs
+++ b/test/OpenTelemetry.Instrumentation.AspNet.Tests/HttpInListenerTests.cs
@@ -151,6 +151,7 @@ public void AspNetRequestsAreCollectedSuccessfully(
                 {
                     options.Filter = httpContext =>
                     {
+                        Assert.True(Activity.Current.IsAllDataRequested);
                         if (string.IsNullOrEmpty(filter))
                         {
                             return true;
@@ -341,6 +342,7 @@ private static Action<Activity, string, object> GetEnrichmentAction(Status statu
 
             enrichAction = (activity, method, obj) =>
             {
+                Assert.True(activity.IsAllDataRequested);
                 switch (method)
                 {
                     case "OnStartActivity":