3 files changed, 96 insertions, 35 deletions

diff --git a/src/Network/BitTorrent.hs b/src/Network/BitTorrent.hs
index 85571470..b9dc39eb 100644
--- a/src/Network/BitTorrent.hs
+++ b/src/Network/BitTorrent.hs
@@ -29,8 +29,10 @@ module Network.BitTorrent
        , awaitEvent, yieldEvent
        ) where
 
+import Control.Concurrent
 import Control.Exception
 import Control.Monad
+import Control.Monad.Reader
 
 import Network
 
@@ -41,9 +43,8 @@ import Network.BitTorrent.Exchange.Protocol
 import Network.BitTorrent.Tracker
 
 
-
 -- discover should hide tracker and DHT communication under the hood
--- thus we can obtain unified interface
+-- thus we can obtain an unified interface
 
 discover :: SwarmSession -> P2P () -> IO ()
 discover swarm action = do
@@ -58,14 +59,14 @@ discover swarm action = do
 
   putStrLn "lookup peers"
   withTracker progress conn $ \tses -> do
+    putStrLn "get peer list "
     forever $ do
       addr <- getPeerAddr tses
-      putStrLn "connecting to peer"
-      handle handler (withPeer swarm addr action)
-
-  where
-    handler :: IOException -> IO ()
-    handler _ = return ()
+      putStrLn "connect to peer"
+      spawnP2P swarm addr $ do
+        liftIO $ putStrLn "run p2p session"
+        action
+      putStrLn "connected"
 
 listener :: SwarmSession -> P2P () -> IO PortNumber
 listener _ _ = do
diff --git a/src/Network/BitTorrent/Exchange.hs b/src/Network/BitTorrent/Exchange.hs
index 3235a626..978e86db 100644
--- a/src/Network/BitTorrent/Exchange.hs
+++ b/src/Network/BitTorrent/Exchange.hs
@@ -17,12 +17,14 @@ module Network.BitTorrent.Exchange
          -- * Event
        , Event(..)
 
-       , P2P, withPeer
+       , P2P
+       , runP2P, spawnP2P
        , awaitEvent, yieldEvent
        ) where
 
 import Control.Applicative
 import Control.Exception
+import Control.Concurrent
 import Control.Lens
 import Control.Monad.Reader
 import Control.Monad.Trans.Resource
@@ -54,8 +56,8 @@ data Event = Available Bitfield
 
 type PeerWire = ConduitM Message Message IO
 
-runConduit :: Socket -> PeerWire () -> IO ()
-runConduit sock p2p =
+runPeerWire :: Socket -> PeerWire () -> IO ()
+runPeerWire sock p2p =
   sourceSocket sock   $=
     conduitGet S.get  $=
       forever p2p     $=
@@ -81,7 +83,6 @@ yieldMessage msg = P2P $ ReaderT $ \se -> do
   liftIO $ print $ "sent:" <+> ppMessage msg
   liftIO $ updateOutcoming se
 
-
 peerWant :: P2P Bitfield
 peerWant   = BF.difference <$> getClientBF  <*> use bitfield
 
@@ -272,18 +273,32 @@ checkPiece = undefined
 -- |
 --   Exceptions:
 --
---     * SessionException: is visible with one peer session. Use this
---     exception to terminate P2P session, but not the swarm session.
+--     * SessionException: is visible only within one peer
+--     session. Use this exception to terminate P2P session, but not
+--     the swarm session.
 --
 newtype P2P a = P2P {
-    runP2P :: ReaderT PeerSession PeerWire a
+    unP2P :: ReaderT PeerSession PeerWire a
   } deriving ( Functor, Applicative, Monad
              , MonadIO, MonadThrow, MonadActive
              , MonadReader PeerSession
              )
 
-withPeer :: SwarmSession -> PeerAddr -> P2P () -> IO ()
-withPeer se addr p2p =
+runSession :: SwarmSession -> PeerAddr -> P2P () -> IO ()
+runSession  se addr p2p =
   withPeerSession se addr $ \(sock, pses) -> do
-    handle putSessionException $
-      runConduit sock (runReaderT (runP2P p2p) pses)
+    runPeerWire sock (runReaderT (unP2P p2p) pses)
+
+-- | Run P2P session in the current thread. Normally you don't need this
+-- function in client application.
+runP2P :: SwarmSession -> PeerAddr -> P2P () -> IO ()
+runP2P se addr p2p = waitVacancy se $ runSession se addr p2p
+
+-- | Run P2P session in forked thread. Might be used in listener or
+-- some other loop. Note that this function may block while waiting
+-- for a vacant place: use forkIO and runP2P instead.
+spawnP2P :: SwarmSession -> PeerAddr -> P2P () -> IO ThreadId
+spawnP2P se addr p2p = do
+  enterSwarm se
+  forkIO $ do
+    runSession se addr p2p `finally` leaveSwarm se
diff --git a/src/Network/BitTorrent/Internal.hs b/src/Network/BitTorrent/Internal.hs
index 0355c147..5c9efd1f 100644
--- a/src/Network/BitTorrent/Internal.hs
+++ b/src/Network/BitTorrent/Internal.hs
@@ -31,6 +31,7 @@ module Network.BitTorrent.Internal
          -- * Swarm
        , SwarmSession(SwarmSession, torrentMeta, clientSession)
        , newLeacher, newSeeder
+       , enterSwarm, leaveSwarm , waitVacancy
 
          -- * Peer
        , PeerSession(PeerSession, connectedPeerAddr
@@ -83,7 +84,7 @@ import Data.Torrent
 import Network.BitTorrent.Extension
 import Network.BitTorrent.Peer
 import Network.BitTorrent.Exchange.Protocol as BT
-
+import Network.BitTorrent.Tracker.Protocol as BT
 
 
 -- | 'Progress' contains upload/download/left stats about
@@ -106,9 +107,10 @@ startProgress = Progress 0 0
 -----------------------------------------------------------------------}
 
 -- | In one application we could have many clients with difference
--- ID's and enabled extensions.
+-- ID's and different enabled extensions.
 data ClientSession = ClientSession {
-    -- | Our peer ID used in handshaked and discovery mechanism.
+    -- | Our peer ID used in handshaked and discovery mechanism. The
+    -- clientPeerID is unique 'ClientSession' identifier.
     clientPeerID      :: PeerID
 
     -- | Extensions we should try to use. Hovewer some particular peer
@@ -116,7 +118,14 @@ data ClientSession = ClientSession {
     -- 'PeerSession'.
   , allowedExtensions :: [Extension]
 
+    -- | Semaphor used to bound number of active P2P sessions.
+  , activeThreads     :: QSemN
+
+    -- | Max number of active connections.
+  , maxActive         :: Int
+
   , swarmSessions     :: TVar (Set SwarmSession)
+
   , eventManager      :: EventManager
   , currentProgress   :: TVar  Progress
   }
@@ -130,8 +139,11 @@ instance Ord ClientSession where
 getCurrentProgress :: MonadIO m => ClientSession -> m Progress
 getCurrentProgress = liftIO . readTVarIO . currentProgress
 
-newClient :: [Extension] -> IO ClientSession
-newClient exts = do
+newClient :: Int              -- ^ Maximum count of active P2P Sessions.
+          -> [Extension]      -- ^ Extensions allowed to use.
+          -> IO ClientSession
+
+newClient n exts = do
   mgr <- Ev.new
   -- TODO kill this thread when leave client
   _   <- forkIO $ loop mgr
@@ -139,6 +151,8 @@ newClient exts = do
   ClientSession
     <$> newPeerID
     <*> pure exts
+    <*> newQSemN n
+    <*> pure n
     <*> newTVarIO S.empty
     <*> pure mgr
     <*> newTVarIO (startProgress 0)
@@ -153,6 +167,10 @@ data SwarmSession = SwarmSession {
     torrentMeta       :: Torrent
   , clientSession     :: ClientSession
 
+    -- | Represent count of peers we _currently_ can connect to in the
+    -- swarm. Used to bound number of concurrent threads.
+  , vacantPeers       :: QSemN
+
     -- | Modify this carefully updating global progress.
   , clientBitfield    :: TVar  Bitfield
   , connectedPeers    :: TVar (Set PeerSession)
@@ -164,20 +182,28 @@ instance Eq SwarmSession where
 instance Ord SwarmSession where
   compare = comparing (tInfoHash . torrentMeta)
 
-newSwarmSession :: Bitfield -> ClientSession -> Torrent -> IO SwarmSession
-newSwarmSession bf cs @ ClientSession {..} t @ Torrent {..}
+newSwarmSession :: Int -> Bitfield -> ClientSession -> Torrent
+                -> IO SwarmSession
+newSwarmSession n bf cs @ ClientSession {..} t @ Torrent {..}
   = SwarmSession <$> pure t
                  <*> pure cs
+                 <*> newQSemN n
                  <*> newTVarIO bf
                  <*> newTVarIO S.empty
 
 newSeeder :: ClientSession -> Torrent -> IO SwarmSession
 newSeeder cs t @ Torrent {..}
-  = newSwarmSession (haveAll (pieceCount tInfo)) cs t
+  = newSwarmSession defSeederConns (haveAll (pieceCount tInfo)) cs t
 
 newLeacher :: ClientSession -> Torrent -> IO SwarmSession
 newLeacher cs t @ Torrent {..}
-  = newSwarmSession (haveNone (pieceCount tInfo)) cs t
+  = newSwarmSession defLeacherConns (haveNone (pieceCount tInfo)) cs t
+
+defSeederConns :: Int
+defSeederConns = defaultUnchokeSlots
+
+defLeacherConns :: Int
+defLeacherConns = defaultNumWant
 
 --isLeacher :: SwarmSession -> IO Bool
 --isLeacher = undefined
@@ -190,6 +216,22 @@ haveDone ix =
     writeTVar (have ix bf)
     currentProgress
 -}
+
+enterSwarm :: SwarmSession -> IO ()
+enterSwarm SwarmSession {..} = do
+  waitQSemN (activeThreads clientSession) 1
+  waitQSemN vacantPeers 1
+
+leaveSwarm :: SwarmSession -> IO ()
+leaveSwarm SwarmSession {..} = do
+  signalQSemN vacantPeers 1
+  signalQSemN (activeThreads clientSession) 1
+
+waitVacancy :: SwarmSession -> IO () -> IO ()
+waitVacancy se =
+  bracket (enterSwarm se) (const (leaveSwarm se))
+                  . const
+
 {-----------------------------------------------------------------------
     Peer session
 -----------------------------------------------------------------------}
@@ -263,16 +305,18 @@ sessionError msg
 
 -- TODO check if it connected yet peer
 withPeerSession :: SwarmSession -> PeerAddr
-                -> ((Socket, PeerSession) -> IO a)
-                -> IO a
+                -> ((Socket, PeerSession) -> IO ())
+                -> IO ()
 
 withPeerSession ss @ SwarmSession {..} addr
-    = bracket openSession closeSession
+    = handle isSessionException . bracket openSession closeSession
   where
     openSession = do
-      let caps = encodeExts $ allowedExtensions $ clientSession
-      let pid  = clientPeerID $ clientSession
-      let chs  = Handshake defaultBTProtocol caps (tInfoHash torrentMeta) pid
+      let caps  = encodeExts $ allowedExtensions $ clientSession
+      let ihash = tInfoHash torrentMeta
+      let pid   = clientPeerID $ clientSession
+      let chs   = Handshake defaultBTProtocol caps ihash pid
+
       sock <- connectToPeer addr
       phs  <- handshake sock chs `onException` close sock
 
@@ -292,7 +336,8 @@ withPeerSession ss @ SwarmSession {..} addr
            }
       return (sock, ps)
 
-    closeSession = close . fst
+    closeSession (sock, _) = do
+      close sock
 
 getPieceCount :: (MonadReader PeerSession m) => m PieceCount
 getPieceCount = asks (pieceCount . tInfo . torrentMeta . swarmSession)