From 2ac5e2bb8a48dcb3038307338a04914ba708edc2 Mon Sep 17 00:00:00 2001
From: David Fowler <davidfowl@gmail.com>
Date: Mon, 14 Jun 2021 08:18:45 -0700
Subject: [PATCH] Set the async local just before execution (#54133)

* Set the async local just before execution.
- Subscribing to DiagnosticListener.AllListeners replays all created DiagnosticListener instances. Because of this, we need to set the async local just before the execution of the entry point so that we only collect the events that are relevant to the call. Right now, it's also firing with the async local set pre-maturely.
- Wrote a concurrency test to make sure it's safe to instantiate the factory in parallel.
---
 .../src/HostFactoryResolver.cs                     |  8 ++++----
 .../tests/HostFactoryResolverTests.cs              | 23 ++++++++++++++++++++++
 .../NoSpecialEntryPointPatternHangs/Program.cs     |  2 +-
 3 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/src/libraries/Microsoft.Extensions.HostFactoryResolver/src/HostFactoryResolver.cs b/src/libraries/Microsoft.Extensions.HostFactoryResolver/src/HostFactoryResolver.cs
index 6dc54c7..a49a7b7 100644
--- a/src/libraries/Microsoft.Extensions.HostFactoryResolver/src/HostFactoryResolver.cs
+++ b/src/libraries/Microsoft.Extensions.HostFactoryResolver/src/HostFactoryResolver.cs
@@ -194,10 +194,6 @@ namespace Microsoft.Extensions.Hosting
 
             public object CreateHost()
             {
-                // Set the async local to the instance of the HostingListener so we can filter events that
-                // aren't scoped to this execution of the entry point.
-                _currentListener.Value = this;
-
                 using var subscription = DiagnosticListener.AllListeners.Subscribe(this);
 
                 // Kick off the entry point on a new thread so we don't block the current one
@@ -208,6 +204,10 @@ namespace Microsoft.Extensions.Hosting
 
                     try
                     {
+                        // Set the async local to the instance of the HostingListener so we can filter events that
+                        // aren't scoped to this execution of the entry point.
+                        _currentListener.Value = this;
+
                         var parameters = _entryPoint.GetParameters();
                         if (parameters.Length == 0)
                         {
diff --git a/src/libraries/Microsoft.Extensions.HostFactoryResolver/tests/HostFactoryResolverTests.cs b/src/libraries/Microsoft.Extensions.HostFactoryResolver/tests/HostFactoryResolverTests.cs
index 567b73e..8f6cf79 100644
--- a/src/libraries/Microsoft.Extensions.HostFactoryResolver/tests/HostFactoryResolverTests.cs
+++ b/src/libraries/Microsoft.Extensions.HostFactoryResolver/tests/HostFactoryResolverTests.cs
@@ -6,6 +6,7 @@ using System;
 using System.Diagnostics.CodeAnalysis;
 using System.Reflection;
 using System.Threading;
+using System.Threading.Tasks;
 using Xunit;
 
 namespace Microsoft.Extensions.Hosting.Tests
@@ -250,5 +251,27 @@ namespace Microsoft.Extensions.Hosting.Tests
             Assert.NotNull(factory);
             Assert.IsAssignableFrom<IServiceProvider>(factory(Array.Empty<string>()));
         }
+
+        [ConditionalFact(typeof(PlatformDetection), nameof(PlatformDetection.IsThreadingSupported))]
+        [DynamicDependency(DynamicallyAccessedMemberTypes.All, typeof(NoSpecialEntryPointPattern.Program))]
+        public void NoSpecialEntryPointPatternCanRunInParallel()
+        {
+            var factory = HostFactoryResolver.ResolveServiceProviderFactory(typeof(NoSpecialEntryPointPattern.Program).Assembly, s_WaitTimeout);
+            Assert.NotNull(factory);
+
+            var tasks = new Task<IServiceProvider>[30];
+            int index = 0;
+            for (int i = 0; i < tasks.Length; i++)
+            {
+                tasks[index++] = Task.Run(() => factory(Array.Empty<string>()));
+            }
+
+            Task.WaitAll(tasks);
+
+            foreach (var t in tasks)
+            {
+                Assert.IsAssignableFrom<IServiceProvider>(t.Result);
+            }
+        }
     }
 }
diff --git a/src/libraries/Microsoft.Extensions.HostFactoryResolver/tests/NoSpecialEntryPointPatternHangs/Program.cs b/src/libraries/Microsoft.Extensions.HostFactoryResolver/tests/NoSpecialEntryPointPatternHangs/Program.cs
index e8183fc..52ac9d4 100644
--- a/src/libraries/Microsoft.Extensions.HostFactoryResolver/tests/NoSpecialEntryPointPatternHangs/Program.cs
+++ b/src/libraries/Microsoft.Extensions.HostFactoryResolver/tests/NoSpecialEntryPointPatternHangs/Program.cs
@@ -9,7 +9,7 @@ namespace NoSpecialEntryPointPatternHangs
     {
         public static void Main(string[] args)
         {
-            Console.ReadLine();
+            System.Threading.Thread.Sleep(-1);
         }
     }
 }
-- 
2.7.4