~ Isolate MVar operations.

2 files changed, 20 insertions, 16 deletions

diff --git a/src/Network/BitTorrent/Discovery.hs b/src/Network/BitTorrent/Discovery.hs
index 770ae818..1b9e7b26 100644
--- a/src/Network/BitTorrent/Discovery.hs
+++ b/src/Network/BitTorrent/Discovery.hs
@@ -41,16 +41,14 @@ discover swarm @ SwarmSession {..} action = {-# SCC discover #-} do
 
 startListener :: ClientSession -> PortNumber -> IO ()
 startListener cs @ ClientSession {..} port =
-  putMVar peerListener =<< startService port (listener cs (error "listener"))
+  startService peerListener port $ listener cs (error "listener")
 
 startDHT :: ClientSession -> PortNumber -> IO ()
-startDHT ClientSession {..} nodePort = do
-    maybe failure start =<< tryTakeMVar peerListener
+startDHT ClientSession {..} nodePort = withRunning peerListener failure start
   where
     start ClientService {..} = do
       ses  <- newNodeSession servPort
-      serv <- startService nodePort (dhtServer ses)
-      putMVar nodeListener serv
+      startService nodeListener nodePort (dhtServer ses)
 
     failure = throwIO $ userError msg
     msg = "unable to start DHT server: peer listener is not running"
diff --git a/src/Network/BitTorrent/Internal.lhs b/src/Network/BitTorrent/Internal.lhs
index 8461a841..7edcc977 100644
--- a/src/Network/BitTorrent/Internal.lhs
+++ b/src/Network/BitTorrent/Internal.lhs
@@ -24,6 +24,8 @@
 >          Progress(..), startProgress
 >
 >        , ClientService(..)
+>        , startService
+>        , withRunning
 >
 >          -- * Client
 >        , ClientSession ( ClientSession
@@ -33,8 +35,6 @@
 >        , withClientSession
 >        , listenerPort, dhtPort
 >
->        , startService
->
 >        , ThreadCount
 >        , defaultThreadCount
 >
@@ -96,7 +96,7 @@
 > import Control.Concurrent.STM
 > import Control.Concurrent.MSem as MSem
 > import Control.Lens
-> import Control.Monad (when, forever)
+> import Control.Monad (when, forever, (>=>))
 > import Control.Exception
 > import Control.Monad.Trans
 
@@ -304,11 +304,20 @@ so we can abstract out into ClientService:
 >   , servThread :: !ThreadId
 >   } deriving Show
 
-> startService :: PortNumber -> (PortNumber -> IO ()) -> IO ClientService
-> startService port m = ClientService port <$> forkIO (m port)
+> startService :: MVar ClientService -> PortNumber -> (PortNumber -> IO ()) -> IO ()
+> startService s port m = do
+>     stopService s
+>     putMVar s =<< spawn
+>   where
+>     spawn = ClientService port <$> forkIO (m port)
 
-> stopService :: ClientService -> IO ()
-> stopService ClientService {..} = killThread servThread
+> stopService :: MVar ClientService -> IO ()
+> stopService = tryTakeMVar >=> maybe (return ()) (killThread . servThread)
+
+Service A might depend on service B.
+
+> withRunning :: MVar ClientService -> IO () -> (ClientService -> IO ()) -> IO ()
+> withRunning dep failure action = tryTakeMVar dep >>= maybe failure action
 
 Client Sessions
 ------------------------------------------------------------------------
@@ -436,10 +445,7 @@ Retrieving client info
 
 > closeClientSession :: ClientSession -> IO ()
 > closeClientSession ClientSession {..} =
->     maybeStop (tryTakeMVar peerListener) `finally`
->     maybeStop (tryTakeMVar nodeListener)
->   where
->     maybeStop m = maybe (return ()) stopService =<< m
+>     stopService nodeListener `finally` stopService peerListener
 
 > withClientSession :: SessionCount -> [Extension] -> (ClientSession -> IO ()) -> IO ()
 > withClientSession c es = bracket (openClientSession c es) closeClientSession