71 lines
2.7 KiB
Haskell
71 lines
2.7 KiB
Haskell
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
|
|
| not shouldListen = 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
|
|
|
|
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{..}
|
|
where shouldListen = has (_jobsAcceptOffload . only True) jMode
|