1 files changed, 31 insertions, 16 deletions

diff --git a/Utility/LockPool/STM.hs b/Utility/LockPool/STM.hs
index 1dc30b09b..d1ee0dbaf 100644
--- a/Utility/LockPool/STM.hs
+++ b/Utility/LockPool/STM.hs
@@ -15,8 +15,12 @@ module Utility.LockPool.STM (
 	tryTakeLock,
 	getLockStatus,
 	releaseLock,
+	CloseLockFile,
+	registerCloseLockFile,
 ) where
 
+import Utility.Monad
+
 import System.IO.Unsafe (unsafePerformIO)
 import qualified Data.Map.Strict as M
 import Control.Concurrent.STM
@@ -36,7 +40,9 @@ type LockHandle = TMVar (LockPool, LockFile)
 
 type LockCount = Integer
 
-data LockStatus = LockStatus LockMode LockCount
+data LockStatus = LockStatus LockMode LockCount CloseLockFile
+
+type CloseLockFile = IO ()
 
 -- This TMVar is normally kept full.
 type LockPool = TMVar (M.Map LockFile LockStatus)
@@ -59,11 +65,11 @@ waitTakeLock :: LockPool -> LockFile -> LockMode -> STM LockHandle
 waitTakeLock pool file mode = do
 	m <- takeTMVar pool
 	v <- case M.lookup file m of
-		Just (LockStatus mode' n)
+		Just (LockStatus mode' n closelockfile)
 			| mode == LockShared && mode' == LockShared ->
-				return $ LockStatus mode (succ n)
+				return $ LockStatus mode (succ n) closelockfile
 			| n > 0 -> retry -- wait for lock
-		_ -> return $ LockStatus mode 1
+		_ -> return $ LockStatus mode 1 noop
 	putTMVar pool (M.insert file v m)
 	newTMVar (pool, file)
 
@@ -74,6 +80,16 @@ tryTakeLock pool file mode =
 		`orElse`
 	return Nothing
 
+-- Call after waitTakeLock or tryTakeLock, to register a CloseLockFile
+-- action to run when releasing the lock.
+registerCloseLockFile :: LockPool -> LockFile -> CloseLockFile -> STM ()
+registerCloseLockFile pool file closelockfile = do
+	m <- takeTMVar pool
+	putTMVar pool (M.update go file m)
+  where
+	go (LockStatus mode n closelockfile') = Just $
+		LockStatus mode n (closelockfile' >> closelockfile)
+
 -- Checks if a lock is being held. If it's held by the current process,
 -- runs the getdefault action; otherwise runs the checker action.
 --
@@ -87,7 +103,7 @@ getLockStatus pool file getdefault checker = do
 	v <- atomically $ do
 		m <- takeTMVar pool
 		let threadlocked = case M.lookup file m of
-			Just (LockStatus _ n) | n > 0 -> True
+			Just (LockStatus _ n _) | n > 0 -> True
 			_ -> False
 		if threadlocked
 			then do
@@ -99,25 +115,24 @@ getLockStatus pool file getdefault checker = do
 		Just restore -> bracket_ (return ()) restore checker
 
 -- Only runs action to close underlying lock file when this is the last
--- user of the lock, and when the handle has not already been closed.
+-- user of the lock, and when the lock has not already been closed.
 --
--- Note that the lock pool is left empty while the closelockfile action
+-- Note that the lock pool is left empty while the CloseLockFile action
 -- is run, to avoid race with another thread trying to open the same lock
 -- file.
-releaseLock :: LockHandle -> IO () -> IO ()
-releaseLock h closelockfile = go =<< atomically (tryTakeTMVar h)
+releaseLock :: LockHandle -> IO ()
+releaseLock h = go =<< atomically (tryTakeTMVar h)
   where
 	go (Just (pool, file)) = do
-		(m, unused) <- atomically $ do
+		(m, closelockfile) <- atomically $ do
 			m <- takeTMVar pool
 			return $ case M.lookup file m of
-				Just (LockStatus mode n)
-					| n == 1 -> (M.delete file m, True)
+				Just (LockStatus mode n closelockfile)
+					| n == 1 -> (M.delete file m, closelockfile)
 					| otherwise ->
-						(M.insert file (LockStatus mode (pred n)) m, False)
-				Nothing -> (m, True)
-		when unused
-			closelockfile
+						(M.insert file (LockStatus mode (pred n) closelockfile) m, noop)
+				Nothing -> (m, noop)
+		closelockfile
 		atomically $ putTMVar pool m
 	-- The LockHandle was already closed.
 	go Nothing = return ()