~ Minor changes.

3 files changed, 33 insertions, 21 deletions

diff --git a/src/Network/BitTorrent.hs b/src/Network/BitTorrent.hs
index 8c3189d3..c166b1b1 100644
--- a/src/Network/BitTorrent.hs
+++ b/src/Network/BitTorrent.hs
@@ -46,8 +46,7 @@ import System.Torrent.Storage
 -- | Client session with default parameters. Use it for testing only.
 withDefaultClient :: PortNumber -> PortNumber -> (ClientSession -> IO ()) -> IO ()
 withDefaultClient dhtPort listPort action = do
-  withClientSession defaultThreadCount defaultExtensions $ \client -> do
-    action client
+  withClientSession defaultThreadCount defaultExtensions listPort dhtPort action
 
 {-----------------------------------------------------------------------
     Torrent management
diff --git a/src/Network/BitTorrent/Sessions.hs b/src/Network/BitTorrent/Sessions.hs
index 2e51fde6..9a1d0c6a 100644
--- a/src/Network/BitTorrent/Sessions.hs
+++ b/src/Network/BitTorrent/Sessions.hs
@@ -48,7 +48,7 @@ module Network.BitTorrent.Sessions
        , discover
        ) where
 
-import Prelude hiding (mapM_)
+import Prelude hiding (mapM_, elem)
 
 import Control.Applicative
 import Control.Concurrent
@@ -61,6 +61,7 @@ import Control.Monad.Trans
 import Data.IORef
 import Data.Map as M
 import Data.HashMap.Strict as HM
+import Data.Foldable as F
 import Data.Set as S
 
 import Data.Serialize hiding (get)
@@ -141,14 +142,13 @@ startDHT ClientSession {..} nodePort = withRunning peerListener failure start
 
 -- | Create a new client session. The data passed to this function are
 -- usually loaded from configuration file.
-openClientSession :: SessionCount     -- ^ Maximum count of active P2P Sessions.
-                  -> [Extension]      -- ^ Extensions allowed to use.
-                  -> IO ClientSession -- ^ Client with unique peer ID.
-openClientSession n exts = do
+openClientSession :: SessionCount -> [Extension] -> PortNumber -> PortNumber -> IO ClientSession
+openClientSession n exts listenerPort _ = do
   mgr <- Ev.new
   -- TODO kill this thread when leave client
   _   <- forkIO $ loop mgr
-  ClientSession
+
+  cs <- ClientSession
     <$> genPeerId
     <*> pure exts
     <*> newEmptyMVar
@@ -160,13 +160,21 @@ openClientSession n exts = do
     <*> newTVarIO (startProgress 0)
     <*> newTVarIO HM.empty
 
+  startListener cs listenerPort
+  return cs
+
 closeClientSession :: ClientSession -> IO ()
-closeClientSession ClientSession {..} =
-  stopService nodeListener `finally` stopService peerListener
--- TODO stop all swarm sessions
+closeClientSession ClientSession {..} = do
+  stopService nodeListener
+  stopService peerListener
+
+  sws <- readTVarIO swarmSessions
+  forM_ sws closeSwarmSession
 
-withClientSession :: SessionCount -> [Extension] -> (ClientSession -> IO ()) -> IO ()
-withClientSession c es = bracket (openClientSession c es) closeClientSession
+withClientSession :: SessionCount -> [Extension]
+                  -> PortNumber -> PortNumber
+                  -> (ClientSession -> IO ()) -> IO ()
+withClientSession c es l d = bracket (openClientSession c es l d) closeClientSession
 
 -- | Get current global progress of the client. This value is usually
 -- shown to a user.
@@ -222,6 +230,14 @@ discover swarm @ SwarmSession {..} action = {-# SCC discover #-} do
         initiatePeerSession swarm addr $ \conn ->
           runP2P conn action
 
+registerSwarmSession :: SwarmSession -> IO ()
+registerSwarmSession = undefined
+
+unregisterSwarmSession :: SwarmSession -> IO ()
+unregisterSwarmSession SwarmSession {..} =
+  atomically $ modifyTVar (swarmSessions clientSession) $
+    M.delete $ tInfoHash torrentMeta
+
 newSwarmSession :: Int -> Bitfield -> ClientSession -> Torrent
                 -> IO SwarmSession
 newSwarmSession n bf cs @ ClientSession {..} t @ Torrent {..}
@@ -246,11 +262,6 @@ closeSwarmSession se @ SwarmSession {..} = do
   -- TODO the order is important!
   closeStorage storage
 
-unregisterSwarmSession :: SwarmSession -> IO ()
-unregisterSwarmSession SwarmSession {..} =
-  atomically $ modifyTVar (swarmSessions clientSession) $
-  M.delete $ tInfoHash torrentMeta
-
 getSwarm :: ClientSession -> InfoHash -> IO SwarmSession
 getSwarm cs @ ClientSession {..} ih = do
   ss <- readTVarIO swarmSessions
diff --git a/src/Network/BitTorrent/Sessions/Types.lhs b/src/Network/BitTorrent/Sessions/Types.lhs
index f94dbfa6..3f9c6db1 100644
--- a/src/Network/BitTorrent/Sessions/Types.lhs
+++ b/src/Network/BitTorrent/Sessions/Types.lhs
@@ -73,10 +73,12 @@ Thread layout
 
 When client session created 2 new threads appear:
 
-  * DHT listener - replies to DHT requests;
-
   * Peer listener - accept new P2P connection initiated by other
-peers.
+peers;
+
+  * Tracker announcer - announce that the peer have this torrent.
+
+  * OPTIONAL: DHT listener - replies to DHT requests;
 
 When swarn session created 3 new threads appear: