switch to TMChans so they can be closed when a connection is Done

1 files changed, 12 insertions, 9 deletions

diff --git a/WebSockets.hs b/WebSockets.hs
index ea6e251..f3712a9 100644
--- a/WebSockets.hs
+++ b/WebSockets.hs
@@ -19,6 +19,7 @@ import SessionID
 
 import Network.WebSockets hiding (Message)
 import Control.Concurrent.STM
+import Control.Concurrent.STM.TMChan
 import Control.Concurrent.Async
 import Control.Exception
 import GHC.Generics (Generic)
@@ -59,7 +60,7 @@ clientApp
 	:: Mode
 	-> (sent -> LogMessage)
 	-> (LogMessage -> Maybe received)
-	-> (TChan sent -> TChan received -> SessionID -> IO a)
+	-> (TMChan sent -> TMChan received -> SessionID -> IO a)
 	-> ClientApp a
 clientApp mode mksent filterreceived a conn = do
 	-- Ping every 30 seconds to avoid timeouts caused by proxies etc.
@@ -73,19 +74,19 @@ clientApp mode mksent filterreceived a conn = do
 		_ -> protocolError conn "Did not get expected Ready message from server"
   where
 	setup = do
-		schan <- newTChanIO
-		rchan <- newTChanIO
+		schan <- newTMChanIO
+		rchan <- newTMChanIO
 		sthread <- async $ relayFromSocket conn $ \v ->
 			case filterreceived v of
 				Nothing -> return ()
-				Just r -> atomically $ writeTChan rchan r
+				Just r -> atomically $ writeTMChan rchan r
 		rthread <- async $ relayToSocket conn $
-			Just . mksent <$> atomically (readTChan schan)
+			fmap mksent <$> atomically (readTMChan schan)
 		return (schan, rchan, sthread, rthread)
 	cleanup (_, _, sthread, rthread) = do
 		sendBinaryData conn Done
-		cancel sthread
-		cancel rthread
+		() <- wait sthread
+		wait rthread
 	go sid (schan, rchan, _, _) = a schan rchan sid
 
 relayFromSocket :: Connection -> (LogMessage -> IO ()) -> IO ()
@@ -97,7 +98,9 @@ relayFromSocket conn sender = go
 			LogMessage msg -> do
 				sender msg
 				go
-			Done -> return ()
+			Done -> do
+				print "GOT DONE"
+				return ()
 			WireProtocolError e -> protocolError conn e
 			_ -> protocolError conn "Protocol error"
 
@@ -107,7 +110,7 @@ relayToSocket conn getter = go
 	go = do
 		mmsg <- getter
 		case mmsg of
-			Nothing -> go
+			Nothing -> return ()
 			Just msg -> do
 				sendBinaryData conn (LogMessage msg)
 				go