Use DisposeAsync with scopes and hosts

JasperFx · Sep 25, 2024 · 0432255 · 0432255
1 parent e5205d4
commit 0432255
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 18 deletions.
diff --git a/src/Oakton/DependencyInjectionCommandCreator.cs b/src/Oakton/DependencyInjectionCommandCreator.cs
@@ -34,29 +34,24 @@ public object CreateModel(Type modelType)
 
 internal class WrappedOaktonCommand : IOaktonCommand
 {
-    private readonly IServiceScope _scope;
+    private readonly AsyncServiceScope _scope;
     private readonly IOaktonCommand _inner;
 
     public WrappedOaktonCommand(IServiceProvider provider, Type commandType)
     {
-        _scope = provider.CreateScope();
+        _scope = provider.CreateAsyncScope();
         _inner = (IOaktonCommand)_scope.ServiceProvider.GetRequiredService(commandType);
     }
 
     public Type InputType => _inner.InputType;
     public UsageGraph Usages => _inner.Usages;
     public async Task<bool> Execute(object input)
     {
-        try
+        await using (_scope)
         {
             // Execute your actual command
             return await _inner.Execute(input);
         }
-        finally
-        {
-            // Make sure the entire scope is disposed
-            _scope.SafeDispose();
-        }
     }
 }
 

diff --git a/src/Oakton/HostWrapperCommand.cs b/src/Oakton/HostWrapperCommand.cs
@@ -24,15 +24,29 @@ public HostWrapperCommand(IOaktonCommand inner, Func<IHost> hostSource, Property
     public UsageGraph Usages => _inner.Usages;
     public async Task<bool> Execute(object input)
     {
-        using var host = _hostSource();
-        using var scope = host.Services.CreateScope();
-        foreach (var prop in _props)
+        var host = _hostSource();
+        try
         {
-            var serviceType = prop.PropertyType;
-            var service = scope.ServiceProvider.GetRequiredService(serviceType);
-            prop.SetValue(_inner, service);
-        }
+            await using var scope = host.Services.CreateAsyncScope();
+            foreach (var prop in _props)
+            {
+                var serviceType = prop.PropertyType;
+                var service = scope.ServiceProvider.GetRequiredService(serviceType);
+                prop.SetValue(_inner, service);
+            }
 
-        return await _inner.Execute(input);
+            return await _inner.Execute(input);
+        }
+        finally
+        {
+            if (host is IAsyncDisposable ad)
+            {
+                await ad.DisposeAsync();
+            }
+            else
+            {
+                host.Dispose();
+            }
+        }
     }
 }
diff --git a/src/Oakton/HostedCommandExtensions.cs b/src/Oakton/HostedCommandExtensions.cs
@@ -67,7 +67,7 @@ public static async Task<int> RunHostedOaktonCommandsAsync(this IHost host, stri
     {
         try
         {
-            using var scope = host.Services.CreateScope();
+            await using var scope = host.Services.CreateAsyncScope();
             var options = scope.ServiceProvider.GetRequiredService<IOptions<OaktonOptions>>().Value;
             args = ApplyArgumentDefaults(args, options);
 
@@ -91,7 +91,14 @@ public static async Task<int> RunHostedOaktonCommandsAsync(this IHost host, stri
         }
         finally
         {
-            host.SafeDispose();
+            if (host is IAsyncDisposable ad)
+            {
+                await ad.DisposeAsync();
+            }
+            else
+            {
+                host.Dispose();
+            }
         }
     }