Merge pull request #167 from edsko/edsko/receiver-done

Ensure sender notices when receiver has terminated

Author: kazu-yamamoto

Network/HTTP2/Client/Run.hs

@@ -151,8 +151,15 @@ runH2 conf ctx runClient = do
         er <- race runReceiver runClient
         case er of
             Right r -> return r
-            -- never reached because runReceiver throws an exception to exit.
-            Left () -> throwIO ConnectionIsClosed
+            Left err -> throwIO err
+
+    -- When 'runClientReceiver' terminates, it is important we give the sender
+    -- a chance to terminate cleanly also (it's possible the client terminated
+    -- but there are still some messages in the queue to be sent).
+    --
+    -- If the client terminated successfully, we ignore any other errors in the
+    -- sender (indeed, any exception here might simply be that the background
+    -- threads were cancelled /because/ the client terminated).
     runAll = snd <$> concurrently runSender runClientReceiver
 
 makeStream

Network/HTTP2/H2/Context.hs

@@ -90,7 +90,7 @@ data Context = Context
     , mySockAddr         :: SockAddr
     , peerSockAddr       :: SockAddr
     , threadManager      :: T.ThreadManager
-    , receiverDone       :: TVar Bool
+    , receiverDone       :: TVar (Maybe SomeException)
     , workersDone        :: STM Bool
     }
 {- FOURMOLU_ENABLE -}
@@ -138,7 +138,7 @@ newContext roleInfo Config{..} cacheSiz connRxWS mySettings timmgr mdone = do
     let mySockAddr   = confMySockAddr
     let peerSockAddr = confPeerSockAddr
     threadManager   <- T.newThreadManager timmgr
-    receiverDone    <- newTVarIO False
+    receiverDone    <- newTVarIO Nothing
     let workersDone = fromMaybe (T.isAllGone threadManager) mdone
     return Context{..}
   where

Network/HTTP2/H2/Receiver.hs

@@ -19,6 +19,7 @@ import qualified Data.ByteString.Char8 as C8
 import qualified Data.ByteString.Short as Short
 import qualified Data.ByteString.UTF8 as UTF8
 import Data.IORef
+import Data.Void
 import Network.Control
 import Network.HTTP.Semantics
 import qualified System.ThreadManager as T
@@ -45,28 +46,35 @@ headerFragmentLimit = 51200 -- 50K
 
 ----------------------------------------------------------------
 
-frameReceiver :: Context -> Config -> IO ()
-frameReceiver ctx conf@Config{..} =
-    (switch `E.catch` handler)
-        `E.finally` atomically
-            (writeTVar (receiverDone ctx) True)
+frameReceiver :: Context -> Config -> IO E.SomeException
+frameReceiver ctx@Context{receiverDone} conf@Config{..} =
+    E.mask $ \unmask -> do
+        mErr <- E.try $ unmask switch
+        case mErr of
+            Left err -> do
+                atomically $ writeTVar receiverDone $ Just err
+                return err
+            Right x -> do
+                absurd x -- We only terminate due to exceptions
   where
-    handler ConnectionIsClosed = return ()
-    handler e = E.throwIO e
+    switch :: IO Void
     switch = do
         labelMe "H2 receiver"
         tid <- myThreadId
         if confReadNTimeout
             then
                 loop1
             else
-                void $
-                    T.withHandle (threadManager ctx) (E.throwTo tid ConnectionIsTimeout) loop2
+                T.withHandle (threadManager ctx) (E.throwTo tid ConnectionIsTimeout) loop2
+
+    loop1 :: IO Void
     loop1 = do
         hd <- confReadN frameHeaderLength -- throwing an exception on timeout
         when (BS.null hd) $ E.throwIO ConnectionIsClosed
         processFrame ctx conf $ decodeFrameHeader hd
         loop1
+
+    loop2 :: T.Handle -> IO Void
     loop2 th = do
         -- If 'confReadN' is timeouted, 'ConnectionIsTimeout' is thrown
         -- to destroy the thread trees.

Network/HTTP2/H2/Sender.hs

@@ -16,7 +16,6 @@ import Foreign.Ptr (minusPtr, plusPtr)
 import Network.ByteOrder
 import Network.HTTP.Semantics.Client
 import Network.HTTP.Semantics.IO
-import System.ThreadManager
 
 import Imports
 import Network.HPACK (setLimitForEncoding, toTokenHeaderTable)
@@ -59,38 +58,42 @@ updatePeerSettings Context{peerSettings, oddStreamTable, evenStreamTable} peerAl
     updateAllStreamTxFlow siz strms =
         forM_ strms $ \strm -> increaseStreamWindowSize strm siz
 
-checkDone :: Context -> Int -> IO Bool
+checkDone :: Context -> Int -> IO (Maybe E.SomeException)
 checkDone Context{..} 0 = atomically $ do
     isEmptyC <- isEmptyTQueue controlQ
     isEmptyO <- isEmptyTQueue outputQ
     if not isEmptyC || not isEmptyO
         then
-            return False
+            return Nothing
         else do
-            gone <- isAllGone threadManager
-            unless gone retry
-            done <- readTVar receiverDone
-            unless done retry
-            return True
-checkDone _ _ = return False
+            recv <- readTVar receiverDone
+            case recv of
+                Just done ->
+                    return $ Just done
+                _otherwise ->
+                    retry
+checkDone _ _ = return Nothing
 
-frameSender :: Context -> Config -> IO ()
+frameSender :: Context -> Config -> IO E.SomeException
 frameSender
     ctx@Context{outputQ, controlQ, encodeDynamicTable, outputBufferLimit}
     Config{..} = do
         labelMe "H2 sender"
-        loop 0
+        loop 0 `E.catch` return
       where
         ----------------------------------------------------------------
-        loop :: Offset -> IO ()
+        loop :: Offset -> IO E.SomeException
         loop off = do
-            done <- checkDone ctx off
-            unless done $ do
-                x <- atomically $ dequeue off
-                case x of
-                    C ctl -> flushN off >> control ctl >> loop 0
-                    O out -> outputAndSync out off >>= flushIfNecessary >>= loop
-                    Flush -> flushN off >> loop 0
+            mDone <- checkDone ctx off
+            case mDone of
+                Just done ->
+                    return done
+                Nothing -> do
+                    x <- atomically $ dequeue off
+                    case x of
+                        C ctl -> flushN off >> control ctl >> loop 0
+                        O out -> outputAndSync out off >>= flushIfNecessary >>= loop
+                        Flush -> flushN off >> loop 0
 
         -- Flush the connection buffer to the socket, where the first 'n' bytes of
         -- the buffer are filled.

Network/HTTP2/H2/Stream.hs

@@ -77,12 +77,15 @@ readStreamState Stream{streamState} = readIORef streamState
 
 closeAllStreams
     :: TVar OddStreamTable -> TVar EvenStreamTable -> Maybe SomeException -> IO ()
-closeAllStreams ovar evar mErr' = do
+closeAllStreams ovar evar mErr = do
     ostrms <- clearOddStreamTable ovar
     mapM_ finalize ostrms
     estrms <- clearEvenStreamTable evar
     mapM_ finalize estrms
   where
+    -- We treat /every/ exception, including 'ConectionIsClosed', as abnormal
+    -- termination: we should only report a clean termination when we receive an
+    -- explicit @END_STREAM@ frame.
     finalize strm = do
         st <- readStreamState strm
         void $ tryPutMVar (streamInput strm) err
@@ -92,14 +95,6 @@ closeAllStreams ovar evar mErr' = do
             _otherwise ->
                 return ()
 
-    mErr :: Maybe SomeException
-    mErr = case mErr' of
-        Just e
-            | Just ConnectionIsClosed <- fromException e ->
-                Nothing
-        _otherwise ->
-            mErr'
-
     err :: Either SomeException a
     err = Left $ fromMaybe (toException ConnectionIsClosed) mErr
 

Network/HTTP2/Server/Run.hs

@@ -3,9 +3,8 @@
 
 module Network.HTTP2.Server.Run where
 
-import Control.Concurrent.Async (concurrently_)
+import Control.Concurrent.Async
 import Control.Concurrent.STM
-import qualified Control.Exception as E
 import Imports
 import Network.Control (defaultMaxData)
 import Network.HTTP.Semantics.IO
@@ -128,10 +127,8 @@ runH2 conf ctx = do
         runReceiver = frameReceiver ctx conf
         runSender = frameSender ctx conf
         runBackgroundThreads = do
-            er <- E.try $ concurrently_ runReceiver runSender
-            case er of
-                Right () -> return ()
-                Left e -> closureServer conf ctx e
+            e <- snd <$> concurrently runReceiver runSender
+            closureServer conf ctx e
     T.stopAfter mgr runBackgroundThreads $ \res ->
         closeAllStreams (oddStreamTable ctx) (evenStreamTable ctx) res