envrace.hs

{-# LANGUAGE BlockArguments #-}

module Main where

import System.Environment
import Control.Monad
import Control.Concurrent
import Control.Exception.Base

{- Inspiration:

   withProgArgv :: [String] -> IO a -> IO a
   withProgArgv new_args act = do
     pName <- System.Environment.getProgName
     existing_args <- System.Environment.getArgs
     bracket_ (setProgArgv new_args)
              (setProgArgv (pName:existing_args))
              act
-}

-- Strictly control the environment
-- Test normal bracketed environment setting in a single thread
-- Test environment setting in background threads
-- Test environment setting between interleaved threads

main :: IO ()
main = do
  let w = withEnv3
  printEnv "m1"
  w [("a","1")] do
    printEnv "m2"
  printEnv "m3"
  fork do
    printEnv "m4"
    w [("b","2")] do
      printEnv "m5"
    printEnv "m6"
  printEnv "m7"
  race w

blocker :: IO (MVar ())
blocker = newEmptyMVar

unblock :: MVar () -> IO ()
unblock m = putMVar m ()

{- *Main> race withEnv2
  t2a > []
  t1 > [("ra","1")]
  t2b > [("ra","1")]
-}

race :: ([(String, String)] -> IO () -> IO ()) -> IO ()
race w = do
  t1End                   <- blocker
  t2End                   <- blocker
  afterT1writeBeforeExits <- blocker
  letT1Exit               <- blocker
  letT1Start              <- blocker

  _t1 <- forkIO do
    w [("ra","1")] do
      takeMVar letT1Start
      printEnv "t1"
      unblock afterT1writeBeforeExits
      takeMVar letT1Exit
    unblock t1End

  _t2 <- forkIO do
    printEnv "t2a"
    unblock letT1Start
    takeMVar afterT1writeBeforeExits
    printEnv "t2b"
    unblock letT1Exit
    unblock t2End

  takeMVar t1End
  takeMVar t2End

printEnv :: String -> IO ()
printEnv s = do
  e <- filter relevant <$> getEnvironment
  putStr (s ++ " > ")
  print e

relevant :: (String, String) -> Bool
relevant (s, _) = s `elem` (words "a b c ra rb rc")

fork :: IO a -> IO a
fork e = do
  m <- newEmptyMVar
  _ <- forkIO do
    r <- e
    putMVar m r
  takeMVar m

injectEnvs :: [(String, String)] -> IO ()
injectEnvs = mapM_ (uncurry setEnv)

setEnvs :: [(String, String)] -> IO ()
setEnvs e = do
  e' <- getEnvironment
  mapM_ (uncurry setEnv) e
  mapM_ (foo (map fst e)) (map fst e')

  where foo a b = unless (b `elem` a) (unsetEnv b)

{- *Main> main
  []
  [("a","1")]
  [("a","1")]
  [("a","1")]
  [("b","2")]
  [("b","2")]
  [("b","2")]
-}

withEnv :: [(String, String)] -> IO a -> IO a
withEnv e m = do
  setEnvs e
  m

{- *Main> main
  []
  [("a","1")]
  []
  []
  [("b","2")]
  []
  []
-}

withEnv2 :: [(String, String)] -> IO a -> IO a
withEnv2 e m =
  bracket
    getEnvironment
    setEnvs
    (\_ -> setEnvs e >> m)

{- *Main> main
  []
  [("a","1")]
  []
  []
  [("b","2")]
  []
  []
-}

withEnv3 :: [(String, String)] -> IO a -> IO a
withEnv3 e m =
  bracket
    getEnvironment
    setEnvs
    (\_ -> injectEnvs e >> m)