Fix validation on build (#87354)

mapogolions · web-flow · commit c06d77a74f64 · 2023-06-28T09:39:29.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
@@ -316,13 +316,9 @@ void AddCallSite(ServiceCallSite callSite, int index)
                         callSitesByIndex.Add(new(index, callSite));
                     }
                 }
-
-                ResultCache resultCache = ResultCache.None;
-                if (cacheLocation == CallSiteResultCacheLocation.Scope || cacheLocation == CallSiteResultCacheLocation.Root)
-                {
-                    resultCache = new ResultCache(cacheLocation, callSiteKey);
-                }
-
+                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);
             }
             finally
diff --git a/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/CallSiteValidator.cs b/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/CallSiteValidator.cs
@@ -10,22 +10,23 @@ namespace Microsoft.Extensions.DependencyInjection.ServiceLookup
     internal sealed class CallSiteValidator: CallSiteVisitor<CallSiteValidator.CallSiteValidatorState, Type?>
     {
         // Keys are services being resolved via GetService, values - first scoped service in their call site tree
-        private readonly ConcurrentDictionary<Type, Type> _scopedServices = new ConcurrentDictionary<Type, Type>();
+        private readonly ConcurrentDictionary<ServiceCacheKey, Type> _scopedServices = new ConcurrentDictionary<ServiceCacheKey, Type>();
 
         public void ValidateCallSite(ServiceCallSite callSite)
         {
             Type? scoped = VisitCallSite(callSite, default);
             if (scoped != null)
             {
-                _scopedServices[callSite.ServiceType] = scoped;
+                _scopedServices[callSite.Cache.Key] = scoped;
             }
         }
 
-        public void ValidateResolution(Type serviceType, IServiceScope scope, IServiceScope rootScope)
+        public void ValidateResolution(ServiceCallSite callSite, IServiceScope scope, IServiceScope rootScope)
         {
             if (ReferenceEquals(scope, rootScope)
-                && _scopedServices.TryGetValue(serviceType, out Type? scopedService))
+                && _scopedServices.TryGetValue(callSite.Cache.Key, out Type? scopedService))
             {
+                Type serviceType = callSite.ServiceType;
                 if (serviceType == scopedService)
                 {
                     throw new InvalidOperationException(
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)
+        public ConstantCallSite(Type serviceType, object? defaultValue) : base(ResultCache.None(serviceType))
         {
             _serviceType = serviceType ?? throw new ArgumentNullException(nameof(serviceType));
             if (defaultValue != null && !serviceType.IsInstanceOfType(defaultValue))
diff --git a/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/ResultCache.cs b/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/ResultCache.cs
@@ -8,7 +8,11 @@ namespace Microsoft.Extensions.DependencyInjection.ServiceLookup
 {
     internal struct ResultCache
     {
-        public static ResultCache None { get; } = new ResultCache(CallSiteResultCacheLocation.None, ServiceCacheKey.Empty);
+        public static ResultCache None(Type serviceType)
+        {
+            var cacheKey = new ServiceCacheKey(serviceType, 0);
+            return new ResultCache(CallSiteResultCacheLocation.None, cacheKey);
+        }
 
         internal ResultCache(CallSiteResultCacheLocation lifetime, ServiceCacheKey cacheKey)
         {
diff --git a/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/ServiceCacheKey.cs b/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceLookup/ServiceCacheKey.cs
@@ -8,8 +8,6 @@ namespace Microsoft.Extensions.DependencyInjection.ServiceLookup
 {
     internal readonly struct ServiceCacheKey : IEquatable<ServiceCacheKey>
     {
-        public static ServiceCacheKey Empty { get; } = new ServiceCacheKey(null, 0);
-
         /// <summary>
         /// Type of service being cached
         /// </summary>
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)
+        public ServiceProviderCallSite() : base(ResultCache.None(typeof(IServiceProvider)))
         {
         }
 
diff --git a/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceProvider.cs b/src/libraries/Microsoft.Extensions.DependencyInjection/src/ServiceProvider.cs
@@ -120,9 +120,9 @@ private void OnCreate(ServiceCallSite callSite)
             _callSiteValidator?.ValidateCallSite(callSite);
         }
 
-        private void OnResolve(Type serviceType, IServiceScope scope)
+        private void OnResolve(ServiceCallSite callSite, IServiceScope scope)
         {
-            _callSiteValidator?.ValidateResolution(serviceType, scope, Root);
+            _callSiteValidator?.ValidateResolution(callSite, scope, Root);
         }
 
         internal object? GetService(Type serviceType, ServiceProviderEngineScope serviceProviderEngineScope)
@@ -133,8 +133,6 @@ private void OnResolve(Type serviceType, IServiceScope scope)
             }
 
             Func<ServiceProviderEngineScope, object?> realizedService = _realizedServices.GetOrAdd(serviceType, _createServiceAccessor);
-            OnResolve(serviceType, serviceProviderEngineScope);
-            DependencyInjectionEventSource.Log.ServiceResolved(this, serviceType);
             var result = realizedService.Invoke(serviceProviderEngineScope);
             System.Diagnostics.Debug.Assert(result is null || CallSiteFactory.IsService(serviceType));
             return result;
@@ -173,10 +171,20 @@ private void ValidateService(ServiceDescriptor descriptor)
                 if (callSite.Cache.Location == CallSiteResultCacheLocation.Root)
                 {
                     object? value = CallSiteRuntimeResolver.Instance.Resolve(callSite, Root);
-                    return scope => value;
+                    return scope =>
+                    {
+                        DependencyInjectionEventSource.Log.ServiceResolved(this, serviceType);
+                        return value;
+                    };
                 }
 
-                return _engine.RealizeService(callSite);
+                Func<ServiceProviderEngineScope, object?> realizedService = _engine.RealizeService(callSite);
+                return scope =>
+                {
+                    OnResolve(callSite, scope);
+                    DependencyInjectionEventSource.Log.ServiceResolved(this, serviceType);
+                    return realizedService(scope);
+                };
             }
 
             return _ => null;
diff --git a/src/libraries/Microsoft.Extensions.DependencyInjection/tests/DI.Tests/ServiceLookup/CallSiteFactoryTest.cs b/src/libraries/Microsoft.Extensions.DependencyInjection/tests/DI.Tests/ServiceLookup/CallSiteFactoryTest.cs
@@ -786,17 +786,10 @@ public void CreateCallSite_EnumberableCachedAtLowestLevel(ServiceDescriptor[] de
             var callSite = factory(typeof(IEnumerable<FakeService>));
 
             var expectedLocation = (CallSiteResultCacheLocation)expectedCacheLocation;
-            Assert.Equal(expectedLocation, callSite.Cache.Location);
 
-            if (expectedLocation != CallSiteResultCacheLocation.None)
-            {
-                Assert.Equal(0, callSite.Cache.Key.Slot);
-                Assert.Equal(typeof(IEnumerable<FakeService>), callSite.Cache.Key.Type);
-            }
-            else
-            {
-                Assert.Equal(ResultCache.None, callSite.Cache);
-            }
+            Assert.Equal(expectedLocation, callSite.Cache.Location);
+            Assert.Equal(0, callSite.Cache.Key.Slot);
+            Assert.Equal(typeof(IEnumerable<FakeService>), callSite.Cache.Key.Type);
         }
 
         [ConditionalFact(typeof(PlatformDetection), nameof(PlatformDetection.IsThreadingSupported))]
diff --git a/src/libraries/Microsoft.Extensions.DependencyInjection/tests/DI.Tests/ServiceProviderValidationTests.cs b/src/libraries/Microsoft.Extensions.DependencyInjection/tests/DI.Tests/ServiceProviderValidationTests.cs
@@ -2,6 +2,8 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 
 using System;
+using System.Collections.Generic;
+using System.Linq;
 using Microsoft.Extensions.DependencyInjection.Specification.Fakes;
 using Xunit;
 
@@ -97,6 +99,49 @@ public void GetService_Throws_WhenGetServiceForScopedServiceIsCalledOnRootViaTra
             Assert.Equal($"Cannot resolve '{typeof(IFoo)}' from root provider because it requires scoped service '{typeof(IBar)}'.", exception.Message);
         }
 
+        [Theory]
+        [InlineData(true)]
+        [InlineData(false)]
+        public void GetService_DoesNotThrow_WhenGetServiceForPolymorphicServiceIsCalledOnRoot_AndTheLastOneIsNotScoped(bool validateOnBuild)
+        {
+            // Arrange
+            var serviceCollection = new ServiceCollection();
+            serviceCollection.AddScoped<IBar, Bar>();
+            serviceCollection.AddTransient<IBar, Bar3>();
+            using var serviceProvider = serviceCollection.BuildServiceProvider(new ServiceProviderOptions
+            {
+                ValidateScopes = true,
+                ValidateOnBuild = validateOnBuild
+            });
+
+            // Act
+            var actual = serviceProvider.GetService<IBar>();
+
+            // Assert
+            Assert.IsType<Bar3>(actual);
+        }
+
+        [Fact]
+        public void ScopeValidation_ShouldBeAbleToDistingushGenericCollections_WhenGetServiceIsCalledOnRoot()
+        {
+            // Arrange
+            var serviceCollection = new ServiceCollection();
+            serviceCollection.AddTransient<IBar, Bar>();
+            serviceCollection.AddScoped<IBar, Bar3>();
+
+            serviceCollection.AddTransient<IBaz, Baz>();
+            serviceCollection.AddTransient<IBaz, Baz2>();
+
+            // Act
+            using var serviceProvider = serviceCollection.BuildServiceProvider(validateScopes: true);
+            Assert.Throws<InvalidOperationException>(() => serviceProvider.GetService<IEnumerable<IBar>>());
+            var actual = serviceProvider.GetService<IEnumerable<IBaz>>();
+
+            // Assert
+            Assert.IsType<Baz>(actual.First());
+            Assert.IsType<Baz2>(actual.Last());
+        }
+
         [Fact]
         public void GetService_DoesNotThrow_WhenScopeFactoryIsInjectedIntoSingleton()
         {
@@ -206,13 +251,18 @@ private class Bar : IBar
         {
         }
 
+
         private class Bar2 : IBar
         {
             public Bar2(IBaz baz)
             {
             }
         }
 
+        private class Bar3 : IBar
+        {
+        }
+
         private interface IBaz
         {
         }
@@ -221,6 +271,10 @@ private class Baz : IBaz
         {
         }
 
+        private class Baz2 : IBaz
+        {
+        }
+
         private class BazRecursive : IBaz
         {
             public BazRecursive(IBaz baz)

Original file line number	Diff line number	Diff line change
`@@ -316,13 +316,9 @@ void AddCallSite(ServiceCallSite callSite, int index)`
`316`	`316`	`callSitesByIndex.Add(new(index, callSite));`
`317`	`317`	`}`
`318`	`318`	`}`
`319`		`-`
`320`		`- ResultCache resultCache = ResultCache.None;`
`321`		`- if (cacheLocation == CallSiteResultCacheLocation.Scope \|\| cacheLocation == CallSiteResultCacheLocation.Root)`
`322`		`- {`
`323`		`- resultCache = new ResultCache(cacheLocation, callSiteKey);`
`324`		`- }`
`325`		`-`
	`319`	`+ ResultCache resultCache = (cacheLocation == CallSiteResultCacheLocation.Scope \|\| cacheLocation == CallSiteResultCacheLocation.Root)`
	`320`	`+ ? new ResultCache(cacheLocation, callSiteKey)`
	`321`	`+ : new ResultCache(CallSiteResultCacheLocation.None, callSiteKey);`
`326`	`322`	`return _callSiteCache[callSiteKey] = new IEnumerableCallSite(resultCache, itemType, callSites);`
`327`	`323`	`}`
`328`	`324`	`finally`
Original file line number	Diff line number	Diff line change
`@@ -10,22 +10,23 @@ namespace Microsoft.Extensions.DependencyInjection.ServiceLookup`
`10`	`10`	`internal sealed class CallSiteValidator: CallSiteVisitor<CallSiteValidator.CallSiteValidatorState, Type?>`
`11`	`11`	`{`
`12`	`12`	`// Keys are services being resolved via GetService, values - first scoped service in their call site tree`
`13`		`- private readonly ConcurrentDictionary<Type, Type> _scopedServices = new ConcurrentDictionary<Type, Type>();`
	`13`	`+ private readonly ConcurrentDictionary<ServiceCacheKey, Type> _scopedServices = new ConcurrentDictionary<ServiceCacheKey, Type>();`
`14`	`14`
`15`	`15`	`public void ValidateCallSite(ServiceCallSite callSite)`
`16`	`16`	`{`
`17`	`17`	`Type? scoped = VisitCallSite(callSite, default);`
`18`	`18`	`if (scoped != null)`
`19`	`19`	`{`
`20`		`- _scopedServices[callSite.ServiceType] = scoped;`
	`20`	`+ _scopedServices[callSite.Cache.Key] = scoped;`
`21`	`21`	`}`
`22`	`22`	`}`
`23`	`23`
`24`		`- public void ValidateResolution(Type serviceType, IServiceScope scope, IServiceScope rootScope)`
	`24`	`+ public void ValidateResolution(ServiceCallSite callSite, IServiceScope scope, IServiceScope rootScope)`
`25`	`25`	`{`
`26`	`26`	`if (ReferenceEquals(scope, rootScope)`
`27`		`- && _scopedServices.TryGetValue(serviceType, out Type? scopedService))`
	`27`	`+ && _scopedServices.TryGetValue(callSite.Cache.Key, out Type? scopedService))`
`28`	`28`	`{`
	`29`	`+ Type serviceType = callSite.ServiceType;`
`29`	`30`	`if (serviceType == scopedService)`
`30`	`31`	`{`
`31`	`32`	`throw new InvalidOperationException(`
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)`
	`13`	`+ public ConstantCallSite(Type serviceType, object? defaultValue) : base(ResultCache.None(serviceType))`
`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
`@@ -8,7 +8,11 @@ namespace Microsoft.Extensions.DependencyInjection.ServiceLookup`
`8`	`8`	`{`
`9`	`9`	`internal struct ResultCache`
`10`	`10`	`{`
`11`		`- public static ResultCache None { get; } = new ResultCache(CallSiteResultCacheLocation.None, ServiceCacheKey.Empty);`
	`11`	`+ public static ResultCache None(Type serviceType)`
	`12`	`+ {`
	`13`	`+ var cacheKey = new ServiceCacheKey(serviceType, 0);`
	`14`	`+ return new ResultCache(CallSiteResultCacheLocation.None, cacheKey);`
	`15`	`+ }`
`12`	`16`
`13`	`17`	`internal ResultCache(CallSiteResultCacheLocation lifetime, ServiceCacheKey cacheKey)`
`14`	`18`	`{`
Original file line number	Diff line number	Diff line change
`@@ -8,8 +8,6 @@ namespace Microsoft.Extensions.DependencyInjection.ServiceLookup`
`8`	`8`	`{`
`9`	`9`	`internal readonly struct ServiceCacheKey : IEquatable<ServiceCacheKey>`
`10`	`10`	`{`
`11`		`- public static ServiceCacheKey Empty { get; } = new ServiceCacheKey(null, 0);`
`12`		`-`
`13`	`11`	`/// <summary>`
`14`	`12`	`/// Type of service being cached`
`15`	`13`	`/// </summary>`
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)`
	`10`	`+ public ServiceProviderCallSite() : base(ResultCache.None(typeof(IServiceProvider)))`
`11`	`11`	`{`
`12`	`12`	`}`
`13`	`13`
Original file line number	Diff line number	Diff line change
`@@ -120,9 +120,9 @@ private void OnCreate(ServiceCallSite callSite)`
`120`	`120`	`_callSiteValidator?.ValidateCallSite(callSite);`
`121`	`121`	`}`
`122`	`122`
`123`		`- private void OnResolve(Type serviceType, IServiceScope scope)`
	`123`	`+ private void OnResolve(ServiceCallSite callSite, IServiceScope scope)`
`124`	`124`	`{`
`125`		`- _callSiteValidator?.ValidateResolution(serviceType, scope, Root);`
	`125`	`+ _callSiteValidator?.ValidateResolution(callSite, scope, Root);`
`126`	`126`	`}`
`127`	`127`
`128`	`128`	`internal object? GetService(Type serviceType, ServiceProviderEngineScope serviceProviderEngineScope)`
`@@ -133,8 +133,6 @@ private void OnResolve(Type serviceType, IServiceScope scope)`
`133`	`133`	`}`
`134`	`134`
`135`	`135`	`Func<ServiceProviderEngineScope, object?> realizedService = _realizedServices.GetOrAdd(serviceType, _createServiceAccessor);`
`136`		`- OnResolve(serviceType, serviceProviderEngineScope);`
`137`		`- DependencyInjectionEventSource.Log.ServiceResolved(this, serviceType);`
`138`	`136`	`var result = realizedService.Invoke(serviceProviderEngineScope);`
`139`	`137`	`System.Diagnostics.Debug.Assert(result is null \|\| CallSiteFactory.IsService(serviceType));`
`140`	`138`	`return result;`
`@@ -173,10 +171,20 @@ private void ValidateService(ServiceDescriptor descriptor)`
`173`	`171`	`if (callSite.Cache.Location == CallSiteResultCacheLocation.Root)`
`174`	`172`	`{`
`175`	`173`	`object? value = CallSiteRuntimeResolver.Instance.Resolve(callSite, Root);`
`176`		`- return scope => value;`
	`174`	`+ return scope =>`
	`175`	`+ {`
	`176`	`+ DependencyInjectionEventSource.Log.ServiceResolved(this, serviceType);`
	`177`	`+ return value;`
	`178`	`+ };`
`177`	`179`	`}`
`178`	`180`
`179`		`- return _engine.RealizeService(callSite);`
	`181`	`+ Func<ServiceProviderEngineScope, object?> realizedService = _engine.RealizeService(callSite);`
	`182`	`+ return scope =>`
	`183`	`+ {`
	`184`	`+ OnResolve(callSite, scope);`
	`185`	`+ DependencyInjectionEventSource.Log.ServiceResolved(this, serviceType);`
	`186`	`+ return realizedService(scope);`
	`187`	`+ };`
`180`	`188`	`}`
`181`	`189`
`182`	`190`	`return _ => null;`