feat(jobs): batch job offloading

BREAKING CHANGE: Job offloading

src/Handler/Admin/Crontab.hs

@@ -17,6 +17,8 @@ import qualified Data.Text.Lazy.Builder as Text.Builder
 import qualified Data.HashSet as HashSet
 import qualified Data.HashMap.Strict as HashMap
 
+import qualified Data.UUID as UUID
+
 
 deriveJSON defaultOptions
   { constructorTagModifier = camelToPathPiece' 1
@@ -32,9 +34,12 @@ getAdminCrontabR = do
 
   let mCrontab = mCrontab' <&> _2 %~ filter (hasn't $ _3 . _MatchNone)
 
+  instanceId <- getsYesod appInstanceID
+
   selectRep $ do
     provideRep $ do
       crontabBearer <- runMaybeT . hoist runDB $ do
+        guardM $ hasGlobalGetParam GetGenerateToken
         uid <- MaybeT maybeAuthId
         guardM . lift . existsBy $ UniqueUserGroupMember UserGroupCrontab uid
 
@@ -49,6 +54,10 @@ getAdminCrontabR = do
             <section>
               <pre .token>
                 #{toPathPiece t}
+          <section>
+            <dl .deflist>
+              <dt .deflist__dt>_{MsgInstanceId}
+              <dd .deflist__dd .uuid>#{UUID.toText instanceId}
           <section>
             $maybe (genTime, crontab) <- mCrontab
               <p>

src/Handler/Metrics.hs

@@ -26,6 +26,7 @@ getMetricsR = selectRep $ do
       samples <- sortBy metricSort <$> collectMetrics
 
       metricsBearer <- runMaybeT . hoist runDB $ do
+        guardM $ hasGlobalGetParam GetGenerateToken
         uid <- MaybeT maybeAuthId
         guardM . lift . existsBy $ UniqueUserGroupMember UserGroupMetrics uid
 

src/Jobs.hs

@@ -10,6 +10,7 @@ module Jobs
 import Import hiding (StateT)
 import Jobs.Types as Types hiding (JobCtl(JobCtlQueue))
 import Jobs.Queue
+import Jobs.Offload
 import Jobs.Crontab
 
 import qualified Data.Conduit.List as C
@@ -105,6 +106,7 @@ handleJobs foundation@UniWorX{..}
       jobShutdown <- liftIO newEmptyTMVarIO
       jobCurrentCrontab <- liftIO $ newTVarIO Nothing
       jobHeldLocks <- liftIO $ newTVarIO Set.empty
+      jobOffload <- liftIO newEmptyTMVarIO
       registerJobHeldLocksCount jobHeldLocks
       registerJobWorkerQueueDepth appJobState
       atomically $ putTMVar appJobState JobState
@@ -155,7 +157,9 @@ manageJobPool foundation@UniWorX{..} unmask = shutdownOnException $ \routeExc ->
       atomically . asum $
         [ spawnMissingWorkers
         , reapDeadWorkers
-        ] ++ maybe [] (\(cTime, delay) -> [return () <$ waitDelay delay, transferJobs cTime]) transferInfo ++
+        ] ++ maybe [] (\(cTime, delay) -> [return () <$ waitDelay delay, transferJobs cTime]) transferInfo
+          ++ maybeToList (manageOffloadHandler <$> mkJobOffloadHandler (appDatabaseConf appSettings') (appJobMode appSettings'))
+          ++
         [ terminateGracefully terminate'
         ]
   where
@@ -286,6 +290,27 @@ manageJobPool foundation@UniWorX{..} unmask = shutdownOnException $ \routeExc ->
 
       return $ $logWarnS "JobPoolManager" [st|Moved #{tshow (olength movePairs)} long-unadressed jobs from #{tshow (olength senders)} senders to #{tshow (olength receivers)} receivers|]
 
+    manageOffloadHandler :: (ReaderT UniWorX m JobOffloadHandler) -> STM (ContT () m ())
+    manageOffloadHandler spawn = do
+      shouldTerminate' <- readTMVar appJobState >>= fmap not . isEmptyTMVar . jobShutdown
+      guard $ not shouldTerminate'
+
+      JobContext{jobOffload} <- jobContext <$> readTMVar appJobState
+      cOffload <- tryReadTMVar jobOffload
+
+      let respawn = do
+            nOffload <- lift $ runReaderT spawn foundation
+            atomically $ do
+              putTMVar jobOffload nOffload
+              whenIsJust cOffload $ \pOffload -> do
+                pOutgoing <- readTVar $ jobOffloadOutgoing pOffload
+                modifyTVar (jobOffloadOutgoing nOffload) (pOutgoing <>)
+
+      respawn <$ case cOffload of
+        Nothing -> return ()
+        Just JobOffloadHandler{..} -> waitSTM jobOffloadHandler
+      
+
 stopJobCtl :: MonadUnliftIO m => UniWorX -> m ()
 -- ^ Stop all worker threads currently running
 stopJobCtl UniWorX{appJobState} = do
@@ -471,7 +496,16 @@ handleJobs' wNum = C.mapM_ $ \jctl -> hoist delimitInternalState . withJobWorker
       $logDebugS logIdent "JobCtlQueue..."
       lift $ queueJob' job
       $logInfoS logIdent "JobCtlQueue"
-    handleCmd (JobCtlPerform jId) = handle handleQueueException . jLocked jId $ \(Entity _ j@QueuedJob{..}) -> lift $ do
+    handleCmd (JobCtlPerform jId) = do
+      jMode <- getsYesod $ view _appJobMode
+      case jMode of
+        JobsLocal{} -> performLocal
+        JobsOffload -> performOffload
+      where
+        performOffload = hoist atomically $ do
+          JobOffloadHandler{..} <- lift . readTMVar =<< asks jobOffload
+          lift $ modifyTVar jobOffloadOutgoing (`snoc` jId)
+        performLocal = handle handleQueueException . jLocked jId $ \(Entity _ j@QueuedJob{..}) -> lift $ do
           content <- case fromJSON queuedJobContent of
             Aeson.Success c -> return c
             Aeson.Error t -> do

src/Jobs/Crontab.hs

@@ -27,6 +27,29 @@ determineCrontab :: DB (Crontab JobCtl)
 determineCrontab = execWriterT $ do
   UniWorX{ appSettings' = AppSettings{..} } <- getYesod
 
+  whenIsJust appJobCronInterval $ \interval ->
+    tell $ HashMap.singleton
+      JobCtlDetermineCrontab
+      Cron
+        { cronInitial = CronAsap
+        , cronRepeat = CronRepeatScheduled CronAsap
+        , cronRateLimit = interval
+        , cronNotAfter = Right CronNotScheduled
+        }
+
+  tell . flip foldMap universeF $ \kind ->
+    case appHealthCheckInterval kind of
+      Just int -> HashMap.singleton
+        (JobCtlGenerateHealthReport kind)
+        Cron
+          { cronInitial = CronAsap
+          , cronRepeat = CronRepeatScheduled CronAsap
+          , cronRateLimit = int
+          , cronNotAfter = Right CronNotScheduled
+          }
+      Nothing  -> mempty
+
+  when (is _JobsLocal appJobMode) $ do
     case appJobFlushInterval of
       Just interval -> tell $ HashMap.singleton
         JobCtlFlush
@@ -38,16 +61,6 @@ determineCrontab = execWriterT $ do
           }
       Nothing -> return ()
 
-  whenIsJust appJobCronInterval $ \interval ->
-    tell $ HashMap.singleton
-      JobCtlDetermineCrontab
-      Cron
-        { cronInitial = CronAsap
-        , cronRepeat = CronRepeatScheduled CronAsap
-        , cronRateLimit = interval
-        , cronNotAfter = Right CronNotScheduled
-        }
-
     oldestInvitationMUTC <- lift $ preview (_head . _entityVal . _invitationExpiresAt . _Just) <$> selectList [InvitationExpiresAt !=. Nothing] [Asc InvitationExpiresAt, LimitTo 1]
     whenIsJust oldestInvitationMUTC $ \oldestInvUTC -> tell $ HashMap.singleton
       (JobCtlQueue JobPruneInvitations)
@@ -119,18 +132,6 @@ determineCrontab = execWriterT $ do
           , cronNotAfter    = Right CronNotScheduled
           }
 
-  tell . flip foldMap universeF $ \kind ->
-    case appHealthCheckInterval kind of
-      Just int -> HashMap.singleton
-        (JobCtlGenerateHealthReport kind)
-        Cron
-          { cronInitial = CronAsap
-          , cronRepeat = CronRepeatScheduled CronAsap
-          , cronRateLimit = int
-          , cronNotAfter = Right CronNotScheduled
-          }
-      Nothing  -> mempty
-
     let newyear = cronCalendarAny
           { cronDayOfYear  = cronMatchOne 1
           }
@@ -472,7 +473,7 @@ determineCrontab = execWriterT $ do
           , cronRateLimit = appNotificationRateLimit
           , cronNotAfter  = maybe (Right CronNotScheduled) (Right . CronTimestamp . utcToLocalTimeTZ appTZ) $ nBot =<< minimumOf (folded . _entityVal . _allocationStaffAllocationTo . to NTop . filtered (> NTop (Just staffAllocationFrom))) allocs
           }
-  iforM (allocationTimes AllocationRegisterTo) $ \registerTo allocs' -> do
+    iforM_ (allocationTimes AllocationRegisterTo) $ \registerTo allocs' -> do
       let allocs = flip filter allocs' $ \(Entity _ Allocation{..}) -> maybe True (> registerTo) allocationStaffAllocationTo
       tell $ HashMap.singleton
         (JobCtlQueue $ JobQueueNotification NotificationAllocationUnratedApplications{ nAllocations = setOf (folded . _entityKey) allocs })

src/Jobs/Offload.hs

@@ -0,0 +1,70 @@
+module Jobs.Offload
+  ( mkJobOffloadHandler
+  ) where
+
+import Import hiding (bracket, js)
+import Jobs.Types
+import Jobs.Queue
+
+import qualified Database.PostgreSQL.Simple as PG
+import qualified Database.PostgreSQL.Simple.Types as PG
+import qualified Database.PostgreSQL.Simple.Notification as PG
+
+import Database.Persist.Postgresql (PostgresConf, pgConnStr)
+
+import Data.Text.Encoding (decodeUtf8')
+
+import UnliftIO.Exception (bracket)
+
+
+jobOffloadChannel :: Text
+jobOffloadChannel = "job-offload"
+
+mkJobOffloadHandler :: forall m.
+                       ( MonadResource m
+                       , MonadUnliftIO m
+                       , MonadThrow m, MonadReader UniWorX m
+                       , MonadLogger m
+                       )
+                    => PostgresConf -> JobMode
+                    -> Maybe (m JobOffloadHandler)
+mkJobOffloadHandler dbConf jMode
+  | is _JobsLocal jMode, hasn't (_jobsAcceptOffload . only True) jMode = Nothing
+  | otherwise = Just $ do
+      jobOffloadOutgoing <- newTVarIO mempty
+      jobOffloadHandler <- allocateAsync . bracket (liftIO . PG.connectPostgreSQL $ pgConnStr dbConf) (liftIO . PG.close) $ \pgConn -> do
+        myPid <- liftIO $ PG.getBackendPID pgConn
+        let shouldListen = has (_jobsAcceptOffload . only True) jMode
+
+        when shouldListen $
+          void . liftIO $ PG.execute pgConn "LISTEN ?" (PG.Only $ PG.Identifier jobOffloadChannel)
+  
+        foreverBreak $ \(($ ()) -> terminate) -> do
+          UniWorX{appJobState} <- ask
+          shouldTerminate <- atomically $ readTMVar appJobState >>= fmap not . isEmptyTMVar . jobShutdown
+          when shouldTerminate terminate
+
+          let
+            getInput = do
+              n@PG.Notification{..} <- liftIO $ PG.getNotification pgConn
+              if | notificationPid == myPid || notificationChannel /= (encodeUtf8 jobOffloadChannel) -> getInput
+                 | otherwise -> return n
+            getOutput = atomically $ do
+              jQueue <- readTVar jobOffloadOutgoing
+              case jQueue of
+                j :< js -> j <$ writeTVar jobOffloadOutgoing js
+                _other  -> mzero
+
+          io <- lift $ if
+            | shouldListen -> getInput `race` getOutput
+            | otherwise    -> Right <$> getOutput
+
+          case io of
+            Left PG.Notification{..}
+              | Just jId <- fromPathPiece =<< either (const Nothing) Just (decodeUtf8' notificationData)
+                -> writeJobCtl $ JobCtlPerform jId
+              | otherwise
+                -> $logErrorS "JobOffloadHandler" $ "Could not parse incoming notification data: " <> tshow notificationData
+            Right jId -> void . liftIO $ PG.execute pgConn "NOTIFY ?, ?" (PG.Identifier jobOffloadChannel, encodeUtf8 $ toPathPiece jId)
+        
+      return JobOffloadHandler{..}

src/Jobs/Types.hs

@@ -9,7 +9,7 @@ module Jobs.Types
   , classifyJobCtl
   , YesodJobDB
   , JobHandler(..), _JobHandlerAtomic, _JobHandlerException
-  , JobContext(..)
+  , JobOffloadHandler(..), JobContext(..)
   , JobState(..), _jobWorkers, _jobWorkerName, _jobContext, _jobPoolManager, _jobCron, _jobShutdown, _jobCurrentCrontab
   , jobWorkerNames
   , JobWorkerState(..), _jobWorkerJobCtl, _jobWorkerJob
@@ -238,10 +238,16 @@ showWorkerId = tshow . hashUnique . jobWorkerUnique
 newWorkerId :: MonadIO m => m JobWorkerId
 newWorkerId = JobWorkerId <$> liftIO newUnique
 
+data JobOffloadHandler = JobOffloadHandler
+  { jobOffloadHandler :: Async ()
+  , jobOffloadOutgoing :: TVar (Seq QueuedJobId)
+  }
+
 data JobContext = JobContext
   { jobCrontab :: TVar (Crontab JobCtl)
   , jobConfirm :: TVar (HashMap JobCtl (NonEmpty (TMVar (Maybe SomeException))))
   , jobHeldLocks :: TVar (Set QueuedJobId)
+  , jobOffload :: TMVar JobOffloadHandler
   }
 
 

src/Settings.hs

@@ -206,9 +206,16 @@ data AppSettings = AppSettings
     , appInitialInstanceID      :: Maybe (Either FilePath UUID)
     , appRibbon                 :: Maybe Text
 
+    , appJobMode                :: JobMode
+
     , appMemcacheAuth           :: Bool
     } deriving Show
 
+data JobMode = JobsLocal { jobsAcceptOffload :: Bool }
+             | JobsOffload
+  deriving (Eq, Ord, Read, Show, Generic, Typeable)
+  deriving anyclass (Hashable)
+
 data ApprootScope = ApprootUserGenerated | ApprootDefault
   deriving (Eq, Ord, Read, Show, Enum, Bounded, Generic, Typeable)
   deriving anyclass (Universe, Finite, Hashable)
@@ -342,6 +349,11 @@ deriveFromJSON defaultOptions
   { fieldLabelModifier = camelToPathPiece' 2
   } ''UserDefaultConf
 
+deriveJSON defaultOptions
+  { fieldLabelModifier = camelToPathPiece' 1
+  , constructorTagModifier = camelToPathPiece' 1
+  } ''JobMode
+
 instance FromJSON LdapConf where
     parseJSON = withObject "LdapConf" $ \o -> do
         ldapTls <- o .:? "tls"
@@ -596,6 +608,8 @@ instance FromJSON AppSettings where
 
         appMemcacheAuth <- o .:? "memcache-auth" .!= False
 
+        appJobMode <- o .:? "job-mode" .!= JobsLocal True
+
         return AppSettings{..}
 
 makeClassy_ ''AppSettings

src/Utils.hs

@@ -70,6 +70,7 @@ import Control.Monad.Writer.Class (MonadWriter(..))
 import Control.Monad.Catch
 import Control.Monad.Morph (hoist)
 import Control.Monad.Fail
+import Control.Monad.Trans.Cont (ContT, evalContT, callCC)
 
 import Language.Haskell.TH
 import Language.Haskell.TH.Instances ()
@@ -943,6 +944,11 @@ forever' :: Monad m
          -> m b
 forever' start cont = cont start >>= flip forever' cont
 
+foreverBreak :: Monad m
+             => ((r -> ContT r m b) -> ContT r m a)
+             -> m r
+foreverBreak cont = evalContT . callCC $ forever . cont
+
 
 --------------
 -- Foldable --

src/Utils/Lens.hs

@@ -5,6 +5,7 @@
 module Utils.Lens ( module Utils.Lens ) where
 
 import Import.NoModel
+import Settings
 import Model
 import Model.Rating
 import qualified ClassyPrelude.Yesod as Yesod (HasHttpManager(..))
@@ -272,6 +273,9 @@ makePrisms ''AllocationPriority
 makePrisms ''RoomReference
 makeLenses_ ''RoomReference
 
+makePrisms ''JobMode
+makeLenses_ ''JobMode
+
 -- makeClassy_ ''Load
 
 --------------------------

src/Utils/Parameters.hs

@@ -32,6 +32,7 @@ data GlobalGetParam = GetLang
                     | GetDownload
                     | GetError
                     | GetSelectTable
+                    | GetGenerateToken
   deriving (Eq, Ord, Enum, Bounded, Read, Show, Generic)
   deriving anyclass (Universe, Finite)