{-# LANGUAGE TemplateHaskell, DeriveDataTypeable, DeriveGeneric #-}
{-# LANGUAGE RecordWildCards #-}
import Control.Distributed.Process hiding (mask, finally)
import Control.Distributed.Process.Closure
import Control.Concurrent.Async
import Control.Monad.IO.Class
import Control.Monad
import Text.Printf
import Control.Concurrent
import GHC.Generics (Generic)
import Data.Binary
import Data.Typeable
import Network
import System.IO
import Control.Concurrent.STM
import qualified Data.Map as Map
import Data.Map (Map)
import qualified Data.Foldable as F
import Control.Exception
import ConcurrentUtils
import DistribUtils
-- ---------------------------------------------------------------------------
-- Data structures and initialisation
-- <<Client
type ClientName = String
data Client
= ClientLocal LocalClient
| ClientRemote RemoteClient
data RemoteClient = RemoteClient
{ remoteName :: ClientName
, clientHome :: ProcessId
}
data LocalClient = LocalClient
{ localName :: ClientName
, clientHandle :: Handle
, clientKicked :: TVar (Maybe String)
, clientSendChan :: TChan Message
}
clientName :: Client -> ClientName
clientName (ClientLocal c) = localName c
clientName (ClientRemote c) = remoteName c
newLocalClient :: ClientName -> Handle -> STM LocalClient
newLocalClient name handle = do
c <- newTChan
k <- newTVar Nothing
return LocalClient { localName = name
, clientHandle = handle
, clientSendChan = c
, clientKicked = k
}
-- >>
-- <<Message
data Message = Notice String
| Tell ClientName String
| Broadcast ClientName String
| Command String
deriving (Typeable, Generic)
instance Binary Message
-- >>
-- <<PMessage
data PMessage
= MsgServers [ProcessId]
| MsgSend ClientName Message
| MsgBroadcast Message
| MsgKick ClientName ClientName
| MsgNewClient ClientName ProcessId
| MsgClientDisconnected ClientName ProcessId
deriving (Typeable, Generic)
instance Binary PMessage
-- >>
-- <<Server
data Server = Server
{ clients :: TVar (Map ClientName Client)
, proxychan :: TChan (Process ())
, servers :: TVar [ProcessId]
, spid :: ProcessId
}
newServer :: [ProcessId] -> Process Server
newServer pids = do
pid <- getSelfPid
liftIO $ do
s <- newTVarIO pids
c <- newTVarIO Map.empty
o <- newTChanIO
return Server { clients = c, servers = s, proxychan = o, spid = pid }
-- >>
-- -----------------------------------------------------------------------------
-- Basic operations
-- <<sendLocal
sendLocal :: LocalClient -> Message -> STM ()
sendLocal LocalClient{..} msg = writeTChan clientSendChan msg
-- >>
-- <<sendRemote
sendRemote :: Server -> ProcessId -> PMessage -> STM ()
sendRemote Server{..} pid pmsg = writeTChan proxychan (send pid pmsg)
-- >>
-- <<sendMessage
sendMessage :: Server -> Client -> Message -> STM ()
sendMessage server (ClientLocal client) msg =
sendLocal client msg
sendMessage server (ClientRemote client) msg =
sendRemote server (clientHome client) (MsgSend (remoteName client) msg)
-- >>
-- <<sendToName
sendToName :: Server -> ClientName -> Message -> STM Bool
sendToName server@Server{..} name msg = do
clientmap <- readTVar clients
case Map.lookup name clientmap of
Nothing -> return False
Just client -> sendMessage server client msg >> return True
-- >>
-- <<sendRemoteAll
sendRemoteAll :: Server -> PMessage -> STM ()
sendRemoteAll server@Server{..} pmsg = do
pids <- readTVar servers
mapM_ (\pid -> sendRemote server pid pmsg) pids
-- >>
-- <<broadcastLocal
broadcastLocal :: Server -> Message -> STM ()
broadcastLocal server@Server{..} msg = do
clientmap <- readTVar clients
mapM_ sendIfLocal (Map.elems clientmap)
where
sendIfLocal (ClientLocal c) = sendLocal c msg
sendIfLocal (ClientRemote _) = return ()
-- >>
-- <<broadcast
broadcast :: Server -> Message -> STM ()
broadcast server@Server{..} msg = do
sendRemoteAll server (MsgBroadcast msg)
broadcastLocal server msg
-- >>
-- <<tell
tell :: Server -> LocalClient -> ClientName -> String -> IO ()
tell server@Server{..} LocalClient{..} who msg = do
ok <- atomically $ sendToName server who (Tell localName msg)
if ok
then return ()
else hPutStrLn clientHandle (who ++ " is not connected.")
-- >>
-- <<kick
kick :: Server -> ClientName -> ClientName -> STM ()
kick server@Server{..} who by = do
clientmap <- readTVar clients
case Map.lookup who clientmap of
Nothing ->
void $ sendToName server by (Notice $ who ++ " is not connected")
Just (ClientLocal victim) -> do
writeTVar (clientKicked victim) $ Just ("by " ++ by)
void $ sendToName server by (Notice $ "you kicked " ++ who)
Just (ClientRemote victim) -> do
sendRemote server (clientHome victim) (MsgKick who by)
-- >>
-- -----------------------------------------------------------------------------
-- Handle a local client
talk :: Server -> Handle -> IO ()
talk server@Server{..} handle = do
hSetNewlineMode handle universalNewlineMode
-- Swallow carriage returns sent by telnet clients
hSetBuffering handle LineBuffering
readName
where
-- <<readName
readName = do
hPutStrLn handle "What is your name?"
name <- hGetLine handle
if null name
then readName
else mask $ \restore -> do
client <- atomically $ newLocalClient name handle
ok <- atomically $ checkAddClient server (ClientLocal client)
if not ok
then restore $ do
hPrintf handle
"The name %s is in use, please choose another\n" name
readName
else do
atomically $ sendRemoteAll server (MsgNewClient name spid)
restore (runClient server client)
`finally` disconnectLocalClient server name
-- >>
checkAddClient :: Server -> Client -> STM Bool
checkAddClient server@Server{..} client = do
clientmap <- readTVar clients
let name = clientName client
if Map.member name clientmap
then return False
else do writeTVar clients (Map.insert name client clientmap)
broadcastLocal server $ Notice $ name ++ " has connected"
return True
deleteClient :: Server -> ClientName -> STM ()
deleteClient server@Server{..} name = do
modifyTVar' clients $ Map.delete name
broadcastLocal server $ Notice $ name ++ " has disconnected"
disconnectLocalClient :: Server -> ClientName -> IO ()
disconnectLocalClient server@Server{..} name = atomically $ do
deleteClient server name
sendRemoteAll server (MsgClientDisconnected name spid)
-- <<runClient
runClient :: Server -> LocalClient -> IO ()
runClient serv@Server{..} client@LocalClient{..} = do
race server receive
return ()
where
receive = forever $ do
msg <- hGetLine clientHandle
atomically $ sendLocal client (Command msg)
server = join $ atomically $ do
k <- readTVar clientKicked
case k of
Just reason -> return $
hPutStrLn clientHandle $ "You have been kicked: " ++ reason
Nothing -> do
msg <- readTChan clientSendChan
return $ do
continue <- handleMessage serv client msg
when continue $ server
-- >>
-- <<handleMessage
handleMessage :: Server -> LocalClient -> Message -> IO Bool
handleMessage server client@LocalClient{..} message =
case message of
Notice msg -> output $ "*** " ++ msg
Tell name msg -> output $ "*" ++ name ++ "*: " ++ msg
Broadcast name msg -> output $ "<" ++ name ++ ">: " ++ msg
Command msg ->
case words msg of
["/kick", who] -> do
atomically $ kick server who localName
return True
"/tell" : who : what -> do
tell server client who (unwords what)
return True
["/quit"] ->
return False
('/':_):_ -> do
hPutStrLn clientHandle $ "Unrecognised command: " ++ msg
return True
_ -> do
atomically $ broadcast server $ Broadcast localName msg
return True
where
output s = do hPutStrLn clientHandle s; return True
-- >>
-- -----------------------------------------------------------------------------
-- Main server
-- <<socketListener
socketListener :: Server -> Int -> IO ()
socketListener server port = withSocketsDo $ do
sock <- listenOn (PortNumber (fromIntegral port))
printf "Listening on port %d\n" port
forever $ do
(handle, host, port) <- accept sock
printf "Accepted connection from %s: %s\n" host (show port)
forkFinally (talk server handle)
(\_ -> hClose handle)
-- >>
-- <<proxy
proxy :: Server -> Process ()
proxy Server{..} = forever $ join $ liftIO $ atomically $ readTChan proxychan
-- >>
-- <<chatServer
chatServer :: Int -> Process ()
chatServer port = do
server <- newServer []
liftIO $ forkIO (socketListener server port) -- <1>
spawnLocal (proxy server) -- <2>
forever $ do m <- expect; handleRemoteMessage server m -- <3>
-- >>
-- <<handleRemoteMessage
handleRemoteMessage :: Server -> PMessage -> Process ()
handleRemoteMessage server@Server{..} m = liftIO $ atomically $
case m of
MsgServers pids -> writeTVar servers (filter (/= spid) pids) -- <1>
MsgSend name msg -> void $ sendToName server name msg -- <2>
MsgBroadcast msg -> broadcastLocal server msg -- <2>
MsgKick who by -> kick server who by -- <2>
MsgNewClient name pid -> do -- <3>
ok <- checkAddClient server (ClientRemote (RemoteClient name pid))
when (not ok) $
sendRemote server pid (MsgKick name "SYSTEM")
MsgClientDisconnected name pid -> do -- <4>
clientmap <- readTVar clients
case Map.lookup name clientmap of
Nothing -> return ()
Just (ClientRemote (RemoteClient _ pid')) | pid == pid' ->
deleteClient server name
Just _ ->
return ()
-- >>
remotable ['chatServer]
-- <<main
port :: Int
port = 44444
master :: [NodeId] -> Process ()
master peers = do
let run nid port = do
say $ printf "spawning on %s" (show nid)
spawn nid ($(mkClosure 'chatServer) port)
pids <- zipWithM run peers [port+1..]
mypid <- getSelfPid
let all_pids = mypid : pids
mapM_ (\pid -> send pid (MsgServers all_pids)) all_pids
chatServer port
main = distribMain master Main.__remoteTable
-- >>