1 files changed, 135 insertions, 0 deletions
diff --git a/src/Propellor/Property/Concurrent.hs b/src/Propellor/Property/Concurrent.hs
new file mode 100644
index 00000000..e69dc17d
--- /dev/null
+++ b/src/Propellor/Property/Concurrent.hs
@@ -0,0 +1,135 @@
+{-# 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 a1 a2 = do
+		n <- liftIO getNumProcessors
+		withCapabilities n $
+			concurrentSatisfy a1 a2
+
+-- | 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' <- actionMessageOn hn
+					(getDesc p)
+					(getSatisfy p)
+				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 ()