@@ -56,11 +56,10 @@ import GHC.IO (IO (IO))
56import Control.Concurrent (myThreadId)
56import Control.Concurrent (myThreadId)
57import qualified Control.Exception as E
57import qualified Control.Exception as E
58import GHC.Weak (Weak (..))
58import GHC.Weak (Weak (..))
59−import Network.Wai.Handler.Warp.IORef (IORef)
60−import qualified Network.Wai.Handler.Warp.IORef as I
61import System.Mem.Weak (deRefWeak)
59import System.Mem.Weak (deRefWeak)
62import Data.Typeable (Typeable)
60import Data.Typeable (Typeable)
63import Control.Reaper
61import Control.Reaper
62+import qualified Data.Atomics.Counter.Unboxed as C
65----------------------------------------------------------------
64----------------------------------------------------------------
@@ -71,12 +70,14 @@ type Manager = Reaper [Handle] Handle
71type TimeoutAction = IO ()
70type TimeoutAction = IO ()
73-- | A handle used by 'Manager'
72-- | A handle used by 'Manager'
74−data Handle = Handle TimeoutAction (IORef State)
73+data Handle = Handle TimeoutAction {-# UNPACK #-} !C.AtomicCounter
75+-- Four states for the AtomicCounter:
76−data State = Active -- Manager turns it to Inactive.
77+-- 1: Active -- Manager turns it to Inactive.
77− | Inactive -- Manager removes it with timeout action.
78+-- 2: Inactive -- Manager removes it with timeout action.
78− | Paused -- Manager does not change it.
79+-- 3: Paused -- Manager does not change it.
79− | Canceled -- Manager removes it without timeout action.
80+-- 4: Canceled -- Manager removes it without timeout action.
81----------------------------------------------------------------
82----------------------------------------------------------------
@@ -89,17 +90,15 @@ initialize timeout = mkReaper defaultReaperSettings
91 prune m@(Handle onTimeout iactive) = do
92 prune m@(Handle onTimeout iactive) = do
92− state <- I.atomicModifyIORef' iactive (\x -> (inactivate x, x))
93+ -- Try to change from active to inactive
94+ (wasActive, newState) <- C.casCounter iactive 1 2
93− case state of
95+ case newState of
94− Inactive -> do
96+ 2 | not wasActive -> do -- inactive
95 onTimeout `E.catch` ignoreAll
97 onTimeout `E.catch` ignoreAll
96 return Nothing
98 return Nothing
97− Canceled -> return Nothing
99+ 4 -> return Nothing -- canceled
98 _ -> return $ Just m
100 _ -> return $ Just m
100− inactivate Active = Inactive
103----------------------------------------------------------------
102----------------------------------------------------------------
105-- | Stopping timeout manager.
104-- | Stopping timeout manager.
@@ -116,7 +115,7 @@ ignoreAll _ = return ()
116-- | Registering a timeout action.
115-- | Registering a timeout action.
117register :: Manager -> TimeoutAction -> IO Handle
116register :: Manager -> TimeoutAction -> IO Handle
118register mgr onTimeout = do
117register mgr onTimeout = do
119− iactive <- I.newIORef Active
118+ iactive <- C.newCounter 1
120 let h = Handle onTimeout iactive
119 let h = Handle onTimeout iactive
121 reaperAdd mgr h
120 reaperAdd mgr h
@@ -153,17 +152,17 @@ mkWeakThreadId t@(ThreadId t#) = IO $ \s ->
153-- | Setting the state to active.
152-- | Setting the state to active.
154-- 'Manager' turns active to inactive repeatedly.
153-- 'Manager' turns active to inactive repeatedly.
155tickle :: Handle -> IO ()
154tickle :: Handle -> IO ()
156−tickle (Handle _ iactive) = I.writeIORef iactive Active
155+tickle (Handle _ iactive) = C.writeCounter iactive 1
158-- | Setting the state to canceled.
157-- | Setting the state to canceled.
159-- 'Manager' eventually removes this without timeout action.
158-- 'Manager' eventually removes this without timeout action.
160cancel :: Handle -> IO ()
159cancel :: Handle -> IO ()
161−cancel (Handle _ iactive) = I.writeIORef iactive Canceled
160+cancel (Handle _ iactive) = C.writeCounter iactive 4
163-- | Setting the state to paused.
162-- | Setting the state to paused.
164-- 'Manager' does not change the value.
163-- 'Manager' does not change the value.
165pause :: Handle -> IO ()
164pause :: Handle -> IO ()
166−pause (Handle _ iactive) = I.writeIORef iactive Paused
165+pause (Handle _ iactive) = C.writeCounter iactive 3
168-- | Setting the paused state to active.
167-- | Setting the paused state to active.
169-- This is an alias to 'tickle'.
168-- This is an alias to 'tickle'.