250 lines
9.7 KiB
Haskell
250 lines
9.7 KiB
Haskell
{-# OPTIONS_GHC -fno-warn-unused-top-binds #-}
|
|
{-# OPTIONS_GHC -fno-warn-orphans #-}
|
|
{-# OPTIONS_GHC -fno-warn-deprecations #-}
|
|
|
|
module Load
|
|
( main
|
|
) where
|
|
|
|
import "uniworx" Import hiding (Option(..), Normal, responseBody)
|
|
|
|
import System.Console.GetOpt
|
|
|
|
import qualified Data.Text as Text
|
|
|
|
import qualified Data.Map.Strict as Map
|
|
|
|
import Data.Random.Normal
|
|
import qualified Control.Monad.Random.Class as Random
|
|
import System.Random (RandomGen)
|
|
|
|
import System.Exit (exitWith, ExitCode(..))
|
|
import System.IO (hPutStrLn)
|
|
|
|
import UnliftIO.Concurrent (threadDelay)
|
|
|
|
import System.Clock (getTime, Clock(Monotonic))
|
|
import qualified System.Clock as Clock
|
|
|
|
import Network.URI
|
|
|
|
import qualified Data.ByteString as BS
|
|
import qualified Data.ByteString.Char8 as CBS
|
|
import qualified Data.Char as Char (isSpace)
|
|
|
|
import Network.Wreq
|
|
import Network.Wreq.Session (Session)
|
|
import qualified Network.Wreq.Session as Session
|
|
import Network.HTTP.Client.TLS (tlsManagerSettings)
|
|
|
|
|
|
instance (a ~ b, Monad m) => Monoid (Kleisli m a b) where
|
|
mempty = Kleisli return
|
|
|
|
instance (a ~ b, Monad m) => Semigroup (Kleisli m a b) where
|
|
Kleisli f <> Kleisli g = Kleisli $ f <=< g
|
|
|
|
|
|
data Normal k = Normal
|
|
{ dAvg :: k
|
|
, dRelDev :: Centi
|
|
} deriving (Eq, Ord, Read, Show, Generic, Typeable)
|
|
|
|
sampleN :: (Random.MonadSplit g m, RandomGen g) => (k -> Centi -> k) -> Normal k -> m k
|
|
sampleN scale Normal{..}
|
|
| dRelDev == 0 = return dAvg
|
|
| otherwise = do
|
|
gen <- Random.getSplit
|
|
let (realToFrac -> r, _) = normal' (1, realToFrac dRelDev :: Double) gen
|
|
return $ dAvg `scale` r
|
|
|
|
instance PathPiece k => PathPiece (Normal k) where
|
|
toPathPiece Normal{dRelDev = MkFixed perc, dAvg}
|
|
| perc == 0 = toPathPiece dAvg
|
|
| otherwise = toPathPiece dAvg <> ";" <> toPathPiece perc <> "%"
|
|
fromPathPiece t
|
|
| (avg, relDev') <- Text.breakOn ";" t
|
|
, Just relDev <- Text.stripSuffix "%" =<< Text.stripPrefix ";" relDev'
|
|
= Normal <$> fromPathPiece avg <*> (MkFixed <$> fromPathPiece relDev)
|
|
| otherwise
|
|
= Normal <$> fromPathPiece t <*> pure 0
|
|
|
|
scaleDiffTime :: DiffTime -> Centi -> DiffTime
|
|
scaleDiffTime (diffTimeToPicoseconds -> ps) s = picosecondsToDiffTime . round $ s * fromIntegral ps
|
|
|
|
sampleNDiffTime :: (Random.MonadSplit g m, RandomGen g) => Normal DiffTime -> m DiffTime
|
|
sampleNDiffTime = sampleN scaleDiffTime
|
|
|
|
|
|
instance PathPiece DiffTime where
|
|
toPathPiece = toPathPiece . MkFixed @E12 . diffTimeToPicoseconds
|
|
fromPathPiece t = fromPathPiece t <&> \(MkFixed ps :: Pico) -> picosecondsToDiffTime ps
|
|
|
|
|
|
data LoadSimulation
|
|
= LoadSheetDownload
|
|
| LoadSheetSubmission
|
|
deriving (Eq, Ord, Read, Show, Enum, Bounded, Generic, Typeable)
|
|
deriving anyclass (Universe, Finite)
|
|
|
|
nullaryPathPiece ''LoadSimulation $ camelToPathPiece' 1
|
|
|
|
data LoadOptions = LoadOptions
|
|
{ loadSimulations :: Map LoadSimulation SimulationOptions
|
|
, loadBaseURI :: URI
|
|
, loadToken :: Maybe Jwt
|
|
, loadTerm :: TermId, loadSchool :: SchoolId, loadCourse :: CourseShorthand, loadSheet :: SheetName
|
|
} deriving (Eq, Ord, Show, Generic, Typeable)
|
|
|
|
instance Default LoadOptions where
|
|
def = LoadOptions
|
|
{ loadSimulations = Map.empty
|
|
, loadBaseURI = error "No BaseURI given"
|
|
, loadToken = Nothing
|
|
, loadTerm = error "No term given", loadSchool = error "No school given", loadCourse = error "No course given", loadSheet = error "No sheet given"
|
|
}
|
|
|
|
data SimulationOptions = SimulationOptions
|
|
{ simParallel :: Natural
|
|
, simDelay, simDuration :: Normal DiffTime
|
|
} deriving (Eq, Ord, Show, Generic, Typeable)
|
|
|
|
instance Default SimulationOptions where
|
|
def = SimulationOptions
|
|
{ simParallel = 1
|
|
, simDelay = Normal 0 0
|
|
, simDuration = Normal 10 0
|
|
}
|
|
|
|
|
|
data SimulationContext = SimulationContext
|
|
{ loadOptions :: LoadOptions
|
|
, simulationOptions :: SimulationOptions
|
|
, targetDuration :: DiffTime
|
|
, runtime :: forall m. MonadIO m => m DiffTime
|
|
}
|
|
|
|
|
|
makeLenses_ ''LoadOptions
|
|
makeLenses_ ''SimulationOptions
|
|
makeLenses_ ''SimulationContext
|
|
|
|
|
|
_MapF :: (Finite k, Ord k) => Iso' (Map k v) (k -> Maybe v)
|
|
_MapF = iso (flip Map.lookup) (\f -> Map.fromList $ mapMaybe (\k -> (k, ) <$> f k) universeF)
|
|
|
|
|
|
argsDescr :: [OptDescr (Kleisli IO LoadOptions LoadOptions)]
|
|
argsDescr
|
|
= [ Option ['n', 'p'] ["number", "parallel"] (ReqArg (\(splitArg -> (cloneIndexedTraversal -> f, arg)) -> Kleisli $ return . over f (set _simParallel arg)) "NATURAL") "Number of simulations to run in parallel"
|
|
, Option ['r'] ["run"] (ReqArg (\(ppArg -> sim) -> Kleisli $ return . over (_loadSimulations . at sim) (<|> Just def)) "SIMULATION") "Run the given Simulation"
|
|
, Option ['d'] ["duration"] (ReqArg (\(splitArg -> (cloneIndexedTraversal -> f, arg)) -> Kleisli $ return . over f (set _simDuration arg)) "DURATION") "Try to run each simulation to take up the given duration"
|
|
, Option ['w', 's'] ["wait", "delay", "stagger"] (ReqArg (\(splitArg -> (cloneIndexedTraversal -> f, arg)) -> Kleisli $ return . over f (set _simDelay arg)) "DURATION") "Wait the given time before starting each simulation"
|
|
, Option ['b', 'u'] ["base", "uri"] (ReqArg (\uriStr -> let uri = fromMaybe (error $ "Could not parse URI: " <> uriStr) $ parseURI uriStr in Kleisli $ return . set _loadBaseURI uri ) "URI") "Base URI"
|
|
, Option ['t'] ["token"] (ReqArg (Kleisli . loadTokenFile) "FILE") "File containing bearer token"
|
|
, Option [] ["tid", "term"] (ReqArg (\(ppArg -> tid) -> Kleisli $ return . set _loadTerm tid) "TERM") "TermId"
|
|
, Option [] ["ssh", "school"] (ReqArg (\(ppArg -> ssh) -> Kleisli $ return . set _loadSchool ssh) "SCHOOL") "SchoolId"
|
|
, Option [] ["csh", "course"] (ReqArg (\(ppArg -> csh) -> Kleisli $ return . set _loadCourse csh) "COURSE") "CourseName"
|
|
, Option [] ["shn", "sheet"] (ReqArg (\(ppArg -> shn) -> Kleisli $ return . set _loadSheet shn) "SHEET") "SheetName"
|
|
]
|
|
where
|
|
splitArg :: PathPiece p => String -> (AnIndexedTraversal' LoadSimulation LoadOptions SimulationOptions, p)
|
|
splitArg (Text.pack -> t)
|
|
| (ref, arg) <- Text.breakOn ":" t
|
|
, let refs = Text.splitOn "," ref
|
|
sArg = Text.stripPrefix ":" arg
|
|
, Just refs' <- if | is _Just sArg -> mapM fromPathPiece refs
|
|
| otherwise -> Just []
|
|
, Just arg' <- fromPathPiece $ fromMaybe ref sArg
|
|
= (, arg') $ if
|
|
| null refs' -> _loadSimulations . itraversed
|
|
| otherwise -> _loadSimulations . _MapF . itraversed . indices (`elem` refs') . iplens (fromMaybe def) (const Just)
|
|
| otherwise
|
|
= terror $ "Invalid option argument: " <> t
|
|
|
|
ppArg :: PathPiece p => String -> p
|
|
ppArg (Text.pack -> a) = fromMaybe (terror $ "Invalid option argument: " <> a) $ fromPathPiece a
|
|
|
|
loadTokenFile :: FilePath -> LoadOptions -> IO LoadOptions
|
|
loadTokenFile fp pOpts = do
|
|
token <- Jwt . CBS.filter (not . Char.isSpace) <$> BS.readFile fp
|
|
return $ pOpts & _loadToken ?~ token
|
|
|
|
main :: IO ()
|
|
main = do
|
|
args <- map unpack <$> getArgs
|
|
case getOpt Permute argsDescr args of
|
|
(kl, [], []) -> do
|
|
cfg <- over (mapped . _loadSimulations) (Map.filter $ (> 0) . simParallel) . (`runKleisli` def) . getDual $ foldMap Dual kl
|
|
if | not . Map.null $ loadSimulations cfg
|
|
-> imapM_ (\sim simOpts -> runReaderT (runSimulation sim) (cfg & _loadSimulations . at sim .~ Nothing, simOpts)) $ loadSimulations cfg
|
|
| otherwise -> do
|
|
hPutStrLn stderr $ usageInfo "uniworxload" argsDescr
|
|
exitWith $ ExitFailure 2
|
|
(_, _, errs) -> do
|
|
forM_ errs $ hPutStrLn stderr
|
|
hPutStrLn stderr $ usageInfo "uniworxload" argsDescr
|
|
exitWith $ ExitFailure 2
|
|
|
|
runSimulation :: LoadSimulation -> ReaderT (LoadOptions, SimulationOptions) IO ()
|
|
runSimulation sim = do
|
|
p <- view $ _2 . _simParallel
|
|
replicateConcurrently_ (fromIntegral p) $ do
|
|
d <- view $ _2 . _simDelay
|
|
d' <- sampleNDiffTime d
|
|
|
|
dur <- view $ _2 . _simDuration
|
|
tDuration <- sampleNDiffTime dur
|
|
|
|
let MkFixed us = realToFrac d' :: Micro
|
|
threadDelay $ fromInteger us
|
|
|
|
cTime <- liftIO $ getTime Monotonic
|
|
let running :: forall m. MonadIO m => m DiffTime
|
|
running = do
|
|
cTime' <- liftIO $ getTime Monotonic
|
|
let diff = MkFixed . Clock.toNanoSecs $ cTime' - cTime :: Nano
|
|
MkFixed ps = realToFrac diff :: Pico
|
|
return $ picosecondsToDiffTime ps
|
|
|
|
withReaderT (\(lO, sO) -> SimulationContext lO sO tDuration running) $ runSimulation' sim
|
|
|
|
|
|
delayRemaining :: (MonadReader SimulationContext m, MonadIO m, RealFrac r) => r -> m ()
|
|
delayRemaining p = do
|
|
total <- asks targetDuration
|
|
cTime <- join $ asks runtime
|
|
let remaining = MkFixed . diffTimeToPicoseconds $ total - cTime :: Pico
|
|
MkFixed us = realToFrac $ realToFrac remaining * p :: Micro
|
|
threadDelay $ fromInteger us
|
|
|
|
|
|
runSimulation' :: LoadSimulation -> ReaderT SimulationContext IO ()
|
|
runSimulation' LoadSheetDownload = do
|
|
session <- newLoadSession
|
|
uri <- sheetZipURI
|
|
resp <- liftIO . Session.get session $ uriToString id uri mempty
|
|
print . length $ resp ^. responseBody
|
|
runSimulation' other = terror $ "Not implemented: " <> tshow other
|
|
|
|
|
|
newLoadSession :: ReaderT SimulationContext IO Session
|
|
newLoadSession = do
|
|
LoadOptions{..} <- asks loadOptions
|
|
|
|
let withToken = case loadToken of
|
|
Nothing -> id
|
|
Just (Jwt bs) -> (:) $ traceShowId (hAuthorization, "Bearer " <> bs)
|
|
|
|
|
|
liftIO . Session.newSessionControl (Just mempty) $ tlsManagerSettings
|
|
{ managerModifyRequest = \req -> return $ req { requestHeaders = withToken $ requestHeaders req }
|
|
}
|
|
|
|
sheetZipURI :: ReaderT SimulationContext IO URI
|
|
sheetZipURI = do
|
|
LoadOptions{..} <- asks loadOptions
|
|
let zipURI = nullURI { uriPath = unpack . Text.intercalate "/" $ "." : zipPath }
|
|
where (zipPath, _) = renderRoute . CSheetR loadTerm loadSchool loadCourse loadSheet $ SZipR SheetExercise
|
|
return $ zipURI `relativeTo` loadBaseURI
|