@@ -1,7 +1,4 @@
1−{-# LANGUAGE CPP #-}
21{-# LANGUAGE DeriveDataTypeable #-}
3−{-# LANGUAGE MagicHash #-}
4−{-# LANGUAGE UnboxedTuples #-} -- for GHC 7.4 or earlier
52
63module Network.Wai.Handler.Warp.Timeout (
74 -- ** Types
@@ -11,6 +8,7 @@ module Network.Wai.Handler.Warp.Timeout (
118 -- ** Manager
129 , initialize
1310 , stopManager
11+ , killManager
1412 , withManager
1513 -- ** Registration
1614 , register
@@ -24,25 +22,12 @@ module Network.Wai.Handler.Warp.Timeout (
2422 , TimeoutThread (..)
2523 ) where
2624
27−#ifndef MIN_VERSION_base
28−#define MIN_VERSION_base(x,y,z) 1
29−#endif
30−
31−#if MIN_VERSION_base(4,6,0)
32−import Control.Concurrent (mkWeakThreadId, ThreadId)
33−#else
34−import GHC.Conc (ThreadId(..))
35−import GHC.Exts (mkWeak#)
36−import GHC.IO (IO (IO))
37−#endif
3825import Control.Concurrent (myThreadId)
3926import qualified Control.Exception as E
40−import GHC.Weak (Weak (..))
27+import Control.Reaper
28+import Data.Typeable (Typeable)
4129import Network.Wai.Handler.Warp.IORef (IORef)
4230import qualified Network.Wai.Handler.Warp.IORef as I
43−import System.Mem.Weak (deRefWeak)
44−import Data.Typeable (Typeable)
45−import Control.Reaper
4631
4732----------------------------------------------------------------
4833
@@ -53,7 +38,7 @@ type Manager = Reaper [Handle] Handle
5338type TimeoutAction = IO ()
5439
5540-- | A handle used by 'Manager'
56−data Handle = Handle TimeoutAction (IORef State)
41+data Handle = Handle !(IORef TimeoutAction) !(IORef State)
5742
5843data State = Active -- Manager turns it to Inactive.
5944 | Inactive -- Manager removes it with timeout action.
@@ -70,10 +55,11 @@ initialize timeout = mkReaper defaultReaperSettings
7055 , reaperDelay = timeout
7156 }
7257 where
73− prune m@(Handle onTimeout iactive) = do
74− state <- I.atomicModifyIORef' iactive (\x -> (inactivate x, x))
58+ prune m@(Handle actionRef stateRef) = do
59+ state <- I.atomicModifyIORef' stateRef (\x -> (inactivate x, x))
7560 case state of
7661 Inactive -> do
62+ onTimeout <- I.readIORef actionRef
7763 onTimeout `E.catch` ignoreAll
7864 return Nothing
7965 Canceled -> return Nothing
@@ -84,38 +70,46 @@ initialize timeout = mkReaper defaultReaperSettings
8470
8571----------------------------------------------------------------
8672
87−-- | Stopping timeout manager.
73+-- | Stopping timeout manager with onTimeout fired.
8874stopManager :: Manager -> IO ()
8975stopManager mgr = E.mask_ (reaperStop mgr >>= mapM_ fire)
9076 where
91− fire (Handle onTimeout _) = onTimeout `E.catch` ignoreAll
77+ fire (Handle actionRef _) = do
78+ onTimeout <- I.readIORef actionRef
79+ onTimeout `E.catch` ignoreAll
9280
9381ignoreAll :: E.SomeException -> IO ()
9482ignoreAll _ = return ()
9583
84+-- | Killing timeout manager immediately without firing onTimeout.
85+killManager :: Manager -> IO ()
86+killManager mgr = E.mask_ (reaperStop mgr >>= mapM_ fire)
87+ where
88+ fire (Handle actionRef _) = do
89+ onTimeout <- I.readIORef actionRef
90+ onTimeout `E.catch` ignoreAll
91+
9692----------------------------------------------------------------
9793
9894-- | Registering a timeout action.
9995register :: Manager -> TimeoutAction -> IO Handle
10096register mgr onTimeout = do
101− iactive <- I.newIORef Active
102− let h = Handle onTimeout iactive
97+ actionRef <- I.newIORef onTimeout
98+ stateRef <- I.newIORef Active
99+ let h = Handle actionRef stateRef
103100 reaperAdd mgr h
104101 return h
105102
106103-- | Registering a timeout action of killing this thread.
107104registerKillThread :: Manager -> IO Handle
108105registerKillThread m = do
109− wtid <- myThreadId >>= mkWeakThreadId
110− register m $ killIfExist wtid
111−
112−-- If ThreadId is hold referred by a strong reference,
113−-- it leaks even after the thread is killed.
114−-- So, let's use a weak reference so that CG can throw ThreadId away.
115−-- deRefWeak checks if ThreadId referenced by the weak reference
116−-- exists. If exists, it means that the thread is alive.
117−killIfExist :: Weak ThreadId -> TimeoutAction
118−killIfExist wtid = deRefWeak wtid >>= maybe (return ()) (`E.throwTo` TimeoutThread)
106+ -- If we hold ThreadId, the stack and data of the thread is leaked.
107+ -- If we hold Weak ThreadId, the stack is released. However, its
108+ -- data is still leaked probably because of a bug of GHC.
109+ -- So, let's just use ThreadId and release ThreadId by
110+ -- overriding the timeout action by "cancel".
111+ tid <- myThreadId
112+ register m $ E.throwTo tid TimeoutThread
119113
120114data TimeoutThread = TimeoutThread
121115 deriving Typeable
@@ -123,29 +117,24 @@ instance E.Exception TimeoutThread
123117instance Show TimeoutThread where
124118 show TimeoutThread = "Thread killed by Warp's timeout reaper"
125119
126−#if !MIN_VERSION_base(4,6,0)
127−mkWeakThreadId :: ThreadId -> IO (Weak ThreadId)
128−mkWeakThreadId t@(ThreadId t#) = IO $ \s ->
129− case mkWeak# t# t Nothing s of
130− (# s1, w #) -> (# s1, Weak w #)
131−#endif
132−
133120----------------------------------------------------------------
134121
135122-- | Setting the state to active.
136123-- 'Manager' turns active to inactive repeatedly.
137124tickle :: Handle -> IO ()
138−tickle (Handle _ iactive) = I.writeIORef iactive Active
125+tickle (Handle _ stateRef) = I.writeIORef stateRef Active
139126
140127-- | Setting the state to canceled.
141128-- 'Manager' eventually removes this without timeout action.
142129cancel :: Handle -> IO ()
143−cancel (Handle _ iactive) = I.writeIORef iactive Canceled
130+cancel (Handle actionRef stateRef) = do
131+ I.writeIORef actionRef (return ()) -- ensuring to release ThreadId
132+ I.writeIORef stateRef Canceled
144133
145134-- | Setting the state to paused.
146135-- 'Manager' does not change the value.
147136pause :: Handle -> IO ()
148−pause (Handle _ iactive) = I.writeIORef iactive Paused
137+pause (Handle _ stateRef) = I.writeIORef stateRef Paused
149138
150139-- | Setting the paused state to active.
151140-- This is an alias to 'tickle'.