Add AsyncDisposable support (#1005)

src/DependencyInjection/DI/src/Microsoft.Extensions.DependencyInjection.csproj

@@ -2,12 +2,13 @@
 
   <PropertyGroup>
     <Description>Default implementation of dependency injection for Microsoft.Extensions.DependencyInjection.</Description>
-    <TargetFrameworks>netcoreapp2.0;net461;netstandard2.0</TargetFrameworks>
+    <TargetFrameworks>netcoreapp3.0;net461;netstandard2.0</TargetFrameworks>
     <GenerateDocumentationFile>true</GenerateDocumentationFile>
     <PackageTags>dependencyinjection;di</PackageTags>
 
     <ILEmitBackend Condition="$(TargetFramework) != 'netstandard2.0'">True</ILEmitBackend>
     <DefineConstants Condition="'$(ILEmitBackend)' == 'True'">$(DefineConstants);IL_EMIT</DefineConstants>
+    <DefineConstants Condition="'$(TargetFramework)' == 'netcoreapp3.0'">$(DefineConstants);DISPOSE_ASYNC</DefineConstants>
 
     <!-- Debug IL generation -->
     <ILEmitBackendSaveAssemblies>False</ILEmitBackendSaveAssemblies>

src/DependencyInjection/DI/src/Resources.resx

@@ -156,4 +156,7 @@
   <data name="ImplementationTypeCantBeConvertedToServiceType" xml:space="preserve">
     <value>Implementation type '{0}' can't be converted to service type '{1}'</value>
   </data>
+  <data name="AsyncDisposableServiceDispose" xml:space="preserve">
+    <value>'{0}' type only implements IAsyncDisposable. Use DisposeAsync to dispose the container.</value>
+  </data>
 </root>

src/DependencyInjection/DI/src/ServiceLookup/Expressions/ExpressionResolverBuilder.cs

@@ -147,24 +147,21 @@ protected override Expression VisitIEnumerable(IEnumerableCallSite callSite, obj
 
         protected override Expression VisitDisposeCache(ServiceCallSite callSite, object context)
         {
-            var implType = callSite.ImplementationType;
             // Elide calls to GetCaptureDisposable if the implementation type isn't disposable
             return TryCaptureDisposable(
-                implType,
+                callSite,
                 ScopeParameter,
                 VisitCallSiteMain(callSite, context));
         }
 
-        private Expression TryCaptureDisposable(Type implType, ParameterExpression scope, Expression service)
+        private Expression TryCaptureDisposable(ServiceCallSite callSite, ParameterExpression scope, Expression service)
         {
-            if (implType != null &&
-                !typeof(IDisposable).GetTypeInfo().IsAssignableFrom(implType.GetTypeInfo()))
+            if (!callSite.CaptureDisposable)
             {
                 return service;
             }
 
-            return Expression.Invoke(GetCaptureDisposable(scope),
-                service);
+            return Expression.Invoke(GetCaptureDisposable(scope), service);
         }
 
         protected override Expression VisitConstructor(ConstructorCallSite callSite, object context)
@@ -221,7 +218,7 @@ private Expression BuildScopedExpression(ServiceCallSite callSite)
                 keyExpression,
                 resolvedVariable);
 
-            var captureDisposible = TryCaptureDisposable(callSite.ImplementationType, ScopeParameter, VisitCallSiteMain(callSite, null));
+            var captureDisposible = TryCaptureDisposable(callSite, ScopeParameter, VisitCallSiteMain(callSite, null));
 
             var assignExpression = Expression.Assign(
                 resolvedVariable,

src/DependencyInjection/DI/src/ServiceLookup/ILEmit/ILEmitResolverBuilder.cs

@@ -132,15 +132,16 @@ private GeneratedMethod BuildTypeNoCache(ServiceCallSite callSite)
 
         protected override object VisitDisposeCache(ServiceCallSite transientCallSite, ILEmitResolverBuilderContext argument)
         {
-            // RuntimeScope.CaptureDisposables([create value])
-            var shouldCapture = BeginCaptureDisposable(transientCallSite.ImplementationType, argument);
-
-            VisitCallSiteMain(transientCallSite, argument);
-
-            if (shouldCapture)
+            if (transientCallSite.CaptureDisposable)
             {
+                BeginCaptureDisposable(argument);
+                VisitCallSiteMain(transientCallSite, argument);
                 EndCaptureDisposable(argument);
             }
+            else
+            {
+                VisitCallSiteMain(transientCallSite, argument);
+            }
             return null;
         }
 
@@ -353,8 +354,9 @@ private ILEmitResolverBuilderRuntimeContext GenerateMethodBody(ServiceCallSite c
                 VisitCallSiteMain(callSite, context);
                 Stloc(context.Generator, resultLocal.LocalIndex);
 
-                if (BeginCaptureDisposable(callSite.ImplementationType, context))
+                if (callSite.CaptureDisposable)
                 {
+                    BeginCaptureDisposable(context);
                     Ldloc(context.Generator, resultLocal.LocalIndex);
                     EndCaptureDisposable(context);
                     // Pop value returned by CaptureDisposable off the stack
@@ -408,17 +410,9 @@ private ILEmitResolverBuilderRuntimeContext GenerateMethodBody(ServiceCallSite c
             };
         }
 
-        private static bool BeginCaptureDisposable(Type implType, ILEmitResolverBuilderContext argument)
+        private static void BeginCaptureDisposable(ILEmitResolverBuilderContext argument)
         {
-            var shouldCapture = implType == null || typeof(IDisposable).GetTypeInfo().IsAssignableFrom(implType.GetTypeInfo());
-
-            if (shouldCapture)
-            {
-                // context
-                argument.Generator.Emit(OpCodes.Ldarg_1);
-            }
-
-            return shouldCapture;
+            argument.Generator.Emit(OpCodes.Ldarg_1);
         }
 
         private static void EndCaptureDisposable(ILEmitResolverBuilderContext argument)

src/DependencyInjection/DI/src/ServiceLookup/IServiceProviderEngine.cs

@@ -5,7 +5,10 @@
 
 namespace Microsoft.Extensions.DependencyInjection.ServiceLookup
 {
-    internal interface IServiceProviderEngine : IDisposable, IServiceProvider
+    internal interface IServiceProviderEngine : IServiceProvider, IDisposable
+#if DISPOSE_ASYNC
+        , IAsyncDisposable
+#endif
     {
         IServiceScope RootScope { get; }
         void ValidateService(ServiceDescriptor descriptor);

src/DependencyInjection/DI/src/ServiceLookup/ServiceCallSite.cs

@@ -2,6 +2,7 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
+using System.Reflection;
 
 namespace Microsoft.Extensions.DependencyInjection.ServiceLookup
 {
@@ -19,5 +20,18 @@ protected ServiceCallSite(ResultCache cache)
         public abstract Type ImplementationType { get; }
         public abstract CallSiteKind Kind { get; }
         public ResultCache Cache { get; }
+
+        public bool CaptureDisposable
+        {
+            get
+            {
+                return ImplementationType == null ||
+                       typeof(IDisposable).GetTypeInfo().IsAssignableFrom(ImplementationType.GetTypeInfo())
+#if DISPOSE_ASYNC
+                       || typeof(IAsyncDisposable).GetTypeInfo().IsAssignableFrom(ImplementationType.GetTypeInfo())
+#endif
+                    ;
+            }
+        }
     }
 }

src/DependencyInjection/DI/src/ServiceLookup/ServiceProviderEngine.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Threading.Tasks;
 
 namespace Microsoft.Extensions.DependencyInjection.ServiceLookup
 {
@@ -68,6 +69,14 @@ public void Dispose()
             Root.Dispose();
         }
 
+#if DISPOSE_ASYNC
+        public ValueTask DisposeAsync()
+        {
+            _disposed = true;
+            return Root.DisposeAsync();
+        }
+#endif
+
         internal object GetService(Type serviceType, ServiceProviderEngineScope serviceProviderEngineScope)
         {
             if (_disposed)

src/DependencyInjection/DI/src/ServiceLookup/ServiceProviderEngineScope.cs

@@ -3,6 +3,9 @@
 
 using System;
 using System.Collections.Generic;
+using System.Diagnostics;
+using System.Threading.Tasks;
+using Microsoft.Extensions.Internal;
 
 namespace Microsoft.Extensions.DependencyInjection.ServiceLookup
 {
@@ -11,7 +14,7 @@ internal class ServiceProviderEngineScope : IServiceScope, IServiceProvider
         // For testing only
         internal Action<object> _captureDisposableCallback;
 
-        private List<IDisposable> _disposables;
+        private List<object> _disposables;
 
         private bool _disposed;
 
@@ -36,51 +39,95 @@ public object GetService(Type serviceType)
 
         public IServiceProvider ServiceProvider => this;
 
-        public void Dispose()
+        internal object CaptureDisposable(object service)
         {
+            Debug.Assert(!_disposed);
+
+            _captureDisposableCallback?.Invoke(service);
+
+            if (ReferenceEquals(this, service) ||
+               !(service is IDisposable
+#if DISPOSE_ASYNC
+                || service is IAsyncDisposable
+#endif
+                ))
+            {
+                return service;
+            }
+
             lock (ResolvedServices)
             {
-                if (_disposed)
+                if (_disposables == null)
                 {
-                    return;
+                    _disposables = new List<object>();
                 }
 
-                _disposed = true;
-                if (_disposables != null)
+                _disposables.Add(service);
+            }
+            return service;
+        }
+
+        public void Dispose()
+        {
+            var toDispose = BeginDispose();
+
+            if (toDispose != null)
+            {
+                for (var i = toDispose.Count - 1; i >= 0; i--)
                 {
-                    for (var i = _disposables.Count - 1; i >= 0; i--)
+                    if (toDispose[i] is IDisposable disposable)
                     {
-                        var disposable = _disposables[i];
                         disposable.Dispose();
                     }
-
-                    _disposables.Clear();
+                    else
+                    {
+                        throw new InvalidOperationException(Resources.FormatAsyncDisposableServiceDispose(TypeNameHelper.GetTypeDisplayName(toDispose[i])));
+                    }
                 }
-
-                ResolvedServices.Clear();
             }
         }
 
-        internal object CaptureDisposable(object service)
+#if DISPOSE_ASYNC
+        public async ValueTask DisposeAsync()
         {
-            _captureDisposableCallback?.Invoke(service);
+            var toDispose = BeginDispose();
 
-            if (!ReferenceEquals(this, service))
+            if (toDispose != null)
             {
-                if (service is IDisposable disposable)
+                for (var i = toDispose.Count - 1; i >= 0; i--)
                 {
-                    lock (ResolvedServices)
+                    var disposable = toDispose[i];
+                    if (disposable is IAsyncDisposable asyncDisposable)
                     {
-                        if (_disposables == null)
-                        {
-                            _disposables = new List<IDisposable>();
-                        }
-
-                        _disposables.Add(disposable);
+                        await asyncDisposable.DisposeAsync();
+                    }
+                    else
+                    {
+                        ((IDisposable)disposable).Dispose();
                     }
                 }
             }
-            return service;
+        }
+#endif
+
+        private List<object> BeginDispose()
+        {
+            List<object> toDispose;
+            lock (ResolvedServices)
+            {
+                if (_disposed)
+                {
+                    return null;
+                }
+
+                _disposed = true;
+                toDispose = _disposables;
+                _disposables = null;
+
+                ResolvedServices.Clear();
+            }
+
+            return toDispose;
         }
     }
-}
+}

src/DependencyInjection/DI/src/ServiceProvider.cs

@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Generic;
 using System.Runtime.CompilerServices;
+using System.Threading.Tasks;
 using Microsoft.Extensions.DependencyInjection.ServiceLookup;
 
 namespace Microsoft.Extensions.DependencyInjection
@@ -12,6 +13,9 @@ namespace Microsoft.Extensions.DependencyInjection
     /// The default IServiceProvider.
     /// </summary>
     public sealed class ServiceProvider : IServiceProvider, IDisposable, IServiceProviderEngineCallback
+#if DISPOSE_ASYNC
+        , IAsyncDisposable
+#endif
     {
         private readonly IServiceProviderEngine _engine;
 
@@ -92,7 +96,10 @@ internal ServiceProvider(IEnumerable<ServiceDescriptor> serviceDescriptors, Serv
         public object GetService(Type serviceType) => _engine.GetService(serviceType);
 
         /// <inheritdoc />
-        public void Dispose() => _engine.Dispose();
+        public void Dispose()
+        {
+            _engine.Dispose();
+        }
 
         void IServiceProviderEngineCallback.OnCreate(ServiceCallSite callSite)
         {
@@ -103,5 +110,13 @@ void IServiceProviderEngineCallback.OnResolve(Type serviceType, IServiceScope sc
         {
             _callSiteValidator.ValidateResolution(serviceType, scope, _engine.RootScope);
         }
+
+#if DISPOSE_ASYNC
+        /// <inheritdoc/>
+        public ValueTask DisposeAsync()
+        {
+            return _engine.DisposeAsync();
+        }
+#endif
     }
 }