From 0432255ae09c89500253c38ca32d573d9b061978 Mon Sep 17 00:00:00 2001
From: Artur <artur@dorochowicz.com>
Date: Tue, 24 Sep 2024 18:57:55 +0000
Subject: [PATCH] Use DisposeAsync with scopes and hosts

---
 .../DependencyInjectionCommandCreator.cs      | 11 ++-----
 src/Oakton/HostWrapperCommand.cs              | 30 ++++++++++++++-----
 src/Oakton/HostedCommandExtensions.cs         | 11 +++++--
 3 files changed, 34 insertions(+), 18 deletions(-)

diff --git a/src/Oakton/DependencyInjectionCommandCreator.cs b/src/Oakton/DependencyInjectionCommandCreator.cs
index 97826235..555952cf 100644
--- a/src/Oakton/DependencyInjectionCommandCreator.cs
+++ b/src/Oakton/DependencyInjectionCommandCreator.cs
@@ -34,12 +34,12 @@ 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);
     }
 
@@ -47,16 +47,11 @@ public WrappedOaktonCommand(IServiceProvider provider, Type commandType)
     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
index abbaaa59..b7bcea27 100644
--- 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();
+            }
+        }
     }
 }
\ No newline at end of file
diff --git a/src/Oakton/HostedCommandExtensions.cs b/src/Oakton/HostedCommandExtensions.cs
index 4702e2a7..27698532 100644
--- 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();
+            }
         }
     }