call decCounter when StreamingFinished is received

kazu-yamamoto · Sep 14, 2023 · 287e8f6 · 287e8f6
1 parent 94948b3
commit 287e8f6
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 7 deletions.
diff --git a/Network/HTTP2/Arch/Sender.hs b/Network/HTTP2/Arch/Sender.hs
@@ -438,7 +438,9 @@ runStreamBuilder buf0 room0 takeQ = loop buf0 room0 0
                     B.More  _ writer  -> return (True,  total', False, LOne writer)
                     B.Chunk bs writer -> return (True,  total', False, LTwo bs writer)
             Just StreamingFlush       -> return (True,  total,  True,  LZero)
-            Just StreamingFinished    -> return (False, total,  True,  LZero)
+            Just (StreamingFinished dec) -> do
+                dec
+                return (False, total,  True,  LZero)
 
 fillBufStream :: Leftover -> IO (Maybe StreamingChunk) -> DynaNext
 fillBufStream leftover0 takeQ buf0 siz0 lim0 = do

diff --git a/Network/HTTP2/Arch/Types.hs b/Network/HTTP2/Arch/Types.hs
@@ -284,7 +284,7 @@ data Control = CFinish HTTP2Error
 
 ----------------------------------------------------------------
 
-data StreamingChunk = StreamingFinished
+data StreamingChunk = StreamingFinished (IO ())
                     | StreamingFlush
                     | StreamingBuilder Builder
 

diff --git a/Network/HTTP2/Client/Run.hs b/Network/HTTP2/Client/Run.hs
@@ -101,9 +101,7 @@ sendRequest ctx@Context{..} mgr scheme auth (Request req) processResponse = do
                     writeTBQueue tbq (StreamingBuilder b)
                     writeTVar tbqNonEmpty True
                 flush  = atomically $ writeTBQueue tbq StreamingFlush
-                finished = do
-                    atomically $ writeTBQueue tbq StreamingFinished
-                    decCounter mgr
+                finished = atomically $ writeTBQueue tbq $ StreamingFinished (decCounter mgr)
             incCounter mgr
             strmbdy unmask push flush `finally` finished
         atomically $ do

diff --git a/Network/HTTP2/Server/Worker.hs b/Network/HTTP2/Server/Worker.hs
@@ -141,8 +141,9 @@ response wc@WorkerConf{..} mgr th tconf strm (Request req) (Response rsp) pps =
             atomically $ writeTBQueue tbq (StreamingBuilder b)
             T.resume th
           flush  = atomically $ writeTBQueue tbq StreamingFlush
-      strmbdy push flush
-      atomically $ writeTBQueue tbq StreamingFinished
+          finished = atomically $ writeTBQueue tbq $ StreamingFinished (decCounter mgr)
+      incCounter mgr
+      strmbdy push flush `E.finally` finished
       -- 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).