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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
|
{-# LANGUAGE DeriveDataTypeable #-}
module System.TimeManager (
-- ** Types
Manager
, TimeoutAction
, Handle
-- ** Manager
, initialize
, stopManager
, killManager
, withManager
, withManager'
-- ** Registration
, register
, registerKillThread
-- ** Control
, tickle
, cancel
, pause
, resume
-- ** Exceptions
, TimeoutThread (..)
) where
import Control.Concurrent (myThreadId)
import qualified UnliftIO.Exception as E
import Control.Reaper
import Data.Typeable (Typeable)
import Data.IORef (IORef)
import qualified Data.IORef as I
----------------------------------------------------------------
-- | A timeout manager
type Manager = Reaper [Handle] Handle
-- | An action to be performed on timeout.
type TimeoutAction = IO ()
-- | A handle used by 'Manager'
data Handle = Handle !(IORef TimeoutAction) !(IORef State)
data State = Active -- Manager turns it to Inactive.
| Inactive -- Manager removes it with timeout action.
| Paused -- Manager does not change it.
| Canceled -- Manager removes it without timeout action.
----------------------------------------------------------------
-- | Creating timeout manager which works every N micro seconds
-- where N is the first argument.
initialize :: Int -> IO Manager
initialize timeout = mkReaper defaultReaperSettings
{ reaperAction = mkListAction prune
, reaperDelay = timeout
}
where
prune m@(Handle actionRef stateRef) = do
state <- I.atomicModifyIORef' stateRef (\x -> (inactivate x, x))
case state of
Inactive -> do
onTimeout <- I.readIORef actionRef
onTimeout `E.catch` ignoreAll
return Nothing
Canceled -> return Nothing
_ -> return $ Just m
inactivate Active = Inactive
inactivate x = x
----------------------------------------------------------------
-- | Stopping timeout manager with onTimeout fired.
stopManager :: Manager -> IO ()
stopManager mgr = E.mask_ (reaperStop mgr >>= mapM_ fire)
where
fire (Handle actionRef _) = do
onTimeout <- I.readIORef actionRef
onTimeout `E.catch` ignoreAll
ignoreAll :: E.SomeException -> IO ()
ignoreAll _ = return ()
-- | Killing timeout manager immediately without firing onTimeout.
killManager :: Manager -> IO ()
killManager = reaperKill
----------------------------------------------------------------
-- | Registering a timeout action.
register :: Manager -> TimeoutAction -> IO Handle
register mgr onTimeout = do
actionRef <- I.newIORef onTimeout
stateRef <- I.newIORef Active
let h = Handle actionRef stateRef
reaperAdd mgr h
return h
-- | Registering a timeout action of killing this thread.
registerKillThread :: Manager -> TimeoutAction -> IO Handle
registerKillThread m onTimeout = do
-- If we hold ThreadId, the stack and data of the thread is leaked.
-- If we hold Weak ThreadId, the stack is released. However, its
-- data is still leaked probably because of a bug of GHC.
-- So, let's just use ThreadId and release ThreadId by
-- overriding the timeout action by "cancel".
tid <- myThreadId
-- First run the timeout action in case the child thread is masked.
register m $ onTimeout `E.finally` E.throwTo tid TimeoutThread
data TimeoutThread = TimeoutThread
deriving Typeable
instance E.Exception TimeoutThread where
toException = E.asyncExceptionToException
fromException = E.asyncExceptionFromException
instance Show TimeoutThread where
show TimeoutThread = "Thread killed by timeout manager"
----------------------------------------------------------------
-- | Setting the state to active.
-- 'Manager' turns active to inactive repeatedly.
tickle :: Handle -> IO ()
tickle (Handle _ stateRef) = I.writeIORef stateRef Active
-- | Setting the state to canceled.
-- 'Manager' eventually removes this without timeout action.
cancel :: Handle -> IO ()
cancel (Handle actionRef stateRef) = do
I.writeIORef actionRef (return ()) -- ensuring to release ThreadId
I.writeIORef stateRef Canceled
-- | Setting the state to paused.
-- 'Manager' does not change the value.
pause :: Handle -> IO ()
pause (Handle _ stateRef) = I.writeIORef stateRef Paused
-- | Setting the paused state to active.
-- This is an alias to 'tickle'.
resume :: Handle -> IO ()
resume = tickle
----------------------------------------------------------------
-- | Call the inner function with a timeout manager.
-- 'stopManager' is used after that.
withManager :: Int -- ^ timeout in microseconds
-> (Manager -> IO a)
-> IO a
withManager timeout f = E.bracket (initialize timeout)
stopManager
f
-- | Call the inner function with a timeout manager.
-- 'killManager' is used after that.
withManager' :: Int -- ^ timeout in microseconds
-> (Manager -> IO a)
-> IO a
withManager' timeout f = E.bracket (initialize timeout)
killManager
f
|
