@@ -1,4 +1,5 @@
11{-# LANGUAGE RankNTypes, FlexibleInstances, FlexibleContexts, PartialTypeSignatures, UndecidableInstances #-}
2+{-# LANGUAGE ExistentialQuantification, RecordWildCards, ScopedTypeVariables #-}
23
34module OddJobs.Job
45 (
@@ -12,6 +13,7 @@ module OddJobs.Job
1213 -- $config
1314 , Config(..)
1415 , ConcurrencyControl(..)
16+ , ResourceCfg(..)
1517
1618 -- * Creating/scheduling jobs
1719 --
@@ -19,6 +21,9 @@ module OddJobs.Job
1921 , createJob
2022 , scheduleJob
2123
24+ , createJobWithResources
25+ , scheduleJobWithResources
26+
2227 -- * @Job@ and associated data-types
2328 --
2429 -- $dataTypes
@@ -31,6 +36,9 @@ module OddJobs.Job
3136 , Seconds(..)
3237 , JobErrHandler(..)
3338 , AllJobTypes(..)
39+ , ResourceId(..)
40+ , FunctionName
41+ , RunnerEnv(..)
3442
3543 -- ** Structured logging
3644 --
@@ -45,6 +53,7 @@ module OddJobs.Job
4553 , jobEventListener
4654 , jobPoller
4755 , jobPollingSql
56+ , pollRunJob
4857 , JobRunner
4958 , HasJobRunner (..)
5059
@@ -56,6 +65,7 @@ module OddJobs.Job
5665 , runJobNowIO
5766 , unlockJobIO
5867 , cancelJobIO
68+ , killJobIO
5969 , jobDbColumns
6070 , jobPollingIO
6171 , concatJobDbColumns
@@ -78,7 +88,7 @@ import Data.Pool(Pool)
7888import Data.Text as T
7989import Database.PostgreSQL.Simple as PGS
8090import Database.PostgreSQL.Simple.Notification
81−import UnliftIO.Async
91+import UnliftIO.Async hiding (poll)
8292import UnliftIO.Concurrent (threadDelay, myThreadId)
8393import Data.String
8494import System.Posix.Process (getProcessID)
@@ -89,7 +99,8 @@ import Control.Monad.Logger as MLogger (LogLevel(..), LogStr, toLogStr)
8999import UnliftIO.IORef
90100import UnliftIO.Exception ( SomeException(..), try, catch, finally
91101 , catchAny, bracket, Exception(..), throwIO
92− , catches, Handler(..), mask_, throwString
102+ , catches, Handler(..), mask_, onException
103+ , throwString
93104 )
94105import Data.Proxy
95106import Control.Monad.Trans.Control
@@ -100,14 +111,16 @@ import Data.Aeson hiding (Success)
100111import qualified Data.Aeson as Aeson
101112import qualified Data.Aeson.Types as Aeson (Parser, parseMaybe)
102113import Data.String.Conv (StringConv(..), toS)
103−import Data.Functor (void)
104−import Control.Monad (forever)
114+import Data.Functor ((<&>), void)
115+import Control.Monad (forever, forM_, join)
105116import Data.Maybe (isNothing, maybe, fromMaybe, listToMaybe, mapMaybe)
106117import Data.Either (either)
107118import Control.Monad.Reader
108119import GHC.Generics
120+import Data.Map (Map)
109121import qualified Data.HashMap.Strict as HM
110122import qualified Data.List as DL
123+import qualified Data.Map as DM
111124import qualified Data.ByteString as BS
112125import qualified Data.ByteString.Lazy as BSL
113126import System.FilePath (FilePath)
@@ -141,6 +154,7 @@ class (MonadUnliftIO m, MonadBaseControl IO m) => HasJobRunner m where
141154 getConcurrencyControl :: m ConcurrencyControl
142155 log :: LogLevel -> LogEvent -> m ()
143156 getDefaultJobTimeout :: m Seconds
157+ onJobTimeout :: Job -> m ()
144158 getDefaultRetryBackoff :: Int -> m Seconds
145159
146160-- $logging
@@ -158,7 +172,7 @@ class (MonadUnliftIO m, MonadBaseControl IO m) => HasJobRunner m where
158172
159173data RunnerEnv = RunnerEnv
160174 { envConfig :: !Config
161− , envJobThreadsRef :: !(IORef [Async ()])
175+ , envJobThreadsRef :: !(IORef (Map JobId (Async ())))
162176 }
163177
164178type RunnerM = ReaderT RunnerEnv IO
@@ -192,6 +206,9 @@ instance HasJobRunner RunnerM where
192206 liftIO $ loggerFn logLevel logEvent
193207
194208 getDefaultJobTimeout = asks (cfgDefaultJobTimeout . envConfig)
209+ onJobTimeout job = do
210+ fn <- asks (cfgOnJobTimeout . envConfig)
211+ logCallbackErrors (jobId job) "onJobTimeout" $ liftIO $ fn job
195212
196213 getDefaultRetryBackoff attempts = do
197214 retryFn <- asks (cfgDefaultRetryBackoff . envConfig)
@@ -204,7 +221,7 @@ instance HasJobRunner RunnerM where
204221startJobRunner :: Config -> IO ()
205222startJobRunner jm = do
206223 traceM "startJobRunner top"
207− r <- newIORef []
224+ r <- newIORef DM.empty
208225 let monitorEnv = RunnerEnv
209226 { envConfig = jm
210227 , envJobThreadsRef = r
@@ -305,7 +322,7 @@ saveJob j = do
305322saveJobIO :: Connection -> TableName -> Job -> IO Job
306323saveJobIO conn tname Job{jobRunAt, jobStatus, jobPayload, jobLastError, jobAttempts, jobLockedBy, jobLockedAt, jobId} = do
307324 rs <- PGS.query conn saveJobQuery
308− ( tname
325+ ( tname
309326 , jobRunAt
310327 , jobStatus
311328 , jobPayload
@@ -332,7 +349,7 @@ deleteJobIO conn tname jid = do
332349runJobNowIO :: Connection -> TableName -> JobId -> IO (Maybe Job)
333350runJobNowIO conn tname jid = do
334351 t <- getCurrentTime
335− updateJobHelper tname conn (Queued, [Queued, Retry, Failed], Just t, jid)
352+ updateJobHelper tname conn (Queued, [Queued, Retry, Failed, Cancelled], Just t, jid)
336353
337354-- | TODO: First check in all job-runners if this job is still running, or not,
338355-- and somehow send an uninterruptibleCancel to that thread.
@@ -344,7 +361,11 @@ unlockJobIO conn tname jid = do
344361
345362cancelJobIO :: Connection -> TableName -> JobId -> IO (Maybe Job)
346363cancelJobIO conn tname jid =
347− 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)
348369
349370updateJobHelper :: TableName
350371 -> Connection
@@ -362,30 +383,37 @@ updateJobHelper tname conn (newStatus, existingStates, mRunAt, jid) =
362383data TimeoutException = TimeoutException deriving (Eq, Show)
363384instance Exception TimeoutException
364385
386+-- | this spawns the job on a a seperate thread, and
387+-- also spawns a monitoring thread, whichever finishes
388+-- first kills the other.
365389runJobWithTimeout :: (HasJobRunner m)
366390 => Seconds
367391 -> Job
368392 -> m ()
369−runJobWithTimeout timeoutSec job = do
393+runJobWithTimeout timeoutSec job@Job{jobId} = do
370394 threadsRef <- envJobThreadsRef <$> getRunnerEnv
371395 jobRunner_ <- getJobRunner
372396
373397 a <- async $ liftIO $ jobRunner_ job
374398
375− _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+
376404 -- liftIO $ putStrLn $ "Threads: " <> show x
377405 log LevelDebug $ LogText $ toS $ "Spawned job in " <> show (asyncThreadId a)
378406
379407 t <- async $ do
380408 delaySeconds timeoutSec
381− uninterruptibleCancel a
382409 throwIO TimeoutException
383410
384411 void $ finally
385412 (waitEitherCancel a t)
386− (atomicModifyIORef' threadsRef $ \threads -> (DL.delete a threads, ()))
413+ (atomicModifyIORef' threadsRef $ \threads -> (DM.delete jobId threads, ()))
387414
388415
416+-- | runs a job, blocks for as long as it's in progress
389417runJob :: (HasJobRunner m) => JobId -> m ()
390418runJob jid = do
391419 findJobById jid >>= \case
@@ -394,7 +422,8 @@ runJob jid = do
394422 startTime <- liftIO getCurrentTime
395423 lockTimeout <- getDefaultJobTimeout
396424 log LevelInfo $ LogJobStart job
397− flip catches [Handler $ timeoutHandler job startTime, Handler $ exceptionHandler job startTime] $ do
425+ flip catch (exceptionHandler job startTime) $ do
426+ onJobStart job
398427 runJobWithTimeout lockTimeout job
399428 endTime <- liftIO getCurrentTime
400429 shouldDeleteJob <- deleteSuccessfulJobs
@@ -406,7 +435,6 @@ runJob jid = do
406435 onJobSuccess newJob
407436 pure ()
408437 where
409− timeoutHandler job startTime (e :: TimeoutException) = retryOrFail (toException e) job startTime
410438 exceptionHandler job startTime (e :: SomeException) = retryOrFail (toException e) job startTime
411439 retryOrFail e job@Job{jobAttempts} startTime = do
412440 endTime <- liftIO getCurrentTime
@@ -424,16 +452,39 @@ runJob jid = do
424452 , jobRunAt=addUTCTime (fromIntegral $ unSeconds backoffInSeconds) t
425453 }
426454 case fromException e :: Maybe TimeoutException of
427− Nothing -> log logLevel $ LogJobFailed newJob e failureMode runTime
428− Just _ -> log logLevel $ LogJobTimeout newJob
429−
430− let tryHandler (JobErrHandler handler) res = case fromException e of
431− Nothing -> res
432− Just e_ -> void $ handler e_ newJob failureMode
433− handlers <- onJobFailed
434− 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+
435467 pure ()
436468
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+
437488-- TODO: This might have a resource leak.
438489restartUponCrash :: (HasJobRunner m, Show a) => Text -> m a -> m ()
439490restartUponCrash name_ action = do
@@ -459,9 +510,11 @@ jobMonitor = do
459510 traceM "jobMonitor top"
460511 a1 <- async $ restartUponCrash "Job poller" jobPoller
461512 a2 <- async $ restartUponCrash "Job event listener" jobEventListener
513+ a3 <- async $ restartUponCrash "Job Kill poller" killJobPoller
462514 traceM "jobMonitor after async"
463− finally (void $ waitAnyCatch [a1, a2]) $ do
515+ finally (void $ waitAnyCatch [a1, a2, a3]) $ do
464516 log LevelInfo (LogText "Stopping jobPoller and jobEventListener threads.")
517+ cancel a3
465518 cancel a2
466519 cancel a1
467520 log LevelInfo (LogText "Waiting for jobs to complete.")
@@ -469,13 +522,32 @@ jobMonitor = do
469522
470523-- | Ref: 'jobPoller'
471524jobPollingSql :: Query
472−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"
473545
474546waitForJobs :: (HasJobRunner m)
475547 => m ()
476548waitForJobs = do
477− threadsRef <- envJobThreadsRef <$> getRunnerEnv
478− readIORef threadsRef >>= \case
549+ curJobs <- getRunnerEnv >>= (readIORef . envJobThreadsRef) <&> DM.elems
550+ case curJobs of
479551 [] -> log LevelInfo $ LogText "All job-threads exited"
480552 as -> do
481553 tid <- myThreadId
@@ -484,15 +556,23 @@ waitForJobs = do
484556 delaySeconds (Seconds 1)
485557 waitForJobs
486558
559+data ConcurrencyAction
560+ = DontPoll
561+ | PollAny
562+ | PollWithResources ResourceCfg
563+
487564getConcurrencyControlFn :: (HasJobRunner m)
488− => m (m Bool)
565+ => m (Connection -> m ConcurrencyAction)
489566getConcurrencyControlFn = getConcurrencyControl >>= \case
490− UnlimitedConcurrentJobs -> pure $ pure True
491− MaxConcurrentJobs maxJobs -> pure $ do
492− curJobs <- getRunnerEnv >>= (readIORef . envJobThreadsRef)
493− pure $ DL.length curJobs < maxJobs
494− 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
495573
574+ where
575+ pollIf cond = if cond then PollAny else DontPoll
496576
497577jobPollingIO :: Connection -> String -> TableName -> Seconds -> IO [Only JobId]
498578jobPollingIO pollerDbConn processName tname lockTimeout = do
@@ -525,32 +605,104 @@ jobPollingIO pollerDbConn processName tname lockTimeout = do
525605jobPoller :: (HasJobRunner m) => m ()
526606jobPoller = do
527607 processName <- liftIO jobWorkerName
528− pool <- getDbPool
529− tname <- getTableName
530− lockTimeout <- getDefaultJobTimeout
531− traceM "jobPoller just before log statement"
532608 log LevelInfo $ LogText $ toS $ "Starting the job monitor via DB polling with processName=" <> processName
533609 concurrencyControlFn <- getConcurrencyControlFn
534− withResource pool $ \pollerDbConn -> forever $ concurrencyControlFn >>= \case
535− False -> do
536− log LevelWarn $ LogText "NOT polling the job queue due to concurrency control"
537− -- If we can't run any jobs ATM, relax and wait for resources to free up
538− delayAction
539− True -> do
540− join
541− (mask_ $ do
542− log LevelDebug $ LogText $ toS $ "[" <> processName <> "] Polling the job queue.."
543− r <- liftIO $ jobPollingIO pollerDbConn processName tname lockTimeout
544− case r of
545− -- When we don't have any jobs to run, we can relax a bit...
546− [] -> pure delayAction
547−
548− -- When we find a job to run, fork and try to find the next job without any delay...
549− [Only (jid :: JobId)] -> do
550− void $ async $ runJob jid
551− pure noDelayAction
552−
553− x -> error $ "WTF just happened? I was supposed to get only a single row, but got: " ++ show x)
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
699+ pure noDelayAction
700+
701+ x ->
702+ error $ "I was supposed to get only a single row, but got: " ++ show x
703+
704+ withResource pool (forever . pollJobToKill)
705+
554706 where
555707 delayAction = delaySeconds =<< getPollingInterval
556708 noDelayAction = pure ()
@@ -566,9 +718,17 @@ jobEventListener = do
566718 jwName <- liftIO jobWorkerName
567719 concurrencyControlFn <- getConcurrencyControlFn
568720
569− let tryLockingJob jid = do
721+ let tryLockingJob jid mResCfg = withDbConnection $ \conn -> do
570722 let q = "UPDATE ? SET status=?, locked_at=now(), locked_by=?, attempts=attempts+1 WHERE id=? AND status in ? RETURNING id"
571− 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
572732 [] -> do
573733 log LevelDebug $ LogText $ toS $ "Job was locked by someone else before I could start. Skipping it. JobId=" <> show jid
574734 pure Nothing
@@ -580,29 +740,34 @@ jobEventListener = do
580740 forever $ do
581741 log LevelDebug $ LogText "[LISTEN/NOTIFY] Event loop"
582742 notif <- liftIO $ getNotification monitorDbConn
583− concurrencyControlFn >>= \case
584− False -> log LevelWarn $ LogText "Received job event, but ignoring it due to concurrency control"
585− True -> do
586− let pload = notificationData notif
587− log LevelDebug $ LogText $ toS $ "NOTIFY | " <> show pload
588− case eitherDecode $ toS pload of
589− Left e -> log LevelError $ LogText $ toS $ "Unable to decode notification payload received from Postgres. Payload=" <> show pload <> " Error=" <> show e
590−
591− -- Checking if job needs to be fired immediately AND it is not already
592− -- taken by some othe thread, by the time it got to us
593− Right (v :: Value) -> case Aeson.parseMaybe parser v of
594− Nothing -> log LevelError $ LogText $ toS $ "Unable to extract id/run_at/locked_at from " <> show pload
595− Just (jid, runAt_, mLockedAt_) -> do
596− t <- liftIO getCurrentTime
597− if (runAt_ <= t) && isNothing mLockedAt_
598− then do log LevelDebug $ LogText $ toS $ "Job needs needs to be run immediately. Attempting to fork in background. JobId=" <> show jid
599− void $ async $ do
600− -- Let's try to lock the job first... it is possible that it has already
601− -- been picked up by the poller by the time we get here.
602− tryLockingJob jid >>= \case
603− Nothing -> pure ()
604− Just lockedJid -> runJob lockedJid
605− 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)
606771 where
607772 parser :: Value -> Aeson.Parser (JobId, UTCTime, Maybe UTCTime)
608773 parser = withObject "expecting an object to parse job.run_at and job.locked_at" $ \o -> do
@@ -616,6 +781,12 @@ jobEventListener = do
616781createJobQuery :: PGS.Query
617782createJobQuery = "INSERT INTO ? (run_at, status, payload, last_error, attempts, locked_at, locked_by) VALUES (?, ?, ?, ?, ?, ?, ?) RETURNING " <> concatJobDbColumns
618783
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+
619790-- $createJobs
620791--
621792-- Ideally you'd want to create wrappers for 'createJob' and 'scheduleJob' in
@@ -659,8 +830,51 @@ scheduleJob conn tname payload runAt = do
659830 case rs of
660831 [] -> Prelude.error . (<> "Not expecting a blank result set when creating a job. Query=") <$> queryFormatter
661832 [r] -> pure r
662− _ -> Prelude.error . (<> "Not expecting multiple rows when creating a single job. Query=") <$> queryFormatter
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
663865
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
868+ [r] -> pure r
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)
874+
875+ PGS.commit conn
876+
877+ pure job
664878
665879-- getRunnerEnv :: (HasJobRunner m) => m RunnerEnv
666880-- getRunnerEnv = ask
@@ -715,4 +929,3 @@ fetchAllJobRunners :: (MonadIO m)
715929 -> m [JobRunnerName]
716930fetchAllJobRunners UIConfig{uicfgTableName, uicfgDbPool} = liftIO $ withResource uicfgDbPool $ \conn -> do
717931 mapMaybe fromOnly <$> PGS.query conn "select distinct locked_by from ?" (Only uicfgTableName)
718−