Merge pull request #1740 from stevemao/cleanup

update the examples to show how to cleanup resources once user discon…

yesod-websockets/chat-with-multiple-channels.hs

@@ -1,4 +1,4 @@
-{-# LANGUAGE QuasiQuotes, TemplateHaskell, TypeFamilies, OverloadedStrings #-}
+{-# LANGUAGE QuasiQuotes, TemplateHaskell, TypeFamilies, OverloadedStrings, ScopedTypeVariables #-}
 import Yesod.Core
 import Yesod.WebSockets
 import qualified Data.Text.Lazy as TL
@@ -11,8 +11,9 @@ import Data.Monoid ((<>))
 import Control.Concurrent.STM.Lifted
 import Data.Text (Text)
 import qualified Data.Map as M
+import UnliftIO.Exception (try, SomeException)
 
-data App = App (TVar (M.Map Text (TChan Text)))
+data App = App (TVar (M.Map Text (TChan Text, Int)))
 
 instance Yesod App
 
@@ -20,6 +21,15 @@ mkYesod "App" [parseRoutes|
 / HomeR GET
 |]
 
+cleanupChannel :: (Eq a1, Num a1) => Maybe (a2, a1) -> Maybe (a2, a1)
+cleanupChannel Nothing = Nothing
+cleanupChannel (Just (writeChan, 1)) = Nothing
+cleanupChannel (Just c) = Just c
+
+userJoinedChannel :: Num b => Maybe (a, b) -> Maybe (a, b)
+userJoinedChannel Nothing = Nothing
+userJoinedChannel (Just (writeChan, numUsers)) = Just (writeChan, numUsers + 1)
+
 chatApp :: WebSocketsT Handler ()
 chatApp = do
     sendTextData ("Welcome to the chat server, please enter your name." :: Text)
@@ -37,18 +47,28 @@ chatApp = do
     writeChan <- atomically $ case maybeChan of
                                 Nothing -> do
                                     chan <- newBroadcastTChan
-                                    writeTVar channelMapTVar $ M.insert channelId chan channelMap
+                                    writeTVar channelMapTVar $ M.insert channelId (chan, 1) channelMap
                                     return chan
-                                Just writeChan -> return writeChan
+                                Just (writeChan, _) -> do
+                                    writeTVar channelMapTVar $ M.alter userJoinedChannel channelId channelMap
+                                    return writeChan
 
     readChan <- atomically $ do
         writeTChan writeChan $ name <> " has joined the chat"
         dupTChan writeChan
-    race_
+    (e :: Either SomeException ()) <- try $ race_
         (forever $ atomically (readTChan readChan) >>= sendTextData)
         (sourceWS $$ mapM_C (\msg ->
             atomically $ writeTChan writeChan $ name <> ": " <> msg))
 
+    atomically $ case e of
+        Left _ -> do
+            -- clean up your resources when user disconnects here
+            let newChannelMap = M.alter cleanupChannel channelId channelMap
+            writeTVar channelMapTVar newChannelMap
+            writeTChan writeChan $ name <> " has left the chat"
+        Right () -> return ()
+
 getHomeR :: Handler Html
 getHomeR = do
     webSockets chatApp

yesod-websockets/chat-with-timeout-control.hs

@@ -1,4 +1,4 @@
-{-# LANGUAGE QuasiQuotes, TemplateHaskell, TypeFamilies, OverloadedStrings #-}
+{-# LANGUAGE QuasiQuotes, TemplateHaskell, TypeFamilies, OverloadedStrings, ScopedTypeVariables #-}
 import Yesod.Core
 import Yesod.WebSockets
 import qualified Data.Text.Lazy as TL
@@ -10,6 +10,7 @@ import Conduit
 import Data.Monoid ((<>))
 import Control.Concurrent.STM.Lifted
 import Data.Text (Text)
+import UnliftIO.Exception (try, SomeException)
 
 data App = App (TChan Text)
 
@@ -28,11 +29,15 @@ chatApp = do
     readChan <- atomically $ do
         writeTChan writeChan $ name <> " has joined the chat"
         dupTChan writeChan
-    race_
+    (e :: Either SomeException ()) <- try $ race_
         (forever $ atomically (readTChan readChan) >>= sendTextData)
         (sourceWS $$ mapM_C (\msg ->
             atomically $ writeTChan writeChan $ name <> ": " <> msg))
 
+    atomically $ case e of
+        Left _ -> writeTChan writeChan $ name <> " has left the chat"
+        Right () -> return ()
+
 getHomeR :: Handler Html
 getHomeR = do
     webSockets chatApp

yesod-websockets/chat.hs

@@ -1,4 +1,4 @@
-{-# LANGUAGE QuasiQuotes, TemplateHaskell, TypeFamilies, OverloadedStrings #-}
+{-# LANGUAGE QuasiQuotes, TemplateHaskell, TypeFamilies, OverloadedStrings, ScopedTypeVariables #-}
 import Yesod.Core
 import Yesod.WebSockets
 import qualified Data.Text.Lazy as TL
@@ -10,6 +10,7 @@ import Conduit
 import Data.Monoid ((<>))
 import Control.Concurrent.STM.Lifted
 import Data.Text (Text)
+import UnliftIO.Exception (try, SomeException)
 
 data App = App (TChan Text)
 
@@ -28,11 +29,15 @@ chatApp = do
     readChan <- atomically $ do
         writeTChan writeChan $ name <> " has joined the chat"
         dupTChan writeChan
-    race_
+    (e :: Either SomeException ()) <- try $ race_
         (forever $ atomically (readTChan readChan) >>= sendTextData)
         (sourceWS $$ mapM_C (\msg ->
             atomically $ writeTChan writeChan $ name <> ": " <> msg))
 
+    atomically $ case e of
+        Left _ -> writeTChan writeChan $ name <> " has left the chat"
+        Right () -> return ()
+
 getHomeR :: Handler Html
 getHomeR = do
     webSockets chatApp