Merge pull request yesodweb#1015 from kazu-yamamoto/with-handle

withHandle

time-manager/ChangeLog.md

@@ -1,5 +1,9 @@
 # ChangeLog for time-manager
 
+## 0.1.3
+
+* Providing `withHandle` and `withHandleKillThread`.
+
 ## 0.1.2
 
 * Holding `Weak ThreadId` to prevent thread leak again

time-manager/System/TimeManager.hs

@@ -14,16 +14,20 @@ module System.TimeManager (
     withManager,
     withManager',
 
-    -- ** Registration
-    register,
-    registerKillThread,
+    -- ** Registering a timeout action
+    withHandle,
+    withHandleKillThread,
 
     -- ** Control
     tickle,
-    cancel,
     pause,
     resume,
 
+    -- ** Low level
+    register,
+    registerKillThread,
+    cancel,
+
     -- ** Exceptions
     TimeoutThread (..),
 ) where
@@ -96,6 +100,22 @@ killManager = reaperKill
 
 ----------------------------------------------------------------
 
+-- | Registering a timeout action and unregister its handle
+--   when the body action is finished.
+withHandle :: Manager -> TimeoutAction -> (Handle -> IO a) -> IO a
+withHandle mgr onTimeout action =
+    E.bracket (register mgr onTimeout) cancel action
+
+-- | Registering a timeout action of killing this thread and
+--   unregister its handle when the body action is killed or finished.
+withHandleKillThread :: Manager -> TimeoutAction -> (Handle -> IO ()) -> IO ()
+withHandleKillThread mgr onTimeout action =
+    E.handle handler $ E.bracket (registerKillThread mgr onTimeout) cancel action
+  where
+    handler TimeoutThread = return ()
+
+----------------------------------------------------------------
+
 -- | Registering a timeout action.
 register :: Manager -> TimeoutAction -> IO Handle
 register mgr !onTimeout = do
@@ -105,6 +125,34 @@ register mgr !onTimeout = do
     reaperAdd mgr h
     return h
 
+-- | Removing the 'Handle' from the 'Manager' immediately.
+cancel :: Handle -> IO ()
+cancel (Handle mgr _ stateRef) = do
+    _ <- reaperModify mgr filt
+    return ()
+  where
+    -- It's very important that this function forces the whole workload so we
+    -- don't retain old handles, otherwise disasterous leaks occur.
+    filt [] = []
+    filt (h@(Handle _ _ stateRef') : hs)
+        | stateRef == stateRef' = hs
+        | otherwise =
+            let !hs' = filt hs
+             in h : hs'
+
+----------------------------------------------------------------
+
+-- | The asynchronous exception thrown if a thread is registered via
+-- 'registerKillThread'.
+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"
+
 -- | Registering a timeout action of killing this thread.
 --   'TimeoutThread' is thrown to the thread which called this
 --   function on timeout. Catch 'TimeoutThread' if you don't
@@ -121,40 +169,13 @@ registerKillThread m onTimeout = do
                 Nothing -> return ()
                 Just tid' -> E.throwTo tid' TimeoutThread
 
--- | The asynchronous exception thrown if a thread is registered via
--- 'registerKillThread'.
-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
 
--- | Removing the 'Handle' from the 'Manager' immediately.
-cancel :: Handle -> IO ()
-cancel (Handle mgr _ stateRef) = do
-    _ <- reaperModify mgr filt
-    return ()
-  where
-    -- It's very important that this function forces the whole workload so we
-    -- don't retain old handles, otherwise disasterous leaks occur.
-    filt [] = []
-    filt (h@(Handle _ _ stateRef') : hs)
-        | stateRef == stateRef' =
-            hs
-        | otherwise =
-            let !hs' = filt hs
-             in h : hs'
-
 -- | Setting the state to paused.
 --   'Manager' does not change the value.
 pause :: Handle -> IO ()

time-manager/time-manager.cabal

@@ -1,5 +1,5 @@
 Name:                time-manager
-Version:             0.1.2
+Version:             0.1.3
 Synopsis:            Scalable timer
 License:             MIT
 License-file:        LICENSE

warp/Network/Wai/Handler/Warp/HTTP2.hs

@@ -106,10 +106,6 @@ http2server label settings ii transport addr app h2req0 aux0 response = do
             logResponse req st msiz
             mapM_ (logPushPromise req) pps
         Left e
-          -- killed by the local worker manager
-          | Just E.ThreadKilled  <- E.fromException e -> return ()
-          -- killed by the local timeout manager
-          | Just T.TimeoutThread <- E.fromException e -> return ()
           | isAsyncException e -> E.throwIO e
           | otherwise -> do
             S.settingsOnException settings (Just req) e

warp/Network/Wai/Handler/Warp/Run.hs

@@ -355,7 +355,7 @@ fork set mkConn addr app counter ii = settingsFork set $ \unmask -> do
 
     -- We need to register a timeout handler for this thread, and
     -- cancel that handler as soon as we exit.
-    serve unmask (conn, transport) = E.bracket register cancel $ \th -> do
+    serve unmask (conn, transport) = T.withHandleKillThread (timeoutManager ii) (return ()) $ \th -> do
         -- We now have fully registered a connection close handler in
         -- the case of all exceptions, so it is safe to once again
         -- allow async exceptions.
@@ -368,9 +368,6 @@ fork set mkConn addr app counter ii = settingsFork set $ \unmask -> do
                 -- Actually serve this connection.  bracket with closeConn
                 -- above ensures the connection is closed.
                 when goingon $ serveConnection conn ii th addr transport set app
-      where
-        register = T.registerKillThread (timeoutManager ii) (return ())
-        cancel = T.cancel
 
     onOpen adr = increase counter >> settingsOnOpen set adr
     onClose adr _ = decrease counter >> settingsOnClose set adr

warp/warp.cabal

@@ -102,7 +102,7 @@ library
         stm >=2.3,
         streaming-commons >=0.1.10,
         text,
-        time-manager >=0.1 && <0.2,
+        time-manager >=0.1.3 && <0.2,
         vault >=0.3,
         wai >=3.2.4 && <3.3,
         word8