cc05b64 · odd-jobs · git.iko.soy

git.iko.soy/odd-jobs

on master help sign in

Merge pull request #80 from SupercedeTech/resource-tables-pr

cc05b64Isaac van Bakel4 years ago

src/OddJobs/Job.hs+169−51

compare with master

unifiedsplit

@@ -1,4 +1,5 @@

1{-# LANGUAGE RankNTypes, FlexibleInstances, FlexibleContexts, PartialTypeSignatures, UndecidableInstances #-}

2+{-# LANGUAGE ExistentialQuantification, RecordWildCards, ScopedTypeVariables #-}

3module OddJobs.Job

4module OddJobs.Job

4 (

5 (

@@ -12,6 +13,7 @@ module OddJobs.Job

12 -- $config

13 -- $config

13 , Config(..)

14 , Config(..)

14 , ConcurrencyControl(..)

15 , ConcurrencyControl(..)

16+ , ResourceCfg(..)

16 -- * Creating/scheduling jobs

18 -- * Creating/scheduling jobs

17 --

19 --

@@ -19,6 +21,9 @@ module OddJobs.Job

19 , createJob

21 , createJob

20 , scheduleJob

22 , scheduleJob

24+ , createJobWithResources

25+ , scheduleJobWithResources

26+

22 -- * @Job@ and associated data-types

27 -- * @Job@ and associated data-types

23 --

28 --

24 -- $dataTypes

29 -- $dataTypes

@@ -31,6 +36,8 @@ module OddJobs.Job

31 , Seconds(..)

36 , Seconds(..)

32 , JobErrHandler(..)

37 , JobErrHandler(..)

33 , AllJobTypes(..)

38 , AllJobTypes(..)

39+ , ResourceId(..)

40+ , FunctionName

35 -- ** Structured logging

42 -- ** Structured logging

36 --

43 --

@@ -78,7 +85,7 @@ import Data.Pool(Pool)

78import Data.Text as T

85import Data.Text as T

79import Database.PostgreSQL.Simple as PGS

86import Database.PostgreSQL.Simple as PGS

80import Database.PostgreSQL.Simple.Notification

87import Database.PostgreSQL.Simple.Notification

81−import UnliftIO.Async

88+import UnliftIO.Async hiding (poll)

82import UnliftIO.Concurrent (threadDelay, myThreadId)

89import UnliftIO.Concurrent (threadDelay, myThreadId)

83import Data.String

90import Data.String

84import System.Posix.Process (getProcessID)

91import System.Posix.Process (getProcessID)

@@ -89,7 +96,8 @@ import Control.Monad.Logger as MLogger (LogLevel(..), LogStr, toLogStr)

89import UnliftIO.IORef

96import UnliftIO.IORef

90import UnliftIO.Exception ( SomeException(..), try, catch, finally

97import UnliftIO.Exception ( SomeException(..), try, catch, finally

91 , catchAny, bracket, Exception(..), throwIO

98 , catchAny, bracket, Exception(..), throwIO

92− , catches, Handler(..), mask_, throwString

99+ , catches, Handler(..), mask_, onException

100+ , throwString

93 )

101 )

94import Data.Proxy

102import Data.Proxy

95import Control.Monad.Trans.Control

103import Control.Monad.Trans.Control

@@ -309,7 +317,7 @@ saveJob j = do

309saveJobIO :: Connection -> TableName -> Job -> IO Job

317saveJobIO :: Connection -> TableName -> Job -> IO Job

310saveJobIO conn tname Job{jobRunAt, jobStatus, jobPayload, jobLastError, jobAttempts, jobLockedBy, jobLockedAt, jobId} = do

318saveJobIO conn tname Job{jobRunAt, jobStatus, jobPayload, jobLastError, jobAttempts, jobLockedBy, jobLockedAt, jobId} = do

311 rs <- PGS.query conn saveJobQuery

319 rs <- PGS.query conn saveJobQuery

312− ( tname

320+ ( tname

313 , jobRunAt

321 , jobRunAt

314 , jobStatus

322 , jobStatus

315 , jobPayload

323 , jobPayload

@@ -476,7 +484,18 @@ jobMonitor = do

476

484

477-- | Ref: 'jobPoller'

485-- | Ref: 'jobPoller'

478jobPollingSql :: Query

486jobPollingSql :: Query

487+jobPollingSql =

488+ "update ? set status = ?, locked_at = ?, locked_by = ?, attempts=attempts+1\

479−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"

489+ \ WHERE id in (select id from ? where (run_at<=? AND ((status in ?) OR (status = ? and locked_at<?))) \

490+ \ ORDER BY attempts ASC, run_at ASC LIMIT 1 FOR UPDATE) RETURNING id"

491+

492+jobPollingWithResourceSql :: Query

493+jobPollingWithResourceSql =

494+ " UPDATE ? SET status = ?, locked_at = ?, locked_by = ?, attempts = attempts + 1 \

495+ \ WHERE id in (select id from ? where (run_at<=? AND ((status in ?) OR (status = ? and locked_at<?))) \

496+ \ AND ?(id) \

497+ \ ORDER BY attempts ASC, run_at ASC LIMIT 1) \

498+ \ RETURNING id"

480

499

481waitForJobs :: (HasJobRunner m)

500waitForJobs :: (HasJobRunner m)

482 => m ()

501 => m ()

@@ -491,15 +510,23 @@ waitForJobs = do

491 delaySeconds (Seconds 1)

510 delaySeconds (Seconds 1)

492 waitForJobs

511 waitForJobs

493

512

513+data ConcurrencyAction

514+ = DontPoll

515+ | PollAny

516+ | PollWithResources ResourceCfg

517+

494getConcurrencyControlFn :: (HasJobRunner m)

518getConcurrencyControlFn :: (HasJobRunner m)

495− => m (m Bool)

519+ => m (Connection -> m ConcurrencyAction)

496getConcurrencyControlFn = getConcurrencyControl >>= \case

520getConcurrencyControlFn = getConcurrencyControl >>= \case

497− UnlimitedConcurrentJobs -> pure $ pure True

521+ UnlimitedConcurrentJobs -> pure $ const $ pure PollAny

498− MaxConcurrentJobs maxJobs -> pure $ do

522+ MaxConcurrentJobs maxJobs -> pure $ const $ do

499 curJobs <- getRunnerEnv >>= (readIORef . envJobThreadsRef)

523 curJobs <- getRunnerEnv >>= (readIORef . envJobThreadsRef)

500− pure $ DL.length curJobs < maxJobs

524+ pure $ pollIf $ DL.length curJobs < maxJobs

525+ ResourceLimits resCfg -> pure $ const $ pure $ PollWithResources resCfg

501− DynamicConcurrency fn -> pure $ liftIO fn

526+ DynamicConcurrency fn -> pure $ const $ pollIf <$> liftIO fn

502

527

528+ where

529+ pollIf cond = if cond then PollAny else DontPoll

503

530

504jobPollingIO :: Connection -> String -> TableName -> Seconds -> IO [Only JobId]

531jobPollingIO :: Connection -> String -> TableName -> Seconds -> IO [Only JobId]

505jobPollingIO pollerDbConn processName tname lockTimeout = do

532jobPollingIO pollerDbConn processName tname lockTimeout = do

@@ -537,27 +564,55 @@ jobPoller = do

537 lockTimeout <- getDefaultJobTimeout

564 lockTimeout <- getDefaultJobTimeout

538 traceM "jobPoller just before log statement"

565 traceM "jobPoller just before log statement"

539 log LevelInfo $ LogText $ toS $ "Starting the job monitor via DB polling with processName=" <> processName

566 log LevelInfo $ LogText $ toS $ "Starting the job monitor via DB polling with processName=" <> processName

567+

568+ let poll conn mResCfg = join $ mask_ $ do

569+ log LevelDebug $ LogText $ toS $ "[" <> processName <> "] Polling the job queue.."

570+ t <- liftIO getCurrentTime

571+ r <- case mResCfg of

572+ Nothing -> liftIO $

573+ PGS.query conn jobPollingSql

574+ ( tname

575+ , Locked

576+ , t

577+ , processName

578+ , tname

579+ , t

580+ , In [Queued, Retry]

581+ , Locked

582+ , addUTCTime (fromIntegral $ negate $ unSeconds lockTimeout) t)

583+ Just ResourceCfg{..} -> liftIO $

584+ PGS.query conn jobPollingWithResourceSql

585+ ( tname

586+ , Locked

587+ , t

588+ , processName

589+ , tname

590+ , t

591+ , In [Queued, Retry]

592+ , Locked

593+ , addUTCTime (fromIntegral $ negate $ unSeconds lockTimeout) t

594+ , resCfgCheckResourceFunction

595+ )

596+ case r of

597+ -- When we don't have any jobs to run, we can relax a bit...

598+ [] -> pure delayAction

599+

600+ -- When we find a job to run, fork and try to find the next job without any delay...

601+ [Only (jid :: JobId)] -> do

602+ void $ async $ runJob jid

603+ pure noDelayAction

604+

605+ x -> error $ "WTF just happened? I was supposed to get only a single row, but got: " ++ show x

606+

540 concurrencyControlFn <- getConcurrencyControlFn

607 concurrencyControlFn <- getConcurrencyControlFn

541− withResource pool $ \pollerDbConn -> forever $ concurrencyControlFn >>= \case

608+ withResource pool $ \pollerDbConn -> forever $ concurrencyControlFn pollerDbConn >>= \case

542− False -> do

609+ DontPoll -> do

543 log LevelWarn $ LogText "NOT polling the job queue due to concurrency control"

610 log LevelWarn $ LogText "NOT polling the job queue due to concurrency control"

544 -- If we can't run any jobs ATM, relax and wait for resources to free up

611 -- If we can't run any jobs ATM, relax and wait for resources to free up

545 delayAction

612 delayAction

546− True -> do

613+ PollAny -> poll pollerDbConn Nothing

547− join

614+ PollWithResources resCfg -> poll pollerDbConn (Just resCfg)

548− (mask_ $ do

615+

549− log LevelDebug $ LogText $ toS $ "[" <> processName <> "] Polling the job queue.."

550− r <- liftIO $ jobPollingIO pollerDbConn processName tname lockTimeout

551− case r of

552− -- When we don't have any jobs to run, we can relax a bit...

553− [] -> pure delayAction

554−

555− -- When we find a job to run, fork and try to find the next job without any delay...

556− [Only (jid :: JobId)] -> do

557− void $ async $ runJob jid

558− pure noDelayAction

559−

560− x -> error $ "WTF just happened? I was supposed to get only a single row, but got: " ++ show x)

561 where

616 where

562 delayAction = delaySeconds =<< getPollingInterval

617 delayAction = delaySeconds =<< getPollingInterval

563 noDelayAction = pure ()

618 noDelayAction = pure ()

@@ -573,9 +628,17 @@ jobEventListener = do

573 jwName <- liftIO jobWorkerName

628 jwName <- liftIO jobWorkerName

574 concurrencyControlFn <- getConcurrencyControlFn

629 concurrencyControlFn <- getConcurrencyControlFn

575

630

576− let tryLockingJob jid = do

631+ let tryLockingJob jid mResCfg = withDbConnection $ \conn -> do

577 let q = "UPDATE ? SET status=?, locked_at=now(), locked_by=?, attempts=attempts+1 WHERE id=? AND status in ? RETURNING id"

632 let q = "UPDATE ? SET status=?, locked_at=now(), locked_by=?, attempts=attempts+1 WHERE id=? AND status in ? RETURNING id"

633+ qWithResources =

634+ "UPDATE ? SET status=?, locked_at=now(), locked_by=?, attempts=attempts+1 \

635+ \ WHERE id=? AND status in ? AND ?(id) RETURNING id"

636+ result <- case mResCfg of

578− withDbConnection (\conn -> liftIO $ PGS.query conn q (tname, Locked, jwName, jid, In [Queued, Retry])) >>= \case

637+ Nothing -> liftIO $ PGS.query conn q (tname, Locked, jwName, jid, In [Queued, Retry])

638+ Just ResourceCfg{..} -> liftIO $ PGS.query conn qWithResources

639+ (tname, Locked, jwName, jid, In [Queued, Retry], resCfgCheckResourceFunction)

640+

641+ case result of

579 [] -> do

642 [] -> do

580 log LevelDebug $ LogText $ toS $ "Job was locked by someone else before I could start. Skipping it. JobId=" <> show jid

643 log LevelDebug $ LogText $ toS $ "Job was locked by someone else before I could start. Skipping it. JobId=" <> show jid

581 pure Nothing

644 pure Nothing

@@ -587,29 +650,34 @@ jobEventListener = do

587 forever $ do

650 forever $ do

588 log LevelDebug $ LogText "[LISTEN/NOTIFY] Event loop"

651 log LevelDebug $ LogText "[LISTEN/NOTIFY] Event loop"

589 notif <- liftIO $ getNotification monitorDbConn

652 notif <- liftIO $ getNotification monitorDbConn

590− concurrencyControlFn >>= \case

653+

591− False -> log LevelWarn $ LogText "Received job event, but ignoring it due to concurrency control"

592− True -> do

593− let pload = notificationData notif

654+ let pload = notificationData notif

655+ runNotifWithFilter :: HasJobRunner m => Maybe ResourceCfg -> m ()

656+ runNotifWithFilter mResCfg = do

594− log LevelDebug $ LogText $ toS $ "NOTIFY | " <> show pload

657+ log LevelDebug $ LogText $ toS $ "NOTIFY | " <> show pload

595− case eitherDecode $ toS pload of

658+ case eitherDecode $ toS pload of

596− Left e -> log LevelError $ LogText $ toS $ "Unable to decode notification payload received from Postgres. Payload=" <> show pload <> " Error=" <> show e

659+ Left e -> log LevelError $ LogText $ toS $ "Unable to decode notification payload received from Postgres. Payload=" <> show pload <> " Error=" <> show e

597−

660+

598− -- Checking if job needs to be fired immediately AND it is not already

661+ -- Checking if job needs to be fired immediately AND it is not already

599− -- taken by some othe thread, by the time it got to us

662+ -- taken by some othe thread, by the time it got to us

600− Right (v :: Value) -> case Aeson.parseMaybe parser v of

663+ Right (v :: Value) -> case Aeson.parseMaybe parser v of

601− Nothing -> log LevelError $ LogText $ toS $ "Unable to extract id/run_at/locked_at from " <> show pload

664+ Nothing -> log LevelError $ LogText $ toS $ "Unable to extract id/run_at/locked_at from " <> show pload

602− Just (jid, runAt_, mLockedAt_) -> do

665+ Just (jid, runAt_, mLockedAt_) -> do

603− t <- liftIO getCurrentTime

666+ t <- liftIO getCurrentTime

604− if (runAt_ <= t) && isNothing mLockedAt_

667+ if (runAt_ <= t) && isNothing mLockedAt_

605− then do log LevelDebug $ LogText $ toS $ "Job needs needs to be run immediately. Attempting to fork in background. JobId=" <> show jid

668+ then do log LevelDebug $ LogText $ toS $ "Job needs needs to be run immediately. Attempting to fork in background. JobId=" <> show jid

606− void $ async $ do

669+ void $ async $ do

607− -- Let's try to lock the job first... it is possible that it has already

670+ -- Let's try to lock the job first... it is possible that it has already

608− -- been picked up by the poller by the time we get here.

671+ -- been picked up by the poller by the time we get here.

609− tryLockingJob jid >>= \case

672+ tryLockingJob jid mResCfg >>= \case

610− Nothing -> pure ()

673+ Nothing -> pure ()

611− Just lockedJid -> runJob lockedJid

674+ Just lockedJid -> runJob lockedJid

612− else log LevelDebug $ LogText $ toS $ "Job is either for future, or is already locked. Skipping. JobId=" <> show jid

675+ else log LevelDebug $ LogText $ toS $ "Job is either for future, is already locked, or would violate concurrency constraints. Skipping. JobId=" <> show jid

676+

677+ concurrencyControlFn monitorDbConn >>= \case

678+ DontPoll -> log LevelWarn $ LogText "Received job event, but ignoring it due to concurrency control"

679+ PollAny -> runNotifWithFilter Nothing

680+ PollWithResources resCfg -> runNotifWithFilter (Just resCfg)

613 where

681 where

614 parser :: Value -> Aeson.Parser (JobId, UTCTime, Maybe UTCTime)

682 parser :: Value -> Aeson.Parser (JobId, UTCTime, Maybe UTCTime)

615 parser = withObject "expecting an object to parse job.run_at and job.locked_at" $ \o -> do

683 parser = withObject "expecting an object to parse job.run_at and job.locked_at" $ \o -> do

@@ -623,6 +691,12 @@ jobEventListener = do

623createJobQuery :: PGS.Query

691createJobQuery :: PGS.Query

624createJobQuery = "INSERT INTO ? (run_at, status, payload, last_error, attempts, locked_at, locked_by) VALUES (?, ?, ?, ?, ?, ?, ?) RETURNING " <> concatJobDbColumns

692createJobQuery = "INSERT INTO ? (run_at, status, payload, last_error, attempts, locked_at, locked_by) VALUES (?, ?, ?, ?, ?, ?, ?) RETURNING " <> concatJobDbColumns

625

693

694+ensureResource :: PGS.Query

695+ensureResource = "INSERT INTO ? (id, usage_limit) VALUES (?, ?) ON CONFLICT DO NOTHING"

696+

697+registerResourceUsage :: PGS.Query

698+registerResourceUsage = "INSERT INTO ? (job_id, resource_id, usage) VALUES (?, ?, ?)"

699+

626-- $createJobs

700-- $createJobs

627--

701--

628-- Ideally you'd want to create wrappers for 'createJob' and 'scheduleJob' in

702-- Ideally you'd want to create wrappers for 'createJob' and 'scheduleJob' in

@@ -668,6 +742,50 @@ scheduleJob conn tname payload runAt = do

668 [r] -> pure r

742 [r] -> pure r

669 _ -> Prelude.error . (<> "Not expecting multiple rows when creating a single job. Query=") <$> queryFormatter

743 _ -> Prelude.error . (<> "Not expecting multiple rows when creating a single job. Query=") <$> queryFormatter

670

744

745+type ResourceList = [(ResourceId, Int)]

746+

747+createJobWithResources

748+ :: ToJSON p

749+ => Connection

750+ -> TableName

751+ -> ResourceCfg

752+ -> p

753+ -> ResourceList

754+ -> IO Job

755+createJobWithResources conn tname resCfg payload resources = do

756+ t <- getCurrentTime

757+ scheduleJobWithResources conn tname resCfg payload resources t

758+

759+scheduleJobWithResources

760+ :: ToJSON p

761+ => Connection

762+ -> TableName

763+ -> ResourceCfg

764+ -> p

765+ -> ResourceList

766+ -> UTCTime

767+ -> IO Job

768+scheduleJobWithResources conn tname ResourceCfg{..} payload resources runAt = do

769+ -- We insert everything in a single transaction to delay @NOTIFY@ calls,

770+ -- so a job isn't picked up before its resources are inserted.

771+ PGS.begin conn

772+ let args = ( tname, runAt, Queued, toJSON payload, Nothing :: Maybe Value, 0 :: Int, Nothing :: Maybe Text, Nothing :: Maybe Text )

773+ queryFormatter = toS <$> PGS.formatQuery conn createJobQuery args

774+ rs <- PGS.query conn createJobQuery args

775+

776+ job <- flip onException (PGS.rollback conn) $ case rs of

777+ [] -> Prelude.error . (<> "Not expecting a blank result set when creating a job. Query=") <$> queryFormatter

778+ [r] -> pure r

779+ _ -> Prelude.error . (<> "Not expecting multiple rows when creating a single job. Query=") <$> queryFormatter

780+

781+ forM_ resources $ \(resourceId, usage) -> do

782+ void $ PGS.execute conn ensureResource (resCfgResourceTable, rawResourceId resourceId, resCfgDefaultLimit)

783+ void $ PGS.execute conn registerResourceUsage (resCfgUsageTable, jobId job, rawResourceId resourceId, usage)

784+

785+ PGS.commit conn

786+

787+ pure job

788+

671

789

672-- getRunnerEnv :: (HasJobRunner m) => m RunnerEnv

790-- getRunnerEnv :: (HasJobRunner m) => m RunnerEnv

673-- getRunnerEnv = ask

791-- getRunnerEnv = ask

src/OddJobs/Migrations.hs+56

@@ -1,3 +1,4 @@

1+{-# LANGUAGE RecordWildCards #-}

1module OddJobs.Migrations

2module OddJobs.Migrations

2 ( module OddJobs.Migrations

3 ( module OddJobs.Migrations

3 , module OddJobs.Types

4 , module OddJobs.Types

@@ -72,3 +73,58 @@ createJobTable conn tname = void $ do

72 fnName = PGS.Identifier $ "notify_job_monitor_for_" <> getTnameTxt tname

73 fnName = PGS.Identifier $ "notify_job_monitor_for_" <> getTnameTxt tname

73 trgName = PGS.Identifier $ "trg_notify_job_monitor_for_" <> getTnameTxt tname

74 trgName = PGS.Identifier $ "trg_notify_job_monitor_for_" <> getTnameTxt tname

74 getTnameTxt (PGS.QualifiedIdentifier _ tname') = tname'

75 getTnameTxt (PGS.QualifiedIdentifier _ tname') = tname'

76+

77+createResourceTableQuery :: Query

78+createResourceTableQuery = "CREATE TABLE IF NOT EXISTS ?" <>

79+ "( id text primary key" <>

80+ ", usage_limit int not null" <>

81+ ")";

82+

83+createUsageTableQuery :: Query

84+createUsageTableQuery = "CREATE TABLE IF NOT EXISTS ?" <>

85+ "( job_id serial not null REFERENCES ? ON DELETE CASCADE" <>

86+ ", resource_id text not null REFERENCES ? ON DELETE CASCADE" <>

87+ ", usage int not null" <>

88+ ", PRIMARY KEY (job_id, resource_id)" <>

89+ ");"

90+

91+createUsageFunction :: Query

92+createUsageFunction = "CREATE OR REPLACE FUNCTION ?(resourceId text) RETURNS int as $$" <>

93+ " SELECT coalesce(sum(usage), 0) FROM ? AS j INNER JOIN ? AS jr ON j.id = jr.job_id " <>

94+ " WHERE jr.resource_id = $1 AND j.status = ? " <>

95+ " $$ LANGUAGE SQL;"

96+

97+createCheckResourceFunction :: Query

98+createCheckResourceFunction = "CREATE OR REPLACE FUNCTION ?(jobId int) RETURNS bool as $$" <>

99+ " SELECT coalesce(bool_and(?(resource.id) + job_resource.usage <= resource.usage_limit), true) FROM " <>

100+ " ? AS job_resource INNER JOIN ? AS resource ON job_resource.resource_id = resource.id " <>

101+ " WHERE job_resource.job_id = $1" <>

102+ " $$ LANGUAGE SQL;"

103+

104+createResourceTables

105+ :: Connection

106+ -> TableName -- ^ Name of the jobs table

107+ -> ResourceCfg

108+ -> IO ()

109+createResourceTables conn jobTableName ResourceCfg{..} = do

110+ void $ PGS.execute conn createResourceTableQuery (PGS.Only resCfgResourceTable)

111+ void $ PGS.execute conn createUsageTableQuery

112+ ( resCfgUsageTable

113+ , jobTableName

114+ , resCfgResourceTable

115+ )

116+ void $ PGS.execute conn createUsageFunction

117+ ( usageFnName

118+ , jobTableName

119+ , resCfgUsageTable

120+ , Locked

121+ )

122+ void $ PGS.execute conn createCheckResourceFunction

123+ ( resCfgCheckResourceFunction

124+ , usageFnName

125+ , resCfgUsageTable

126+ , resCfgResourceTable

127+ )

128+ where

129+ usageFnName = PGS.Identifier $ "calculate_usage_for_resource_from_" <> getTnameTxt resCfgUsageTable

130+ getTnameTxt (PGS.QualifiedIdentifier _ tname') = tname'

src/OddJobs/Types.hs+34

@@ -109,6 +109,30 @@ data FailureMode

109-- by 'JobErrHandler' and 'cfgOnJobFailed'.

110data JobErrHandler = forall a e . (Exception e) => JobErrHandler (e -> Job -> FailureMode -> IO a)

111

112+type FunctionName = PGS.Identifier

113+

114+data ResourceCfg = ResourceCfg

115+ { resCfgResourceTable :: TableName

116+ -- ^ Table to use for tracking resources and their limits. Both this and

117+ -- 'resCfgUsageTable' should be created by 'OddJobs.Migrations.createResourceTables'.

118+

119+ , resCfgUsageTable :: TableName

120+ -- ^ Table to use for tracking how jobs use resources.

121+

122+ , resCfgCheckResourceFunction :: FunctionName

123+ -- ^ Name of the function that checks that the resources required to run a

124+ -- job are all available (i.e. that the total usage of each resource, plus

125+ -- the usage the job needs to run, does not exceed the resource limit). The

126+ -- function should have the signature @(int) RETURNS bool@, and return @TRUE@

127+ -- if the job with the given ID has its resources available.

128+

129+ , resCfgDefaultLimit :: Int

130+ -- ^ When a job requires a resource not already in 'resCfgResourceTable',

131+ -- what should its limit be set to?

132+ } deriving (Show)

133+

134+newtype ResourceId = ResourceId { rawResourceId :: Text }

135+ deriving (Show)

112

136

113-- | __Note:__ Please read the section on [controlling

137-- | __Note:__ Please read the section on [controlling

114-- concurrency](https://www.haskelltutorials.com/odd-jobs/guide.html#controlling-concurrency)

138-- concurrency](https://www.haskelltutorials.com/odd-jobs/guide.html#controlling-concurrency)

@@ -125,6 +149,15 @@ data ConcurrencyControl

125 -- such as, CPU, memory, disk IOPS, or network bandwidth.

149 -- such as, CPU, memory, disk IOPS, or network bandwidth.

126 | UnlimitedConcurrentJobs

150 | UnlimitedConcurrentJobs

127

151

152+ -- | Limit jobs according to their access to resources, as tracked in the DB

153+ -- according to the 'ResourceCfg'.

154+ --

155+ -- __Warning:__ without sufficient limits, this can easily hit the same problems

156+ -- as 'UnlimitedConcurrentJobs' where jobs are able to exhaust system resources.

157+ -- It is therefore recommended that all jobs in your system use /some/ DB-tracked

158+ -- resource.

159+ | ResourceLimits ResourceCfg

160+

128 -- | Use this to dynamically determine if the next job should be picked-up, or

161 -- | Use this to dynamically determine if the next job should be picked-up, or

129 -- not. This is useful to write custom-logic to determine whether a limited

162 -- not. This is useful to write custom-logic to determine whether a limited

130 -- resource is below a certain usage threshold (eg. CPU usage is below 80%).

163 -- resource is below a certain usage threshold (eg. CPU usage is below 80%).

@@ -135,6 +168,7 @@ instance Show ConcurrencyControl where

135 show cc = case cc of

168 show cc = case cc of

136 MaxConcurrentJobs n -> "MaxConcurrentJobs " <> show n

169 MaxConcurrentJobs n -> "MaxConcurrentJobs " <> show n

137 UnlimitedConcurrentJobs -> "UnlimitedConcurrentJobs"

170 UnlimitedConcurrentJobs -> "UnlimitedConcurrentJobs"

171+ ResourceLimits cfg -> "ResourceLimits " <> show cfg

138 DynamicConcurrency _ -> "DynamicConcurrency (IO Bool)"

172 DynamicConcurrency _ -> "DynamicConcurrency (IO Bool)"

139

173

140type JobId = Int

174type JobId = Int

test/Test.hs+166−2

@@ -1,6 +1,8 @@

1{-# LANGUAGE FlexibleInstances, NamedFieldPuns, DeriveGeneric, FlexibleContexts, TypeFamilies, StandaloneDeriving, RankNTypes #-}

2{-# LANGUAGE CPP #-}

3+{-# LANGUAGE PartialTypeSignatures #-}

4+{-# LANGUAGE RecordWildCards #-}

3−{-# OPTIONS_GHC -Wno-missing-signatures #-}

5+{-# OPTIONS_GHC -Wno-missing-signatures -Wno-partial-type-signatures #-}

4module Test where

6module Test where

6import Test.Tasty as Tasty

8import Test.Tasty as Tasty

@@ -103,6 +105,7 @@ tests appPool jobPool = testGroup "All tests"

103 , testOnJobSuccess appPool jobPool

105 , testOnJobSuccess appPool jobPool

104 , testOnJobTimeout appPool jobPool

106 , testOnJobTimeout appPool jobPool

105 ]

107 ]

108+ , testResourceLimitedScheduling appPool jobPool

106 ]

109 ]

107 -- , testGroup "property tests" [ testEverything appPool jobPool

110 -- , testGroup "property tests" [ testEverything appPool jobPool

108 -- -- , propFilterJobs appPool jobPool

111 -- -- , propFilterJobs appPool jobPool

@@ -238,13 +241,31 @@ ensureJobId conn tname jid = Job.findJobByIdIO conn tname jid >>= \case

238 Nothing -> error $ "Not expecting job to be deleted. JobId=" <> show jid

241 Nothing -> error $ "Not expecting job to be deleted. JobId=" <> show jid

239 Just j -> pure j

242 Just j -> pure j

240

243

241−-- withRandomTable :: (MonadIO m) => Pool Connection -> (Job.TableName -> m a) -> m a

244+withRandomTable :: _ => Pool Connection -> (Job.TableName -> m a) -> m a

242withRandomTable jobPool action = do

245withRandomTable jobPool action = do

243 (tname :: Job.TableName) <- liftIO (fromString . ("jobs_" <>) <$> replicateM 10 (R.randomRIO ('a', 'z')))

246 (tname :: Job.TableName) <- liftIO (fromString . ("jobs_" <>) <$> replicateM 10 (R.randomRIO ('a', 'z')))

244 finally

247 finally

245 (Pool.withResource jobPool (\conn -> liftIO $ Migrations.createJobTable conn tname) >> action tname)

248 (Pool.withResource jobPool (\conn -> liftIO $ Migrations.createJobTable conn tname) >> action tname)

246 (Pool.withResource jobPool $ \conn -> liftIO $ void $ PGS.execute conn "drop table if exists ?" (Only tname))

249 (Pool.withResource jobPool $ \conn -> liftIO $ void $ PGS.execute conn "drop table if exists ?" (Only tname))

247

250

251+withRandomResourceTables :: _ => Int -> Pool Connection -> Job.TableName -> (Job.ResourceCfg -> m a) -> m a

252+withRandomResourceTables defaultLimit jobPool tname action = do

253+ resCfgResourceTable <- liftIO $ fromString . ("resources_" <>) <$> replicateM 10 (R.randomRIO ('a', 'z'))

254+ resCfgUsageTable <- liftIO $ fromString . ("usage_" <>) <$> replicateM 10 (R.randomRIO ('a', 'z'))

255+ resCfgCheckResourceFunction <- liftIO $ fromString . ("check_resource_fn_" <>) <$> replicateM 10 (R.randomRIO ('a', 'z'))

256+ let resCfg = Job.ResourceCfg

257+ { Job.resCfgDefaultLimit = defaultLimit

258+ , ..

259+ }

260+

261+ finally

262+ (do

263+ Pool.withResource jobPool $ \conn -> liftIO $ Migrations.createResourceTables conn tname resCfg

264+ action resCfg)

265+ (Pool.withResource jobPool $ \conn -> liftIO $ void $ PGS.execute conn

266+ "drop function if exists ?; drop table if exists ?; drop table if exists ?"

267+ (resCfgCheckResourceFunction, resCfgUsageTable, resCfgResourceTable))

268+

248testMaxAttempts :: Int

269testMaxAttempts :: Int

249testMaxAttempts = 3

270testMaxAttempts = 3

250

271

@@ -438,6 +459,149 @@ testRetryBackoff appPool jobPool = testCase "retry backoff" $ do

438 modifyRetryBackoff cfg

459 modifyRetryBackoff cfg

439 = cfg { Job.cfgDefaultRetryBackoff = pure . backoff }

460 = cfg { Job.cfgDefaultRetryBackoff = pure . backoff }

440

461

462+testResourceLimitedScheduling appPool jobPool = testGroup "concurrency control by resources"

463+ [ testCase "resources control how jobs run" $ setup $ \tname resCfg logRef conn -> do

464+ let firstResource = [(Job.ResourceId "first", 1)]

465+ secondResource = [(Job.ResourceId "second", 1)]

466+

467+ Job{jobId = firstJob} <- Job.createJobWithResources conn tname resCfg (PayloadSucceed 10) firstResource

468+ Job{jobId = secondJob} <- Job.createJobWithResources conn tname resCfg (PayloadSucceed 10) firstResource

469+ Job{jobId = thirdJob} <- Job.createJobWithResources conn tname resCfg (PayloadSucceed 10) secondResource

470+

471+ delaySeconds $ Job.defaultPollingInterval + Seconds 2

472+

473+ assertJobIdStatus conn tname logRef "Job was created first - it should be running" Job.Locked firstJob

474+ assertJobIdStatus conn tname logRef "Job uses same resources as first - it shouldn't be running" Job.Queued secondJob

475+ assertJobIdStatus conn tname logRef "Job uses different resources - it should be running" Job.Locked thirdJob

476+

477+ , testCase "resource-less jobs run normally" $ setup $ \tname resCfg logRef conn -> do

478+ Job{jobId = firstJob} <- Job.createJobWithResources conn tname resCfg (PayloadSucceed 10) []

479+ Job{jobId = secondJob} <- Job.createJobWithResources conn tname resCfg (PayloadSucceed 10) []

480+

481+ delaySeconds $ Job.defaultPollingInterval + Seconds 2

482+

483+ assertJobIdStatus conn tname logRef "First job uses no resources - it should be running" Job.Locked firstJob

484+ assertJobIdStatus conn tname logRef "Second job uses no resources - it should be running" Job.Locked secondJob

485+

486+ , testCase "jobs can hold multiple resources" $ setup $ \tname resCfg logRef conn -> do

487+ let firstResource = [(Job.ResourceId "first", 1)]

488+ secondResource = [(Job.ResourceId "second", 1)]

489+

490+ Job{jobId = firstJob} <- Job.createJobWithResources conn tname resCfg (PayloadSucceed 10) (firstResource <> secondResource)

491+ Job{jobId = secondJob} <- Job.createJobWithResources conn tname resCfg (PayloadSucceed 10) firstResource

492+ Job{jobId = thirdJob} <- Job.createJobWithResources conn tname resCfg (PayloadSucceed 10) secondResource

493+

494+ delaySeconds $ Job.defaultPollingInterval + Seconds 2

495+

496+ assertJobIdStatus conn tname logRef "Job was created first - it should be running" Job.Locked firstJob

497+ assertJobIdStatus conn tname logRef "Job uses same resources as first - it shouldn't be running" Job.Queued secondJob

498+ assertJobIdStatus conn tname logRef "Job uses same resources as first - it shouldn't be running" Job.Queued thirdJob

499+

500+ , testCase "resources are freed when jobs succeed" $ setup $ \tname resCfg logRef conn -> do

501+ let firstResource = [(Job.ResourceId "first", 1)]

502+

503+ Job{jobId = firstJob} <- Job.createJobWithResources conn tname resCfg (PayloadSucceed 10) firstResource

504+ Job{jobId = secondJob} <- Job.createJobWithResources conn tname resCfg (PayloadSucceed 10) firstResource

505+

506+ delaySeconds $ Job.defaultPollingInterval + Seconds 2

507+

508+ assertJobIdStatus conn tname logRef "Job was created first - it should be running" Job.Locked firstJob

509+ assertJobIdStatus conn tname logRef "Job uses same resources as first - it shouldn't be running" Job.Queued secondJob

510+

511+ delaySeconds $ (Job.defaultPollingInterval * 2) + Seconds 8

512+

513+ assertJobIdStatus conn tname logRef "First job should have succeeded by now" Job.Success firstJob

514+ assertJobIdStatus conn tname logRef "Second job should be running after first job suceeded" Job.Locked secondJob

515+

516+ , testCase "resources are freed when jobs fail" $ setup $ \tname resCfg logRef conn -> do

517+ let firstResource = [(Job.ResourceId "first", 1)]

518+

519+ Job{jobId = firstJob} <- Job.createJobWithResources conn tname resCfg (PayloadAlwaysFail 10) firstResource

520+ Job{jobId = secondJob} <- Job.createJobWithResources conn tname resCfg (PayloadSucceed 10) firstResource

521+

522+ delaySeconds $ Job.defaultPollingInterval + Seconds 2

523+

524+ assertJobIdStatus conn tname logRef "Job was created first - it should be running" Job.Locked firstJob

525+ assertJobIdStatus conn tname logRef "Job uses same resources as first - it shouldn't be running" Job.Queued secondJob

526+

527+ delaySeconds $ (Job.defaultPollingInterval * 2) + Seconds 8

528+

529+ assertJobIdStatus conn tname logRef "First job should have failed by now" Job.Failed firstJob

530+ assertJobIdStatus conn tname logRef "Second job should be running after first job failed" Job.Locked secondJob

531+

532+ , testCase "resources allow for multiple usage" $ setup $ \tname resCfg logRef conn -> do

533+ let firstResource = [(Job.ResourceId "first", 1)]

534+ resCfg' = resCfg { Job.resCfgDefaultLimit = 2 }

535+

536+ Job{jobId = firstJob} <- Job.createJobWithResources conn tname resCfg' (PayloadSucceed 10) firstResource

537+ Job{jobId = secondJob} <- Job.createJobWithResources conn tname resCfg' (PayloadSucceed 10) firstResource

538+

539+ delaySeconds $ Job.defaultPollingInterval + Seconds 2

540+

541+ assertJobIdStatus conn tname logRef "First job uses some resources - it should be running" Job.Locked firstJob

542+ assertJobIdStatus conn tname logRef "Second job uses some resources - it should be running" Job.Locked secondJob

543+

544+ , testCase "resource usage can differ by job" $ setup $ \tname resCfg logRef conn -> do

545+ let firstResource n = [(Job.ResourceId "first", n)]

546+ resCfg' = resCfg { Job.resCfgDefaultLimit = 5 }

547+

548+ Job{jobId = firstJob} <- Job.createJobWithResources conn tname resCfg' (PayloadSucceed 10) (firstResource 4)

549+ Job{jobId = secondJob} <- Job.createJobWithResources conn tname resCfg' (PayloadSucceed 10) (firstResource 3)

550+ Job{jobId = thirdJob} <- Job.createJobWithResources conn tname resCfg' (PayloadSucceed 10) (firstResource 2)

551+

552+ delaySeconds $ Job.defaultPollingInterval + Seconds 2

553+

554+ assertJobIdStatus conn tname logRef "Job was created first - it should be running" Job.Locked firstJob

555+ assertJobIdStatus conn tname logRef "Job doesn't fit alongside first job - it shouldn't be running" Job.Queued secondJob

556+ assertJobIdStatus conn tname logRef "Job doesn't fit alongside first job - it shouldn't be running" Job.Queued thirdJob

557+

558+ delaySeconds $ (Job.defaultPollingInterval * 2) + Seconds 8

559+

560+ assertJobIdStatus conn tname logRef "First job should have succeeded by now" Job.Success firstJob

561+ assertJobIdStatus conn tname logRef "Second job should be running after first job suceeded" Job.Locked secondJob

562+ assertJobIdStatus conn tname logRef "Third job fits alongside second job - it should be running after first job suceeded" Job.Locked thirdJob

563+

564+ , testCase "resource limits permanently block too-big jobs" $ setup $ \tname resCfg logRef conn -> do

565+ let firstResource = [(Job.ResourceId "first", 2)]

566+

567+ Job{jobId = firstJob} <- Job.createJobWithResources conn tname resCfg (PayloadSucceed 10) firstResource

568+

569+ delaySeconds $ Job.defaultPollingInterval + Seconds 2

570+

571+ assertJobIdStatus conn tname logRef "Job requires too many resources - shouldn't be running" Job.Queued firstJob

572+

573+ , testCase "resources with 0 limit block jobs" $ setup' 0 $ \tname resCfg logRef conn -> do

574+ let firstResource = [(Job.ResourceId "first", 1)]

575+

576+ Job{jobId = firstJob} <- Job.createJobWithResources conn tname resCfg (PayloadSucceed 10) firstResource

577+

578+ delaySeconds $ Job.defaultPollingInterval + Seconds 2

579+

580+ assertJobIdStatus conn tname logRef "Job can't access 0-limited resources - shouldn't be running" Job.Queued firstJob

581+

582+ , testCase "resources with negative limit block jobs" $ setup' (-5) $ \tname resCfg logRef conn -> do

583+ let firstResource = [(Job.ResourceId "first", 1)]

584+

585+ Job{jobId = firstJob} <- Job.createJobWithResources conn tname resCfg (PayloadSucceed 10) firstResource

586+

587+ delaySeconds $ Job.defaultPollingInterval + Seconds 2

588+

589+ assertJobIdStatus conn tname logRef "Job can't access negative-limited resources - shouldn't be running" Job.Queued firstJob

590+ ]

591+ where

592+ setup = setup' 1

593+

594+ setup' defaultLimit action =

595+ withRandomTable jobPool $ \tname ->

596+ withRandomResourceTables defaultLimit jobPool tname $ \resCfg ->

597+ withNamedJobMonitor tname jobPool (cfgFn resCfg) $ \logRef -> do

598+ Pool.withResource appPool $ \conn -> action tname resCfg logRef conn

599+

600+ cfgFn resCfg cfg = cfg

601+ { Job.cfgDefaultMaxAttempts = 1 -- Simplifies some tests where we have a failing job

602+ , Job.cfgConcurrencyControl = Job.ResourceLimits resCfg

603+ }

604+

441data JobEvent = JobStart

605data JobEvent = JobStart

442 | JobRetry

606 | JobRetry

443 | JobSuccess

607 | JobSuccess