AsyncState can be used across different service providers

src/Libraries/Microsoft.Extensions.AsyncState/AsyncState.cs

@@ -14,7 +14,7 @@ internal sealed class AsyncState : IAsyncState
 {
     private static readonly AsyncLocal<AsyncStateHolder> _asyncContextCurrent = new();
     private static readonly ObjectPool<List<object?>> _featuresPool = PoolFactory.CreatePool(new FeaturesPooledPolicy());
-    private int _contextCount;
+    private static int _contextCount;
 
     public void Initialize()
     {
@@ -104,7 +104,9 @@ internal static void EnsureCount(List<object?> features, int count)
         }
     }
 
+#pragma warning disable CA1822 // Member 'ContextCount' does not access instance data and can be marked as static.
     internal int ContextCount => Volatile.Read(ref _contextCount);
+#pragma warning restore CA1822 // Member 'ContextCount' does not access instance data and can be marked as static.
 
     private sealed class AsyncStateHolder
     {

test/Libraries/Microsoft.Extensions.AsyncState.Tests/AssemblyInfo.cs

@@ -0,0 +1,6 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using Xunit;
+
+[assembly: CollectionBehavior(DisableTestParallelization = true)]

test/Libraries/Microsoft.Extensions.AsyncState.Tests/AsyncContextTests.cs

@@ -13,12 +13,13 @@ public class AsyncContextTests
     public async Task CreateAsyncContext_BeforeInitialize()
     {
         var state = new AsyncState();
+        var initialContextCount = state.ContextCount;
         var context1 = new AsyncContext<IThing>(state);
         var context2 = new AsyncContext<IThing>(state);
         var obj1 = new Thing();
         var obj2 = new Thing();
 
-        Assert.Equal(2, state.ContextCount);
+        Assert.Equal(2, state.ContextCount - initialContextCount);
         state.Initialize();
 
         await Task.Run(() => context1.Set(obj1));
@@ -36,11 +37,12 @@ public async Task CreateAsyncContext_AfterInitialize()
     {
         var state = new AsyncState();
         state.Initialize();
+        var initialContextCount = state.ContextCount;
 
         var context1 = new AsyncContext<IThing>(state);
         var context2 = new AsyncContext<IThing>(state);
 
-        Assert.Equal(2, state.ContextCount);
+        Assert.Equal(2, state.ContextCount - initialContextCount);
 
         var obj1 = new Thing();
 
@@ -58,14 +60,14 @@ await Task.Run(() =>
     public async Task CreateAsyncContext_BeforeAndAfterInitialize()
     {
         var state = new AsyncState();
-
+        var initialContextCount = state.ContextCount;
         var context1 = new AsyncContext<IThing>(state);
         state.Initialize();
         var context2 = new AsyncContext<IThing>(state);
         var obj1 = new Thing();
         var obj2 = new Thing();
 
-        Assert.Equal(2, state.ContextCount);
+        Assert.Equal(2, state.ContextCount - initialContextCount);
 
         await Task.Run(() =>
         {
@@ -85,12 +87,13 @@ await Task.Run(() =>
     public async Task Tryget_BeforeAndAfterInitialize()
     {
         var state = new AsyncState();
+        var initialContextCount = state.ContextCount;
         var context1 = new AsyncContext<IThing>(state);
         Assert.False(context1.TryGet(out _));
         state.Initialize();
         var obj1 = new Thing();
 
-        Assert.Equal(1, state.ContextCount);
+        Assert.Equal(1, state.ContextCount - initialContextCount);
 
         await Task.Run(() =>
         {
@@ -110,12 +113,12 @@ await Task.Run(() =>
     public async Task CreateAsyncContextInAsync_AfterInitialize()
     {
         var state = new AsyncState();
-
+        var initialContextCount = state.ContextCount;
         var context1 = new AsyncContext<IThing>(state);
         var obj1 = new Thing();
 
         state.Initialize();
-        Assert.Equal(1, state.ContextCount);
+        Assert.Equal(1, state.ContextCount - initialContextCount);
 
         await Task.Run(async () =>
         {
@@ -181,6 +184,7 @@ public void Reset_DoesNotDisposeObjects()
     public async Task TwoAsyncFlows_WithDiffrentAsyncStates()
     {
         var state = new AsyncState();
+        var initialContextCount = state.ContextCount;
 
         var task1 = Task.Run(async () =>
         {
@@ -215,13 +219,14 @@ await Task.Run(async () =>
         });
 
         await Task.WhenAll(task1, task2);
-        Assert.Equal(2, state.ContextCount);
+        Assert.Equal(2, state.ContextCount - initialContextCount);
     }
 
     [Fact]
     public async Task IndependentAsyncFlows_WithSameAsyncState()
     {
         var state = new AsyncState();
+        var initialContextCount = state.ContextCount;
         var context = new AsyncContext<IThing>(state);
 
         Func<Task?> setAsyncState = async () =>
@@ -251,6 +256,6 @@ await Task.Run(() =>
 
         await Task.WhenAll(tasks);
 
-        Assert.Equal(1, state.ContextCount);
+        Assert.Equal(1, state.ContextCount - initialContextCount);
     }
 }

test/Libraries/Microsoft.Extensions.AsyncState.Tests/AsyncStateTests.cs

@@ -6,6 +6,7 @@
 using System.Diagnostics.CodeAnalysis;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.Extensions.DependencyInjection;
 using Xunit;
 
 namespace Microsoft.Extensions.AsyncState.Test;
@@ -195,15 +196,16 @@ public async Task GettingAsyncContextDoesNotFlowIfExecutionContextDoesNotFlow()
     public void RegisterContextCorrectly()
     {
         var asyncState = new AsyncState();
+        var initialContextCount = asyncState.ContextCount;
 
         var c1 = asyncState.RegisterAsyncContext();
-        Assert.Equal(0, c1.Index);
+        Assert.Equal(0, c1.Index - initialContextCount);
         var c2 = asyncState.RegisterAsyncContext();
-        Assert.Equal(1, c2.Index);
+        Assert.Equal(1, c2.Index - initialContextCount);
         var c3 = asyncState.RegisterAsyncContext();
-        Assert.Equal(2, c3.Index);
+        Assert.Equal(2, c3.Index - initialContextCount);
 
-        Assert.Equal(3, asyncState.ContextCount);
+        Assert.Equal(3, asyncState.ContextCount - initialContextCount);
     }
 
     [Fact]
@@ -229,4 +231,28 @@ public void EnsureCount_WhenCountEqualWithExpected()
         AsyncState.EnsureCount(l, 5);
         Assert.Equal(5, l.Count);
     }
+
+    [Fact]
+    public void AsyncStateCanBeUsedInDifferentServiceProviders()
+    {
+        var expectedValue1 = new Tuple<double>(3.14);
+        var expectedValue2 = new Tuple<int>(42);
+        var spOne = CreateAsyncState<Tuple<double>>();
+        var spTwo = CreateAsyncState<Tuple<int>>();
+        spOne.Set(expectedValue1);
+        spTwo.Set(expectedValue2);
+        var value1 = spOne.Get();
+        var value2 = spTwo.Get();
+
+        Assert.Equal(expectedValue1, value1);
+        Assert.Equal(expectedValue2, value2);
+
+        static IAsyncContext<T> CreateAsyncState<T>()
+            where T : notnull
+        {
+            var services = new ServiceCollection().AddAsyncState().BuildServiceProvider();
+            services.GetRequiredService<IAsyncState>().Initialize();
+            return services.GetRequiredService<IAsyncContext<T>>();
+        }
+    }
 }

test/Libraries/Microsoft.Extensions.AsyncState.Tests/Microsoft.Extensions.AsyncState.Tests.csproj

@@ -4,6 +4,10 @@
     <Description>Unit tests for Microsoft.Extensions.AsyncState.</Description>
   </PropertyGroup>
 
+  <ItemGroup>
+    <PackageReference Include="Microsoft.Extensions.DependencyInjection" />
+  </ItemGroup>
+
   <ItemGroup>
     <ProjectReference Include="..\..\..\src\Libraries\Microsoft.Extensions.AsyncState\Microsoft.Extensions.AsyncState.csproj" ProjectUnderTest="true" />
   </ItemGroup>