Pool the underlying list and dictionary in scopes. (#50463)

- This change pools a set of scopes assuming they are short lived.
- One breaking change is that after disposal, pooled scopes will throw if services are accessed afterwards on the scope.
- Modified test to throw after dispose

Author: davidfowl

src/libraries/Microsoft.Extensions.DependencyInjection/src/ScopePool.cs

@@ -0,0 +1,74 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Diagnostics;
+using System.Threading;
+using Microsoft.Extensions.DependencyInjection.ServiceLookup;
+
+namespace Microsoft.Extensions.DependencyInjection
+{
+    internal class ScopePool
+    {
+        // Modest number to re-use. We only really care about reuse for short lived scopes
+        private const int MaxQueueSize = 128;
+
+        private int _count;
+        private readonly ConcurrentQueue<State> _queue = new();
+
+        public State Rent()
+        {
+            if (_queue.TryDequeue(out State state))
+            {
+                Interlocked.Decrement(ref _count);
+                return state;
+            }
+            return new State(this);
+        }
+
+        public bool Return(State state)
+        {
+            if (Interlocked.Increment(ref _count) > MaxQueueSize)
+            {
+                Interlocked.Decrement(ref _count);
+                return false;
+            }
+
+            state.Clear();
+            _queue.Enqueue(state);
+            return true;
+        }
+
+        public class State
+        {
+            private readonly ScopePool _pool;
+
+            public IDictionary<ServiceCacheKey, object> ResolvedServices { get; }
+            public List<object> Disposables { get; set; }
+
+            public State(ScopePool pool = null)
+            {
+                _pool = pool;
+                // When pool is null, we're in the global scope which doesn't need pooling.
+                // To reduce lock contention for singletons upon resolve we use a concurrent dictionary.
+                ResolvedServices = pool == null ? new ConcurrentDictionary<ServiceCacheKey, object>() : new Dictionary<ServiceCacheKey, object>();
+            }
+
+            internal void Clear()
+            {
+                // This should only get called from the pool
+                Debug.Assert(_pool != null);
+                // REVIEW: Should we trim excess here as well?
+                ResolvedServices.Clear();
+                Disposables?.Clear();
+            }
+
+            public bool Return()
+            {
+                return _pool?.Return(this) ?? false;
+            }
+        }
+    }
+}

src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/CallSiteRuntimeResolver.cs

@@ -90,15 +90,15 @@ protected override object VisitScopeCache(ServiceCallSite callSite, RuntimeResol
         private object VisitCache(ServiceCallSite callSite, RuntimeResolverContext context, ServiceProviderEngineScope serviceProviderEngine, RuntimeResolverLock lockType)
         {
             bool lockTaken = false;
-            IDictionary<ServiceCacheKey, object> resolvedServices = serviceProviderEngine.ResolvedServices;
+            object sync = serviceProviderEngine.Sync;
 
             // Taking locks only once allows us to fork resolution process
             // on another thread without causing the deadlock because we
             // always know that we are going to wait the other thread to finish before
             // releasing the lock
             if ((context.AcquiredLocks & lockType) == 0)
             {
-                Monitor.Enter(resolvedServices, ref lockTaken);
+                Monitor.Enter(sync, ref lockTaken);
             }
 
             try
@@ -109,7 +109,7 @@ private object VisitCache(ServiceCallSite callSite, RuntimeResolverContext conte
             {
                 if (lockTaken)
                 {
-                    Monitor.Exit(resolvedServices);
+                    Monitor.Exit(sync);
                 }
             }
         }
@@ -123,7 +123,7 @@ private object ResolveService(ServiceCallSite callSite, RuntimeResolverContext c
             // For scoped: takes a dictionary as both a resolution lock and a dictionary access lock.
             Debug.Assert(
                 (lockType == RuntimeResolverLock.Root && resolvedServices is ConcurrentDictionary<ServiceCacheKey, object>) ||
-                (lockType == RuntimeResolverLock.Scope && Monitor.IsEntered(resolvedServices)));
+                (lockType == RuntimeResolverLock.Scope && Monitor.IsEntered(serviceProviderEngine.Sync)));
 
             if (resolvedServices.TryGetValue(callSite.Cache.Key, out object resolved))
             {

src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/Expressions/ExpressionResolverBuilder.cs

@@ -26,12 +26,19 @@ internal sealed class ExpressionResolverBuilder : CallSiteVisitor<object, Expres
         private static readonly ParameterExpression ScopeParameter = Expression.Parameter(typeof(ServiceProviderEngineScope));
 
         private static readonly ParameterExpression ResolvedServices = Expression.Variable(typeof(IDictionary<ServiceCacheKey, object>), ScopeParameter.Name + "resolvedServices");
+        private static readonly ParameterExpression Sync = Expression.Variable(typeof(object), ScopeParameter.Name + "sync");
         private static readonly BinaryExpression ResolvedServicesVariableAssignment =
             Expression.Assign(ResolvedServices,
                 Expression.Property(
                     ScopeParameter,
                     typeof(ServiceProviderEngineScope).GetProperty(nameof(ServiceProviderEngineScope.ResolvedServices), BindingFlags.Instance | BindingFlags.NonPublic)));
 
+        private static readonly BinaryExpression SyncVariableAssignment =
+            Expression.Assign(Sync,
+                Expression.Property(
+                    ScopeParameter,
+                    typeof(ServiceProviderEngineScope).GetProperty(nameof(ServiceProviderEngineScope.Sync), BindingFlags.Instance | BindingFlags.NonPublic)));
+
         private static readonly ParameterExpression CaptureDisposableParameter = Expression.Parameter(typeof(object));
         private static readonly LambdaExpression CaptureDisposable = Expression.Lambda(
                     Expression.Call(ScopeParameter, CaptureDisposableMethodInfo, CaptureDisposableParameter),
@@ -96,8 +103,9 @@ private Expression<Func<ServiceProviderEngineScope, object>> BuildExpression(Ser
             {
                 return Expression.Lambda<Func<ServiceProviderEngineScope, object>>(
                     Expression.Block(
-                        new[] { ResolvedServices },
+                        new[] { ResolvedServices, Sync },
                         ResolvedServicesVariableAssignment,
+                        SyncVariableAssignment,
                         BuildScopedExpression(callSite)),
                     ScopeParameter);
             }
@@ -213,7 +221,6 @@ protected override Expression VisitScopeCache(ServiceCallSite callSite, object c
         // Move off the main stack
         private Expression BuildScopedExpression(ServiceCallSite callSite)
         {
-
             ConstantExpression keyExpression = Expression.Constant(
                 callSite.Cache.Key,
                 typeof(ServiceCacheKey));
@@ -257,9 +264,10 @@ private Expression BuildScopedExpression(ServiceCallSite callSite)
             // The C# compiler would copy the lock object to guard against mutation.
             // We don't, since we know the lock object is readonly.
             ParameterExpression lockWasTaken = Expression.Variable(typeof(bool), "lockWasTaken");
+            ParameterExpression sync = Sync;
 
-            MethodCallExpression monitorEnter = Expression.Call(MonitorEnterMethodInfo, resolvedServices, lockWasTaken);
-            MethodCallExpression monitorExit = Expression.Call(MonitorExitMethodInfo, resolvedServices);
+            MethodCallExpression monitorEnter = Expression.Call(MonitorEnterMethodInfo, sync, lockWasTaken);
+            MethodCallExpression monitorExit = Expression.Call(MonitorExitMethodInfo, sync);
 
             BlockExpression tryBody = Expression.Block(monitorEnter, blockExpression);
             ConditionalExpression finallyBody = Expression.IfThen(lockWasTaken, monitorExit);

src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/ILEmit/ILEmitResolverBuilder.cs

@@ -14,6 +14,9 @@ internal sealed class ILEmitResolverBuilder : CallSiteVisitor<ILEmitResolverBuil
         private static readonly MethodInfo ResolvedServicesGetter = typeof(ServiceProviderEngineScope).GetProperty(
             nameof(ServiceProviderEngineScope.ResolvedServices), BindingFlags.Instance | BindingFlags.NonPublic).GetMethod;
 
+        private static readonly MethodInfo ScopeLockGetter = typeof(ServiceProviderEngineScope).GetProperty(
+            nameof(ServiceProviderEngineScope.Sync), BindingFlags.Instance | BindingFlags.NonPublic).GetMethod;
+
         private static readonly FieldInfo FactoriesField = typeof(ILEmitResolverBuilderRuntimeContext).GetField(nameof(ILEmitResolverBuilderRuntimeContext.Factories));
         private static readonly FieldInfo ConstantsField = typeof(ILEmitResolverBuilderRuntimeContext).GetField(nameof(ILEmitResolverBuilderRuntimeContext.Constants));
         private static readonly MethodInfo GetTypeFromHandleMethod = typeof(Type).GetMethod(nameof(Type.GetTypeFromHandle));
@@ -319,6 +322,7 @@ private ILEmitResolverBuilderRuntimeContext GenerateMethodBody(ServiceCallSite c
             {
                 LocalBuilder cacheKeyLocal = context.Generator.DeclareLocal(typeof(ServiceCacheKey));
                 LocalBuilder resolvedServicesLocal = context.Generator.DeclareLocal(typeof(IDictionary<ServiceCacheKey, object>));
+                LocalBuilder syncLocal = context.Generator.DeclareLocal(typeof(object));
                 LocalBuilder lockTakenLocal = context.Generator.DeclareLocal(typeof(bool));
                 LocalBuilder resultLocal = context.Generator.DeclareLocal(typeof(object));
 
@@ -339,8 +343,15 @@ private ILEmitResolverBuilderRuntimeContext GenerateMethodBody(ServiceCallSite c
                 // Store resolved services
                 Stloc(context.Generator, resolvedServicesLocal.LocalIndex);
 
-                // Load resolvedServices
-                Ldloc(context.Generator, resolvedServicesLocal.LocalIndex);
+                // scope
+                context.Generator.Emit(OpCodes.Ldarg_1);
+                // .Sync
+                context.Generator.Emit(OpCodes.Callvirt, ScopeLockGetter);
+                // Store syncLocal
+                Stloc(context.Generator, syncLocal.LocalIndex);
+
+                // Load syncLocal
+                Ldloc(context.Generator, syncLocal.LocalIndex);
                 // Load address of lockTaken
                 context.Generator.Emit(OpCodes.Ldloca_S, lockTakenLocal.LocalIndex);
                 // Monitor.Enter
@@ -388,8 +399,8 @@ private ILEmitResolverBuilderRuntimeContext GenerateMethodBody(ServiceCallSite c
                 Ldloc(context.Generator, lockTakenLocal.LocalIndex);
                 // return if not
                 context.Generator.Emit(OpCodes.Brfalse, returnLabel);
-                // Load resolvedServices
-                Ldloc(context.Generator, resolvedServicesLocal.LocalIndex);
+                // Load syncLocal
+                Ldloc(context.Generator, syncLocal.LocalIndex);
                 // Monitor.Exit
                 context.Generator.Emit(OpCodes.Call, ExpressionResolverBuilder.MonitorExitMethodInfo);
 

src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/ServiceProviderEngine.cs

@@ -25,8 +25,11 @@ protected ServiceProviderEngine(IEnumerable<ServiceDescriptor> serviceDescriptor
             CallSiteFactory.Add(typeof(IServiceProvider), new ServiceProviderCallSite());
             CallSiteFactory.Add(typeof(IServiceScopeFactory), new ServiceScopeFactoryCallSite());
             RealizedServices = new ConcurrentDictionary<Type, Func<ServiceProviderEngineScope, object>>();
+            ScopePool = new ScopePool();
         }
 
+        internal ScopePool ScopePool { get; }
+
         internal ConcurrentDictionary<Type, Func<ServiceProviderEngineScope, object>> RealizedServices { get; }
 
         internal CallSiteFactory CallSiteFactory { get; }

src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/ServiceProviderEngineScope.cs

@@ -2,9 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
-using System.Collections.Concurrent;
 using System.Collections.Generic;
-using System.Diagnostics;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Internal;
 
@@ -15,20 +13,23 @@ internal sealed class ServiceProviderEngineScope : IServiceScope, IServiceProvid
         // For testing only
         internal Action<object> _captureDisposableCallback;
 
-        private List<object> _disposables;
-
         private bool _disposed;
-        private readonly object _disposelock = new object();
+        private ScopePool.State _state;
+
+        // This lock protects state on the scope, in particular, for the root scope, it protects
+        // the list of disposable entries only, since ResolvedServices is a concurrent dictionary.
+        // For other scopes, it protects ResolvedServices and the list of disposables
+        private readonly object _scopeLock = new object();
 
         public ServiceProviderEngineScope(ServiceProviderEngine engine, bool isRoot = false)
         {
             Engine = engine;
-
-            // To reduce lock contention for singletons upon resolve we use a concurrent dictionary.
-            ResolvedServices = isRoot ? new ConcurrentDictionary<ServiceCacheKey, object>() : new Dictionary<ServiceCacheKey, object>();
+            _state = isRoot ? new ScopePool.State() : engine.ScopePool.Rent();
         }
 
-        internal IDictionary<ServiceCacheKey, object> ResolvedServices { get; }
+        internal IDictionary<ServiceCacheKey, object> ResolvedServices => _state?.ResolvedServices ?? ScopeDisposed();
+
+        internal object Sync => _scopeLock;
 
         public ServiceProviderEngine Engine { get; }
 
@@ -53,7 +54,7 @@ internal object CaptureDisposable(object service)
                 return service;
             }
 
-            lock (_disposelock)
+            lock (_scopeLock)
             {
                 if (_disposed)
                 {
@@ -66,16 +67,15 @@ internal object CaptureDisposable(object service)
                         // sync over async, for the rare case that an object only implements IAsyncDisposable and may end up starving the thread pool.
                         Task.Run(() => ((IAsyncDisposable)service).DisposeAsync().AsTask()).GetAwaiter().GetResult();
                     }
+
                     ThrowHelper.ThrowObjectDisposedException();
                 }
 
-                if (_disposables == null)
-                {
-                    _disposables = new List<object>();
-                }
+                _state.Disposables ??= new List<object>();
 
-                _disposables.Add(service);
+                _state.Disposables.Add(service);
             }
+
             return service;
         }
 
@@ -97,6 +97,8 @@ public void Dispose()
                     }
                 }
             }
+
+            ClearState();
         }
 
         public ValueTask DisposeAsync()
@@ -115,7 +117,7 @@ public ValueTask DisposeAsync()
                             ValueTask vt = asyncDisposable.DisposeAsync();
                             if (!vt.IsCompletedSuccessfully)
                             {
-                                return Await(i, vt, toDispose);
+                                return Await(this, i, vt, toDispose);
                             }
 
                             // If its a IValueTaskSource backed ValueTask,
@@ -134,9 +136,11 @@ public ValueTask DisposeAsync()
                 }
             }
 
+            ClearState();
+
             return default;
 
-            static async ValueTask Await(int i, ValueTask vt, List<object> toDispose)
+            static async ValueTask Await(ServiceProviderEngineScope scope, int i, ValueTask vt, List<object> toDispose)
             {
                 await vt.ConfigureAwait(false);
                 // vt is acting on the disposable at index i,
@@ -155,30 +159,52 @@ static async ValueTask Await(int i, ValueTask vt, List<object> toDispose)
                         ((IDisposable)disposable).Dispose();
                     }
                 }
+
+                scope.ClearState();
+            }
+        }
+
+        private IDictionary<ServiceCacheKey, object> ScopeDisposed()
+        {
+            ThrowHelper.ThrowObjectDisposedException();
+            return null;
+        }
+
+        private void ClearState()
+        {
+            // We lock here since ResolvedServices is always accessed in the scope lock, this means we'll never
+            // try to return to the pool while somebody is trying to access ResolvedServices.
+            lock (_scopeLock)
+            {
+                // ResolvedServices is never cleared for singletons because there might be a compilation running in background
+                // trying to get a cached singleton service. If it doesn't find it
+                // it will try to create a new one which will result in an ObjectDisposedException.
+
+                // Dispose the state, which will end up attempting to return the state pool.
+                // This will return false if the pool is full or if this state object is the root scope
+                if (_state.Return())
+                {
+                    _state = null;
+                }
             }
         }
 
         private List<object> BeginDispose()
         {
-            List<object> toDispose;
-            lock (_disposelock)
+            lock (_scopeLock)
             {
                 if (_disposed)
                 {
                     return null;
                 }
 
+                // We've transitioned to the disposed state, so future calls to
+                // CaptureDisposable will immediately dispose the object.
+                // No further changes to _state.Disposables, are allowed.
                 _disposed = true;
-                toDispose = _disposables;
-                _disposables = null;
 
-                // Not clearing ResolvedServices here because there might be a compilation running in background
-                // trying to get a cached singleton service instance and if it won't find
-                // it it will try to create a new one tripping the Debug.Assert in CaptureDisposable
-                // and leaking a Disposable object in Release mode
+                return _state.Disposables;
             }
-
-            return toDispose;
         }
     }
 }