Streamlined caching for collection registrations.

Author: tillig

src/Autofac/Features/Collections/CollectionRegistrationSource.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT License. See LICENSE in the project root for license information.
 
 using System.Collections;
+using System.Collections.Concurrent;
 using System.Linq.Expressions;
 using System.Reflection;
 using System.Runtime.CompilerServices;
@@ -191,35 +192,21 @@ private static Func<int, IList> GenerateArrayFactory(Type elementType)
         var cache = GetKeyedRegistrationCache(registry);
         var registryStamp = GetRegistryStamp(registry);
 
-        lock (cache)
-        {
-            if (cache.TryGetValue(elementType, out var entry) && entry.Stamp == registryStamp)
-            {
-                return entry.Registrations;
-            }
-
-            var registrations = BuildKeyedRegistrationList(registry, elementType);
-            cache[elementType] = new KeyedRegistrationCacheEntry(registryStamp, registrations);
-            return registrations;
-        }
+        return cache.GetOrUpdate(
+            elementType,
+            registryStamp,
+            () => BuildKeyedRegistrationList(registry, elementType));
     }
 
     private static IReadOnlyList<ServiceRegistration> GetStandardRegistrationList(IComponentRegistry registry, Service elementTypeService)
     {
         var cache = GetStandardRegistrationCache(registry);
         var registryStamp = GetRegistryStamp(registry);
 
-        lock (cache)
-        {
-            if (cache.TryGetValue(elementTypeService, out var entry) && entry.Stamp == registryStamp)
-            {
-                return entry.Registrations;
-            }
-
-            var registrations = BuildStandardRegistrationList(registry, elementTypeService);
-            cache[elementTypeService] = new StandardRegistrationCacheEntry(registryStamp, registrations);
-            return registrations;
-        }
+        return cache.GetOrUpdate(
+            elementTypeService,
+            registryStamp,
+            () => BuildStandardRegistrationList(registry, elementTypeService));
     }
 
     private static IReadOnlyList<ServiceRegistration> BuildStandardRegistrationList(IComponentRegistry registry, Service elementTypeService)
@@ -358,14 +345,19 @@ private static IList BuildCollection(
         return output;
     }
 
-    private static Dictionary<Type, KeyedRegistrationCacheEntry> GetKeyedRegistrationCache(IComponentRegistry registry)
+    private static KeyedRegistrationCacheStore.RegistryCache GetKeyedRegistrationCache(IComponentRegistry registry)
         => KeyedCacheStore.GetOrCreateRegistryCache(registry);
 
-    private static Dictionary<Service, StandardRegistrationCacheEntry> GetStandardRegistrationCache(IComponentRegistry registry)
+    private static StandardRegistrationCacheStore.RegistryCache GetStandardRegistrationCache(IComponentRegistry registry)
         => StandardCacheStore.GetOrCreateRegistryCache(registry);
 
     private static int GetRegistryStamp(IComponentRegistry registry)
     {
+        if (registry is IRegistrationDiagnostics diagnostics)
+        {
+            return diagnostics.RegistrationsVersion;
+        }
+
         if (registry.Registrations is ICollection<IComponentRegistration> collection)
         {
             return collection.Count;
@@ -414,11 +406,11 @@ private sealed class KeyedRegistrationCacheStore : IReflectionCache
 
         public ReflectionCacheUsage Usage { get; } = ReflectionCacheUsage.All;
 
-        public Dictionary<Type, KeyedRegistrationCacheEntry> GetOrCreateRegistryCache(IComponentRegistry registry)
+        public RegistryCache GetOrCreateRegistryCache(IComponentRegistry registry)
         {
             var cache = _registries.GetValue(registry, static _ => new RegistryCache());
             RegisterCacheReference(cache);
-            return cache.Entries;
+            return cache;
         }
 
         public void Clear()
@@ -506,37 +498,43 @@ private void CleanupDeadRegistryReferences()
         /// <summary>
         /// Holds the cached keyed registrations for a single registry instance.
         /// </summary>
-        private sealed class RegistryCache
+        public sealed class RegistryCache
         {
+            private readonly ConcurrentDictionary<Type, KeyedRegistrationCacheEntry> _entries = new();
             private int _registrationState;
 
-            public Dictionary<Type, KeyedRegistrationCacheEntry> Entries { get; } = new();
-
-            public void ClearAll()
+            public IReadOnlyList<(Service Service, ServiceRegistration Registration)> GetOrUpdate(
+                Type elementType,
+                int stamp,
+                Func<IReadOnlyList<(Service Service, ServiceRegistration Registration)>> valueFactory)
             {
-                lock (Entries)
+                if (_entries.TryGetValue(elementType, out var existing) && existing.Stamp == stamp)
                 {
-                    Entries.Clear();
+                    return existing.Registrations;
                 }
+
+                var updated = _entries.AddOrUpdate(
+                    elementType,
+                    _ => new KeyedRegistrationCacheEntry(stamp, valueFactory()),
+                    (_, current) => current.Stamp == stamp ? current : new KeyedRegistrationCacheEntry(stamp, valueFactory()));
+
+                return updated.Registrations;
             }
 
-            public void Clear(ReflectionCacheClearPredicate predicate)
+            public void ClearAll()
             {
-                lock (Entries)
-                {
-                    if (Entries.Count == 0)
-                    {
-                        return;
-                    }
+                _entries.Clear();
+            }
 
-                    var assemblySet = new HashSet<Assembly>();
-                    var keysToRemove = Entries.Keys
-                        .Where(key => predicate(key, TypeAssemblyReferenceProvider.GetAllReferencedAssemblies(key, assemblySet)))
-                        .ToArray();
+            public void Clear(ReflectionCacheClearPredicate predicate)
+            {
+                var assemblySet = new HashSet<Assembly>();
 
-                    foreach (var key in keysToRemove)
+                foreach (var key in _entries.Keys)
+                {
+                    if (predicate(key, TypeAssemblyReferenceProvider.GetAllReferencedAssemblies(key, assemblySet)))
                     {
-                        Entries.Remove(key);
+                        _entries.TryRemove(key, out _);
                     }
                 }
             }
@@ -561,11 +559,11 @@ private sealed class StandardRegistrationCacheStore : IReflectionCache
 
         public ReflectionCacheUsage Usage { get; } = ReflectionCacheUsage.All;
 
-        public Dictionary<Service, StandardRegistrationCacheEntry> GetOrCreateRegistryCache(IComponentRegistry registry)
+        public RegistryCache GetOrCreateRegistryCache(IComponentRegistry registry)
         {
             var cache = _registries.GetValue(registry, static _ => new RegistryCache());
             RegisterCacheReference(cache);
-            return cache.Entries;
+            return cache;
         }
 
         public void Clear()
@@ -649,37 +647,43 @@ private void CleanupDeadRegistryReferences()
             }
         }
 
-        private sealed class RegistryCache
+        public sealed class RegistryCache
         {
+            private readonly ConcurrentDictionary<Service, StandardRegistrationCacheEntry> _entries = new();
             private int _registrationState;
 
-            public Dictionary<Service, StandardRegistrationCacheEntry> Entries { get; } = new();
-
-            public void ClearAll()
+            public IReadOnlyList<ServiceRegistration> GetOrUpdate(
+                Service service,
+                int stamp,
+                Func<IReadOnlyList<ServiceRegistration>> valueFactory)
             {
-                lock (Entries)
+                if (_entries.TryGetValue(service, out var existing) && existing.Stamp == stamp)
                 {
-                    Entries.Clear();
+                    return existing.Registrations;
                 }
+
+                var updated = _entries.AddOrUpdate(
+                    service,
+                    _ => new StandardRegistrationCacheEntry(stamp, valueFactory()),
+                    (_, current) => current.Stamp == stamp ? current : new StandardRegistrationCacheEntry(stamp, valueFactory()));
+
+                return updated.Registrations;
             }
 
-            public void Clear(ReflectionCacheClearPredicate predicate)
+            public void ClearAll()
             {
-                lock (Entries)
-                {
-                    if (Entries.Count == 0)
-                    {
-                        return;
-                    }
+                _entries.Clear();
+            }
 
-                    var assemblySet = new HashSet<Assembly>();
-                    var keysToRemove = Entries.Keys
-                        .Where(service => ShouldRemove(service, predicate, assemblySet))
-                        .ToArray();
+            public void Clear(ReflectionCacheClearPredicate predicate)
+            {
+                var assemblySet = new HashSet<Assembly>();
 
-                    foreach (var key in keysToRemove)
+                foreach (var key in _entries.Keys)
+                {
+                    if (ShouldRemove(key, predicate, assemblySet))
                     {
-                        Entries.Remove(key);
+                        _entries.TryRemove(key, out _);
                     }
                 }
             }