advisory_locks.md

Using PostgreSQL Advisory Locks in Haskell

Recently two different administrative users of a Haskell app I maintain started a long-running background job at the same time. This resulted in some Bad Things happening which I would like to prevent in the future!

So I want a lock. This app runs on at least two servers at any given time, so it needs to be a distributed lock, and for Reasons (not relevant to this blogpost) I can't use our Redis cache. So the PostgreSQL database it is!

My first thought is that I could just make a table of "locks", write a row to the table, if the row is there, the job is locked. That probably would have worked just fine! But then I did a quick Google to see what other options were out there and learned about advisory locks. And these seem perfect for what I want to do.

PostgreSQL advisory locks

From the PostgreSQL documentation:

While a flag stored in a table could be used for the same purpose, advisory locks are faster, avoid table bloat, and are automatically cleaned up by the server at the end of the session.

This blog post is a great rundown of how these locks work.

So I got coding. And here's what my lock function in the Haskell app looks like:

withLock :: Pool Connection -> Text -> IO a -> IO a
withLock connPool lockName func =
  withResource connPool $ \conn ->
    withTransaction conn $ do
      let key = hash lockName
      locked <- query conn "SELECT pg_try_advisory_xact_lock(?)" (Only key)
      case listToMaybe locked of
        Just (Only True) -> func
        _                -> throw (LockException lockName)

Line-by-line

Let's go through it line-by-line:

withResource connPool $ \conn ->
  withTransaction conn $ do

We use withResource to get a connection from the database connection pool, then we use postgresql-simple's withTransaction to wrap the rest of the function in a database transaction.

That sounds a little scary at first. Does this mean everything that happens inside the lock is happening in the same transaction? So, if we're charging 500 credit cards and recording 500 receipts, and the 499th charge has an error, is it going to roll back 498 receipts!? That would be terrible!! 😱

Luckily that is not what happens! Only database calls using this connection conn will be a part of the transaction. We don't pass conn as an argument into func, so we know that database calls within that function are going to grab their own connections from the pool. Anything that those connections do will be outside this transaction.

let key = hash lockName

Next we hash the lockName to turn it into a integer that we can give to pg_try_advisory_xact_lock.

locked <- query conn "SELECT pg_try_advisory_xact_lock(?)" (Only key)

Now we'll try to take the lock. The postgresql-simple library's query function takes a connection to the database, a query, and an n-tuple of arguments for the query. In this case there's only one, but how does the compiler distinguish a 1-tuple from just... a value? The postgresql-simple has us wrap it in this Only constructor to make the types work.

We know that pg_try_advisory_xact_lock returns just a single boolean, but as far as postgresql-simple is concerned, any query can return a list of rows containing a list or tuple of columns. But we just want to know if that single result is true or false.

That's where listToMaybe locked comes in. listToMaybe on a list of "rows" will give us Just the first element or Nothing. So then we have Maybe the single result we're looking for.

We only want to call the function inside the lock if the query returned at least one row (it's Just something), if there was Only one field returned from the database, and if the lock succeeded (the result is True). Otherwise we need to do something else.

_                -> throw (LockException lockName)

Here, we're throwing a custom LockException. This will let us handle just these errors if we want to let the user know that a job is already in progress.

Here's what the definition of LockException looks like:

newtype LockException = LockException Text deriving Show
instance Exception LockException

Testing

I'm not going to attempt automated tests for this function, because these sorts of tests (trying to force two things to run exactly at once) are just way too flaky.

So what kind of test can I write for this to make sure it works? Here's what I did to test as I developed:

testLock :: Pool Connection -> IO ()
testLock conn = do
  forkIO $
    catch (withLock conn "test lock" $ do
      putStrLn "✅ Successfully took first lock ✅"
      threadDelay 1000000)
    (\(LockException _) -> putStrLn "🚫 Unable to take first lock 🚫")
  forkIO $
    catch (withLock conn "test lock" $ do
      putStrLn "🚫 Successfully took second lock 🚫")
    (\(LockException _) -> putStrLn "✅ Unable to take second lock ✅")
  forkIO $
    catch (withLock conn "a different lock" $ do
      threadDelay 1000000 -- this delay is just to prevent overlapping output
      putStrLn "✅ Successfully took third lock ✅")
    (\(LockException _) -> putStrLn "🚫 Unable to take third lock 🚫")
  threadDelay 1500000
  forkIO $
    catch (withLock conn "test lock" $ do
      putStrLn "✅ Successfully took fourth lock ✅")
    (\(LockException _) -> putStrLn "🚫 Unable to take fourth lock 🚫")
  return ()

We need to use forkIO to create two threads that can run more or less at once. They need to take a little time, so I add a threadDelay inside the first job to make sure they'll both try to take the lock. Next I'll make sure I can take a lock that uses a different key. Finally, we'll wait for the first job to finish and try to take a third lock.

This results in:

✅ Successfully took first lock ✅
✅ Unable to take second lock ✅
✅ Successfully took fourth lock ✅
✅ Successfully took third lock ✅

Yay!

What do you think?

I would love to know others' thoughts on this approach and how I can make it better. I'd also love to answer any questions.

Finally...

I'm a developer (and soon-to-be worker-owner) at Position Development, a worker-run software development company. If you would like to support our work, please let others know about us and consider working with us on your next project!