Shutdown socket gracefully. Simplify server cleanup.

README.md

@@ -2,7 +2,7 @@
 
 [![Build Status](https://secure.travis-ci.org/k0001/network-simple.png)](http://travis-ci.org/k0001/network-simple)
 
-Haskell library abstracring over simple network sockets usage patterns.
+Haskell library abstracting over simple network sockets usage patterns.
 Currently, only TCP sockets are supported. Support for UDP and Unix
 sockets is planned for future versions.
 

network-simple.cabal

@@ -1,5 +1,5 @@
 name:                network-simple
-version:             0.4.0.3
+version:             0.4.0.4
 homepage:            https://github.com/k0001/network-simple
 bug-reports:         https://github.com/k0001/network-simple/issues
 license:             BSD3
@@ -16,7 +16,7 @@ description:
   usage patterns.
   .
   See the @changelog.md@ file in the source distribution to learn about any
-  important changes between version.
+  important changes between versions.
 extra-source-files:
   changelog.md
   README.md
@@ -38,4 +38,5 @@ library
                    , bytestring   (>=0.9.2.1 && <0.11)
                    , transformers (>=0.2 && <0.5)
                    -- Packages not in The Haskell Platform
+                   , slave-thread
                    , exceptions   (>=0.6 && <0.9)

src/Network/Simple/TCP.hs

@@ -56,18 +56,19 @@ module Network.Simple.TCP (
   , NS.SockAddr
   ) where
 
-import           Control.Concurrent             (ThreadId, forkIO)
-import qualified Control.Exception              as E
-import qualified Control.Monad.Catch            as C
+import           Control.Concurrent (ThreadId)
+import qualified Control.Exception as E
 import           Control.Monad
-import           Control.Monad.IO.Class         (MonadIO(liftIO))
-import qualified Data.ByteString                as BS
-import qualified Data.ByteString.Lazy           as BSL
-import           Data.List                      (partition)
-import qualified Network.Socket                 as NS
+import qualified Control.Monad.Catch as C
+import           Control.Monad.IO.Class (MonadIO(liftIO))
+import qualified Data.ByteString as BS
+import qualified Data.ByteString.Lazy as BSL
+import           Data.List (partition)
 import           Network.Simple.Internal
-import qualified Network.Socket.ByteString      as NSB
+import qualified Network.Socket as NS
+import qualified Network.Socket.ByteString as NSB
 import qualified Network.Socket.ByteString.Lazy as NSBL
+import           SlaveThread (forkFinally)
 
 --------------------------------------------------------------------------------
 -- $tcp-101
@@ -156,7 +157,9 @@ connect host port = C.bracket (connectSock host port)
 -- in case of exceptions.
 --
 -- Note: This function performs 'listen' and 'acceptFork', so you don't need to
--- perform those manually.
+-- perform those manually. 'acceptFork' uses 'forkThread' from 'SlaveThread'.
+-- If you start your server with 'forkFinally' or 'fork' from this package it
+-- will close all incoming connections for you on exit.
 serve
   :: MonadIO m
   => HostPreference   -- ^Preferred host to bind.
@@ -179,7 +182,7 @@ serve hp port k = liftIO $ do
 -- If you prefer to acquire and close the socket yourself, then use 'bindSock',
 -- 'closeSock' and the 'NS.listen' function from "Network.Socket" instead.
 --
--- Note: 'N.maxListenQueue' is tipically 128, which is too small for high
+-- Note: 'N.maxListenQueue' is typically 128, which is too small for high
 -- performance servers. So, we use the maximum between 'N.maxListenQueue' and
 -- 2048 as the default size of the listening queue. The 'NS.NoDelay' and
 -- 'NS.ReuseAddr' options are set on the socket.
@@ -228,9 +231,8 @@ acceptFork
   -> m ThreadId
 acceptFork lsock k = liftIO $ do
     conn@(csock,_) <- NS.accept lsock
-    forkFinally (k conn)
-                (\ea -> do silentCloseSock csock
-                           either E.throwIO return ea)
+    forkFinally (silentCloseSock csock) (k conn)
+
 {-# INLINABLE acceptFork #-}
 
 --------------------------------------------------------------------------------
@@ -292,11 +294,13 @@ bindSock hp port = liftIO $ do
 
 -- | Close the 'NS.Socket'.
 closeSock :: MonadIO m => NS.Socket -> m ()
-closeSock = liftIO .
+closeSock soc =
+  liftIO (NS.shutdown soc NS.ShutdownBoth >> close' soc)
+  where
 #if MIN_VERSION_network(2,4,0)
-    NS.close
+    close' = NS.close
 #else
-    NS.sClose
+    close' = NS.sClose
 #endif
 {-# INLINE closeSock #-}
 
@@ -357,14 +361,6 @@ prioritize p = uncurry (++) . partition p
 
 --------------------------------------------------------------------------------
 
--- | 'Control.Concurrent.forkFinally' was introduced in base==4.6.0.0. We'll use
--- our own version here for a while, until base==4.6.0.0 is widely establised.
-forkFinally :: IO a -> (Either E.SomeException a -> IO ()) -> IO ThreadId
-forkFinally action and_then =
-    E.mask $ \restore ->
-        forkIO $ E.try (restore action) >>= and_then
-
-
 -- | Like 'closeSock', except it swallows all 'IOError' exceptions.
 silentCloseSock :: MonadIO m => NS.Socket -> m ()
 silentCloseSock sock = liftIO $ do