Label all threads and make sure they are killable

Labeling all threads ensures we can identify them when analyzing a GHC
event log file using ghc-events-analyze.

I also solve the current bug that all threads created by snap are
unkillable when you start the server in a masked state. It is solved by
using:

  forkIOWithUnmask $ \unmask -> unmask ...

so that asynchronous exception can always be delivered to the forked
thread.

snap-server.cabal

@@ -82,7 +82,8 @@ Library
     Snap.Internal.Http.Server.HttpPort,
     Snap.Internal.Http.Server.SimpleBackend,
     Snap.Internal.Http.Server.TimeoutManager,
-    Snap.Internal.Http.Server.TLS
+    Snap.Internal.Http.Server.TLS,
+    Control.Concurrent.Extended
 
   build-depends:
     attoparsec                >= 0.10     && < 0.12,

src/Control/Concurrent/Extended.hs

@@ -0,0 +1,84 @@
+{-# LANGUAGE RankNTypes #-}
+
+-- | Handy functions that should really be merged into
+-- Control.Concurrent itself.
+module Control.Concurrent.Extended
+    ( forkIOLabeled
+    , forkIOLabeledWithUnmask
+
+    , forkOnLabeled
+    , forkOnLabeledWithUnmask
+    ) where
+
+import Control.Exception
+import Control.Concurrent
+import GHC.Conc.Sync (labelThread)
+
+-- | Sparks off a new thread using 'forkIO' to run the given IO
+-- computation, but first labels the thread with the given label
+-- (using 'labelThread').
+--
+-- The implementation makes sure that asynchronous exceptions are
+-- masked until the given computation is executed. This ensures the
+-- thread will always be labeled which guarantees you can always
+-- easily find it in the GHC event log.
+--
+-- Note that the given computation is executed in the masked state of
+-- the calling thread.
+--
+-- Returns the 'ThreadId' of the newly created thread.
+forkIOLabeled :: String -- ^ Thread label
+              -> IO ()
+              -> IO ThreadId
+forkIOLabeled label m =
+    mask $ \restore -> forkIO $ do
+      tid <- myThreadId
+      labelThread tid label
+      restore m
+
+-- | Like 'forkIOLabeled', but lets you specify on which capability
+-- (think CPU) the thread should run.
+forkOnLabeled :: String -- ^ Thread label
+              -> Int    -- ^ Capability
+              -> IO ()
+              -> IO ThreadId
+forkOnLabeled label cap m =
+    mask $ \restore -> forkOn cap $ do
+      tid <- myThreadId
+      labelThread tid label
+      restore m
+
+-- | Sparks off a new thread using 'forkIOWithUnmask' to run the given
+-- IO computation, but first labels the thread with the given label
+-- (using 'labelThread').
+--
+-- The implementation makes sure that asynchronous exceptions are
+-- masked until the given computation is executed. This ensures the
+-- thread will always be labeled which guarantees you can always
+-- easily find it in the GHC event log.
+--
+-- Like 'forkIOWithUnmask', the given computation is given a function
+-- to unmask asynchronous exceptions. See the documentation of that
+-- function for the motivation.
+--
+-- Returns the 'ThreadId' of the newly created thread.
+forkIOLabeledWithUnmask :: String -- ^ Thread label
+                        -> ((forall a. IO a -> IO a) -> IO ())
+                        -> IO ThreadId
+forkIOLabeledWithUnmask label m =
+    mask_ $ forkIOWithUnmask $ \unmask -> do
+      tid <- myThreadId
+      labelThread tid label
+      m unmask
+
+-- | Like 'forkIOLabeledWithUnmask', but lets you specify on which
+-- capability (think CPU) the thread should run.
+forkOnLabeledWithUnmask :: String -- ^ Thread label
+                        -> Int    -- ^ Capability
+                        -> ((forall a. IO a -> IO a) -> IO ())
+                        -> IO ThreadId
+forkOnLabeledWithUnmask label cap m =
+    mask_ $ forkOnWithUnmask cap $ \unmask -> do
+      tid <- myThreadId
+      labelThread tid label
+      m unmask

src/Snap/Internal/Http/Server/SimpleBackend.hs

@@ -16,6 +16,7 @@ module Snap.Internal.Http.Server.SimpleBackend
 import           Control.Monad.Trans
 
 import           Control.Concurrent hiding (yield)
+import           Control.Concurrent.Extended (forkOnLabeledWithUnmask)
 import           Control.Exception
 import           Control.Monad
 import           Data.ByteString (ByteString)
@@ -24,7 +25,7 @@ import           Data.ByteString.Internal (c2w)
 import           Data.Maybe
 import           Foreign hiding (new)
 import           Foreign.C
-import           GHC.Conc (labelThread, forkOnIO)
+import           GHC.Conc (labelThread)
 import           Network.Socket
 #if !MIN_VERSION_base(4,6,0)
 import           Prelude hiding (catch)
@@ -85,8 +86,10 @@ newLoop :: Int
 newLoop defaultTimeout sockets handler elog cpu = do
     tmgr       <- TM.initialize defaultTimeout getCurrentDateTime
     exit       <- newEmptyMVar
-    accThreads <- forM sockets $ \p -> forkOnIO cpu $
-                  acceptThread defaultTimeout handler tmgr elog cpu p exit
+    accThreads <- forM sockets $ \p -> do
+      let label = "snap-server: " ++ show p ++ " on capability: " ++ show cpu
+      forkOnLabeledWithUnmask label cpu $ \unmask ->
+        unmask $ acceptThread defaultTimeout handler tmgr elog cpu p exit
 
     return $! EventLoopCpu cpu accThreads tmgr exit
 
@@ -115,7 +118,10 @@ acceptThread defaultTimeout handler tmgr elog cpu sock exitMVar =
         (s,addr) <- accept $ Listen.listenSocket sock
         setSocketOption s NoDelay 1
         debug $ "acceptThread: accepted connection from remote: " ++ show addr
-        _ <- forkOnIO cpu (go s addr `catches` cleanup)
+        let label = "snap-server: connection from remote: " ++ show addr
+                    ++ " on socket: " ++ show (fdSocket s)
+        _ <- forkOnLabeledWithUnmask label cpu $ \unmask ->
+               unmask $ go s addr `catches` cleanup
         return ()
 
     loop = do
@@ -156,7 +162,6 @@ runSession defaultTimeout handler tmgr lsock sock addr = do
     curId <- myThreadId
 
     debug $ "Backend.withConnection: running session: " ++ show addr
-    labelThread curId $ "connHndl " ++ show fd
 
     (rport,rhost) <- getAddress addr
     (lport,lhost) <- getSocketName sock >>= getAddress

src/Snap/Internal/Http/Server/TimeoutManager.hs

@@ -14,6 +14,7 @@ module Snap.Internal.Http.Server.TimeoutManager
 
 ------------------------------------------------------------------------------
 import           Control.Concurrent
+import           Control.Concurrent.Extended (forkIOLabeledWithUnmask)
 import           Control.Exception
 import           Control.Monad
 import           Data.IORef
@@ -116,7 +117,8 @@ initialize defaultTimeout getTime = do
 
     let tm = TimeoutManager defaultTimeout getTime conns inact mp mthr
 
-    thr <- forkIO $ managerThread tm
+    thr <- forkIOLabeledWithUnmask "snap-server: timeout manager" $ \unmask ->
+             unmask $ managerThread tm
     putMVar mthr thr
     return tm
 

src/System/FastLogger.hs

@@ -18,6 +18,7 @@ module System.FastLogger
 import           Blaze.ByteString.Builder
 import           Blaze.ByteString.Builder.Char.Utf8
 import           Control.Concurrent
+import           Control.Concurrent.Extended (forkIOLabeledWithUnmask)
 import           Control.Exception
 import           Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as S
@@ -74,7 +75,8 @@ newLoggerWithCustomErrorFunction errAction fp = do
 
     let lg = Logger q dw fp th errAction
 
-    tid <- forkIO $ loggingThread lg
+    tid <- forkIOLabeledWithUnmask "snap-server: logging" $ \unmask ->
+             unmask $ loggingThread lg
     putMVar th tid
 
     return lg