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?

4 comments
  1. 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.

    • John said:

      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.

    • John said:

      OMG, I’m embarrassed to say I didn’t know about reversed()! You’re right!

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Google photo

You are commenting using your Google account. Log Out /  Change )

Twitter picture

You are commenting using your Twitter account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s

This site uses Akismet to reduce spam. Learn how your comment data is processed.

%d bloggers like this: