diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright Alexander Thiemann (c) 2017 + +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 Alexander Thiemann 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/bench/Bench.hs b/bench/Bench.hs new file mode 100644 --- /dev/null +++ b/bench/Bench.hs @@ -0,0 +1,43 @@ +module Main where + +import Data.PowerQueue +import Data.PowerQueue.Backend.LevelMem + +import Control.Monad +import Criterion.Main +import System.IO.Temp +import qualified Data.Serialize as S + +binEncoding :: S.Serialize a => JobEncoding a +binEncoding = + JobEncoding + { j_encode = S.encode + , j_decode = S.decode + } + +nopWorker :: QueueBackend () -> Queue () +nopWorker be = + newQueue be $ newQueueWorker $ \() -> pure JOk + +testCfg :: S.Serialize a => FilePath -> LevelMemCfg a +testCfg fp = + LevelMemCfg + { lmc_storageDir = fp + , lmc_maxQueueSize = 200000 + , lmc_jobEncoding = binEncoding + , lmc_inProgressRecovery = IpRestart + } + +main :: IO () +main = + withSystemTempDirectory "lmsbenchXXX" $ \tempDir -> + withLevelMem (testCfg tempDir) $ \lm -> + do let q = nopWorker (newLevelMemBackend lm) + putStrLn "Filling queue with 100k entries" + replicateM_ 100000 $ enqueueJob () q + putStrLn "Ready." + defaultMain + [ bgroup "enqueue dequeue performance" + [ bench "enqueue + dequeue" $ nfIO $ enqueueJob () q >> workStep q + ] + ] diff --git a/powerqueue-levelmem.cabal b/powerqueue-levelmem.cabal new file mode 100644 --- /dev/null +++ b/powerqueue-levelmem.cabal @@ -0,0 +1,69 @@ +name: powerqueue-levelmem +version: 0.1.0.0 +synopsis: A high performance in memory and LevelDB backend for powerqueue +description: A high performance in memory and LevelDB backend for powerqueue +homepage: https://github.com/agrafix/powerqueue#readme +license: BSD3 +license-file: LICENSE +author: Alexander Thiemann +maintainer: mail@athiemann.net +copyright: 2017 Alexander Thiemann +category: Web +build-type: Simple +cabal-version: >=1.10 + +library + hs-source-dirs: src + exposed-modules: + Data.PowerQueue.Backend.LevelMem + build-depends: + base >= 4.7 && < 5 + , powerqueue >= 0.1 + , stm >= 2.4 + , stm-containers >= 0.2 + , unagi-chan >= 0.4 + , focus >= 0.1 + , leveldb-haskell >= 0.6 + , async >= 2.1 + , cereal >= 0.5 + , bytestring >= 0.10 + , dlist >= 0.8 + , list-t >= 1 + , filepath >= 1.4 + default-language: Haskell2010 + +test-suite powerqueue-levelmem-test + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Spec.hs + other-modules: Data.PowerQueue.Backend.LevelMemSpec + build-depends: + base + , powerqueue-levelmem + , powerqueue + , hspec >= 2.2 + , async + , temporary + , cereal + ghc-options: -threaded -rtsopts -with-rtsopts=-N + extra-libraries: stdc++, leveldb, snappy + default-language: Haskell2010 + +benchmark powerqueue-levelmem-bench + type: exitcode-stdio-1.0 + ghc-options: -O2 -threaded -rtsopts -with-rtsopts=-N + extra-libraries: stdc++, leveldb, snappy + hs-source-dirs: bench + default-language: Haskell2010 + main-is: Bench.hs + build-depends: + base + , powerqueue + , powerqueue-levelmem + , criterion >= 1.1 + , temporary + , cereal + +source-repository head + type: git + location: https://github.com/agrafix/powerqueue diff --git a/src/Data/PowerQueue/Backend/LevelMem.hs b/src/Data/PowerQueue/Backend/LevelMem.hs new file mode 100644 --- /dev/null +++ b/src/Data/PowerQueue/Backend/LevelMem.hs @@ -0,0 +1,308 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE BangPatterns #-} +module Data.PowerQueue.Backend.LevelMem + ( -- * Data container + withLevelMem, LevelMemCfg(..), InProgressCfg(..), JobEncoding(..), LevelMem + , JobStatus(..), getJobStatusMap, getJobStatus + -- * Actual backend for powerqueue + , newLevelMemBackend + ) +where + +import Control.Concurrent.Async +import Control.Concurrent.STM +import Control.Exception +import Control.Monad +import Data.Maybe +import Data.Monoid +import Data.PowerQueue +import Data.Word +import System.FilePath +import qualified Control.Concurrent.Chan.Unagi.Bounded as C +import qualified Data.ByteString as BS +import qualified Data.DList as DL +import qualified Data.Serialize as S +import qualified Database.LevelDB.Base as L +import qualified ListT as L +import qualified STMContainers.Map as SMap +import qualified STMContainers.Set as SSet + +type Chan a = (C.InChan a, C.OutChan a) +type JobIdx = Word64 + +data QueueAction j + = QaEnqueue !JobIdx !j + | QaDequeue !JobIdx + | QaConfirm !JobIdx + | QaRollback !JobIdx + +databaseOps :: (j -> BS.ByteString) -> QueueAction j -> [L.BatchOp] +databaseOps serJob qa = + case qa of + QaEnqueue idx job -> + [ L.Put (mKey idx) (serJob job) + , L.Put (qKey idx) "" + ] + QaDequeue idx -> + [ L.Del (qKey idx) + , L.Put (pKey idx) "" + ] + QaConfirm idx -> + [ L.Del (pKey idx) + , L.Del (mKey idx) + ] + QaRollback idx -> + [ L.Del (pKey idx) + , L.Put (qKey idx) "" + ] + + +type DbKeyFun = Word64 -> BS.ByteString + +qKey :: DbKeyFun +qKey idx = "q_" <> S.encode idx + +mKey :: DbKeyFun +mKey idx = "map_" <> S.encode idx + +pKey :: DbKeyFun +pKey idx = "p_" <> S.encode idx + +data Terminate + = Continue + | Terminate + +databaseWorkerStep :: + L.DB -> (j -> BS.ByteString) -> TVar Terminate -> TQueue (QueueAction j) -> IO Bool +databaseWorkerStep db serJob termV q = + do (dbOps, goAgain) <- + atomically $ + do op <- tryReadTQueue q + term <- readTVar termV + case (op, term) of + (Just o, _) -> + (,) + <$> (concatMap (databaseOps serJob) . DL.toList <$> drainLoop (DL.singleton o)) + <*> pure True + (Nothing, Continue) -> retry + (Nothing, Terminate) -> pure ([], False) + L.write db (L.WriteOptions True) dbOps + pure goAgain + where + drainLoop !accum = + do mOp <- tryReadTQueue q + case mOp of + Nothing -> pure accum + Just nextOp -> drainLoop (DL.snoc accum nextOp) + +databaseWorker :: L.DB -> (j -> BS.ByteString) -> TVar Terminate -> TQueue (QueueAction j) -> IO () +databaseWorker db serJob termV q = + do goAgain <- databaseWorkerStep db serJob termV q + if goAgain + then databaseWorker db serJob termV q + else pure () + +databaseRecovery :: + InProgressCfg + -> TVar JobIdx + -> SMap.Map JobIdx j + -> Chan JobIdx + -> SSet.Set JobIdx + -> (BS.ByteString -> Either String j) + -> L.DB + -> IO () +databaseRecovery inPCfg jobIdxV jobMap jobQueue progressSet parseJob db = + L.withIter db L.defaultReadOptions $ \it -> + do populateMap it + populateSet it 2 "p_" $ \jobIdx -> + case inPCfg of + IpRecover -> atomically $ SSet.insert jobIdx progressSet + IpRestart -> C.writeChan (fst jobQueue) jobIdx + IpForget -> atomically $ SMap.delete jobIdx jobMap + populateSet it 2 "q_" $ C.writeChan (fst jobQueue) + where + populateMap it = + do L.iterSeek it "map_" + let workLoop !maxIdx = + do k <- L.iterKey it + v <- L.iterValue it + case (,) <$> k <*> v of + Just (kBs, vBs) + | BS.take 4 kBs == "map_" -> + case (,) <$> S.decode (BS.drop 4 kBs) <*> parseJob vBs of + Left err -> fail err + Right (jobIdx, job) -> + do atomically $ SMap.insert job jobIdx jobMap + L.iterNext it + workLoop (if jobIdx > maxIdx then jobIdx else maxIdx) + _ -> pure maxIdx + jobIdx <- workLoop 0 + atomically $ writeTVar jobIdxV (jobIdx + 1) + populateSet it dropPrefix prefix onVal = + do L.iterSeek it prefix + let fetchLoop = + do k <- L.iterKey it + case k of + Just bs | BS.take dropPrefix bs == prefix -> + case S.decode (BS.drop dropPrefix bs) of + Left _ -> pure () + Right ok -> + do void $ onVal ok + L.iterNext it + fetchLoop + _ -> pure () + fetchLoop + +-- | Binary encoding of a single job. Note that it is highly recommended to use a +-- backwards compatible decoder, otherwise the persistent state can not be read. +-- You could use safecopy or an appropriate cereal, binary, aeson or other decoding. +data JobEncoding j + = JobEncoding + { j_encode :: j -> BS.ByteString + , j_decode :: BS.ByteString -> Either String j + } + +-- | Behavoir for in progress jobs after loading the state from disk on launch +data InProgressCfg + = IpRecover + -- ^ mark the job as in progress + | IpRestart + -- ^ add the job to the queue + | IpForget + -- ^ discard the job + +data LevelMemCfg j + = LevelMemCfg + { lmc_storageDir :: !FilePath + -- ^ directory for persistence + , lmc_maxQueueSize :: !Int + -- ^ maximum size of queue, succeeding enqueues will block + , lmc_jobEncoding :: !(JobEncoding j) + -- ^ binary encoding of jobs + , lmc_inProgressRecovery :: !InProgressCfg + -- ^ how should in progress jobs be handled after a restart while loading from disk + } + +-- | Create a new data container for in memory job tracking and leveldb disk persistence. +-- Provide a 'FilePath' for storing the data, an 'Int' as maximum queue size and +-- a 'JobEncoding' for individual job encoding. +withLevelMem :: LevelMemCfg j -> (LevelMem j -> IO a) -> IO a +withLevelMem LevelMemCfg{..} action = + L.withDB (lmc_storageDir "queue.db") opts $ \dbHandle -> + do queueKey <- newTVarIO 0 + jobMap <- SMap.newIO + jobQueue <- C.newChan lmc_maxQueueSize + progressSet <- SSet.newIO + databaseRecovery + lmc_inProgressRecovery queueKey jobMap jobQueue progressSet + (j_decode lmc_jobEncoding) dbHandle + persistQueue <- newTQueueIO + termVar <- newTVarIO Continue + let alloc = async $ databaseWorker dbHandle (j_encode lmc_jobEncoding) termVar persistQueue + dealloc aHandle = + do atomically $ writeTVar termVar Terminate + wait aHandle + bracket alloc dealloc $ \_ -> + action + LevelMem + { lm_queueKey = queueKey + , lm_jobs = jobMap + , lm_queue = jobQueue + , lm_inProgress = progressSet + , lm_persistQueue = persistQueue + } + where + opts = L.defaultOptions { L.createIfMissing = True } + +-- | The data container for tracking jobs and persisting them to disk +data LevelMem j + = LevelMem + { lm_queueKey :: !(TVar JobIdx) + , lm_jobs :: !(SMap.Map JobIdx j) + , lm_queue :: !(Chan JobIdx) + , lm_inProgress :: !(SSet.Set JobIdx) + , lm_persistQueue :: !(TQueue (QueueAction j)) + } + +data JobStatus + = JQueued + -- ^ job is enqueued, not being worked on + | JInProgress + -- ^ job currently being worked on + deriving (Show, Eq) + +-- | Get a snapshot of the current state of all jobs. +getJobStatusMap :: LevelMem j -> IO [(j, JobStatus)] +getJobStatusMap LevelMem{..} = + atomically $ L.fold folder [] $ SMap.stream lm_jobs + where + folder jobs (k, v) = + do progress <- SSet.lookup k lm_inProgress + pure ((v, if progress then JInProgress else JQueued) : jobs) + +-- | Get the job status for a job. Note that this is a potentially +-- expensive operation as all known jobs must be traversed. A more efficient +-- handling would be to call 'getJobStatusMap' once beforce launching the +-- worker and then tracking status manually within the worker. +getJobStatus :: Eq j => j -> LevelMem j -> IO (Maybe JobStatus) +getJobStatus job LevelMem{..} = + -- not implemented using 'getJobStatusMap' to allow early termination + atomically $ L.foldMaybe folder Nothing $ SMap.stream lm_jobs + where + folder st (k, v) + | isJust st = pure Nothing + | v /= job = pure (Just st) + | otherwise = + do progress <- SSet.lookup k lm_inProgress + pure $ Just $ Just $ if progress then JInProgress else JQueued + +-- | Create a queue backend from 'LevelMem' +newLevelMemBackend :: LevelMem j -> QueueBackend j +newLevelMemBackend levelMem = + QueueBackend + { qb_lift = id + , qb_enqueue = flip enqueue levelMem + , qb_dequeue = dequeue levelMem + , qb_confirm = flip ack levelMem + , qb_rollback = flip nack levelMem + } + +enqueue :: j -> LevelMem j -> IO Bool +enqueue job lm@LevelMem{..} = + do ix <- + atomically $ + do key <- readTVar lm_queueKey + SMap.insert job key lm_jobs + writeTVar lm_queueKey (key+1) + writeTQueue lm_persistQueue (QaEnqueue key job) + return key + ok <- C.tryWriteChan (fst lm_queue) ix + unless ok $ ack ix lm + pure ok + +dequeue :: LevelMem j -> IO (JobIdx, j) +dequeue lm@LevelMem{..} = + do jobIdx <- C.readChan (snd lm_queue) + mJob <- + atomically $ + do SSet.insert jobIdx lm_inProgress + writeTQueue lm_persistQueue (QaDequeue jobIdx) + SMap.lookup jobIdx lm_jobs + case mJob of + Nothing -> dequeue lm + Just j -> pure (jobIdx, j) + +ack :: JobIdx -> LevelMem j -> IO () +ack jobIdx LevelMem{..} = + atomically $ + do SSet.delete jobIdx lm_inProgress + SMap.delete jobIdx lm_jobs + writeTQueue lm_persistQueue (QaConfirm jobIdx) + +nack :: JobIdx -> LevelMem j -> IO () +nack jobIdx LevelMem{..} = + do atomically $ + do SSet.delete jobIdx lm_inProgress + writeTQueue lm_persistQueue (QaRollback jobIdx) + C.writeChan (fst lm_queue) jobIdx diff --git a/test/Data/PowerQueue/Backend/LevelMemSpec.hs b/test/Data/PowerQueue/Backend/LevelMemSpec.hs new file mode 100644 --- /dev/null +++ b/test/Data/PowerQueue/Backend/LevelMemSpec.hs @@ -0,0 +1,90 @@ +module Data.PowerQueue.Backend.LevelMemSpec + ( spec ) +where + +import Data.PowerQueue +import Data.PowerQueue.Backend.LevelMem + +import Control.Concurrent.Async +import Control.Monad +import Data.IORef +import Data.List +import System.IO.Temp +import Test.Hspec +import qualified Data.Serialize as S + +dummyRefWorker :: (a -> b -> b) -> IORef b -> QueueWorker a +dummyRefWorker f ioRef = + newQueueWorker $ \v -> atomicModifyIORef' ioRef $ \i -> (f v i, JOk) + +dummyRefQueue :: QueueBackend a -> (a -> b -> b) -> IORef b -> Queue a +dummyRefQueue be f ioRef = + newQueue be (dummyRefWorker f ioRef) + +dummyCounterQueue :: QueueBackend () -> IORef Int -> Queue () +dummyCounterQueue be = dummyRefQueue be (\_ i -> i + 1) + +dummyReverseAccumQueue :: QueueBackend Int -> IORef [Int] -> Queue Int +dummyReverseAccumQueue be = dummyRefQueue be (:) + +binEncoding :: S.Serialize a => JobEncoding a +binEncoding = + JobEncoding + { j_encode = S.encode + , j_decode = S.decode + } + +withTmpDir :: (FilePath -> IO a) -> IO a +withTmpDir = withSystemTempDirectory "lmstestsXXX" + +testCfg :: S.Serialize a => FilePath -> LevelMemCfg a +testCfg fp = + LevelMemCfg + { lmc_storageDir = fp + , lmc_maxQueueSize = 100 + , lmc_jobEncoding = binEncoding + , lmc_inProgressRecovery = IpRestart + } + +spec :: Spec +spec = + describe "Level Mem Backend" $ + do it "providing 100 jobs and stepping 100 times works" $ + withTmpDir $ \tmpDir -> + withLevelMem (testCfg tmpDir) $ \lm -> + do ref <- newIORef 0 + let queue = dummyCounterQueue (newLevelMemBackend lm) ref + replicateM_ 100 $ enqueueJob () queue + replicateM_ 100 $ workStep queue + result <- readIORef ref + result `shouldBe` 100 + it "synchronous recovery works" $ + withTmpDir $ \tmpDir -> + do ref <- newIORef 0 + withLevelMem (testCfg tmpDir) $ \lm -> + do let queue = dummyCounterQueue (newLevelMemBackend lm) ref + replicateM_ 100 $ enqueueJob () queue + replicateM_ 50 $ workStep queue + withLevelMem (testCfg tmpDir) $ \lm -> + do statusMap <- getJobStatusMap lm + length statusMap `shouldBe` 50 + let queue = dummyCounterQueue (newLevelMemBackend lm) ref + replicateM_ 50 $ workStep queue + result <- readIORef ref + result `shouldBe` 100 + it "asynchronous working and recovery works" $ + withTmpDir $ \tmpDir -> + do ref <- newIORef [] + withLevelMem (testCfg tmpDir) $ \lm -> + do let queue = dummyReverseAccumQueue (newLevelMemBackend lm) ref + adder <- async $ forConcurrently_ [1..100] $ \idx -> enqueueJob idx queue + worker <- async $ replicateConcurrently_ 10 $ replicateM_ 5 $ workStep queue + wait adder + wait worker + withLevelMem (testCfg tmpDir) $ \lm -> + do statusMap <- getJobStatusMap lm + length statusMap `shouldBe` 50 + let queue = dummyReverseAccumQueue (newLevelMemBackend lm) ref + replicateM_ 50 $ workStep queue + result <- sort <$> readIORef ref + result `shouldBe` [1..100] diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,7 @@ +import Test.Hspec +import qualified Data.PowerQueue.Backend.LevelMemSpec + +main :: IO () +main = + hspec $ + Data.PowerQueue.Backend.LevelMemSpec.spec