module Control.Concurrent.Capataz.Internal.Util where
import Protolude
import Control.Concurrent.STM (STM, atomically, retry)
import Control.Concurrent.STM.TVar (TVar, readTVar, writeTVar)
import Data.IORef (atomicModifyIORef', readIORef)
import qualified Data.Text as T
import Data.Time.Clock (getCurrentTime)
import qualified Data.HashMap.Strict as HashMap
import GHC.Conc (labelThread)
import Control.Concurrent.Capataz.Internal.Types
getTidNumber :: ThreadId -> Maybe Text
getTidNumber tid = case T.words $ show tid of
(_:tidNumber:_) -> Just tidNumber
_ -> Nothing
setProcessThreadName :: WorkerId -> WorkerName -> IO ()
setProcessThreadName workerId workerName = do
tid <- myThreadId
let workerIdentifier =
T.unpack workerName <> "_" <> show workerId <> "_" <> maybe
""
T.unpack
(getTidNumber tid)
labelThread tid workerIdentifier
getProcessId :: Process -> ProcessId
getProcessId process = case process of
WorkerProcess Worker { workerId } -> workerId
SupervisorProcess Supervisor { supervisorId } -> supervisorId
fetchProcess :: SupervisorEnv -> ProcessId -> IO (Maybe Process)
fetchProcess SupervisorEnv { supervisorProcessMap } processId = do
processMap <- readIORef supervisorProcessMap
case HashMap.lookup processId processMap of
Just process -> return $ Just process
_ -> return Nothing
appendProcessToMap :: SupervisorEnv -> Process -> IO ()
appendProcessToMap SupervisorEnv { supervisorProcessMap } process =
atomicModifyIORef' supervisorProcessMap
(\processMap -> (appendProcess processMap, ()))
where
appendProcess = HashMap.alter (const $ Just process) (getProcessId process)
removeProcessFromMap :: SupervisorEnv -> ProcessId -> IO ()
removeProcessFromMap SupervisorEnv { supervisorProcessMap } processId =
atomicModifyIORef'
supervisorProcessMap
( \processMap -> maybe (processMap, ())
(const (HashMap.delete processId processMap, ()))
(HashMap.lookup processId processMap)
)
resetProcessMap :: SupervisorEnv -> (ProcessMap -> ProcessMap) -> IO ()
resetProcessMap SupervisorEnv { supervisorProcessMap } processMapFn =
atomicModifyIORef' supervisorProcessMap
(\processMap -> (processMapFn processMap, ()))
readProcessMap :: SupervisorEnv -> IO ProcessMap
readProcessMap SupervisorEnv { supervisorProcessMap } =
readIORef supervisorProcessMap
sortProcessesByTerminationOrder
:: ProcessTerminationOrder -> ProcessMap -> [Process]
sortProcessesByTerminationOrder terminationOrder processMap =
case terminationOrder of
OldestFirst -> workers
NewestFirst -> reverse workers
where
processCreationTime (WorkerProcess Worker { workerCreationTime }) =
workerCreationTime
processCreationTime (SupervisorProcess Supervisor { supervisorCreationTime })
= supervisorCreationTime
workers = sortBy (comparing processCreationTime) (HashMap.elems processMap)
readSupervisorStatusSTM :: TVar SupervisorStatus -> STM SupervisorStatus
readSupervisorStatusSTM statusVar = do
status <- readTVar statusVar
if status == Initializing then retry else return status
readSupervisorStatus :: SupervisorEnv -> IO SupervisorStatus
readSupervisorStatus SupervisorEnv { supervisorStatusVar } =
atomically $ readTVar supervisorStatusVar
writeSupervisorStatus :: SupervisorEnv -> SupervisorStatus -> IO ()
writeSupervisorStatus SupervisorEnv { supervisorId, supervisorName, supervisorStatusVar, notifyEvent } newSupervisorStatus
= do
prevSupervisorStatus <- atomically $ do
prevStatus <- readTVar supervisorStatusVar
writeTVar supervisorStatusVar newSupervisorStatus
return prevStatus
eventTime <- getCurrentTime
notifyEvent SupervisorStatusChanged
{ supervisorId = supervisorId
, supervisorName = supervisorName
, prevSupervisorStatus
, newSupervisorStatus
, eventTime
}
sendControlMsg :: SupervisorEnv -> ControlAction -> IO ()
sendControlMsg SupervisorEnv { supervisorNotify } ctrlMsg =
supervisorNotify (ControlAction ctrlMsg)
sendSyncControlMsg
:: SupervisorEnv
-> (IO () -> ControlAction)
-> IO ()
sendSyncControlMsg SupervisorEnv { supervisorNotify } mkCtrlMsg = do
result <- newEmptyMVar
supervisorNotify (ControlAction $ mkCtrlMsg (putMVar result ()))
takeMVar result
capatazOptionsToSupervisorOptions :: CapatazOptions -> SupervisorOptions
capatazOptionsToSupervisorOptions CapatazOptions {..} =
SupervisorOptions {supervisorName = "capataz-root-supervisor", ..}
toParentSupervisorEnv :: SupervisorEnv -> ParentSupervisorEnv
toParentSupervisorEnv SupervisorEnv { supervisorId, supervisorName, supervisorNotify, notifyEvent }
= ParentSupervisorEnv {..}