content: {-# LANGUAGE BangPatterns, DeriveDataTypeable, OverloadedStrings,
      RecordWildCards, ScopedTypeVariables #-}
  
  module Main (main) where
  
  import Control.Applicative ((<$>))
  import Control.Concurrent.MVar
  import Control.DeepSeq (rnf)
  import Control.Exception (bracket, catch, evaluate, finally)
  import Control.Monad (forM_, unless)
  import Control.Monad.Trans.Resource (ResourceT)
  import Data.Aeson ((.=), encode, object)
  import Data.Char (toLower)
  import Data.Maybe (catMaybes)
  import Data.Text (Text, pack)
  import Data.Text.Encoding (encodeUtf8)
  import Data.Text.Lazy.Builder (toLazyText)
  import Data.Time.Clock.POSIX (getPOSIXTime)
  import Data.Version (showVersion)
  import Network.HTTP.LoadTest (NetworkError(..), Req(..))
  import Network.HTTP.LoadTest.Analysis (analyseBasic, analyseFull)
  import Network.HTTP.LoadTest.Environment (environment)
  import Network.HTTP.LoadTest.Report
  import Network.Socket (withSocketsDo)
  import Paths_pronk (version)
  import Prelude hiding (catch)
  import System.CPUTime (getCPUTime)
  import System.Console.CmdArgs
  import System.Exit (ExitCode(ExitFailure), exitWith)
  import System.IO (Handle, IOMode(..), hClose, hPutStrLn, openFile, stderr, stdout)
  import qualified Data.Aeson.Generic as G
  import qualified Data.ByteString.Char8 as B
  import qualified Data.ByteString.Lazy.Char8 as BL
  import qualified Data.Text.Format as T
  import qualified Data.Text.Lazy.IO as TL
  import qualified Network.HTTP.Conduit as E
  import qualified Network.HTTP.LoadTest as LoadTest
  
  data Args = Args {
        concurrency :: Int
      , method :: Maybe String
      , num_requests :: Int
      , requests_per_second :: Double
      , timeout :: Double
      , url :: String
  
      , from_file :: Maybe FilePath
      , literal :: Maybe String
  
      , bootstrap :: Bool
      , dump_events :: Maybe FilePath
      , output :: Maybe FilePath
      , template :: FilePath
      , json :: Maybe FilePath
      } deriving (Eq, Show, Typeable, Data)
  
  pronk_version :: String
  pronk_version = showVersion version
  
  defaultArgs :: Args
  defaultArgs = Args {
                  concurrency = 1
                  &= groupname "Load testing"
                  &= help "Number of requests to issue concurrently"
                , method = def &= typ "METHOD"
                  &= help "HTTP method to use (GET, POST, ...)"
                , num_requests = 1
                  &= help "Total number of requests to issue"
                , requests_per_second = def
                  &= help "Maximum request rate to sustain"
                , timeout = 60 &= typ "SECS"
                  &= help "Time to wait before killing a connection"
                , url = def &= argPos 0 &= typ "URL"
  
                -- --------------------------------------------------
                , from_file = def &= typ "FILE"
                  &= groupname "Supplying a request body"
                  &= help "Use file contents as request body"
                , literal = def &= typ "STRING"
                  &= help "Use given text as request body"
  
                -- --------------------------------------------------
                , bootstrap = def
                  &= groupname "Analysis of results"
                  &= help "Statistically robust analysis of results"
                , dump_events = def &= typ "FILE"
                  &= help "Save raw events in CSV format"
                , output = def &= typ "FILE"
                  &= help "Write report to named file"
                , template = "report.tpl" &= typ "FILE"
                  &= help "Use the given report template"
                , json = def &= typ "FILE"
                  &= help "Save analysis in JSON format"
                } &= verbosity
                  &= summary ("Pronk " ++ pronk_version ++
                              " - a modern HTTP load tester")
  
  fromArgs :: Args -> E.Request (ResourceT IO) -> LoadTest.Config
  fromArgs Args{..} req =
      LoadTest.Config {
        LoadTest.concurrency = concurrency
      , LoadTest.numRequests = num_requests
      , LoadTest.requestsPerSecond = requests_per_second
      , LoadTest.timeout = timeout
      , LoadTest.request = Req req
      }
  
  main :: IO ()
  main = withSocketsDo $ do
    as@Args{..} <- cmdArgs $ defaultArgs &= program "pronk"
    validateArgs as
    cfg <- fromArgs as <$> createRequest as
    (run,time) <- timed "tested" $ LoadTest.run cfg
    case run of
      Left [NetworkError err] -> fatal (show err)
      Left errs -> do
        T.hprint stderr "Errors:\n" ()
        forM_ errs $ \(NetworkError err) -> T.hprint stderr "  {}\n" [show err]
        exitWith (ExitFailure 1)
      Right results -> do
        whenNormal $ T.print "analysing results\n" ()
        (analysis,_) <- timed "analysed" $ do
                        r <- if bootstrap
                             then Right <$> analyseFull results time
                             else return . Left $ analyseBasic results time
                        evaluate $ rnf r
                        return r
        env <- environment
        let dump = object [ "config" .= cfg
                          , "environment" .= env
                          , "analysis" .= G.toJSON analysis ]
        maybeWriteFile json $ \h -> BL.hPut h (BL.append (encode dump) "\n")
        maybeWriteFile dump_events $ \h ->
            TL.hPutStr h . toLazyText . csvEvents $ results
        maybeWriteFile output $ \h -> either (writeReport template h time)
                                             (writeReport template h time) analysis
        whenNormal $ do
          reportEvents stdout results
          either (reportBasic stdout) (reportFull whenLoud stdout)
                 analysis
  
  maybeWriteFile :: Maybe FilePath -> (Handle -> IO ()) -> IO ()
  maybeWriteFile (Just "-") act = act stdout
  maybeWriteFile (Just p)   act = bracket (openFile p WriteMode) hClose act
  maybeWriteFile _          _   = return ()
  
  validateArgs :: Args -> IO ()
  validateArgs Args{..} = do
    let p !? what | p         = Nothing
                  | otherwise = Just $ "Argument to " ++ what
        infix 1 !?
        problems = catMaybes [
           concurrency > 0 !? "--concurrency must be positive"
         , num_requests > 0 !? "--num-requests must be positive"
         , requests_per_second >= 0 !? "--requests-per-second cannot be negative"
         , timeout >= 0 !? "--timeout cannot be negative"
         ]
    forM_ problems $ hPutStrLn stderr . ("Error: " ++)
    unless (null problems) $ exitWith (ExitFailure 1)
  
  createRequest :: Args -> IO (E.Request (ResourceT IO))
  createRequest Args{..} = do
    req0 <- E.parseUrl url `catch` \(e::E.HttpException) ->
            fatal $ "could not parse URL - " ++
                  case e of
                    E.InvalidUrlException _ s -> map toLower s
                    _ -> show e
    let check Nothing       = return "POST"
        check (Just "POST") = return "POST"
        check (Just "PUT")  = return "PUT"
        check _      = fatal "only POST or PUT may have a body"
        req = req0 { E.redirectCount = 0, E.checkStatus = \_ _ -> Nothing }
    case (from_file, literal) of
      (Nothing,Nothing) -> return req { E.method = maybe "GET" B.pack method }
      (Just f,Nothing) -> do
        s <- B.readFile f
        meth <- check method
        return req { E.method = meth
                   , E.requestBody = E.RequestBodyBS s }
      (Nothing,Just s) -> do
        meth <- check method
        return req { E.method = meth
                   , E.requestBody = E.RequestBodyBS . encodeUtf8 . pack $ s
                   }
      _ -> do
        hPutStrLn stderr "Error: --literal and --from-file are mutually exclusive"
        exitWith (ExitFailure 1)
  
  timed :: Text -> IO a -> IO (a,Double)
  timed desc act = do
    t <- newEmptyMVar
    startCPU <- getCPUTime
    startWall <- getPOSIXTime
    ret <- act `finally` do
      endCPU <- getCPUTime
      endWall <- getPOSIXTime
      let elapsedCPU  = fromIntegral (endCPU - startCPU) / 1e12
          elapsedWall = realToFrac $ endWall - startWall
          ratio = elapsedCPU / elapsedWall
      whenNormal $
        -- Try to work around the 64-bit Mac getCPUTime bug
        -- http://hackage.haskell.org/trac/ghc/ticket/4970
        if ratio > 0 && ratio < 32
        then T.print "{} in {} ({}% CPU)\n"
                 (desc, buildTime 4 elapsedWall,
                  T.fixed 1 $ 100 * elapsedCPU / elapsedWall)
        else T.print "{} in {}\n"
                 (desc, buildTime 4 elapsedWall)
      putMVar t elapsedWall
    ((,) ret) <$> takeMVar t
  
  fatal :: String -> IO a
  fatal e = do
    T.hprint stderr "Error: {}\n" (T.Only e)
    exitWith (ExitFailure 1)
  