    -- | Code for ensuring cleanup actions are run.
    module General.Cleanup(
        Cleanup, withCleanup, addCleanup
        ) where

    import Control.Exception as E
    import Control.Monad
    import qualified Data.HashMap.Strict as Map
    import Data.Function
    import Data.IORef
    import Data.List

    data S = S {unique :: !Int, items :: Map.HashMap Int (IO ())}

    newtype Cleanup = Cleanup (IORef S)

    -- | Run with some cleanup scope. Regardless of exceptions/threads, all 'addCleanup' actions
    --   will be run by the time it exits. The 'addCleanup' actions will be run in reverse order.
    withCleanup :: (Cleanup -> IO a) -> IO a
    withCleanup act = do
        ref <- newIORef $ S 0 Map.empty
        act (Cleanup ref) `finally` do
            items <- atomicModifyIORef ref $ \s -> (s{items=Map.empty}, items s)
            mapM_ snd $ sortBy (compare `on` negate . fst) $ Map.toList items

    -- | Add a cleanup action to a 'Cleanup' scope. If the return action is not run by the time
    --   'withCleanup' terminates then it will be run then. The argument 'Bool' is 'True' to say
    --   run the action, 'False' to say ignore the action (and never run it).
    addCleanup :: Cleanup -> IO () -> IO (Bool -> IO ())
    addCleanup (Cleanup ref) act = atomicModifyIORef ref $ \s -> let i = unique s in
        (,) (S (unique s + 1) (Map.insert i act $ items s)) $ \b ->
            join $ atomicModifyIORef ref $ \s -> case Map.lookup i $ items s of
                Nothing -> (s, return ())
                Just act -> (s{items = Map.delete i $ items s}, when b act)