Staged builtin rules

shake.cabal

@@ -101,6 +101,7 @@ library
         primitive,
         process >= 1.1,
         random,
+        stm,
         time,
         transformers >= 0.2,
         unordered-containers >= 0.2.7,
@@ -197,6 +198,7 @@ library
         General.Makefile
         General.Pool
         General.Process
+        General.RLock
         General.Template
         General.Thread
         General.Timing
@@ -227,6 +229,7 @@ executable shake
         primitive,
         process >= 1.1,
         random,
+        stm,
         time,
         transformers >= 0.2,
         unordered-containers >= 0.2.7,
@@ -318,6 +321,7 @@ executable shake
         General.Makefile
         General.Pool
         General.Process
+        General.RLock
         General.Template
         General.Thread
         General.Timing
@@ -351,6 +355,7 @@ test-suite shake-test
         process >= 1.1,
         QuickCheck >= 2.0,
         random,
+        stm,
         time,
         transformers >= 0.2,
         unordered-containers >= 0.2.7,
@@ -446,6 +451,7 @@ test-suite shake-test
         General.Makefile
         General.Pool
         General.Process
+        General.RLock
         General.Template
         General.Thread
         General.Timing

src/Development/Shake/Internal/Core/Build.hs

@@ -102,8 +102,8 @@ buildOne global@Global{..} stack database i k r = case addStack i k stack of
     Right stack -> Later $ \continue -> do
         setIdKeyStatus global database i k (Running (NoShow continue) r)
         let go = buildRunMode global stack database r
-        fromLater go $ \mode -> liftIO $ addPool PoolStart globalPool $
-            runKey global stack k r mode $ \res -> do
+        fromLater go $ \mode -> liftIO $
+            fromLater(runKey global stack k r mode) $ \res -> mask_ $ do
                 runLocked database $ do
                     let val = fmap runValue res
                     res <- liftIO $ getKeyValueFromId database i
@@ -184,45 +184,49 @@ runKey
     -> Key -- The key to build
     -> Maybe (Result BS.ByteString) -- A previous result, or Nothing if never been built before
     -> RunMode -- True if any of the children were dirty
-    -> Capture (Either SomeException (RunResult (Result (Value, BS_Store))))
+    -> Wait IO (Either SomeException (RunResult (Result (Value, BS_Store))))
         -- Either an error, or a (the produced files, the result).
-runKey global@Global{globalOptions=ShakeOptions{..},..} stack k r mode continue = do
+runKey global@Global{globalOptions=ShakeOptions{..},..} stack k r mode = do
     let tk = typeKey k
     BuiltinRule{..} <- case Map.lookup tk globalRules of
         Nothing -> throwM $ errorNoRuleToBuildType tk (Just $ show k) Nothing
         Just r -> pure r
 
     let s = (newLocal stack shakeVerbosity){localBuiltinVersion = builtinVersion}
-    time <- offsetTime
-    runAction global s (do
-        res <- builtinRun k (fmap result r) mode
-        liftIO $ evaluate $ rnf res
-
-        -- completed, now track anything required afterwards
-        when (runChanged res `elem` [ChangedRecomputeSame,ChangedRecomputeDiff]) $ do
-            -- if the users code didn't run you don't have to check anything (we assume builtin rules are correct)
-            globalRuleFinished k
-            producesCheck
-
-        Action $ fmap (res,) getRW) $ \case
-            Left e ->
-                continue . Left . toException =<< shakeException global stack e
-            Right (RunResult{..}, Local{..})
-                | runChanged == ChangedNothing || runChanged == ChangedStore, Just r <- r ->
-                    continue $ Right $ RunResult runChanged runStore (r{result = mkResult runValue runStore})
-                | otherwise -> do
-                    dur <- time
-                    let (cr, c) | Just r <- r, runChanged == ChangedRecomputeSame = (ChangedRecomputeSame, changed r)
-                                | otherwise = (ChangedRecomputeDiff, globalStep)
-                    continue $ Right $ RunResult cr runStore Result
-                        {result = mkResult runValue runStore
-                        ,changed = c
-                        ,built = globalStep
-                        ,depends = flattenDepends localDepends
-                        ,execution = doubleToFloat $ dur - localDiscount
-                        ,traces = flattenTraces localTraces}
-            where
-                mkResult value store = (value, if globalOneShot then BS.empty else store)
+    time <- liftIO offsetTime
+    let followUp = \case
+                Left e ->
+                    Left . toException <$> shakeException global stack e
+                Right (RunResult{..}, Local{..})
+                    | runChanged == ChangedNothing || runChanged == ChangedStore, Just r <- r ->
+                        pure $ Right $ RunResult runChanged runStore (r{result = mkResult runValue runStore})
+                    | otherwise -> do
+                        dur <- liftIO time
+                        let (cr, c) | Just r <- r, runChanged == ChangedRecomputeSame = (ChangedRecomputeSame, changed r)
+                                    | otherwise = (ChangedRecomputeDiff, globalStep)
+                        pure $ Right $ RunResult cr runStore Result
+                            {result = mkResult runValue runStore
+                            ,changed = c
+                            ,built = globalStep
+                            ,depends = flattenDepends localDepends
+                            ,execution = doubleToFloat $ dur - localDiscount
+                            ,traces = flattenTraces localTraces}
+                where
+                    mkResult value store = (value, if globalOneShot then BS.empty else store)
+    stage1 <- liftIO $ try $ builtinRun k (fmap result r) mode
+    case stage1 of
+        Left e -> Now . Left . toException =<< liftIO (shakeException global stack e)
+        Right (BuiltinRunChangedNothing done) ->
+            liftIO $ followUp (Right (RunResult ChangedNothing (result $ fromJust r) done, s))
+        Right (BuiltinRunMore more) -> Later $ \continue -> liftIO $ addPool PoolStart globalPool $ runAction global s (do
+            res <- more
+            liftIO $ evaluate $ rnf res
+            -- completed, now track anything required afterwards
+            when (runChanged res `elem` [ChangedRecomputeSame,ChangedRecomputeDiff]) $ do
+                -- if the users code didn't run you don't have to check anything (we assume builtin rules are correct)
+                globalRuleFinished k
+                producesCheck
+            Action $ fmap (res,) getRW) (followUp >=> continue)
 
 ---------------------------------------------------------------------
 -- USER key/value WRAPPERS

src/Development/Shake/Internal/Core/Database.hs

@@ -14,7 +14,7 @@ import General.Intern(Id, Intern)
 import Development.Shake.Classes
 import qualified Data.HashMap.Strict as Map
 import qualified General.Intern as Intern
-import Control.Concurrent.Extra
+import General.RLock as RLock
 import Control.Monad.IO.Class
 import qualified General.Ids as Ids
 import Control.Monad.Fail
@@ -25,15 +25,15 @@ newtype Locked a = Locked (IO a)
     deriving (Functor, Applicative, Monad, MonadIO, MonadFail)
 
 runLocked :: DatabasePoly k v -> Locked b -> IO b
-runLocked db (Locked act) = withLock (lock db) act
+runLocked db (Locked act) = RLock.with (lock db) act
 
 
 -- | Invariant: The database does not have any cycles where a Key depends on itself.
 --   Everything is mutable. intern and status must form a bijecttion.
 --   There may be dangling Id's as a result of version changes.
 --   Lock is used to prevent any torn updates
 data DatabasePoly k v = Database
-    {lock :: Lock
+    {lock :: RLock
     ,intern :: IORef (Intern k) -- ^ Key |-> Id mapping
     ,status :: Ids.Ids (k, v) -- ^ Id |-> (Key, Status) mapping
     ,journal :: Id -> k -> v -> IO () -- ^ Record all changes to status
@@ -50,7 +50,7 @@ createDatabase
 createDatabase status journal vDefault = do
     xs <- Ids.toList status
     intern <- newIORef $ Intern.fromList [(k, i) | (i, (k,_)) <- xs]
-    lock <- newLock
+    lock <- RLock.new
     pure Database{..}
 
 

src/Development/Shake/Internal/Core/Rules.hs

@@ -6,7 +6,7 @@
 
 module Development.Shake.Internal.Core.Rules(
     Rules, SRules(..), runRules,
-    RuleResult, addBuiltinRule, addBuiltinRuleEx,
+    RuleResult, addBuiltinRule, addBuiltinRuleStaged, addBuiltinRuleEx,
     noLint, noIdentity,
     getShakeOptionsRules,
     getUserRuleInternal, getUserRuleOne, getUserRuleList, getUserRuleMaybe,
@@ -240,21 +240,31 @@ type family RuleResult key -- = value
 addBuiltinRule
     :: (RuleResult key ~ value, ShakeValue key, Typeable value, NFData value, Show value, Partial)
     => BuiltinLint key value -> BuiltinIdentity key value -> BuiltinRun key value -> Rules ()
-addBuiltinRule = withFrozenCallStack $ addBuiltinRuleInternal $ BinaryOp
+addBuiltinRule lint check run = addBuiltinRuleStaged lint check (builtinRun' run)
+
+addBuiltinRuleStaged
+    :: (RuleResult key ~ value, ShakeValue key, Typeable value, NFData value, Show value, Partial)
+    => BuiltinLint key value -> BuiltinIdentity key value -> BuiltinRun' key value -> Rules ()
+addBuiltinRuleStaged = withFrozenCallStack $ addBuiltinRuleInternal $ BinaryOp
     (putEx . Bin.toLazyByteString . execPut . put)
     (runGet get . LBS.fromChunks . pure)
 
 addBuiltinRuleEx
     :: (RuleResult key ~ value, ShakeValue key, BinaryEx key, Typeable value, NFData value, Show value, Partial)
     => BuiltinLint key value -> BuiltinIdentity key value -> BuiltinRun key value -> Rules ()
-addBuiltinRuleEx = addBuiltinRuleInternal $ BinaryOp putEx getEx
-
+addBuiltinRuleEx = addBuiltinRuleInternal' $ BinaryOp putEx getEx
 
 -- | Unexpected version of 'addBuiltinRule', which also lets me set the 'BinaryOp'.
-addBuiltinRuleInternal
+addBuiltinRuleInternal'
     :: (RuleResult key ~ value, ShakeValue key, Typeable value, NFData value, Show value, Partial)
     => BinaryOp key -> BuiltinLint key value -> BuiltinIdentity key value -> BuiltinRun key value -> Rules ()
-addBuiltinRuleInternal binary lint check (run :: BuiltinRun key value) = do
+addBuiltinRuleInternal' binary lint check run =
+    addBuiltinRuleInternal binary lint check (builtinRun' run)
+
+addBuiltinRuleInternal
+    :: (RuleResult key ~ value, ShakeValue key, Typeable value, NFData value, Show value, Partial)
+    => BinaryOp key -> BuiltinLint key value -> BuiltinIdentity key value -> BuiltinRun' key value -> Rules ()
+addBuiltinRuleInternal binary lint check (run :: BuiltinRun' key value) = do
     let k = Proxy :: Proxy key
     let lint_ k v = lint (fromKey k) (fromValue v)
     let check_ k v = check (fromKey k) (fromValue v)

src/Development/Shake/Internal/Core/Types.hs

@@ -1,8 +1,10 @@
 {-# LANGUAGE GeneralizedNewtypeDeriving, ScopedTypeVariables, DeriveDataTypeable, ViewPatterns #-}
 {-# LANGUAGE ExistentialQuantification, DeriveFunctor, RecordWildCards, FlexibleInstances #-}
 
+{-# LANGUAGE RankNTypes #-}
 module Development.Shake.Internal.Core.Types(
-    BuiltinRun, BuiltinLint, BuiltinIdentity,
+    BuiltinRun, BuiltinRun', BuiltinLint, BuiltinIdentity,
+    BuiltinRunResult(..), builtinRun',
     RunMode(..), RunResult(..), RunChanged(..),
     UserRule(..), UserRuleVersioned(..), userRuleSize,
     BuiltinRule(..), Global(..), Local(..), Action(..), runAction, addDiscount,
@@ -350,12 +352,26 @@ enumerateDepends d = f d []
 -- * @newStore@, the new value to store in the database, which will be passed in next time as @oldStore@.
 --
 -- * @newValue@, the result that 'Development.Shake.Rule.apply' will return when asked for the given @key@.
+type BuiltinRun' key value
+    = key
+    -> Maybe BS.ByteString
+    -> RunMode
+    -> IO (BuiltinRunResult value)
+
+data BuiltinRunResult value
+    = BuiltinRunChangedNothing !value
+    | BuiltinRunMore !(Action (RunResult value))
+  deriving Functor
+
 type BuiltinRun key value
     = key
     -> Maybe BS.ByteString
     -> RunMode
     -> Action (RunResult value)
 
+builtinRun' :: BuiltinRun k v -> BuiltinRun' k v
+builtinRun' run k bs m = pure $ BuiltinRunMore $ run k bs m
+
 -- | The action performed by @--lint@ for a given @key@/@value@ pair.
 --   At the end of the build the lint action will be called for each @key@ that was built this run,
 --   passing the @value@ it produced. Return 'Nothing' to indicate the value has not changed and
@@ -378,7 +394,7 @@ type BuiltinIdentity key value = key -> value -> Maybe BS.ByteString
 data BuiltinRule = BuiltinRule
     {builtinLint :: BuiltinLint Key Value
     ,builtinIdentity :: BuiltinIdentity Key Value
-    ,builtinRun :: BuiltinRun Key Value
+    ,builtinRun :: BuiltinRun' Key Value
     ,builtinKey :: BinaryOp Key
     ,builtinVersion :: Ver
     ,builtinLocation :: String

src/Development/Shake/Rule.hs

@@ -16,8 +16,9 @@ module Development.Shake.Rule(
 
     -- * Defining builtin rules
     -- | Functions and types for defining new types of Shake rules.
-    addBuiltinRule,
-    BuiltinLint, noLint, BuiltinIdentity, noIdentity, BuiltinRun, RunMode(..), RunChanged(..), RunResult(..),
+    addBuiltinRule, addBuiltinRuleStaged,
+    BuiltinLint, noLint, BuiltinIdentity, noIdentity,
+    BuiltinRun, BuiltinRunResult(..), RunMode(..), RunChanged(..), RunResult(..),
     -- * Calling builtin rules
     -- | Wrappers around calling Shake rules. In general these should be specialised to a builtin rule.
     apply, apply1,

src/General/RLock.hs

@@ -0,0 +1,39 @@
+{-# LANGUAGE LambdaCase #-}
+module General.RLock (RLock, new, acquire, release, with) where
+
+import Control.Applicative
+import Control.Concurrent
+import Control.Concurrent.STM
+import Control.Exception
+
+-- | A reentrant lock inspired by the one in the concurrent-extra package, to
+--   work around https://github.com/basvandijk/concurrent-extra/issues/20
+--   This implementation uses a single 'TVar' and therefore it is not *fair*
+newtype RLock = RLock {_rlock :: TVar State}
+
+data State
+  = Locked !ThreadId
+  | Unlocked
+
+new :: IO RLock
+new = RLock <$> newTVarIO Unlocked
+
+acquire :: RLock -> IO Bool
+acquire (RLock tv) = do
+  tid <- myThreadId
+  atomically $ do
+    readTVar tv >>= \case
+      Locked tid'
+        | tid == tid' ->
+            return False
+        | otherwise -> retry
+      Unlocked -> do
+          writeTVar tv $! Locked tid
+          return True
+
+release :: RLock -> Bool -> IO ()
+release (RLock tv) True  = atomically $ writeTVar tv Unlocked
+release _          False = return ()
+
+with :: RLock -> IO a -> IO a
+with rl act = bracket (acquire rl) (release rl) (const act)