-- | Allows restarting the currently running executable.
--
-- Works with binaries and `runhaskell`/`runghc`. In `ghci` all functions are no-ops.
module System.SelfRestart
  ( -- $usewithcare
    selfRestart
  , forkSelfRestartExePoll
  , forkSelfRestartExePollWithAction
  ) where

import Control.Applicative
import Control.Concurrent
import Control.Monad
import System.Posix.Process (executeFile)
import System.Directory (getModificationTime)
import System.Environment.Executable (getExecutablePath, getScriptPath, ScriptPath (..))
import GHC.Environment (getFullArgs)

-- $usewithcare
-- USE WITH CARE. This module uses `getScriptPath` internally, which is /experimental/.
-- Please review it.

-- | Restart the currently running executable.
selfRestart :: IO ()
selfRestart :: IO ()
selfRestart = do
  FilePath
realExePath <- IO FilePath
getExecutablePath
  ScriptPath
scriptPath  <- IO ScriptPath
getScriptPath
  [FilePath]
fullArgs    <- IO [FilePath]
getFullArgs

  case ScriptPath
scriptPath of
    Executable FilePath
p -> FilePath -> [FilePath] -> IO ()
forall {a}. FilePath -> [FilePath] -> IO a
restart FilePath
p           [FilePath]
fullArgs
    RunGHC FilePath
_     -> FilePath -> [FilePath] -> IO ()
forall {a}. FilePath -> [FilePath] -> IO a
restart FilePath
realExePath [FilePath]
fullArgs
    ScriptPath
Interactive  -> () -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()

  where
    restart :: FilePath -> [FilePath] -> IO a
restart FilePath
exePath [FilePath]
args = FilePath
-> Bool -> [FilePath] -> Maybe [(FilePath, FilePath)] -> IO a
forall a.
FilePath
-> Bool -> [FilePath] -> Maybe [(FilePath, FilePath)] -> IO a
executeFile FilePath
exePath Bool
False [FilePath]
args Maybe [(FilePath, FilePath)]
forall a. Maybe a
Nothing


-- | Forks a thread to check poll the own binary's modification time every,
-- and restarts, and restarts it when it gets changed.
--
-- Waits the given number of seconds between checks.
forkSelfRestartExePoll :: Double -> IO (Maybe ThreadId)
forkSelfRestartExePoll :: Double -> IO (Maybe ThreadId)
forkSelfRestartExePoll Double
pollSeconds = Double -> IO () -> IO (Maybe ThreadId)
forkSelfRestartExePollWithAction Double
pollSeconds (() -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ())


-- | Like `forkSelfRestartExePoll`, but you can run an action just before the restart
-- happens.
forkSelfRestartExePollWithAction :: Double -> IO () -> IO (Maybe ThreadId)
forkSelfRestartExePollWithAction :: Double -> IO () -> IO (Maybe ThreadId)
forkSelfRestartExePollWithAction Double
pollSeconds IO ()
shutdownAction = do
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
pollTimeoutUs Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath -> IO ()
forall a. HasCallStack => FilePath -> a
error (FilePath -> IO ()) -> FilePath -> IO ()
forall a b. (a -> b) -> a -> b
$ FilePath
"forkSelfRestartExePollWithAction: Bad poll timeout: " FilePath -> FilePath -> FilePath
forall a. [a] -> [a] -> [a]
++ Double -> FilePath
forall a. Show a => a -> FilePath
show Double
pollSeconds

  FilePath
realExePath <- IO FilePath
getExecutablePath
  ScriptPath
scriptPath  <- IO ScriptPath
getScriptPath
  [FilePath]
fullArgs    <- IO [FilePath]
getFullArgs

  case ScriptPath
scriptPath of
    Executable FilePath
p -> ThreadId -> Maybe ThreadId
forall a. a -> Maybe a
Just (ThreadId -> Maybe ThreadId) -> IO ThreadId -> IO (Maybe ThreadId)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FilePath -> FilePath -> [FilePath] -> IO ThreadId
spawn FilePath
p FilePath
p           [FilePath]
fullArgs
    RunGHC FilePath
p     -> ThreadId -> Maybe ThreadId
forall a. a -> Maybe a
Just (ThreadId -> Maybe ThreadId) -> IO ThreadId -> IO (Maybe ThreadId)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FilePath -> FilePath -> [FilePath] -> IO ThreadId
spawn FilePath
p FilePath
realExePath [FilePath]
fullArgs
    ScriptPath
Interactive  -> Maybe ThreadId -> IO (Maybe ThreadId)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe ThreadId
forall a. Maybe a
Nothing

  where
    pollTimeoutUs :: Int
pollTimeoutUs = Double -> Int
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor (Double -> Int) -> Double -> Int
forall a b. (a -> b) -> a -> b
$ Double
pollSeconds Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
1000000

    spawn :: FilePath -> FilePath -> [FilePath] -> IO ThreadId
spawn FilePath
watchPath FilePath
exePath [FilePath]
args = do
      UTCTime
mtimeStarted <- FilePath -> IO UTCTime
getModificationTime FilePath
watchPath
      IO () -> IO ThreadId
forkIO (IO () -> IO ThreadId) -> IO () -> IO ThreadId
forall a b. (a -> b) -> a -> b
$
        IO () -> IO ()
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
          Int -> IO ()
threadDelay Int
pollTimeoutUs
          UTCTime
mtime <- FilePath -> IO UTCTime
getModificationTime FilePath
watchPath
          Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (UTCTime
mtime UTCTime -> UTCTime -> Bool
forall a. Ord a => a -> a -> Bool
> UTCTime
mtimeStarted) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
            IO ()
shutdownAction
            FilePath
-> Bool -> [FilePath] -> Maybe [(FilePath, FilePath)] -> IO ()
forall a.
FilePath
-> Bool -> [FilePath] -> Maybe [(FilePath, FilePath)] -> IO a
executeFile FilePath
exePath Bool
False [FilePath]
args Maybe [(FilePath, FilePath)]
forall a. Maybe a
Nothing