content: {-
  Copyright   : (c) Jean-Christophe Mincke, 2016-2017
  
  This Source Code Form is subject to the terms of the Mozilla Public
  License, v. 2.0. If a copy of the MPL was not distributed with this
  file, You can obtain one at http://mozilla.org/MPL/2.0/.
  -}
  
  
  
  {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE ImpredicativeTypes #-}
  {-# LANGUAGE KindSignatures #-}
  {-# LANGUAGE MultiParamTypeClasses #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE TemplateHaskell #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  
  
  module Control.Distributed.Blast.Runner.CloudHaskell
  (
    slaveProcess
    , runRec
    , RpcConfig (..)
    , MasterConfig (..)
    , SlaveConfig (..)
  )
  where
  
  import            Control.Concurrent (threadDelay)
  import            Control.Concurrent.Async
  import            Control.Concurrent.STM.TChan
  import            Control.DeepSeq
  import            Control.Monad.IO.Class
  import            Control.Monad.Logger
  import            Control.Monad.STM
  
  import            Control.Distributed.Process hiding (newChan, send)
  import            Control.Distributed.Process.Backend.SimpleLocalnet
  import            Control.Distributed.Process.Extras.Internal.Types  (ExitReason(..))
  import            Control.Distributed.Process.Extras.Time (Delay(..))
  import            Control.Distributed.Process.ManagedProcess as DMP hiding (runProcess, stop)
  
  import            Data.Binary
  import qualified  Data.List as L
  import qualified  Data.Map as M
  import qualified  Data.Serialize as S
  import            Data.Typeable
  
  import            GHC.Generics (Generic)
  
  import            Control.Distributed.Blast
  import            Control.Distributed.Blast.Distributed.Interface
  
  
  -- | General configuration.
  data RpcConfig = MkRpcConfig {
    commonConfig :: Config          -- ^ Blast configuration.
    , masterConfig :: MasterConfig  -- ^ Specific configuration for master.
    , slaveConfig :: SlaveConfig    -- ^ Specific configuration for slaves.
    }
  
  -- | Master configuration.
  data MasterConfig = MkMasterConfig {
    masterLogger :: forall m a. (MonadIO m) => LoggingT m a -> m a  -- ^ Logger.
    }
  
  -- | Master configuration.
  data SlaveConfig = MkSlaveConfig {
    slaveLogger :: forall m a. (MonadIO m) => LoggingT m a -> m a   -- ^ Logger.
    }
  
  data RpcRequestControl =
    RpcRequestControlStop
  
  data RpcResponseControl =
    RpcResponseControlError String
    |RpcResponseControlStopped
  
  data SlaveInfo = MkSlaveInfo {
    _slaveIndex :: Int
    , _slaveNodeId :: NodeId
    , _slaveRequestChannel :: TChan (Either RpcRequestControl SlaveRequest)
    , _slaveResponseChannel :: TChan (Either RpcResponseControl SlaveResponse)
    }
  
  
  data SlaveControl =
    SlaveCtrlStop
    deriving (Typeable, Generic, Show)
  
  instance Binary SlaveControl
  
  
  
  -- | Defines the CloudHaskell closure.
  slaveProcess :: forall a b . (S.Serialize a, Typeable a, Typeable b) =>
                  IO RpcConfig -> JobDesc a b -> Int -> Process ()
  slaveProcess configurator jobDesc@(MkJobDesc {..}) slaveIdx = do
    (MkRpcConfig config _ (MkSlaveConfig {..})) <- liftIO configurator
    liftIO $ putStrLn $ "starting slave process: " ++ show slaveIdx
    let slaveContext = makeSlaveContext config slaveIdx jobDesc
  
    let (server::ProcessDefinition (SlaveContext (LoggingT IO) a b)) = defaultProcess {
        apiHandlers = [handleCall (handle slaveLogger)]
      , exitHandlers = [handleExit exitHandler]
      , shutdownHandler = shutdownHandler
      , unhandledMessagePolicy = Drop
      }
    serve slaveContext (\ls -> return $ InitOk ls Infinity) server
    where
    handle logger ls req = do
      (resp, ls') <- liftIO $ logger $ runCommand req ls
      let resp' = force resp
      replyWith resp' (ProcessContinue ls')
    exitHandler _ _ () = do
      liftIO $ putStrLn "slave exit"
      return $ ProcessStop ExitShutdown
    shutdownHandler _ _ = do
      liftIO $ putStrLn "slave shutdown"
      return $ ()
  
  data RpcState a = MkRpcState {
    rpcSlaves :: M.Map Int SlaveInfo
    -- not sure we should store it there since it is in the job desc
    , rpcSeed :: Maybe a
    , statefullSlaveMode :: Bool
    }
  
  
  rpcCall :: forall a. (S.Serialize a ) => RpcState a -> Int -> SlaveRequest -> IO SlaveResponse
  rpcCall (MkRpcState {..}) slaveIdx request = do
      let (MkSlaveInfo {..}) = rpcSlaves M.! slaveIdx
      atomically $ writeTChan _slaveRequestChannel (Right request)
  
      respE <- atomically $ readTChan _slaveResponseChannel
      case respE of
        Right resp -> return resp
        Left (RpcResponseControlError err) -> error ("Error in CloudHaskell RPC: " ++ err)
        Left RpcResponseControlStopped -> error "should not reach"
  
  
  instance (S.Serialize a) => CommandClass RpcState a where
    isStatefullSlave (MkRpcState{ statefullSlaveMode = mode }) = mode
    getNbSlaves (MkRpcState {..}) = M.size rpcSlaves
  
    send rpc@(MkRpcState {..}) slaveId req = do
      r <- rpcCall rpc slaveId req
      return $ Right r
  
    stop (MkRpcState {..}) = do
      let slaveInfos = M.elems rpcSlaves
      mapM_ (\(MkSlaveInfo {..}) -> do
        atomically $ writeTChan _slaveRequestChannel (Left RpcRequestControlStop)
        _ <- atomically $ readTChan _slaveResponseChannel
        -- todo add error management
        return ()
        ) slaveInfos
  
    setSeed rpc@(MkRpcState {..}) a = do
      let rpc' = rpc {rpcSeed = Just a}
      resetAll rpc'
      return rpc'
      where
      resetAll aRpc = do
        let nbSlaves = getNbSlaves aRpc
        let slaveIds = [0 .. nbSlaves - 1]
        let req = resetCommand (S.encode a)
        _ <- mapConcurrently (\slaveId -> send aRpc slaveId req) slaveIds
        return ()
  
  
  startClientRpc :: forall a b. (S.Serialize a, S.Serialize b, CommandClass RpcState a) =>
    RpcConfig
    -> JobDesc a b
    -> (Int -> Closure (Process()))
    -> (a -> b -> IO ())
    -> Backend
    -> [NodeId]
    -> Process ()
  startClientRpc (MkRpcConfig config (MkMasterConfig logger) _) theJobDesc slaveClosure k backend _ = do
    loop 0 theJobDesc
    where
    mkSlaveInfo i nodeId = do
      requestChan <- newTChanIO
      responseChannel <- newTChanIO
      return $ (i, MkSlaveInfo i nodeId requestChan responseChannel)
    findSlaveNodes = do
      selfNode <- getSelfNode
      nodeIds <- liftIO $ findPeers backend 1000000
      return $ nodeIds L.\\ [selfNode]
    loop :: Int -> JobDesc a b -> Process ()
    loop n (jobDesc@MkJobDesc {..}) = do
      nodeIds <- findSlaveNodes
      case nodeIds of
        [] -> do
          liftIO $ putStrLn "No node found, retrying"
          liftIO $ threadDelay 5000000
          loop n jobDesc
        _ -> do
          liftIO $ putStrLn ("Nodes found: " ++ show nodeIds)
          slaveInfos <- liftIO $ mapM (\(i, nodeId) -> mkSlaveInfo i nodeId) $ L.zip [0..] nodeIds
          let slaveInfoMap = M.fromList slaveInfos
          -- create processes that handle RPC
          mapM_ (\slaveInfo -> spawnLocal (startOneClientRpc slaveInfo slaveClosure)) $ M.elems slaveInfoMap
          let rpcState = MkRpcState slaveInfoMap Nothing (statefullSlaves config)
          (a, b) <- liftIO $ do logger $ runComputation config rpcState jobDesc
          liftIO $ stop rpcState
          a' <- liftIO $ reportingAction a b
          case shouldStop seed a' b of
            True -> liftIO $ k a b
            False -> do let jobDesc' = jobDesc {seed = a'}
                        liftIO $ putStrLn "iteration finished"
                        loop (n+1) jobDesc'
  
  
  
  startOneClientRpc :: SlaveInfo -> (Int -> Closure (Process ())) -> Process ()
  startOneClientRpc (MkSlaveInfo {..}) slaveClosure  = do
    slavePid <- spawn _slaveNodeId (slaveClosure _slaveIndex)
    catchExit
      (localProcess 0 slavePid)
      (\slavePid' () -> do
        liftIO $ putStrLn ("stopping slave from handler:" ++ show slavePid)
        exit slavePid' ()
       -- shutdown slavePid
        )
    where
    localProcess :: Int -> ProcessId -> Process ()
    localProcess nbError slavePid = do
      requestE <- liftIO $ atomically $ readTChan _slaveRequestChannel
      case requestE of
        Right request -> do
          (respE::Either ExitReason SlaveResponse) <- safeCall slavePid request
          case respE of
            Right resp -> do
              liftIO $ atomically $ writeTChan _slaveResponseChannel $ Right resp
              localProcess 0 slavePid
            Left e | nbError < 10 -> do
              liftIO $ putStrLn ("error: "++show e)
              liftIO $ threadDelay 5000000
              -- todo deprecated, fix me
              liftIO $ atomically $ unGetTChan _slaveRequestChannel requestE
              newSlavePid <- spawn _slaveNodeId (slaveClosure _slaveIndex)
              localProcess (nbError+1) newSlavePid
            Left e  -> do
              liftIO $ atomically $ writeTChan _slaveResponseChannel $ Left $ RpcResponseControlError $ show e
              localProcess nbError slavePid
        Left RpcRequestControlStop -> do
                liftIO $ putStrLn ("stopping slave in a controlled way:" ++ show slavePid)
                exit slavePid ()
                liftIO $ atomically $ writeTChan _slaveResponseChannel $ Left $ RpcResponseControlStopped
                liftIO $ putStrLn "Terminating client process"
  
  
  -- | Run the computation on CloudHaskell
  runRec :: forall a b. (S.Serialize a, S.Serialize b) =>
    RemoteTable                         -- ^ CloudHaskell remote table.
    -> RpcConfig                        -- ^ Configuration.
    -> [String]                         -- ^ Command line arguments
    -> JobDesc a b                      -- ^ Job description.
    -> (Int -> Closure (Process()))     -- ^ Function that takes a slave index and returns the CloudHaskell closure to execute on that slave (see "slaveProcess").
    -> (a -> b -> IO ())                -- ^ A continuation that is called when the computation ends.
    -> IO ()
  runRec rtable rpcConfig args jobDesc slaveClosure k = do
    case args of
      ["master", host, port] -> do
        backend <- initializeBackend host port rtable
        startMaster backend (startClientRpc rpcConfig jobDesc slaveClosure k backend)
        putStrLn ("End")
      ["slave", host, port] -> do
        backend <- initializeBackend host port rtable
        startSlave backend
      _ -> putStrLn ("Bad args: " ++ show args)
  
  