@@ -56,11 +56,10 @@ import GHC.IO (IO (IO))
5656import Control.Concurrent (myThreadId)
5757import qualified Control.Exception as E
5858import GHC.Weak (Weak (..))
59−import Network.Wai.Handler.Warp.IORef (IORef)
60−import qualified Network.Wai.Handler.Warp.IORef as I
6159import System.Mem.Weak (deRefWeak)
6260import Data.Typeable (Typeable)
6361import Control.Reaper
62+import qualified Data.Atomics.Counter.Unboxed as C
6463
6564----------------------------------------------------------------
6665
@@ -71,12 +70,14 @@ type Manager = Reaper [Handle] Handle
7170type TimeoutAction = IO ()
7271
7372-- | A handle used by 'Manager'
74−data Handle = Handle TimeoutAction (IORef State)
73+data Handle = Handle TimeoutAction {-# UNPACK #-} !C.AtomicCounter
7574
76−data State = Active -- Manager turns it to Inactive.
77− | Inactive -- Manager removes it with timeout action.
78− | Paused -- Manager does not change it.
79− | Canceled -- Manager removes it without timeout action.
75+-- Four states for the AtomicCounter:
76+--
77+-- 1: Active -- Manager turns it to Inactive.
78+-- 2: Inactive -- Manager removes it with timeout action.
79+-- 3: Paused -- Manager does not change it.
80+-- 4: Canceled -- Manager removes it without timeout action.
8081
8182----------------------------------------------------------------
8283
@@ -89,17 +90,15 @@ initialize timeout = mkReaper defaultReaperSettings
8990 }
9091 where
9192 prune m@(Handle onTimeout iactive) = do
92− state <- I.atomicModifyIORef' iactive (\x -> (inactivate x, x))
93− case state of
94− Inactive -> do
93+ -- Try to change from active to inactive
94+ (wasActive, newState) <- C.casCounter iactive 1 2
95+ case newState of
96+ 2 | not wasActive -> do -- inactive
9597 onTimeout `E.catch` ignoreAll
9698 return Nothing
97− Canceled -> return Nothing
99+ 4 -> return Nothing -- canceled
98100 _ -> return $ Just m
99101
100− inactivate Active = Inactive
101− inactivate x = x
102−
103102----------------------------------------------------------------
104103
105104-- | Stopping timeout manager.
@@ -116,7 +115,7 @@ ignoreAll _ = return ()
116115-- | Registering a timeout action.
117116register :: Manager -> TimeoutAction -> IO Handle
118117register mgr onTimeout = do
119− iactive <- I.newIORef Active
118+ iactive <- C.newCounter 1
120119 let h = Handle onTimeout iactive
121120 reaperAdd mgr h
122121 return h
@@ -153,17 +152,17 @@ mkWeakThreadId t@(ThreadId t#) = IO $ \s ->
153152-- | Setting the state to active.
154153-- 'Manager' turns active to inactive repeatedly.
155154tickle :: Handle -> IO ()
156−tickle (Handle _ iactive) = I.writeIORef iactive Active
155+tickle (Handle _ iactive) = C.writeCounter iactive 1
157156
158157-- | Setting the state to canceled.
159158-- 'Manager' eventually removes this without timeout action.
160159cancel :: Handle -> IO ()
161−cancel (Handle _ iactive) = I.writeIORef iactive Canceled
160+cancel (Handle _ iactive) = C.writeCounter iactive 4
162161
163162-- | Setting the state to paused.
164163-- 'Manager' does not change the value.
165164pause :: Handle -> IO ()
166−pause (Handle _ iactive) = I.writeIORef iactive Paused
165+pause (Handle _ iactive) = C.writeCounter iactive 3
167166
168167-- | Setting the paused state to active.
169168-- This is an alias to 'tickle'.