diff --git a/postgresql-simple-queue.cabal b/postgresql-simple-queue.cabal --- a/postgresql-simple-queue.cabal +++ b/postgresql-simple-queue.cabal @@ -1,81 +1,104 @@ -name: postgresql-simple-queue -version: 0.5.1.1 -synopsis: A PostgreSQL backed queue -description: - This module utilize PostgreSQL to implement a durable queue for efficently processing arbitrary payloads which can be represented as JSON. - . - Typically a producer would enqueue a new payload as part of larger database - transaction - . - > createAccount userRecord = do - > 'runDBTSerializable' $ do - > createUserDB userRecord - > 'enqueueDB' "queue_schema" $ makeVerificationEmail userRecord - . - In another thread or process, the consumer would drain the queue. - . - > forever $ do - > -- Attempt get a payload or block until one is available - > payload <- lock "queue" conn - > - > -- Perform application specifc parsing of the payload value - > case fromJSON $ pValue payload of - > Success x -> sendEmail x -- Perform application specific processing - > Error err -> logErr err - > - > -- Remove the payload from future processing - > dequeue "queue" conn $ pId payload - > - > To support multiple queues in the same database, the API expects a table name string - > to determine which queue tables to use. -homepage: https://github.com/jfischoff/postgresql-queue#readme -license: BSD3 -license-file: LICENSE -author: Jonathan Fischoff -maintainer: jonathangfischoff@gmail.com -copyright: 2017 Jonathan Fischoff -category: Web -build-type: Simple -extra-source-files: README.md -cabal-version: >=1.10 +-- This file has been generated from package.yaml by hpack version 0.17.1. +-- +-- see: https://github.com/sol/hpack +name: postgresql-simple-queue +version: 1.0.0 +synopsis: A PostgreSQL backed queue +description: This module utilize PostgreSQL to implement a durable queue for efficently processing arbitrary payloads which can be represented as JSON. + . + Typically a producer would enqueue a new payload as part of larger database + transaction + . + > createAccount userRecord = do + > 'runDBTSerializable' $ do + > createUserDB userRecord + > 'enqueueDB' "queue_schema" $ makeVerificationEmail userRecord + . + In another thread or process, the consumer would drain the queue. + . + > forever $ do + > -- Attempt get a payload or block until one is available + > payload <- lock "queue" conn + > + > -- Perform application specifc parsing of the payload value + > case fromJSON $ pValue payload of + > Success x -> sendEmail x -- Perform application specific processing + > Error err -> logErr err + > + > -- Remove the payload from future processing + > dequeue "queue" conn $ pId payload + > + > To support multiple queues in the same database, the API expects a table name string + > to determine which queue tables to use. +homepage: https://github.com/jfischoff/postgresql-queue#readme +bug-reports: https://github.com/jfischoff/postgresql-queue/issues +license: BSD3 +license-file: LICENSE +author: Jonathan Fischoff +maintainer: jonathangfischoff@gmail.com +copyright: 2017 Jonathan Fischoff +category: Web +build-type: Simple +cabal-version: >= 1.10 + +extra-source-files: + README.md + +source-repository head + type: git + location: https://github.com/jfischoff/postgresql-queue + library - hs-source-dirs: src - exposed-modules: Database.PostgreSQL.Simple.Queue - , Database.PostgreSQL.Simple.Queue.Migrate - build-depends: base >= 4.7 && < 5 - , postgresql-simple - , pg-transact - , aeson - , time - , transformers - , random - , text - , monad-control - , exceptions - , bytestring - default-language: Haskell2010 + hs-source-dirs: + src + exposed-modules: + Database.PostgreSQL.Simple.Queue + Database.PostgreSQL.Simple.Queue.Migrate + other-modules: + Paths_postgresql_simple_queue + build-depends: + base >=4.7 && <5 + , time + , transformers + , random + , text + , monad-control + , exceptions + , postgresql-simple + , pg-transact + , aeson + , bytestring + , stm + default-language: Haskell2010 ghc-options: -Wall -Wno-unused-do-bind test-suite unit-tests - type: exitcode-stdio-1.0 - hs-source-dirs: test + type: exitcode-stdio-1.0 + hs-source-dirs: + test main-is: Spec.hs - other-modules: Database.PostgreSQL.Simple.QueueSpec - build-depends: base - , postgresql-simple-queue - , hspec - , hspec-discover - , postgresql-simple - , pg-transact - , bytestring - , aeson - , hspec-expectations-lifted - , hspec-pg-transact - , async + other-modules: + Database.PostgreSQL.Simple.QueueSpec + build-depends: + base >=4.7 && <5 + , time + , transformers + , random + , text + , monad-control + , exceptions + , postgresql-simple + , pg-transact + , aeson + , bytestring + , stm + , postgresql-simple-queue + , hspec + , hspec-discover + , hspec-expectations-lifted + , hspec-pg-transact + , async + , split ghc-options: -Wall -Wno-unused-do-bind -O2 -threaded -rtsopts -with-rtsopts=-N - default-language: Haskell2010 - -source-repository head - type: git - location: https://github.com/jfischoff/postgresql-queue + default-language: Haskell2010 diff --git a/src/Database/PostgreSQL/Simple/Queue.hs b/src/Database/PostgreSQL/Simple/Queue.hs --- a/src/Database/PostgreSQL/Simple/Queue.hs +++ b/src/Database/PostgreSQL/Simple/Queue.hs @@ -54,16 +54,14 @@ , Payload (..) -- * DB API , enqueueDB - , tryLockDB - , unlockDB , dequeueDB + , withPayloadDB , getCountDB -- * IO API , enqueue - , tryLock - , lock - , unlock + , tryDequeue , dequeue + , withPayload , getCount ) where import Control.Monad @@ -71,7 +69,6 @@ import Data.Aeson import Data.Function import Data.Int -import Data.Maybe import Data.Text (Text) import Data.Time import Database.PostgreSQL.Simple (Connection, Only (..)) @@ -86,6 +83,9 @@ import Database.PostgreSQL.Transact import Data.Monoid import Data.String +import Control.Monad.IO.Class +import Data.Maybe + ------------------------------------------------------------------------------- --- Types ------------------------------------------------------------------------------- @@ -98,18 +98,32 @@ instance ToRow PayloadId where toRow = toRow . Only --- | A 'Payload' can exist in three states in the queue, 'Enqueued', 'Locked' --- and 'Dequeued'. A 'Payload' starts in the 'Enqueued' state and is 'Locked' +-- The fundemental record stored in the queue. The queue is a single table +-- and each row consists of a 'Payload' +data Payload = Payload + { pId :: PayloadId + , pValue :: Value + -- ^ The JSON value of a payload + , pState :: State + , pAttempts :: Int + , pCreatedAt :: UTCTime + , pModifiedAt :: UTCTime + } deriving (Show, Eq) + +instance FromRow Payload where + fromRow = Payload <$> field <*> field <*> field <*> field <*> field <*> field + +-- | A 'Payload' can exist in three states in the queue, 'Enqueued', +-- and 'Dequeued'. A 'Payload' starts in the 'Enqueued' state and is locked -- so some sort of process can occur with it, usually something in 'IO'. -- Once the processing is complete, the `Payload' is moved the 'Dequeued' -- state, which is the terminal state. -data State = Enqueued | Locked | Dequeued +data State = Enqueued | Dequeued deriving (Show, Eq, Ord, Enum, Bounded) instance ToField State where toField = toField . \case Enqueued -> "enqueued" :: Text - Locked -> "locked" Dequeued -> "dequeued" -- Converting from enumerations is annoying :( @@ -120,26 +134,12 @@ Nothing -> returnError UnexpectedNull f "state can't be NULL" Just y' -> case y' of "enqueued" -> return Enqueued - "locked" -> return Locked "dequeued" -> return Dequeued x -> returnError ConversionFailed f (show x) else returnError Incompatible f $ "Expect type name to be state but it was " ++ show n --- The fundemental record stored in the queue. The queue is a single table --- and each row consists of a 'Payload' -data Payload = Payload - { pId :: PayloadId - , pValue :: Value - -- ^ The JSON value of a payload - , pState :: State - , pCreatedAt :: UTCTime - , pModifiedAt :: UTCTime - } deriving (Show, Eq) - -instance FromRow Payload where - fromRow = Payload <$> field <*> field <*> field <*> field <*> field ------------------------------------------------------------------------------- --- DB API ------------------------------------------------------------------------------- @@ -161,57 +161,81 @@ @ -} enqueueDB :: String -> Value -> DB PayloadId -enqueueDB schemaName value = +enqueueDB schemaName value = enqueueWithDB schemaName value 0 + +enqueueWithDB :: String -> Value -> Int -> DB PayloadId +enqueueWithDB schemaName value attempts = fmap head $ query (withSchema schemaName $ [sql| NOTIFY |] <> " " <> notifyName schemaName <> ";" <> [sql| - INSERT INTO payloads (value) - VALUES (?) + INSERT INTO payloads (attempts, value) + VALUES (?, ?) RETURNING id;|] ) - $ Only value + $ (attempts, value) -{-| Return a the oldest 'Payload' in the 'Enqueued' state, or 'Nothing' - if there are no payloads. This function is not necessarily useful by - itself, since there are not many use cases where it needs to be combined - with other transactions. See 'tryLock' the IO API version, or for a - blocking version utilizing PostgreSQL's NOTIFY and LISTEN, see 'lock' --} -tryLockDB :: String -> DB (Maybe Payload) -tryLockDB schemaName = fmap listToMaybe $ query_ $ withSchema schemaName +retryDB :: String -> Value -> Int -> DB PayloadId +retryDB schemaName value attempts = enqueueWithDB schemaName value $ attempts + 1 + +-- | Transition a 'Payload' to the 'Dequeued' state. +dequeueDB :: String -> DB (Maybe Payload) +dequeueDB schemaName = fmap listToMaybe $ query_ $ withSchema schemaName [sql| UPDATE payloads - SET state='locked' + SET state='dequeued' WHERE id in ( SELECT id FROM payloads WHERE state='enqueued' - ORDER BY created_at ASC + ORDER BY modified_at ASC FOR UPDATE SKIP LOCKED LIMIT 1 ) - RETURNING id, value, state, created_at, modified_at + RETURNING id, value, state, attempts, created_at, modified_at |] -{-| Transition a 'Payload' from the 'Locked' state to the 'Enqueued' state. - Useful for responding to asynchronous exceptions during a unexpected - shutdown. In general the IO API version, 'unlock', is probably more - useful. The DB version is provided for completeness. +{-| + +Attempt to get a payload and process it. If the function passed in throws an exception +return it on the left side of the `Either`. Re-add the payload up to some passed in +maximum. Return `Nothing` is the `payloads` table is empty otherwise the result is an `a` +from the payload ingesting function. + -} -unlockDB :: String -> PayloadId -> DB () -unlockDB schemaName payloadId = void $ execute (withSchema schemaName - [sql| UPDATE payloads - SET state='enqueued' - WHERE id=? AND state='locked' - |]) - payloadId +withPayloadDB :: String + -- ^ schema + -> Int + -- ^ retry count + -> (Payload -> IO a) + -- ^ payload processing function + -> DB (Either SomeException (Maybe a)) +withPayloadDB schemaName retryCount f + = query_ + ( withSchema schemaName + $ [sql| + SELECT id, value, state, attempts, created_at, modified_at + FROM payloads + WHERE state='enqueued' + ORDER BY modified_at ASC + FOR UPDATE SKIP LOCKED + LIMIT 1 + |] + ) + >>= \case + [] -> return $ return Nothing + [payload@Payload {..}] -> do + execute [sql| UPDATE payloads SET state='dequeued' WHERE id = ? |] pId --- | Transition a 'Payload' to the 'Dequeued' state. -dequeueDB :: String -> PayloadId -> DB () -dequeueDB schemaName payloadId = void $ execute (withSchema schemaName - [sql| UPDATE payloads - SET state='dequeued' - WHERE id=? - |]) - payloadId + -- Retry on failure up to retryCount + handle (\e -> when (pAttempts < retryCount) + (void $ retryDB schemaName pValue pAttempts) + >> return (Left e) + ) + $ Right . Just <$> liftIO (f payload) + xs -> return + $ Left + $ toException + $ userError + $ "LIMIT is 1 but got more than one row: " + ++ show xs -- | Get the number of rows in the 'Enqueued' state. getCountDB :: String -> DB Int64 @@ -229,47 +253,59 @@ enqueue :: String -> Connection -> Value -> IO PayloadId enqueue schemaName conn value = runDBT (enqueueDB schemaName value) ReadCommitted conn -{-| Return a the oldest 'Payload' in the 'Enqueued' state or 'Nothing' - if there are no payloads. For a blocking version utilizing PostgreSQL's - NOTIFY and LISTEN, see 'lock'. This functions runs 'tryLockDB' as a - 'Serializable' transaction. --} -tryLock :: String -> Connection -> IO (Maybe Payload) -tryLock schemaName conn = runDBT (tryLockDB schemaName) ReadCommitted conn - +-- Block until a payload notification is fired. Fired during insertion. notifyPayload :: String -> Connection -> IO () notifyPayload schemaName conn = do Notification {..} <- getNotification conn unless (notificationChannel == notifyName schemaName) $ notifyPayload schemaName conn -{-| Return the oldest 'Payload' in the 'Enqueued' state or block until a - payload arrives. This function utilizes PostgreSQL's LISTEN and NOTIFY - functionality to avoid excessively polling of the DB while - waiting for new payloads, without scarficing promptness. +{-| Return a the oldest 'Payload' in the 'Enqueued' state or 'Nothing' + if there are no payloads. For a blocking version utilizing PostgreSQL's + NOTIFY and LISTEN, see 'dequeue'. This functions runs 'dequeueDb' as a + 'ReadCommitted' transaction. + + See `withPayload' for an alternative interface that will automatically return + the payload to the 'Enqueued' state if an exception occurs. -} -lock :: String -> Connection -> IO Payload -lock schemaName conn = bracket_ +tryDequeue :: String -> Connection -> IO (Maybe Payload) +tryDequeue schemaName conn = runDBT (dequeueDB schemaName) ReadCommitted conn + +-- | Transition a 'Payload' to the 'Dequeued' state. his functions runs +-- 'dequeueDB' as a 'Serializable' transaction. +dequeue :: String -> Connection -> IO Payload +dequeue schemaName conn = bracket_ (Simple.execute_ conn $ "LISTEN " <> notifyName schemaName) (Simple.execute_ conn $ "UNLISTEN " <> notifyName schemaName) $ fix $ \continue -> do - m <- tryLock schemaName conn + m <- tryDequeue schemaName conn case m of Nothing -> do notifyPayload schemaName conn continue Just x -> return x -{-| Transition a 'Payload' from the 'Locked' state to the 'Enqueued' state. - Useful for responding to asynchronous exceptions during a unexpected - shutdown. For a DB API version see 'unlockDB' +{-| Return the oldest 'Payload' in the 'Enqueued' state or block until a + payload arrives. This function utilizes PostgreSQL's LISTEN and NOTIFY + functionality to avoid excessively polling of the DB while + waiting for new payloads, without scarficing promptness. -} -unlock :: String -> Connection -> PayloadId -> IO () -unlock schemaName conn x = runDBT (unlockDB schemaName x) ReadCommitted conn - --- | Transition a 'Payload' to the 'Dequeued' state. his functions runs --- 'dequeueDB' as a 'Serializable' transaction. -dequeue :: String -> Connection -> PayloadId -> IO () -dequeue schemaName conn x = runDBT (dequeueDB schemaName x) ReadCommitted conn +withPayload :: String + -> Connection + -> Int + -- ^ retry count + -> (Payload -> IO a) + -> IO (Either SomeException a) +withPayload schemaName conn retryCount f = bracket_ + (Simple.execute_ conn $ "LISTEN " <> notifyName schemaName) + (Simple.execute_ conn $ "UNLISTEN " <> notifyName schemaName) + $ fix + $ \continue -> runDBT (withPayloadDB schemaName retryCount f) ReadCommitted conn + >>= \case + Left x -> return $ Left x + Right Nothing -> do + notifyPayload schemaName conn + continue + Right (Just x) -> return $ Right x {-| Get the number of rows in the 'Enqueued' state. This function runs 'getCountDB' in a 'ReadCommitted' transaction. diff --git a/src/Database/PostgreSQL/Simple/Queue/Migrate.hs b/src/Database/PostgreSQL/Simple/Queue/Migrate.hs --- a/src/Database/PostgreSQL/Simple/Queue/Migrate.hs +++ b/src/Database/PostgreSQL/Simple/Queue/Migrate.hs @@ -36,34 +36,39 @@ $$ language 'plpgsql'; @ + -} migrate :: String -> Connection -> IO () migrate schemaName conn = void $ execute_ conn $ "CREATE SCHEMA IF NOT EXISTS " <> fromString schemaName <> ";" <> "SET search_path TO " <> fromString schemaName <> ";" <> [sql| - CREATE TYPE state_t AS ENUM ('enqueued', 'locked', 'dequeued'); - CREATE TABLE payloads - ( id SERIAL PRIMARY KEY - , value jsonb NOT NULL - , state state_t NOT NULL DEFAULT 'enqueued' - , created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT clock_timestamp() - , modified_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT clock_timestamp() - ); - - CREATE INDEX state_idx ON payloads (state); + CREATE TYPE state_t AS ENUM ('enqueued', 'dequeued'); - CREATE OR REPLACE FUNCTION update_row_modified_function_() - RETURNS TRIGGER - AS + CREATE OR REPLACE FUNCTION update_row_modified_function() + RETURNS TRIGGER AS $$ BEGIN - -- ASSUMES the table has a column named exactly "modified_at". - -- Fetch date-time of actual current moment from clock, - -- rather than start of statement or start of transaction. NEW.modified_at = clock_timestamp(); RETURN NEW; END; $$ language 'plpgsql'; + + CREATE TABLE payloads + ( id BIGSERIAL PRIMARY KEY + , value jsonb NOT NULL + , attempts int NOT NULL DEFAULT 0 + , state state_t NOT NULL DEFAULT 'enqueued' + , created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT clock_timestamp() + , modified_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT clock_timestamp() + ); + + CREATE INDEX active_created_at_idx ON payloads (created_at) + WHERE (state = 'enqueued'); + + CREATE TRIGGER payloads_modified + BEFORE UPDATE ON payloads + FOR EACH ROW EXECUTE PROCEDURE update_row_modified_function(); + |] diff --git a/test/Database/PostgreSQL/Simple/QueueSpec.hs b/test/Database/PostgreSQL/Simple/QueueSpec.hs --- a/test/Database/PostgreSQL/Simple/QueueSpec.hs +++ b/test/Database/PostgreSQL/Simple/QueueSpec.hs @@ -2,18 +2,23 @@ {-# LANGUAGE RecordWildCards #-} module Database.PostgreSQL.Simple.QueueSpec (spec, main) where import Control.Concurrent +import Control.Concurrent.STM import Control.Concurrent.Async import Control.Monad import Data.Aeson import Data.Function -import Data.IORef import Data.List import Database.PostgreSQL.Simple.Queue import Database.PostgreSQL.Simple.Queue.Migrate import Test.Hspec (Spec, hspec, it) import Test.Hspec.Expectations.Lifted import Test.Hspec.DB +import Control.Monad.Catch +import Control.Monad.IO.Class +import Data.List.Split +import Data.Either + main :: IO () main = hspec spec @@ -23,75 +28,124 @@ spec :: Spec spec = describeDB (migrate schemaName) "Database.Queue" $ do itDB "empty locks nothing" $ - tryLockDB schemaName `shouldReturn` Nothing + (either throwM return =<< (withPayloadDB schemaName 8 return)) + `shouldReturn` Nothing itDB "empty gives count 0" $ getCountDB schemaName `shouldReturn` 0 - itDB "enqueues/tryLocks/unlocks" $ do - payloadId <- enqueueDB schemaName $ String "Hello" - getCountDB schemaName `shouldReturn` 1 + it "enqueuesDB/withPayloadDB" $ \conn -> do + runDB conn $ do + payloadId <- enqueueDB schemaName $ String "Hello" + getCountDB schemaName `shouldReturn` 1 - Just Payload {..} <- tryLockDB schemaName - getCountDB schemaName `shouldReturn` 0 + either throwM return =<< withPayloadDB schemaName 8 (\(Payload {..}) -> do + pId `shouldBe` payloadId + pValue `shouldBe` String "Hello" + ) - pId `shouldBe` payloadId - pValue `shouldBe` String "Hello" - tryLockDB schemaName `shouldReturn` Nothing + -- read committed but still 0. I don't depend on this but I want to see if it + -- stays like this. + getCountDB schemaName `shouldReturn` 0 - unlockDB schemaName pId - getCountDB schemaName `shouldReturn` 1 + runDB conn $ getCountDB schemaName `shouldReturn` 0 - itDB "tryLocks/dequeues" $ do - Just Payload {..} <- tryLockDB schemaName - getCountDB schemaName `shouldReturn` 0 + it "enqueuesDB/withPayloadDB/retries" $ \conn -> do + runDB conn $ do + void $ enqueueDB schemaName $ String "Hello" + getCountDB schemaName `shouldReturn` 1 - dequeueDB schemaName pId `shouldReturn` () - tryLockDB schemaName `shouldReturn` Nothing + xs <- replicateM 7 $ withPayloadDB schemaName 8 (\(Payload {..}) -> + throwM $ userError "not enough tries" + ) - it "enqueues and dequeues concurrently tryLock" $ \testDB -> do + all isLeft xs `shouldBe` True + + either throwM return =<< withPayloadDB schemaName 8 (\(Payload {..}) -> do + pAttempts `shouldBe` 7 + pValue `shouldBe` String "Hello" + ) + + runDB conn $ getCountDB schemaName `shouldReturn` 0 + + it "enqueuesDB/withPayloadDB/timesout" $ \conn -> do + runDB conn $ do + void $ enqueueDB schemaName $ String "Hello" + getCountDB schemaName `shouldReturn` 1 + + xs <- replicateM 2 $ withPayloadDB schemaName 1 (\(Payload {..}) -> + throwM $ userError "not enough tries" + ) + + all isLeft xs `shouldBe` True + + runDB conn $ getCountDB schemaName `shouldReturn` 0 + + it "selects the oldest first" $ \conn -> do + runDB conn $ do + payloadId0 <- enqueueDB schemaName $ String "Hello" + liftIO $ threadDelay 1000000 + + payloadId1 <- enqueueDB schemaName $ String "Hi" + + getCountDB schemaName `shouldReturn` 2 + + either throwM return =<< withPayloadDB schemaName 8 (\(Payload {..}) -> do + pId `shouldBe` payloadId0 + pValue `shouldBe` String "Hello" + ) + + either throwM return =<< withPayloadDB schemaName 8 (\(Payload {..}) -> do + pId `shouldBe` payloadId1 + pValue `shouldBe` String "Hi" + ) + + runDB conn $ getCountDB schemaName `shouldReturn` 0 + + it "enqueues and dequeues concurrently withPayload" $ \testDB -> do let withPool' = withPool testDB elementCount = 1000 :: Int expected = [0 .. elementCount - 1] - ref <- newIORef [] + ref <- newTVarIO [] - loopThreads <- replicateM 10 $ async $ fix $ \next -> do - mpayload <- withPool' $ tryLock schemaName - case mpayload of - Nothing -> next - Just Payload {..} -> do - lastCount <- atomicModifyIORef ref - $ \xs -> (pValue : xs, length xs + 1) - withPool' $ \c -> dequeue schemaName c pId - when (lastCount < elementCount) next + loopThreads <- replicateM 35 $ async $ withPool' $ \c -> fix $ \next -> do + lastCount <- either throwM return <=< withPayload schemaName c 0 $ \(Payload {..}) -> do + atomically $ do + xs <- readTVar ref + writeTVar ref $ pValue : xs + return $ length xs + 1 - -- Fork a hundred threads and enqueue an index - forM_ [0 .. elementCount - 1] $ \i -> - forkIO $ void $ withPool' $ \c -> enqueue schemaName c $ toJSON i + when (lastCount < elementCount) next + forM_ (chunksOf (elementCount `div` 11) expected) $ \xs -> forkIO $ void $ withPool' $ \c -> + forM_ xs $ \i -> enqueue schemaName c $ toJSON i + waitAnyCancel loopThreads - Just decoded <- mapM (decode . encode) <$> readIORef ref + xs <- atomically $ readTVar ref + let Just decoded = mapM (decode . encode) xs sort decoded `shouldBe` sort expected - it "enqueues and dequeues concurrently lock" $ \testDB -> do + it "enqueues and dequeues concurrently dequeue" $ \testDB -> do let withPool' = withPool testDB elementCount = 1000 :: Int expected = [0 .. elementCount - 1] - ref <- newIORef [] + ref <- newTVarIO [] - loopThreads <- replicateM 10 $ async $ fix $ \next -> do - Payload {..} <- withPool' $ lock schemaName - lastCount <- atomicModifyIORef ref - $ \xs -> (pValue : xs, length xs + 1) - withPool' $ \c -> dequeue schemaName c pId + loopThreads <- replicateM 35 $ async $ withPool' $ \c -> fix $ \next -> do + Payload {..} <- dequeue schemaName c + lastCount <- atomically $ do + xs <- readTVar ref + writeTVar ref $ pValue : xs + return $ length xs + 1 + when (lastCount < elementCount) next - -- Fork a hundred threads and enqueue an index - forM_ [0 .. elementCount - 1] $ \i -> forkIO $ void $ withPool' $ \c -> - enqueue schemaName c $ toJSON i + forM_ (chunksOf (elementCount `div` 11) expected) $ \xs -> forkIO $ void $ withPool' $ \c -> + forM_ xs $ \i -> enqueue schemaName c $ toJSON i waitAnyCancel loopThreads - Just decoded <- mapM (decode . encode) <$> readIORef ref + xs <- atomically $ readTVar ref + let Just decoded = mapM (decode . encode) xs sort decoded `shouldBe` sort expected