Fix shared lock file FD leak.

This fixes behavior in this situation:

	l1 <- lockShared Nothing "lck"
	l2 <- lockShared Nothing "lck"
	dropLock l1
	dropLock l2

Before, the lock was dropped upon the second dropLock call, but the fd
remained open, and would never be closed while the program was running.

Fixed by a rather round-about method, but it should work well enough.

It would have been simpler to open open the shared lock once, and not open
it again in the second call to lockShared. But, that's difficult to do
atomically.

This also affects Windows and PID locks, not just posix locks.

In the case of pid locks, multiple calls to waitLock within the same
process are allowed because the side lock is locked using a posix lock,
and so multiple exclusive locks can be taken in the same process. So,
this change fixes a similar problem with pid locks.

	l1 <- waitLock (Seconds 1) "lck"
	l2 <- waitLock (Seconds 1) "lck"
	dropLock l1
	dropLock l2

Here the l2 side lock fd remained open but not locked,
although the pid lock file was removed. After this change, the second
dropLock will close both fds to the side lock, and delete the pidlock.

5 files changed, 84 insertions, 59 deletions

diff --git a/Utility/LockPool/LockHandle.hs b/Utility/LockPool/LockHandle.hs
index 68c979b5d..5697b89cb 100644
--- a/Utility/LockPool/LockHandle.hs
+++ b/Utility/LockPool/LockHandle.hs
@@ -7,7 +7,16 @@
 
 {-# LANGUAGE CPP #-}
 
-module Utility.LockPool.LockHandle where
+module Utility.LockPool.LockHandle (
+	LockHandle,
+	FileLockOps(..),
+	dropLock,
+#ifndef mingw32_HOST_OS
+	checkSaneLock,
+#endif
+	makeLockHandle,
+	tryMakeLockHandle,
+) where
 
 import qualified Utility.LockPool.STM as P
 #ifndef mingw32_HOST_OS
@@ -17,10 +26,7 @@ import Utility.LockPool.STM (LockFile)
 import Control.Concurrent.STM
 import Control.Exception
 
-data LockHandle = LockHandle
-	{ poolHandle :: P.LockHandle
-	, fileLockOps :: FileLockOps
-	}
+data LockHandle = LockHandle P.LockHandle FileLockOps
 
 data FileLockOps = FileLockOps
 	{ fDropLock :: IO ()
@@ -30,7 +36,7 @@ data FileLockOps = FileLockOps
 	}
 
 dropLock :: LockHandle -> IO ()
-dropLock h = P.releaseLock (poolHandle h) (fDropLock (fileLockOps h))
+dropLock (LockHandle ph _) = P.releaseLock ph
 
 #ifndef mingw32_HOST_OS
 checkSaneLock :: LockFile -> LockHandle -> IO Bool
@@ -40,26 +46,30 @@ checkSaneLock lockfile (LockHandle _ flo) = fCheckSaneLock flo lockfile
 -- Take a lock, by first updating the lock pool, and then taking the file
 -- lock. If taking the file lock fails for any reason, take care to
 -- release the lock in the lock pool.
-makeLockHandle :: STM P.LockHandle -> IO FileLockOps -> IO LockHandle
-makeLockHandle pa fa = bracketOnError setup cleanup go
+makeLockHandle :: P.LockPool -> LockFile -> (P.LockPool -> LockFile -> STM P.LockHandle) -> (LockFile -> IO FileLockOps) -> IO LockHandle
+makeLockHandle pool file pa fa = bracketOnError setup cleanup go
   where
-	setup = atomically pa
-	cleanup ph = P.releaseLock ph (return ())
-	go ph = do
-		fo <- fa
-		return $ LockHandle ph fo
+	setup = atomically (pa pool file)
+	cleanup ph = P.releaseLock ph
+	go ph = mkLockHandle pool file ph =<< fa file
 
-tryMakeLockHandle :: STM (Maybe P.LockHandle) -> IO (Maybe FileLockOps) -> IO (Maybe LockHandle)
-tryMakeLockHandle pa fa = bracketOnError setup cleanup go
+tryMakeLockHandle :: P.LockPool -> LockFile -> (P.LockPool -> LockFile -> STM (Maybe P.LockHandle)) -> (LockFile -> IO (Maybe FileLockOps)) -> IO (Maybe LockHandle)
+tryMakeLockHandle pool file pa fa = bracketOnError setup cleanup go
   where
-	setup = atomically pa
+	setup = atomically (pa pool file)
 	cleanup Nothing = return ()
-	cleanup (Just ph) = P.releaseLock ph (return ())
+	cleanup (Just ph) = P.releaseLock ph
 	go Nothing = return Nothing
 	go (Just ph) = do
-		mfo <- fa
+		mfo <- fa file
 		case mfo of
 			Nothing -> do
 				cleanup (Just ph)
 				return Nothing
-			Just fo -> return $ Just $ LockHandle ph fo
+			Just fo -> Just <$> mkLockHandle pool file ph fo
+
+mkLockHandle :: P.LockPool -> LockFile -> P.LockHandle -> FileLockOps -> IO LockHandle
+mkLockHandle pool file ph fo = do
+	atomically $ P.registerCloseLockFile pool file (fDropLock fo)
+	return $ LockHandle ph fo
+	
diff --git a/Utility/LockPool/PidLock.hs b/Utility/LockPool/PidLock.hs
index 2b3ee67f9..26ed96f3c 100644
--- a/Utility/LockPool/PidLock.hs
+++ b/Utility/LockPool/PidLock.hs
@@ -32,17 +32,17 @@ import Prelude
 
 -- Takes a pid lock, blocking until the lock is available or the timeout.
 waitLock :: Seconds -> LockFile -> IO LockHandle
-waitLock timeout file = makeLockHandle
+waitLock timeout file = makeLockHandle P.lockPool file
 	-- LockShared for STM lock, because a pid lock can be the top-level
 	-- lock with various other STM level locks gated behind it.
-	(P.waitTakeLock P.lockPool file LockShared)
-	(mk <$> F.waitLock timeout file)
+	(\p f -> P.waitTakeLock p f LockShared)
+	(\f -> mk <$> F.waitLock timeout f)
 
 -- Tries to take a pid lock, but does not block.
 tryLock :: LockFile -> IO (Maybe LockHandle)
-tryLock file = tryMakeLockHandle
-	(P.tryTakeLock P.lockPool file LockShared)
-	(fmap mk <$> F.tryLock file)
+tryLock file = tryMakeLockHandle P.lockPool file
+	(\p f -> P.tryTakeLock p f LockShared)
+	(\f -> fmap mk <$> F.tryLock f)
 
 checkLocked :: LockFile -> IO (Maybe Bool)
 checkLocked file = P.getLockStatus P.lockPool file
diff --git a/Utility/LockPool/Posix.hs b/Utility/LockPool/Posix.hs
index a77ed8f01..2c0b7c78e 100644
--- a/Utility/LockPool/Posix.hs
+++ b/Utility/LockPool/Posix.hs
@@ -33,27 +33,27 @@ import Prelude
 
 -- Takes a shared lock, blocking until the lock is available.
 lockShared :: Maybe FileMode -> LockFile -> IO LockHandle
-lockShared mode file = makeLockHandle
-	(P.waitTakeLock P.lockPool file LockShared)
-	(mk <$> F.lockShared mode file)
+lockShared mode file = makeLockHandle P.lockPool file
+	(\p f -> P.waitTakeLock p f LockShared)
+	(\f -> mk <$> F.lockShared mode f)
 
 -- Takes an exclusive lock, blocking until the lock is available.
 lockExclusive :: Maybe FileMode -> LockFile -> IO LockHandle
-lockExclusive mode file = makeLockHandle
-	(P.waitTakeLock P.lockPool file LockExclusive)
-	(mk <$> F.lockExclusive mode file)
+lockExclusive mode file = makeLockHandle P.lockPool file
+	(\p f -> P.waitTakeLock p f LockExclusive)
+	(\f -> mk <$> F.lockExclusive mode f)
 
 -- Tries to take a shared lock, but does not block.
 tryLockShared :: Maybe FileMode -> LockFile -> IO (Maybe LockHandle)
-tryLockShared mode file = tryMakeLockHandle
-	(P.tryTakeLock P.lockPool file LockShared)
-	(fmap mk <$> F.tryLockShared mode file)
+tryLockShared mode file = tryMakeLockHandle P.lockPool file
+	(\p f -> P.tryTakeLock p f LockShared)
+	(\f -> fmap mk <$> F.tryLockShared mode f)
 
 -- Tries to take an exclusive lock, but does not block.
 tryLockExclusive :: Maybe FileMode -> LockFile -> IO (Maybe LockHandle)
-tryLockExclusive mode file = tryMakeLockHandle
-	(P.tryTakeLock P.lockPool file LockExclusive)
-	(fmap mk <$> F.tryLockExclusive mode file)
+tryLockExclusive mode file = tryMakeLockHandle P.lockPool file
+	(\p f -> P.tryTakeLock p f LockExclusive)
+	(\f -> fmap mk <$> F.tryLockExclusive mode f)
 
 -- Returns Nothing when the file doesn't exist, for cases where
 -- that is different from it not being locked.
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 ()
diff --git a/Utility/LockPool/Windows.hs b/Utility/LockPool/Windows.hs
index 2641ac37d..9001b313f 100644
--- a/Utility/LockPool/Windows.hs
+++ b/Utility/LockPool/Windows.hs
@@ -22,9 +22,9 @@ import Utility.LockPool.STM (LockFile, LockMode(..))
 {- Tries to lock a file with a shared lock, which allows other processes to
  - also lock it shared. Fails if the file is exclusively locked. -}
 lockShared :: LockFile -> IO (Maybe LockHandle)
-lockShared file = tryMakeLockHandle
-	(P.tryTakeLock P.lockPool file LockShared)
-	(fmap mk <$> F.lockShared file)
+lockShared file = tryMakeLockHandle P.lockPool file
+	(\p f -> P.tryTakeLock p f LockShared)
+	(\f -> fmap mk <$> F.lockShared f)
 
 {- Tries to take an exclusive lock on a file. Fails if another process has
  - a shared or exclusive lock.
@@ -33,9 +33,9 @@ lockShared file = tryMakeLockHandle
  - read or write by any other process. So for advisory locking of a file's
  - content, a separate LockFile should be used. -}
 lockExclusive :: LockFile -> IO (Maybe LockHandle)
-lockExclusive file = tryMakeLockHandle
-	(P.tryTakeLock P.lockPool file LockExclusive)
-	(fmap mk <$> F.lockExclusive file)
+lockExclusive file = tryMakeLockHandle P.lockPool file
+	(\p -> P.tryTakeLock f LockExclusive)
+	(\f -> fmap mk <$> F.lockExclusive f)
 
 {- If the initial lock fails, this is a BUSY wait, and does not
  - guarentee FIFO order of waiters. In other news, Windows is a POS. -}