Thread-safe mysql-simple usage

Database.MySQL.hs

module Database.MySQL
  ( Connected
  , withConnection
  , connected
  ) where

import Control.Concurrent (MVar, newMVar, newEmptyMVar, putMVar, takeMVar, tryPutMVar, modifyMVar_)
import Control.Concurrent.Async (withAsyncBound, withAsync, wait, waitEither)
import Control.Exception (SomeException, SomeAsyncException, fromException, bracket, throwIO, tryJust, tryJust)
import Control.Monad (unless, forever)
import System.IO.Unsafe (unsafePerformIO)

import qualified Database.MySQL.Base as M (initLibrary, initThread, endThread)
import qualified Database.MySQL.Simple as M (Connection, ConnectInfo, connect, close)

-- Run an action in the connection thread.
newtype Connected = Connected
  { withConnection :: forall a. (M.Connection -> IO a) -> IO a
  }

{-# NOINLINE initialised #-}
initialised :: MVar Bool
initialised = unsafePerformIO (newMVar False)

-- | Initialise the C library only once.
initialise :: IO ()
initialise = modifyMVar_ initialised $ \done -> do
  unless done M.initLibrary
  return True

-- | Thread-safe conenction to MySQL.
connected :: M.ConnectInfo -> (Connected -> IO a) -> IO a
connected cinfo f = do
  initialise
  -- tryPutMVar at the end avoids BlockedIndefinitelyOnMVar.
  bracket newEmptyMVar (`tryPutMVar` const (return ())) $ \varAction -> do
    let c = Connected (schedule varAction)
    withAsyncBoundThrow (runner varAction) (f c)
  where
  -- schedule an action to be executed. MVar fairness guarantees FIFO of scheduled actions.
  schedule :: MVar (M.Connection -> IO ()) -> (forall a. (M.Connection -> IO a) -> IO a)
  schedule varAction act = do
    varResult <- newEmptyMVar :: IO (MVar (Either SomeException a))
    let action conn = putMVar varResult =<< try_ (act conn)
    putMVar varAction action
    r <- takeMVar varResult
    case r of
      Right v -> return v
      Left err -> throwIO err

  -- run actions in bound thread.
  runner :: MVar (M.Connection -> IO ()) -> IO ()
  runner varAction =
    bracket open close $ \conn -> do
    forever $ do
      act <- takeMVar varAction
      act conn

  open = do
    M.initThread
    M.connect cinfo

  close conn = do
    M.close conn
    M.endThread

-- | Catch only exceptions thrown by the action itself.
try_ :: IO a -> IO (Either SomeException a)
try_ = tryJust f
  where
  f ex =
    case fromException ex of
      Just (_ :: SomeAsyncException) -> Nothing
      Nothing -> Just ex

-- | Version of withAsyncBound which throws if background task throws.
withAsyncBoundThrow :: IO a -> IO b -> IO b
withAsyncBoundThrow left right =
  withAsyncBound right $ \r ->
  withAsync left $ \l -> do
  e <- waitEither l r
  case e of
    Left _ -> wait r
    Right v -> return v

I use it as such:

import Database.MySQL (connected, withConnection)
...
    connected connInfo $ \c -> do
      r <- withConnection c $ \conn -> query conn q params
      ...