Fix missing ServiceCallSite.Key causing an unkeyed cache entry to be overwritten by a keyed instance (#113343)

IliaShuliatikov · web-flow · commit 23e0798cbaaa · 2025-03-14T09:32:33.000-05:00
diff --git a/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/CallSiteFactory.cs b/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/CallSiteFactory.cs
@@ -362,7 +362,7 @@ void AddCallSite(ServiceCallSite callSite, int index)
                 ResultCache resultCache = (cacheLocation == CallSiteResultCacheLocation.Scope || cacheLocation == CallSiteResultCacheLocation.Root)
                     ? new ResultCache(cacheLocation, callSiteKey)
                     : new ResultCache(CallSiteResultCacheLocation.None, callSiteKey);
-                return _callSiteCache[callSiteKey] = new IEnumerableCallSite(resultCache, itemType, callSites);
+                return _callSiteCache[callSiteKey] = new IEnumerableCallSite(resultCache, itemType, callSites, serviceIdentifier.ServiceKey);
             }
             finally
             {
@@ -415,7 +415,7 @@ private static CallSiteResultCacheLocation GetCommonCacheLocation(CallSiteResult
                 var lifetime = new ResultCache(descriptor.Lifetime, serviceIdentifier, slot);
                 if (descriptor.HasImplementationInstance())
                 {
-                    callSite = new ConstantCallSite(descriptor.ServiceType, descriptor.GetImplementationInstance());
+                    callSite = new ConstantCallSite(descriptor.ServiceType, descriptor.GetImplementationInstance(), descriptor.ServiceKey);
                 }
                 else if (!descriptor.IsKeyedService && descriptor.ImplementationFactory != null)
                 {
@@ -433,7 +433,6 @@ private static CallSiteResultCacheLocation GetCommonCacheLocation(CallSiteResult
                 {
                     throw new InvalidOperationException(SR.InvalidServiceDescriptor);
                 }
-                callSite.Key = descriptor.ServiceKey;
 
                 return _callSiteCache[callSiteKey] = callSite;
             }
@@ -512,7 +511,7 @@ private ConstructorCallSite CreateConstructorCallSite(
                     ParameterInfo[] parameters = constructor.GetParameters();
                     if (parameters.Length == 0)
                     {
-                        return new ConstructorCallSite(lifetime, serviceIdentifier.ServiceType, constructor);
+                        return new ConstructorCallSite(lifetime, serviceIdentifier.ServiceType, constructor, serviceIdentifier.ServiceKey);
                     }
 
                     parameterCallSites = CreateArgumentCallSites(
@@ -522,7 +521,7 @@ private ConstructorCallSite CreateConstructorCallSite(
                         parameters,
                         throwIfCallSiteNotFound: true)!;
 
-                    return new ConstructorCallSite(lifetime, serviceIdentifier.ServiceType, constructor, parameterCallSites);
+                    return new ConstructorCallSite(lifetime, serviceIdentifier.ServiceType, constructor, parameterCallSites, serviceIdentifier.ServiceKey);
                 }
 
                 Array.Sort(constructors,
@@ -586,7 +585,7 @@ private ConstructorCallSite CreateConstructorCallSite(
                 else
                 {
                     Debug.Assert(parameterCallSites != null);
-                    return new ConstructorCallSite(lifetime, serviceIdentifier.ServiceType, bestConstructor, parameterCallSites);
+                    return new ConstructorCallSite(lifetime, serviceIdentifier.ServiceType, bestConstructor, parameterCallSites, serviceIdentifier.ServiceKey);
                 }
             }
             finally
diff --git a/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/ConstantCallSite.cs b/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/ConstantCallSite.cs
@@ -10,7 +10,7 @@ internal sealed class ConstantCallSite : ServiceCallSite
         private readonly Type _serviceType;
         internal object? DefaultValue => Value;
 
-        public ConstantCallSite(Type serviceType, object? defaultValue) : base(ResultCache.None(serviceType))
+        public ConstantCallSite(Type serviceType, object? defaultValue, object? serviceKey = null) : base(ResultCache.None(serviceType), serviceKey)
         {
             _serviceType = serviceType ?? throw new ArgumentNullException(nameof(serviceType));
             if (defaultValue != null && !serviceType.IsInstanceOfType(defaultValue))
diff --git a/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/ConstructorCallSite.cs b/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/ConstructorCallSite.cs
@@ -11,11 +11,11 @@ internal sealed class ConstructorCallSite : ServiceCallSite
         internal ConstructorInfo ConstructorInfo { get; }
         internal ServiceCallSite[] ParameterCallSites { get; }
 
-        public ConstructorCallSite(ResultCache cache, Type serviceType, ConstructorInfo constructorInfo) : this(cache, serviceType, constructorInfo, Array.Empty<ServiceCallSite>())
+        public ConstructorCallSite(ResultCache cache, Type serviceType, ConstructorInfo constructorInfo, object? serviceKey) : this(cache, serviceType, constructorInfo, Array.Empty<ServiceCallSite>(), serviceKey)
         {
         }
 
-        public ConstructorCallSite(ResultCache cache, Type serviceType, ConstructorInfo constructorInfo, ServiceCallSite[] parameterCallSites) : base(cache)
+        public ConstructorCallSite(ResultCache cache, Type serviceType, ConstructorInfo constructorInfo, ServiceCallSite[] parameterCallSites, object? serviceKey) : base(cache, serviceKey)
         {
             if (!serviceType.IsAssignableFrom(constructorInfo.DeclaringType))
             {
diff --git a/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/FactoryCallSite.cs b/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/FactoryCallSite.cs
@@ -9,13 +9,13 @@ internal sealed class FactoryCallSite : ServiceCallSite
     {
         public Func<IServiceProvider, object> Factory { get; }
 
-        public FactoryCallSite(ResultCache cache, Type serviceType, Func<IServiceProvider, object> factory) : base(cache)
+        public FactoryCallSite(ResultCache cache, Type serviceType, Func<IServiceProvider, object> factory) : base(cache, null)
         {
             Factory = factory;
             ServiceType = serviceType;
         }
 
-        public FactoryCallSite(ResultCache cache, Type serviceType, object serviceKey, Func<IServiceProvider, object, object> factory) : base(cache)
+        public FactoryCallSite(ResultCache cache, Type serviceType, object serviceKey, Func<IServiceProvider, object, object> factory) : base(cache, serviceKey)
         {
             Factory = sp => factory(sp, serviceKey);
             ServiceType = serviceType;
diff --git a/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/IEnumerableCallSite.cs b/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/IEnumerableCallSite.cs
@@ -13,7 +13,7 @@ internal sealed class IEnumerableCallSite : ServiceCallSite
         internal Type ItemType { get; }
         internal ServiceCallSite[] ServiceCallSites { get; }
 
-        public IEnumerableCallSite(ResultCache cache, Type itemType, ServiceCallSite[] serviceCallSites) : base(cache)
+        public IEnumerableCallSite(ResultCache cache, Type itemType, ServiceCallSite[] serviceCallSites, object? serviceKey = null) : base(cache, serviceKey)
         {
             Debug.Assert(!ServiceProvider.VerifyAotCompatibility || !itemType.IsValueType, "If VerifyAotCompatibility=true, an IEnumerableCallSite should not be created with a ValueType.");
 
diff --git a/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/ServiceCallSite.cs b/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/ServiceCallSite.cs
@@ -10,17 +10,18 @@ namespace Microsoft.Extensions.DependencyInjection.ServiceLookup
     /// </summary>
     internal abstract class ServiceCallSite
     {
-        protected ServiceCallSite(ResultCache cache)
+        protected ServiceCallSite(ResultCache cache, object? key)
         {
             Cache = cache;
+            Key = key;
         }
 
         public abstract Type ServiceType { get; }
         public abstract Type? ImplementationType { get; }
         public abstract CallSiteKind Kind { get; }
         public ResultCache Cache { get; }
         public object? Value { get; set; }
-        public object? Key { get; set; }
+        public object? Key { get; }
 
         public bool CaptureDisposable =>
             ImplementationType == null ||
diff --git a/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/ServiceProviderCallSite.cs b/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/ServiceProviderCallSite.cs
@@ -7,7 +7,7 @@ namespace Microsoft.Extensions.DependencyInjection.ServiceLookup
 {
     internal sealed class ServiceProviderCallSite : ServiceCallSite
     {
-        public ServiceProviderCallSite() : base(ResultCache.None(typeof(IServiceProvider)))
+        public ServiceProviderCallSite() : base(ResultCache.None(typeof(IServiceProvider)), null)
         {
         }
 
diff --git a/src/libraries/Microsoft.Extensions.DependencyInjection/tests/DI.Tests/CallSiteTests.cs b/src/libraries/Microsoft.Extensions.DependencyInjection/tests/DI.Tests/CallSiteTests.cs
@@ -17,6 +17,11 @@ internal static ServiceCallSite GetCallSite(this CallSiteFactory callSiteFactory
         {
             return callSiteFactory.GetCallSite(ServiceIdentifier.FromServiceType(type), callSiteChain);
         }
+
+        internal static ServiceCallSite GetKeyedCallSite(this CallSiteFactory callSiteFactory, Type type, object? serviceKey, CallSiteChain callSiteChain)
+        {
+            return callSiteFactory.GetCallSite(new ServiceIdentifier(serviceKey, type), callSiteChain);
+        }
     }
 
     public class CallSiteTests
@@ -300,6 +305,61 @@ public void CallSiteFactoryResolvesIEnumerableOfOpenGenericServiceAfterResolving
             Assert.Equal(typeof(FakeOpenGenericService<int>), implementationTypes[1]);
         }
 
+        [Fact]
+        public void ServiceCallSite_ShouldHaveKey_WhenResolvingKeyedService()
+        {
+            // Arrange
+            IServiceCollection services = new ServiceCollection();
+
+            services.Add(ServiceDescriptor.Transient(typeof(SomeService), typeof(SomeService)));
+            services.Add(ServiceDescriptor.KeyedTransient(typeof(SomeService), "someKey", typeof(SomeOtherService)));
+
+            using var serviceProvider = services.BuildServiceProvider();
+
+            // Act
+            var callSite = serviceProvider.CallSiteFactory.GetKeyedCallSite(typeof(SomeService), "someKey", new CallSiteChain());
+
+            // Assert
+            Assert.NotNull(callSite.Key);
+        }
+
+        [Fact]
+        public void ServiceCallSite_ShouldHaveKey_WhenResolvingKeyedClosedImplementationOfOpenGenericService()
+        {
+            // Arrange
+            IServiceCollection services = new ServiceCollection();
+
+            services.Add(ServiceDescriptor.Transient(typeof(IGenericService<>), typeof(UnkeyedGenericService<>)));
+            services.Add(ServiceDescriptor.KeyedTransient(typeof(IGenericService<>), "someKey", typeof(PrimaryKeyedGenericService<>)));
+
+            using var serviceProvider = services.BuildServiceProvider();
+
+            // Act
+            var callSite = serviceProvider.CallSiteFactory.GetKeyedCallSite(typeof(IGenericService<object>), "someKey", new CallSiteChain());
+
+            // Assert
+            Assert.NotNull(callSite.Key);
+        }
+
+        [Fact]
+        public void ServiceCallSite_ShouldHaveKey_WhenResolvingKeyedIEnumerableOfClosedImplementationOfOpenGenericService()
+        {
+            // Arrange
+            IServiceCollection services = new ServiceCollection();
+
+            services.Add(ServiceDescriptor.Transient(typeof(IGenericService<>), typeof(UnkeyedGenericService<>)));
+            services.Add(ServiceDescriptor.KeyedTransient(typeof(IGenericService<>), "someKey", typeof(PrimaryKeyedGenericService<>)));
+            services.Add(ServiceDescriptor.KeyedTransient(typeof(IGenericService<>), "someKey", typeof(SecondaryKeyedGenericService<>)));
+
+            using var serviceProvider = services.BuildServiceProvider();
+
+            // Act
+            var callSite = serviceProvider.CallSiteFactory.GetKeyedCallSite(typeof(IEnumerable<IGenericService<object>>), "someKey", new CallSiteChain());
+
+            // Assert
+            Assert.NotNull(callSite.Key);
+        }
+
         private class FakeIntService : IFakeOpenGenericService<int>
         {
             public int Value => 0;
@@ -395,6 +455,18 @@ public void Dispose()
             }
         }
 
+        private interface IGenericService<T>;
+
+        private class UnkeyedGenericService<T> : IGenericService<T>;
+
+        private class PrimaryKeyedGenericService<T> : IGenericService<T>;
+
+        private class SecondaryKeyedGenericService<T> : IGenericService<T>;
+
+        private class SomeService;
+
+        private class SomeOtherService : SomeService;
+
         private static object Invoke(ServiceCallSite callSite, ServiceProviderEngineScope scope)
         {
             return CallSiteRuntimeResolver.Instance.Resolve(callSite, scope);

Original file line number	Diff line number	Diff line change
`@@ -362,7 +362,7 @@ void AddCallSite(ServiceCallSite callSite, int index)`
`362`	`362`	`ResultCache resultCache = (cacheLocation == CallSiteResultCacheLocation.Scope \|\| cacheLocation == CallSiteResultCacheLocation.Root)`
`363`	`363`	`? new ResultCache(cacheLocation, callSiteKey)`
`364`	`364`	`: new ResultCache(CallSiteResultCacheLocation.None, callSiteKey);`
`365`		`- return _callSiteCache[callSiteKey] = new IEnumerableCallSite(resultCache, itemType, callSites);`
	`365`	`+ return _callSiteCache[callSiteKey] = new IEnumerableCallSite(resultCache, itemType, callSites, serviceIdentifier.ServiceKey);`
`366`	`366`	`}`
`367`	`367`	`finally`
`368`	`368`	`{`
`@@ -415,7 +415,7 @@ private static CallSiteResultCacheLocation GetCommonCacheLocation(CallSiteResult`
`415`	`415`	`var lifetime = new ResultCache(descriptor.Lifetime, serviceIdentifier, slot);`
`416`	`416`	`if (descriptor.HasImplementationInstance())`
`417`	`417`	`{`
`418`		`- callSite = new ConstantCallSite(descriptor.ServiceType, descriptor.GetImplementationInstance());`
	`418`	`+ callSite = new ConstantCallSite(descriptor.ServiceType, descriptor.GetImplementationInstance(), descriptor.ServiceKey);`
`419`	`419`	`}`
`420`	`420`	`else if (!descriptor.IsKeyedService && descriptor.ImplementationFactory != null)`
`421`	`421`	`{`
`@@ -433,7 +433,6 @@ private static CallSiteResultCacheLocation GetCommonCacheLocation(CallSiteResult`
`433`	`433`	`{`
`434`	`434`	`throw new InvalidOperationException(SR.InvalidServiceDescriptor);`
`435`	`435`	`}`
`436`		`- callSite.Key = descriptor.ServiceKey;`
`437`	`436`
`438`	`437`	`return _callSiteCache[callSiteKey] = callSite;`
`439`	`438`	`}`
`@@ -512,7 +511,7 @@ private ConstructorCallSite CreateConstructorCallSite(`
`512`	`511`	`ParameterInfo[] parameters = constructor.GetParameters();`
`513`	`512`	`if (parameters.Length == 0)`
`514`	`513`	`{`
`515`		`- return new ConstructorCallSite(lifetime, serviceIdentifier.ServiceType, constructor);`
	`514`	`+ return new ConstructorCallSite(lifetime, serviceIdentifier.ServiceType, constructor, serviceIdentifier.ServiceKey);`
`516`	`515`	`}`
`517`	`516`
`518`	`517`	`parameterCallSites = CreateArgumentCallSites(`
`@@ -522,7 +521,7 @@ private ConstructorCallSite CreateConstructorCallSite(`
`522`	`521`	`parameters,`
`523`	`522`	`throwIfCallSiteNotFound: true)!;`
`524`	`523`
`525`		`- return new ConstructorCallSite(lifetime, serviceIdentifier.ServiceType, constructor, parameterCallSites);`
	`524`	`+ return new ConstructorCallSite(lifetime, serviceIdentifier.ServiceType, constructor, parameterCallSites, serviceIdentifier.ServiceKey);`
`526`	`525`	`}`
`527`	`526`
`528`	`527`	`Array.Sort(constructors,`
`@@ -586,7 +585,7 @@ private ConstructorCallSite CreateConstructorCallSite(`
`586`	`585`	`else`
`587`	`586`	`{`
`588`	`587`	`Debug.Assert(parameterCallSites != null);`
`589`		`- return new ConstructorCallSite(lifetime, serviceIdentifier.ServiceType, bestConstructor, parameterCallSites);`
	`588`	`+ return new ConstructorCallSite(lifetime, serviceIdentifier.ServiceType, bestConstructor, parameterCallSites, serviceIdentifier.ServiceKey);`
`590`	`589`	`}`
`591`	`590`	`}`
`592`	`591`	`finally`
Original file line number	Diff line number	Diff line change
`@@ -10,7 +10,7 @@ internal sealed class ConstantCallSite : ServiceCallSite`
`10`	`10`	`private readonly Type _serviceType;`
`11`	`11`	`internal object? DefaultValue => Value;`
`12`	`12`
`13`		`- public ConstantCallSite(Type serviceType, object? defaultValue) : base(ResultCache.None(serviceType))`
	`13`	`+ public ConstantCallSite(Type serviceType, object? defaultValue, object? serviceKey = null) : base(ResultCache.None(serviceType), serviceKey)`
`14`	`14`	`{`
`15`	`15`	`_serviceType = serviceType ?? throw new ArgumentNullException(nameof(serviceType));`
`16`	`16`	`if (defaultValue != null && !serviceType.IsInstanceOfType(defaultValue))`
Original file line number	Diff line number	Diff line change
`@@ -11,11 +11,11 @@ internal sealed class ConstructorCallSite : ServiceCallSite`
`11`	`11`	`internal ConstructorInfo ConstructorInfo { get; }`
`12`	`12`	`internal ServiceCallSite[] ParameterCallSites { get; }`
`13`	`13`
`14`		`- public ConstructorCallSite(ResultCache cache, Type serviceType, ConstructorInfo constructorInfo) : this(cache, serviceType, constructorInfo, Array.Empty<ServiceCallSite>())`
	`14`	`+ public ConstructorCallSite(ResultCache cache, Type serviceType, ConstructorInfo constructorInfo, object? serviceKey) : this(cache, serviceType, constructorInfo, Array.Empty<ServiceCallSite>(), serviceKey)`
`15`	`15`	`{`
`16`	`16`	`}`
`17`	`17`
`18`		`- public ConstructorCallSite(ResultCache cache, Type serviceType, ConstructorInfo constructorInfo, ServiceCallSite[] parameterCallSites) : base(cache)`
	`18`	`+ public ConstructorCallSite(ResultCache cache, Type serviceType, ConstructorInfo constructorInfo, ServiceCallSite[] parameterCallSites, object? serviceKey) : base(cache, serviceKey)`
`19`	`19`	`{`
`20`	`20`	`if (!serviceType.IsAssignableFrom(constructorInfo.DeclaringType))`
`21`	`21`	`{`
Original file line number	Diff line number	Diff line change
`@@ -9,13 +9,13 @@ internal sealed class FactoryCallSite : ServiceCallSite`
`9`	`9`	`{`
`10`	`10`	`public Func<IServiceProvider, object> Factory { get; }`
`11`	`11`
`12`		`- public FactoryCallSite(ResultCache cache, Type serviceType, Func<IServiceProvider, object> factory) : base(cache)`
	`12`	`+ public FactoryCallSite(ResultCache cache, Type serviceType, Func<IServiceProvider, object> factory) : base(cache, null)`
`13`	`13`	`{`
`14`	`14`	`Factory = factory;`
`15`	`15`	`ServiceType = serviceType;`
`16`	`16`	`}`
`17`	`17`
`18`		`- public FactoryCallSite(ResultCache cache, Type serviceType, object serviceKey, Func<IServiceProvider, object, object> factory) : base(cache)`
	`18`	`+ public FactoryCallSite(ResultCache cache, Type serviceType, object serviceKey, Func<IServiceProvider, object, object> factory) : base(cache, serviceKey)`
`19`	`19`	`{`
`20`	`20`	`Factory = sp => factory(sp, serviceKey);`
`21`	`21`	`ServiceType = serviceType;`
Original file line number	Diff line number	Diff line change
`@@ -13,7 +13,7 @@ internal sealed class IEnumerableCallSite : ServiceCallSite`
`13`	`13`	`internal Type ItemType { get; }`
`14`	`14`	`internal ServiceCallSite[] ServiceCallSites { get; }`
`15`	`15`
`16`		`- public IEnumerableCallSite(ResultCache cache, Type itemType, ServiceCallSite[] serviceCallSites) : base(cache)`
	`16`	`+ public IEnumerableCallSite(ResultCache cache, Type itemType, ServiceCallSite[] serviceCallSites, object? serviceKey = null) : base(cache, serviceKey)`
`17`	`17`	`{`
`18`	`18`	`Debug.Assert(!ServiceProvider.VerifyAotCompatibility \|\| !itemType.IsValueType, "If VerifyAotCompatibility=true, an IEnumerableCallSite should not be created with a ValueType.");`
`19`	`19`
Original file line number	Diff line number	Diff line change
`@@ -10,17 +10,18 @@ namespace Microsoft.Extensions.DependencyInjection.ServiceLookup`
`10`	`10`	`/// </summary>`
`11`	`11`	`internal abstract class ServiceCallSite`
`12`	`12`	`{`
`13`		`- protected ServiceCallSite(ResultCache cache)`
	`13`	`+ protected ServiceCallSite(ResultCache cache, object? key)`
`14`	`14`	`{`
`15`	`15`	`Cache = cache;`
	`16`	`+ Key = key;`
`16`	`17`	`}`
`17`	`18`
`18`	`19`	`public abstract Type ServiceType { get; }`
`19`	`20`	`public abstract Type? ImplementationType { get; }`
`20`	`21`	`public abstract CallSiteKind Kind { get; }`
`21`	`22`	`public ResultCache Cache { get; }`
`22`	`23`	`public object? Value { get; set; }`
`23`		`- public object? Key { get; set; }`
	`24`	`+ public object? Key { get; }`
`24`	`25`
`25`	`26`	`public bool CaptureDisposable =>`
`26`	`27`	`ImplementationType == null \|\|`
Original file line number	Diff line number	Diff line change
`@@ -7,7 +7,7 @@ namespace Microsoft.Extensions.DependencyInjection.ServiceLookup`
`7`	`7`	`{`
`8`	`8`	`internal sealed class ServiceProviderCallSite : ServiceCallSite`
`9`	`9`	`{`
`10`		`- public ServiceProviderCallSite() : base(ResultCache.None(typeof(IServiceProvider)))`
	`10`	`+ public ServiceProviderCallSite() : base(ResultCache.None(typeof(IServiceProvider)), null)`
`11`	`11`	`{`
`12`	`12`	`}`
`13`	`13`