From 0e200d6c95736bd6f23e03cf18d498bb74dacaa7 Mon Sep 17 00:00:00 2001
From: Blake Niemyjski <bniemyjski@gmail.com>
Date: Wed, 2 Oct 2024 11:25:16 -0500
Subject: [PATCH] Added Async Disposable Action

---
 .../Utility/AsyncDisposableAction.cs          | 32 +++++++++++++++++++
 src/Foundatio/Utility/DisposableAction.cs     | 11 +++----
 2 files changed, 36 insertions(+), 7 deletions(-)
 create mode 100644 src/Foundatio/Utility/AsyncDisposableAction.cs

diff --git a/src/Foundatio/Utility/AsyncDisposableAction.cs b/src/Foundatio/Utility/AsyncDisposableAction.cs
new file mode 100644
index 00000000..cefa2d0f
--- /dev/null
+++ b/src/Foundatio/Utility/AsyncDisposableAction.cs
@@ -0,0 +1,32 @@
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Foundatio.Utility;
+
+/// <summary>
+/// A class that will call an <see cref="Func{TResult}"/> when Disposed.
+/// </summary>
+public sealed class AsyncDisposableAction : IAsyncDisposable
+{
+    private Func<Task> _exitTask;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="DisposableAction"/> class.
+    /// </summary>
+    /// <param name="exitTask">The exit action.</param>
+    public AsyncDisposableAction(Func<Task> exitTask)
+    {
+        _exitTask = exitTask;
+    }
+
+    /// <summary>
+    /// Performs application-defined tasks associated with freeing, releasing, or resetting unmanaged resources.
+    /// </summary>
+    async ValueTask IAsyncDisposable.DisposeAsync()
+    {
+        var exitAction = Interlocked.Exchange(ref _exitTask, null);
+        if (exitAction is not null)
+            await _exitTask().AnyContext();
+    }
+}
diff --git a/src/Foundatio/Utility/DisposableAction.cs b/src/Foundatio/Utility/DisposableAction.cs
index fe6f1bd4..aea8bd95 100644
--- a/src/Foundatio/Utility/DisposableAction.cs
+++ b/src/Foundatio/Utility/DisposableAction.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Threading;
 
 namespace Foundatio.Utility;
 
@@ -7,8 +8,7 @@ namespace Foundatio.Utility;
 /// </summary>
 public sealed class DisposableAction : IDisposable
 {
-    private readonly Action _exitAction;
-    private bool _disposed;
+    private Action _exitAction;
 
     /// <summary>
     /// Initializes a new instance of the <see cref="DisposableAction"/> class.
@@ -24,10 +24,7 @@ public DisposableAction(Action exitAction)
     /// </summary>
     void IDisposable.Dispose()
     {
-        if (_disposed)
-            return;
-
-        _exitAction();
-        _disposed = true;
+        var exitAction = Interlocked.Exchange(ref _exitAction, null);
+        exitAction?.Invoke();
     }
 }