introducing forkManaged

Network/QUIC/Client/Reader.hs

@@ -15,6 +15,7 @@ import Data.List (intersect)
 import Network.Socket (Socket, close, getSocketName)
 import qualified Network.Socket.ByteString as NSB
 
+import Network.QUIC.Common
 import Network.QUIC.Connection
 import Network.QUIC.Connector
 import Network.QUIC.Crypto
@@ -30,6 +31,7 @@ import Network.QUIC.Types
 -- | readerClient dies when the socket is closed.
 readerClient :: Socket -> Connection -> IO ()
 readerClient s0 conn = handleLogUnit logAction $ do
+    labelMe "readerClient"
     wait
     loop
   where
@@ -146,5 +148,5 @@ rebind conn microseconds = do
     newSock <- natRebinding peersa
     oldSock <- setSocket conn newSock
     let reader = readerClient newSock conn
-    forkIO reader >>= addReader conn
+    forkManaged conn reader
     fire conn microseconds $ close oldSock

Network/QUIC/Client/Run.hs

@@ -7,7 +7,6 @@ module Network.QUIC.Client.Run (
     migrate,
 ) where
 
-import Control.Concurrent
 import Control.Concurrent.Async
 import qualified Control.Exception as E
 import qualified Network.Socket as NS
@@ -58,7 +57,7 @@ run conf client = NS.withSocketsDo $ do
 runClient :: ClientConfig -> (Connection -> IO a) -> Bool -> VersionInfo -> IO a
 runClient conf client0 isICVN verInfo = do
     E.bracket open clse $ \(ConnRes conn myAuthCIDs reader) -> do
-        forkIO reader >>= addReader conn
+        forkManaged conn reader
         let conf' =
                 conf
                     { ccParameters =

Network/QUIC/Connection/Misc.hs

@@ -21,7 +21,7 @@ module Network.QUIC.Connection.Misc (
     delayedAck,
     resetDealyedAck,
     setMaxPacketSize,
-    addReader,
+    forkManaged,
     killReaders,
     addResource,
     freeResources,
@@ -35,6 +35,8 @@ module Network.QUIC.Connection.Misc (
 
 import Control.Concurrent
 import qualified Control.Exception as E
+import qualified Data.IntMap as IntMap
+import qualified Data.Map.Strict as Map
 import Network.Socket (Socket)
 import System.Mem.Weak
 
@@ -161,12 +163,33 @@ freeResources Connection{..} =
 addReader :: Connection -> ThreadId -> IO ()
 addReader Connection{..} tid = do
     wtid <- mkWeakThreadId tid
-    atomicModifyIORef'' readers $ \m -> do
-        m
-        deRefWeak wtid >>= mapM_ killThread
+    let n = fromThreadId tid
+    atomicModifyIORef'' readers $ Map.insert n wtid
+
+delReader :: Connection -> ThreadId -> IO ()
+delReader Connection{..} tid = do
+    wtid <- mkWeakThreadId tid
+    let n = fromThreadId tid
+    atomicModifyIORef'' readers $ Map.delete n
+
+forkManaged :: Connection -> IO () -> IO ()
+forkManaged conn action = void $ forkIO $ do
+    E.bracket setup clean $ \_ -> action
+  where
+    setup = do
+        tid <- myThreadId
+        addReader conn tid
+        return tid
+    clean = delReader conn
 
 killReaders :: Connection -> IO ()
-killReaders Connection{..} = join $ readIORef readers
+killReaders Connection{..} = do
+    wtids <- readIORef readers
+    forM_ wtids $ \wtid -> do
+        mtid <- deRefWeak wtid
+        case mtid of
+            Nothing -> return ()
+            Just tid -> killThread tid
 
 ----------------------------------------------------------------
 

Network/QUIC/Connection/Types.hs

@@ -20,6 +20,7 @@ import Foreign.Ptr (nullPtr)
 import Network.Control (Rate, RxFlow, TxFlow, newRate, newRxFlow, newTxFlow)
 import Network.Socket (Cmsg, SockAddr, Socket)
 import Network.TLS.QUIC
+import System.Mem.Weak (Weak)
 
 import Network.QUIC.Config
 import Network.QUIC.Connector
@@ -202,7 +203,7 @@ data Connection = Connection
     -- Manage
     , connRecvQ :: RecvQ
     , connSocket :: IORef Socket
-    , readers :: IORef (IO ())
+    , readers :: IORef (Map Word64 (Weak ThreadId))
     , mainThreadId :: ThreadId
     , controlRate :: Rate
     , -- Info
@@ -306,7 +307,7 @@ newConnection rl myparams verInfo myAuthCIDs peerAuthCIDs debugLog qLog hooks sr
     ecrptBuf <- mallocBytes bufsiz
     dcrptBuf <- mallocBytes bufsiz
     Connection connstate debugLog qLog hooks send recv recvQ sref
-        <$> newIORef (return ())
+        <$> newIORef Map.empty
         <*> myThreadId
         <*> newRate
         -- Info

Network/QUIC/Imports.hs

@@ -25,6 +25,7 @@ module Network.QUIC.Imports (
 #endif
     atomicModifyIORef'',
     copyBS,
+    fromThreadId,
 ) where
 
 import Control.Applicative
@@ -49,6 +50,10 @@ import Network.ByteOrder
 import Network.QUIC.Utils
 import Numeric
 
+#if __GLASGOW_HASKELL__ >= 908
+import GHC.Conc.Sync (fromThreadId)
+#endif
+
 -- | All internal byte sequences.
 --   `ByteString` should be used for FFI related stuff.
 type Bytes = ShortByteString
@@ -71,3 +76,8 @@ copyBS dst (PS fptr off len) = withForeignPtr fptr $ \src0 -> do
     let src = src0 `plusPtr` off
     copyBytes dst src len
     return len
+
+#if __GLASGOW_HASKELL__ < 908
+fromThreadId :: ThreadId -> Word64
+fromThreadId tid = read (drop 9 $ show tid)
+#endif