A simple TCP server

Submitted by mrd on Fri, 01/19/2007 - 10:19pm.

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)

Module Network is the simple networking library, presenting a Handle-based interface.

> 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)

A simple 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 orElse and 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.