Recover rxflow on closing

Network/HTTP2/Client/Run.hs

@@ -85,7 +85,9 @@ run cconf@ClientConfig{..} conf client = do
     clientCore ctx mgr req processResponse = do
         strm <- sendRequest ctx mgr scheme authority req
         rsp <- getResponse strm
-        processResponse rsp
+        x <- processResponse rsp
+        adjustRxWindow ctx strm
+        return x
     runClient ctx mgr = do
         x <- client (clientCore ctx mgr) $ aux ctx
         waitCounter0 mgr
@@ -205,14 +207,15 @@ sendStreaming Context{..} mgr req sid newstrm strmbdy = do
     forkManagedUnmask mgr $ \unmask -> do
         decrementedCounter <- newIORef False
         let decCounterOnce = do
-              alreadyDecremented <- atomicModifyIORef decrementedCounter $ \b -> (True, b)
-              unless alreadyDecremented $ decCounter mgr
-        let iface = OutBodyIface {
-                outBodyUnmask = unmask
-              , outBodyPush = \b -> atomically $ writeTBQueue tbq (StreamingBuilder b Nothing)
-              , outBodyPushFinal = \b -> atomically $ writeTBQueue tbq (StreamingBuilder b (Just decCounterOnce))
-              , outBodyFlush = atomically $ writeTBQueue tbq StreamingFlush
-              }
+                alreadyDecremented <- atomicModifyIORef decrementedCounter $ \b -> (True, b)
+                unless alreadyDecremented $ decCounter mgr
+        let iface =
+                OutBodyIface
+                    { outBodyUnmask = unmask
+                    , outBodyPush = \b -> atomically $ writeTBQueue tbq (StreamingBuilder b Nothing)
+                    , outBodyPushFinal = \b -> atomically $ writeTBQueue tbq (StreamingBuilder b (Just decCounterOnce))
+                    , outBodyFlush = atomically $ writeTBQueue tbq StreamingFlush
+                    }
             finished = atomically $ writeTBQueue tbq $ StreamingFinished decCounterOnce
         incCounter mgr
         strmbdy iface `finally` finished

Network/HTTP2/H2/Receiver.hs

@@ -190,11 +190,12 @@ processState (Open _ (NoBody tbl@(_, reqvt))) ctx@Context{..} strm@Stream{stream
     return False
 
 -- Transition (process2)
-processState (Open hcl (HasBody tbl@(_, reqvt))) ctx@Context{..} strm@Stream{streamInput} _streamId = do
+processState (Open hcl (HasBody tbl@(_, reqvt))) ctx@Context{..} strm@Stream{streamInput, streamRxQ} _streamId = do
     let mcl = fst <$> (getFieldValue tokenContentLength reqvt >>= C8.readInt)
     bodyLength <- newIORef 0
     tlr <- newIORef Nothing
     q <- newTQueueIO
+    writeIORef streamRxQ $ Just q
     setStreamState ctx strm $ Open hcl (Body q mcl bodyLength tlr)
     -- FLOW CONTROL: WINDOW_UPDATE 0: recv: announcing my limit properly
     -- FLOW CONTROL: WINDOW_UPDATE: recv: announcing my limit properly
@@ -612,17 +613,9 @@ stream x FrameHeader{streamId} _ _ _ _ =
 ----------------------------------------------------------------
 
 -- | Type for input streaming.
-data Source
-    = Source
-        (Int -> IO ())
-        (TQueue (Either E.SomeException (ByteString, Bool)))
-        (IORef ByteString)
-        (IORef Bool)
-
-mkSource
-    :: TQueue (Either E.SomeException (ByteString, Bool))
-    -> (Int -> IO ())
-    -> IO Source
+data Source = Source (Int -> IO ()) RxQ (IORef ByteString) (IORef Bool)
+
+mkSource :: RxQ -> (Int -> IO ()) -> IO Source
 mkSource q inform = Source inform q <$> newIORef "" <*> newIORef False
 
 readSource :: Source -> IO (ByteString, Bool)

Network/HTTP2/H2/Stream.hs

@@ -52,6 +52,7 @@ newOddStream sid txwin rxwin =
         <*> newEmptyMVar
         <*> newTVarIO (newTxFlow txwin)
         <*> newIORef (newRxFlow rxwin)
+        <*> newIORef Nothing
 
 newEvenStream :: StreamId -> WindowSize -> WindowSize -> IO Stream
 newEvenStream sid txwin rxwin =
@@ -60,6 +61,7 @@ newEvenStream sid txwin rxwin =
         <*> newEmptyMVar
         <*> newTVarIO (newTxFlow txwin)
         <*> newIORef (newRxFlow rxwin)
+        <*> newIORef Nothing
 
 ----------------------------------------------------------------
 

Network/HTTP2/H2/Types.hs

@@ -149,12 +149,15 @@ instance Show StreamState where
 
 ----------------------------------------------------------------
 
+type RxQ = TQueue (Either E.SomeException (ByteString, Bool))
+
 data Stream = Stream
     { streamNumber :: StreamId
     , streamState :: IORef StreamState
     , streamInput :: MVar (Either SomeException InpObj) -- Client only
     , streamTxFlow :: TVar TxFlow
     , streamRxFlow :: IORef RxFlow
+    , streamRxQ :: IORef (Maybe RxQ)
     }
 
 instance Show Stream where

Network/HTTP2/H2/Window.hs

@@ -1,8 +1,9 @@
 {-# LANGUAGE NamedFieldPuns #-}
-{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE BangPatterns #-}
 
 module Network.HTTP2.H2.Window where
 
+import qualified Data.ByteString as BS
 import Data.IORef
 import Network.Control
 import qualified UnliftIO.Exception as E
@@ -79,3 +80,24 @@ informWindowUpdate Context{controlQ, rxFlow} Stream{streamNumber, streamRxFlow}
         let frame = windowUpdateFrame streamNumber ws
             cframe = CFrames Nothing [frame]
         enqueueControl controlQ cframe
+
+adjustRxWindow :: Context -> Stream -> IO ()
+adjustRxWindow ctx stream@Stream{streamRxQ} = do
+    mq <- readIORef streamRxQ
+    case mq of
+        Nothing -> return ()
+        Just q -> do
+            len <- readQ q
+            informWindowUpdate ctx stream len
+  where
+    readQ q = atomically $ loop 0
+      where
+        loop !total = do
+            meb <- tryReadTQueue q
+            case meb of
+                Just (Right (bs, _)) -> loop (total + BS.length bs)
+                Just le@(Left _) -> do
+                    -- reserving HTTP2Error
+                    writeTQueue q le
+                    return total
+                _ -> return total

Network/HTTP2/Server/Run.hs

@@ -49,7 +49,7 @@ run sconf@ServerConfig{numberOfWorkers} conf server = do
     when ok $ do
         (ctx, mgr) <- setup sconf conf
         let wc = fromContext ctx
-        setAction mgr $ worker wc mgr server
+        setAction mgr $ worker ctx wc mgr server
         replicateM_ numberOfWorkers $ spawnAction mgr
         runH2 conf ctx mgr
 

Network/HTTP2/Server/Worker.hs

@@ -1,4 +1,3 @@
-{-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE PatternGuards #-}
 {-# LANGUAGE RecordWildCards #-}
@@ -15,7 +14,6 @@ import Network.HTTP.Semantics.IO
 import Network.HTTP.Semantics.Server
 import Network.HTTP.Semantics.Server.Internal
 import Network.HTTP.Types
-import Network.Socket (SockAddr)
 import qualified System.TimeManager as T
 import UnliftIO.Exception (SomeException (..))
 import qualified UnliftIO.Exception as E
@@ -33,8 +31,6 @@ data WorkerConf a = WorkerConf
     , workerCleanup :: a -> IO ()
     , isPushable :: IO Bool
     , makePushStream :: a -> PushPromise -> IO (StreamId, a)
-    , mySockAddr :: SockAddr
-    , peerSockAddr :: SockAddr
     }
 
 fromContext :: Context -> WorkerConf Stream
@@ -54,8 +50,6 @@ fromContext ctx@Context{..} =
             (_, newstrm) <- openEvenStreamWait ctx
             let pid = streamNumber pstrm
             return (pid, newstrm)
-        , mySockAddr = mySockAddr
-        , peerSockAddr = peerSockAddr
         }
 
 ----------------------------------------------------------------
@@ -162,8 +156,8 @@ response wc@WorkerConf{..} mgr th tconf strm (Request req) (Response rsp) pps =
     (_, reqvt) = inpObjHeaders req
 
 -- | Worker for server applications.
-worker :: WorkerConf a -> Manager -> Server -> Action
-worker wc@WorkerConf{..} mgr server = do
+worker :: Context -> WorkerConf Stream -> Manager -> Server -> Action
+worker ctx@Context{..} wc@WorkerConf{..} mgr server = do
     sinfo <- newStreamInfo
     tcont <- newThreadContinue
     timeoutKillThread mgr $ go sinfo tcont
@@ -178,7 +172,9 @@ worker wc@WorkerConf{..} mgr server = do
             T.resume th
             T.tickle th
             let aux = Aux th mySockAddr peerSockAddr
-            server (Request req') aux $ response wc mgr th tcont strm (Request req')
+            r <- server (Request req') aux $ response wc mgr th tcont strm (Request req')
+            adjustRxWindow ctx strm
+            return r
         cont1 <- case ex of
             Right () -> return True
             Left e@(SomeException _)