{-# LANGUAGE FlexibleContexts #-}

-- | Propellor properties can be made to run concurrently, using this
-- module. This can speed up propellor, at the expense of using more CPUs
-- and other resources.
--
-- It's up to you to make sure that properties that you make run concurrently
-- don't implicitly depend on one-another. The worst that can happen
-- though, is that propellor fails to ensure some of the properties,
-- and tells you what went wrong.
--
-- Another potential problem is that output of concurrent properties could
-- interleave into a scrambled mess. This is mostly prevented; all messages
-- output by propellor are concurrency safe, including `errorMessage`,
-- `infoMessage`, etc. However, if you write a property that directly
-- uses `print` or `putStrLn`, you can still experience this problem.
--
-- Similarly, when properties run external commands, the command's output
-- can be a problem for concurrency. No need to worry;
-- `Propellor.Property.Cmd.createProcess` is concurrent output safe
-- (it actually uses `Propellor.Message.createProcessConcurrent`), and
-- everything else in propellor that runs external commands is built on top
-- of that. Of course, if you import System.Process and use it in a
-- property, you can bypass that and shoot yourself in the foot.
--
-- Finally, anything that directly accesses the tty can bypass
-- these protections. That's sometimes done for eg, password prompts.
-- A well-written property should avoid running interactive commands
-- anyway.

module Propellor.Property.Concurrent (
	concurrently,
	concurrentList,
	props,
	getNumProcessors,
	concurrentSatisfy,
) where

import Propellor.Base
import Propellor.Types.Core
import Propellor.Types.MetaTypes

import Control.Concurrent
import qualified Control.Concurrent.Async as A
import GHC.Conc (getNumProcessors)
import Control.Monad.RWS.Strict

-- | Ensures two properties concurrently.
--
-- >	& foo `concurrently` bar
--
-- To ensure three properties concurrently, just use this combinator twice:
--
-- >	& foo `concurrently` bar `concurrently` baz
concurrently
	:: (IsProp p1, IsProp p2, Combines p1 p2, IsProp (CombinedType p1 p2))
	=> p1
	-> p2
	-> CombinedType p1 p2
concurrently p1 p2 = (combineWith go go p1 p2)
	`describe` d
  where
	d = getDesc p1 ++ " `concurrently` " ++ getDesc p2
	-- Increase the number of capabilities right up to the number of
	-- processors, so that A `concurrently` B `concurrently` C
	-- runs all 3 properties on different processors when possible.
	go (Just a1) (Just a2) = Just $ do
		n <- liftIO getNumProcessors
		withCapabilities n $
			concurrentSatisfy a1 a2
	go (Just a1) Nothing = Just a1
	go Nothing (Just a2) = Just a2
	go Nothing Nothing = Nothing

-- | Ensures all the properties in the list, with a specified amount of
-- concurrency.
-- 
-- > concurrentList (pure 2) "demo" $ props
-- >	& foo
-- >	& bar
-- >	& baz
--
-- The above example will run foo and bar concurrently, and once either of
-- those 2 properties finishes, will start running baz.
concurrentList :: SingI metatypes => IO Int -> Desc -> Props (MetaTypes metatypes) -> Property (MetaTypes metatypes)
concurrentList getn d (Props ps) = property d go `addChildren` ps
  where
	go = do
		n <- liftIO getn
		withCapabilities n $
			startworkers n =<< liftIO (newMVar ps)
	startworkers n q
		| n < 1 = return NoChange
		| n == 1 = worker q NoChange
		| otherwise = 
			worker q NoChange
				`concurrentSatisfy`
			startworkers (n-1) q
	worker q r = do
		v <- liftIO $ modifyMVar q $ \v -> case v of
			[] -> return ([], Nothing)
			(p:rest) -> return (rest, Just p)
		case v of
			Nothing -> return r
			Just p -> do
				hn <- asks hostName
				r' <- case getSatisfy p of
					Nothing -> return NoChange
					Just a -> actionMessageOn hn (getDesc p) a
				worker q (r <> r')

-- | Run an action with the number of capabiities increased as necessary to
-- allow running on the specified number of cores.
--
-- Never increases the number of capabilities higher than the actual number
-- of processors.
withCapabilities :: Int -> Propellor a -> Propellor a
withCapabilities n a = bracket setup cleanup (const a)
  where
	setup = do
		np <- liftIO getNumProcessors
		let n' = min n np
		c <- liftIO getNumCapabilities
		when (n' > c) $ 
			liftIO $ setNumCapabilities n'
		return c
	cleanup = liftIO . setNumCapabilities

-- | Running Propellor actions concurrently.
concurrentSatisfy :: Propellor Result -> Propellor Result -> Propellor Result
concurrentSatisfy a1 a2 = do
	h <- ask
	((r1, w1), (r2, w2)) <- liftIO $
		runp a1 h `A.concurrently` runp a2 h
	tell (w1 <> w2)
	return (r1 <> r2)
  where
	runp a h = evalRWST (runWithHost (catchPropellor a)) h ()