Fix IAsyncDisposable handling on .NET 6 and double disposing

erisonliang · Jan 10, 2022 · 0ffaaef · 0ffaaef
1 parent 57bf1ca
commit 0ffaaef
Show file tree

Hide file tree

Showing 3 changed files with 99 additions and 15 deletions.
diff --git a/src/Cocona.Core/Command/Dispatcher/CoconaCommandDispatcher.cs b/src/Cocona.Core/Command/Dispatcher/CoconaCommandDispatcher.cs
@@ -44,12 +44,14 @@ public async ValueTask<int> DispatchAsync(CommandResolverResult commandResolverR
 
                 // Activate a command type.
                 var commandInstance = default(object);
+                var shouldCleanup = false;
                 if (matchedCommand.Target is not null)
                 {
                     commandInstance = matchedCommand.Target;
                 }
                 else if (matchedCommand.CommandType.GetConstructors().Any() && !matchedCommand.Method.IsStatic)
                 {
+                    shouldCleanup = true;
                     commandInstance = _activator.GetServiceOrCreateInstance(_serviceProvider, matchedCommand.CommandType);
                     if (commandInstance == null) throw new InvalidOperationException($"Unable to activate command type '{matchedCommand.CommandType.FullName}'");
                 }
@@ -66,16 +68,19 @@ public async ValueTask<int> DispatchAsync(CommandResolverResult commandResolverR
                 }
                 finally
                 {
-                    switch (commandInstance)
+                    if (shouldCleanup)
                     {
-#if NET5_0 || NETSTANDARD2_1
-                        case IAsyncDisposable asyncDisposable:
-                            await asyncDisposable.DisposeAsync().ConfigureAwait(false);
-                            break;
+                        switch (commandInstance)
+                        {
+#if NET5_0_OR_GREATER || NETSTANDARD2_1
+                            case IAsyncDisposable asyncDisposable:
+                                await asyncDisposable.DisposeAsync().ConfigureAwait(false);
+                                break;
 #endif
-                        case IDisposable disposable:
-                            disposable.Dispose();
-                            break;
+                            case IDisposable disposable:
+                                disposable.Dispose();
+                                break;
+                        }
                     }
                 }
             }

diff --git a/src/Cocona.Core/Command/Dispatcher/Middlewares/CoconaCommandInvokeMiddleware.cs b/src/Cocona.Core/Command/Dispatcher/Middlewares/CoconaCommandInvokeMiddleware.cs
@@ -55,13 +55,6 @@ public override async ValueTask<int> DispatchAsync(CommandDispatchContext ctx)
                 ExceptionDispatchInfo.Capture(ex.InnerException).Throw();
                 return 1; // NOTE: This statement is unreachable.
             }
-            finally
-            {
-                if (ctx.CommandTarget is IDisposable disposable)
-                {
-                    disposable.Dispose();
-                }
-            }
         }
     }
 }
diff --git a/test/Cocona.Test/Command/CommandDispatcher/DispatcherTest.cs b/test/Cocona.Test/Command/CommandDispatcher/DispatcherTest.cs
@@ -382,6 +382,92 @@ public async Task StaticSimpleSingleCommandDispatch()
             TestCommandStatic.Log[0].Should().Be($"{nameof(TestCommandStatic.Test)}:option0 -> alice");
         }
 
+        [Fact]
+        public async Task CommandInstance_Dispose_After_Dispatch()
+        {
+            var services = CreateDefaultServices<TestCommand_Dispose_After_Dispatch>(new string[] { });
+            services.AddSingleton<DisposeCounter>();
+            var serviceProvider = services.BuildServiceProvider();
+
+            var dispatcher = serviceProvider.GetService<ICoconaCommandDispatcher>();
+            var resolvedCommand = serviceProvider.GetRequiredService<ICoconaCommandResolver>().ParseAndResolve(
+                serviceProvider.GetRequiredService<ICoconaCommandProvider>().GetCommandCollection(),
+                serviceProvider.GetRequiredService<ICoconaCommandLineArgumentProvider>().GetArguments()
+            );
+            var result = await dispatcher.DispatchAsync(resolvedCommand);
+            serviceProvider.GetRequiredService<DisposeCounter>().Count.Should().Be(1);
+        }
+
+        public class DisposeCounter
+        {
+            public int Count { get; set; }
+        }
+
+        public class TestCommand_Dispose_After_Dispatch : IDisposable
+        {
+            private readonly DisposeCounter _counter;
+
+            public TestCommand_Dispose_After_Dispatch(DisposeCounter counter)
+            {
+                _counter = counter;
+            }
+
+            public void Hello() {}
+
+            void IDisposable.Dispose()
+            {
+                if (_counter.Count > 0)
+                {
+                    throw new InvalidOperationException("Dispose should be called only once.");
+                }
+
+                _counter.Count++;
+            }
+        }
+
+#if NET5_0_OR_GREATER || NETSTANDARD2_1
+        [Fact]
+        public async Task CommandInstance_DisposeAsync_After_Dispatch()
+        {
+            var services = CreateDefaultServices<TestCommand_DisposeAsync_After_Dispatch>(new string[] { });
+            services.AddSingleton<DisposeCounter>();
+            var serviceProvider = services.BuildServiceProvider();
+
+            var dispatcher = serviceProvider.GetService<ICoconaCommandDispatcher>();
+            var resolvedCommand = serviceProvider.GetRequiredService<ICoconaCommandResolver>().ParseAndResolve(
+                serviceProvider.GetRequiredService<ICoconaCommandProvider>().GetCommandCollection(),
+                serviceProvider.GetRequiredService<ICoconaCommandLineArgumentProvider>().GetArguments()
+            );
+            var result = await dispatcher.DispatchAsync(resolvedCommand);
+
+            serviceProvider.GetRequiredService<DisposeCounter>().Count.Should().Be(1);
+        }
+
+        public class TestCommand_DisposeAsync_After_Dispatch : IAsyncDisposable
+        {
+            private readonly DisposeCounter _counter;
+
+            public TestCommand_DisposeAsync_After_Dispatch(DisposeCounter counter)
+            {
+                _counter = counter;
+            }
+
+            public void Hello() { }
+
+            ValueTask IAsyncDisposable.DisposeAsync()
+            {
+                if (_counter.Count > 0)
+                {
+                    throw new InvalidOperationException("DisposeAsync should be called only once.");
+                }
+
+                _counter.Count++;
+
+                return default;
+            }
+        }
+#endif
+
         public class NoCommand
         { }