Context managers and unforced errors

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?