{-# LANGUAGE ApplicativeDo       #-}
{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE RecordWildCards     #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Hyperion.Main where

import           Control.Monad             (unless)
import           Data.Maybe                (isJust)
import           Hyperion.Cluster          (Cluster, ClusterEnv (..),
                                            ProgramInfo (..), runCluster,
                                            runDBWithProgramInfo)
import           Hyperion.Command          (Worker (..), workerOpts)
import           Hyperion.Config           (HyperionConfig (..), newClusterEnv)
import qualified Hyperion.Database         as DB
import           Hyperion.HoldServer       (withHoldServer, newHoldMap)
import qualified Hyperion.Log              as Log
import           Hyperion.Remote           (addressToNodeId,
                                            initWorkerRemoteTable,
                                            runProcessLocalWithRT, worker)
import           Options.Applicative
import           System.Console.Concurrent (withConcurrentOutput)
import           System.Directory          (removeFile)
import           System.Environment        (getEnvironment)
import           System.FilePath.Posix     ((<.>))
import           System.Posix.Process      (getProcessID)

-- | The type for command-line options to 'hyperionMain'. Here @a@ is
-- the type for program-specific options.  In practice we want @a@ to
-- be an instance of 'Show'
data HyperionOpts a =
    HyperionMaster a      -- ^ Constructor for the case of a master
                          -- process, holds program-specific options
  | HyperionWorker Worker -- ^ Constructor for the case of a worker
                          -- process, holds 'Worker' which is parsed
                          -- by 'workerOpts'


-- | Main command-line option parser for 'hyperionMain'.  Returns a
-- 'Parser' that supports commands "worker" and "master", and uses
-- 'workerOpts' or the supplied parser, respectively, to parse the
-- remaining options
hyperionOpts
  :: Parser a -- ^ 'Parser' for program-specific options
  -> Parser (HyperionOpts a)
hyperionOpts :: Parser a -> Parser (HyperionOpts a)
hyperionOpts Parser a
programOpts = Mod CommandFields (HyperionOpts a) -> Parser (HyperionOpts a)
forall a. Mod CommandFields a -> Parser a
subparser (Mod CommandFields (HyperionOpts a) -> Parser (HyperionOpts a))
-> Mod CommandFields (HyperionOpts a) -> Parser (HyperionOpts a)
forall a b. (a -> b) -> a -> b
$ [Mod CommandFields (HyperionOpts a)]
-> Mod CommandFields (HyperionOpts a)
forall a. Monoid a => [a] -> a
mconcat
  [ String
-> ParserInfo (HyperionOpts a)
-> Mod CommandFields (HyperionOpts a)
forall a. String -> ParserInfo a -> Mod CommandFields a
command String
"worker" (ParserInfo (HyperionOpts a) -> Mod CommandFields (HyperionOpts a))
-> ParserInfo (HyperionOpts a)
-> Mod CommandFields (HyperionOpts a)
forall a b. (a -> b) -> a -> b
$
    Parser (HyperionOpts a)
-> InfoMod (HyperionOpts a) -> ParserInfo (HyperionOpts a)
forall a. Parser a -> InfoMod a -> ParserInfo a
info (Parser (HyperionOpts a -> HyperionOpts a)
forall a. Parser (a -> a)
helper Parser (HyperionOpts a -> HyperionOpts a)
-> Parser (HyperionOpts a) -> Parser (HyperionOpts a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Worker -> HyperionOpts a
forall a. Worker -> HyperionOpts a
HyperionWorker (Worker -> HyperionOpts a)
-> Parser Worker -> Parser (HyperionOpts a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser Worker
workerOpts)) (InfoMod (HyperionOpts a) -> ParserInfo (HyperionOpts a))
-> InfoMod (HyperionOpts a) -> ParserInfo (HyperionOpts a)
forall a b. (a -> b) -> a -> b
$
    String -> InfoMod (HyperionOpts a)
forall a. String -> InfoMod a
progDesc String
"Run a worker process. Usually this is run automatically."
  , String
-> ParserInfo (HyperionOpts a)
-> Mod CommandFields (HyperionOpts a)
forall a. String -> ParserInfo a -> Mod CommandFields a
command String
"master" (ParserInfo (HyperionOpts a) -> Mod CommandFields (HyperionOpts a))
-> ParserInfo (HyperionOpts a)
-> Mod CommandFields (HyperionOpts a)
forall a b. (a -> b) -> a -> b
$
    Parser (HyperionOpts a)
-> InfoMod (HyperionOpts a) -> ParserInfo (HyperionOpts a)
forall a. Parser a -> InfoMod a -> ParserInfo a
info (Parser (HyperionOpts a -> HyperionOpts a)
forall a. Parser (a -> a)
helper Parser (HyperionOpts a -> HyperionOpts a)
-> Parser (HyperionOpts a) -> Parser (HyperionOpts a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (a -> HyperionOpts a
forall a. a -> HyperionOpts a
HyperionMaster (a -> HyperionOpts a) -> Parser a -> Parser (HyperionOpts a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser a
programOpts)) (InfoMod (HyperionOpts a) -> ParserInfo (HyperionOpts a))
-> InfoMod (HyperionOpts a) -> ParserInfo (HyperionOpts a)
forall a b. (a -> b) -> a -> b
$
    String -> InfoMod (HyperionOpts a)
forall a. String -> InfoMod a
progDesc String
"Run a master process"
  ]

-- | Same as 'hyperionOpts' but with added @--help@ option and wrapped
-- into 'ParserInfo' (by adding program description).  This now can be
-- used in 'execParser' from "Options.Applicative".
opts :: Parser a -> ParserInfo (HyperionOpts a)
opts :: Parser a -> ParserInfo (HyperionOpts a)
opts Parser a
programOpts = Parser (HyperionOpts a)
-> InfoMod (HyperionOpts a) -> ParserInfo (HyperionOpts a)
forall a. Parser a -> InfoMod a -> ParserInfo a
info (Parser (HyperionOpts a -> HyperionOpts a)
forall a. Parser (a -> a)
helper Parser (HyperionOpts a -> HyperionOpts a)
-> Parser (HyperionOpts a) -> Parser (HyperionOpts a)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Parser a -> Parser (HyperionOpts a)
forall a. Parser a -> Parser (HyperionOpts a)
hyperionOpts Parser a
programOpts) InfoMod (HyperionOpts a)
forall a. InfoMod a
fullDesc

-- | 'hyperionMain' produces an @'IO' ()@ action that runs @hyperion@ and can be
-- assigned to @main@. It performs the following actions
--
--  1. If command-line arguments start with command @master@ then
--
--      - Uses the supplied parser to parse the remaining options into type @a@
--      - Uses the supplied function to extract 'HyperionConfig' from @a@
--      - The data in 'HyperionConfig' is then used for all following actions
--      - Depending on 'HyperionConfig', extra actions may be performed, see 'newClusterEnv'.
--      - Starts a log in 'stderr', and then redirects it to a file
--      - Starts a hold server from "Hyperion.HoldServer"
--      - Uses 'DB.setupKeyValTable' to setup a "Hyperion.Database.KeyValMap" in the program database
--      - Runs the supplied @'Cluster' ()@ action
--      - Cleans up the copy of the executable, if exists (see 'newClusterEnv').
--
--  2. If command-line arguments start with command @worker@ then
--
--      - Extracts 'Worker' from the rest of the command-line args.
--      - Logs 'ServiceId' of the worker and the system environment to worker log file.
--      - Runs @'worker' (...) :: 'Process' ()@ that connects to the master and waits for a
--        'Hyperion.Remote.ShutDown' message (see 'worker' for details).
--        While waiting, the master can run computations on the node. Low-level
--        functions for this are implemented in "Hyperion.Remote", and some
--        higher-level functions in "Hyperion.HasWorkers"
hyperionMain
  :: Show a
  => Parser a
  -> (a -> HyperionConfig)
  -> (a -> Cluster ())
  -> IO ()
hyperionMain :: Parser a -> (a -> HyperionConfig) -> (a -> Cluster ()) -> IO ()
hyperionMain Parser a
programOpts a -> HyperionConfig
mkHyperionConfig a -> Cluster ()
clusterProgram = IO () -> IO ()
forall (m :: * -> *) a. (MonadIO m, MonadMask m) => m a -> m a
withConcurrentOutput (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
  ParserInfo (HyperionOpts a) -> IO (HyperionOpts a)
forall a. ParserInfo a -> IO a
execParser (Parser a -> ParserInfo (HyperionOpts a)
forall a. Parser a -> ParserInfo (HyperionOpts a)
opts Parser a
programOpts) IO (HyperionOpts a) -> (HyperionOpts a -> IO ()) -> IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
  HyperionWorker Worker{String
Text
ServiceId
workerLogFile :: Worker -> String
workerService :: Worker -> ServiceId
workerMasterAddress :: Worker -> Text
workerLogFile :: String
workerService :: ServiceId
workerMasterAddress :: Text
..} -> do
    String -> IO ()
Log.redirectToFile String
workerLogFile
    Text -> ServiceId -> IO ()
forall a (m :: * -> *). (Show a, MonadIO m) => Text -> a -> m ()
Log.info Text
"Starting service" ServiceId
workerService
    Text -> [(String, String)] -> IO ()
forall a (m :: * -> *). (Show a, MonadIO m) => Text -> a -> m ()
Log.info Text
"Environment" ([(String, String)] -> IO ()) -> IO [(String, String)] -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO [(String, String)]
getEnvironment
    let masterNid :: NodeId
masterNid = Text -> NodeId
addressToNodeId Text
workerMasterAddress
    RemoteTable -> Process () -> IO ()
forall a. RemoteTable -> Process a -> IO a
runProcessLocalWithRT
      (Maybe NodeId -> RemoteTable
initWorkerRemoteTable (NodeId -> Maybe NodeId
forall a. a -> Maybe a
Just NodeId
masterNid))
      (NodeId -> ServiceId -> Process ()
worker NodeId
masterNid ServiceId
workerService)
  HyperionMaster a
args -> do
    let hyperionConfig :: HyperionConfig
hyperionConfig = a -> HyperionConfig
mkHyperionConfig a
args
    HoldMap
holdMap <- IO HoldMap
newHoldMap
    HoldMap -> (Int -> IO ()) -> IO ()
forall a. HoldMap -> (Int -> IO a) -> IO a
withHoldServer HoldMap
holdMap ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
holdPort -> do
      (clusterEnv :: ClusterEnv
clusterEnv@ClusterEnv{Int
LockMap
SbatchOptions
Pool
ProgramInfo
SbatchOptions -> ProgramInfo -> WorkerLauncher JobId
clusterLockMap :: ClusterEnv -> LockMap
clusterDatabaseRetries :: ClusterEnv -> Int
clusterDatabasePool :: ClusterEnv -> Pool
clusterJobOptions :: ClusterEnv -> SbatchOptions
clusterProgramInfo :: ClusterEnv -> ProgramInfo
clusterWorkerLauncher :: ClusterEnv -> SbatchOptions -> ProgramInfo -> WorkerLauncher JobId
clusterLockMap :: LockMap
clusterDatabaseRetries :: Int
clusterDatabasePool :: Pool
clusterJobOptions :: SbatchOptions
clusterProgramInfo :: ProgramInfo
clusterWorkerLauncher :: SbatchOptions -> ProgramInfo -> WorkerLauncher JobId
..}, String
hyperionExecutable) <- HyperionConfig -> HoldMap -> Int -> IO (ClusterEnv, String)
newClusterEnv HyperionConfig
hyperionConfig HoldMap
holdMap Int
holdPort
      let progId :: ProgramId
progId = ProgramInfo -> ProgramId
programId ProgramInfo
clusterProgramInfo
          masterLogFile :: String
masterLogFile = ProgramInfo -> String
programLogDir ProgramInfo
clusterProgramInfo String -> String -> String
<.> String
"log"
      ProcessID
pid <- IO ProcessID
getProcessID
      let logMasterInfo :: IO ()
logMasterInfo = do
            Text -> ProgramId -> IO ()
forall a (m :: * -> *). (Show a, MonadIO m) => Text -> a -> m ()
Log.info Text
"Program id" ProgramId
progId
            Text -> ProcessID -> IO ()
forall a (m :: * -> *). (Show a, MonadIO m) => Text -> a -> m ()
Log.info Text
"Process id" ProcessID
pid
            Text -> a -> IO ()
forall a (m :: * -> *). (Show a, MonadIO m) => Text -> a -> m ()
Log.info Text
"Program arguments" a
args
            Text -> String -> IO ()
forall a (m :: * -> *). (Show a, MonadIO m) => Text -> a -> m ()
Log.info Text
"Using database" (ProgramInfo -> String
programDatabase ProgramInfo
clusterProgramInfo)
            Text -> Int -> IO ()
forall a (m :: * -> *). (Show a, MonadIO m) => Text -> a -> m ()
Log.info Text
"Running hold server on port" Int
holdPort
      Text -> IO ()
forall (m :: * -> *). MonadIO m => Text -> m ()
Log.rawText Text
"--------------------------------------------------------------------------------\n"
      IO ()
logMasterInfo
      Text -> String -> IO ()
forall a (m :: * -> *). (Show a, MonadIO m) => Text -> a -> m ()
Log.info Text
"Logging to" String
masterLogFile
      IO ()
Log.flush
      String -> IO ()
Log.redirectToFile String
masterLogFile
      IO ()
logMasterInfo
      ProgramInfo -> ReaderT DatabaseConfig IO () -> IO ()
forall a. ProgramInfo -> ReaderT DatabaseConfig IO a -> IO a
runDBWithProgramInfo ProgramInfo
clusterProgramInfo ReaderT DatabaseConfig IO ()
forall (m :: * -> *) env.
(MonadIO m, MonadReader env m, HasDB env, MonadCatch m) =>
m ()
DB.setupKeyValTable
      ClusterEnv -> Cluster () -> IO ()
forall a. ClusterEnv -> Cluster a -> IO a
runCluster ClusterEnv
clusterEnv (a -> Cluster ()
clusterProgram a
args)
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Maybe String -> Bool
forall a. Maybe a -> Bool
isJust (HyperionConfig -> Maybe String
hyperionCommand HyperionConfig
hyperionConfig)) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
        String -> IO ()
removeFile String
hyperionExecutable
      Text -> ProgramId -> IO ()
forall a (m :: * -> *). (Show a, MonadIO m) => Text -> a -> m ()
Log.info Text
"Finished" ProgramId
progId