[Haskell] Help needed interrupting accepting a network connection

I'd like to write a server accepting incoming network connections that can be gracefully shutdown. When the server is asked to shutdown, it should stop accepting new connections, finish processing any current connections, and then terminate. Clients can retry if they attempt to make a connection and the connection is refused. This allows the server to restart seamlessly: any existing connections are not interrupted, and clients will see at most a pause while the server restarts. I am using the model from Simon Marlow's Haskell Web Server (as updated by Björn Bringert and available at http://www.cs.chalmers.se/~bringert/darcs/hws/): spawning a lightweight Haskell child thread for each client connection. In Control.Exception, I see that operations such as "accept" are interruptible by exceptions thrown to the thread, so I can interrupt an accept with a dynamic exception.

-- create a datatype to use to interrupt the accept data ExitGracefully = ExitGracefully deriving Typeable

I want to control when I'm paying attention to the ExitGracefully exception. I don't want to get the exception when I'm in the middle of updating a data structure, just in a few controlled points such as when I'm in an accept. Reading further in Control.Exception, I see that I can use "block" to put off receiving the exception generally, but accept is an interruptible operation so I don't need to do anything more to get the exception inside of the accept.

block ( ... result <- catchDyn (do (clientSocket, addr) <- accept sock return $ Just clientSocket) (\ (e :: ExitGracefully) -> return Nothing)

Typing the exception "e" as an "ExitGracefully" tells catchDyn that I only need to catch exceptions of that type. If the thread has been thrown a ExitGracefully, "result" will be Nothing, but if accept returned with a client connection, "result" will be Just the clientSocket.

case result of

Nothing -> do { putStrLn "accept loop exiting"; putMVar acceptLoopDone () }

Just clientSocket ->

So far so good. I also want to keep track of when the threads spawned to handle the client connections are finished, so I use the code from the "Terminating the program" section of the Control.Concurrent documentation to keep a list of MVar's indicating when the child threads are done:

childDone <- newEmptyMVar childDoneList <- takeMVar childrenDone putMVar childrenDone (childDone : childDoneList)

then I fork a child thread to handle the connection:

clientHandle <- socketToHandle clientSocket ReadWriteMode forkIO $ handleConnection childDone clientHandle

"handleConnection" runs inside the child thread, communicating with the client. When done, it closes the clientHandle, and does an "putMVar childDone ()" to say that it done. Except that, whoops, the "takeMVar" in the accept thread code which updates the childrenDone" MVar is also interruptible. So now I'm getting an interruption right where I don't want it, when I'm updating my data structure. Only the accept thread is thrown the ExitGracefully exception, so one thought I had was that I could move those three lines which update the childrenDone MVar into the child thread. But this introduces a race condition: as the server was shutting down, it could look at the childrenDone list and see that it was empty, before the child thread had a chance to start running and update the data structure to say that there was another child that needed to be waited for. Or, updating the childrenDone MVar could be done in its own thread, which again would protect it from the ExitGracefully exception... except that how would the accept thread wait for that thread... except by using an MVar? Oops, again. Any ideas? For reference sake here's the complete implementation. (This code is in the public domain... in case it would be useful to anyone else). Thank you, Cat

-- A ConnectionHandler is a function which handles an incoming -- client connection. The handler is run in its own thread, and is -- passed a handle to the client socket. The handler does whatever -- communication it wants to do with the client, and when it returns, -- the client socket handle is closed and the thread terminates. -- A list of active handlers is kept, and the client connection is -- also marked as finished when the handler returns.

type ConnectionHandler = Handle -> IO ()

example_connection_handler :: ConnectionHandler

example_connection_handler handle = do hPutStrLn handle "Hello." hPutStrLn handle "Goodbye."

type ChildrenDone = MVar [MVar ()]

data ExitGracefully = ExitGracefully deriving Typeable

waitForChildren :: ChildrenDone -> IO ()

waitForChildren childrenDone = do cs <- takeMVar childrenDone case cs of [] -> return () m:ms -> do putMVar childrenDone ms takeMVar m waitForChildren childrenDone

shutdownServer :: MVar () -> ChildrenDone -> ThreadId -> IO ()

shutdownServer acceptLoopDone childrenDone acceptThreadId = do throwDynTo acceptThreadId ExitGracefully takeMVar acceptLoopDone waitForChildren childrenDone return ()

acceptConnections :: MVar () -> ChildrenDone -> ConnectionHandler -> Socket -> IO ()

acceptConnections acceptLoopDone childrenDone connectionHandler sock = do block (acceptConnections' acceptLoopDone childrenDone connectionHandler sock)

acceptConnections' acceptLoopDone childrenDone connectionHandler sock = do

result <- catchDyn (do (clientSocket, addr) <- accept sock return $ Just clientSocket) (\ (e :: ExitGracefully) -> return Nothing)

case result of

Nothing -> do { putStrLn "accept loop exiting"; putMVar acceptLoopDone () }

Just clientSocket -> do clientHandle <- socketToHandle clientSocket ReadWriteMode childDone <- newEmptyMVar childDoneList <- takeMVar childrenDone putMVar childrenDone (childDone : childDoneList) forkIO $ handleConnection childDone connectionHandler clientHandle acceptConnections' acceptLoopDone childrenDone connectionHandler sock

handleConnection childDone connectionHandler clientHandle = do Exception.catch (connectionHandler clientHandle `finally` do { hClose clientHandle; putMVar childDone () })

-- TODO we'll want to do something better when -- connectionHandler throws an exception, but -- for now we'll at least display the exception. (\e -> do { putStrLn $ show e; return () })