Watch out for this (async) generator cleanup pitfall in Python

Key Takeaways

hello and welcome back to M coding I'm James Murphy today we're talking about generators specifically generators and async generators that used the width statement or a try finally and a common Pitfall that you might fall into suppose we have a generator that needs some kind of resource it could be something like talking to a database where you need to start a transaction and then commit at the end of it or you need to grab a lock and then release the lock at the end of it or you need to open a file read some data from the file and then close the file at the end of it for rtion purposes this resource is just a dummy resource that has a name and prints out its own name on cleanup so that we can see when the cleanup happens but in real life it could be basically anything the normal way that we would use a generator like this is just to Loop over it so we would see yield zero and then got zero yield one and then got one yield two and then got two then we would exit the loop we would see the resource cleanup happen and then we would see the print statement after the loop and if we actually run it that's indeed what we see 0 1 2 clean up and then the stuff after the loop in this case the resource got cleaned up because we ran the generator past the point of when resource clean up would happen within the generator we exited the for Loop we exited the width block and we reached the end of the generator but here's a question for you does the resource still get cleaned up if we don't run the generator to completion we either stop it with a break or there's an exception that gets thrown and we never actually run the generator to the point where it would exit the WID block take a moment to think about it will this code still print the cleanup message answer coming in 3 2 1 interestingly yes the cleanup still happens we never got to iteration two in our for Loop but there's the cleanup message this behavior is possible because when a generator is garbage collected it automatically has its close method called the close method actually throws an exception inside the generator and then continues to run it until completion so instead of continuing with the for Loop it will exit the for Loop and exit the width statement which is why the cleanup happens we can actually observe this Behavior by surrounding the yield statement in a tri accept that catches a base exception and then printing it out before re raising it this allows us to kind of spy on the exception being raised but still raisee it so that it can do its job now when we run the code we see Loop zero Loop one and then indeed our Tri accept catches this generator exit exception but here's the tricky part and the source of the common Pitfall remember what I said causes this cleanup to actually happen it's the garbage collection of the generator object that triggers the cleanup in cpython if there are no more references to an object it'll get cleaned up immediately so in this case because we know garbage collection happens here we know that our resource will be cleaned up before anything that happens after the loop but what if we did some refactoring and pulled the generator out into a variable since we still maintain a reference to it even after the loop the generator won't be garbage collected here but rather at the end of the main function therefore our resource doesn't get cleaned up at the same place it gets cleaned up at the end of main this could already be a problem if the stuff after the loop expected that resource to have been cleaned up like if it was holding a lock and the stuff afterwards also needs to use the lock even if we didn't Factor things out into a separate variable like we did here other implementations of python don't guarantee the same behavior with garbage collection and they may not garbage collect something immediately even if its ref count is zero this non-determinism in generator resource cleanup is especially bad when you have nested resources just the way the code looks programmers generally tend to assume that inner resources will get cleaned up before outer resources but in this case that's not true the outer resource is cleaned up at the end of the width block but the inner resource isn't cleaned up until the generator is garbage collected which happens after the WID block at the end of main you can guarantee the cleanup happens earlier by manually triggering it by calling the generator's close function technically though this still isn't good enough because something in here could throw an exception and we'd miss that close call in order to guarantee clean up of the inner resources before the outer ones you can use something like contextlib closing which is a context manager that will just call that generator close function with this approach the place where the generator is garbage collected becomes irrelevant to the cleanup of the resources but be warned code in the wild very rarely actually does this it's super annoying to have an extra width statement around every for Loop that uses a generator so in the wild you typically just see people depending on the garbage collection behavior of cpython remember in cpython which is the python that 99% of people people are using the one that you download from python.org if you just don't keep a reference to this generator then it will be garbage collected and the cleanup of the inner thing will happen before the outer thing it's just very delicate and Flaky code where if someone comes along and does a slight refactor it might break and for normal generators that's kind of a good enough solution and the end of the story just don't hold on to a reference to a generator or if you do use context li. closing to ensure it gets cleaned up but if you happen to be using async generators then the problem gets much worse let's go back up to our generator and make this an async def function we could be using an async width block or width block or Tri finally it doesn't really matter here whether the resource is async or not what's going to make this worse is just that the generator is async and then back in main let's go ahead and make main an async depth function and use an async for Loop to Loop over the elements in the generator and we now run our main function using async i.run so this is the async equivalent of the problem we had before and the generator is garbage collected at this point but when we run the code we unfortunately see that the outer resource is cleaned up before the inner one the inner one does still get cleaned up in a similar way as in the synchronous case a special exception gets thrown inside the generator and then it's run until completion which causes the WID block to eventually exit but for some reason in this async case the garbage collection isn't triggering the resource to be cleaned up just after the for Loop like it was in the synchronous case and as you could probably guess the delaying has something to do with it being async we can get a little hint about what's really going on here by reaching into the internals of the event Loop that's running and looking at what the ready tasks are this print statement just before the break is sort of our Baseline then right after the break it will go to this print statement which is outside the for Loop the first one is just before the generator is garbage collected and the second one is just after the generator is garbage collected when we run the code and see the before and after we see that there's this extra thing in the after and it doesn't take a huge leap to figure out that yes this is the task that closes the generator and there in lies the Crux of the problem an async generator is async but the garbage collector is not async so even though the generator is garbage collected here the best the garbage collector could do is schedule the task that causes the generator to be closed it can't actually await that task though if we manually sleep a few times we give the event Loop enough time to run any pending tasks now when we run the code we see that indeed the inner one gets cleaned up before the outer one which is what we want so how do we get around this obviously I'm not suggesting that you add sleep statements after every for Loop similar to the synchronous case the most correct thing to do that almost nobody actually does is to surround any async for Loop using a generator in an async with block where you use context Libs a closing all this does is await the generator's a close method on the way out just like with the synchronous case this makes the point where the generator is garbage collected irrelevant to when the cleanup happens but this is kind of unsatisfying for a number of reasons first off I mean good luck convinc everyone to do this in the first place but it's also unsatisfying because it breaks a level of abstraction a closing only works for async generators not arbitrary async iterables so immediately I now have to keep track of whether this is an async generator or just an async iterable that's something else and second since the pragmatic thing to do is not to put an async width around every async for Loop but only to put it around the ones that need it then in addition to keeping track of whether or not this is an async generator you also need to know about the implementation of the generator and whether or not that implementation uses any Tri finales width statements or async width statements not to mention the additional refactoring cost because remember whenever a generator uses a tri finally a width statement or an asyn width statement then you need to wrap all of its usages in an async WID block but now all those callers have an async WID block so if any of them are generators you now have to wrap all their colors in asyn blocks and this process can get out of hand very quickly the real solution is of course to completely ban asnc for just kidding but yeah unfortunately this is sort of just the State of Affairs right now if you care about asnc generator clean up you just have to use a closing there is a pep Number 533 that proposes a solution to this though the idea is pretty simple basically just add two new dunders a iter close and iter close and change the meaning of for Loops have an async for Loop await a it close at the end of the for loop after a regular exit breaking or an exception and have an normal for Loop call inter close again when you either Exit normally exit with a break or exit with an exception this is a pretty good idea and if python was being written for the very first time then this might have made it in but as you can imagine changing the semantics of iteration is probably a no-go at this point not to mention how hard it would be to implement and it would be a breaking change to any code that intends to use the current semantics where it has a break statement but then continues to iterate over the same generator or iterable afterwards in this case we open a file process some header lines in the file and then process the rest of the lines in a different way so we would break and then continue processing the rest of the lines technically open doesn't return a generator so lines in this case is not a generator but just imagine you've written your own kind of open wrapper that uses a generator in that case it's perfectly reasonable that the developer intended it to be this way that they didn't want lines to be closed after the first for Loop so anyway I don't think this a itose and itose pep will ever be added into python aside from outright Banning the use of async generators in your code which I guess technically is a fine solution the best thing to do is to just be aware of it and add this to your list of quirks of python that you need to watch out for once again I'm James Murphy and this is M coding I'm available for code reviews Contracting and Consulting check out my company at M coding. and thank you to my patrons and donors for helping to support the channel don't forget to leave a comment subscribe and slap that like button in odd number times see you in the next one