Added notes to the docs about thread safety of implementations of IAsyncState, IAsyncContext<T> and IAsyncLocalContext<T>

src/Libraries/Microsoft.AspNetCore.AsyncState/AsyncStateHttpContextExtensions.cs

@@ -16,6 +16,7 @@ public static class AsyncStateHttpContextExtensions
     /// <summary>
     /// Adds default implementations for <see cref="IAsyncState"/>, <see cref="IAsyncContext{T}"/>, and <see cref="IAsyncLocalContext{T}"/> services,
     /// scoped to the lifetime of <see cref="AspNetCore.Http.HttpContext"/> instances.
+    /// Please note that implementations of these interfaces are not thread safe.
     /// </summary>
     /// <param name="services">The <see cref="IServiceCollection"/> to add the service to.</param>
     /// <returns>The value of <paramref name="services"/>.</returns>

src/Libraries/Microsoft.AspNetCore.AsyncState/README.md

@@ -4,6 +4,9 @@ This provides the ability to store and retrieve state objects that flow with the
 
 The lifetime of the shared data is controlled automatically and will be the same as of `HttpContext`.
 
+> [!NOTE]
+> Please note, the implementation of `IAsyncContext<T>` provided by this library is not thread-safe.
+
 ## Install the package
 
 From the command-line:

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

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Collections.Generic;
 using System.Threading;
 using Microsoft.Extensions.ObjectPool;
 using Microsoft.Shared.Diagnostics;
@@ -12,7 +13,7 @@ namespace Microsoft.Extensions.AsyncState;
 internal sealed class AsyncState : IAsyncState
 {
     private static readonly AsyncLocal<AsyncStateHolder> _asyncContextCurrent = new();
-    private static readonly ObjectPool<Features> _featuresPool = PoolFactory.CreatePool(new FeaturesPooledPolicy());
+    private static readonly ObjectPool<List<object?>> _featuresPool = PoolFactory.CreatePool(new FeaturesPooledPolicy());
     private int _contextCount;
 
     public void Initialize()
@@ -21,12 +22,12 @@ public void Initialize()
 
         // Use an object indirection to hold the AsyncContext in the AsyncLocal,
         // so it can be cleared in all ExecutionContexts when its cleared.
-        var asyncStateHolder = new AsyncStateHolder
+        var features = new AsyncStateHolder
         {
             Features = _featuresPool.Get()
         };
 
-        _asyncContextCurrent.Value = asyncStateHolder;
+        _asyncContextCurrent.Value = features;
     }
 
     public void Reset()
@@ -59,7 +60,9 @@ public bool TryGet(AsyncStateToken token, out object? value)
             return false;
         }
 
-        value = _asyncContextCurrent.Value.Features.Get(token.Index);
+        EnsureCount(_asyncContextCurrent.Value.Features, token.Index + 1);
+
+        value = _asyncContextCurrent.Value.Features[token.Index];
         return true;
     }
 
@@ -83,14 +86,28 @@ public void Set(AsyncStateToken token, object? value)
             Throw.InvalidOperationException("Context is not initialized");
         }
 
-        _asyncContextCurrent.Value.Features.Set(token.Index, value);
+        EnsureCount(_asyncContextCurrent.Value.Features, token.Index + 1);
+
+        _asyncContextCurrent.Value.Features[token.Index] = value;
+    }
+
+    internal static void EnsureCount(List<object?> features, int count)
+    {
+#if NET6_0_OR_GREATER
+        features.EnsureCapacity(count);
+#endif
+        var difference = count - features.Count;
+
+        for (int i = 0; i < difference; i++)
+        {
+            features.Add(null);
+        }
     }
 
     internal int ContextCount => Volatile.Read(ref _contextCount);
 
     private sealed class AsyncStateHolder
     {
-        public Features? Features { get; set; }
+        public List<object?>? Features { get; set; }
     }
-
 }

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

@@ -15,6 +15,7 @@ public static class AsyncStateExtensions
 {
     /// <summary>
     /// Adds default implementations for <see cref="IAsyncState"/>, <see cref="IAsyncContext{T}"/>, and <see cref="IAsyncLocalContext{T}"/> services.
+    /// Please note that implementations of these interfaces are not thread safe.
     /// </summary>
     /// <param name="services">The dependency injection container to add the implementations to.</param>
     /// <returns>The value of <paramref name="services"/>.</returns>

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

@@ -1,22 +1,27 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Collections.Generic;
 using Microsoft.Extensions.ObjectPool;
 
 namespace Microsoft.Extensions.AsyncState;
 
-internal sealed class FeaturesPooledPolicy : IPooledObjectPolicy<Features>
+internal sealed class FeaturesPooledPolicy : IPooledObjectPolicy<List<object?>>
 {
     /// <inheritdoc/>
-    public Features Create()
+    public List<object?> Create()
     {
-        return new Features();
+        return [];
     }
 
     /// <inheritdoc/>
-    public bool Return(Features obj)
+    public bool Return(List<object?> obj)
     {
-        obj.Clear();
+        for (int i = 0; i < obj.Count; i++)
+        {
+            obj[i] = null;
+        }
+
         return true;
     }
 }

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

@@ -8,6 +8,7 @@ namespace Microsoft.Extensions.AsyncState;
 
 /// <summary>
 /// Provides access to the current async context.
+/// Some implementations of this interface may not be thread safe.
 /// </summary>
 /// <typeparam name="T">The type of the asynchronous state.</typeparam>
 [SuppressMessage("Naming", "CA1716:Identifiers should not match keywords", Justification = "Getter and setter throw exceptions.")]

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

@@ -7,6 +7,7 @@ namespace Microsoft.Extensions.AsyncState;
 
 /// <summary>
 /// Provides access to the current async context stored outside of the HTTP pipeline.
+/// Some implementations of this interface may not be thread safe.
 /// </summary>
 /// <typeparam name="T">The type of the asynchronous state.</typeparam>
 /// <remarks>This type is intended for internal use. Use <see cref="IAsyncContext{T}"/> instead.</remarks>

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

@@ -9,6 +9,7 @@ namespace Microsoft.Extensions.AsyncState;
 
 /// <summary>
 /// Encapsulates all information within the asynchronous flow in an <see cref="AsyncLocal{T}"/> variable.
+/// Some implementations of this interface may not be thread safe.
 /// </summary>
 [SuppressMessage("Naming", "CA1716:Identifiers should not match keywords", Justification = "Getter and setter throw exceptions.")]
 public interface IAsyncState

src/Libraries/Microsoft.Extensions.AsyncState/README.md

@@ -1,12 +1,14 @@
 # Microsoft.Extensions.AsyncState
 
 This provides the ability to store and retrieve objects that flow with the current asynchronous context.
-
 It has a few advantages over using the [`AsyncLocal<T>`](https://learn.microsoft.com/dotnet/api/system.threading.asynclocal-1) class directly:
 - By abstracting the way the ambient data is stored we can use more optimized implementations, for instance when using ASP.NET Core, without exposing these components.
 - Improves the performance by minimizing the number of `AsyncLocal<T>` instances required when multiple objects are shared.
 - Provides a way to manage the lifetime of the ambient data objects.
 
+> [!NOTE]
+> Please note, the implementations of `IAsyncState` and `IAsyncContext<T>` are not thread-safe.
+
 ## Install the package
 
 From the command-line:

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

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Collections.Generic;
 using System.Diagnostics.CodeAnalysis;
 using System.Threading;
 using System.Threading.Tasks;
@@ -204,4 +205,28 @@ public void RegisterContextCorrectly()
 
         Assert.Equal(3, asyncState.ContextCount);
     }
+
+    [Fact]
+    public void EnsureCount_IncreasesCountCorrectly()
+    {
+        var l = new List<object?>();
+        AsyncState.EnsureCount(l, 5);
+        Assert.Equal(5, l.Count);
+    }
+
+    [Fact]
+    public void EnsureCount_WhenCountLessThanExpected()
+    {
+        var l = new List<object?>(new object?[5]);
+        AsyncState.EnsureCount(l, 2);
+        Assert.Equal(5, l.Count);
+    }
+
+    [Fact]
+    public void EnsureCount_WhenCountEqualWithExpected()
+    {
+        var l = new List<object?>(new object?[5]);
+        AsyncState.EnsureCount(l, 5);
+        Assert.Equal(5, l.Count);
+    }
 }

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

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Collections.Generic;
 using Xunit;
 
 namespace Microsoft.Extensions.AsyncState.Test;
@@ -12,22 +13,20 @@ public void Return_ShouldBeTrue()
     {
         var policy = new FeaturesPooledPolicy();
 
-        Assert.True(policy.Return(new Features()));
+        Assert.True(policy.Return([]));
     }
 
     [Fact]
     public void Return_ShouldNullList()
     {
         var policy = new FeaturesPooledPolicy();
 
-        var features = policy.Create();
-        features.Set(0, string.Empty);
-        features.Set(1, Array.Empty<int>());
-        features.Set(2, new object());
+        var list = policy.Create();
+        list.Add(string.Empty);
+        list.Add(Array.Empty<int>());
+        list.Add(new object());
 
-        Assert.True(policy.Return(features));
-        Assert.Null(features.Get(0));
-        Assert.Null(features.Get(1));
-        Assert.Null(features.Get(2));
+        Assert.True(policy.Return(list));
+        Assert.All(list, el => Assert.Null(el));
     }
 }