A simple TCP server
A simple TCP server which accepts multiple clients and echos input text back to all those connected. It uses threads to manage multiple handles, and Software Transactional Memory to pass messages.
This text is literate Haskell, and has been tested with ghc 6.6 on Linux/x86. Type annotations are included for didactic purposes.
> module Main where > import Prelude hiding (catch)
Network is the simple networking
library, presenting a
> import Network (listenOn, accept, sClose, Socket, > withSocketsDo, PortID(..)) > import System.IO > import System.Environment (getArgs) > import Control.Exception (finally, catch) > import Control.Concurrent > import Control.Concurrent.STM > import Control.Monad (forM, filterM, liftM, when)
main to parse a port number from
the command line, and fire up the server socket.
> main = withSocketsDo $ do > [portStr] <- getArgs > let port = fromIntegral (read portStr :: Int) > servSock <- listenOn $ PortNumber port > putStrLn $ "listening on: " ++ show port > start servSock `finally` sClose servSock
> start servSock = do > acceptChan <- atomically newTChan > forkIO $ acceptLoop servSock acceptChan > mainLoop servSock acceptChan 
acceptLoop manages the server socket, accepting
connections, starting client threads, and forwarding the relevant
information about them over the channel so the
main loop can multiplex it all together.
> type Client = (TChan String, Handle) > > acceptLoop :: Socket -> TChan Client -> IO () > acceptLoop servSock chan = do > (cHandle, host, port) <- accept servSock > cChan <- atomically newTChan > cTID <- forkIO $ clientLoop cHandle cChan > atomically $ writeTChan chan (cChan, cHandle) > acceptLoop servSock chan
As before, each client gets a loop which reads from the handle and pumps the data right into a channel. However, this time, exception handling is done per-thread; if a client disconnects we just want the thread to die silently. A more clever implementation might have a more structured channel which allows it to indicate when the client disconnects.
> clientLoop :: Handle -> TChan String -> IO () > clientLoop handle chan = > listenLoop (hGetLine handle) chan > `catch` (const $ return ()) > `finally` hClose handle > listenLoop :: IO a -> TChan a -> IO () > listenLoop act chan = > sequence_ (repeat (act >>= atomically . writeTChan chan))
STM conveniently allows composition of actions which makes custom tailoring of library code a snap. Here, I've added an additional action to check the status of the acceptChan along with all the clients. The acceptChan has a different type than any of the client channels, so I separate it from the others using an Either data type for simplicity. `fmap` here acts very much like (.), the functional composition operator.
> mainLoop :: Socket -> TChan Client -> [Client] -> IO () > mainLoop servSock acceptChan clients = do > r <- atomically $ (Left `fmap` readTChan acceptChan) > `orElse` > (Right `fmap` tselect clients) > case r of > Left (ch,h) -> do > putStrLn "new client" > mainLoop servSock acceptChan $ (ch,h):clients > Right (line,_) -> do > putStrLn $ "data: " ++ line
In addition to sending the data out to every client, this loop catches any errors from writing to handles and excludes that client from the list.
> clients' <- forM clients $ > \(ch,h) -> do > hPutStrLn h line > hFlush h > return [(ch,h)] > `catch` const (hClose h >> return ) > let dropped = length $ filter null clients' > when (dropped > 0) $ > putStrLn ("clients lost: " ++ show dropped) > mainLoop servSock acceptChan $ concat clients'
tselect is a function which multiplexes any number of
TChans. It will return the data from whichever TChan it can read,
along with the "key" value that can be supplied in the pair.
This takes advantage of the STM combinators
retry by applying them to a list of actions
constructed around the TChans.
> tselect :: [(TChan a, t)] -> STM (a, t) > tselect = foldl orElse retry > . map (\(ch, ty) -> (flip (,) ty) `fmap` readTChan ch)
This code demonstrates a basic TCP server as well as a more
generally applicable function
tselect. It serves
as an example of the strength and simplicity of the Software
Transactional Memory model, and of network IO in Haskell.