fix connection closing

Now when the user quits, the developer also exits.

1 files changed, 22 insertions, 15 deletions

diff --git a/WebSockets.hs b/WebSockets.hs
index f3712a9..d8d43e7 100644
--- a/WebSockets.hs
+++ b/WebSockets.hs
@@ -30,6 +30,7 @@ import qualified Data.Text as T
 import qualified Data.ByteString.Lazy as L
 import Data.List
 import Data.Monoid
+import Control.Monad
 
 -- | Enable compression.
 connectionOptions :: ConnectionOptions
@@ -76,17 +77,25 @@ clientApp mode mksent filterreceived a conn = do
 	setup = do
 		schan <- newTMChanIO
 		rchan <- newTMChanIO
-		sthread <- async $ relayFromSocket conn $ \v ->
-			case filterreceived v of
-				Nothing -> return ()
-				Just r -> atomically $ writeTMChan rchan r
-		rthread <- async $ relayToSocket conn $
-			fmap mksent <$> atomically (readTMChan schan)
+		sthread <- async $ relayToSocket conn mksent $
+			atomically (readTMChan schan)
+		rthread <- async $ do
+			relayFromSocket conn $ \v -> do
+				case filterreceived v of
+					Nothing -> return ()
+					Just r -> atomically $ writeTMChan rchan r
+			-- Server sent Done, so close channels.
+			atomically $ do
+				closeTMChan schan
+				closeTMChan rchan
 		return (schan, rchan, sthread, rthread)
-	cleanup (_, _, sthread, rthread) = do
+	cleanup (schan, _, sthread, rthread) = do
 		sendBinaryData conn Done
-		() <- wait sthread
-		wait rthread
+		atomically $ closeTMChan schan
+		-- Wait for any more data from the server.
+		-- These often die with a ConnectionClosed.
+		void $ waitCatch sthread
+		void $ waitCatch rthread
 	go sid (schan, rchan, _, _) = a schan rchan sid
 
 relayFromSocket :: Connection -> (LogMessage -> IO ()) -> IO ()
@@ -98,21 +107,19 @@ relayFromSocket conn sender = go
 			LogMessage msg -> do
 				sender msg
 				go
-			Done -> do
-				print "GOT DONE"
-				return ()
+			Done -> return ()
 			WireProtocolError e -> protocolError conn e
 			_ -> protocolError conn "Protocol error"
 
-relayToSocket :: Connection -> (IO (Maybe LogMessage)) -> IO ()
-relayToSocket conn getter = go
+relayToSocket :: Connection -> (received -> LogMessage) -> IO (Maybe received) -> IO ()
+relayToSocket conn mksent getter = go
   where
 	go = do
 		mmsg <- getter
 		case mmsg of
 			Nothing -> return ()
 			Just msg -> do
-				sendBinaryData conn (LogMessage msg)
+				sendBinaryData conn $ LogMessage $ mksent msg
 				go
 
 -- | Framing protocol used over a websocket connection.