From 4649449206edae4992e09419334e736726fd2c8e Mon Sep 17 00:00:00 2001
From: Liudmila Molkova <lmolkova@microsoft.com>
Date: Mon, 15 Apr 2019 18:48:56 -0700
Subject: [PATCH] Call DiagnosticListener.OnActivityImport before starting
 Activity to support sampling better (#9340)

* use OnActivityImport and HttpContext in the extra argument
---
 .../Internal/HostingApplicationDiagnostics.cs |  2 ++
 .../Hosting/test/HostingApplicationTests.cs   | 30 +++++++++++++++++--
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/Hosting/Hosting/src/Internal/HostingApplicationDiagnostics.cs b/src/Hosting/Hosting/src/Internal/HostingApplicationDiagnostics.cs
index 82d825ea2b..bf09fa4dcb 100644
--- a/src/Hosting/Hosting/src/Internal/HostingApplicationDiagnostics.cs
+++ b/src/Hosting/Hosting/src/Internal/HostingApplicationDiagnostics.cs
@@ -273,6 +273,8 @@ namespace Microsoft.AspNetCore.Hosting.Internal
                 }
             }
 
+            _diagnosticListener.OnActivityImport(activity, httpContext);
+
             if (_diagnosticListener.IsEnabled(ActivityStartKey))
             {
                 _diagnosticListener.StartActivity(activity, new { HttpContext = httpContext });
diff --git a/src/Hosting/Hosting/test/HostingApplicationTests.cs b/src/Hosting/Hosting/test/HostingApplicationTests.cs
index 002421a90d..dbb1cb2d23 100644
--- a/src/Hosting/Hosting/test/HostingApplicationTests.cs
+++ b/src/Hosting/Hosting/test/HostingApplicationTests.cs
@@ -11,8 +11,6 @@ using Microsoft.AspNetCore.Hosting.Internal;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Http.Features;
 using Microsoft.Extensions.Logging;
-using Microsoft.Extensions.ObjectPool;
-using Microsoft.Extensions.Options;
 using Moq;
 using Xunit;
 
@@ -304,6 +302,34 @@ namespace Microsoft.AspNetCore.Hosting.Tests
             Assert.Contains(Activity.Current.Baggage, pair => pair.Key == "Key2" && pair.Value == "value2");
         }
 
+        [Fact]
+        public void ActivityOnExportHookIsCalled()
+        {
+            var diagnosticSource = new DiagnosticListener("DummySource");
+            var hostingApplication = CreateApplication(out var features, diagnosticSource: diagnosticSource);
+
+            bool onActivityImportCalled = false;
+            diagnosticSource.Subscribe(
+                observer: new CallbackDiagnosticListener(pair => { }),
+                isEnabled: (s, o, _) => true,
+                onActivityImport: (activity, context) =>
+                {
+                    onActivityImportCalled = true;
+                    Assert.Null(Activity.Current);
+                    Assert.Equal("Microsoft.AspNetCore.Hosting.HttpRequestIn", activity.OperationName);
+                    Assert.NotNull(context);
+                    Assert.IsAssignableFrom<HttpContext>(context);
+
+                    activity.ActivityTraceFlags = ActivityTraceFlags.Recorded;
+                });
+
+            hostingApplication.CreateContext(features);
+
+            Assert.True(onActivityImportCalled);
+            Assert.NotNull(Activity.Current);
+            Assert.True(Activity.Current.Recorded);
+        }
+
 
         private static void AssertProperty<T>(object o, string name)
         {