Merge pull request yesodweb#1013 from kazu-yamamoto/rethrow-async-exc…

…eption

Rethrow async exception

auto-update/Control/Reaper.hs

@@ -207,18 +207,22 @@ reaper settings@ReaperSettings{..} stateRef tidRef = do
     !merge <- reaperAction wl
     -- Merging the left jobs and new jobs.
     -- If there is no jobs, this thread finishes.
-    next <- atomicModifyIORef' stateRef (check merge)
-    next
+    cont <- atomicModifyIORef' stateRef (check merge)
+    if cont
+        then
+            reaper settings stateRef tidRef
+        else
+            writeIORef tidRef Nothing
   where
     swapWithEmpty NoReaper = error "Control.Reaper.reaper: unexpected NoReaper (1)"
     swapWithEmpty (Workload wl) = (Workload reaperEmpty, wl)
 
     check _ NoReaper = error "Control.Reaper.reaper: unexpected NoReaper (2)"
     check merge (Workload wl)
         -- If there is no job, reaper is terminated.
-        | reaperNull wl' = (NoReaper, writeIORef tidRef Nothing)
+        | reaperNull wl' = (NoReaper, False)
         -- If there are jobs, carry them out.
-        | otherwise = (Workload wl', reaper settings stateRef tidRef)
+        | otherwise = (Workload wl', True)
       where
         wl' = merge wl
 

time-manager/System/TimeManager.hs

@@ -28,12 +28,13 @@ module System.TimeManager (
     TimeoutThread (..),
 ) where
 
-import Control.Concurrent (myThreadId)
+import Control.Concurrent (myThreadId, mkWeakThreadId)
 import qualified Control.Exception as E
 import Control.Reaper
 import Data.IORef (IORef)
 import qualified Data.IORef as I
 import Data.Typeable (Typeable)
+import GHC.Weak (deRefWeak)
 
 ----------------------------------------------------------------
 
@@ -107,14 +108,14 @@ register mgr !onTimeout = do
 -- | Registering a timeout action of killing this thread.
 registerKillThread :: Manager -> TimeoutAction -> IO Handle
 registerKillThread m onTimeout = do
-    -- If we hold ThreadId, the stack and data of the thread is leaked.
-    -- If we hold Weak ThreadId, the stack is released. However, its
-    -- data is still leaked probably because of a bug of GHC.
-    -- So, let's just use ThreadId and release ThreadId by
-    -- overriding the timeout action by "cancel".
     tid <- myThreadId
+    wtid <- mkWeakThreadId tid
     -- First run the timeout action in case the child thread is masked.
-    register m $ onTimeout `E.finally` E.throwTo tid TimeoutThread
+    register m $ onTimeout `E.finally` do
+        mtid <- deRefWeak wtid
+        case mtid of
+          Nothing -> return ()
+          Just tid' -> E.throwTo tid' TimeoutThread
 
 data TimeoutThread = TimeoutThread
     deriving (Typeable)

warp-tls/Network/Wai/Handler/WarpTLS.hs

@@ -335,6 +335,17 @@ mkConn tlsset set s params = do
 
 ----------------------------------------------------------------
 
+isAsyncException :: Exception e => e -> Bool
+isAsyncException e =
+    case E.fromException (E.toException e) of
+        Just (E.SomeAsyncException _) -> True
+        Nothing -> False
+
+throughAsync :: IO a -> SomeException -> IO a
+throughAsync action (SomeException e)
+  | isAsyncException e = E.throwIO e
+  | otherwise          = action
+
 httpOverTls
     :: TLS.TLSParams params
     => TLSSettings
@@ -360,13 +371,13 @@ httpOverTls TLSSettings{..} set s bs0 params =
         case mconn of
           Nothing -> throwIO IncompleteHeaders
           Just conn -> return conn
-    wrappedRecvN recvN n = handle (\(SomeException _) -> mempty) $ recvN n
+    wrappedRecvN recvN n = handle (throughAsync (return "")) $ recvN n
     backend recvN =
         TLS.Backend
             { TLS.backendFlush = return ()
 #if MIN_VERSION_network(3,1,1)
             , TLS.backendClose =
-                gracefulClose s 5000 `E.catch` \(SomeException _) -> return ()
+                gracefulClose s 5000 `E.catch` throughAsync (return ())
 #else
             , TLS.backendClose = close s
 #endif

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

@@ -122,6 +122,7 @@ http1server settings ii conn transport app addr th istatus src =
         | Just NoKeepAliveRequest <- fromException e = return ()
         -- No valid request
         | Just (BadFirstLine _) <- fromException e = return ()
+        | isAsyncException e = throwIO e
         | otherwise = do
             _ <-
                 sendErrorResponse
@@ -204,6 +205,7 @@ processRequest settings ii conn app th istatus src req mremainingRef idxhdr next
         Right ResponseReceived -> return ()
         Left (e :: SomeException)
             | Just (ExceptionInsideResponseBody e') <- fromException e -> throwIO e'
+            | isAsyncException e -> throwIO e
             | otherwise -> do
                 keepAlive <- sendErrorResponse settings ii conn th istatus req e
                 settingsOnException settings (Just req) e

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

@@ -110,6 +110,7 @@ http2server label settings ii transport addr app h2req0 aux0 response = do
           | 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
             let ersp = S.settingsOnExceptionResponse settings e
@@ -152,7 +153,7 @@ wrappedRecvN th slowlorisSize readN bufsize = do
     return bs
   where
     handler :: E.SomeException -> IO ByteString
-    handler _ = return ""
+    handler = throughAsync (return "")
 
 -- connClose must not be called here since Run:fork calls it
 goaway :: Connection -> H2.ErrorCodeId -> ByteString -> IO ()

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

@@ -10,9 +10,12 @@ module Network.Wai.Handler.Warp.Imports (
     module Data.Word,
     module Data.Maybe,
     module Numeric,
+    throughAsync,
+    isAsyncException,
 ) where
 
 import Control.Applicative
+import Control.Exception
 import Control.Monad
 import Data.Bits
 import Data.ByteString.Internal (ByteString (..))
@@ -23,3 +26,14 @@ import Data.Monoid
 import Data.Ord
 import Data.Word
 import Numeric
+
+isAsyncException :: Exception e => e -> Bool
+isAsyncException e =
+    case fromException (toException e) of
+        Just (SomeAsyncException _) -> True
+        Nothing -> False
+
+throughAsync :: IO a -> SomeException -> IO a
+throughAsync action (SomeException e)
+    | isAsyncException e = throwIO e
+    | otherwise = action

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

@@ -81,7 +81,7 @@ socketConnection _ s = do
                             else settingsGracefulCloseTimeout1 set
                 if tm == 0
                     then close s
-                    else gracefulClose s tm `E.catch` \(E.SomeException _) -> return ()
+                    else gracefulClose s tm `E.catch` throughAsync (return ())
 #else
             , connClose = close s
 #endif
@@ -179,7 +179,7 @@ runSettingsSocket set@Settings{settingsAccept = accept'} socket app = do
         (s, sa) <- accept' socket
         setSocketCloseOnExec s
         -- NoDelay causes an error for AF_UNIX.
-        setSocketOption s NoDelay 1 `E.catch` \(E.SomeException _) -> return ()
+        setSocketOption s NoDelay 1 `E.catch` throughAsync (return ())
         conn <- socketConnection set s
         return (conn, sa)
 
@@ -369,7 +369,7 @@ fork set mkConn addr app counter ii = settingsFork set $ \unmask -> do
                 -- above ensures the connection is closed.
                 when goingon $ serveConnection conn ii th addr transport set app
       where
-        register = T.registerKillThread (timeoutManager ii) (connClose conn)
+        register = T.registerKillThread (timeoutManager ii) (return ())
         cancel = T.cancel
 
     onOpen adr = increase counter >> settingsOnOpen set adr

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

@@ -9,15 +9,15 @@
 
 module Network.Wai.Handler.Warp.Settings where
 
-import Control.Exception (SomeException, fromException)
+import Control.Exception (SomeException(..), fromException, throw)
 import qualified Data.ByteString.Builder as Builder
 import qualified Data.ByteString.Char8 as C8
 import Data.Streaming.Network (HostPreference)
 import qualified Data.Text as T
 import qualified Data.Text.IO as TIO
 import Data.Version (showVersion)
 import GHC.IO (IO (IO), unsafeUnmask)
-import GHC.IO.Exception (AsyncException (ThreadKilled), IOErrorType (..))
+import GHC.IO.Exception (IOErrorType (..))
 import GHC.Prim (fork#)
 import qualified Network.HTTP.Types as H
 import Network.Socket (SockAddr, Socket, accept)
@@ -228,12 +228,11 @@ defaultSettings =
 -- Since 2.1.3
 defaultShouldDisplayException :: SomeException -> Bool
 defaultShouldDisplayException se
-    | Just ThreadKilled <- fromException se = False
     | Just (_ :: InvalidRequest) <- fromException se = False
     | Just (ioeGetErrorType -> et) <- fromException se
     , et == ResourceVanished || et == InvalidArgument =
         False
-    | Just TimeoutThread <- fromException se = False
+    | isAsyncException se = False
     | otherwise = True
 
 -- | Printing an exception to standard error
@@ -255,6 +254,7 @@ defaultOnException _ e =
 -- Since 3.2.27
 defaultOnExceptionResponse :: SomeException -> Response
 defaultOnExceptionResponse e
+    | isAsyncException e = throw e
     | Just PayloadTooLarge <- fromException e =
         responseLBS
             H.status413