diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c) 2015, Scrive + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Scrive nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/consumers.cabal b/consumers.cabal new file mode 100644 --- /dev/null +++ b/consumers.cabal @@ -0,0 +1,52 @@ +name: consumers +version: 0.1 +synopsis: Concurrent PostgreSQL data consumers +description: Library for setting up concurrent consumers of data stored inside PostgreSQL database in a simple, declarative manner. +homepage: https://github.com/scrive/consumers +license: BSD3 +license-file: LICENSE +author: Scrive +maintainer: Andrzej Rybczak +category: Concurrency, Database +build-type: Simple +cabal-version: >=1.10 + +Source-repository head + Type: git + Location: https://github.com/scrive/consumers + +library + exposed-modules: Database.PostgreSQL.Consumers, + Database.PostgreSQL.Consumers.Config, + Database.PostgreSQL.Consumers.Consumer, + Database.PostgreSQL.Consumers.Components, + Database.PostgreSQL.Consumers.Utils + + build-depends: base <5, + containers, + exceptions, + hpqtypes >=1.4, + lifted-base, + lifted-threads, + log, + monad-control, + mtl, + stm, + time, + transformers-base + + hs-source-dirs: src + + ghc-options: -O2 -Wall -funbox-strict-fields + + default-language: Haskell2010 + default-extensions: + FlexibleContexts, + GeneralizedNewtypeDeriving, + OverloadedStrings, + RankNTypes, + RecordWildCards, + ScopedTypeVariables, + TupleSections, + TypeFamilies, + UndecidableInstances diff --git a/src/Database/PostgreSQL/Consumers.hs b/src/Database/PostgreSQL/Consumers.hs new file mode 100644 --- /dev/null +++ b/src/Database/PostgreSQL/Consumers.hs @@ -0,0 +1,9 @@ +module Database.PostgreSQL.Consumers ( + runConsumer + , module Database.PostgreSQL.Consumers.Config + , module Database.PostgreSQL.Consumers.Utils + ) where + +import Database.PostgreSQL.Consumers.Components +import Database.PostgreSQL.Consumers.Config +import Database.PostgreSQL.Consumers.Utils diff --git a/src/Database/PostgreSQL/Consumers/Components.hs b/src/Database/PostgreSQL/Consumers/Components.hs new file mode 100644 --- /dev/null +++ b/src/Database/PostgreSQL/Consumers/Components.hs @@ -0,0 +1,339 @@ +module Database.PostgreSQL.Consumers.Components ( + runConsumer + , spawnListener + , spawnMonitor + , spawnDispatcher + ) where + +import Control.Applicative +import Control.Concurrent.Lifted +import Control.Concurrent.STM hiding (atomically) +import Control.Exception (AsyncException(ThreadKilled)) +import Control.Monad +import Control.Monad.Base +import Control.Monad.Catch +import Control.Monad.Trans +import Control.Monad.Trans.Control +import Data.Function +import Data.Int +import Data.Monoid +import Data.Monoid.Utils +import Database.PostgreSQL.PQTypes +import Log +import Log.Class.Instances () +import qualified Control.Concurrent.STM as STM +import qualified Control.Concurrent.Thread.Lifted as T +import qualified Data.Foldable as F +import qualified Data.Map.Strict as M + +import Database.PostgreSQL.Consumers.Config +import Database.PostgreSQL.Consumers.Consumer +import Database.PostgreSQL.Consumers.Utils + +-- | Run the consumer. The purpose of the returned monadic +-- action is to wait for currently processed jobs and clean up. +-- This function is best used in conjunction with 'finalize' to +-- seamlessly handle the finalization. +runConsumer + :: (MonadBaseControl IO m, MonadLog m, MonadMask m, Eq idx, Show idx, ToSQL idx) + => ConsumerConfig m idx job + -> ConnectionSource + -> m (m ()) +runConsumer cc cs = do + semaphore <- newMVar () + runningJobsInfo <- liftBase $ newTVarIO M.empty + runningJobs <- liftBase $ newTVarIO 0 + + cid <- registerConsumer cc cs + localData ["consumer_id" .= show cid] $ do + listener <- spawnListener cc cs semaphore + monitor <- localDomain "monitor" $ spawnMonitor cc cs cid + dispatcher <- localDomain "dispatcher" $ spawnDispatcher cc cs cid semaphore runningJobsInfo runningJobs + return . localDomain "finalizer" $ do + stopExecution listener + stopExecution dispatcher + waitForRunningJobs runningJobsInfo runningJobs + stopExecution monitor + unregisterConsumer cc cs cid + where + waitForRunningJobs runningJobsInfo runningJobs = do + initialJobs <- liftBase $ readTVarIO runningJobsInfo + (`fix` initialJobs) $ \loop jobsInfo -> do + -- If jobs are still running, display info about them. + when (not $ M.null jobsInfo) $ do + logInfo "Waiting for running jobs" $ object [ + "job_id" .= showJobsInfo jobsInfo + ] + join . atomically $ do + jobs <- readTVar runningJobs + if jobs == 0 + then return $ return () + else do + newJobsInfo <- readTVar runningJobsInfo + -- If jobs info didn't change, wait for it to change. + -- Otherwise loop so it either displays the new info + -- or exits if there are no jobs running anymore. + if (newJobsInfo == jobsInfo) + then retry + else return $ loop newJobsInfo + where + showJobsInfo = M.foldr (\idx acc -> show idx : acc) [] + +-- | Spawn a thread that generates signals for the +-- dispatcher to probe the database for incoming jobs. +spawnListener + :: (MonadBaseControl IO m, MonadMask m) + => ConsumerConfig m idx job + -> ConnectionSource + -> MVar () + -> m ThreadId +spawnListener cc cs semaphore = forkP "listener" $ case ccNotificationChannel cc of + Just chan -> runDBT cs noTs . bracket_ (listen chan) (unlisten chan) . forever $ do + -- If there are many notifications, we need to collect them + -- as soon as possible, because they are stored in memory by + -- libpq. They are also not squashed, so we perform the + -- squashing ourselves with the help of MVar (). + void . getNotification $ ccNotificationTimeout cc + lift signalDispatcher + Nothing -> forever $ do + liftBase . threadDelay $ ccNotificationTimeout cc + signalDispatcher + where + signalDispatcher = do + liftBase $ tryPutMVar semaphore () + + noTs = def { + tsAutoTransaction = False + } + +-- | Spawn a thread that monitors working consumers +-- for activity and periodically updates its own. +spawnMonitor + :: (MonadBaseControl IO m, MonadLog m, MonadMask m) + => ConsumerConfig m idx job + -> ConnectionSource + -> ConsumerID + -> m ThreadId +spawnMonitor ConsumerConfig{..} cs cid = forkP "monitor" . forever $ do + runDBT cs ts $ do + -- Update last_activity of the consumer. + ok <- runSQL01 $ smconcat [ + "UPDATE" <+> raw ccConsumersTable + , "SET last_activity = now()" + , "WHERE id =" cid + , " AND name =" unRawSQL ccJobsTable + ] + if ok + then logInfo_ "Activity of the consumer updated" + else do + logInfo_ $ "Consumer is not registered" + throwM ThreadKilled + -- Freeing jobs locked by inactive consumers needs to happen + -- exactly once, otherwise it's possible to free it twice, after + -- it was already marked as reserved by other consumer, so let's + -- run it in serializable transaction. + (inactiveConsumers, freedJobs) <- runDBT cs tsSerializable $ do + -- Delete all inactive (assumed dead) consumers and get their ids. + runSQL_ $ smconcat [ + "DELETE FROM" <+> raw ccConsumersTable + , " WHERE last_activity +" iminutes 1 <+> "<= now()" + , " AND name =" unRawSQL ccJobsTable + , " RETURNING id::bigint" + ] + inactive :: [Int64] <- fetchMany runIdentity + -- Reset reserved jobs manually, do not rely + -- on the foreign key constraint to do its job. + freed <- if null inactive + then return 0 + else runSQL $ smconcat [ + "UPDATE" <+> raw ccJobsTable + , "SET reserved_by = NULL" + , "WHERE reserved_by = ANY(" Array1 inactive <+> ")" + ] + return (length inactive, freed) + when (inactiveConsumers > 0) $ do + logInfo "Unregistered inactive consumers" $ object [ + "inactive_consumers" .= inactiveConsumers + ] + when (freedJobs > 0) $ do + logInfo "Freed locked jobs" $ object [ + "freed_jobs" .= freedJobs + ] + liftBase . threadDelay $ 30 * 1000000 -- wait 30 seconds + where + tsSerializable = ts { + tsIsolationLevel = Serializable + } + +-- | Spawn a thread that reserves and processes jobs. +spawnDispatcher + :: forall m idx job. (MonadBaseControl IO m, MonadLog m, MonadMask m, Show idx, ToSQL idx) + => ConsumerConfig m idx job + -> ConnectionSource + -> ConsumerID + -> MVar () + -> TVar (M.Map ThreadId idx) + -> TVar Int + -> m ThreadId +spawnDispatcher ConsumerConfig{..} cs cid semaphore runningJobsInfo runningJobs = + forkP "dispatcher" . forever $ do + void $ takeMVar semaphore + loop 1 + where + loop :: Int -> m () + loop limit = do + (batch, batchSize) <- reserveJobs limit + when (batchSize > 0) $ do + logInfo "Processing batch" $ object [ + "batch_size" .= batchSize + ] + -- Update runningJobs before forking so that we can + -- adjust maxBatchSize appropriately later. We also + -- need to mask asynchronous exceptions here as we + -- rely on correct value of runningJobs to perform + -- graceful termination. + mask $ \restore -> do + atomically $ modifyTVar' runningJobs (+batchSize) + let subtractJobs = atomically $ do + modifyTVar' runningJobs (subtract batchSize) + void . forkP "batch processor" . (`finally` subtractJobs) . restore $ do + mapM startJob batch >>= mapM joinJob >>= updateJobs + + when (batchSize == limit) $ do + maxBatchSize <- atomically $ do + jobs <- readTVar runningJobs + when (jobs >= ccMaxRunningJobs) retry + return $ ccMaxRunningJobs - jobs + loop $ min maxBatchSize (2*limit) + + reserveJobs :: Int -> m ([job], Int) + reserveJobs limit = runDBT cs ts $ do + n <- runSQL $ smconcat [ + "UPDATE" <+> raw ccJobsTable <+> "SET" + , " reserved_by =" cid + , ", attempts = CASE" + , " WHEN finished_at IS NULL THEN attempts + 1" + , " ELSE 1" + , " END" + , "WHERE id IN (" <> reservedJobs <> ")" + , "RETURNING" <+> mintercalate ", " ccJobSelectors + ] + -- Decode lazily as we want the transaction to be as short as possible. + (, n) . F.toList . fmap ccJobFetcher <$> queryResult + where + reservedJobs :: SQL + reservedJobs = smconcat [ + "SELECT id FROM" <+> raw ccJobsTable + -- Converting id to text and hashing it may seem silly, + -- especially when we're dealing with integers in the first + -- place, but even in such case the overhead is small enough + -- (converting 100k integers to text and hashing them takes + -- around 15 ms on i7) to be worth the generality. + -- Also: after PostgreSQL 9.5 is released, we can use SELECT + -- FOR UPDATE SKIP LOCKED instead of advisory locks (see + -- http://michael.otacoo.com/postgresql-2/postgres-9-5-feature-highlight-skip-locked-row-level/ + -- for more details). Also, note that even if IDs of two + -- pending jobs produce the same hash, it just means that + -- in the worst case they will be processed by the same consumer. + , "WHERE pg_try_advisory_xact_lock(" unRawSQL ccJobsTable <> "::regclass::integer, hashtext(id::text))" + , " AND reserved_by IS NULL" + , " AND run_at IS NOT NULL" + , " AND run_at <= now()" + , "LIMIT" limit + , "FOR UPDATE" + ] + + -- | Spawn each job in a separate thread. + startJob :: job -> m (job, m (T.Result Result)) + startJob job = do + (_, joinFork) <- mask $ \restore -> T.fork $ do + tid <- myThreadId + bracket_ (registerJob tid) (unregisterJob tid) . restore $ do + ccProcessJob job + return (job, joinFork) + where + registerJob tid = atomically $ do + modifyTVar' runningJobsInfo . M.insert tid $ ccJobIndex job + unregisterJob tid = atomically $ do + modifyTVar' runningJobsInfo $ M.delete tid + + -- | Wait for all the jobs and collect their results. + joinJob :: (job, m (T.Result Result)) -> m (idx, Result) + joinJob (job, joinFork) = joinFork >>= \eres -> case eres of + Right result -> return (ccJobIndex job, result) + Left ex -> do + action <- ccOnException ex job + logAttention "Unexpected exception caught while processing job" $ object [ + "job_id" .= show (ccJobIndex job) + , "exception" .= show ex + , "action" .= show action + ] + return (ccJobIndex job, Failed action) + + -- | Update status of the jobs. + updateJobs :: [(idx, Result)] -> m () + updateJobs results = runDBT cs ts $ do + runSQL_ $ smconcat [ + "WITH removed AS (" + , " DELETE FROM" <+> raw ccJobsTable + , " WHERE id = ANY(" Array1 deletes <+> ")" + , ")" + , "UPDATE" <+> raw ccJobsTable <+> "SET" + , " reserved_by = NULL" + , ", run_at = CASE" + , " WHEN FALSE THEN run_at" + , smconcat $ M.foldrWithKey retryToSQL [] retries + , " ELSE NULL" -- processed + , " END" + , ", finished_at = CASE" + , " WHEN id = ANY(" Array1 successes <+> ") THEN now()" + , " ELSE NULL" + , " END" + , "WHERE id = ANY(" Array1 (map fst updates) <+> ")" + ] + where + retryToSQL (Left int) ids = + ("WHEN id = ANY(" Array1 ids <+> ") THEN now() +" int :) + retryToSQL (Right time) ids = + ("WHEN id = ANY(" Array1 ids <+> ") THEN" time :) + + retries = foldr step M.empty $ map f updates + where + f (idx, result) = case result of + Ok action -> (idx, action) + Failed action -> (idx, action) + + step (idx, action) iretries = case action of + MarkProcessed -> iretries + RerunAfter int -> M.insertWith (++) (Left int) [idx] iretries + RerunAt time -> M.insertWith (++) (Right time) [idx] iretries + Remove -> error "updateJobs: Remove should've been filtered out" + + successes = foldr step [] updates + where + step (idx, Ok _) acc = idx : acc + step (_, Failed _) acc = acc + + (deletes, updates) = foldr step ([], []) results + where + step job@(idx, result) (ideletes, iupdates) = case result of + Ok Remove -> (idx : ideletes, iupdates) + Failed Remove -> (idx : ideletes, iupdates) + _ -> (ideletes, job : iupdates) + +---------------------------------------- + +ts :: TransactionSettings +ts = def { + -- PostgreSQL doesn't seem to handle very high amount of + -- concurrent transactions that modify multiple rows in + -- the same table well (see updateJobs) and sometimes (very + -- rarely though) ends up in a deadlock. It doesn't matter + -- much though, we just restart the transaction in such case. + tsRestartPredicate = Just . RestartPredicate + $ \e _ -> qeErrorCode e == DeadlockDetected + || qeErrorCode e == SerializationFailure +} + +atomically :: MonadBase IO m => STM a -> m a +atomically = liftBase . STM.atomically diff --git a/src/Database/PostgreSQL/Consumers/Config.hs b/src/Database/PostgreSQL/Consumers/Config.hs new file mode 100644 --- /dev/null +++ b/src/Database/PostgreSQL/Consumers/Config.hs @@ -0,0 +1,100 @@ +{-# LANGUAGE ExistentialQuantification #-} +module Database.PostgreSQL.Consumers.Config ( + Action(..) + , Result(..) + , ConsumerConfig(..) + ) where + +import Control.Exception (SomeException) +import Data.Time +import Database.PostgreSQL.PQTypes + +-- | Action to take after a job was processed. +data Action + = MarkProcessed + | RerunAfter Interval + | RerunAt UTCTime + | Remove + deriving (Eq, Ord, Show) + +-- | Result of processing a job. +data Result = Ok Action | Failed Action + deriving (Eq, Ord, Show) + +-- | Config of a consumer. +data ConsumerConfig m idx job = forall row. FromRow row => ConsumerConfig { +-- | Name of a database table where jobs are stored. The table needs to have +-- the following columns in order to be suitable for acting as a job queue: +-- +-- * id - represents ID of the job. Needs to be a primary key of the type +-- convertible to text, not nullable. +-- +-- * run_at - represents the time at which the job will be processed. Needs +-- to be nullable, of the type comparable with now() (TIMESTAMPTZ is recommended). +-- Note: a job with run_at set to NULL is never picked for processing. Useful +-- for storing already processed/expired jobs for debugging purposes. +-- +-- * finished_at - represents the time at which job processing was finished. +-- Needs to be nullable, of the type you can assign now() to (TIMESTAMPTZ is +-- recommended). NULL means that the job was either never processed or that +-- it was started and failed at least once. +-- +-- * reserved_by - represents ID of the consumer that currently processes the +-- job. Needs to be nullable, of the type corresponding to id in the table +-- 'ccConsumersTable'. It's recommended (though not neccessary) to make it +-- a foreign key referencing id in 'ccConsumersTable' with ON DELETE SET NULL. +-- +-- * attempts - represents number of job processing attempts made so far. Needs +-- to be not nullable, of type INTEGER. Initial value of a fresh job should be +-- 0, therefore it makes sense to make the column default to 0. +-- + ccJobsTable :: !(RawSQL ()) +-- | Name of a database table where registered consumers are stored. The table +-- itself needs to have the following columns: +-- +-- * id - represents ID of a consumer. Needs to be a primary key of the type +-- SERIAL or BIGSERIAL (recommended). +-- +-- * name - represents jobs table of the consumer. Needs to be not nullable, +-- of type TEXT. Allows for tracking consumers of multiple queues with one +-- table. Set to 'ccJobsTable'. +-- +-- * last_activity - represents the last registered activity of the consumer. +-- It's updated periodically by all currently running consumers every 30 +-- seconds to prove that they are indeed running. They also check for the +-- registered consumers that didn't update their status for a minute. If any +-- such consumers are found, they are presumed to be not working and all the +-- jobs reserved by them are released. This prevents the situation where +-- a consumer with reserved jobs silently fails (e.g. because of a hard crash) +-- and these jobs stay locked forever, yet are never processed. +-- +, ccConsumersTable :: !(RawSQL ()) +-- | Fields needed to be selected from the jobs table in order to assemble a job. +, ccJobSelectors :: ![SQL] +-- | Function that transforms the list of fields into a job. +, ccJobFetcher :: !(row -> job) +-- | Selector for taking out job ID from the job object. +, ccJobIndex :: !(job -> idx) +-- | Notification channel used for listening for incoming jobs. If set +-- to 'Nothing', no listening is performed and jobs are selected from +-- the database every 'ccNotificationTimeout' microseconds. +, ccNotificationChannel :: !(Maybe Channel) +-- | Timeout of listening for incoming jobs, in microseconds. The consumer +-- checks between the timeouts if there are any jobs in the database that +-- needs to be processed, so even if 'ccNotificationChannel' is 'Just', you +-- need to set it to a reasonable number if the jobs you process may fail +-- and are retried later, as there is no way to signal with a notification +-- that a job will need to be performed e.g. in 5 minutes. However, if +-- 'ccNotificationChannel' is 'Just' and jobs are never retried, you can +-- set it to -1, then listening will never timeout. Otherwise it needs to +-- be a positive number. +, ccNotificationTimeout :: !Int +-- | Maximum amount of jobs that can be processed in parallel. +, ccMaxRunningJobs :: !Int +-- | Function that processes a job. +, ccProcessJob :: !(job -> m Result) +-- | Action taken if job processing function throws an exception. +-- Note that if this action throws an exception, the consumer goes +-- down, so it's best to ensure that it doesn't throw. +, ccOnException :: !(SomeException -> job -> m Action) +} diff --git a/src/Database/PostgreSQL/Consumers/Consumer.hs b/src/Database/PostgreSQL/Consumers/Consumer.hs new file mode 100644 --- /dev/null +++ b/src/Database/PostgreSQL/Consumers/Consumer.hs @@ -0,0 +1,73 @@ +module Database.PostgreSQL.Consumers.Consumer ( + ConsumerID + , registerConsumer + , unregisterConsumer + ) where + +import Control.Applicative +import Control.Monad.Base +import Control.Monad.Catch +import Data.Int +import Data.Monoid +import Data.Monoid.Utils +import Database.PostgreSQL.PQTypes + +import Database.PostgreSQL.Consumers.Config + +-- | ID of a consumer. +newtype ConsumerID = ConsumerID Int64 + deriving (Eq, Ord, PQFormat) +instance FromSQL ConsumerID where + type PQBase ConsumerID = PQBase Int64 + fromSQL mbase = ConsumerID <$> fromSQL mbase +instance ToSQL ConsumerID where + type PQDest ConsumerID = PQDest Int64 + toSQL (ConsumerID n) = toSQL n + +instance Show ConsumerID where + showsPrec p (ConsumerID n) = showsPrec p n + +-- | Register consumer in the consumers table, +-- so that it can reserve jobs using acquired ID. +registerConsumer + :: (MonadBase IO m, MonadMask m) + => ConsumerConfig n idx job + -> ConnectionSource + -> m ConsumerID +registerConsumer ConsumerConfig{..} cs = runDBT cs ts $ do + runSQL_ $ smconcat [ + "INSERT INTO" <+> raw ccConsumersTable + , "(name, last_activity) VALUES (" unRawSQL ccJobsTable <> ", now())" + , "RETURNING id" + ] + fetchOne runIdentity + where + ts = def { + tsAutoTransaction = False + } + +-- | Unregister consumer with a given ID. +unregisterConsumer + :: (MonadBase IO m, MonadMask m) + => ConsumerConfig n idx job + -> ConnectionSource + -> ConsumerID + -> m () +unregisterConsumer ConsumerConfig{..} cs wid = runDBT cs ts $ do + -- Free tasks manually in case there is no + -- foreign key constraint on reserved_by, + runSQL_ $ smconcat [ + "UPDATE" <+> raw ccJobsTable + , " SET reserved_by = NULL" + , " WHERE reserved_by =" wid + ] + runSQL_ $ smconcat [ + "DELETE FROM " <+> raw ccConsumersTable + , "WHERE id =" wid + , " AND name =" unRawSQL ccJobsTable + ] + where + ts = def { + tsRestartPredicate = Just . RestartPredicate + $ \e _ -> qeErrorCode e == DeadlockDetected + } diff --git a/src/Database/PostgreSQL/Consumers/Utils.hs b/src/Database/PostgreSQL/Consumers/Utils.hs new file mode 100644 --- /dev/null +++ b/src/Database/PostgreSQL/Consumers/Utils.hs @@ -0,0 +1,73 @@ +module Database.PostgreSQL.Consumers.Utils ( + finalize + , ThrownFrom(..) + , stopExecution + , forkP + , gforkP + ) where + +import Control.Concurrent.Lifted +import Control.Monad.Base +import Control.Monad.Catch +import Control.Monad.Trans.Control +import Data.Typeable +import qualified Control.Concurrent.Thread.Group.Lifted as TG +import qualified Control.Concurrent.Thread.Lifted as T +import qualified Control.Exception.Lifted as E + +-- | Runs an action that returns a finalizer and performs it at the end. +finalize :: (MonadMask m, MonadBase IO m) => m (m ()) -> m a -> m a +finalize m action = do + finalizer <- newEmptyMVar + flip finally (tryTakeMVar finalizer >>= maybe (return ()) id) $ do + putMVar finalizer =<< m + action + +---------------------------------------- + +-- | Exception thrown to a thread to stop its execution. +-- All exceptions other than 'StopExecution' thrown to +-- threads spawned by 'forkP' and 'gforkP' are propagated +-- back to the parent thread. +data StopExecution = StopExecution + deriving (Show, Typeable) +instance Exception StopExecution + +-- | Exception thrown from a child thread. +data ThrownFrom = ThrownFrom String SomeException + deriving (Show, Typeable) +instance Exception ThrownFrom + +-- | Stop execution of a thread. +stopExecution :: MonadBase IO m => ThreadId -> m () +stopExecution = flip throwTo StopExecution + +---------------------------------------- + +-- | Modified version of 'fork' that propagates +-- thrown exceptions to the parent thread. +forkP :: MonadBaseControl IO m => String -> m () -> m ThreadId +forkP = forkImpl fork + +-- | Modified version of 'TG.fork' that propagates +-- thrown exceptions to the parent thread. +gforkP :: MonadBaseControl IO m + => TG.ThreadGroup + -> String + -> m () + -> m (ThreadId, m (T.Result ())) +gforkP = forkImpl . TG.fork + +---------------------------------------- + +forkImpl :: MonadBaseControl IO m + => (m () -> m a) + -> String + -> m () + -> m a +forkImpl ffork tname m = E.mask $ \release -> do + parent <- myThreadId + ffork $ release m `E.catches` [ + E.Handler $ \StopExecution -> return () + , E.Handler $ (throwTo parent . ThrownFrom tname) + ]