diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,12 @@ # Revision history for network-run +## 0.8.0 + +A server told to stop returns only after its connections are done. + +* Wait for the connections once the server is told to stop + [#21](https://github.com/kazu-yamamoto/network-run/pull/21) + ## 0.7.1 * Sockets this library owns and has not handed to anyone are closed with diff --git a/Network/Run/Core.hs b/Network/Run/Core.hs --- a/Network/Run/Core.hs +++ b/Network/Run/Core.hs @@ -21,6 +21,9 @@ ServerSettings (..), defaultServerSettings, gcloseWith, + Connections, + newConnections, + waitConnections, forkWith, forkConnection, forkDatagram, @@ -224,10 +227,11 @@ -- nothing, and the server is then stopped the other way, by an -- asynchronous exception to the thread running it. -- - -- __The server returns without waiting for the connections it is - -- serving.__ Stopping ends the accepting, not the handlers: a - -- caller which has to see them out has to do that itself, before - -- letting the process go. + -- __The server returns once the connections it is serving are + -- done.__ Stopping ends the accepting, not the handlers: they have + -- to end on their own, and a caller which cannot wait for them puts + -- a limit around the server. A server stopped by an asynchronous + -- exception does not wait. -- -- This is installed once per server, and a TCP and a UDP server -- share nothing but these settings, so running both with one @@ -350,21 +354,36 @@ ---------------------------------------------------------------- +-- | The threads a server has forked and which have not finished. +newtype Connections = Connections (TVar Int) + +newConnections :: IO Connections +newConnections = Connections <$> newTVarIO 0 + +-- | Waiting until every thread forked has finished. +waitConnections :: Connections -> IO () +waitConnections (Connections var) = atomically $ readTVar var >>= check . (== 0) + -- | Forking a thread for an accepted socket. An exception which -- escapes the action is reported, and the socket is closed by the --- given closer in either case. +-- given closer in either case. The thread is counted from before it +-- is forked until after the socket is closed. forkWith :: ServerSettings + -> Connections -> (Socket -> IO ()) -> Socket -> SockAddr -> IO a -> IO () -forkWith set closer sock peer action = void $ forkFinally action finish +forkWith set (Connections var) closer sock peer action = E.mask $ \restore -> do + atomically $ modifyTVar' var (+ 1) + void (forkIO $ E.try (restore action) >>= finish) `E.onException` done where -- The socket must be closed even if the hook throws, which it does -- when an asynchronous exception arrives while it is running. - finish er = reporting er `E.finally` closing + finish er = (reporting er `E.finally` closing) `E.finally` done + done = atomically $ modifyTVar' var (subtract 1) reporting (Right _) = return () reporting (Left se) = report set (Just peer) se @@ -375,9 +394,11 @@ onCloseError e = report set (Just peer) $ E.toException e -- | 'forkWith' closing the socket gracefully. For TCP. -forkConnection :: ServerSettings -> Socket -> SockAddr -> IO a -> IO () -forkConnection set = forkWith set (gcloseWith set) +forkConnection + :: ServerSettings -> Connections -> Socket -> SockAddr -> IO a -> IO () +forkConnection set conns = forkWith set conns (gcloseWith set) -- | 'forkWith' closing the socket immediately. For UDP. -forkDatagram :: ServerSettings -> Socket -> SockAddr -> IO a -> IO () -forkDatagram set = forkWith set close +forkDatagram + :: ServerSettings -> Connections -> Socket -> SockAddr -> IO a -> IO () +forkDatagram set conns = forkWith set conns close diff --git a/Network/Run/TCP.hs b/Network/Run/TCP.hs --- a/Network/Run/TCP.hs +++ b/Network/Run/TCP.hs @@ -54,8 +54,10 @@ -> IO () runTCPServerWithSettings set mhost port server = do addr <- resolve Stream mhost port [AI_PASSIVE] NE.head + -- Closed before the connections are waited for, so that the port is + -- free while they finish; the bracket closes it if that is not reached. E.bracket (openTCPServerSocket addr) unsafeClose $ \sock -> - runTCPServerWithSocketAndSettings set sock server + serve set (unsafeClose sock) sock server -- | Running a TCP server on a given listen socket. runTCPServerWithSocket @@ -75,21 +77,29 @@ -> (Socket -> IO a) -- ^ Called for each incoming connection, in a new thread -> IO () -runTCPServerWithSocketAndSettings set sock server = do +runTCPServerWithSocketAndSettings set = serve set (return ()) + +-- | Running the accept loop, then the given action, then waiting for +-- the connections. +serve :: ServerSettings -> IO () -> Socket -> (Socket -> IO a) -> IO () +serve set stopped sock server = do listener <- makeListener set sock - loop listener + conns <- newConnections + loop listener conns + stopped + waitConnections conns where -- Accepting inside the acquire: once it has returned, the socket -- is closed if anything goes wrong before the thread owning it has -- been forked. - loop listener = do + loop listener conns = do more <- E.bracketOnError (acceptOrStop listener) (mapM_ (unsafeClose . fst)) $ \maccepted -> case maccepted of Nothing -> return False Just (conn, peer) -> do - forkConnection set conn peer (labelMe "TCP server" >> server conn) + forkConnection set conns conn peer (labelMe "TCP server" >> server conn) return True - when more $ loop listener + when more $ loop listener conns ---------------------------------------------------------------- diff --git a/Network/Run/TCP/Timeout.hs b/Network/Run/TCP/Timeout.hs --- a/Network/Run/TCP/Timeout.hs +++ b/Network/Run/TCP/Timeout.hs @@ -59,8 +59,10 @@ -> IO () runTCPServerWithSettings set tm mhost port server = do addr <- resolve Stream mhost port [AI_PASSIVE] NE.head + -- Closed before the connections are waited for, so that the port is + -- free while they finish; the bracket closes it if that is not reached. E.bracket (openTCPServerSocket addr) unsafeClose $ \sock -> - runTCPServerWithSocketAndSettings set tm sock server + serve set (unsafeClose sock) tm sock server -- | Running a TCP server on a given listen socket. runTCPServerWithSocket @@ -82,21 +84,31 @@ -- ^ A listening socket created by 'openTCPServerSocket'. -> TimeoutServer () -> IO () -runTCPServerWithSocketAndSettings set tm sock server = do +runTCPServerWithSocketAndSettings set = serve set (return ()) + +-- | Running the accept loop, then the given action, then waiting for +-- the connections, which use the manager until they are done. +serve + :: ServerSettings -> IO () -> Int -> Socket -> TimeoutServer () -> IO () +serve set stopped tm sock server = do listener <- makeListener set sock - T.withManager (tm * 1000000) $ \mgr -> loop listener mgr + conns <- newConnections + T.withManager (tm * 1000000) $ \mgr -> do + loop listener conns mgr + stopped + waitConnections conns where -- Accepting inside the acquire: once it has returned, the socket -- is closed if anything goes wrong before the thread owning it has -- been forked. - loop listener mgr = do + loop listener conns mgr = do more <- E.bracketOnError (acceptOrStop listener) (mapM_ (unsafeClose . fst)) $ \maccepted -> case maccepted of Nothing -> return False Just (conn, peer) -> do - forkConnection set conn peer (runServer mgr conn) + forkConnection set conns conn peer (runServer mgr conn) return True - when more $ loop listener mgr + when more $ loop listener conns mgr runServer mgr conn = do labelMe "TCP timeout server" T.withHandleKillThread mgr (return ()) $ \th -> server mgr th conn diff --git a/Network/Run/UDP.hs b/Network/Run/UDP.hs --- a/Network/Run/UDP.hs +++ b/Network/Run/UDP.hs @@ -12,9 +12,10 @@ resolve, ) where -import Control.Concurrent (forkIO) +import Control.Concurrent (forkIO, newEmptyMVar, putMVar, takeMVar) import Control.Concurrent.STM import qualified Control.Exception as E +import Control.Monad (forM) import Data.ByteString (ByteString) import qualified Data.List.NonEmpty as NE import Network.Socket @@ -51,7 +52,9 @@ :: [HostName] -> ServiceName -> (Socket -> ByteString -> IO ()) -> IO () runUDPServerFork = runUDPServerForkWithSettings defaultServerSettings --- | 'runUDPServerFork' with the given settings. +-- | 'runUDPServerFork' with the given settings. Once stopped, it +-- returns when the loops of all the hosts and the threads they forked +-- are done. runUDPServerForkWithSettings :: ServerSettings -> [HostName] @@ -64,30 +67,37 @@ -- server, which is all of these loops. stopping <- newTVarIO False settingsInstallShutdownHandler set $ atomically $ writeTVar stopping True - mapM_ (forkIO . run stopping) hs - run stopping h + conns <- newConnections + dones <- forM hs $ \host -> do + done <- newEmptyMVar + -- What escapes is left to the thread, as before, to be printed. + _ <- forkIO $ run stopping conns host `E.finally` putMVar done () + return done + run stopping conns h + mapM_ takeMVar dones + waitConnections conns where - run stopping host = do + run stopping conns host = do labelMe $ "UDP server for " ++ host - runUDPServer (Just host) port $ \lsock -> loop stopping lsock + runUDPServer (Just host) port $ \lsock -> loop stopping conns lsock -- The receive and the stop variable are waited for at once, so -- nothing has to be closed under this thread and nothing has to be -- thrown into it. - loop stopping lsock = do + loop stopping conns lsock = do r <- recvFromWith lsock 2048 $ check =<< readTVar stopping case r of Left () -> return () Right (bs0, peeraddr) -> do -- Everything below is per-datagram work. A failure here -- must not take the entire server down. - dispatch peeraddr bs0 `E.catch` onDispatchError peeraddr - loop stopping lsock + dispatch conns peeraddr bs0 `E.catch` onDispatchError peeraddr + loop stopping conns lsock onDispatchError peeraddr e = report set (Just peeraddr) $ E.toException (e :: E.IOException) - dispatch peeraddr bs0 = case familyOf peeraddr of + dispatch conns peeraddr bs0 = case familyOf peeraddr of -- Neither IPv4 nor IPv6. Just drop the datagram. Nothing -> return () Just family -> do @@ -103,7 +113,7 @@ -- its finalizer. E.bracketOnError (openServerSocket addr) unsafeClose $ \s -> do connect s peeraddr - forkDatagram set s peeraddr $ + forkDatagram set conns s peeraddr $ labelMe "UDP server" >> server s bs0 familyOf SockAddrInet{} = Just AF_INET diff --git a/network-run.cabal b/network-run.cabal --- a/network-run.cabal +++ b/network-run.cabal @@ -1,6 +1,6 @@ cabal-version: 2.0 name: network-run -version: 0.7.1 +version: 0.8.0 license: BSD3 license-file: LICENSE maintainer: kazu@iij.ad.jp diff --git a/test/CoreSpec.hs b/test/CoreSpec.hs --- a/test/CoreSpec.hs +++ b/test/CoreSpec.hs @@ -102,7 +102,8 @@ (set, getReports) <- collecting withDummySocket $ \sock -> do closed <- newEmptyMVar - forkWith set (\_ -> putMVar closed ()) sock fakePeer $ + conns <- newConnections + forkWith set conns (\_ -> putMVar closed ()) sock fakePeer $ E.throwIO $ userError "boom" takeMVar closed `shouldReturn` () @@ -116,16 +117,35 @@ } withDummySocket $ \sock -> do closed <- newEmptyMVar - forkWith set (\_ -> putMVar closed ()) sock fakePeer $ + conns <- newConnections + forkWith set conns (\_ -> putMVar closed ()) sock fakePeer $ E.throwIO $ userError "boom" r <- timeout 1000000 $ takeMVar closed r `shouldBe` Just () + it "counts the thread until its socket is closed" $ limited $ do + conns <- newConnections + release <- newEmptyMVar + closed <- newIORef False + withDummySocket $ \sock -> + forkWith + defaultServerSettings + conns + (\_ -> writeIORef closed True) + sock + fakePeer $ + takeMVar release + timeout 200000 (waitConnections conns) `shouldReturn` Nothing + putMVar release () + timeout 1000000 (waitConnections conns) `shouldReturn` Just () + readIORef closed `shouldReturn` True + it "reports a failure of the closer" $ limited $ do (set, getReports) <- collecting + conns <- newConnections withDummySocket $ \sock -> - forkWith set (\_ -> ioError $ userError "close failed") sock fakePeer $ + forkWith set conns (\_ -> ioError $ userError "close failed") sock fakePeer $ return () reports <- waitFor 1 getReports case reports of diff --git a/test/TCPSpec.hs b/test/TCPSpec.hs --- a/test/TCPSpec.hs +++ b/test/TCPSpec.hs @@ -10,6 +10,7 @@ import Data.IORef import Data.List (nub) import qualified Data.List.NonEmpty as NE +import Data.Maybe (isJust) import GHC.IO.Exception (IOErrorType (InvalidArgument), ioe_errno) import Network.Socket import Network.Socket.ByteString @@ -102,6 +103,55 @@ Just (Left e) -> expectationFailure $ "the accept loop threw: " ++ show e Nothing -> expectationFailure "the accept loop did not stop" + + it "waits for its connections once told to stop" $ limited $ do + told <- newEmptyMVar + done <- newEmptyMVar + let set = defaultServerSettings{settingsInstallShutdownHandler = putMVar told} + withListenSocket $ \lsock port -> do + void $ + forkFinally + (runTCPServerWithSocketAndSettings set lsock echo) + (putMVar done) + stop <- takeMVar told + client port $ \sock -> do + sendAll sock "hello" + _ <- recv sock 1024 + stop + isJust <$> timeout 300000 (readMVar done) `shouldReturn` False + r <- timeout 2000000 $ takeMVar done + fmap (either (Left . show) Right) r `shouldBe` Just (Right ()) + + it "frees its port while waiting for its connections" $ limited $ do + told <- newEmptyMVar + done <- newEmptyMVar + port <- freeTCPPort + let set = defaultServerSettings{settingsInstallShutdownHandler = putMVar told} + void $ + forkFinally + (runTCPServerWithSettings set (Just loopback) (show port) echo) + (putMVar done) + stop <- takeMVar told + threadDelay 200000 + client port $ \sock -> do + sendAll sock "hello" + _ <- recv sock 1024 + stop + threadDelay 200000 + addr <- resolve Stream (Just loopback) (show port) [AI_PASSIVE] NE.head + E.bracket (openTCPServerSocket addr) close $ \_ -> return () + void $ timeout 2000000 $ takeMVar done + + it "does not wait for its connections when killed" $ limited $ do + done <- newEmptyMVar + withListenSocket $ \lsock port -> do + tid <- + forkFinally (runTCPServerWithSocket lsock echo) (putMVar done) + client port $ \sock -> do + sendAll sock "hello" + _ <- recv sock 1024 + killThread tid + isJust <$> timeout 1000000 (takeMVar done) `shouldReturn` True it "listens on a single address family" $ limited $ do port <- freeTCPPort