lock pools to work around non-concurrency/composition safety of POSIX fcntl

Utility/LockPool/STM.hs

@@ -0,0 +1,125 @@
+{- STM implementation of lock pools.
+ -
+ - Copyright 2015 Joey Hess <id@joeyh.name>
+ -
+ - License: BSD-2-clause
+ -}
+
+module Utility.LockPool.STM (
+	LockPool,
+	lockPool,
+	LockFile,
+	LockMode(..),
+	LockHandle,
+	waitTakeLock,
+	tryTakeLock,
+	getLockStatus,
+	releaseLock,
+) where
+
+import System.IO.Unsafe (unsafePerformIO)
+import qualified Data.Map.Strict as M
+import Control.Concurrent
+import Control.Concurrent.STM
+import Control.Applicative
+import Control.Exception
+import Control.Monad
+import Data.Maybe
+
+type LockFile = FilePath
+
+data LockMode = LockExclusive | LockShared
+	deriving (Eq)
+
+-- This TMVar is full when the handle is open, and is emptied when it's
+-- closed.
+type LockHandle = TMVar (LockPool, LockFile)
+
+type LockCount = Integer
+
+data LockStatus = LockStatus LockMode LockCount
+
+-- This TMVar is normally kept full.
+type LockPool = TMVar (M.Map LockFile LockStatus)
+
+-- A shared global variable for the lockPool. Avoids callers needing to
+-- maintain state for this implementation detail.
+lockPool :: LockPool
+lockPool = unsafePerformIO (newTMVarIO M.empty)
+{-# NOINLINE lockPool #-}
+
+-- Updates the LockPool, blocking as necessary if another thread is holding
+-- a conflicting lock.
+-- 
+-- Note that when a shared lock is held, an exclusive lock will block.
+-- While that blocking is happening, another call to this function to take
+-- the same shared lock should not be blocked on the exclusive lock.
+-- Keeping the whole Map in a TMVar accomplishes this, at the expense of
+-- sometimes retrying after unrelated changes in the map.
+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)
+			| mode == LockShared && mode' == LockShared ->
+				return $ LockStatus mode (succ n)
+			| n > 0 -> retry -- wait for lock
+		_ -> return $ LockStatus mode 1
+	putTMVar pool (M.insert file v m)
+	newTMVar (pool, file)
+
+-- Avoids blocking if another thread is holding a conflicting lock.
+tryTakeLock :: LockPool -> LockFile -> LockMode -> STM (Maybe LockHandle)
+tryTakeLock pool file mode =
+	(Just <$> waitTakeLock pool file mode)
+		`orElse`
+	return Nothing
+
+-- Checks if a lock is being held. If it's held by the current process,
+-- runs the getdefault action; otherwise runs the checker action.
+--
+-- Note that the lock pool is left empty while the checker action is run.
+-- This allows checker actions that open/close files, and so would be in
+-- danger of conflicting with existing locks. Since the lock pool is
+-- kept empty, anything that attempts to take a lock will block,
+-- avoiding that race.
+getLockStatus :: LockPool -> LockFile -> IO v -> IO (Maybe v) -> IO (Maybe v)
+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
+			_ -> False
+		if threadlocked
+			then do
+				putTMVar pool m
+				return Nothing
+			else return $ Just $ atomically $ putTMVar pool m
+	case v of
+		Nothing -> Just <$> getdefault
+		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.
+--
+-- 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)
+  where
+	go (Just (pool, file)) = do
+		(m, unused) <- atomically $ do
+			m <- takeTMVar pool
+			return $ case M.lookup file m of
+				Just (LockStatus mode n)
+					| n == 1 -> (M.delete file m, True)
+					| otherwise ->
+						(M.insert file (LockStatus mode (pred n)) m, False)
+				Nothing -> (m, True)
+		when unused
+			closelockfile
+		atomically $ putTMVar pool m
+	-- The LockHandle was already closed.
+	go Nothing = return ()