1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
|
{- 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 locks created at the same time this is
-- running. With the lock pool 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 ()
|
