content: module Test.Stress.Driver
    ( driver
    ) where
  
  import Control.Concurrent
  import Control.Concurrent.Async
  import Control.Exception
  import Control.Monad
  import Control.Monad.Catch (ExitCase(..), generalBracket)
  import Data.IORef
  import System.Environment
  import System.Exit
  import System.Process
  import System.Random
  
  import Test.Stress.Common
  import Test.Stress.Driver.Summary
  
  -------------------------------------------------------------------------------
  -- Top-level
  -------------------------------------------------------------------------------
  
  -- | Run the automatic stress test suite
  --
  -- See documentation in the source repository for details.
  driver :: Bool -> Bool -> Maybe FilePath -> Int -> IO ()
  driver v genCharts mwd duration = do
      putStrLn $
        "Running stress test driver for " ++ show duration ++ " seconds..."
      exitCodeRef <- newIORef ExitSuccess
      bracket
        ( do
            runningServers <- mapM (forkComponent v genCharts mwd) servers
            runningClients <- mapM (forkComponent v genCharts mwd) clients
            return $ runningServers ++ runningClients
        )
        (\running -> do
          say v "(driver) stopping all components"
          cancelMany running
          say v "(driver) stopped all components"
        )
        (\running -> do
          result <- race (threadDelay (duration * 1_000_000)) (waitAnyCatch running)
          case result of
            Right (_, Left e)
              | Just (TestFailure _) <- fromException e -> do
                writeIORef exitCodeRef (ExitFailure 1)
                throwIO e
              | otherwise ->
                return ()
            _ -> return ()
        )
        `catch`
          \case
            e | Just f@TestFailure{} <- fromException e -> do
                  putStrLn $ "stress test failed: " ++ show f
              | otherwise -> do
                  say v $ "(driver) exiting cleanly, got exception: " ++ show e
  
      -- At this point, the heap profiles should all be written in the working
      -- directory. We convert them each to charts of total heap usage over time
      -- and combine into a summary document
      when genCharts $
        createSummaryPlots v mwd
      putStrLn "Done"
      exitCode <- readIORef exitCodeRef
      exitWith exitCode
  
  -------------------------------------------------------------------------------
  -- Auxiliary
  -------------------------------------------------------------------------------
  
  data Component = Component {
        -- | Client or server?
        componentType   :: ClientServer
  
        -- | What port to bind/connect to
      , componentPort   :: Int
  
        -- | Use TLS?
      , componentSecure :: Bool
  
        -- | Should the component stay running indefinitely?
        --
        -- If 'False', component will be killed and restarted at random intervals
        -- and it will not be configured to write a heap profile.
      , componentStable :: Bool
  
        -- | Heap size limit in Megabytes (e.g. @Just 15@ becomes @-M15m@)
        --
        -- Set to 'Nothing' for no limit.
      , componentLimit  :: Maybe Int
  
        -- | Name
        --
        -- Determines the name of the heap profile file that will be written out.
      , componentName   :: String
      }
    deriving (Show)
  
  data ClientServer =
        Client {
            -- | Which compression should be used?
            --
            -- We don't currently use this in the driver to avoid running too many
            -- clients, but it is left here as an option.
            clientCompr :: Maybe String
  
            -- | Specify which calls should be executed
          , clientFlags :: [String]
          }
      | Server
    deriving (Show, Eq)
  
  newtype TestFailure = TestFailure String
    deriving (Show)
    deriving anyclass Exception
  
  cmd :: Bool -> ClientServer -> [String]
  cmd v t = mconcat [
        [ "-v" | v ]
      , case t of
          Client mcompr flags ->
            mconcat [
                [ "client" ]
              , [ "--"++compr | Just compr <- [mcompr] ]
              , flags
              ]
          Server ->
            [ "server" ]
      ]
  
  forkComponent :: Bool -> Bool -> Maybe FilePath -> Component -> IO (Async ())
  forkComponent v genCharts mwd c = do
      say v $ "(driver) forking component " ++ show c
      async $
        runComponent v genCharts mwd c `catch`
          \case
            (e :: SomeException)
              | Just AsyncCancelled <- fromException e ->
                  say v $ "(driver) cancelled forked component " ++ show c
              | otherwise -> do
                  say v $ "(driver) component exited with " ++ show e
                  throwIO e
  
  runComponent :: Bool -> Bool -> Maybe FilePath -> Component -> IO ()
  runComponent v genCharts mwd c@Component{..} = do
      say' $ "starting " ++ componentName
      exe <- getExecutablePath
      let cp = proc exe (mconcat [
              cmd v componentType
            , [ "--port=" ++ show componentPort ]
            , [ "--secure" | componentSecure ]
            , if componentType == Server then
                [ "+RTS", "-N", "-RTS" ]
              else
                []
            , filter (const $ componentStable && genCharts) [
                  "+RTS"
                , "-l"
                , "-hT"
                , "-ol" ++ componentName ++ ".eventlog"
                , "-RTS"
                ]
            , case componentLimit of
                Just limit -> [
                    "+RTS"
                  , "-M" ++ show limit ++ "m"
                  , "-RTS"
                  ]
                Nothing ->
                  []
            ])
      (k, ec) <-
        generalBracket
          (do
            (_, _, _, ph) <-
              createProcess_
                ("runComponent (" ++ componentName ++ ")")
                cp {cwd = mwd}
            return ph
          )
          (\ph ec -> do
            terminateProcess ph
            return ec
          )
          watchComponent
      case ec of
        ExitCaseSuccess _ -> k
        ExitCaseException e
          | Just (TestFailure _) <- fromException e ->
            throwIO e
          | otherwise ->
            k
        ExitCaseAbort -> throwIO $ TestFailure $ componentName ++ " aborted"
    where
      watchComponent :: ProcessHandle -> IO (IO ())
      watchComponent ph =
          if componentStable then do
            say' $ "watching " ++ componentName
            ec <- waitForProcess ph
            case ec of
              ExitFailure _ -> do
                say' $ "unexpected ExitFailure from " ++ show componentName
                throwIO $ TestFailure componentName
              ExitSuccess -> do
                say' $ componentName ++ " exited successfully, rerunning"
                return $ runComponent v genCharts mwd c
          else do
            d <- randomRIO (16_000_000, 20_000_000)
            say' $
              "waiting " ++ show d ++ "µs before killing " ++ show componentName
            threadDelay d
            mec <- getProcessExitCode ph
            case mec of
              Just (ExitFailure _) -> do
                say' $ "unexpected ExitFailure from " ++ show componentName
                throwIO $ TestFailure componentName
              _ -> do
                say' $ "terminating " ++ show componentName
                terminateProcess ph
                d' <- randomRIO (200_000, 500_000)
                say' $
                  "waiting " ++ show d ++ "µs before restarting " ++
                  show componentName
                threadDelay d'
                return $ runComponent v genCharts mwd c
  
      say' :: String -> IO ()
      say' = say v . ("(driver) " ++)
  
  servers :: [Component]
  servers = [
        Component {
            componentType   = Server
          , componentPort   = 50000
          , componentSecure = False
          , componentStable = False
          , componentLimit  = Just 400
          , componentName   = "server-unstable-insecure"
          }
      , Component {
            componentType   = Server
          , componentPort   = 50001
          , componentSecure = True
          , componentStable = False
          , componentLimit  = Just 400
          , componentName   = "server-unstable-secure"
          }
      , Component {
            componentType   = Server
          , componentPort   = 50002
          , componentSecure = False
          , componentStable = True
          , componentLimit  = Just 400
          , componentName   = "server-stable-insecure"
          }
      , Component {
            componentType   = Server
          , componentPort   = 50003
          , componentSecure = True
          , componentStable = True
          , componentLimit  = Just 400
          , componentName   = "server-stable-secure"
          }
      ]
  
  clients :: [Component]
  clients = [
        Component {
          componentType   = Client Nothing flags
        , componentPort   = snd portSecurity
        , componentSecure = fst portSecurity
        , componentStable = stable
        , componentLimit  = Just 60
        , componentName = mconcat [
              "client-"
            , show (snd portSecurity) ++ "-"
            , if stable then "stable-" else "unstable-"
            , if fst portSecurity then "secure-" else "insecure-"
            , nameStr
            ]
        }
      | (flags, nameStr) <- [
            ( [ indefinitely "--num-connections"
              , "--non-streaming"
              ]
            , "non-streaming-many-connections"
            )
          , ( [ indefinitely "--num-calls"
              , "--non-streaming"
              ]
            , "non-streaming-many-calls"
            )
          , ( [ indefinitely "--client-streaming"
              ]
            , "client-stream"
            )
          , ( [ indefinitely "--num-calls"
              , "--client-streaming=10"
              ]
            , "client-stream-many-calls"
            )
          , ( [ indefinitely "--server-streaming"
              ]
            , "server-stream"
            )
          , ( [ indefinitely "--num-calls"
              , "--server-streaming=10"
              ]
            , "server-stream-many-calls"
            )
          , ( [ indefinitely "--bidi-streaming"
              ]
            , "bidi-stream"
            )
          , ( [ indefinitely "--num-calls"
              , "--bidi-streaming=10"
              ]
            , "bidi-stream-many-calls"
            )
          ]
      , portSecurity <- [
            (False, 50000)
          , (False, 50002)
          , (True , 50001)
          , (True , 50003)
          ]
      , stable <- [True, False]
      ]
    where
      indefinitely :: String -> String
      indefinitely = (++ "=100000000")
  