[Haskell] Help needed interrupting accepting a network connection

[Cat Dancer]

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