@@ -65,6 +65,7 @@ module OddJobs.Job
6565 , runJobNowIO
6666 , unlockJobIO
6767 , cancelJobIO
68+ , killJobIO
6869 , jobDbColumns
6970 , jobPollingIO
7071 , concatJobDbColumns
@@ -110,14 +111,16 @@ import Data.Aeson hiding (Success)
110111import qualified Data.Aeson as Aeson
111112import qualified Data.Aeson.Types as Aeson (Parser, parseMaybe)
112113import Data.String.Conv (StringConv(..), toS)
113−import Data.Functor (void)
114+import Data.Functor ((<&>), void)
114115import Control.Monad (forever, forM_, join)
115116import Data.Maybe (isNothing, maybe, fromMaybe, listToMaybe, mapMaybe)
116117import Data.Either (either)
117118import Control.Monad.Reader
118119import GHC.Generics
120+import Data.Map (Map)
119121import qualified Data.HashMap.Strict as HM
120122import qualified Data.List as DL
123+import qualified Data.Map as DM
121124import qualified Data.ByteString as BS
122125import qualified Data.ByteString.Lazy as BSL
123126import System.FilePath (FilePath)
@@ -169,7 +172,7 @@ class (MonadUnliftIO m, MonadBaseControl IO m) => HasJobRunner m where
169172
170173data RunnerEnv = RunnerEnv
171174 { envConfig :: !Config
172− , envJobThreadsRef :: !(IORef [Async ()])
175+ , envJobThreadsRef :: !(IORef (Map JobId (Async ())))
173176 }
174177
175178type RunnerM = ReaderT RunnerEnv IO
@@ -218,7 +221,7 @@ instance HasJobRunner RunnerM where
218221startJobRunner :: Config -> IO ()
219222startJobRunner jm = do
220223 traceM "startJobRunner top"
221− r <- newIORef []
224+ r <- newIORef DM.empty
222225 let monitorEnv = RunnerEnv
223226 { envConfig = jm
224227 , envJobThreadsRef = r
@@ -346,7 +349,7 @@ deleteJobIO conn tname jid = do
346349runJobNowIO :: Connection -> TableName -> JobId -> IO (Maybe Job)
347350runJobNowIO conn tname jid = do
348351 t <- getCurrentTime
349− updateJobHelper tname conn (Queued, [Queued, Retry, Failed], Just t, jid)
352+ updateJobHelper tname conn (Queued, [Queued, Retry, Failed, Cancelled], Just t, jid)
350353
351354-- | TODO: First check in all job-runners if this job is still running, or not,
352355-- and somehow send an uninterruptibleCancel to that thread.
@@ -358,7 +361,11 @@ unlockJobIO conn tname jid = do
358361
359362cancelJobIO :: Connection -> TableName -> JobId -> IO (Maybe Job)
360363cancelJobIO conn tname jid =
361− 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)
362369
363370updateJobHelper :: TableName
364371 -> Connection
@@ -383,13 +390,17 @@ runJobWithTimeout :: (HasJobRunner m)
383390 => Seconds
384391 -> Job
385392 -> m ()
386−runJobWithTimeout timeoutSec job = do
393+runJobWithTimeout timeoutSec job@Job{jobId} = do
387394 threadsRef <- envJobThreadsRef <$> getRunnerEnv
388395 jobRunner_ <- getJobRunner
389396
390397 a <- async $ liftIO $ jobRunner_ job
391398
392− _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+
393404 -- liftIO $ putStrLn $ "Threads: " <> show x
394405 log LevelDebug $ LogText $ toS $ "Spawned job in " <> show (asyncThreadId a)
395406
@@ -399,7 +410,7 @@ runJobWithTimeout timeoutSec job = do
399410
400411 void $ finally
401412 (waitEitherCancel a t)
402− (atomicModifyIORef' threadsRef $ \threads -> (DL.delete a threads, ()))
413+ (atomicModifyIORef' threadsRef $ \threads -> (DM.delete jobId threads, ()))
403414
404415
405416-- | runs a job, blocks for as long as it's in progress
@@ -455,6 +466,25 @@ runJob jid = do
455466
456467 pure ()
457468
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+
458488-- TODO: This might have a resource leak.
459489restartUponCrash :: (HasJobRunner m, Show a) => Text -> m a -> m ()
460490restartUponCrash name_ action = do
@@ -480,9 +510,11 @@ jobMonitor = do
480510 traceM "jobMonitor top"
481511 a1 <- async $ restartUponCrash "Job poller" jobPoller
482512 a2 <- async $ restartUponCrash "Job event listener" jobEventListener
513+ a3 <- async $ restartUponCrash "Job Kill poller" killJobPoller
483514 traceM "jobMonitor after async"
484− finally (void $ waitAnyCatch [a1, a2]) $ do
515+ finally (void $ waitAnyCatch [a1, a2, a3]) $ do
485516 log LevelInfo (LogText "Stopping jobPoller and jobEventListener threads.")
517+ cancel a3
486518 cancel a2
487519 cancel a1
488520 log LevelInfo (LogText "Waiting for jobs to complete.")
@@ -503,11 +535,19 @@ jobPollingWithResourceSql =
503535 \ ORDER BY attempts ASC, run_at ASC LIMIT 1) \
504536 \ RETURNING id"
505537
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"
545+
506546waitForJobs :: (HasJobRunner m)
507547 => m ()
508548waitForJobs = do
509− threadsRef <- envJobThreadsRef <$> getRunnerEnv
510− readIORef threadsRef >>= \case
549+ curJobs <- getRunnerEnv >>= (readIORef . envJobThreadsRef) <&> DM.elems
550+ case curJobs of
511551 [] -> log LevelInfo $ LogText "All job-threads exited"
512552 as -> do
513553 tid <- myThreadId
@@ -526,7 +566,7 @@ getConcurrencyControlFn :: (HasJobRunner m)
526566getConcurrencyControlFn = getConcurrencyControl >>= \case
527567 UnlimitedConcurrentJobs -> pure $ const $ pure PollAny
528568 MaxConcurrentJobs maxJobs -> pure $ const $ do
529− curJobs <- getRunnerEnv >>= (readIORef . envJobThreadsRef)
569+ curJobs <- getRunnerEnv >>= (readIORef . envJobThreadsRef) <&> DM.elems
530570 pure $ pollIf $ DL.length curJobs < maxJobs
531571 ResourceLimits resCfg -> pure $ const $ pure $ PollWithResources resCfg
532572 DynamicConcurrency fn -> pure $ const $ pollIf <$> liftIO fn
@@ -630,6 +670,43 @@ pollRunJob processName mResCfg = do
630670 delayAction = delaySeconds =<< getPollingInterval
631671 noDelayAction = pure ()
632672
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+
706+ where
707+ delayAction = delaySeconds =<< getPollingInterval
708+ noDelayAction = pure ()
709+
633710-- | Uses PostgreSQL's LISTEN/NOTIFY to be immediately notified of newly created
634711-- jobs.
635712jobEventListener :: (HasJobRunner m)