diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright Li Meng Jun (c) 2017 + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Li Meng Jun nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,3 @@ +# periodic-server + +Periodic task system haskell server. diff --git a/app/periodicd.hs b/app/periodicd.hs new file mode 100644 --- /dev/null +++ b/app/periodicd.hs @@ -0,0 +1,125 @@ +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} + +module Main + ( main + ) where + + +import Control.Monad (when) +import Data.Maybe (fromMaybe) +import Data.String (fromString) +import Metro.SocketServer (socketServer) +import Metro.TP.TLS (makeServerParams', tlsConfig) +import Metro.TP.WebSockets (serverConfig) +import Metro.TP.XOR (xorConfig) +import Metro.Utils (setupLog) +import Periodic.Server (startServer) +import Periodic.Server.Persist (Persist, PersistConfig) +import Periodic.Server.Persist.PSQL (PSQL) +import Periodic.Server.Persist.SQLite (SQLite) +import System.Environment (getArgs, lookupEnv) +import System.Exit (exitSuccess) +import System.Log (Priority (..)) + +data Options = Options + { host :: String + , xorFile :: FilePath + , storePath :: FilePath + , useTls :: Bool + , useWs :: Bool + , certKey :: FilePath + , cert :: FilePath + , caStore :: FilePath + , showHelp :: Bool + , logLevel :: Priority + } + +options :: Maybe String -> Maybe String -> Options +options h f = Options + { host = fromMaybe "unix:///tmp/periodic.sock" h + , xorFile = fromMaybe "" f + , storePath = "data.sqlite" + , useTls = False + , useWs = False + , certKey = "server-key.pem" + , cert = "server.pem" + , caStore = "ca.pem" + , showHelp = False + , logLevel = ERROR + } + +parseOptions :: [String] -> Options -> Options +parseOptions [] opt = opt +parseOptions ("-H":x:xs) opt = parseOptions xs opt { host = x } +parseOptions ("--host":x:xs) opt = parseOptions xs opt { host = x } +parseOptions ("--xor":x:xs) opt = parseOptions xs opt { xorFile = x } +parseOptions ("-p":x:xs) opt = parseOptions xs opt { storePath = x } +parseOptions ("--path":x:xs) opt = parseOptions xs opt { storePath = x } +parseOptions ("-h":xs) opt = parseOptions xs opt { showHelp = True } +parseOptions ("--help":xs) opt = parseOptions xs opt { showHelp = True } +parseOptions ("--tls":xs) opt = parseOptions xs opt { useTls = True } +parseOptions ("--ws":xs) opt = parseOptions xs opt { useWs = True } +parseOptions ("--cert-key":x:xs) opt = parseOptions xs opt { certKey = x } +parseOptions ("--cert":x:xs) opt = parseOptions xs opt { cert = x } +parseOptions ("--ca":x:xs) opt = parseOptions xs opt { caStore = x } +parseOptions ("--log":x:xs) opt = parseOptions xs opt { logLevel = read x } +parseOptions (_:xs) opt = parseOptions xs opt + +printHelp :: IO () +printHelp = do + putStrLn "periodicd - Periodic task system server" + putStrLn "" + putStrLn "Usage: periodicd [--host|-H HOST] [--path|-p PATH] [--xor FILE|--ws|--tls [--hostname HOSTNAME] [--cert-key FILE] [--cert FILE] [--ca FILE]]" + putStrLn "" + putStrLn "Available options:" + putStrLn " -H --host Socket path [$PERIODIC_PORT]" + putStrLn " eg: tcp://:5000 (optional: unix:///tmp/periodic.sock) " + putStrLn " -p --path State store path (optional: data.sqlite)" + putStrLn " eg: file://data.sqlite" + putStrLn " eg: postgres://host='127.0.0.1' port=5432 dbname='periodicd' user='postgres' password=''" + putStrLn " --xor XOR Transport encode file [$XOR_FILE]" + putStrLn " --tls Use tls transport" + putStrLn " --ws Use websockets transport" + putStrLn " --cert-key Private key associated" + putStrLn " --cert Public certificate (X.509 format)" + putStrLn " --ca Server will use these certificates to validate clients" + putStrLn " --log Set log level DEBUG INFO NOTICE WARNING ERROR CRITICAL ALERT EMERGENCY (optional: ERROR)" + putStrLn " -h --help Display help message" + putStrLn "" + putStrLn "Version: v1.1.7.0" + putStrLn "" + exitSuccess + +main :: IO () +main = do + h <- lookupEnv "PERIODIC_PORT" + f <- lookupEnv "XOR_FILE" + + opts@Options {..} <- flip parseOptions (options h f) <$> getArgs + + when showHelp printHelp + + setupLog logLevel + + if take 11 storePath == "postgres://" then + run opts (fromString $ drop 11 storePath :: PersistConfig PSQL) + else if take 7 storePath == "file://" then + run opts (fromString $ drop 7 storePath :: PersistConfig SQLite) + else + run opts (fromString storePath :: PersistConfig SQLite) + +run :: Persist db => Options -> PersistConfig db -> IO () +run Options {useTls = True, ..} config = do + prms <- makeServerParams' cert [] certKey caStore + startServer config (tlsConfig prms) (socketServer host) + +run Options {useWs = True, host} config = + startServer config serverConfig (socketServer host) + +run Options {xorFile = "", host} config = + startServer config id (socketServer host) + +run Options {xorFile = f, host} config = + startServer config (xorConfig f) (socketServer host) diff --git a/periodic-server.cabal b/periodic-server.cabal new file mode 100644 --- /dev/null +++ b/periodic-server.cabal @@ -0,0 +1,71 @@ +name: periodic-server +version: 1.1.7.1 +synopsis: Periodic task system haskell server +description: Periodic task system haskell server and command peridoicd. +homepage: https://github.com/Lupino/haskell-periodic/tree/master/periodic-server#readme +license: BSD3 +license-file: LICENSE +author: Li Meng Jun +maintainer: lmjubuntu@gmail.com +copyright: MIT +category: System,Web +build-type: Simple +extra-source-files: README.md +cabal-version: >=1.10 + +library + hs-source-dirs: src + exposed-modules: Periodic.Server.Scheduler + , Periodic.Server.GrabQueue + , Periodic.Server.FuncStat + , Periodic.Server.Persist + , Periodic.Server.Persist.SQLite + , Periodic.Server.Persist.PSQL + , Periodic.Server.Client + , Periodic.Server.Types + , Periodic.Server + build-depends: base >= 4.7 && < 5 + , periodic-common + , bytestring + , network + , unordered-containers + , psqueues + , filepath + , hslogger + , byteable + , stm + , async + , binary + , transformers + , mtl + + , direct-sqlite + , unliftio + , metro + , entropy + + , postgresql-simple + , base64-bytestring + , resource-pool + default-language: Haskell2010 + +executable periodicd + hs-source-dirs: app + main-is: periodicd.hs + ghc-options: -threaded -rtsopts -with-rtsopts=-N + build-depends: base + , periodic-server + , periodic-common + , bytestring + , unliftio + , hslogger + , metro + , metro-socket + , metro-transport-xor + , metro-transport-tls + , metro-transport-websockets + default-language: Haskell2010 + +source-repository head + type: git + location: https://github.com/Lupino/haskell-periodic diff --git a/src/Periodic/Server.hs b/src/Periodic/Server.hs new file mode 100644 --- /dev/null +++ b/src/Periodic/Server.hs @@ -0,0 +1,68 @@ +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Periodic.Server + ( startServer + ) where + + +import Metro (NodeMode (..), SessionMode (..)) +import Metro.Class (Servable (STP), + Transport (TransportConfig)) +import qualified Metro.Class as S (Servable (ServerConfig)) +import Metro.Conn (receive, runConnT, send) +import Metro.Server (initServerEnv, runServerT, + setDefaultSessionTimeout, + setKeepalive, setNodeMode, + setOnNodeLeave, setServerName, + setSessionMode, stopServerT) +import qualified Metro.Server as M (ServerEnv, startServer) +import Periodic.IOList (newIOList, toList) +import Periodic.Node (sessionGen) +import Periodic.Server.Client (handleSessionT) +import Periodic.Server.Persist (Persist, PersistConfig) +import Periodic.Server.Scheduler (failJob, initSchedEnv, removeFunc, + runSchedT, shutdown, startSchedT) +import Periodic.Server.Types (ClientConfig (..), Command, + ServerCommand (Data)) +import Periodic.Types (ClientType, Msgid, Nid (..), Packet, + getClientType, regPacketRES) +import System.Entropy (getEntropy) +import UnliftIO (MonadUnliftIO) + +type ServerEnv serv = + M.ServerEnv serv ClientConfig Nid Msgid (Packet Command) + +startServer + :: (Servable serv, Transport tp, Persist db, MonadUnliftIO m) + => PersistConfig db + -> (TransportConfig (STP serv) -> TransportConfig tp) + -> S.ServerConfig serv + -> m () +startServer dbconfig mk config = do + sEnv <- fmap mapEnv . initServerEnv config sessionGen mk $ \_ connEnv -> do + (_ :: ClientType) <- getClientType <$> runConnT connEnv receive + nid <- getEntropy 4 + runConnT connEnv $ send (regPacketRES $ Data nid) + wFuncList <- newIOList + wJobQueue <- newIOList + return $ Just (Nid nid, ClientConfig {..}) + + schedEnv <- initSchedEnv dbconfig $ runServerT sEnv stopServerT + + setOnNodeLeave sEnv $ \_ ClientConfig {..} -> + runSchedT schedEnv $ do + mapM_ failJob =<< toList wJobQueue + mapM_ removeFunc =<< toList wFuncList + + runSchedT schedEnv $ do + startSchedT + M.startServer sEnv handleSessionT + shutdown + where mapEnv :: ServerEnv serv tp -> ServerEnv serv tp + mapEnv = + setNodeMode Multi + . setSessionMode SingleAction + . setKeepalive 500 + . setDefaultSessionTimeout 100 + . setServerName "Periodic" diff --git a/src/Periodic/Server/Client.hs b/src/Periodic/Server/Client.hs new file mode 100644 --- /dev/null +++ b/src/Periodic/Server/Client.hs @@ -0,0 +1,140 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} + +module Periodic.Server.Client + ( ClientT + , handleSessionT + ) where + +import Control.Monad (unless, when) +import Control.Monad.Trans.Class (lift) +import Data.Binary (encode) +import Data.Byteable (toBytes) +import qualified Data.ByteString.Char8 as B (intercalate) +import Data.ByteString.Lazy (toStrict) +import Metro.Class (Transport) +import Metro.Conn (fromConn) +import qualified Metro.Conn as Conn +import Metro.Session (env, getSessionEnv1, receive, + send) +import Periodic.IOList (delete, elem, insert) +import Periodic.Node +import Periodic.Server.Persist (Persist) +import Periodic.Server.Scheduler +import Periodic.Server.Types +import qualified Periodic.Types.ClientCommand as CC +import Periodic.Types.Internal (ConfigKey (..)) +import Periodic.Types.Job (getFuncName, initJob) +import Periodic.Types.Packet (getPacketData, packetRES) +import qualified Periodic.Types.WorkerCommand as WC +import Prelude hiding (elem) +import System.Log.Logger (errorM) +import UnliftIO + + +type ClientT db tp m = NodeT ClientConfig Command tp (SchedT db tp m) + +handleClientSessionT + :: (MonadUnliftIO m, Persist db, Transport tp) + => CC.ClientCommand -> SessionT ClientConfig Command tp (SchedT db tp m) () +handleClientSessionT (CC.SubmitJob job) = do + lift $ pushJob job + send $ packetRES Success +handleClientSessionT (CC.RunJob job) = do + preR <- lift $ lookupPrevResult job + case preR of + Just v -> send $ packetRES $ Data v + Nothing -> do + c <- lift . canRun $ getFuncName job + if c then do + lift $ prepareWait job + lift $ pushJob job + state <- fromConn Conn.statusTVar + w <- lift $ waitResult state job + send . packetRES $ Data w + else send $ packetRES NoWorker + +handleClientSessionT CC.Status = do + stats <- lift $ map toBytes <$> status + send . packetRES . Data $ B.intercalate "\n" stats + +handleClientSessionT CC.Ping = send $ packetRES Pong + +handleClientSessionT (CC.DropFunc fn) = do + lift $ dropFunc fn + send $ packetRES Success + +handleClientSessionT (CC.RemoveJob fn jn) = do + lift $ removeJob $ initJob fn jn + send $ packetRES Success +handleClientSessionT CC.Shutdown = lift shutdown + +handleClientSessionT (CC.ConfigGet (ConfigKey key)) = do + v <- lift $ getConfigInt key + send $ packetRES $ Config v + +handleClientSessionT (CC.ConfigSet (ConfigKey key) v) = do + lift $ setConfigInt key v + send $ packetRES Success + +handleClientSessionT CC.Dump = send =<< lift (packetRES . Data . toStrict . encode <$> dumpJob) + +handleClientSessionT (CC.Load jobs) = do + lift $ mapM_ pushJob jobs + send $ packetRES Success + +handleWorkerSessionT + :: (MonadUnliftIO m, Persist db, Transport tp) + => ClientConfig -> WC.WorkerCommand -> SessionT ClientConfig Command tp (SchedT db tp m) () +handleWorkerSessionT ClientConfig {..} WC.GrabJob = do + env0 <- getSessionEnv1 + lift $ pushGrab wFuncList wJobQueue env0 +handleWorkerSessionT ClientConfig {..} (WC.WorkDone jh w) = do + lift $ doneJob jh w + delete wJobQueue jh +handleWorkerSessionT ClientConfig {..} (WC.WorkFail jh) = do + lift $ failJob jh + delete wJobQueue jh +handleWorkerSessionT ClientConfig {..} (WC.SchedLater jh l s) = do + lift $ schedLaterJob jh l s + delete wJobQueue jh +handleWorkerSessionT ClientConfig {..} WC.Sleep = send $ packetRES Noop +handleWorkerSessionT ClientConfig {..} WC.Ping = send $ packetRES Pong +handleWorkerSessionT ClientConfig {..} (WC.CanDo fn) = do + has <- elem wFuncList fn + unless has $ do + lift $ addFunc fn + insert wFuncList fn +handleWorkerSessionT ClientConfig {..} (WC.CantDo fn) = do + has <- elem wFuncList fn + when has $ do + lift $ removeFunc fn + delete wFuncList fn +handleWorkerSessionT ClientConfig {..} (WC.Broadcast fn) = do + has <- elem wFuncList fn + unless has $ do + lift $ broadcastFunc fn True + insert wFuncList fn +handleWorkerSessionT _ (WC.Acquire n c jh) = do + r <- lift $ acquireLock n c jh + send $ packetRES $ Acquired r +handleWorkerSessionT _ (WC.Release n jh) = lift $ releaseLock n jh + +handleSessionT + :: (MonadUnliftIO m, Persist db, Transport tp) + => SessionT ClientConfig Command tp (SchedT db tp m) () +handleSessionT = do + mcmd <- receive + case mcmd of + Nothing -> do + liftIO $ errorM "Periodic.Server.Client" "Client error" + fromConn Conn.close -- close client + Just pkt -> + case getPacketData pkt of + CC cmd -> handleClientSessionT cmd + WC cmd -> do + env0 <- env + handleWorkerSessionT env0 cmd diff --git a/src/Periodic/Server/FuncStat.hs b/src/Periodic/Server/FuncStat.hs new file mode 100644 --- /dev/null +++ b/src/Periodic/Server/FuncStat.hs @@ -0,0 +1,40 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} + +module Periodic.Server.FuncStat + ( FuncStat (..) + , funcStat + , FuncStatList + ) where + +import Data.Byteable +import qualified Data.ByteString.Char8 as B (intercalate, pack) +import Data.Int (Int64) +import Metro.IOHashMap (IOHashMap) +import Periodic.Types (FuncName (..)) + +data FuncStat = FuncStat + { sSchedAt :: Int64 + , sWorker :: Int64 + , sJob :: Int64 + , sRunning :: Int64 + , sLocking :: Int64 + , sFuncName :: FuncName + , sBroadcast :: Bool + } + +instance Byteable FuncStat where + toBytes FuncStat{..} = B.intercalate "," + [ unFN sFuncName + , B.pack $ show sWorker + , B.pack $ show sJob + , B.pack $ show sRunning + , B.pack $ show sLocking + , B.pack $ show sSchedAt + ] + +type FuncStatList = IOHashMap FuncName FuncStat + +funcStat :: FuncName -> FuncStat +funcStat sFuncName = FuncStat + { sSchedAt = 0, sWorker = 0, sJob = 0, sRunning = 0, sLocking = 0, sBroadcast = False, .. } diff --git a/src/Periodic/Server/GrabQueue.hs b/src/Periodic/Server/GrabQueue.hs new file mode 100644 --- /dev/null +++ b/src/Periodic/Server/GrabQueue.hs @@ -0,0 +1,75 @@ +{-# LANGUAGE RecordWildCards #-} + +module Periodic.Server.GrabQueue + ( GrabQueue + , newGrabQueue + , GrabItem (..) + , pushAgent + , popAgentSTM + , popAgentList + ) where + +import Control.Arrow ((&&&)) +import Control.Monad (unless) +import Control.Monad.STM (STM, retry) +import Metro.Session (ident) +import Periodic.IOList (IOList, append, delete, deleteSTM, elem, + elemSTM, newIOList, toList, toListSTM) +import Periodic.Node (Nid) +import Periodic.Server.Types (CSEnv) +import Periodic.Types (FuncName, JobHandle, Msgid) +import Prelude hiding (elem) +import UnliftIO (MonadIO (..)) + +data GrabItem tp = GrabItem + { gFuncList :: IOList FuncName + , gAgent :: CSEnv tp + , gJobQueue :: IOList JobHandle + } + +instance Eq (GrabItem tp) where + (==) = eqGrabItem + +key :: GrabItem tp -> (Nid, Msgid) +key GrabItem{gAgent = a} = ident a + +eqGrabItem :: GrabItem tp -> GrabItem tp -> Bool +eqGrabItem a b = key a == key b + +type GrabQueue tp = IOList (GrabItem tp) + +newGrabQueue :: MonadIO m => m (GrabQueue tp) +newGrabQueue = newIOList + +pushAgent :: MonadIO m => GrabQueue tp -> IOList FuncName -> IOList JobHandle -> CSEnv tp -> m () +pushAgent q gFuncList gJobQueue gAgent = do + has <- elem q i + unless has $ append q i + where i = GrabItem {..} + +popAgentSTM :: GrabQueue tp -> FuncName -> STM (IOList JobHandle, CSEnv tp) +popAgentSTM q n = do + item <- go =<< toListSTM q + deleteSTM q item + return (gJobQueue item, gAgent item) + + where go :: [GrabItem tp] -> STM (GrabItem tp) + go [] = retry + go (x:xs) = do + has <- elemSTM (gFuncList x) n + if has then return x + else go xs + +popAgentList :: MonadIO m => GrabQueue tp -> FuncName -> m [(IOList JobHandle, CSEnv tp)] +popAgentList q n = do + items <- go =<< toList q + mapM_ (delete q) items + pure $ map (gJobQueue &&& gAgent) items + + where go :: MonadIO m => [GrabItem tp] -> m [GrabItem tp] + go [] = return [] + go (x:xs) = do + has <- elem (gFuncList x) n + xs' <- go xs + if has then pure (x:xs') + else pure xs' diff --git a/src/Periodic/Server/Persist.hs b/src/Periodic/Server/Persist.hs new file mode 100644 --- /dev/null +++ b/src/Periodic/Server/Persist.hs @@ -0,0 +1,39 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeFamilies #-} + +module Periodic.Server.Persist + ( Persist (..) + , State (..) + ) where + +import Control.Exception (Exception) +import Data.Int (Int64) +import Periodic.Types.Job (FuncName, Job, JobName) +import Prelude hiding (foldr, lookup) + +data State = Pending + | Running + | Locking + +class (Exception (PersistException db)) => Persist db where + data PersistConfig db + data PersistException db + newPersist :: PersistConfig db -> IO db + member :: db -> State -> FuncName -> JobName -> IO Bool + lookup :: db -> State -> FuncName -> JobName -> IO (Maybe Job) + insert :: db -> State -> FuncName -> JobName -> Job -> IO () + delete :: db -> FuncName -> JobName -> IO () + size :: db -> State -> FuncName -> IO Int64 + foldr :: forall a . db -> State -> (Job -> a -> a) -> a -> IO a + foldrPending :: forall a . db -> Int64 -> [FuncName] -> (Job -> a -> a) -> a -> IO a + foldrLocking :: forall a . db -> Int -> FuncName -> (Job -> a -> a) -> a -> IO a + dumpJob :: db -> IO [Job] + configGet :: db -> String -> IO (Maybe Int) + configSet :: db -> String -> Int -> IO () + insertFuncName :: db -> FuncName -> IO () + removeFuncName :: db -> FuncName -> IO () + funcList :: db -> IO [FuncName] + minSchedAt :: db -> FuncName -> IO Int64 diff --git a/src/Periodic/Server/Persist/PSQL.hs b/src/Periodic/Server/Persist/PSQL.hs new file mode 100644 --- /dev/null +++ b/src/Periodic/Server/Persist/PSQL.hs @@ -0,0 +1,324 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE TypeFamilies #-} + +module Periodic.Server.Persist.PSQL + ( PSQL + ) where + +import Control.Monad (void) +import Data.Binary (decodeOrFail) +import Data.Byteable (toBytes) +import Data.ByteString (ByteString) +import Data.ByteString.Base64 (decode, encode) +import Data.ByteString.Lazy (fromStrict) +import qualified Data.Foldable as F (foldrM) +import Data.Int (Int64) +import Data.List (intercalate) +import Data.Maybe (fromMaybe, isJust, listToMaybe) +import Data.Pool (Pool, createPool, withResource) +import Data.String (IsString (..)) +import Database.PostgreSQL.Simple +import Periodic.Server.Persist (Persist (PersistConfig, PersistException), + State (..)) +import qualified Periodic.Server.Persist as Persist +import Periodic.Types.Job (FuncName (..), Job, JobName (..), + getSchedAt) +import Prelude hiding (foldr, lookup) +import System.Log.Logger (errorM, infoM) +import UnliftIO (Exception, SomeException, Typeable) + +stateName :: State -> ByteString +stateName Pending = "0" +stateName Running = "1" +stateName Locking = "2" + +newtype PSQL = PSQL (Pool Connection) + +numStripes = 1 +idleTime = 10 +maxResources = 10 + +instance Persist PSQL where + data PersistConfig PSQL = PSQLPath ByteString + data PersistException PSQL = PSQLException SomeException deriving (Show, Typeable) + + newPersist (PSQLPath path) = do + infoM "Periodic.Server.Persist.PSQL" ("PSQL connected " ++ show path) + pool <- createPool (connectPostgreSQL path) close numStripes idleTime maxResources + withResource pool $ \conn -> + withTransaction conn $ do + createConfigTable conn + createJobTable conn + createFuncTable conn + allPending pool + return $ PSQL pool + + member (PSQL pool) = doMember pool + lookup (PSQL pool) = doLookup pool + insert (PSQL pool) = doInsert pool + delete (PSQL pool) = doDelete pool + size (PSQL pool) = doSize pool + foldr (PSQL pool) = doFoldr pool + foldrPending (PSQL pool) = doFoldrPending pool + foldrLocking (PSQL pool) = doFoldrLocking pool + dumpJob (PSQL pool) = doDumpJob pool + configSet (PSQL pool) = doConfigSet pool + configGet (PSQL pool) = doConfigGet pool + insertFuncName (PSQL pool) = doInsertFuncName pool + removeFuncName (PSQL pool) = doRemoveFuncName pool + funcList (PSQL pool) = doFuncList pool + minSchedAt (PSQL pool) = doMinSchedAt pool Pending + +instance Exception (PersistException PSQL) + +instance IsString (PersistConfig PSQL) where + fromString = PSQLPath . fromString + +newtype TableName = TableName String + deriving (Show) + +instance IsString TableName where + fromString = TableName + +getTableName :: TableName -> String +getTableName (TableName name) = name + +newtype Column = Column { unColumn :: String } + deriving (Show) + +instance IsString Column where + fromString = Column + +type Columns = [Column] + +columnsToString :: Columns -> String +columnsToString = intercalate ", " . map unColumn + +createTable :: TableName -> Columns -> Connection -> IO Int64 +createTable tn cols conn = execute_ conn sql + where sql = fromString $ concat + [ "CREATE TABLE IF NOT EXISTS ", getTableName tn, " (" + , columnsToString cols + , ")" + ] + +newtype IndexName = IndexName String + deriving (Show) + +instance IsString IndexName where + fromString = IndexName + +getOnlyDefault :: FromRow (Only a) => a -> [Only a] -> a +getOnlyDefault a = maybe a fromOnly . listToMaybe + +insertOrUpdate :: ToRow a => TableName -> Columns -> Columns -> a -> Pool Connection -> IO Int64 +insertOrUpdate tn ucols vcols a pool = withResource pool $ \conn -> execute conn sql a + where cols = ucols ++ vcols + v = replicate (length cols) "?" + + setSql = intercalate ", " $ map appendSet vcols + + appendSet :: Column -> String + appendSet (Column col) | '=' `elem` col = col + | otherwise = col ++ " = excluded." ++ col + + doSql = if null vcols then " DO NOTHING" else " DO UPDATE SET " ++ setSql + + sql = fromString $ concat + [ "INSERT INTO ", getTableName tn + , " (", columnsToString cols, ")" + , " VALUES" + , " (", columnsToString v, ")" + , " ON CONFLICT (", columnsToString ucols, ")" + , doSql + ] + +update :: ToRow a => TableName -> Columns -> String -> a -> Pool Connection -> IO Int64 +update tn cols partSql a pool = withResource pool $ \conn -> execute conn sql a + where setSql = intercalate ", " $ map appendSet cols + whereSql = if null partSql then "" else " WHERE " ++ partSql + sql = fromString $ concat + [ "UPDATE ", getTableName tn + , " SET ", setSql + , whereSql + ] + + appendSet :: Column -> String + appendSet (Column col) | '=' `elem` col = col + | otherwise = col ++ " = ?" + + +selectOne + :: (ToRow a, FromRow b, Show b) + => TableName -> Columns -> String -> a -> Connection -> IO (Maybe b) +selectOne tn cols partSql a conn = listToMaybe <$> query conn sql a + where whereSql = " WHERE " ++ partSql + sql = fromString $ concat + [ "SELECT ", columnsToString cols, " FROM ", getTableName tn + , whereSql + ] + +selectOneOnly + :: (ToRow a, FromRow (Only b), Show b) + => TableName -> Column -> String -> a -> Pool Connection -> IO (Maybe b) +selectOneOnly tn col partSql a pool = + withResource pool $ fmap (fmap fromOnly) . selectOne tn [col] partSql a + +count :: ToRow a => TableName -> String -> a -> Pool Connection -> IO Int64 +count tn partSql a pool = withResource pool $ \conn -> + getOnlyDefault 0 <$> query conn sql a + where whereSql = " WHERE " ++ partSql + sql = fromString $ concat + [ "SELECT count(*) FROM ", getTableName tn, whereSql + ] + +delete :: ToRow a => TableName -> String -> a -> Pool Connection -> IO Int64 +delete tn partSql a pool = withResource pool $ \conn -> execute conn sql a + where whereSql = " WHERE " ++ partSql + sql = fromString $ concat + [ "DELETE FROM ", getTableName tn, whereSql + ] + + +configs :: TableName +configs = "configs" + +jobs :: TableName +jobs = "jobs" + +funcs :: TableName +funcs = "funcs" + +createConfigTable :: Connection -> IO () +createConfigTable = + void . createTable configs + [ "name VARCHAR(256) NOT NULL" + , "value INT DEFAULT 0" + , "CONSTRAINT config_pk PRIMARY KEY (name)" + ] + +createJobTable :: Connection -> IO () +createJobTable = + void . createTable jobs + [ "func VARCHAR(256) NOT NULL" + , "name VARCHAR(256) NOT NULL" + , "value text" + , "state INT DEFAULT 0" + , "sched_at INT DEFAULT 0" + , "CONSTRAINT job_pk PRIMARY KEY (func, name)" + ] + +createFuncTable :: Connection -> IO () +createFuncTable = + void . createTable funcs + [ "func VARCHAR(256) NOT NULL" + , "CONSTRAINT func_pk PRIMARY KEY (func)" + ] + +allPending :: Pool Connection -> IO () +allPending = void . update jobs ["state"] "" (Only (stateName Pending)) + +doLookup :: Pool Connection -> State -> FuncName -> JobName -> IO (Maybe Job) +doLookup pool state fn jn = do + r <- selectOneOnly jobs "value" "func=? AND name=? AND state=?" + (unFN fn, unJN jn, stateName state) pool + case r of + Nothing -> return Nothing + Just bs -> + case decodeJob bs of + Left e -> do + errorM "Periodic.Server.Persist.PSQL" + $ "doLookup error: decode " ++ show bs ++ " " ++ show e + return Nothing + Right job -> return $ Just job + +doMember :: Pool Connection -> State -> FuncName -> JobName -> IO Bool +doMember pool st fn jn = isJust <$> doLookup pool st fn jn + +doInsert :: Pool Connection -> State -> FuncName -> JobName -> Job -> IO () +doInsert pool state fn jn job = do + doInsertFuncName pool fn + void $ insertOrUpdate jobs + ["func", "name"] + ["value", "state", "sched_at"] + (unFN fn, unJN jn, encode $ toBytes job, stateName state, getSchedAt job) + pool + +doInsertFuncName :: Pool Connection -> FuncName -> IO () +doInsertFuncName pool fn = void $ insertOrUpdate funcs ["func"] [] (Only $ unFN fn) pool + +doFoldr :: Pool Connection -> State -> (Job -> a -> a) -> a -> IO a +doFoldr pool state f acc = withResource pool $ \conn -> + fold conn sql (Only $ stateName state) acc (mkFoldFunc f) + where sql = fromString $ "SELECT value FROM " ++ getTableName jobs ++ " WHERE state=?" + +doFoldrPending :: Pool Connection -> Int64 -> [FuncName] -> (Job -> a -> a) -> a -> IO a +doFoldrPending pool ts fns f acc = withResource pool $ \conn -> + F.foldrM (foldFunc conn f) acc fns + + where sql = fromString $ "SELECT value FROM " + ++ getTableName jobs ++ " WHERE func=? AND state=? AND sched_at < ?" + foldFunc :: Connection -> (Job -> a -> a) -> FuncName -> a -> IO a + foldFunc conn f0 fn acc0 = + fold conn sql (unFN fn, stateName Pending, ts) acc0 (mkFoldFunc f0) + +doFoldrLocking :: Pool Connection -> Int -> FuncName -> (Job -> a -> a) -> a -> IO a +doFoldrLocking pool limit fn f acc = withResource pool $ \conn -> + fold conn sql (unFN fn, stateName Locking, limit) acc (mkFoldFunc f) + where sql = fromString $ "SELECT value FROM " ++ getTableName jobs + ++ " WHERE func=? AND state=? ORDER BY sched_at ASC LIMIT ?" + +doDumpJob :: Pool Connection -> IO [Job] +doDumpJob pool = withResource pool $ \conn -> + fold_ conn sql [] (mkFoldFunc (:)) + where sql = fromString $ "SELECT value FROM " ++ getTableName jobs + +doFuncList :: Pool Connection -> IO [FuncName] +doFuncList pool = withResource pool $ \conn -> + map (FuncName . fromOnly) <$> query_ conn sql + where sql = fromString $ "SELECT func FROM " ++ getTableName funcs + +doDelete :: Pool Connection -> FuncName -> JobName -> IO () +doDelete pool fn jn = void $ delete jobs "func=? AND name=?" (unFN fn, unJN jn) pool + +doRemoveFuncName :: Pool Connection -> FuncName -> IO () +doRemoveFuncName pool fn = do + void $ delete jobs "func=?" (Only $ unFN fn) pool + void $ delete funcs "func=?" (Only $ unFN fn) pool + +doMinSchedAt :: Pool Connection -> State -> FuncName -> IO Int64 +doMinSchedAt pool state fn = + fromMaybe 0 + . fromMaybe Nothing + <$> selectOneOnly jobs "min(sched_at)" "func=? AND state=?" + (unFN fn, stateName state) + pool + +doSize :: Pool Connection -> State -> FuncName -> IO Int64 +doSize pool state fn = count jobs "func=? AND state=?" (unFN fn, stateName state) pool + +doConfigSet :: Pool Connection -> String -> Int -> IO () +doConfigSet pool name v = + void $ insertOrUpdate configs ["name"] ["value"] (name, v) pool + +doConfigGet :: Pool Connection -> String -> IO (Maybe Int) +doConfigGet pool name = selectOneOnly configs "value" "name=?" (Only name) pool + +mkFoldFunc :: (Job -> a -> a) -> a -> Only ByteString -> IO a +mkFoldFunc f acc (Only bs) = + case decodeJob bs of + Left e -> do + errorM "Periodic.Server.Persist.PSQL" $ "mkFoldFunc error: decode " ++ show bs ++ " " ++ show e + return acc + Right job -> return $ f job acc + +decodeJob :: ByteString -> Either String Job +decodeJob bs = + case decode bs of + Left e -> Left e + Right bs0 -> + case decodeOrFail (fromStrict bs0) of + Left e -> Left $ show e + Right (_, _, job) -> Right job diff --git a/src/Periodic/Server/Persist/SQLite.hs b/src/Periodic/Server/Persist/SQLite.hs new file mode 100644 --- /dev/null +++ b/src/Periodic/Server/Persist/SQLite.hs @@ -0,0 +1,274 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} + +module Periodic.Server.Persist.SQLite + ( SQLite + ) where + +import Control.Monad (void) +import Data.Binary (decodeOrFail) +import Data.Byteable (toBytes) +import Data.ByteString (ByteString, append) +import qualified Data.ByteString.Char8 as B (pack) +import Data.ByteString.Lazy (fromStrict) +import qualified Data.Foldable as F (foldrM) +import Data.Int (Int64) +import Data.Maybe (isJust, listToMaybe) +import Data.String (IsString (..)) +import Database.SQLite3.Direct +import Periodic.Server.Persist +import Periodic.Types.Job (FuncName (..), Job, JobName (..), + getSchedAt) +import Prelude hiding (foldr, lookup) +import System.Log.Logger (infoM) +import UnliftIO (Exception, Typeable, throwIO) + +stateName :: State -> ByteString +stateName Pending = "0" +stateName Running = "1" +stateName Locking = "2" + +stateName' :: State -> Int64 +stateName' Pending = 0 +stateName' Running = 1 +stateName' Locking = 2 + +newtype SQLite = SQLite Database + +instance Persist SQLite where + data PersistConfig SQLite = SQLitePath Utf8 + data PersistException SQLite = SQLiteException Error deriving (Eq, Show, Typeable) + + newPersist (SQLitePath path) = do + infoM "Periodic.Server.Persist.SQLite" ("SQLite connected " ++ show path) + edb <- open path + case edb of + Left (e, _) -> throwIO $ SQLiteException e + Right db -> do + beginTx db + createConfigTable db + createJobTable db + createFuncTable db + allPending db + commitTx db + return $ SQLite db + + member (SQLite db) = doMember db + lookup (SQLite db) = doLookup db + insert (SQLite db) = doInsert db + delete (SQLite db) = doDelete db + size (SQLite db) = doSize db + foldr (SQLite db) = doFoldr db + foldrPending (SQLite db) = doFoldrPending db + foldrLocking (SQLite db) = doFoldrLocking db + dumpJob (SQLite db) = doDumpJob db + configSet (SQLite db) = doConfigSet db + configGet (SQLite db) = doConfigGet db + insertFuncName (SQLite db) = doInsertFuncName db + removeFuncName (SQLite db) = doRemoveFuncName db + funcList (SQLite db) = doFuncList db + minSchedAt (SQLite db) = doMinSchedAt db Pending + +instance Exception (PersistException SQLite) + +instance IsString (PersistConfig SQLite) where + fromString = SQLitePath . fromString + +beginTx :: Database -> IO () +beginTx db = void $ exec db "BEGIN TRANSACTION" + +commitTx :: Database -> IO () +commitTx db = void $ exec db "COMMIT TRANSACTION" + +rollbackTx :: Database -> IO () +rollbackTx db = void $ exec db "ROLLBACK TRANSACTION" + +createConfigTable :: Database -> IO () +createConfigTable db = void . exec db $ Utf8 $ + "CREATE TABLE IF NOT EXISTS configs (" + `append` "name CHAR(256) NOT NULL," + `append` "value INTEGER DEFAULT 0," + `append` "PRIMARY KEY (name))" + +createJobTable :: Database -> IO () +createJobTable db = void . exec db $ Utf8 $ + "CREATE TABLE IF NOT EXISTS jobs (" + `append` " func CHAR(256) NOT NULL," + `append` " name CHAR(256) NOT NULL," + `append` " value BLOB," + `append` " state INTEGER DEFAULT 0," + `append` " sched_at INTEGER DEFAULT 0," + `append` " PRIMARY KEY (func, name))" + +createFuncTable :: Database -> IO () +createFuncTable db = void . exec db $ Utf8 $ + "CREATE TABLE IF NOT EXISTS funcs (" + `append` " func CHAR(256) NOT NULL," + `append` " PRIMARY KEY (func))" + +allPending :: Database -> IO () +allPending db = void . exec db $ Utf8 "UPDATE jobs SET state=0" + +doLookup :: Database -> State -> FuncName -> JobName -> IO (Maybe Job) +doLookup db state fn jn = + listToMaybe <$> doFoldr_ db sql (bindFnAndJn fn jn) (mkFoldFunc f) [] + where sql = Utf8 $ "SELECT value FROM jobs WHERE func=? AND name=? AND state=" `append` stateName state `append` " LIMIT 1" + f :: Job -> [Job] -> [Job] + f job acc = job : acc + +doMember :: Database -> State -> FuncName -> JobName -> IO Bool +doMember db st fn jn = isJust <$> doLookup db st fn jn + +doInsert :: Database -> State -> FuncName -> JobName -> Job -> IO () +doInsert db state fn jn job = do + execStmt db sql $ \stmt -> do + bindFnAndJn fn jn stmt + void $ bindBlob stmt 3 $ toBytes job + void $ bindInt64 stmt 4 $ stateName' state + void $ bindInt64 stmt 5 $ getSchedAt job + doInsertFuncName db fn + where sql = Utf8 "INSERT OR REPLACE INTO jobs VALUES (?, ?, ?, ?, ?)" + +doInsertFuncName :: Database -> FuncName -> IO () +doInsertFuncName db = execFN db sql + where sql = Utf8 "INSERT OR REPLACE INTO funcs VALUES (?)" + +doFoldr :: Database -> State -> (Job -> a -> a) -> a -> IO a +doFoldr db state f = doFoldr_ db sql (const $ pure ()) (mkFoldFunc f) + where sql = Utf8 $ "SELECT value FROM jobs WHERE state=" `append` stateName state + +doFoldrPending :: Database -> Int64 -> [FuncName] -> (Job -> a -> a) -> a -> IO a +doFoldrPending db ts fns f acc = F.foldrM (foldFunc f) acc fns + where sql = Utf8 $ "SELECT value FROM jobs WHERE func=? AND state=" + <> stateName Pending + <> " AND sched_at<" + <> B.pack (show ts) + + foldFunc :: (Job -> a -> a) -> FuncName -> a -> IO a + foldFunc f0 fn = + doFoldr_ db sql (`bindFN` fn) (mkFoldFunc f0) + +doFoldrLocking :: Database -> Int -> FuncName -> (Job -> a -> a) -> a -> IO a +doFoldrLocking db limit fn f = doFoldr_ db sql (`bindFN` fn) (mkFoldFunc f) + where sql = Utf8 $ "SELECT value FROM jobs WHERE func=? AND state=" + <> stateName Locking + <> " ORDER BY sched_at ASC LIMIT " <> B.pack (show limit) + +doDumpJob :: Database -> IO [Job] +doDumpJob db = doFoldr_ db sql (const $ pure ()) (mkFoldFunc (:)) [] + where sql = Utf8 "SELECT value FROM jobs" + +doFuncList :: Database -> IO [FuncName] +doFuncList db = + doFoldr_ db sql (const $ pure ()) (\fn acc -> FuncName fn : acc) [] + where sql = Utf8 "SELECT func FROM funcs" + +doDelete :: Database -> FuncName -> JobName -> IO () +doDelete db fn jn = execStmt db sql $ bindFnAndJn fn jn + where sql = Utf8 "DELETE FROM jobs WHERE func=? AND name=?" + +doRemoveFuncName :: Database -> FuncName -> IO () +doRemoveFuncName db fn = do + execFN db sql0 fn + execFN db sql1 fn + + where sql0 = Utf8 "DELETE FROM funcs WHERE func=?" + sql1 = Utf8 "DELETE FROM jobs WHERE func=?" + +doMinSchedAt :: Database -> State -> FuncName -> IO Int64 +doMinSchedAt db state fn = queryStmt db sql (`bindFN` fn) stepInt64 + where sql = Utf8 $ "SELECT sched_at FROM jobs WHERE func=? AND state=" `append` stateName state `append` " ORDER BY sched_at ASC LIMIT 1" + +doSize :: Database -> State -> FuncName -> IO Int64 +doSize db state fn = queryStmt db sql (`bindFN` fn) stepInt64 + where sql = Utf8 $ "SELECT COUNT(*) FROM jobs WHERE func=? AND state=" `append` stateName state + +doConfigSet :: Database -> String -> Int -> IO () +doConfigSet db name v = execStmt db sql $ \stmt -> do + void $ bindText stmt 1 $ fromString name + void $ bindInt64 stmt 2 $ fromIntegral v + where sql = Utf8 "INSERT OR REPLACE INTO configs VALUES (?,?)" + +doConfigGet :: Database -> String -> IO (Maybe Int) +doConfigGet db name = queryStmt db sql (\stmt -> void $ bindText stmt 1 $ fromString name) stepMaybeInt + where sql = Utf8 "SELECT value FROM configs WHERE name=?" + +dbError :: String -> IO a +dbError = throwIO . userError . ("Database error: " ++) + +liftEither :: Show a => IO (Either a b) -> IO b +liftEither a = do + er <- a + case er of + (Left e) -> dbError (show e) + (Right r) -> return r +{-# INLINE liftEither #-} + +prepStmt :: Database -> Utf8 -> IO Statement +prepStmt c q = do + r <- prepare c q + case r of + Left e -> dbError (show e) + Right Nothing -> dbError "Statement prep failed" + Right (Just s) -> return s + +bindFN :: Statement -> FuncName -> IO () +bindFN stmt (FuncName fn) = void $ bindBlob stmt 1 fn + +execStmt :: Database -> Utf8 -> (Statement -> IO ()) -> IO () +execStmt db sql bindStmt = do + stmt <- prepStmt db sql + bindStmt stmt + void $ liftEither $ step stmt + void $ finalize stmt + +execFN :: Database -> Utf8 -> FuncName -> IO () +execFN db sql fn = execStmt db sql (`bindFN` fn) + +queryStmt :: Database -> Utf8 -> (Statement -> IO ()) -> (Statement -> IO a) -> IO a +queryStmt db sql bindStmt stepStmt = do + stmt <- prepStmt db sql + bindStmt stmt + ret <- stepStmt stmt + void $ finalize stmt + pure ret + +bindFnAndJn :: FuncName -> JobName -> Statement -> IO () +bindFnAndJn fn (JobName jn) stmt = do + bindFN stmt fn + void $ bindBlob stmt 2 jn + +mkFoldFunc :: (Job -> a -> a) -> ByteString -> a -> a +mkFoldFunc f bs acc = + case decodeOrFail (fromStrict bs) of + Left _ -> acc + Right (_, _, job) -> f job acc + +foldStmt :: (ByteString -> a -> a) -> a -> Statement -> IO a +foldStmt f acc stmt = do + sr <- liftEither $ step stmt + case sr of + Done -> pure acc + Row -> do + bs <- columnBlob stmt 0 + foldStmt f (f bs acc) stmt + +doFoldr_ :: Database -> Utf8 -> (Statement -> IO ()) -> (ByteString -> a -> a) -> a -> IO a +doFoldr_ db sql bindStmt f acc = queryStmt db sql bindStmt $ foldStmt f acc + +stepInt64 :: Statement -> IO Int64 +stepInt64 stmt = do + sr <- liftEither $ step stmt + + case sr of + Done -> pure 0 + Row -> columnInt64 stmt 0 + +stepMaybeInt :: Statement -> IO (Maybe Int) +stepMaybeInt stmt = do + sr <- liftEither $ step stmt + case sr of + Done -> pure Nothing + Row -> Just . fromIntegral <$> columnInt64 stmt 0 diff --git a/src/Periodic/Server/Scheduler.hs b/src/Periodic/Server/Scheduler.hs new file mode 100644 --- /dev/null +++ b/src/Periodic/Server/Scheduler.hs @@ -0,0 +1,913 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} + +module Periodic.Server.Scheduler + ( SchedT + , runSchedT + , SchedEnv + , initSchedEnv + , startSchedT + , pushJob + , pushGrab + , failJob + , doneJob + , schedLaterJob + , acquireLock + , releaseLock + , addFunc + , removeFunc + , broadcastFunc + , dropFunc + , removeJob + , dumpJob + , status + , shutdown + , keepalive + , setConfigInt + , getConfigInt + , prepareWait + , lookupPrevResult + , waitResult + , canRun + ) where + +import Control.Monad (forever, mzero, unless, void, + when) +import Control.Monad.Reader.Class (MonadReader (ask), asks) +import Control.Monad.Trans.Class (MonadTrans, lift) +import Control.Monad.Trans.Maybe (runMaybeT) +import Control.Monad.Trans.Reader (ReaderT (..), runReaderT) +import Data.ByteString (ByteString) +import Data.Foldable (forM_) +import Data.HashPSQ (HashPSQ) +import qualified Data.HashPSQ as PSQ +import Data.Int (Int64) +import qualified Data.List as L (delete) +import Data.Maybe (fromJust, fromMaybe, isJust) +import Metro.Class (Transport) +import Metro.IOHashMap (IOHashMap, newIOHashMap) +import qualified Metro.IOHashMap as FL +import qualified Metro.Lock as L (Lock, new, with) +import Metro.Session (runSessionT1, send, sessionState) +import Metro.Utils (getEpochTime) +import Periodic.IOList (IOList) +import qualified Periodic.IOList as IL +import Periodic.Server.FuncStat +import Periodic.Server.GrabQueue +import Periodic.Server.Persist (Persist, State (..)) +import qualified Periodic.Server.Persist as P +import Periodic.Server.Types (CSEnv) +import Periodic.Types (packetRES) +import Periodic.Types.Internal (LockName) +import Periodic.Types.Job +import Periodic.Types.ServerCommand (ServerCommand (JobAssign)) +import System.Log.Logger (errorM, infoM) +import UnliftIO +import UnliftIO.Concurrent (threadDelay) + +data Action = Add Job + | Remove Job + | Cancel + | PollJob + | TryPoll JobHandle + +data WaitItem = WaitItem + { itemTs :: Int64 + , itemValue :: Maybe ByteString + , itemWait :: Int + } + +-- Cache runJob result +-- expiredAt, Nothing retrySTM +-- expiredAt, Just bs return bs +type WaitList = IOHashMap JobHandle WaitItem + +-- Distributed lock +-- acquired locked +data LockInfo = LockInfo + { acquired :: [JobHandle] + , locked :: [JobHandle] + , maxCount :: Int + } + +type LockList = IOHashMap LockName LockInfo + +data SchedEnv db tp = SchedEnv + { sPollInterval :: TVar Int -- main poll loop every time interval + -- revert process queue loop every time interval + , sRevertInterval :: TVar Int -- revert process queue loop every time interval + -- the task do timeout + , sTaskTimeout :: TVar Int -- the task do timeout + -- max poll batch size + , sMaxBatchSize :: TVar Int -- max poll batch size + -- client or worker keepalive + , sKeepalive :: TVar Int -- client or worker keepalive + -- run job cache expiration + , sExpiration :: TVar Int -- run job cache expiration + -- auto poll job when job done or failed + , sAutoPoll :: TVar Bool -- auto poll job when job done or failed + -- auto poll lock + , sPolled :: TVar Bool -- auto poll lock + , sCleanup :: IO () + , sFuncStatList :: FuncStatList + , sLocker :: L.Lock + , sGrabQueue :: GrabQueue tp + -- sched state, when false sched is exited. + , sAlive :: TVar Bool -- sched state, when false sched is exited. + , sChanList :: TVar [Action] + , sWaitList :: WaitList + , sLockList :: LockList + , sPersist :: db + } + +newtype SchedT db tp m a = SchedT {unSchedT :: ReaderT (SchedEnv db tp) m a} + deriving + ( Functor + , Applicative + , Monad + , MonadTrans + , MonadIO + , MonadReader (SchedEnv db tp) + ) + +instance MonadUnliftIO m => MonadUnliftIO (SchedT db tp m) where + withRunInIO inner = SchedT $ + ReaderT $ \r -> + withRunInIO $ \run -> + inner (run . runSchedT r) + + +type TaskList = IOHashMap JobHandle (Int64, Async ()) + +runSchedT :: SchedEnv db tp -> SchedT db tp m a -> m a +runSchedT schedEnv = flip runReaderT schedEnv . unSchedT + +initSchedEnv :: (MonadUnliftIO m, Persist db) => P.PersistConfig db -> m () -> m (SchedEnv db tp) +initSchedEnv config sC = do + sFuncStatList <- newIOHashMap + sWaitList <- newIOHashMap + sLockList <- newIOHashMap + sLocker <- L.new + sGrabQueue <- newGrabQueue + sAlive <- newTVarIO True + sChanList <- newTVarIO [] + sPollInterval <- newTVarIO 300 + sRevertInterval <- newTVarIO 300 + sTaskTimeout <- newTVarIO 600 + sMaxBatchSize <- newTVarIO 250 + sKeepalive <- newTVarIO 300 + sExpiration <- newTVarIO 300 + sAutoPoll <- newTVarIO False + sPolled <- newTVarIO False + sCleanup <- toIO sC + sPersist <- liftIO $ P.newPersist config + pure SchedEnv{..} + +startSchedT :: (MonadUnliftIO m, Persist db, Transport tp) => SchedT db tp m () +startSchedT = do + liftIO $ infoM "Periodic.Server.Scheduler" "Scheduler started" + SchedEnv{..} <- ask + runTask_ sRevertInterval revertRunningQueue + taskList <- newIOHashMap + runTask_ sPollInterval $ pushChanList PollJob + runTask 0 $ runChanJob taskList + runTask 100 purgeExpired + runTask 60 revertLockingQueue + + loadInt "poll-interval" sPollInterval + loadInt "revert-interval" sRevertInterval + loadInt "timeout" sTaskTimeout + loadInt "keepalive" sKeepalive + loadInt "max-batch-size" sMaxBatchSize + loadInt "expiration" sExpiration + +loadInt :: (MonadIO m, Persist db) => String -> TVar Int -> SchedT db tp m () +loadInt name ref = do + v <- liftIO . flip P.configGet name =<< asks sPersist + case v of + Nothing -> pure () + Just v' -> atomically $ writeTVar ref v' + +saveInt :: (MonadIO m, Persist db) => String -> Int -> TVar Int -> SchedT db tp m () +saveInt name v ref = do + p <- asks sPersist + liftIO $ P.configSet p name v + atomically $ writeTVar ref v + +setConfigInt :: (MonadIO m, Persist db) => String -> Int -> SchedT db tp m () +setConfigInt key val = do + SchedEnv {..} <- ask + case key of + "poll-interval" -> saveInt "poll-interval" val sPollInterval + "revert-interval" -> saveInt "revert-interval" val sRevertInterval + "timeout" -> saveInt "timeout" val sTaskTimeout + "keepalive" -> saveInt "keepalive" val sKeepalive + "max-batch-size" -> saveInt "max-batch-size" val sMaxBatchSize + "expiration" -> saveInt "expiration" val sExpiration + _ -> pure () + +getConfigInt :: (MonadIO m, Persist db) => String -> SchedT db tp m Int +getConfigInt key = do + SchedEnv {..} <- ask + case key of + "poll-interval" -> readTVarIO sPollInterval + "revert-interval" -> readTVarIO sRevertInterval + "timeout" -> readTVarIO sTaskTimeout + "keepalive" -> readTVarIO sKeepalive + "max-batch-size" -> readTVarIO sMaxBatchSize + "expiration" -> readTVarIO sExpiration + _ -> pure 0 + +keepalive :: Monad m => SchedT db tp m (TVar Int) +keepalive = asks sKeepalive + +runTask :: (MonadUnliftIO m) => Int -> SchedT db tp m () -> SchedT db tp m () +runTask d m = flip runTask_ m =<< newTVarIO d + +runTask_ :: (MonadUnliftIO m) => TVar Int -> SchedT db tp m () -> SchedT db tp m () +runTask_ d m = void . async $ do + SchedEnv{..} <- ask + void . runMaybeT . forever $ do + interval <- readTVarIO d + when (interval > 0) $ threadDelay $ interval * 1000 * 1000 + alive <- readTVarIO sAlive + if alive then lift m + else mzero + +runChanJob + :: (MonadUnliftIO m, Persist db, Transport tp) + => TaskList -> SchedT db tp m () +runChanJob taskList = do + cl <- asks sChanList + al <- asks sAlive + acts <- atomically $ do + acts <- readTVar cl + if null acts then do + st <- readTVar al + if st then retrySTM + else pure [] + else do + writeTVar cl [] + pure acts + + mapM_ (doChanJob taskList) acts + + where doChanJob + :: (MonadUnliftIO m, Persist db, Transport tp) + => TaskList -> Action -> SchedT db tp m () + doChanJob tl (Add job) = reSchedJob tl job + doChanJob tl (Remove job) = findTask tl job >>= mapM_ cancel + doChanJob tl Cancel = mapM_ (cancel . snd) =<< FL.elems tl + doChanJob tl PollJob = pollJob tl + doChanJob tl (TryPoll jh) = removeTaskAndTryPoll tl jh + + +pollInterval :: (MonadIO m, Num a) => SchedT db tp m a +pollInterval = fmap fromIntegral . readTVarIO =<< asks sPollInterval + +removeTaskAndTryPoll :: MonadIO m => TaskList -> JobHandle -> SchedT db tp m () +removeTaskAndTryPoll taskList jh = do + FL.delete taskList jh + polled <- asks sPolled + isPolled <- readTVarIO polled + autoPoll <- readTVarIO =<< asks sAutoPoll + when (isPolled && autoPoll) $ do + maxBatchSize <- readTVarIO =<< asks sMaxBatchSize + size <- FL.size taskList + when (size < maxBatchSize) $ do + atomically $ writeTVar polled False + pushChanList PollJob + +pollJob + :: (MonadUnliftIO m, Persist db, Transport tp) + => TaskList -> SchedT db tp m () +pollJob taskList = do + polled <- asks sPolled + atomically $ writeTVar polled False + mapM_ checkPoll =<< FL.toList taskList + stList <- asks sFuncStatList + funcList <- foldr foldFunc [] <$> FL.toList stList + pollJob_ taskList funcList + atomically $ writeTVar polled True + + where foldFunc :: (FuncName, FuncStat) -> [FuncName] -> [FuncName] + foldFunc (_, FuncStat{sWorker=0}) acc = acc + foldFunc (fn, _) acc = fn:acc + + checkPoll + :: (MonadIO m) + => (JobHandle, (Int64, Async ())) -> SchedT db tp m () + checkPoll (jh, (_, w)) = do + r <- poll w + case r of + Just (Right ()) -> FL.delete taskList jh + Just (Left e) -> do + FL.delete taskList jh + liftIO $ errorM "Periodic.Server.Scheduler" ("Poll error: " ++ show e) + Nothing -> do + r0 <- canRun fn + unless r0 $ cancel w + + where (fn, _) = unHandle jh + +pollJob_ + :: (MonadUnliftIO m, Persist db, Transport tp) + => TaskList -> [FuncName] -> SchedT db tp m () +pollJob_ _ [] = pure () +pollJob_ taskList funcList = do + now <- getEpochTime + next <- (+ (100 + now)) <$> pollInterval + handles <- FL.keys taskList + let check job = notElem (getHandle job) handles && (getSchedAt job < next) + + maxBatchSize <- readTVarIO =<< asks sMaxBatchSize + p <- asks sPersist + jobs <- liftIO $ + P.foldrPending p next funcList (foldFunc (maxBatchSize * 2) check now) PSQ.empty + + mapM_ (checkJob taskList) jobs + + autoPoll <- asks sAutoPoll + atomically $ writeTVar autoPoll (length jobs > maxBatchSize) + + where foldFunc :: Int -> (Job -> Bool) -> Int64 -> Job -> HashPSQ JobHandle Int64 Job -> HashPSQ JobHandle Int64 Job + foldFunc s f now job acc | f job = trimPSQ $ PSQ.insert (getHandle job) (now - getSchedAt job) job acc + | otherwise = acc + where trimPSQ :: HashPSQ JobHandle Int64 Job -> HashPSQ JobHandle Int64 Job + trimPSQ q | PSQ.size q > s = trimPSQ $ PSQ.deleteMin q + | otherwise = q + + checkJob + :: (MonadUnliftIO m, Persist db, Transport tp) + => TaskList -> Job -> SchedT db tp m () + checkJob tl job = do + w <- findTask tl job + case w of + Nothing -> do + p <- asks sPersist + isProc <- liftIO $ P.member p Running fn jn + unless isProc $ reSchedJob tl job + Just w0 -> do + r <- canRun fn + unless r $ cancel w0 + where fn = getFuncName job + jn = getName job + +pushChanList :: MonadIO m => Action -> SchedT db tp m () +pushChanList act = do + cl <- asks sChanList + atomically $ do + l <- readTVar cl + writeTVar cl (act:l) + +pushJob :: (MonadIO m, Persist db) => Job -> SchedT db tp m () +pushJob job = do + liftIO $ infoM "Periodic.Server.Scheduler" ("pushJob: " ++ show (getHandle job)) + p <- asks sPersist + isRunning <- liftIO $ P.member p Running fn jn + unless isRunning $ do + job' <- fixedSchedAt job + liftIO $ P.insert p Pending fn jn job' + pushChanList (Add job') + + where fn = getFuncName job + jn = getName job + + +fixedSchedAt :: MonadIO m => Job -> SchedT db tp m Job +fixedSchedAt job = do + now <- getEpochTime + if getSchedAt job < now then + return $ setSchedAt now job + else return job + +reSchedJob :: (MonadUnliftIO m, Persist db, Transport tp) => TaskList -> Job -> SchedT db tp m () +reSchedJob taskList job = do + w <- findTask taskList job + forM_ w cancel + + interval <- (+100) <$> pollInterval + next <- (+ interval) <$> getEpochTime + when (getSchedAt job < next) $ do + r <- canRun $ getFuncName job + c <- check taskList + when (r && c) $ do + w' <- schedJob taskList job + FL.insert taskList (getHandle job) (getSchedAt job, w') + where check :: (MonadIO m) => TaskList -> SchedT db tp m Bool + check tl = do + maxBatchSize <- readTVarIO =<< asks sMaxBatchSize + size <- FL.size tl + if size < maxBatchSize * 2 then return True + else do + lastTask <- findLastTask tl + case lastTask of + Nothing -> return True + Just (sc, jh, w) -> + if sc < getSchedAt job then return False + else do + cancel w + FL.delete taskList jh + return True + +findTask :: MonadIO m => TaskList -> Job -> SchedT db tp m (Maybe (Async ())) +findTask taskList job = fmap snd <$> FL.lookup taskList (getHandle job) + +findLastTask :: MonadIO m => TaskList -> SchedT db tp m (Maybe (Int64, JobHandle, Async ())) +findLastTask tl = atomically $ FL.foldrWithKeySTM tl f Nothing + where f :: JobHandle + -> (Int64, a) + -> Maybe (Int64, JobHandle, a) + -> Maybe (Int64, JobHandle, a) + f jh (sc, t) Nothing = Just (sc, jh, t) + f jh (sc, t) (Just (sc1, jh1, t1)) + | sc > sc1 = Just (sc, jh, t) + | otherwise = Just (sc1, jh1, t1) + +canRun :: MonadIO m => FuncName -> SchedT db tp m Bool +canRun fn = asks sFuncStatList >>= flip canRun_ fn + +canRun_ :: MonadIO m => FuncStatList -> FuncName -> m Bool +canRun_ stList fn = do + st0 <- FL.lookup stList fn + case st0 of + Nothing -> pure False + Just FuncStat{sWorker=0} -> pure False + Just _ -> pure True + +schedJob + :: (MonadUnliftIO m, Persist db, Transport tp) + => TaskList -> Job -> SchedT db tp m (Async ()) +schedJob taskList = async . schedJob_ taskList + +schedJob_ :: (MonadUnliftIO m, Persist db, Transport tp) => TaskList -> Job -> SchedT db tp m () +schedJob_ taskList job = do + SchedEnv{..} <- ask + r <- canRun fn + when r $ do + now <- getEpochTime + when (schedAt > now + 1) . threadDelay . fromIntegral $ (schedAt - now) * 1000000 + FuncStat{..} <- atomically $ do + st <- FL.lookupSTM sFuncStatList fn + case st of + Nothing -> retrySTM + Just FuncStat{sWorker=0} -> retrySTM + Just st' -> pure st' + if sBroadcast then popAgentListThen + else popAgentThen taskList + + where fn = getFuncName job + jn = getName job + schedAt = getSchedAt job + jh = getHandle job + + popAgentThen + :: (MonadUnliftIO m, Persist db, Transport tp) => TaskList -> SchedT db tp m () + popAgentThen tl = do + SchedEnv{..} <- ask + (jq, env0) <- atomically $ popAgentSTM sGrabQueue fn + alive <- runSessionT1 env0 sessionState + if alive then do + IL.insert jq (getHandle job) + nextSchedAt <- getEpochTime + liftIO $ P.insert sPersist Running fn jn $ setSchedAt nextSchedAt job + r <- doSubmitJob env0 + case r of + Left _ -> do + liftIO $ P.insert sPersist Pending fn jn $ setSchedAt nextSchedAt job + IL.delete jq jh + schedJob_ tl job + Right _ -> endSchedJob + else schedJob_ tl job + + popAgentListThen :: (MonadUnliftIO m, Transport tp) => SchedT db tp m () + popAgentListThen = do + SchedEnv{..} <- ask + agents <- popAgentList sGrabQueue fn + mapM_ (doSubmitJob . snd) agents + unless (null agents) endSchedJob -- wait to resched the broadcast job + + doSubmitJob :: (MonadUnliftIO m, Transport tp) => CSEnv tp -> SchedT db tp m (Either SomeException ()) + doSubmitJob agent = do + SchedEnv{..} <- ask + tryAny $ assignJob agent job + + endSchedJob :: MonadIO m => SchedT db tp m () + endSchedJob = pushChanList (TryPoll jh) + +adjustFuncStat :: (MonadIO m, Persist db) => FuncName -> SchedT db tp m () +adjustFuncStat fn = do + SchedEnv{..} <- ask + size <- liftIO $ P.size sPersist Pending fn + sizePQ <- liftIO $ P.size sPersist Running fn + sizeL <- liftIO $ P.size sPersist Locking fn + sc <- liftIO $ P.minSchedAt sPersist fn + + schedAt <- if sc > 0 then pure sc else getEpochTime + + FL.alter sFuncStatList (update (size + sizePQ + sizeL) sizePQ sizeL schedAt) fn + + where update :: Int64 -> Int64 -> Int64 -> Int64 -> Maybe FuncStat -> Maybe FuncStat + update size sizePQ sizeL schedAt st = + Just ((fromMaybe (funcStat fn) st) { sJob = size + , sRunning = sizePQ + , sLocking = sizeL + , sSchedAt = schedAt + }) + +removeJob :: (MonadIO m, Persist db) => Job -> SchedT db tp m () +removeJob job = do + liftIO $ infoM "Periodic.Server.Scheduler" ("removeJob: " ++ show (getHandle job)) + p <- asks sPersist + liftIO $ P.delete p fn jn + + pushChanList (Remove job) + pushResult jh "" + where jn = getName job + fn = getFuncName job + jh = getHandle job + +dumpJob :: (MonadIO m, Persist db) => SchedT db tp m [Job] +dumpJob = liftIO . P.dumpJob =<< asks sPersist + +alterFunc :: (MonadIO m, Persist db) => FuncName -> (Maybe FuncStat -> Maybe FuncStat) -> SchedT db tp m () +alterFunc n f = do + SchedEnv{..} <- ask + FL.alter sFuncStatList f n + liftIO $ P.insertFuncName sPersist n + pushChanList PollJob + +addFunc :: (MonadIO m, Persist db) => FuncName -> SchedT db tp m () +addFunc n = broadcastFunc n False + +broadcastFunc :: (MonadIO m, Persist db) => FuncName -> Bool -> SchedT db tp m () +broadcastFunc n cast = do + liftIO $ infoM "Periodic.Server.Scheduler" (h ++ ": " ++ show n) + alterFunc n updateStat + + where updateStat :: Maybe FuncStat -> Maybe FuncStat + updateStat Nothing = Just ((funcStat n) {sWorker = 1, sBroadcast = cast}) + updateStat (Just fs) = Just (fs { sWorker = sWorker fs + 1, sBroadcast = cast }) + + h = if cast then "broadcastFunc" else "addFunc" + +removeFunc :: (MonadIO m, Persist db) => FuncName -> SchedT db tp m () +removeFunc n = do + liftIO $ infoM "Periodic.Server.Scheduler" ("removeFunc: " ++ show n) + alterFunc n updateStat + + where updateStat :: Maybe FuncStat -> Maybe FuncStat + updateStat Nothing = Just (funcStat n) + updateStat (Just fs) = Just (fs { sWorker = max (sWorker fs - 1) 0 }) + +dropFunc :: (MonadUnliftIO m, Persist db) => FuncName -> SchedT db tp m () +dropFunc n = do + liftIO $ infoM "Periodic.Server.Scheduler" ("dropFunc: " ++ show n) + SchedEnv{..} <- ask + L.with sLocker $ do + st <- FL.lookup sFuncStatList n + case st of + Just FuncStat{sWorker=0} -> do + FL.delete sFuncStatList n + liftIO $ P.removeFuncName sPersist n + _ -> pure () + + pushChanList PollJob + +pushGrab :: MonadIO m => IOList FuncName -> IOList JobHandle -> CSEnv tp -> SchedT db tp m () +pushGrab funcList handleList ag = do + queue <- asks sGrabQueue + pushAgent queue funcList handleList ag + +assignJob :: (MonadUnliftIO m, Transport tp) => CSEnv tp -> Job -> m () +assignJob env0 job = do + liftIO $ infoM "Periodic.Server.Scheduler" ("assignJob: " ++ show (getHandle job)) + runSessionT1 env0 $ send $ packetRES (JobAssign job) + +failJob :: (MonadUnliftIO m, Persist db) => JobHandle -> SchedT db tp m () +failJob jh = do + liftIO $ infoM "Periodic.Server.Scheduler" ("failJob: " ++ show jh) + releaseLock' jh + isWaiting <- existsWaitList jh + if isWaiting then do + removeFromWaitList jh + doneJob jh "" + else do + p <- asks sPersist + job <- liftIO $ P.lookup p Running fn jn + when (isJust job) $ do + nextSchedAt <- getEpochTime + retryJob $ setSchedAt nextSchedAt $ fromJust job + + where (fn, jn) = unHandle jh + +retryJob :: (MonadIO m, Persist db) => Job -> SchedT db tp m () +retryJob job = do + p <- asks sPersist + liftIO $ P.insert p Pending fn jn job + + pushChanList (Add job) + + where fn = getFuncName job + jn = getName job + +doneJob + :: (MonadUnliftIO m, Persist db) + => JobHandle -> ByteString -> SchedT db tp m () +doneJob jh w = do + liftIO $ infoM "Periodic.Server.Scheduler" ("doneJob: " ++ show jh) + releaseLock' jh + p <- asks sPersist + liftIO $ P.delete p fn jn + pushResult jh w + where (fn, jn) = unHandle jh + +schedLaterJob + :: (MonadUnliftIO m, Persist db) + => JobHandle -> Int64 -> Int -> SchedT db tp m () +schedLaterJob jh later step = do + liftIO $ infoM "Periodic.Server.Scheduler" ("schedLaterJob: " ++ show jh) + releaseLock' jh + isWaiting <- existsWaitList jh + if isWaiting then do + removeFromWaitList jh + doneJob jh "" + else do + p <- asks sPersist + job <- liftIO $ P.lookup p Running fn jn + when (isJust job) $ do + let job' = fromJust job + + nextSchedAt <- (+) later <$> getEpochTime + retryJob $ setCount (getCount job' + step) $ setSchedAt nextSchedAt job' + + where (fn, jn) = unHandle jh + +acquireLock + :: (MonadUnliftIO m, Persist db) + => LockName -> Int -> JobHandle -> SchedT db tp m Bool +acquireLock name count jh = do + liftIO $ infoM "Periodic.Server.Scheduler" ("acquireLock: " ++ show name ++ " " ++ show count ++ " " ++ show jh) + locker <- asks sLocker + L.with locker $ do + lockList <- asks sLockList + p <- asks sPersist + j <- liftIO $ P.lookup p Running fn jn + case j of + Nothing -> pure True + Just job -> do + r <- atomically $ do + l <- FL.lookupSTM lockList name + case l of + Nothing -> do + FL.insertSTM lockList name LockInfo + { acquired = [jh] + , locked = [] + , maxCount = count + } + pure True + Just info@LockInfo {..} -> do + let newCount = max maxCount count + if jh `elem` acquired then pure True + else if jh `elem` locked then pure False + else + if length acquired < maxCount then do + FL.insertSTM lockList name info + { acquired = acquired ++ [jh] + , maxCount = newCount + } + pure True + else do + FL.insertSTM lockList name info + { locked = locked ++ [jh] + , maxCount = newCount + } + pure False + + unless r $ liftIO $ P.insert p Locking fn jn job + return r + + where (fn, jn) = unHandle jh + +releaseLock + :: (MonadUnliftIO m, Persist db) + => LockName -> JobHandle -> SchedT db tp m () +releaseLock name jh = do + locker <- asks sLocker + L.with locker $ releaseLock_ name jh + +releaseLock_ + :: (MonadUnliftIO m, Persist db) + => LockName -> JobHandle -> SchedT db tp m () +releaseLock_ name jh = do + liftIO $ infoM "Periodic.Server.Scheduler" ("releaseLock: " ++ show name ++ " " ++ show jh) + p <- asks sPersist + lockList <- asks sLockList + h <- atomically $ do + l <- FL.lookupSTM lockList name + case l of + Nothing -> pure Nothing + Just info@LockInfo {..} -> + if jh `elem` acquired then + case locked of + [] -> do + FL.insertSTM lockList name info + { acquired = L.delete jh acquired + } + pure Nothing + x:xs -> do + FL.insertSTM lockList name info + { acquired = L.delete jh acquired + , locked = xs + } + pure $ Just x + else pure Nothing + + case h of + Nothing -> pure () + Just hh -> do + let (fn, jn) = unHandle hh + j <- liftIO $ P.lookup p Locking fn jn + case j of + Nothing -> releaseLock_ name hh + Just job -> do + liftIO $ P.insert p Pending fn jn job + pushChanList (Add job) + +releaseLock' + :: (MonadUnliftIO m, Persist db) + => JobHandle -> SchedT db tp m () +releaseLock' jh = do + lockList <- asks sLockList + names <- atomically $ FL.foldrWithKeySTM lockList foldFunc [] + mapM_ (`releaseLock` jh) names + + where foldFunc :: LockName -> LockInfo -> [LockName] -> [LockName] + foldFunc n LockInfo {..} acc | jh `elem` acquired = n : acc + | jh `elem` locked = n : acc + | otherwise = acc + +countLock :: MonadUnliftIO m => (LockInfo -> [JobHandle]) -> FuncName -> SchedT db tp m Int +countLock f fn = do + lockList <- asks sLockList + sum . map mapFunc <$> FL.elems lockList + where filterFunc :: JobHandle -> Bool + filterFunc jh = fn0 == fn + where (fn0, _) = unHandle jh + + mapFunc :: LockInfo -> Int + mapFunc = length . filter filterFunc . f + +getMaxLockCount :: MonadUnliftIO m => SchedT db tp m Int +getMaxLockCount = do + lockList <- asks sLockList + maximum . map maxCount <$> FL.elems lockList + + +status :: (MonadIO m, Persist db) => SchedT db tp m [FuncStat] +status = do + mapM_ adjustFuncStat =<< liftIO . P.funcList =<< asks sPersist + FL.elems =<< asks sFuncStatList + +revertRunningQueue :: (MonadUnliftIO m, Persist db) => SchedT db tp m () +revertRunningQueue = do + now <- getEpochTime + tout <- fmap fromIntegral . readTVarIO =<< asks sTaskTimeout + p <- asks sPersist + handles <- liftIO $ P.foldr p Running (foldFunc (check now tout)) [] + mapM_ (failJob . getHandle) handles + + where foldFunc :: (Job -> Bool) -> Job -> [Job] -> [Job] + foldFunc f job acc | f job = job : acc + | otherwise = acc + + check :: Int64 -> Int64 -> Job -> Bool + check now t0 job + | getTimeout job > 0 = getSchedAt job + fromIntegral (getTimeout job) < now + | otherwise = getSchedAt job + fromIntegral t0 < now + +revertLockingQueue :: (MonadUnliftIO m, Persist db) => SchedT db tp m () +revertLockingQueue = mapM_ checkAndReleaseLock =<< liftIO . P.funcList =<< asks sPersist + + where checkAndReleaseLock + :: (MonadUnliftIO m, Persist db) + => FuncName -> SchedT db tp m () + checkAndReleaseLock fn = do + p <- asks sPersist + sizeLocked <- countLock locked fn + sizeAcquired <- countLock acquired fn + liftIO $ infoM "Peridic.Server.Scheduler" + $ "LockInfo " ++ show fn + ++ " Locked:" ++ show sizeLocked + ++ " Acquired:" ++ show sizeAcquired + when (sizeLocked > 0 && sizeAcquired == 0) $ do + count <- getMaxLockCount + handles <- liftIO $ P.foldrLocking p count fn (:) [] + mapM_ pushJob handles + + +purgeExpired :: MonadIO m => SchedT db tp m () +purgeExpired = do + now <- getEpochTime + wl <- asks sWaitList + ex <- fmap fromIntegral . readTVarIO =<< asks sExpiration + atomically $ do + ks <- FL.foldrWithKeySTM wl (foldFunc (check (now - ex))) [] + mapM_ (FL.deleteSTM wl) ks + + where foldFunc :: (WaitItem -> Bool) -> JobHandle -> WaitItem -> [JobHandle] -> [JobHandle] + foldFunc f jh v acc | f v = jh : acc + | otherwise = acc + + check :: Int64 -> WaitItem -> Bool + check t0 item = itemTs item < t0 + +shutdown :: (MonadUnliftIO m) => SchedT db tp m () +shutdown = do + liftIO $ infoM "Periodic.Server.Scheduler" "Scheduler shutdown" + SchedEnv{..} <- ask + pushChanList Cancel + alive <- atomically $ do + t <- readTVar sAlive + writeTVar sAlive False + return t + when alive . void . async $ liftIO sCleanup + +prepareWait :: MonadIO m => Job -> SchedT db tp m () +prepareWait job = pushResult_ updateWL jh + where updateWL :: Int64 -> Maybe WaitItem -> Maybe WaitItem + updateWL now Nothing = Just $ WaitItem {itemTs = now, itemValue = Nothing, itemWait = 1} + updateWL now (Just item) = Just $ item {itemTs = now, itemWait = itemWait item + 1} + + jh = getHandle job + +waitResult :: MonadIO m => TVar Bool -> Job -> SchedT db tp m ByteString +waitResult state job = do + wl <- asks sWaitList + atomically $ do + st <- readTVar state + if st then do + w0 <- FL.lookupSTM wl jh + case w0 of + Nothing -> pure "" + Just item -> + case itemValue item of + Nothing -> retrySTM + Just w1 -> do + + if itemWait item > 1 then + FL.insertSTM wl jh item { itemWait = itemWait item - 1 } + else + FL.deleteSTM wl jh + + pure w1 + + else pure "" + + where jh = getHandle job + +pushResult + :: MonadIO m + => JobHandle -> ByteString -> SchedT db tp m () +pushResult jh w = pushResult_ updateWL jh + where updateWL :: Int64 -> Maybe WaitItem -> Maybe WaitItem + updateWL _ Nothing = Nothing + updateWL now (Just item) = Just item {itemTs=now, itemValue = Just w} + +pushResult_ + :: MonadIO m + => (Int64 -> Maybe WaitItem -> Maybe WaitItem) + -> JobHandle -> SchedT db tp m () +pushResult_ f jh = do + wl <- asks sWaitList + now <- getEpochTime + FL.alter wl (f now) jh + +existsWaitList :: MonadIO m => JobHandle -> SchedT db tp m Bool +existsWaitList jh = do + wl <- asks sWaitList + isJust <$> FL.lookup wl jh + +lookupPrevResult :: MonadIO m => Job -> SchedT db tp m (Maybe ByteString) +lookupPrevResult job = do + wl <- asks sWaitList + r <- FL.lookup wl jh + case r of + Nothing -> pure Nothing + (Just WaitItem {itemValue = Nothing}) -> pure Nothing + (Just WaitItem {itemValue = Just v}) -> pure (Just v) + + where jh = getHandle job + +removeFromWaitList :: MonadIO m => JobHandle -> SchedT db tp m () +removeFromWaitList jh = do + wl <- asks sWaitList + FL.delete wl jh diff --git a/src/Periodic/Server/Types.hs b/src/Periodic/Server/Types.hs new file mode 100644 --- /dev/null +++ b/src/Periodic/Server/Types.hs @@ -0,0 +1,55 @@ +module Periodic.Server.Types + ( Command (..) + , ClientConfig (..) + , CSEnv + , ServerCommand (..) + ) where +import Data.Binary (Binary (..), getWord8) +import Data.Binary.Get (lookAhead) +import Periodic.IOList (IOList) +import Periodic.Node (SessionEnv1) +import qualified Periodic.Types.ClientCommand as CC +import Periodic.Types.Job (FuncName, JobHandle) +import Periodic.Types.ServerCommand (ServerCommand (..)) +import qualified Periodic.Types.WorkerCommand as WC + +data Command = CC CC.ClientCommand + | WC WC.WorkerCommand + +instance Binary Command where + get = do + cmd <- lookAhead getWord8 + case cmd of + 1 -> WC <$> get + 2 -> WC <$> get + 3 -> WC <$> get + 4 -> WC <$> get + 11 -> WC <$> get + -- 9 -> WC <$> get + 7 -> WC <$> get + 8 -> WC <$> get + 21 -> WC <$> get + 27 -> WC <$> get + 28 -> WC <$> get + 13 -> CC <$> get + 14 -> CC <$> get + 9 -> CC <$> get + 15 -> CC <$> get + 17 -> CC <$> get + 20 -> CC <$> get + 22 -> CC <$> get + 23 -> CC <$> get + 18 -> CC <$> get + 19 -> CC <$> get + 25 -> CC <$> get + _ -> error $ "Error Command" ++ show cmd + + put (CC cmd) = put cmd + put (WC cmd) = put cmd + +data ClientConfig = ClientConfig + { wFuncList :: IOList FuncName + , wJobQueue :: IOList JobHandle + } + +type CSEnv = SessionEnv1 ClientConfig Command