tl;dr: Think about exceptions when writing a context manager.
I made a huge unforced error with a context manager at work.
We use Redis distributed locks for system synchronization. I wanted a context manager that acquired n locks, executed protected code, and then released the n locks in reverse order. It would be simple to use:
from common.util import Semaphore, distlock semaphore1 = Semaphore(OwnerDisambiguationUpdate.UPDATE_LOCK) semaphore2 = Semaphore(USMaintenanceFeeUpdate.UPDATE_LOCK) with distlock(semaphore1, semaphore2): do_some_work()
(The Semaphore class does other work with aborting Celery tasks, but that’s not germane here. It’s a Redis distributed lock with extra fanciness.)
I implemented distlock like so:
from contextlib import contextmanager @contextmanager def distlock(*locks): """Execute code within a Redis distributed lock. :param locks: One or more semaphores. They are acquired from 0 -> last, and released from last -> 0 :type locks: Objects with acquire and release methods """ for entry in locks: entry.acquire() yield for index in range(len(locks) - 1, -1, -1): locks[index].release()
Short and sweet, yeah? We acquire the locks, do the work, then release the locks.
This worked great. But over time I noticed that some periodic celery tasks weren’t launching because one or more locks weren’t clear.
The great thing about software development is this: You can put a gun to your head and pull the trigger and not feel the bullet until weeks, months, or years later. I had neglected to note this in the contextlib documentation:
If an unhandled exception occurs in the block, it is reraised inside the generator at the point where the yield occurred. Thus, you can use a try…except…finally statement to trap the error (if any), or ensure that some cleanup takes place.
My silly-ass goof was easy to fix:
@contextmanager def distlock(*locks): """Execute code within a Redis distributed lock. :param locks: One or more semaphores. They are acquired from 0 -> last, and released from last -> 0 :type locks: Objects with acquire and release methods """ for entry in locks: entry.acquire() # Run the protected code. The locks are released even if an # exception happens, which will propagate up the call stack. try: yield finally: for index in range(len(locks) - 1, -1, -1): locks[index].release()
Looks good, yeah?
Have you thought about enforcing canonical locking order inside the context manager? You’d need a comparable representation of Semaphore objects, but that should be easy, I think.
Because of how the code is organized, it already does that without explicitly defining the order somewhere. But that would be easy to add if it gets more complex.
Also, nitpick, but I’d probably use “for entry in reversed(locks):” for the finally block.
OMG, I’m embarrassed to say I didn’t know about reversed()! You’re right!