Adds ping-pong for detecting client timeout.

link.cabal

@@ -17,12 +17,15 @@ library
   hs-source-dirs:      src
   exposed-modules:     Link.Server,
                        Link.Types,
-                       Link.Protocol
+                       Link.Protocol,
+                       Link.Client
+  other-modules:       Link.Util
   default-extensions:  OverloadedStrings, RecordWildCards, ScopedTypeVariables,
                        BangPatterns, TupleSections, NamedFieldPuns
   build-depends:       base >= 4.7 && < 5,
                        network >= 2.6 && < 2.7,
-                       containers >= 0.5 && < 0.6
+                       containers >= 0.5 && < 0.6,
+                       time >= 1.4 && < 1.6
   default-language:    Haskell2010
 
 executable link-exe

src/Link/Client.hs

@@ -1,9 +1,11 @@
 module Link.Client where
 
-import Control.Exception  hiding (handle)
 import Control.Concurrent hiding (forkFinally)
-import Control.Monad      (forever)
+import Control.Exception  hiding (handle)
+import Control.Monad      (void)
+import Data.Time          (getCurrentTime, diffUTCTime)
 import System.IO          (hGetLine)
+import System.Timeout     (timeout)
 import Text.Printf        (printf)
 
 import qualified Data.Map.Strict as Map
@@ -20,29 +22,52 @@ forkFinally action fun =
 race :: IO a -> IO b -> IO (Either a b)
 race ioa iob = do
   m <- newEmptyMVar
-  bracket (forkFinally (fmap Left ioa) (putMVar m)) cancel $ \_ ->
+  bracket (forkFinally (fmap Left ioa) (putMVar m)) killThread $ \_ ->
-    bracket (forkFinally (fmap Right iob) (putMVar m)) cancel $ \_ -> do
+    bracket (forkFinally (fmap Right iob) (putMVar m)) killThread $ \_ -> do
       r <- readMVar m
       case r of
         Left e -> throwIO e
         Right a -> return a
-  where
-    cancel t = throwTo t ThreadKilled
 
 sendMessage :: Message -> Client -> IO ()
 sendMessage message Client {..} = writeChan clientChan message
 
 runClient :: Server -> Client -> IO ()
-runClient Server {..} Client {..} = forever $ do
+runClient Server {..} client@Client {..} = do
-  r <- try $ race readCommand readMessage
+  clientAlive <- newMVar True
-  case r of
+  pingThread <- forkIO $ ping clientAlive
-    Left (e :: SomeException) -> printf "Exception: %s\n" (show e)
+  run clientAlive `finally` killThread pingThread
-    Right cm -> case cm of
-      Left mcommand -> case mcommand of
-        Nothing      -> printf "Could not parse command\n"
-        Just command -> handleCommand command
-      Right message -> handleMessage message
   where
+    pingDelay = 5
+    pingDelayMicros = pingDelay * 1000 * 1000
+
+    ping clientAlive = do
+      sendMessage Ping client
+      threadDelay pingDelayMicros
+      now <- getCurrentTime
+      pongTime <- readMVar clientPongTime
+      if diffUTCTime now pongTime > fromIntegral pingDelay
+        then void $ swapMVar clientAlive False
+        else ping clientAlive
+
+    run clientAlive = do
+      alive <- readMVar clientAlive
+      if not alive
+        then printf "Client timed out: %s\n" (userName clientUser)
+        else do
+          r <-  try . timeout pingDelayMicros $ race readCommand readMessage
+          case r of
+            Left (e :: SomeException) -> printf "Exception: %s\n" (show e)
+            Right g -> case g of
+              Nothing -> run clientAlive
+              Just cm  -> do
+                case cm of
+                  Left mcommand -> case mcommand of
+                    Nothing      -> printf "Could not parse command\n"
+                    Just command -> handleCommand command
+                  Right message -> handleMessage message
+                run clientAlive
+
     readCommand = do
       command <- hGetLine clientHandle
       printf "<%s>: %s\n" (userName clientUser) command
@@ -54,6 +79,9 @@ runClient Server {..} Client {..} = forever $ do
       withMVar serverUsers $ \clientMap ->
         case Map.lookup user clientMap of
           Nothing -> printf "No such user: %s\n" (userName user)
-          Just client -> sendMessage (PrivMsg clientUser msg) client
+          Just client' -> sendMessage (PrivMsg clientUser msg) client'
+    handleCommand Pong = do
+      now <- getCurrentTime
+      void $ swapMVar clientPongTime now
 
     handleMessage = printToHandle clientHandle . formatMessage

src/Link/Protocol.hs

@@ -6,10 +6,12 @@ import Link.Types
 
 parseCommand :: String -> Maybe Message
 parseCommand command = case words command of
+  ["PONG"]                   -> Just Pong
   "PRIVMSG" : userName : msg -> Just $ PrivMsg (User userName) (unwords msg)
-  _  -> Nothing
+  _                          -> Nothing
 
 formatMessage :: Message -> String
 formatMessage (PrivMsg user msg) = printf "PRIVMSG %s %s" (userName user) msg
 formatMessage (NameInUse name)   = printf "NAMEINUSE %s" name
 formatMessage (Connected name)   = printf "CONNECTED %s" name
+formatMessage Ping               = "PING"

src/Link/Server.hs

@@ -3,6 +3,7 @@ module Link.Server where
 import Control.Exception  hiding (handle)
 import Control.Concurrent hiding (forkFinally)
 import Control.Monad      (forever)
+import Data.Time          (getCurrentTime)
 import Network            (withSocketsDo, listenOn, accept, PortID(..))
 import System.IO          (hClose, hSetNewlineMode, hSetBuffering, BufferMode(..),
                            universalNewlineMode, hGetLine, Handle, stdout)
@@ -25,7 +26,7 @@ runServer port = withSocketsDo $ do
   forever $ do
      (handle, host, port') <- accept sock
      printf "Accepted connection from %s: %s\n" host (show port')
-     forkIO (connectClient server handle) `finally` (hClose handle)
+     forkIO $ connectClient server handle `finally` hClose handle
 
 connectClient :: Server -> Handle -> IO ()
 connectClient server handle = do
@@ -55,7 +56,8 @@ checkAddClient Server {..} user@User {..} handle =
       then return (clientMap, Nothing)
       else do
         clientChan <- newChan
-        let client = Client user handle clientChan
+        now <- newMVar =<< getCurrentTime
+        let client = Client user handle clientChan now
         printf "New user connected: %s\n" userName
         return (Map.insert user client clientMap, Just client)
 

src/Link/Types.hs

@@ -1,5 +1,6 @@
 module Link.Types where
 
+import Data.Time (UTCTime)
 import System.IO (Handle)
 import Control.Concurrent (MVar, Chan)
 import qualified Data.Map as Map
@@ -10,9 +11,10 @@ data User = User { userName :: !UserName }
             deriving (Show, Eq, Ord)
 
 data Client = Client {
-                clientUser   :: !User
+                clientUser     :: !User
-              , clientHandle :: !Handle
+              , clientHandle   :: !Handle
-              , clientChan   :: !(Chan Message)
+              , clientChan     :: !(Chan Message)
+              , clientPongTime :: MVar UTCTime
               }
 
 data Server = Server {
@@ -21,5 +23,7 @@ data Server = Server {
 
 data Message = NameInUse UserName
              | Connected UserName
+             | Ping
+             | Pong
              | PrivMsg User String
                deriving (Show, Eq)

src/Link/Util.hs

@@ -4,4 +4,4 @@ import System.IO   (Handle)
 import Text.Printf (hPrintf)
 
 printToHandle :: Handle -> String -> IO ()
-printToHandle handle str = hPrintf handle "%s\n" str
+printToHandle handle = hPrintf handle "%s\n"