Fix validation on build

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

@@ -10,21 +10,22 @@ 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[GetCacheKey(callSite)] = scoped;
             }
         }
 
-        public void ValidateResolution(Type serviceType, IServiceScope scope, IServiceScope rootScope)
+        public void ValidateResolution(ServiceCallSite callSite, IServiceScope scope, IServiceScope rootScope)
         {
+            Type serviceType = callSite.ServiceType;
             if (ReferenceEquals(scope, rootScope)
-                && _scopedServices.TryGetValue(serviceType, out Type? scopedService))
+                && _scopedServices.TryGetValue(GetCacheKey(callSite), out Type? scopedService))
             {
                 if (serviceType == scopedService)
                 {
@@ -97,6 +98,12 @@ public void ValidateResolution(Type serviceType, IServiceScope scope, IServiceSc
 
         protected override Type? VisitFactory(FactoryCallSite factoryCallSite, CallSiteValidatorState state) => null;
 
+        private static ServiceCacheKey GetCacheKey(ServiceCallSite callSite)
+        {
+            return callSite.Cache.Key.Equals(ServiceCacheKey.Empty)
+                ? new ServiceCacheKey(callSite.ServiceType, 0) : callSite.Cache.Key;
+        }
+
         internal struct CallSiteValidatorState
         {
             [DisallowNull]

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.Invoke(scope);
+                };
             }
 
             return _ => null;

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)