@@ -5,6 +5,7 @@
55module Network.Wai.Handler.Warp.Timeout (
66 Manager
77 , Handle
8+ , TimeoutAction
89 , initialize
910 , stopManager
1011 , register
@@ -17,87 +18,117 @@ module Network.Wai.Handler.Warp.Timeout (
1718 , dummyHandle
1819 ) where
1920
20−import System.Mem.Weak (deRefWeak)
2121#if MIN_VERSION_base(4,6,0)
22−import Control.Concurrent (mkWeakThreadId)
22+import Control.Concurrent (mkWeakThreadId, ThreadId)
2323#else
24−import GHC.Weak (Weak (..))
25−import GHC.Conc.Sync (ThreadId (..))
26−import GHC.IO (IO (IO))
24+import Control.Concurrent (ThreadId(..))
2725import GHC.Exts (mkWeak#)
26+import GHC.IO (IO (IO))
2827#endif
28+import GHC.Weak (Weak (..))
2929import Control.Concurrent (forkIO, threadDelay, myThreadId, killThread)
3030import qualified Control.Exception as E
3131import Control.Monad (forever, void)
32+import Data.IORef (IORef)
3233import qualified Data.IORef as I
33−import System.IO.Unsafe (unsafePerformIO)
3434import Data.Typeable (Typeable)
35+import System.IO.Unsafe (unsafePerformIO)
36+import System.Mem.Weak (deRefWeak)
37+
38+----------------------------------------------------------------
3539
3640-- | A timeout manager
37−newtype Manager = Manager (I.IORef [Handle])
41+newtype Manager = Manager (IORef [Handle])
42+
43+-- | An action to be performed on timeout.
44+type TimeoutAction = IO ()
3845
3946-- | A handle used by 'Manager'
40−--
41−-- First field is action to be performed on timeout.
42−data Handle = Handle (IO ()) (I.IORef State)
47+data Handle = Handle TimeoutAction (IORef State)
48+
49+data State = Active -- Manager turns it to Inactive.
50+ | Inactive -- Manager removes it with timeout action.
51+ | Paused -- Manager does not change it.
52+ | Canceled -- Manager removes it without timeout action.
53+
54+----------------------------------------------------------------
4355
4456-- | A dummy @Handle@.
4557dummyHandle :: Handle
4658dummyHandle = Handle (return ()) (unsafePerformIO $ I.newIORef Active)
4759
48−data State = Active | Inactive | Paused | Canceled
60+----------------------------------------------------------------
4961
62+data TimeoutManagerStopped = TimeoutManagerStopped
63+ deriving (Show, Typeable)
64+instance E.Exception TimeoutManagerStopped
65+
66+----------------------------------------------------------------
67+
68+-- | Creating timeout manager which works every N micro seconds
69+-- where N is the first argument.
5070initialize :: Int -> IO Manager
5171initialize timeout = do
5272 ref <- I.newIORef []
5373 void . forkIO $ E.handle ignoreStop $ forever $ do
5474 threadDelay timeout
55− ms <- I.atomicModifyIORef ref (\x -> ([], x))
56− ms' <- go ms id
57− I.atomicModifyIORef ref (\x -> (ms' x, ()))
75+ -- FIXME: isn't mask_ necessary?
76+ old <- I.atomicModifyIORef ref (\x -> ([], x))
77+ merge <- prune old id
78+ I.atomicModifyIORef ref (\new -> (merge new, ()))
5879 return $ Manager ref
5980 where
6081 ignoreStop TimeoutManagerStopped = return ()
6182
62− go [] front = return front
63− go (m@(Handle onTimeout iactive):rest) front = do
64− state <- I.atomicModifyIORef iactive (\x -> (go' x, x))
83+ prune [] front = return front
84+ prune (m@(Handle onTimeout iactive):rest) front = do
85+ state <- I.atomicModifyIORef iactive (\x -> (inactivate x, x))
6586 case state of
6687 Inactive -> do
6788 onTimeout `E.catch` ignoreAll
68− go rest front
69− Canceled -> go rest front
70− _ -> go rest (front . (:) m)
71− go' Active = Inactive
72− go' x = x
89+ prune rest front
90+ Canceled -> prune rest front
91+ _ -> prune rest (front . (:) m)
92+ inactivate Active = Inactive
93+ inactivate x = x
7394
74−data TimeoutManagerStopped = TimeoutManagerStopped
75− deriving (Show, Typeable)
76−instance E.Exception TimeoutManagerStopped
95+----------------------------------------------------------------
7796
7897stopManager :: Manager -> IO ()
7998stopManager (Manager ihandles) = E.mask_ $ do
8099 -- Put an undefined value in the IORef to kill the worker thread (yes, it's
81100 -- a bit of a hack)
82101 !handles <- I.atomicModifyIORef ihandles $ \h -> (E.throw TimeoutManagerStopped, h)
83− mapM_ go handles
102+ mapM_ fire handles
84103 where
85− go (Handle onTimeout _) = onTimeout `E.catch` ignoreAll
104+ fire (Handle onTimeout _) = onTimeout `E.catch` ignoreAll
86105
87106ignoreAll :: E.SomeException -> IO ()
88107ignoreAll _ = return ()
89108
90−register :: Manager -> IO () -> IO Handle
109+----------------------------------------------------------------
110+
111+-- | Registering a timeout action.
112+register :: Manager -> TimeoutAction -> IO Handle
91113register (Manager ref) onTimeout = do
92114 iactive <- I.newIORef Active
93115 let h = Handle onTimeout iactive
94116 I.atomicModifyIORef ref (\x -> (h : x, ()))
95117 return h
96118
119+-- | Registering a timeout action of killing this thread.
97120registerKillThread :: Manager -> IO Handle
98121registerKillThread m = do
99122 wtid <- myThreadId >>= mkWeakThreadId
100− register m $ deRefWeak wtid >>= maybe (return ()) killThread
123+ register m $ killIfExist wtid
124+
125+-- If ThreadId is hold referred by a strong reference,
126+-- it leaks even after the thread is killed.
127+-- So, let's use a weak reference so that CG can throw ThreadId away.
128+-- deRefWeak checks if ThreadId referenced by the weak reference
129+-- exists. If exists, it means that the thread is alive.
130+killIfExist :: Weak ThreadId -> TimeoutAction
131+killIfExist wtid = deRefWeak wtid >>= maybe (return ()) killThread
101132
102133#if !MIN_VERSION_base(4,6,0)
103134mkWeakThreadId :: ThreadId -> IO (Weak ThreadId)
@@ -106,11 +137,29 @@ mkWeakThreadId t@(ThreadId t#) = IO $ \s ->
106137 (# s1, w #) -> (# s1, Weak w #)
107138#endif
108139
109−tickle, pause, resume, cancel :: Handle -> IO ()
140+----------------------------------------------------------------
141+
142+-- | Setting the state to active.
143+-- 'Manager' turns active to inactive repeatedly.
144+tickle :: Handle -> IO ()
110145tickle (Handle _ iactive) = I.writeIORef iactive Active
146+
147+-- | Setting the state to canceled.
148+-- 'Manager' eventually removes this without timeout action.
149+cancel :: Handle -> IO ()
150+cancel (Handle _ iactive) = I.writeIORef iactive Canceled
151+
152+-- | Setting the state to paused.
153+-- 'Manager' does not change the value.
154+pause :: Handle -> IO ()
111155pause (Handle _ iactive) = I.writeIORef iactive Paused
156+
157+-- | Setting the state to active.
158+-- This is an alias to 'ticle'.
159+resume :: Handle -> IO ()
112160resume = tickle
113−cancel (Handle _ iactive) = I.writeIORef iactive Canceled
161+
162+----------------------------------------------------------------
114163
115164-- | Call the inner function with a timeout manager.
116165withManager :: Int -- ^ timeout in microseconds