@@ -1,6 +1,5 @@
1−{-# LANGUAGE RankNTypes, FlexibleInstances, FlexibleContexts, PartialTypeSignatures, TupleSections, DeriveGeneric, UndecidableInstances #-}
2−{-# LANGUAGE ExistentialQuantification #-}
3−{-# LANGUAGE ScopedTypeVariables #-}
1+{-# LANGUAGE RankNTypes, FlexibleInstances, FlexibleContexts, PartialTypeSignatures, UndecidableInstances #-}
2+{-# LANGUAGE ExistentialQuantification, RecordWildCards, ScopedTypeVariables #-}
43
54module OddJobs.Job
65 (
@@ -14,6 +13,7 @@ module OddJobs.Job
1413 -- $config
1514 , Config(..)
1615 , ConcurrencyControl(..)
16+ , ResourceCfg(..)
1717
1818 -- * Creating/scheduling jobs
1919 --
@@ -21,6 +21,9 @@ module OddJobs.Job
2121 , createJob
2222 , scheduleJob
2323
24+ , createJobWithResources
25+ , scheduleJobWithResources
26+
2427 -- * @Job@ and associated data-types
2528 --
2629 -- $dataTypes
@@ -33,6 +36,9 @@ module OddJobs.Job
3336 , Seconds(..)
3437 , JobErrHandler(..)
3538 , AllJobTypes(..)
39+ , ResourceId(..)
40+ , FunctionName
41+ , RunnerEnv(..)
3642
3743 -- ** Structured logging
3844 --
@@ -47,6 +53,7 @@ module OddJobs.Job
4753 , jobEventListener
4854 , jobPoller
4955 , jobPollingSql
56+ , pollRunJob
5057 , JobRunner
5158 , HasJobRunner (..)
5259
@@ -58,6 +65,7 @@ module OddJobs.Job
5865 , runJobNowIO
5966 , unlockJobIO
6067 , cancelJobIO
68+ , killJobIO
6169 , jobDbColumns
6270 , jobPollingIO
6371 , concatJobDbColumns
@@ -80,7 +88,7 @@ import Data.Pool(Pool)
8088import Data.Text as T
8189import Database.PostgreSQL.Simple as PGS
8290import Database.PostgreSQL.Simple.Notification
83−import UnliftIO.Async
91+import UnliftIO.Async hiding (poll)
8492import UnliftIO.Concurrent (threadDelay, myThreadId)
8593import Data.String
8694import System.Posix.Process (getProcessID)
@@ -91,7 +99,8 @@ import Control.Monad.Logger as MLogger (LogLevel(..), LogStr, toLogStr)
9199import UnliftIO.IORef
92100import UnliftIO.Exception ( SomeException(..), try, catch, finally
93101 , catchAny, bracket, Exception(..), throwIO
94− , catches, Handler(..), mask_, throwString
102+ , catches, Handler(..), mask_, onException
103+ , throwString
95104 )
96105import Data.Proxy
97106import Control.Monad.Trans.Control
@@ -102,14 +111,16 @@ import Data.Aeson hiding (Success)
102111import qualified Data.Aeson as Aeson
103112import qualified Data.Aeson.Types as Aeson (Parser, parseMaybe)
104113import Data.String.Conv (StringConv(..), toS)
105−import Data.Functor (void)
106−import Control.Monad (forever)
114+import Data.Functor ((<&>), void)
115+import Control.Monad (forever, forM_, join)
107116import Data.Maybe (isNothing, maybe, fromMaybe, listToMaybe, mapMaybe)
108117import Data.Either (either)
109118import Control.Monad.Reader
110119import GHC.Generics
120+import Data.Map (Map)
111121import qualified Data.HashMap.Strict as HM
112122import qualified Data.List as DL
123+import qualified Data.Map as DM
113124import qualified Data.ByteString as BS
114125import qualified Data.ByteString.Lazy as BSL
115126import System.FilePath (FilePath)
@@ -143,6 +154,7 @@ class (MonadUnliftIO m, MonadBaseControl IO m) => HasJobRunner m where
143154 getConcurrencyControl :: m ConcurrencyControl
144155 log :: LogLevel -> LogEvent -> m ()
145156 getDefaultJobTimeout :: m Seconds
157+ onJobTimeout :: Job -> m ()
146158 getDefaultRetryBackoff :: Int -> m Seconds
147159
148160-- $logging
@@ -160,7 +172,7 @@ class (MonadUnliftIO m, MonadBaseControl IO m) => HasJobRunner m where
160172
161173data RunnerEnv = RunnerEnv
162174 { envConfig :: !Config
163− , envJobThreadsRef :: !(IORef [Async ()])
175+ , envJobThreadsRef :: !(IORef (Map JobId (Async ())))
164176 }
165177
166178type RunnerM = ReaderT RunnerEnv IO
@@ -169,34 +181,37 @@ logCallbackErrors :: (HasJobRunner m) => JobId -> Text -> m () -> m ()
169181logCallbackErrors jid msg action = catchAny action $ \e -> log LevelError $ LogText $ msg <> " Job ID=" <> toS (show jid) <> ": " <> toS (show e)
170182
171183instance HasJobRunner RunnerM where
172− getPollingInterval = cfgPollingInterval . envConfig <$> ask
173− onJobFailed = do { env <- ask; pure (cfgOnJobFailed (envConfig env)) }
184+ getPollingInterval = asks (cfgPollingInterval . envConfig)
185+ onJobFailed = asks (cfgOnJobFailed . envConfig)
174186 onJobSuccess job = do
175− fn <- cfgOnJobSuccess . envConfig <$> ask
187+ fn <- asks (cfgOnJobSuccess . envConfig)
176188 logCallbackErrors (jobId job) "onJobSuccess" $ liftIO $ fn job
177− deleteSuccessfulJobs = cfgDeleteSuccessfulJobs . envConfig <$> ask
189+ deleteSuccessfulJobs = asks (cfgDeleteSuccessfulJobs . envConfig)
178190
179− getJobRunner = cfgJobRunner . envConfig <$> ask
180− getDbPool = cfgDbPool . envConfig <$> ask
181− getTableName = cfgTableName . envConfig <$> ask
191+ getJobRunner = asks (cfgJobRunner . envConfig)
192+ getDbPool = asks (cfgDbPool . envConfig)
193+ getTableName = asks (cfgTableName . envConfig)
182194 onJobStart job = do
183− fn <- cfgOnJobStart . envConfig <$> ask
195+ fn <- asks (cfgOnJobStart . envConfig)
184196 logCallbackErrors (jobId job) "onJobStart" $ liftIO $ fn job
185197
186− getDefaultMaxAttempts = cfgDefaultMaxAttempts . envConfig <$> ask
198+ getDefaultMaxAttempts = asks (cfgDefaultMaxAttempts . envConfig)
187199
188200 getRunnerEnv = ask
189201
190− getConcurrencyControl = (cfgConcurrencyControl . envConfig <$> ask)
202+ getConcurrencyControl = asks (cfgConcurrencyControl . envConfig)
191203
192204 log logLevel logEvent = do
193− loggerFn <- cfgLogger . envConfig <$> ask
205+ loggerFn <- asks (cfgLogger . envConfig)
194206 liftIO $ loggerFn logLevel logEvent
195207
196− getDefaultJobTimeout = cfgDefaultJobTimeout . envConfig <$> ask
208+ getDefaultJobTimeout = asks (cfgDefaultJobTimeout . envConfig)
209+ onJobTimeout job = do
210+ fn <- asks (cfgOnJobTimeout . envConfig)
211+ logCallbackErrors (jobId job) "onJobTimeout" $ liftIO $ fn job
197212
198213 getDefaultRetryBackoff attempts = do
199− retryFn <- cfgDefaultRetryBackoff . envConfig <$> ask
214+ retryFn <- asks (cfgDefaultRetryBackoff . envConfig)
200215 liftIO $ retryFn attempts
201216
202217-- | Start the job-runner in the /current/ thread, i.e. you'll need to use
@@ -206,7 +221,7 @@ instance HasJobRunner RunnerM where
206221startJobRunner :: Config -> IO ()
207222startJobRunner jm = do
208223 traceM "startJobRunner top"
209− r <- newIORef []
224+ r <- newIORef DM.empty
210225 let monitorEnv = RunnerEnv
211226 { envConfig = jm
212227 , envJobThreadsRef = r
@@ -218,7 +233,7 @@ jobWorkerName :: IO String
218233jobWorkerName = do
219234 pid <- getProcessID
220235 hname <- getHostName
221− pure $ hname ++ ":" ++ (show pid)
236+ pure $ hname ++ ":" ++ show pid
222237
223238-- | If you are writing SQL queries where you want to return ALL columns from
224239-- the jobs table it is __recommended__ that you do not issue a @SELECT *@ or
@@ -249,7 +264,7 @@ concatJobDbColumns :: (IsString s, Semigroup s) => s
249264concatJobDbColumns = concatJobDbColumns_ jobDbColumns ""
250265 where
251266 concatJobDbColumns_ [] x = x
252− concatJobDbColumns_ (col:[]) x = x <> col
267+ concatJobDbColumns_ [col] x = x <> col
253268 concatJobDbColumns_ (col:cols) x = concatJobDbColumns_ cols (x <> col <> ", ")
254269
255270
@@ -290,7 +305,7 @@ findJobByIdIO :: Connection -> TableName -> JobId -> IO (Maybe Job)
290305findJobByIdIO conn tname jid = PGS.query conn findJobByIdQuery (tname, jid) >>= \case
291306 [] -> pure Nothing
292307 [j] -> pure (Just j)
293− js -> Prelude.error $ "Not expecting to find multiple jobs by id=" <> (show jid)
308+ _js -> Prelude.error $ "Not expecting to find multiple jobs by id=" <> show jid
294309
295310
296311saveJobQuery :: PGS.Query
@@ -307,7 +322,7 @@ saveJob j = do
307322saveJobIO :: Connection -> TableName -> Job -> IO Job
308323saveJobIO conn tname Job{jobRunAt, jobStatus, jobPayload, jobLastError, jobAttempts, jobLockedBy, jobLockedAt, jobId} = do
309324 rs <- PGS.query conn saveJobQuery
310− ( tname
325+ ( tname
311326 , jobRunAt
312327 , jobStatus
313328 , jobPayload
@@ -318,9 +333,9 @@ saveJobIO conn tname Job{jobRunAt, jobStatus, jobPayload, jobLastError, jobAttem
318333 , jobId
319334 )
320335 case rs of
321− [] -> Prelude.error $ "Could not find job while updating it id=" <> (show jobId)
336+ [] -> Prelude.error $ "Could not find job while updating it id=" <> show jobId
322337 [j] -> pure j
323− js -> Prelude.error $ "Not expecting multiple rows to ber returned when updating job id=" <> (show jobId)
338+ _js -> Prelude.error $ "Not expecting multiple rows to ber returned when updating job id=" <> show jobId
324339
325340deleteJob :: (HasJobRunner m) => JobId -> m ()
326341deleteJob jid = do
@@ -334,26 +349,30 @@ deleteJobIO conn tname jid = do
334349runJobNowIO :: Connection -> TableName -> JobId -> IO (Maybe Job)
335350runJobNowIO conn tname jid = do
336351 t <- getCurrentTime
337− updateJobHelper tname conn (Queued, [Queued, Retry, Failed], Just t, jid)
352+ updateJobHelper tname conn (Queued, [Queued, Retry, Failed, Cancelled], Just t, jid)
338353
339354-- | TODO: First check in all job-runners if this job is still running, or not,
340355-- and somehow send an uninterruptibleCancel to that thread.
341356unlockJobIO :: Connection -> TableName -> JobId -> IO (Maybe Job)
342357unlockJobIO conn tname jid = do
343− fmap listToMaybe $ PGS.query conn q (tname, Retry, jid, In [Locked])
358+ listToMaybe <$> PGS.query conn q (tname, Retry, jid, In [Locked])
344359 where
345360 q = "update ? set status=?, run_at=now(), locked_at=null, locked_by=null where id=? and status in ? returning " <> concatJobDbColumns
346361
347362cancelJobIO :: Connection -> TableName -> JobId -> IO (Maybe Job)
348363cancelJobIO conn tname jid =
349− updateJobHelper tname conn (Failed, [Queued, Retry], Nothing, jid)
364+ updateJobHelper tname conn (Cancelled, [Queued, Retry], Nothing, jid)
365+
366+killJobIO :: Connection -> TableName -> JobId -> IO (Maybe Job)
367+killJobIO conn tname jid =
368+ updateJobHelper tname conn (Cancelled, [Locked], Nothing, jid)
350369
351370updateJobHelper :: TableName
352371 -> Connection
353372 -> (Status, [Status], Maybe UTCTime, JobId)
354373 -> IO (Maybe Job)
355374updateJobHelper tname conn (newStatus, existingStates, mRunAt, jid) =
356− fmap listToMaybe $ PGS.query conn q (tname, newStatus, runAt, jid, PGS.In existingStates)
375+ listToMaybe <$> PGS.query conn q (tname, newStatus, runAt, jid, PGS.In existingStates)
357376 where
358377 q = "update ? set attempts=0, status=?, run_at=? where id=? and status in ? returning " <> concatJobDbColumns
359378 runAt = case mRunAt of
@@ -364,39 +383,47 @@ updateJobHelper tname conn (newStatus, existingStates, mRunAt, jid) =
364383data TimeoutException = TimeoutException deriving (Eq, Show)
365384instance Exception TimeoutException
366385
386+-- | this spawns the job on a a seperate thread, and
387+-- also spawns a monitoring thread, whichever finishes
388+-- first kills the other.
367389runJobWithTimeout :: (HasJobRunner m)
368390 => Seconds
369391 -> Job
370392 -> m ()
371−runJobWithTimeout timeoutSec job = do
393+runJobWithTimeout timeoutSec job@Job{jobId} = do
372394 threadsRef <- envJobThreadsRef <$> getRunnerEnv
373395 jobRunner_ <- getJobRunner
374396
375397 a <- async $ liftIO $ jobRunner_ job
376398
377− x <- atomicModifyIORef' threadsRef $ \threads -> (a:threads, DL.map asyncThreadId (a:threads))
399+ _x <- atomicModifyIORef' threadsRef $ \threads ->
400+ ( DM.insert jobId a threads
401+ , DL.map asyncThreadId $ DM.elems $ DM.insert jobId a threads
402+ )
403+
378404 -- liftIO $ putStrLn $ "Threads: " <> show x
379405 log LevelDebug $ LogText $ toS $ "Spawned job in " <> show (asyncThreadId a)
380406
381407 t <- async $ do
382408 delaySeconds timeoutSec
383− uninterruptibleCancel a
384409 throwIO TimeoutException
385410
386411 void $ finally
387412 (waitEitherCancel a t)
388− (atomicModifyIORef' threadsRef $ \threads -> (DL.delete a threads, ()))
413+ (atomicModifyIORef' threadsRef $ \threads -> (DM.delete jobId threads, ()))
389414
390415
416+-- | runs a job, blocks for as long as it's in progress
391417runJob :: (HasJobRunner m) => JobId -> m ()
392418runJob jid = do
393− (findJobById jid) >>= \case
419+ findJobById jid >>= \case
394420 Nothing -> Prelude.error $ "Could not find job id=" <> show jid
395421 Just job -> do
396422 startTime <- liftIO getCurrentTime
397423 lockTimeout <- getDefaultJobTimeout
398424 log LevelInfo $ LogJobStart job
399− (flip catches) [Handler $ timeoutHandler job startTime, Handler $ exceptionHandler job startTime] $ do
425+ flip catch (exceptionHandler job startTime) $ do
426+ onJobStart job
400427 runJobWithTimeout lockTimeout job
401428 endTime <- liftIO getCurrentTime
402429 shouldDeleteJob <- deleteSuccessfulJobs
@@ -408,7 +435,6 @@ runJob jid = do
408435 onJobSuccess newJob
409436 pure ()
410437 where
411− timeoutHandler job startTime (e :: TimeoutException) = retryOrFail (toException e) job startTime
412438 exceptionHandler job startTime (e :: SomeException) = retryOrFail (toException e) job startTime
413439 retryOrFail e job@Job{jobAttempts} startTime = do
414440 endTime <- liftIO getCurrentTime
@@ -422,27 +448,50 @@ runJob jid = do
422448 newJob <- saveJob job{ jobStatus=newStatus
423449 , jobLockedBy=Nothing
424450 , jobLockedAt=Nothing
425− , jobLastError=(Just $ toJSON $ show e) -- TODO: convert errors to json properly
426− , jobRunAt=(addUTCTime (fromIntegral $ unSeconds backoffInSeconds) t)
451+ , jobLastError=Just $ toJSON $ show e -- TODO: convert errors to json properly
452+ , jobRunAt=addUTCTime (fromIntegral $ unSeconds backoffInSeconds) t
427453 }
428454 case fromException e :: Maybe TimeoutException of
429− Nothing -> log logLevel $ LogJobFailed newJob e failureMode runTime
430− Just _ -> log logLevel $ LogJobTimeout newJob
431−
432− let tryHandler (JobErrHandler handler) res = case fromException e of
433− Nothing -> res
434− Just e_ -> void $ handler e_ newJob failureMode
435− handlers <- onJobFailed
436− liftIO $ void $ Prelude.foldr tryHandler (throwIO e) handlers
455+ Nothing -> do
456+ log logLevel $ LogJobFailed newJob e failureMode runTime
457+ let tryHandler (JobErrHandler handler) res = case fromException e of
458+ Nothing -> res
459+ Just e_ -> void $ handler e_ newJob failureMode
460+ handlers <- onJobFailed
461+ liftIO $ void $ Prelude.foldr tryHandler (throwIO e) handlers
462+
463+ Just _ -> do
464+ log logLevel $ LogJobTimeout newJob
465+ onJobTimeout newJob
466+
437467 pure ()
438468
469+killJob :: (HasJobRunner m) => JobId -> m ()
470+killJob jid = do
471+ threadsRef <- envJobThreadsRef <$> getRunnerEnv
472+ threads <- liftIO $ readIORef threadsRef
473+ mJob <- findJobById jid
474+
475+ case (mJob, jid `DM.lookup` threads) of
476+ (Just job, Just thread) -> do
477+ log LevelInfo $ LogKillJobSuccess job
478+ void $ finally
479+ (uninterruptibleCancel thread)
480+ (atomicModifyIORef' threadsRef $ \threads' -> (DM.delete jid threads', ()))
481+
482+ (Just job, Nothing) -> do
483+ log LevelInfo $ LogKillJobFailed job
484+
485+ (Nothing, _) ->
486+ log LevelError $ LogText $ "Unable to find job in db to kill, jobId = " <> toS (show jid)
487+
439488-- TODO: This might have a resource leak.
440489restartUponCrash :: (HasJobRunner m, Show a) => Text -> m a -> m ()
441490restartUponCrash name_ action = do
442491 traceM "restartUponCrash top"
443492 a <- async action
444493 finally (waitCatch a >>= fn) $ do
445− (log LevelInfo $ LogText $ "Received shutdown: " <> toS name_)
494+ log LevelInfo $ LogText $ "Received shutdown: " <> toS name_
446495 cancel a
447496 where
448497 fn x = do
@@ -461,9 +510,11 @@ jobMonitor = do
461510 traceM "jobMonitor top"
462511 a1 <- async $ restartUponCrash "Job poller" jobPoller
463512 a2 <- async $ restartUponCrash "Job event listener" jobEventListener
513+ a3 <- async $ restartUponCrash "Job Kill poller" killJobPoller
464514 traceM "jobMonitor after async"
465− finally (void $ waitAnyCatch [a1, a2]) $ do
515+ finally (void $ waitAnyCatch [a1, a2, a3]) $ do
466516 log LevelInfo (LogText "Stopping jobPoller and jobEventListener threads.")
517+ cancel a3
467518 cancel a2
468519 cancel a1
469520 log LevelInfo (LogText "Waiting for jobs to complete.")
@@ -471,30 +522,57 @@ jobMonitor = do
471522
472523-- | Ref: 'jobPoller'
473524jobPollingSql :: Query
474−jobPollingSql = "update ? set status = ?, locked_at = ?, locked_by = ?, attempts=attempts+1 WHERE id in (select id from ? where (run_at<=? AND ((status in ?) OR (status = ? and locked_at<?))) ORDER BY attempts ASC, run_at ASC LIMIT 1 FOR UPDATE) RETURNING id"
525+jobPollingSql =
526+ "update ? set status = ?, locked_at = ?, locked_by = ?, attempts=attempts+1\
527+ \ WHERE id in (select id from ? where (run_at<=? AND ((status in ?) OR (status = ? and locked_at<?))) \
528+ \ ORDER BY attempts ASC, run_at ASC LIMIT 1 FOR UPDATE) RETURNING id"
529+
530+jobPollingWithResourceSql :: Query
531+jobPollingWithResourceSql =
532+ " UPDATE ? SET status = ?, locked_at = ?, locked_by = ?, attempts = attempts + 1 \
533+ \ WHERE id in (select id from ? where (run_at<=? AND ((status in ?) OR (status = ? and locked_at<?))) \
534+ \ AND ?(id) \
535+ \ ORDER BY attempts ASC, run_at ASC LIMIT 1) \
536+ \ RETURNING id"
537+
538+-- | Ref: 'killJobPoller'
539+killJobPollingSql :: Query
540+killJobPollingSql =
541+ "UPDATE ? SET locked_at = NULL, locked_by = NULL \
542+ \ WHERE id IN (SELECT id FROM ? WHERE status = ? AND locked_by = ? AND locked_at <= ? \
543+ \ ORDER BY locked_at ASC LIMIT 1 FOR UPDATE \
544+ \ ) RETURNING id"
475545
476546waitForJobs :: (HasJobRunner m)
477547 => m ()
478548waitForJobs = do
479− threadsRef <- envJobThreadsRef <$> getRunnerEnv
480− readIORef threadsRef >>= \case
549+ curJobs <- getRunnerEnv >>= (readIORef . envJobThreadsRef) <&> DM.elems
550+ case curJobs of
481551 [] -> log LevelInfo $ LogText "All job-threads exited"
482552 as -> do
483553 tid <- myThreadId
484554 void $ waitAnyCatch as
485− log LevelDebug $ LogText $ toS $ "Waiting for " <> show (DL.length as) <> " jobs to complete before shutting down. myThreadId=" <> (show tid)
555+ log LevelDebug $ LogText $ toS $ "Waiting for " <> show (DL.length as) <> " jobs to complete before shutting down. myThreadId=" <> show tid
486556 delaySeconds (Seconds 1)
487557 waitForJobs
488558
559+data ConcurrencyAction
560+ = DontPoll
561+ | PollAny
562+ | PollWithResources ResourceCfg
563+
489564getConcurrencyControlFn :: (HasJobRunner m)
490− => m (m Bool)
565+ => m (Connection -> m ConcurrencyAction)
491566getConcurrencyControlFn = getConcurrencyControl >>= \case
492− UnlimitedConcurrentJobs -> pure $ pure True
493− MaxConcurrentJobs maxJobs -> pure $ do
494− curJobs <- getRunnerEnv >>= (readIORef . envJobThreadsRef)
495− pure $ (DL.length curJobs) < maxJobs
496− DynamicConcurrency fn -> pure $ liftIO fn
567+ UnlimitedConcurrentJobs -> pure $ const $ pure PollAny
568+ MaxConcurrentJobs maxJobs -> pure $ const $ do
569+ curJobs <- getRunnerEnv >>= (readIORef . envJobThreadsRef) <&> DM.elems
570+ pure $ pollIf $ DL.length curJobs < maxJobs
571+ ResourceLimits resCfg -> pure $ const $ pure $ PollWithResources resCfg
572+ DynamicConcurrency fn -> pure $ const $ pollIf <$> liftIO fn
497573
574+ where
575+ pollIf cond = if cond then PollAny else DontPoll
498576
499577jobPollingIO :: Connection -> String -> TableName -> Seconds -> IO [Only JobId]
500578jobPollingIO pollerDbConn processName tname lockTimeout = do
@@ -506,9 +584,9 @@ jobPollingIO pollerDbConn processName tname lockTimeout = do
506584 , processName
507585 , tname
508586 , t
509− , (In [Queued, Retry])
587+ , In [Queued, Retry]
510588 , Locked
511− , (addUTCTime (fromIntegral $ negate $ unSeconds lockTimeout) t))
589+ , addUTCTime (fromIntegral $ negate $ unSeconds lockTimeout) t)
512590
513591-- | Executes 'jobPollingSql' every 'cfgPollingInterval' seconds to pick up jobs
514592-- for execution. Uses @UPDATE@ along with @SELECT...FOR UPDATE@ to efficiently
@@ -527,32 +605,104 @@ jobPollingIO pollerDbConn processName tname lockTimeout = do
527605jobPoller :: (HasJobRunner m) => m ()
528606jobPoller = do
529607 processName <- liftIO jobWorkerName
530− pool <- getDbPool
531− tname <- getTableName
532− lockTimeout <- getDefaultJobTimeout
533− traceM "jobPoller just before log statement"
534608 log LevelInfo $ LogText $ toS $ "Starting the job monitor via DB polling with processName=" <> processName
535609 concurrencyControlFn <- getConcurrencyControlFn
536− withResource pool $ \pollerDbConn -> forever $ concurrencyControlFn >>= \case
537− False -> do
538− log LevelWarn $ LogText $ "NOT polling the job queue due to concurrency control"
539− -- If we can't run any jobs ATM, relax and wait for resources to free up
540− delayAction
541− True -> do
542− nextAction <- mask_ $ do
543− log LevelDebug $ LogText $ toS $ "[" <> processName <> "] Polling the job queue.."
544− r <- liftIO $ jobPollingIO pollerDbConn processName tname lockTimeout
545− case r of
546− -- When we don't have any jobs to run, we can relax a bit...
547− [] -> pure delayAction
548−
549− -- When we find a job to run, fork and try to find the next job without any delay...
550− [Only (jid :: JobId)] -> do
551− void $ async $ runJob jid
610+ pool <- getDbPool
611+ forever $ do
612+ concurencyPolicy <- withResource pool concurrencyControlFn
613+ case concurencyPolicy of
614+ DontPoll -> log LevelWarn $ LogText "NOT polling the job queue due to concurrency control"
615+ PollAny -> void $ pollRunJob processName Nothing
616+ PollWithResources resCfg -> void $ pollRunJob processName (Just resCfg)
617+
618+-- | Polls a job and runs it, or executes a delay action if no job was found
619+--
620+-- returns an async to the task that was started, allowing you
621+-- to block until it's finished
622+pollRunJob :: (HasJobRunner m) => String -> Maybe ResourceCfg -> m (Maybe (Async ()))
623+pollRunJob processName mResCfg = do
624+ tname <- getTableName
625+ -- note it's better to use fine grained pool connection,
626+ -- since data.pool already has it's internal resource cache.
627+ -- this way the user can configure how long a connection
628+ -- needs to remain open.
629+ pool <- getDbPool
630+ lockTimeout <- getDefaultJobTimeout
631+ join $ withResource pool $ \pollerDbConn -> mask_ $ do
632+ log LevelDebug $ LogText $ toS $ "[" <> processName <> "] Polling the job queue.."
633+ t <- liftIO getCurrentTime
634+ r <- case mResCfg of
635+ Nothing -> liftIO $
636+ PGS.query pollerDbConn jobPollingSql
637+ ( tname
638+ , Locked
639+ , t
640+ , processName
641+ , tname
642+ , t
643+ , In [Queued, Retry]
644+ , Locked
645+ , addUTCTime (fromIntegral $ negate $ unSeconds lockTimeout) t)
646+ Just ResourceCfg{..} -> liftIO $
647+ PGS.query pollerDbConn jobPollingWithResourceSql
648+ ( tname
649+ , Locked
650+ , t
651+ , processName
652+ , tname
653+ , t
654+ , In [Queued, Retry]
655+ , Locked
656+ , addUTCTime (fromIntegral $ negate $ unSeconds lockTimeout) t
657+ , resCfgCheckResourceFunction
658+ )
659+ case r of
660+ -- When we don't have any jobs to run, we can relax a bit...
661+ [] -> pure (Nothing <$ delayAction)
662+
663+ -- When we find a job to run, fork and try to find the next job without any delay...
664+ [Only (jid :: JobId)] -> do
665+ x <- async $ runJob jid
666+ pure $ Just x <$ noDelayAction
667+
668+ x -> error $ "WTF just happened? I was supposed to get only a single row, but got: " ++ show x
669+ where
670+ delayAction = delaySeconds =<< getPollingInterval
671+ noDelayAction = pure ()
672+
673+-- | Executes 'killJobPollingSql' every 'cfgPollingInterval' seconds to pick up jobs
674+-- that are cancelled and need to be killed. Uses @UPDATE@ along with @SELECT...
675+-- ..FOR UPDATE@ to efficiently find a job that matches /all/ of the following
676+-- conditions:
677+--
678+-- * 'jobStatus' should be 'cancelled'
679+-- * 'jobLockedAt' should be in the past
680+-- * 'jobLockedBy' should be the current job worker name
681+
682+killJobPoller :: (HasJobRunner m) => m ()
683+killJobPoller = do
684+ processName <- liftIO jobWorkerName
685+ pool <- getDbPool
686+ tname <- getTableName
687+
688+ let pollJobToKill conn = join $ mask_ $ do
689+ currentTime <- liftIO getCurrentTime
690+ result <- liftIO $ PGS.query conn killJobPollingSql
691+ (tname, tname, Cancelled, processName, currentTime)
692+
693+ case result of
694+ [] ->
695+ pure delayAction
696+
697+ [Only (jobId :: JobId)] -> do
698+ void $ async $ killJob jobId
552699 pure noDelayAction
553700
554− x -> error $ "WTF just happened? I was supposed to get only a single row, but got: " ++ (show x)
555− nextAction
701+ x ->
702+ error $ "I was supposed to get only a single row, but got: " ++ show x
703+
704+ withResource pool (forever . pollJobToKill)
705+
556706 where
557707 delayAction = delaySeconds =<< getPollingInterval
558708 noDelayAction = pure ()
@@ -568,43 +718,56 @@ jobEventListener = do
568718 jwName <- liftIO jobWorkerName
569719 concurrencyControlFn <- getConcurrencyControlFn
570720
571− let tryLockingJob jid = do
721+ let tryLockingJob jid mResCfg = withDbConnection $ \conn -> do
572722 let q = "UPDATE ? SET status=?, locked_at=now(), locked_by=?, attempts=attempts+1 WHERE id=? AND status in ? RETURNING id"
573− (withDbConnection $ \conn -> (liftIO $ PGS.query conn q (tname, Locked, jwName, jid, In [Queued, Retry]))) >>= \case
723+ qWithResources =
724+ "UPDATE ? SET status=?, locked_at=now(), locked_by=?, attempts=attempts+1 \
725+ \ WHERE id=? AND status in ? AND ?(id) RETURNING id"
726+ result <- case mResCfg of
727+ Nothing -> liftIO $ PGS.query conn q (tname, Locked, jwName, jid, In [Queued, Retry])
728+ Just ResourceCfg{..} -> liftIO $ PGS.query conn qWithResources
729+ (tname, Locked, jwName, jid, In [Queued, Retry], resCfgCheckResourceFunction)
730+
731+ case result of
574732 [] -> do
575733 log LevelDebug $ LogText $ toS $ "Job was locked by someone else before I could start. Skipping it. JobId=" <> show jid
576734 pure Nothing
577735 [Only (_ :: JobId)] -> pure $ Just jid
578− x -> error $ "WTF just happned? Was expecting a single row to be returned, received " ++ (show x)
736+ x -> error $ "WTF just happned? Was expecting a single row to be returned, received " ++ show x
579737
580738 withResource pool $ \monitorDbConn -> do
581− void $ liftIO $ PGS.execute monitorDbConn ("LISTEN ?") (Only $ pgEventName tname)
739+ void $ liftIO $ PGS.execute monitorDbConn "LISTEN ?" (Only $ pgEventName tname)
582740 forever $ do
583741 log LevelDebug $ LogText "[LISTEN/NOTIFY] Event loop"
584742 notif <- liftIO $ getNotification monitorDbConn
585− concurrencyControlFn >>= \case
586− False -> log LevelWarn $ LogText "Received job event, but ignoring it due to concurrency control"
587− True -> do
588− let pload = notificationData notif
589− log LevelDebug $ LogText $ toS $ "NOTIFY | " <> show pload
590− case (eitherDecode $ toS pload) of
591− Left e -> log LevelError $ LogText $ toS $ "Unable to decode notification payload received from Postgres. Payload=" <> show pload <> " Error=" <> show e
592−
593− -- Checking if job needs to be fired immediately AND it is not already
594− -- taken by some othe thread, by the time it got to us
595− Right (v :: Value) -> case (Aeson.parseMaybe parser v) of
596− Nothing -> log LevelError $ LogText $ toS $ "Unable to extract id/run_at/locked_at from " <> show pload
597− Just (jid, runAt_, mLockedAt_) -> do
598− t <- liftIO getCurrentTime
599− if (runAt_ <= t) && (isNothing mLockedAt_)
600− then do log LevelDebug $ LogText $ toS $ "Job needs needs to be run immediately. Attempting to fork in background. JobId=" <> show jid
601− void $ async $ do
602− -- Let's try to lock the job first... it is possible that it has already
603− -- been picked up by the poller by the time we get here.
604− tryLockingJob jid >>= \case
605− Nothing -> pure ()
606− Just lockedJid -> runJob lockedJid
607− else log LevelDebug $ LogText $ toS $ "Job is either for future, or is already locked. Skipping. JobId=" <> show jid
743+
744+ let pload = notificationData notif
745+ runNotifWithFilter :: HasJobRunner m => Maybe ResourceCfg -> m ()
746+ runNotifWithFilter mResCfg = do
747+ log LevelDebug $ LogText $ toS $ "NOTIFY | " <> show pload
748+ case eitherDecode $ toS pload of
749+ Left e -> log LevelError $ LogText $ toS $ "Unable to decode notification payload received from Postgres. Payload=" <> show pload <> " Error=" <> show e
750+
751+ -- Checking if job needs to be fired immediately AND it is not already
752+ -- taken by some othe thread, by the time it got to us
753+ Right (v :: Value) -> case Aeson.parseMaybe parser v of
754+ Nothing -> log LevelError $ LogText $ toS $ "Unable to extract id/run_at/locked_at from " <> show pload
755+ Just (jid, runAt_, mLockedAt_) -> do
756+ t <- liftIO getCurrentTime
757+ if (runAt_ <= t) && isNothing mLockedAt_
758+ then do log LevelDebug $ LogText $ toS $ "Job needs needs to be run immediately. Attempting to fork in background. JobId=" <> show jid
759+ void $ async $ do
760+ -- Let's try to lock the job first... it is possible that it has already
761+ -- been picked up by the poller by the time we get here.
762+ tryLockingJob jid mResCfg >>= \case
763+ Nothing -> pure ()
764+ Just lockedJid -> runJob lockedJid
765+ else log LevelDebug $ LogText $ toS $ "Job is either for future, is already locked, or would violate concurrency constraints. Skipping. JobId=" <> show jid
766+
767+ concurrencyControlFn monitorDbConn >>= \case
768+ DontPoll -> log LevelWarn $ LogText "Received job event, but ignoring it due to concurrency control"
769+ PollAny -> runNotifWithFilter Nothing
770+ PollWithResources resCfg -> runNotifWithFilter (Just resCfg)
608771 where
609772 parser :: Value -> Aeson.Parser (JobId, UTCTime, Maybe UTCTime)
610773 parser = withObject "expecting an object to parse job.run_at and job.locked_at" $ \o -> do
@@ -618,6 +781,12 @@ jobEventListener = do
618781createJobQuery :: PGS.Query
619782createJobQuery = "INSERT INTO ? (run_at, status, payload, last_error, attempts, locked_at, locked_by) VALUES (?, ?, ?, ?, ?, ?, ?) RETURNING " <> concatJobDbColumns
620783
784+ensureResource :: PGS.Query
785+ensureResource = "INSERT INTO ? (id, usage_limit) VALUES (?, ?) ON CONFLICT DO NOTHING"
786+
787+registerResourceUsage :: PGS.Query
788+registerResourceUsage = "INSERT INTO ? (job_id, resource_id, usage) VALUES (?, ?, ?)"
789+
621790-- $createJobs
622791--
623792-- Ideally you'd want to create wrappers for 'createJob' and 'scheduleJob' in
@@ -656,13 +825,56 @@ scheduleJob :: ToJSON p
656825 -> IO Job
657826scheduleJob conn tname payload runAt = do
658827 let args = ( tname, runAt, Queued, toJSON payload, Nothing :: Maybe Value, 0 :: Int, Nothing :: Maybe Text, Nothing :: Maybe Text )
659− queryFormatter = toS <$> (PGS.formatQuery conn createJobQuery args)
828+ queryFormatter = toS <$> PGS.formatQuery conn createJobQuery args
660829 rs <- PGS.query conn createJobQuery args
661830 case rs of
662− [] -> (Prelude.error . (<> "Not expecting a blank result set when creating a job. Query=")) <$> queryFormatter
831+ [] -> Prelude.error . (<> "Not expecting a blank result set when creating a job. Query=") <$> queryFormatter
832+ [r] -> pure r
833+ _ -> Prelude.error . (<> "Not expecting multiple rows when creating a single job. Query=") <$> queryFormatter
834+
835+type ResourceList = [(ResourceId, Int)]
836+
837+createJobWithResources
838+ :: ToJSON p
839+ => Connection
840+ -> TableName
841+ -> ResourceCfg
842+ -> p
843+ -> ResourceList
844+ -> IO Job
845+createJobWithResources conn tname resCfg payload resources = do
846+ t <- getCurrentTime
847+ scheduleJobWithResources conn tname resCfg payload resources t
848+
849+scheduleJobWithResources
850+ :: ToJSON p
851+ => Connection
852+ -> TableName
853+ -> ResourceCfg
854+ -> p
855+ -> ResourceList
856+ -> UTCTime
857+ -> IO Job
858+scheduleJobWithResources conn tname ResourceCfg{..} payload resources runAt = do
859+ -- We insert everything in a single transaction to delay @NOTIFY@ calls,
860+ -- so a job isn't picked up before its resources are inserted.
861+ PGS.begin conn
862+ let args = ( tname, runAt, Queued, toJSON payload, Nothing :: Maybe Value, 0 :: Int, Nothing :: Maybe Text, Nothing :: Maybe Text )
863+ queryFormatter = toS <$> PGS.formatQuery conn createJobQuery args
864+ rs <- PGS.query conn createJobQuery args
865+
866+ job <- flip onException (PGS.rollback conn) $ case rs of
867+ [] -> Prelude.error . (<> "Not expecting a blank result set when creating a job. Query=") <$> queryFormatter
663868 [r] -> pure r
664− _ -> (Prelude.error . (<> "Not expecting multiple rows when creating a single job. Query=")) <$> queryFormatter
869+ _ -> Prelude.error . (<> "Not expecting multiple rows when creating a single job. Query=") <$> queryFormatter
870+
871+ forM_ resources $ \(resourceId, usage) -> do
872+ void $ PGS.execute conn ensureResource (resCfgResourceTable, rawResourceId resourceId, resCfgDefaultLimit)
873+ void $ PGS.execute conn registerResourceUsage (resCfgUsageTable, jobId job, rawResourceId resourceId, usage)
665874
875+ PGS.commit conn
876+
877+ pure job
666878
667879-- getRunnerEnv :: (HasJobRunner m) => m RunnerEnv
668880-- getRunnerEnv = ask
@@ -670,14 +882,14 @@ scheduleJob conn tname payload runAt = do
670882eitherParsePayload :: (FromJSON a)
671883 => Job
672884 -> Either String a
673−eitherParsePayload job =
674− eitherParsePayloadWith parseJSON job
885+eitherParsePayload =
886+ eitherParsePayloadWith parseJSON
675887
676888throwParsePayload :: (FromJSON a)
677889 => Job
678890 -> IO a
679−throwParsePayload job =
680− throwParsePayloadWith parseJSON job
891+throwParsePayload =
892+ throwParsePayloadWith parseJSON
681893
682894eitherParsePayloadWith :: (Aeson.Value -> Aeson.Parser a)
683895 -> Job
@@ -692,7 +904,7 @@ throwParsePayloadWith :: (Aeson.Value -> Aeson.Parser a)
692904 -> Job
693905 -> IO a
694906throwParsePayloadWith parser job =
695− either throwString (pure . Prelude.id) (eitherParsePayloadWith parser job)
907+ either throwString pure (eitherParsePayloadWith parser job)
696908
697909
698910-- | Used by the web\/admin UI to fetch a \"master list\" of all known
@@ -716,5 +928,4 @@ fetchAllJobRunners :: (MonadIO m)
716928 => UIConfig
717929 -> m [JobRunnerName]
718930fetchAllJobRunners UIConfig{uicfgTableName, uicfgDbPool} = liftIO $ withResource uicfgDbPool $ \conn -> do
719− fmap (mapMaybe fromOnly) $ PGS.query conn "select distinct locked_by from ?" (Only uicfgTableName)
720−
931+ mapMaybe fromOnly <$> PGS.query conn "select distinct locked_by from ?" (Only uicfgTableName)