In a previous post, I showed some sample code to implement a UsingAsync method to handle async cleanup. This is a useful pattern but it starts to break down as the number of items to be cleaned up grows beyond a few. It is also a bit unnatural as it relies on delegates to implement the inner block of the using scope.
So let’s look at a different approach which I have dubbed CleanupGuard. Astute readers may find some similarities between this and Andrei Alexandrescu and Petru Marginean’s ScopeGuard but there are several important differences (mostly due to the fact that we don’t have deterministic destructors in C#).
Why would you need CleanupGuard? Often you need to write code that creates various external resources such as child processes, temporary files on disk, or remote database entries. These resources will hang around (perhaps indefinitely) unless you attempt to clean up before you exit, successfully or otherwise.
Writing custom code to deal with this type of cleanup is rather cumbersome. Here is a first cut example that creates files and makes sure to zero them out before returning so that disk space is reclaimed.
public static async Task CreateFilesAsync(Func<string> nextFileName)
{
// Keep track of files we have created, for later cleanup.
List<string> files = new List<string>();
string fileName;
do
{
fileName = nextFileName();
if (fileName != null)
{
files.Add(fileName);
await WriteFileAsync(fileName);
}
}
while (fileName != null);
await CleanupAsync(files);
}
private static async Task CleanupAsync(IList<string> fileNames)
{
foreach (string filename in fileNames)
{
using (FileStream stream = CreateAsyncStream(fileName))
{
// After this, file will be zero length.
await stream.WriteAsync(new byte[0], 0, 0);
}
}
}
This is a noble attempt but it fails to take into account that errors could occur while writing the files. In this case no cleanup would be attempted — bad news! We can fix that by adding an appropriate try/catch block though unfortunately we can’t use a finally since the cleanup code is async.
public static async Task CreateFilesAsync(Func<string> nextFileName)
{
// Keep track of files we have created, for later cleanup.
List<string> files = new List<string>();
Exception exception = null;
try
{
// . . . code snipped for brevity . . .
}
catch (Exception e)
{
exception = e;
}
await CleanupAsync(files);
if (exception != null)
{
// Preserve original stack trace by wrapping exception.
throw new AggregateException(exception);
}
}
That’s better, but there is still another problem — cleanup can also fail, causing us to bail out early and leave files behind. We could attempt to fix that as well (it would involve placing a try/catch inside the cleanup loop and aggregating all exceptions) but it would only help us with this single example. What we would prefer is a simple and lightweight mechanism to handle this type of cleanup without a lot of custom code and explicit exception tracking.
Introducing CleanupGuard! Here is the full source code:
public class CleanupGuard
{
private readonly Stack<Func<Task>> steps;
public CleanupGuard()
{
this.steps = new Stack<Func<Task>>();
}
public void Register(Func<Task> cleanupAsync)
{
this.steps.Push(cleanupAsync);
}
public async Task RunAsync(Func<CleanupGuard, Task> doAsync)
{
List<Exception> exceptions = new List<Exception>();
try
{
await doAsync(this);
}
catch (Exception e)
{
exceptions.Add(e);
}
while (this.steps.Count > 0)
{
Func<Task> cleanupAsync = this.steps.Pop();
try
{
await cleanupAsync();
}
catch (Exception e)
{
exceptions.Add(e);
}
}
if (exceptions.Count > 0)
{
throw new AggregateException(exceptions).Flatten();
}
}
}
Basically, we maintain a Stack of registered async cleanup methods which are simply Func<Task> instances. The RunAsync method launches the user’s delegate, passing a reference to the CleanupGuard itself so the method can register cleanup steps as needed. On success or failure, the cleanup steps are popped and executed — last in, first out. Any exceptions that occur are deferred until the very end, ensuring that we always attempt to do as much cleanup as possible.
Let’s rewrite the example above using CleanupGuard:
public static Task CreateFilesAsync(Func<string> nextFileName)
{
CleanupGuard guard = new CleanupGuard();
return guard.RunAsync(g => CreateFilesInnerAsync(g, nextFileName));
}
private static async Task CreateFilesInnerAsync(CleanupGuard guard, Func<string> nextFileName)
{
string fileName;
do
{
fileName = nextFileName();
if (fileName != null)
{
guard.Register(() => ZeroFileAsync(fileName));
await WriteFileAsync(fileName);
}
}
while (fileName != null);
}
private static async Task ZeroFileAsync(string fileName)
{
using (FileStream stream = CreateAsyncStream(fileName))
{
// After this, file will be zero length.
await stream.WriteAsync(new byte[0], 0, 0);
}
}
Clean, concise, and no more bugs! On any error while creating or cleaning up files, CleanupGuard ensures that cleanup steps will run and that all exceptions are rethrown from RunAsync.
You can look at the GitHub commit history for the CleanupSample project to see the evolution of the code as the unit tests were written. The project includes a sample app showing a use case of creating processes and files and cleaning them up before exiting. The app periodically generates files and launches Notepad to display the files, registering cleanup steps with CleanupGuard along the way. When the method finishes, CleanupGuard springs into action and runs all the user’s cleanup steps, killing processes and zeroing out files. The sample app also demonstrates how one might handle non-async cleanup steps (e.g. killing a process) by essentially blocking and creating a dummy Task result. This is not generally a good practice but sometimes there is no better option.