Introduce forkManaged

Network/HTTP2/Arch/Manager.hs

@@ -8,12 +8,13 @@ module Network.HTTP2.Arch.Manager (
   , setAction
   , stop
   , spawnAction
-  , addMyId
+  , forkManaged
   , deleteMyId
   , timeoutKillThread
   , timeoutClose
   ) where
 
+import Control.Exception
 import Data.Foldable
 import Data.IORef
 import Data.Set (Set)
@@ -77,13 +78,34 @@ stop (Manager q _ _) = atomically $ writeTQueue q Stop
 spawnAction :: Manager -> IO ()
 spawnAction (Manager q _ _) = atomically $ writeTQueue q Spawn
 
+----------------------------------------------------------------
+
+-- | Fork managed thread
+--
+-- This guarantees that the thread ID is added to the manager's queue before
+-- the thread starts, and is removed again when the thread terminates
+-- (normally or abnormally).
+forkManaged :: Manager -> IO () -> IO ()
+forkManaged mgr io =
+    void $ mask_ $ forkIOWithUnmask $ \unmask -> do
+      addMyId mgr
+      r <- unmask io `onException` deleteMyId mgr
+      deleteMyId mgr
+      return r
+
 -- | Adding my thread id to the kill-thread list on stopping.
+--
+-- This is not part of the public API; see 'forkManaged' instead.
 addMyId :: Manager -> IO ()
 addMyId (Manager q _ _) = do
     tid <- myThreadId
     atomically $ writeTQueue q $ Add tid
 
 -- | Deleting my thread id from the kill-thread list on stopping.
+--
+-- This is /only/ necessary when you want to remove the thread's ID from
+-- the manager /before/ the thread terminates (thereby assuming responsibility
+-- for thread cleanup yourself).
 deleteMyId :: Manager -> IO ()
 deleteMyId (Manager q _ _) = do
     tid <- myThreadId

Network/HTTP2/Arch/Queue.hs

@@ -2,23 +2,17 @@
 
 module Network.HTTP2.Arch.Queue where
 
-import UnliftIO.Concurrent (forkIO)
-import UnliftIO.Exception (bracket)
 import UnliftIO.STM
 
-import Imports
 import Network.HTTP2.Arch.Manager
 import Network.HTTP2.Arch.Types
 
 {-# INLINE forkAndEnqueueWhenReady #-}
 forkAndEnqueueWhenReady :: IO () -> TQueue (Output Stream) -> Output Stream -> Manager -> IO ()
-forkAndEnqueueWhenReady wait outQ out mgr = bracket setup teardown $ \_ ->
-    void . forkIO $ do
+forkAndEnqueueWhenReady wait outQ out mgr =
+    forkManaged mgr $ do
         wait
         enqueueOutput outQ out
-  where
-    setup = addMyId mgr
-    teardown _ = deleteMyId mgr
 
 {-# INLINE enqueueOutput #-}
 enqueueOutput :: TQueue (Output Stream) -> Output Stream -> IO ()

Network/HTTP2/Client/Run.hs

@@ -72,9 +72,7 @@ sendRequest ctx@Context{..} mgr scheme auth (Request req) processResponse = do
             OutBodyStreaming strmbdy -> do
                 tbq <- newTBQueueIO 10 -- fixme: hard coding: 10
                 tbqNonMmpty <- newTVarIO False
-                let setup = addMyId mgr
-                let teardown _ = deleteMyId mgr
-                E.bracket setup teardown $ \_ -> void $ forkIO $ do
+                forkManaged mgr $ do
                     let push b = atomically $ do
                             writeTBQueue tbq (StreamingBuilder b)
                             writeTVar tbqNonMmpty True

Network/HTTP2/Server/Worker.hs

@@ -139,6 +139,9 @@ response wc@WorkerConf{..} mgr th tconf strm (Request req) (Response rsp) pps =
           flush  = atomically $ writeTBQueue tbq StreamingFlush
       strmbdy push flush
       atomically $ writeTBQueue tbq StreamingFinished
+      -- Remove the thread's ID from the manager's queue, to ensure the that the
+      -- manager will not terminate it before we are done. (The thread ID was
+      -- added implicitly when the worker was spawned by the manager).
       deleteMyId mgr
   where
     (_,reqvt) = inpObjHeaders req

test/HTTP2/ClientSpec.hs

@@ -63,7 +63,8 @@ runClient sc au hd = runTCPClient host port $ runHTTP2Client
     cliconf = ClientConfig sc au 20
     runHTTP2Client s = E.bracket (allocSimpleConfig s 4096)
                                  freeSimpleConfig
-                                 (\conf -> run cliconf conf client)
+                                 (\conf -> run cliconf conf $ \sendRequest ->
+                                   client sendRequest)
     client sendRequest = do
         let req = requestNoBody methodGet "/" hd
         sendRequest req $ \rsp -> do