git.iko.soy/odd-jobs
on master help sign in
Merge pull request #80 from SupercedeTech/resource-tables-pr

Add resource-based concurrency control

commitcc05b640528697ba37529ee617a031d1acd60c4b authorIsaac van Bakel <isaac@supercede.com> dateThu 16 Feb 2023, 11:13 (4 years ago) parentf921230cec060e
4 files changed+425−53
compare with master
unifiedsplit
src/OddJobs/Job.hs+169−51 src/OddJobs/Migrations.hs+56 src/OddJobs/Types.hs+34 test/Test.hs+166−2
src/OddJobs/Job.hs+169−51
@@ -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,8 @@ module OddJobs.Job
3136 , Seconds(..)
3237 , JobErrHandler(..)
3338 , AllJobTypes(..)
39+ , ResourceId(..)
40+ , FunctionName
3441
3542 -- ** Structured logging
3643 --
@@ -78,7 +85,7 @@ import Data.Pool(Pool)
7885import Data.Text as T
7986import Database.PostgreSQL.Simple as PGS
8087import Database.PostgreSQL.Simple.Notification
81−import UnliftIO.Async
88+import UnliftIO.Async hiding (poll)
8289import UnliftIO.Concurrent (threadDelay, myThreadId)
8390import Data.String
8491import System.Posix.Process (getProcessID)
@@ -89,7 +96,8 @@ import Control.Monad.Logger as MLogger (LogLevel(..), LogStr, toLogStr)
8996import UnliftIO.IORef
9097import UnliftIO.Exception ( SomeException(..), try, catch, finally
9198 , catchAny, bracket, Exception(..), throwIO
92− , catches, Handler(..), mask_, throwString
99+ , catches, Handler(..), mask_, onException
100+ , throwString
93101 )
94102import Data.Proxy
95103import Control.Monad.Trans.Control
@@ -309,7 +317,7 @@ saveJob j = do
309317saveJobIO :: Connection -> TableName -> Job -> IO Job
310318saveJobIO conn tname Job{jobRunAt, jobStatus, jobPayload, jobLastError, jobAttempts, jobLockedBy, jobLockedAt, jobId} = do
311319 rs <- PGS.query conn saveJobQuery
312− ( tname
320+ ( tname
313321 , jobRunAt
314322 , jobStatus
315323 , jobPayload
@@ -476,7 +484,18 @@ jobMonitor = do
476484
477485-- | Ref: 'jobPoller'
478486jobPollingSql :: Query
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"
487+jobPollingSql =
488+ "update ? set status = ?, locked_at = ?, locked_by = ?, attempts=attempts+1\
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"
480499
481500waitForJobs :: (HasJobRunner m)
482501 => m ()
@@ -491,15 +510,23 @@ waitForJobs = do
491510 delaySeconds (Seconds 1)
492511 waitForJobs
493512
513+data ConcurrencyAction
514+ = DontPoll
515+ | PollAny
516+ | PollWithResources ResourceCfg
517+
494518getConcurrencyControlFn :: (HasJobRunner m)
495− => m (m Bool)
519+ => m (Connection -> m ConcurrencyAction)
496520getConcurrencyControlFn = getConcurrencyControl >>= \case
497− UnlimitedConcurrentJobs -> pure $ pure True
498− MaxConcurrentJobs maxJobs -> pure $ do
521+ UnlimitedConcurrentJobs -> pure $ const $ pure PollAny
522+ MaxConcurrentJobs maxJobs -> pure $ const $ do
499523 curJobs <- getRunnerEnv >>= (readIORef . envJobThreadsRef)
500− pure $ DL.length curJobs < maxJobs
501− DynamicConcurrency fn -> pure $ liftIO fn
524+ pure $ pollIf $ DL.length curJobs < maxJobs
525+ ResourceLimits resCfg -> pure $ const $ pure $ PollWithResources resCfg
526+ DynamicConcurrency fn -> pure $ const $ pollIf <$> liftIO fn
502527
528+ where
529+ pollIf cond = if cond then PollAny else DontPoll
503530
504531jobPollingIO :: Connection -> String -> TableName -> Seconds -> IO [Only JobId]
505532jobPollingIO pollerDbConn processName tname lockTimeout = do
@@ -537,27 +564,55 @@ jobPoller = do
537564 lockTimeout <- getDefaultJobTimeout
538565 traceM "jobPoller just before log statement"
539566 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+
540607 concurrencyControlFn <- getConcurrencyControlFn
541− withResource pool $ \pollerDbConn -> forever $ concurrencyControlFn >>= \case
542− False -> do
608+ withResource pool $ \pollerDbConn -> forever $ concurrencyControlFn pollerDbConn >>= \case
609+ DontPoll -> do
543610 log LevelWarn $ LogText "NOT polling the job queue due to concurrency control"
544611 -- If we can't run any jobs ATM, relax and wait for resources to free up
545612 delayAction
546− True -> do
547− join
548− (mask_ $ do
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)
613+ PollAny -> poll pollerDbConn Nothing
614+ PollWithResources resCfg -> poll pollerDbConn (Just resCfg)
615+
561616 where
562617 delayAction = delaySeconds =<< getPollingInterval
563618 noDelayAction = pure ()
@@ -573,9 +628,17 @@ jobEventListener = do
573628 jwName <- liftIO jobWorkerName
574629 concurrencyControlFn <- getConcurrencyControlFn
575630
576− let tryLockingJob jid = do
631+ let tryLockingJob jid mResCfg = withDbConnection $ \conn -> do
577632 let q = "UPDATE ? SET status=?, locked_at=now(), locked_by=?, attempts=attempts+1 WHERE id=? AND status in ? RETURNING id"
578− withDbConnection (\conn -> liftIO $ PGS.query conn q (tname, Locked, jwName, jid, In [Queued, Retry])) >>= \case
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
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
579642 [] -> do
580643 log LevelDebug $ LogText $ toS $ "Job was locked by someone else before I could start. Skipping it. JobId=" <> show jid
581644 pure Nothing
@@ -587,29 +650,34 @@ jobEventListener = do
587650 forever $ do
588651 log LevelDebug $ LogText "[LISTEN/NOTIFY] Event loop"
589652 notif <- liftIO $ getNotification monitorDbConn
590− concurrencyControlFn >>= \case
591− False -> log LevelWarn $ LogText "Received job event, but ignoring it due to concurrency control"
592− True -> do
593− let pload = notificationData notif
594− log LevelDebug $ LogText $ toS $ "NOTIFY | " <> show pload
595− 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
597−
598− -- 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
600− 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
602− Just (jid, runAt_, mLockedAt_) -> do
603− t <- liftIO getCurrentTime
604− 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
606− void $ async $ do
607− -- 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.
609− tryLockingJob jid >>= \case
610− Nothing -> pure ()
611− Just lockedJid -> runJob lockedJid
612− else log LevelDebug $ LogText $ toS $ "Job is either for future, or is already locked. Skipping. JobId=" <> show jid
653+
654+ let pload = notificationData notif
655+ runNotifWithFilter :: HasJobRunner m => Maybe ResourceCfg -> m ()
656+ runNotifWithFilter mResCfg = do
657+ log LevelDebug $ LogText $ toS $ "NOTIFY | " <> show pload
658+ case eitherDecode $ toS pload of
659+ Left e -> log LevelError $ LogText $ toS $ "Unable to decode notification payload received from Postgres. Payload=" <> show pload <> " Error=" <> show e
660+
661+ -- Checking if job needs to be fired immediately AND it is not already
662+ -- taken by some othe thread, by the time it got to us
663+ Right (v :: Value) -> case Aeson.parseMaybe parser v of
664+ Nothing -> log LevelError $ LogText $ toS $ "Unable to extract id/run_at/locked_at from " <> show pload
665+ Just (jid, runAt_, mLockedAt_) -> do
666+ t <- liftIO getCurrentTime
667+ if (runAt_ <= t) && isNothing mLockedAt_
668+ then do log LevelDebug $ LogText $ toS $ "Job needs needs to be run immediately. Attempting to fork in background. JobId=" <> show jid
669+ void $ async $ do
670+ -- Let's try to lock the job first... it is possible that it has already
671+ -- been picked up by the poller by the time we get here.
672+ tryLockingJob jid mResCfg >>= \case
673+ Nothing -> pure ()
674+ Just lockedJid -> runJob lockedJid
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)
613681 where
614682 parser :: Value -> Aeson.Parser (JobId, UTCTime, Maybe UTCTime)
615683 parser = withObject "expecting an object to parse job.run_at and job.locked_at" $ \o -> do
@@ -623,6 +691,12 @@ jobEventListener = do
623691createJobQuery :: PGS.Query
624692createJobQuery = "INSERT INTO ? (run_at, status, payload, last_error, attempts, locked_at, locked_by) VALUES (?, ?, ?, ?, ?, ?, ?) RETURNING " <> concatJobDbColumns
625693
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+
626700-- $createJobs
627701--
628702-- Ideally you'd want to create wrappers for 'createJob' and 'scheduleJob' in
@@ -668,6 +742,50 @@ scheduleJob conn tname payload runAt = do
668742 [r] -> pure r
669743 _ -> Prelude.error . (<> "Not expecting multiple rows when creating a single job. Query=") <$> queryFormatter
670744
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+
671789
672790-- getRunnerEnv :: (HasJobRunner m) => m RunnerEnv
673791-- getRunnerEnv = ask
src/OddJobs/Migrations.hs+56
@@ -1,3 +1,4 @@
1+{-# LANGUAGE RecordWildCards #-}
12module OddJobs.Migrations
23 ( module OddJobs.Migrations
34 , module OddJobs.Types
@@ -72,3 +73,58 @@ createJobTable conn tname = void $ do
7273 fnName = PGS.Identifier $ "notify_job_monitor_for_" <> getTnameTxt tname
7374 trgName = PGS.Identifier $ "trg_notify_job_monitor_for_" <> getTnameTxt tname
7475 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
109109-- by 'JobErrHandler' and 'cfgOnJobFailed'.
110110data JobErrHandler = forall a e . (Exception e) => JobErrHandler (e -> Job -> FailureMode -> IO a)
111111
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)
112136
113137-- | __Note:__ Please read the section on [controlling
114138-- concurrency](https://www.haskelltutorials.com/odd-jobs/guide.html#controlling-concurrency)
@@ -125,6 +149,15 @@ data ConcurrencyControl
125149 -- such as, CPU, memory, disk IOPS, or network bandwidth.
126150 | UnlimitedConcurrentJobs
127151
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+
128161 -- | Use this to dynamically determine if the next job should be picked-up, or
129162 -- not. This is useful to write custom-logic to determine whether a limited
130163 -- resource is below a certain usage threshold (eg. CPU usage is below 80%).
@@ -135,6 +168,7 @@ instance Show ConcurrencyControl where
135168 show cc = case cc of
136169 MaxConcurrentJobs n -> "MaxConcurrentJobs " <> show n
137170 UnlimitedConcurrentJobs -> "UnlimitedConcurrentJobs"
171+ ResourceLimits cfg -> "ResourceLimits " <> show cfg
138172 DynamicConcurrency _ -> "DynamicConcurrency (IO Bool)"
139173
140174type JobId = Int
test/Test.hs+166−2
@@ -1,6 +1,8 @@
11{-# LANGUAGE FlexibleInstances, NamedFieldPuns, DeriveGeneric, FlexibleContexts, TypeFamilies, StandaloneDeriving, RankNTypes #-}
22{-# LANGUAGE CPP #-}
3−{-# OPTIONS_GHC -Wno-missing-signatures #-}
3+{-# LANGUAGE PartialTypeSignatures #-}
4+{-# LANGUAGE RecordWildCards #-}
5+{-# OPTIONS_GHC -Wno-missing-signatures -Wno-partial-type-signatures #-}
46module Test where
57
68import Test.Tasty as Tasty
@@ -103,6 +105,7 @@ tests appPool jobPool = testGroup "All tests"
103105 , testOnJobSuccess appPool jobPool
104106 , testOnJobTimeout appPool jobPool
105107 ]
108+ , testResourceLimitedScheduling appPool jobPool
106109 ]
107110 -- , testGroup "property tests" [ testEverything appPool jobPool
108111 -- -- , propFilterJobs appPool jobPool
@@ -238,13 +241,31 @@ ensureJobId conn tname jid = Job.findJobByIdIO conn tname jid >>= \case
238241 Nothing -> error $ "Not expecting job to be deleted. JobId=" <> show jid
239242 Just j -> pure j
240243
241−-- withRandomTable :: (MonadIO m) => Pool Connection -> (Job.TableName -> m a) -> m a
244+withRandomTable :: _ => Pool Connection -> (Job.TableName -> m a) -> m a
242245withRandomTable jobPool action = do
243246 (tname :: Job.TableName) <- liftIO (fromString . ("jobs_" <>) <$> replicateM 10 (R.randomRIO ('a', 'z')))
244247 finally
245248 (Pool.withResource jobPool (\conn -> liftIO $ Migrations.createJobTable conn tname) >> action tname)
246249 (Pool.withResource jobPool $ \conn -> liftIO $ void $ PGS.execute conn "drop table if exists ?" (Only tname))
247250
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+
248269testMaxAttempts :: Int
249270testMaxAttempts = 3
250271
@@ -438,6 +459,149 @@ testRetryBackoff appPool jobPool = testCase "retry backoff" $ do
438459 modifyRetryBackoff cfg
439460 = cfg { Job.cfgDefaultRetryBackoff = pure . backoff }
440461
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+
441605data JobEvent = JobStart
442606 | JobRetry
443607 | JobSuccess