fradrive/backend/load/Load.hs

370 lines
15 KiB
Haskell

-- SPDX-FileCopyrightText: 2022 Gregor Kleen <gregor.kleen@ifi.lmu.de>,Steffen Jost <jost@tcs.ifi.lmu.de>
--
-- SPDX-License-Identifier: AGPL-3.0-or-later
{-# 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, responseStatus)
import Utils.Form (FormIdentifier(..))
import Handler.Admin.Test.Download (generateDownload', seedNew)
import System.Console.GetOpt
import qualified Data.Text as Text
import qualified Data.Map.Strict as Map
import Data.Ratio ((%))
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.Lazy as Lazy (ByteString)
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.Types (FormValue(..))
import Network.Wreq.Session (Session)
import qualified Network.Wreq.Session as Session
import Network.HTTP.Client.MultipartFormData (partFileRequestBody)
import Network.HTTP.Client.TLS (tlsManagerSettings)
import qualified Text.HTML.Scalpel as Scalpel
import qualified Data.Conduit.Combinators as C
import Data.List (genericLength)
import qualified Control.Retry as Retry
data Normal k = Normal
{ dAvg :: k
, dRelDev :: Centi
} deriving (Eq, Ord, Read, Show, Generic)
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
scaleIntegral :: Integral n => n -> Centi -> n
scaleIntegral n s = round $ toRational n * toRational s
sampleIntegral :: (Random.MonadSplit g m, RandomGen g, Integral n) => Normal n -> m n
sampleIntegral = sampleN scaleIntegral
instance PathPiece DiffTime where
toPathPiece = (toPathPiece :: Pico -> Text) . MkFixed . diffTimeToPicoseconds
fromPathPiece t = fromPathPiece t <&> \(MkFixed ps :: Pico) -> picosecondsToDiffTime ps
data LoadSimulation
= LoadSheetDownload
| LoadSheetSubmission
deriving (Eq, Ord, Read, Show, Enum, Bounded, Generic)
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
, loadUploadChunks :: Normal Natural, loadUploadChunkSize :: Normal Natural
} deriving (Eq, Ord, Show, Generic)
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"
, loadUploadChunks = Normal 48 0.11
, loadUploadChunkSize = Normal (2^16) 0
}
data SimulationOptions = SimulationOptions
{ simParallel :: Natural
, simDelay, simDuration :: Normal DiffTime
} deriving (Eq, Ord, Show, Generic)
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"
, Option [] ["chunks"] (ReqArg (\(ppArg -> cs) -> Kleisli $ return . set _loadUploadChunks cs) "NATURAL") "Number of chunks to upload"
, Option [] ["chunk-size"] (ReqArg (\(ppArg -> cs) -> Kleisli $ return . set _loadUploadChunkSize cs) "NATURAL") "Size of chunks to upload"
]
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
delays <- replicateM (fromIntegral p) $ do
d <- view $ _2 . _simDelay
sampleNDiffTime d
forConcurrently_ ([1..p] `zip` sort delays) $ \(n, d') -> do
begin <- liftIO getCurrentTime
dur <- view $ _2 . _simDuration
tDuration <- sampleNDiffTime dur
let MkFixed us = realToFrac d' :: Micro
threadDelay $ fromInteger us
start <- liftIO getCurrentTime
print ("start", n, diffUTCTime start begin, utctDayTime start)
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
end <- liftIO getCurrentTime
print ("end", n, diffUTCTime start begin, diffUTCTime end start)
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
void . evaluate $! resp
-- print . length $ resp ^. responseBody
runSimulation' LoadSheetSubmission = do
LoadOptions{..} <- asks loadOptions
session <- newLoadSession
let formURI = formURI' `relativeTo` loadBaseURI
where formURI' = nullURI { uriPath = unpack . Text.intercalate "/" $ "." : formPath }
(formPath, _) = renderRoute $ CSheetR loadTerm loadSchool loadCourse loadSheet SubmissionNewR
resp <- liftIO . httpRetry . Session.get session $ uriToString id formURI mempty
void . evaluate $! resp
procStart <- join $ asks runtime
-- Just formData <- return . getFormData FIDsubmission $ resp ^. responseBody
-- Just addButtonData <- return . flip (runFormScraper FIDsubmission) (resp ^. responseBody) $ do
-- let btnSel = "button" Scalpel.@: [Scalpel.hasClass "btn-mass-input-add"]
-- name <- Scalpel.attr "name" btnSel
-- value <- Scalpel.attr "value" btnSel
-- guard $ value == "add__0__0"
-- return $ toStrict name := value
-- let miData = addButtonData : map addEmail formData
-- where addEmail dat@(name := _)
-- | "__add__0__fields__emails" `isSuffixOf` name = name := ("loadtest@example.invalid" :: Text)
-- | otherwise = dat
-- resp2 <- liftIO $ Session.post session (uriToString id formURI mempty) miData
-- Just formData2 <- return . getFormData FIDsubmission $ resp2 ^. responseBody
Just formData2 <- return . getFormData FIDsubmission $ resp ^. responseBody
uploadSeed <- liftIO seedNew
chunkCount <- sampleIntegral loadUploadChunks
chunks <- replicateM (fromIntegral chunkCount) $ sampleIntegral loadUploadChunkSize
simCtx <- ask
let fileUploadPart = requestBodySourceChunked $
yieldMany (zip [0..] chunks)
.| runReaderC simCtx
( C.mapM $ \(ci, cs) ->
fromIntegral cs <$ delayRemaining ((1 % max 1 (genericLength chunks - ci)) :: Rational)
)
.| generateDownload' uploadSeed
-- print $ ala Sum foldMap chunks
Just fileData <- return . flip (runFormScraper FIDsubmission) (resp ^. responseBody) $ do
let fileSel = "input" Scalpel.@: ["type" Scalpel.@= "file"]
name <- Scalpel.attr "name" fileSel
return $ partFileRequestBody (decodeUtf8 $ toStrict name) "loadtest.bin" fileUploadPart
let subData = (:) fileData $ formData2 >>= \(name := (renderFormValue -> value)) -> do
guard $ name /= encodeUtf8 (fileData ^. partName)
return $ partBS (decodeUtf8 name) value
void . evaluate $! subData
procEnd <- join $ asks runtime
print ("proc", procEnd - procStart)
resp3 <- liftIO . httpRetry $ Session.post session (uriToString id formURI mempty) subData
void . evaluate $! resp3
where
httpRetry act = Retry.recovering policy handlers $ \Retry.RetryStatus{..} -> do
putStrLn $ "httpRetry; rsIterNumber = " <> tshow rsIterNumber
act
where policy = Retry.fullJitterBackoff 1e3 & Retry.limitRetriesByCumulativeDelay 10e6
handlers = Retry.skipAsyncExceptions `snoc` Retry.logRetries suggestRetry logRetry
suggestRetry :: forall m. Monad m => SomeException -> m Bool
suggestRetry _ = return True
logRetry :: forall e m.
( Exception e
, MonadIO m
)
=> Bool -- ^ Will retry
-> e
-> Retry.RetryStatus
-> m ()
logRetry shouldRetry err status = liftIO . putStrLn . pack $ Retry.defaultLogMsg shouldRetry err status
-- runSimulation' other = terror $ "Not implemented: " <> tshow other
runFormScraper :: FormIdentifier -> Scalpel.Scraper Lazy.ByteString a -> Lazy.ByteString -> Maybe a
runFormScraper fid innerS = fmap join . flip Scalpel.scrapeStringLike $
fmap listToMaybe . Scalpel.chroots "form" $ do
fid' <- Scalpel.attr "value" $ "input" Scalpel.@: ["name" Scalpel.@= "form-identifier"]
guard $ fid' == encodeUtf8 (fromStrict $ toPathPiece fid)
innerS
getFormData :: FormIdentifier -> Lazy.ByteString -> Maybe [FormParam]
getFormData = flip runFormScraper $
Scalpel.chroots "input" $ do
name <- Scalpel.attr "name" Scalpel.anySelector
value <- Scalpel.attr "value" Scalpel.anySelector <|> pure ""
return $ toStrict name := value
newLoadSession :: ReaderT SimulationContext IO Session
newLoadSession = do
LoadOptions{..} <- asks loadOptions
let withToken = case loadToken of
Nothing -> id
Just (Jwt bs) -> (:) (hAuthorization, "Bearer " <> bs) . filter ((/= hAuthorization) . fst)
liftIO . Session.newSessionControl (Just mempty) $ tlsManagerSettings
{ managerModifyRequest = \req -> return $ req { requestHeaders = withToken $ requestHeaders req }
, managerResponseTimeout = responseTimeoutNone
}
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 -- FIXME: Broken with ApprootUserGenerated
return $ zipURI `relativeTo` loadBaseURI