synching client and sender with STM

kazu-yamamoto · Jul 8, 2024 · 80de8db · 80de8db
1 parent c768049
commit 80de8db
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 23 deletions.
diff --git a/Network/HTTP2/Client/Run.hs b/Network/HTTP2/Client/Run.hs
@@ -88,14 +88,19 @@ run cconf@ClientConfig{..} conf client = do
         x <- processResponse rsp
         adjustRxWindow ctx strm
         return x
-    runClient ctx = do
-        x <- client (clientCore ctx) $ aux ctx
-        waitCounter0 $ threadManager ctx
-        let frame = goawayFrame 0 NoError "graceful closing"
-        mvar <- newMVar ()
-        enqueueControl (controlQ ctx) $ CGoaway frame mvar
-        takeMVar mvar
-        return x
+    runClient ctx = wrapClinet ctx $ client (clientCore ctx) $ aux ctx
+
+wrapClinet :: Context -> IO a -> IO a
+wrapClinet ctx client = do
+    x <- client
+    waitCounter0 $ threadManager ctx
+    let frame = goawayFrame 0 NoError "graceful closing"
+    enqueueControl (controlQ ctx) $ CFrames Nothing [frame]
+    enqueueControl (controlQ ctx) $ CFinish GoAwayIsSent
+    atomically $ do
+        done <- readTVar $ senderDone ctx
+        check done
+    return x
 
 -- | Launching a receiver and a sender.
 runIO :: ClientConfig -> Config -> (ClientIO -> IO (IO a)) -> IO a
@@ -107,8 +112,9 @@ runIO cconf@ClientConfig{..} conf@Config{..} action = do
             return (streamNumber strm, strm)
         get = getResponse
         create = openOddStreamWait ctx
-    runClient <-
-        action $ ClientIO confMySockAddr confPeerSockAddr putR get putB create
+    runClient <- do
+        act <- action $ ClientIO confMySockAddr confPeerSockAddr putR get putB create
+        return $ wrapClinet ctx act
     runH2 conf ctx runClient
 
 getResponse :: Stream -> IO Response

diff --git a/Network/HTTP2/H2/Context.hs b/Network/HTTP2/H2/Context.hs
@@ -87,6 +87,7 @@ data Context = Context
     , mySockAddr :: SockAddr
     , peerSockAddr :: SockAddr
     , threadManager :: Manager
+    , senderDone :: TVar Bool
     }
 
 ----------------------------------------------------------------
@@ -128,6 +129,7 @@ newContext rinfo Config{..} cacheSiz connRxWS settings timmgr =
         <*> return confMySockAddr
         <*> return confPeerSockAddr
         <*> start timmgr
+        <*> newTVarIO False
   where
     rl = case rinfo of
         RIC{} -> Client

diff --git a/Network/HTTP2/H2/Receiver.hs b/Network/HTTP2/H2/Receiver.hs
@@ -60,15 +60,15 @@ frameReceiver ctx@Context{..} conf@Config{..} = do
         -- to destroy the thread trees.
         hd <- confReadN frameHeaderLength
         if BS.null hd
-            then enqueueControl controlQ CFinish
+            then enqueueControl controlQ $ CFinish ConnectionIsTimeout
             else do
                 processFrame ctx conf $ decodeFrameHeader hd
                 loop
 
     sendGoaway se
-        | Just e@ConnectionIsClosed <- E.fromException se = do
+        | Just ConnectionIsClosed <- E.fromException se = do
             waitCounter0 threadManager
-            enqueueControl controlQ $ CFinish e
+            enqueueControl controlQ $ CFinish ConnectionIsClosed
         | Just e@(ConnectionErrorIsReceived _ _ _) <- E.fromException se =
             enqueueControl controlQ $ CFinish e
         | Just e@(ConnectionErrorIsSent err sid msg) <- E.fromException se = do

diff --git a/Network/HTTP2/H2/Sender.hs b/Network/HTTP2/H2/Sender.hs
@@ -7,7 +7,6 @@ module Network.HTTP2.H2.Sender (
     frameSender,
 ) where
 
-import Control.Concurrent.MVar (putMVar)
 import Data.IORef (modifyIORef', readIORef, writeIORef)
 import Data.IntMap.Strict (IntMap)
 import Foreign.Ptr (minusPtr, plusPtr)
@@ -39,6 +38,8 @@ data Switch
 
 wrapException :: E.SomeException -> IO ()
 wrapException se
+    | Just GoAwayIsSent <- E.fromException se = return ()
+    | Just ConnectionIsClosed <- E.fromException se = return ()
     | Just (e :: HTTP2Error) <- E.fromException se = E.throwIO e
     | otherwise = E.throwIO $ BadThingHappen se
 
@@ -65,10 +66,10 @@ updatePeerSettings Context{peerSettings, oddStreamTable, evenStreamTable} peerAl
 
 frameSender :: Context -> Config -> IO ()
 frameSender
-    ctx@Context{outputQ, controlQ, encodeDynamicTable, outputBufferLimit}
+    ctx@Context{outputQ, controlQ, encodeDynamicTable, outputBufferLimit, senderDone}
     Config{..} = do
         labelMe "H2 sender"
-        loop 0 `E.catch` wrapException
+        (loop 0 `E.finally` setSenderDone) `E.catch` wrapException
       where
         ----------------------------------------------------------------
         loop :: Offset -> IO ()
@@ -114,12 +115,6 @@ frameSender
         -- called with off == 0
         control :: Control -> IO ()
         control (CFinish e) = E.throwIO e
-        control (CGoaway bs mvar) = do
-            buf <- copyAll [bs] confWriteBuffer
-            let off = buf `minusPtr` confWriteBuffer
-            flushN off
-            putMVar mvar ()
-            E.throwIO GoAwayIsSent
         control (CFrames ms xs) = do
             buf <- copyAll xs confWriteBuffer
             let off = buf `minusPtr` confWriteBuffer
@@ -373,3 +368,5 @@ frameSender
                     , flags = flag
                     , streamId = sid
                     }
+
+        setSenderDone = atomically $ writeTVar senderDone True
diff --git a/Network/HTTP2/H2/Types.hs b/Network/HTTP2/H2/Types.hs
@@ -188,7 +188,6 @@ data Sync = Done | Cont (IO ()) OutputType
 data Control
     = CFinish HTTP2Error
     | CFrames (Maybe SettingsList) [ByteString]
-    | CGoaway ByteString (MVar ())
 
 ----------------------------------------------------------------