From e0bf6943c9a772bd2fad0f7db540a5a22ab1e0fa Mon Sep 17 00:00:00 2001
From: Benjamin Summers <benjamin@tlon.io>
Date: Thu, 19 Dec 2019 05:16:52 -0800
Subject: [PATCH] Make sure that terminal connections are closed properly on
 disconnects.

---
 pkg/king/lib/King/API.hs      | 11 +++++++++--
 pkg/king/lib/Vere/NounServ.hs | 24 ++++++++++++++----------
 2 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/pkg/king/lib/King/API.hs b/pkg/king/lib/King/API.hs
index cc1acdcff..ce7b26535 100644
--- a/pkg/king/lib/King/API.hs
+++ b/pkg/king/lib/King/API.hs
@@ -113,8 +113,15 @@ serveTerminal env api word =
                 inp <- io $ newTBMChanIO 5
                 out <- io $ newTBMChanIO 5
                 atomically $ sp $ NounServ.mkConn inp out
-                runRIO env $
-                    NounServ.wsConn "NOUNSERV (wsServ) " inp out wsc
+                let doit = runRIO env
+                         $ NounServ.wsConn "NOUNSERV (wsServ) " inp out wsc
+
+                -- If `wai` kills this thread for any reason, the TBMChans
+                -- *need* to be closed. If they are not closed, the
+                -- terminal will not know that they disconnected.
+                finally doit $ atomically $ do
+                    closeTBMChan inp
+                    closeTBMChan out
 
 data BadShip = BadShip Text
   deriving (Show, Exception)
diff --git a/pkg/king/lib/Vere/NounServ.hs b/pkg/king/lib/Vere/NounServ.hs
index 031fff441..728bbc12b 100644
--- a/pkg/king/lib/Vere/NounServ.hs
+++ b/pkg/king/lib/Vere/NounServ.hs
@@ -40,17 +40,20 @@ data Server i o a = Server
 
 --------------------------------------------------------------------------------
 
+withRIOThread ∷ RIO e a → RIO e (Async a)
+withRIOThread act = do
+    env <- ask
+    io $ async $ runRIO env $ act
+
 wsConn :: (FromNoun i, ToNoun o, Show i, Show o, HasLogFunc e)
        => Utf8Builder
        -> TBMChan i -> TBMChan o
        -> WS.Connection
        -> RIO e ()
 wsConn pre inp out wsc = do
-    env <- ask
-
     logWarn (pre <> "(wcConn) Connected!")
 
-    writer <- io $ async $ runRIO env $ forever $ do
+    writer <- withRIOThread $ forever $ do
         logWarn (pre <> "(wsConn) Waiting for data.")
         byt <- io $ toStrict <$> WS.receiveData wsc
         logWarn (pre <> "Got data")
@@ -58,7 +61,7 @@ wsConn pre inp out wsc = do
         logWarn (pre <> "(wsConn) Decoded data, writing to chan")
         atomically $ writeTBMChan inp dat
 
-    reader <- io $ async $ runRIO env $ forever $ do
+    reader <- withRIOThread $ forever $ do
         logWarn (pre <> "Waiting for data from chan")
         atomically (readTBMChan out) >>= \case
             Nothing  -> do
@@ -68,14 +71,15 @@ wsConn pre inp out wsc = do
                 logWarn (pre <> "(wsConn) Got message! " <> displayShow msg)
                 io $ WS.sendBinaryData wsc $ fromStrict $ jamBS $ toNoun msg
 
-    res <- atomically (waitCatchSTM writer <|> waitCatchSTM reader)
+    let cleanup = do
+            atomically (closeTBMChan inp >> closeTBMChan out)
+            cancel writer
+            cancel reader
 
-    logWarn $ displayShow (res :: Either SomeException ())
+    flip finally cleanup $ do
+         res <- atomically (waitCatchSTM writer <|> waitCatchSTM reader)
+         logWarn $ displayShow (res :: Either SomeException ())
 
-    atomically (closeTBMChan inp >> closeTBMChan out)
-
-    cancel writer
-    cancel reader
 
 --------------------------------------------------------------------------------