diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright James Dabbs (c) 2016 + +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 James Dabbs 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/app/Main.hs b/app/Main.hs new file mode 100644 --- /dev/null +++ b/app/Main.hs @@ -0,0 +1,94 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TemplateHaskell #-} +module Main where + +import Keenser + +import Control.Concurrent +import Control.Monad +import Control.Monad.Logger +import Data.Aeson +import qualified Data.ByteString as BS +import Data.Monoid +import qualified Data.Text as T +import Database.Redis +import System.Environment +import System.Log.FastLogger + +s :: Show a => a -> T.Text +s = T.pack . show + +instance MonadLogger IO where + monadLoggerLog loc src lvl msg = void . forkIO . BS.putStr . fromLogStr . defaultLogStr loc src lvl $ toLogStr msg + +count :: Worker IO Int +count = Worker "count" "default" $ \n -> + forM_ [1 .. n] $ \i -> do + $(logInfo) $ s i <> " / " <> s n + sleep 1 + +crash :: Worker IO () +crash = Worker "crash" "default" $ \_ -> do + $(logInfo) "one ..." + sleep 1 + $(logInfo) "two ..." + sleep 1 + $(logInfo) "five ..." + sleep 1 + error "Three, sir" + +nope :: Worker IO Int +nope = Worker "nope" "default" $ \n -> do + $(logError) $ "Noping out " <> s n + error $ "NOPE " ++ show n + +notify :: T.Text -> Middleware IO +notify str _ _ j q run = do + $(logDebug) $ str <> " - starting job " <> s q <> " " <> s (encode j) + run + $(logDebug) $ str <> " - done" + +noop :: Middleware IO +noop _ _ _ _ run = run + +defaults :: [a] -> [a] -> [a] +defaults (_:as) (b:bs) = b : defaults as bs +defaults as [] = as +defaults _ _ = [] + +main :: IO () +main = do + args <- getArgs + let par:counts:crashes:_ = defaults [5,20,0] $ map read args + + putStrLn $ show par ++ " workers, " ++ show counts ++ " passes, " ++ show crashes ++ " fails" + + conn <- connect defaultConnectInfo + conf <- mkConf conn $ do + middleware record + middleware retry + -- middleware $ notify "[Middleware 1] " + -- middleware $ notify "[Middleware 2] " + concurrency par + register count + register crash + register nope + + m <- startProcess conf + + -- enqueueIn 60 m count 3 + -- enqueue m nope 42 + + replicateM_ crashes $ enqueue m crash () + forM_ [1 .. counts] $ enqueue m count + + forkIO . forever $ + getLine >>= handleCmd m + + await m + +handleCmd :: Manager -> String -> IO () +handleCmd m t + | t == "q" = putStrLn "Stopping" >> halt m + | t == "" = return () + | otherwise = putStrLn "Press `q` to quit" diff --git a/keenser.cabal b/keenser.cabal new file mode 100644 --- /dev/null +++ b/keenser.cabal @@ -0,0 +1,79 @@ +name: keenser +version: 0.1.0.0 +synopsis: Initial project template from stack +description: Please see README.md +homepage: https://github.com/jamesdabbs/keenser#readme +license: BSD3 +license-file: LICENSE +author: James Dabbs +maintainer: jamesdabbs@gmail.com +copyright: 2016 +category: Web +build-type: Simple +-- extra-source-files: +cabal-version: >=1.10 + +library + hs-source-dirs: src + ghc-options: -Wall + exposed-modules: Keenser + , Keenser.Middleware + , Keenser.Middleware.Retry + , Keenser.Middleware.Stats + , Keenser.Types + , Keenser.Util + other-modules: Keenser.Import + build-depends: base >= 4.7 && < 5 + , aeson + , attoparsec + , bytestring + , containers + , hedis + , hostname + , old-locale + , lifted-base + , monad-control + , monad-logger + , random + , scientific + , stm + , text + , time + , thyme + , transformers + , transformers-base + , unix + , unordered-containers + , vector-space + default-language: Haskell2010 + default-extensions: LambdaCase + , OverloadedStrings + , RecordWildCards + +executable keenser-exe + hs-source-dirs: app + main-is: Main.hs + ghc-options: -threaded -rtsopts -with-rtsopts=-N + build-depends: base + , keenser + , aeson + , bytestring + , fast-logger + , hedis + , monad-logger + , text + , unix + default-language: Haskell2010 + +test-suite keenser-test + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Spec.hs + build-depends: base + , keenser + ghc-options: -threaded -rtsopts -with-rtsopts=-N + default-language: Haskell2010 + +source-repository head + type: git + location: https://github.com/jamesdabbs/keenser diff --git a/src/Keenser.hs b/src/Keenser.hs new file mode 100644 --- /dev/null +++ b/src/Keenser.hs @@ -0,0 +1,301 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE TemplateHaskell #-} +module Keenser + ( Config + , Configurator + , Manager + , ManagerStatus(..) + , Middleware + , Worker(..) + , await + , checkStatus + , concurrency + , halt + , enqueue + , enqueueAt + , enqueueIn + , mkConf + , middleware + , record + , register + , retry + , sleep + , startProcess + ) where + +import Control.Concurrent (ThreadId) +import Control.Concurrent.Lifted (fork, forkFinally, killThread, threadDelay) +import Control.Concurrent.STM (atomically) +import Control.Concurrent.STM.TVar (newTVarIO, swapTVar, readTVarIO, writeTVar) +import Control.Monad.Logger +import Control.Monad.Trans.State +import Data.Aeson +import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as BSC +import qualified Data.ByteString.Lazy as LBS +import qualified Data.Map as M +import qualified Data.Text as T +import Data.Text.Encoding (decodeUtf8) +import Database.Redis hiding (decode) +import Network.HostName (getHostName) +import System.Posix.Process (getProcessID) +import System.Posix.Signals +import System.Posix.Types (CPid(..)) + +import Debug.Trace + +import Keenser.Import +import Keenser.Middleware +import Keenser.Types +import Keenser.Util + + +redisTimeout :: Integer +redisTimeout = 30 + +heartBeatFreq :: Int +heartBeatFreq = 5 + +-- TODO: check for and handle Redis errors throughout + +mkConf :: Monad m => Connection -> Configurator m -> m (Config m) +mkConf conn conf = execStateT conf Config + { kWorkers = M.empty + , kQueues = ["default"] + , kConcurrency = 25 + , kRedis = conn + , kMiddleware = [] + } + +concurrency :: Monad m => Int -> Configurator m +concurrency n = modify $ \c -> c { kConcurrency = n } + +-- TODO: +-- - verify that worker names are unique +-- - update queue list (or at least warn if there are queues that aren't being covered) +register :: (MonadLogger m, MonadIO m, FromJSON a) => Worker m a -> Configurator m +register Worker{..} = modify $ \c -> c { kWorkers = M.insert workerName w $ kWorkers c } + where + w = Worker workerName workerQueue $ \obj -> + case fromJSON obj of + Data.Aeson.Success job -> workerPerform job + Data.Aeson.Error err -> $(logError) $ "job failed to parse: " <> T.pack (show obj) <> " / " <> T.pack err + + +enqueue :: (ToJSON a, MonadIO m) => Manager -> Worker m a -> a -> m () +enqueue Manager{..} w args = do + job <- mkJob w args + liftIO . void . runRedis managerRedis $ queue job + +enqueueAt :: (ToJSON a, MonadIO m) => UTCTime -> Manager -> Worker m a -> a -> m () +enqueueAt at Manager{..} w args = do + job <- mkJob w args + liftIO . void . runRedis managerRedis $ zadd "schedule" [(timeToDouble at, LBS.toStrict $ encode job)] + +enqueueIn :: (ToJSON a, MonadIO m) => Rational -> Manager -> Worker m a -> a -> m () +enqueueIn snds m w args = do + now <- liftIO getCurrentTime + enqueueAt (snds `secondsFrom` now) m w args + + +-- TODO: +-- - handle the various failure modes here (timeout, can't match worker, can't get work) +-- - if doWork crashes, we should requeue the job, kill this thread, and restart a new processor +-- - track down the intermittent ": hFlush: illegal operation (handle is closed)" error from here +startProcessor :: (MonadBaseControl IO m, MonadIO m, MonadLogger m) + => Config m -> Manager -> m () +startProcessor Config{..} m = repeatUntil (managerStopping m) $ do + let mqs = map (\q -> "queue:" <> q) $ managerQueues m + ejob <- liftIO . runRedis kRedis $ brpop mqs redisTimeout + case ejob of + Right (Just (q, jjob)) -> case dispatch kWorkers jjob of + Just (worker, Object obj, job) -> + runMiddleware + kMiddleware m worker obj q + (workerPerform worker $ jobArgs job) + _ -> $(logError) $ "could not find worker for " <> decodeUtf8 jjob + _ -> $(logDebug) "Nothing to do here" + + +dispatch :: M.Map WorkerName (Worker m Value) -> BS.ByteString -> Maybe (Worker m Value, Value, Job Value) +dispatch workers payload = do + jobj <- parseMaybe json payload + job <- getJob jobj + worker <- M.lookup (jobClass job) workers + return $! (worker, jobj, job) + +getJob :: Value -> Maybe (Job Value) +getJob v = case fromJSON v of + Success j -> Just j + _ -> Nothing + +listenForSignals :: (MonadLogger m, MonadBaseControl IO m, MonadIO m) => Manager -> m () +listenForSignals m@Manager{..} = void . fork . forever $ do + esig <- liftIO . runRedis managerRedis $ brpop [managerIdentity <> "-signals"] redisTimeout + case esig of + Right (Just (_, sig)) -> handleSignal m sig + _ -> return () + +handleSignal :: (MonadLogger m, MonadBaseControl IO m, MonadIO m) => Manager -> BS.ByteString -> m () +handleSignal m sig + | sig == "USR1" = quiet m + | sig == "TERM" = halt m + | otherwise = $(logError) $ "Unrecognized signal: " <> decodeUtf8 sig + +startProcess :: (MonadBaseControl IO m, MonadIO m, MonadLogger m) => Config m -> m Manager +startProcess c@Config{..} = do + m <- mkManager c + startBeat m + listenForSignals m + startPolling m + + forM_ [1 .. kConcurrency] $ \n -> + let name = "processor " <> T.pack (show n) + in forkWatch m name $ startProcessor c m + return $! m + +startPolling :: (MonadBaseControl IO m, MonadLogger m, MonadIO m) => Manager -> m () +startPolling m@Manager{..} = forkWatch m "poller" . forever $ do + forM_ ["schedule", "retry"] $ \q -> do + er <- liftIO $ do + now <- getCurrentTime + runRedis managerRedis $ zrangebyscoreLimit q 0 (timeToDouble now) 0 1 + case er of + Left err -> $(logError) $ "Polling error " <> T.pack (show err) + Right jobs -> forM_ jobs $ liftIO . runRedis managerRedis . requeueFrom q + liftIO $ sleep 5 -- TODO: adaptive poll interval + +requeueFrom :: Queue -> BS.ByteString -> Redis () +requeueFrom q payload = case decode $ LBS.fromStrict payload of + Just job -> do + let q2 = jobQueue (job :: Job Value) + lpush q2 [payload] + zrem q [payload] + return () + _ -> return () + + +mkManager :: MonadIO m => Config m -> m Manager +mkManager Config{..} = liftIO $ do + now <- getCurrentTime + host <- getHostName + CPid pid <- getProcessID + nonce <- randomHex 6 + + let key = BSC.pack $ concat [host, ".", show pid, ".", nonce] + + Manager (BSC.pack host) now pid kConcurrency kQueues [] key kRedis + <$> newTVarIO M.empty + <*> newTVarIO 0 + <*> newTVarIO 0 + <*> newTVarIO False + +clearRedis :: Manager -> Bool -> IO () +clearRedis Manager{..} raise = do + _ <- liftIO . void . runRedis managerRedis $ do + srem "processes" [managerIdentity] + del [managerIdentity] + when raise $ + raiseSignal keyboardSignal + +startBeat :: (MonadBaseControl IO m, MonadLogger m, MonadIO m) => Manager -> m ThreadId +startBeat m@Manager{..} = do + liftIO $ installHandler keyboardSignal (CatchOnce $ clearRedis m True) Nothing + + forkFinally (heartBeat m) $ \e -> do + case e of + Left _ -> $(logError) $ "Heartbeat error " <> T.pack (show e) + _ -> $(logDebug) $ "Stopping heartbeat" + liftIO $ clearRedis m False + +workload :: MonadIO m => Manager -> m [RunningJob] +workload Manager{..} = liftIO $ M.elems <$> readTVarIO managerRunning + +workToRedis :: RunningJob -> (BS.ByteString, BS.ByteString) +workToRedis j = (BSC.pack . trim "ThreadId " . show $ rjThread j, LBS.toStrict $ encode j) + +heartBeat :: MonadIO m => Manager -> m () +heartBeat m@Manager{..} = liftIO $ do + runRedis managerRedis $ do + _ <- del ["queues"] + sadd "processes" [managerIdentity] + sadd "queues" managerQueues + + forever $ do + now <- getCurrentTime + working <- workload m + stopping <- readTVarIO managerQuiet + dones <- atomically $ swapTVar managerComplete 0 + fails <- atomically $ swapTVar managerFailed 0 + + -- For what it's worth, an abort between here and the end of the block could cause us to under-report stats + runRedis managerRedis $ do + hmset managerIdentity + [ ("beat", timestamp now) + , ("info", LBS.toStrict $ encode m) + , ("busy", BSC.pack . show $ length working) + , ("quiet", boolToRedis stopping) + ] + _ <- del [managerIdentity <> ":workers"] + hmset (managerIdentity <> ":workers") $ map workToRedis working + incrby "stat:processed" dones + incrby ("stat:complete:" <> daystamp now) dones + incrby "stat:failed" fails + incrby ("stat:failed:" <> daystamp now) dones + expire managerIdentity 60 + sleep heartBeatFreq + +checkStatus :: MonadIO m => Manager -> Connection -> m ManagerStatus +checkStatus Manager{..} conn = liftIO . runRedis conn $ do + procs <- smembers "processes" + info <- hgetall managerIdentity + done <- Database.Redis.get "stat:processed" + failed <- Database.Redis.get "stat:failed" + qs <- mapM getQueueLength managerQueues + return $! ManagerStatus (d procs []) (d info []) (d' done 0) (d' failed 0) qs + +-- TODO: clean this up +d :: Either Reply a -> a -> a +d (Right a) _ = a +d _ a = a + +d' :: Read a => Either Reply (Maybe BSC.ByteString) -> a -> a +d' (Right (Just a)) _ = read $ BSC.unpack a +d' _ a = a + +getQueueLength :: BS.ByteString -> Redis (BS.ByteString, Integer) +getQueueLength q = do + len <- llen $ "queue:" <> q + return $! (q, d len 0) + +forkWatch :: (MonadLogger m, MonadBaseControl IO m, MonadIO m) => Manager -> T.Text -> m () -> m () +forkWatch m name a = do + $(logInfo) $ name <> " starting." + void $ forkFinally a restart + where + restart (Left e) = do + $(logError) $ name <> " exited " <> T.pack (show e) <> ". Restarting." + forkWatch m name a + restart (Right _) = $(logDebug) $ name <> " exited clean." + +await :: Manager -> IO () +await m = do + repeatUntil (managerStopping m) $ sleep 2 + putStrLn "Shutting down ..." + sleep 2 + +sleep :: Int -> IO () +sleep n = threadDelay $ n * 1000000 + +halt :: (MonadBaseControl IO m, MonadIO m) => Manager -> m () +halt m@Manager{..} = do + quiet m + running <- liftIO $ readTVarIO managerRunning + mapM_ killThread $ M.keys running + +quiet :: MonadIO m => Manager -> m () +quiet Manager{..} = liftIO . atomically $ writeTVar managerQuiet True + +managerStopping :: MonadIO m => Manager -> m Bool +managerStopping = liftIO . readTVarIO . managerQuiet diff --git a/src/Keenser/Import.hs b/src/Keenser/Import.hs new file mode 100644 --- /dev/null +++ b/src/Keenser/Import.hs @@ -0,0 +1,41 @@ +module Keenser.Import + ( module Keenser.Import + ) where + +import Control.Monad as Keenser.Import +import Control.Monad.IO.Class as Keenser.Import (MonadIO, liftIO) +import Control.Monad.Trans.Control as Keenser.Import (MonadBaseControl) +import Control.Monad.Trans.State as Keenser.Import (modify) +import Data.Aeson as Keenser.Import (Value, toJSON, fromJSON, parseJSON, encode, decode) +import Data.Maybe as Keenser.Import (fromMaybe) +import Data.Monoid as Keenser.Import ((<>)) +import Data.Thyme.Clock as Keenser.Import (UTCTime, getCurrentTime) + +import Prelude as Keenser.Import + +import Data.AffineSpace ((.+^)) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Char8 as BSC +import Data.Thyme.Clock (fromSeconds) +import Data.Thyme.Format (formatTime, readTime) +import qualified Data.Scientific as SC +import System.Locale (defaultTimeLocale) + + +timestamp :: UTCTime -> BS.ByteString +timestamp = BSC.pack . formatTime defaultTimeLocale "%s%Q" + +daystamp :: UTCTime -> BS.ByteString +daystamp = BSC.pack . formatTime defaultTimeLocale "%Y-%m-%d" + +timeToJson :: UTCTime -> SC.Scientific +timeToJson = read . BSC.unpack . timestamp + +jsonToTime :: SC.Scientific -> UTCTime +jsonToTime = readTime defaultTimeLocale "%s%Q" . SC.formatScientific SC.Fixed (Just 6) + +timeToDouble :: UTCTime -> Double +timeToDouble = read . formatTime defaultTimeLocale "%s%Q" + +secondsFrom :: Rational -> UTCTime -> UTCTime +secondsFrom n time = time .+^ fromSeconds n diff --git a/src/Keenser/Middleware.hs b/src/Keenser/Middleware.hs new file mode 100644 --- /dev/null +++ b/src/Keenser/Middleware.hs @@ -0,0 +1,27 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TemplateHaskell #-} +module Keenser.Middleware + ( middleware + , record + , retry + , runMiddleware + ) where + +import Data.Aeson + +import Keenser.Types +import Keenser.Import + +import Keenser.Middleware.Retry +import Keenser.Middleware.Stats + +middleware :: Monad m => Middleware m -> Configurator m +middleware m = modify $ \c -> c { kMiddleware = m : kMiddleware c } + +runMiddleware :: Monad m + => [Middleware m] + -> Manager -> Worker m Value -> Object -> Queue + -> m () + -> m () +runMiddleware (t:ts) m w j q start = t m w j q $ runMiddleware ts m w j q start +runMiddleware _ _ _ _ _ start = start diff --git a/src/Keenser/Middleware/Retry.hs b/src/Keenser/Middleware/Retry.hs new file mode 100644 --- /dev/null +++ b/src/Keenser/Middleware/Retry.hs @@ -0,0 +1,62 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE TemplateHaskell #-} +module Keenser.Middleware.Retry + ( retry + ) where + +import Control.Exception.Lifted (SomeException, catch) +import Control.Monad.Logger +import Data.Aeson +import qualified Data.ByteString.Lazy as LBS +import qualified Data.HashMap.Strict as HM +import Database.Redis + +import Keenser.Import +import Keenser.Types + +import qualified Data.Text as T + +retry :: (MonadLogger m, MonadBaseControl IO m, MonadIO m) => Middleware m +retry Manager{..} _ job q inner = catch inner $ \e -> do + (count, ts, rJob) <- nextRetry e job q + void . liftIO $ if count < 10 + then + runRedis managerRedis $ + zadd "retry" [(timeToDouble ts, LBS.toStrict $ encode rJob)] + else do + now <- getCurrentTime + runRedis managerRedis $ do + zadd "dead" [(timeToDouble now, LBS.toStrict $ encode job)] + -- TODO: prune by count and time, customizable + zremrangebyrank "dead" 0 (-1000) + +-- TODO: + rand(30) * (count + 1) to prevent thundering herd +retryTime :: Integer -> UTCTime -> UTCTime +retryTime count start = fromInteger offset `secondsFrom` start + where offset = (count ^ 4) + 15 + +nextRetry :: (MonadLogger m, MonadIO m) + => SomeException -> Object -> Queue -> m (Integer, UTCTime, Object) +nextRetry ex old q = do + now <- liftIO getCurrentTime + + let + -- TODO: I don't love how stringly-typed this direct `Object` manipulation is, + -- but if we're staying consistent w/ Sidekiq's Redis API, we need to allow + -- middleware authors to jam whatever metadata they want on the Jobject + (count, status) = case HM.lookup "retry_count" old >>= mJSON of + Just n -> (n+1, ["retried_at" .= timeToJson now]) + Nothing -> ( 0, ["failed_at" .= timeToJson now]) + updates = HM.fromList $ + [ "queue" .= fromMaybe q (HM.lookup "retry_queue" old >>= mJSON) + , "error_message" .= show ex + , "error_class" .= ("SomeException" :: T.Text) + , "retry_count" .= count + ] ++ status + $(logInfo) $ "Retry number " <> T.pack (show count) + return $! (count, retryTime count now, HM.union updates old) + +mJSON :: FromJSON a => Value -> Maybe a +mJSON v = case fromJSON v of + Success a -> Just a + _ -> Nothing diff --git a/src/Keenser/Middleware/Stats.hs b/src/Keenser/Middleware/Stats.hs new file mode 100644 --- /dev/null +++ b/src/Keenser/Middleware/Stats.hs @@ -0,0 +1,35 @@ +{-# LANGUAGE FlexibleContexts #-} +module Keenser.Middleware.Stats + ( record + ) where + +import Control.Concurrent.Lifted (myThreadId) +import Control.Concurrent.STM (atomically) +import Control.Concurrent.STM.TVar (modifyTVar') +import Control.Exception.Lifted (SomeException, catch, throwIO) +import qualified Data.Map as M + +import Keenser.Import +import Keenser.Types + +record :: (MonadBaseControl IO m, MonadIO m) => Middleware m +record Manager{..} _ job q inner = do + tid <- liftIO myThreadId + now <- liftIO getCurrentTime + let + run = do + inner + liftIO . atomically $ do + modifyTVar' managerRunning $ M.delete tid + modifyTVar' managerComplete (+1) + + recordError :: (MonadBaseControl IO m, MonadIO m) => SomeException -> m () + recordError e = do + liftIO . atomically $ do + modifyTVar' managerRunning $ M.delete tid + modifyTVar' managerFailed (+1) + throwIO e + + work = RunningJob job q now tid + liftIO . atomically . modifyTVar' managerRunning $ M.insert tid work + catch run recordError diff --git a/src/Keenser/Types.hs b/src/Keenser/Types.hs new file mode 100644 --- /dev/null +++ b/src/Keenser/Types.hs @@ -0,0 +1,142 @@ +{-# LANGUAGE TypeSynonymInstances #-} +module Keenser.Types + ( Config(..) + , Configurator + , Job(..) + , JobId + , Manager(..) + , ManagerStatus(..) + , Middleware + , Queue + , RunningJob(..) + , Signal(..) + , Worker(..) + , WorkerName + ) where + +import Control.Concurrent (ThreadId) +import Control.Concurrent.STM.TVar (TVar) +import Control.Monad.Trans.State (StateT) +import Data.Aeson +import qualified Data.ByteString as BS +import Data.Int (Int32) +import qualified Data.Map as M +import qualified Data.Text as T +import Data.Text.Encoding (decodeUtf8, encodeUtf8) +import Database.Redis (Connection) + +import Keenser.Import + +data Config m = Config + { kWorkers :: M.Map WorkerName (Worker m Value) + , kQueues :: [Queue] + , kConcurrency :: Int + , kRedis :: Connection + , kMiddleware :: [Middleware m] + } + +type Configurator m = StateT (Config m) m () + +type JobId = BS.ByteString +data Job a = Job + { jobClass :: WorkerName + , jobArgs :: a + , jobId :: JobId + , jobRetry :: Bool + , jobEnqueuedAt :: UTCTime + , jobQueue :: Queue + } + +data Manager = Manager + { managerHostname :: BS.ByteString + , managerStartedAt :: UTCTime + , managerPid :: Int32 + , managerConcurrency :: Int + , managerQueues :: [Queue] + , managerLabels :: [T.Text] + , managerIdentity :: BS.ByteString + , managerRedis :: Connection + + , managerRunning :: TVar (M.Map ThreadId RunningJob) + , managerComplete :: TVar Integer + , managerFailed :: TVar Integer + , managerQuiet :: TVar Bool + } + +data ManagerStatus = ManagerStatus + { sProcs :: [BS.ByteString] + , sState :: [(BS.ByteString, BS.ByteString)] + , sDone :: Int + , sFailed :: Int + , sQueues :: [(BS.ByteString, Integer)] + } deriving Show + +type Middleware m = Manager -> Worker m Value -> Object -> Queue -> m () -> m () + +type Queue = BS.ByteString + +data RunningJob = RunningJob + { rjPayload :: Object + , rjQueue :: Queue + , rjRunAt :: UTCTime + , rjThread :: ThreadId + } + +data Signal = USR1 | TERM deriving (Show, Read) + +type WorkerName = BS.ByteString +data Worker m a = Worker + { workerName :: WorkerName + , workerQueue :: Queue + , workerPerform :: a -> m () + } + + +instance Functor Job where + fmap f (Job klass args _id retry at q) = Job klass (f args) _id retry at q + +instance ToJSON a => ToJSON (Job a) where + toJSON Job{..} = object + [ "class" .= jobClass + , "args" .= [jobArgs] + , "jid" .= jobId + , "retry" .= jobRetry + , "enqueued_at" .= timeToJson jobEnqueuedAt + , "queue" .= jobQueue + ] + +instance FromJSON a => FromJSON (Job a) where + parseJSON = withObject "job" $ \v -> do + jobClass <- v .: "class" + jobArgs <- head <$> v .: "args" + jobId <- v .: "jid" + jobRetry <- v .: "retry" + jobEnqueuedAt <- jsonToTime <$> v .: "enqueued_at" + jobQueue <- v .: "queue" + return $! Job{..} + + +instance ToJSON Manager where + toJSON Manager{..} = object + [ "hostname" .= managerHostname + , "started_at" .= timeToJson managerStartedAt + , "pid" .= managerPid + -- TODO: tag ~ "app" + , "concurrency" .= managerConcurrency + , "queues" .= managerQueues + -- TODO: labels ~ [] + , "identity" .= managerIdentity + ] + +instance ToJSON RunningJob where + toJSON RunningJob{..} = object + [ "queue" .= rjQueue + , "payload" .= fmap (: []) rjPayload -- since Sidekiq stores *args + , "run_at" .= timeToJson rjRunAt + ] + +instance ToJSON WorkerName where + toJSON = Data.Aeson.String . decodeUtf8 + +instance FromJSON WorkerName where + parseJSON = withText "worker_name" $ return . encodeUtf8 diff --git a/src/Keenser/Util.hs b/src/Keenser/Util.hs new file mode 100644 --- /dev/null +++ b/src/Keenser/Util.hs @@ -0,0 +1,53 @@ +module Keenser.Util where + +import Data.Aeson +import Data.Attoparsec.ByteString (Parser, parseOnly) +import qualified Data.ByteString.Char8 as BSC +import qualified Data.ByteString.Lazy as LBS +import qualified Data.List as L +import Database.Redis (Redis, lpush) +import System.Random (getStdGen, randomRs) + +import Keenser.Import +import Keenser.Types + +queue :: ToJSON a => Job a -> Redis () +queue job = void $ lpush ("queue:" <> jobQueue job) [LBS.toStrict $ encode job] + +randomHex :: Int -> IO String +randomHex n = map (digits !!) . take n . randomRs (0, 15) <$> getStdGen + +digits :: [Char] +digits = ['0' .. '9'] ++ ['A' .. 'F'] + +mkJob :: MonadIO m => Worker m a -> a -> m (Job a) +mkJob Worker{..} args = liftIO $ do + _id <- BSC.pack <$> randomHex 12 + now <- getCurrentTime + return $! Job workerName args _id True now workerQueue + +asJSON :: ToJSON a => a -> Object +asJSON a = case toJSON a of + Object o -> o + _ -> error "FIXME: define Job -> Object mapping directly" + +repeatUntil :: Monad m => m Bool -> m () -> m () +repeatUntil check act = do + done <- check + unless done $ do + act + repeatUntil check act + +rightToMaybe :: Either a b -> Maybe b +rightToMaybe (Right b) = Just b +rightToMaybe _ = Nothing + +trim :: Eq a => [a] -> [a] -> [a] +trim pre corp = fromMaybe corp $ L.stripPrefix pre corp + +parseMaybe :: Parser a -> BSC.ByteString -> Maybe a +parseMaybe p = rightToMaybe . parseOnly p + +boolToRedis :: Bool -> BSC.ByteString +boolToRedis True = "true" +boolToRedis False = "false" diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,2 @@ +main :: IO () +main = putStrLn "Test suite not yet implemented"