From 2f7821eeaf9259e64890c8290882a5213a431866 Mon Sep 17 00:00:00 2001
From: Finley McIlwaine <finley@well-typed.com>
Date: Thu, 12 Sep 2024 16:34:50 -0700
Subject: [PATCH] Avoid `undefined` in client

---
 Network/HTTP2/Client/Run.hs  | 18 ++++++++++++++----
 Network/HTTP2/H2/Receiver.hs |  3 +++
 2 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/Network/HTTP2/Client/Run.hs b/Network/HTTP2/Client/Run.hs
index 0a8cb7d0..ce90e768 100644
--- a/Network/HTTP2/Client/Run.hs
+++ b/Network/HTTP2/Client/Run.hs
@@ -142,7 +142,7 @@ setup ClientConfig{..} conf@Config{..} = do
 
 runH2 :: Config -> Context -> IO a -> IO a
 runH2 conf ctx runClient = do
-    stopAfter mgr (clientResult <$> race runBackgroundThreads runClient) $ \res ->
+    stopAfter mgr runAll $ \res ->
         closeAllStreams (oddStreamTable ctx) (evenStreamTable ctx) res
   where
     mgr = threadManager ctx
@@ -152,9 +152,19 @@ runH2 conf ctx runClient = do
         labelMe "H2 runBackgroundThreads"
         concurrently_ runReceiver runSender
 
-    clientResult :: Either () a -> a
-    clientResult (Left ()) = undefined -- unreachable
-    clientResult (Right a) = a
+    -- Run the background threads and client concurrently. If the client
+    -- finishes first, cancel the background threads. If the background
+    -- threads finish first, wait for the client.
+    runAll = do
+      withAsync runBackgroundThreads $ \runningBackgroundThreads ->
+        withAsync runClient $ \runningClient -> do
+          result <- waitEither runningBackgroundThreads runningClient
+          case result of
+            Right clientResult -> do
+              cancel runningBackgroundThreads
+              return clientResult
+            Left () -> do
+              wait runningClient
 
 sendRequest
     :: Config
diff --git a/Network/HTTP2/H2/Receiver.hs b/Network/HTTP2/H2/Receiver.hs
index 927ecce7..fa4aff06 100644
--- a/Network/HTTP2/H2/Receiver.hs
+++ b/Network/HTTP2/H2/Receiver.hs
@@ -58,6 +58,9 @@ frameReceiver ctx@Context{..} conf@Config{..} = do
                 loop
 
     sendGoaway se
+        | Just GoAwayIsSent <- E.fromException se = do
+            waitCounter0 threadManager
+            enqueueControl controlQ $ CFinish GoAwayIsSent
         | Just ConnectionIsClosed <- E.fromException se = do
             waitCounter0 threadManager
             enqueueControl controlQ $ CFinish ConnectionIsClosed