5 files changed, 113 insertions, 98 deletions

diff --git a/src/Network/BitTorrent.hs b/src/Network/BitTorrent.hs
index 546c7644..24d78e85 100644
--- a/src/Network/BitTorrent.hs
+++ b/src/Network/BitTorrent.hs
@@ -16,6 +16,7 @@ module Network.BitTorrent
 
        , SwarmSession
        , newLeacher, newSeeder
+       , getSessionCount
 
          -- * Discovery
        , discover
@@ -63,16 +64,11 @@ discover swarm action = do
 
   progress <- getCurrentProgress (clientSession swarm)
 
-  putStrLn "lookup peers"
   withTracker progress conn $ \tses -> do
-    putStrLn "get peer list "
     forever $ do
       addr <- getPeerAddr tses
-      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 be9a455b..2d0393c0 100644
--- a/src/Network/BitTorrent/Exchange.hs
+++ b/src/Network/BitTorrent/Exchange.hs
@@ -61,28 +61,73 @@ runPeerWire :: Socket -> PeerWire () -> IO ()
 runPeerWire sock p2p =
   sourceSocket sock   $=
     conduitGet S.get  $=
-      forever p2p     $=
+      p2p             $=
     conduitPut S.put  $$
   sinkSocket sock
 
 awaitMessage :: P2P Message
-awaitMessage = P2P (ReaderT go)
+awaitMessage = P2P (ReaderT (const go))
   where
-    go _ = do
-      liftIO $ putStrLn "trying recv:"
-      mmsg <- await
-      case mmsg of
-        Nothing  -> monadThrow SessionException
-        Just msg -> do
---          liftIO $ updateIncoming se
-          liftIO $ print ("recv:" <+> ppMessage msg)
-          return msg
+    go = await >>= maybe disconnect return
 
 yieldMessage :: Message -> P2P ()
-yieldMessage msg = P2P $ ReaderT $ \se -> do
-  C.yield msg
-  liftIO $ print $ "sent:" <+> ppMessage msg
-  liftIO $ updateOutcoming se
+yieldMessage msg = P2P $ ReaderT $ \se -> C.yield msg
+
+{-----------------------------------------------------------------------
+    P2P monad
+-----------------------------------------------------------------------}
+
+-- |
+--   Exceptions:
+--
+--     * SessionException: is visible only within one peer
+--     session. Use this exception to terminate P2P session, but not
+--     the swarm session.
+--
+newtype P2P a = P2P {
+    unP2P :: ReaderT PeerSession PeerWire a
+  } deriving ( Functor, Applicative, Monad
+             , MonadIO, MonadThrow, MonadActive
+             , MonadReader PeerSession
+             )
+-- TODO instance for MonadFork
+
+runSession :: SwarmSession -> PeerAddr -> P2P () -> IO ()
+runSession  se addr p2p =
+    handle isIOException $
+      withPeerSession se addr $ \(sock, pses) -> do
+        runPeerWire sock (runReaderT (unP2P p2p) pses)
+  where
+    isIOException :: IOException -> IO ()
+    isIOException _ = return ()
+
+-- | 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
+
+{-----------------------------------------------------------------------
+    Exceptions
+-----------------------------------------------------------------------}
+
+disconnect :: MonadThrow m => m a
+disconnect = monadThrow PeerDisconnected
+
+protocolError :: MonadThrow m => Doc -> m a
+protocolError = monadThrow . ProtocolError
+
+{-----------------------------------------------------------------------
+    Helpers
+-----------------------------------------------------------------------}
 
 peerWant :: P2P Bitfield
 peerWant   = BF.difference <$> getClientBF  <*> use bitfield
@@ -116,7 +161,7 @@ requireExtension :: Extension -> P2P ()
 requireExtension required = do
   enabled <- asks enabledExtensions
   unless (required `elem` enabled) $
-    sessionError $ ppExtension required <+> "not enabled"
+    protocolError $ ppExtension required <+> "not enabled"
 
 --    haveMessage bf = do
 --      cbf <- undefined -- liftIO $ readIORef $ clientBitfield swarmSession
@@ -124,6 +169,9 @@ requireExtension required = do
 --        then nextEvent se
 --        else undefined  -- return $ Available diff
 
+{-----------------------------------------------------------------------
+    Exchange
+-----------------------------------------------------------------------}
 
 -- |
 --   +----------+---------+
@@ -266,41 +314,3 @@ yieldEvent (Fragment  blk) = do
 
 checkPiece :: PieceLIx -> {-ByteString -> -} P2P Bool
 checkPiece = undefined
-
-{-----------------------------------------------------------------------
-    P2P monad
------------------------------------------------------------------------}
-
--- |
---   Exceptions:
---
---     * SessionException: is visible only within one peer
---     session. Use this exception to terminate P2P session, but not
---     the swarm session.
---
-newtype P2P a = P2P {
-    unP2P :: ReaderT PeerSession PeerWire a
-  } deriving ( Functor, Applicative, Monad
-             , MonadIO, MonadThrow, MonadActive
-             , MonadReader PeerSession
-             )
--- TODO instance for MonadFork
-
-runSession :: SwarmSession -> PeerAddr -> P2P () -> IO ()
-runSession  se addr p2p =
-  withPeerSession se addr $ \(sock, pses) -> do
-    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/Exchange/Protocol.hs b/src/Network/BitTorrent/Exchange/Protocol.hs
index 5ea104cc..8d42e3a8 100644
--- a/src/Network/BitTorrent/Exchange/Protocol.hs
+++ b/src/Network/BitTorrent/Exchange/Protocol.hs
@@ -164,10 +164,8 @@ defaultHandshake = Handshake defaultBTProtocol defaultReserved
 -- | Handshaking with a peer specified by the second argument.
 handshake :: Socket -> Handshake -> IO Handshake
 handshake sock hs = do
-    putStrLn "send handshake"
     sendAll sock (S.encode hs)
 
-    putStrLn "recv handshake size"
     header <- recv sock 1
     when (B.length header == 0) $
       throw $ userError "Unable to receive handshake."
@@ -175,7 +173,6 @@ handshake sock hs = do
     let protocolLen = B.head header
     let restLen     = handshakeSize protocolLen - 1
 
-    putStrLn "recv handshake body"
     body <- recv sock restLen
     let resp = B.cons protocolLen body
 
@@ -432,9 +429,9 @@ ppMessage msg = text (show msg)
 
 -- |
 data PeerStatus = PeerStatus {
-    _choking    :: Bool
-  , _interested :: Bool
-  }
+    _choking    :: !Bool
+  , _interested :: !Bool
+  } deriving (Show, Eq)
 
 $(makeLenses ''PeerStatus)
 
@@ -443,9 +440,9 @@ instance Default PeerStatus where
 
 -- |
 data SessionStatus = SessionStatus {
-    _clientStatus :: PeerStatus
-  , _peerStatus   :: PeerStatus
-  }
+    _clientStatus :: !PeerStatus
+  , _peerStatus   :: !PeerStatus
+  } deriving (Show, Eq)
 
 $(makeLenses ''SessionStatus)
 
diff --git a/src/Network/BitTorrent/Internal.hs b/src/Network/BitTorrent/Internal.hs
index 918bfed7..db84e879 100644
--- a/src/Network/BitTorrent/Internal.hs
+++ b/src/Network/BitTorrent/Internal.hs
@@ -14,6 +14,7 @@
 --   data should be modified through standalone functions.
 --
 {-# LANGUAGE OverloadedStrings     #-}
+{-# LANGUAGE BangPatterns          #-}
 {-# LANGUAGE RecordWildCards       #-}
 {-# LANGUAGE TemplateHaskell       #-}
 {-# LANGUAGE DeriveDataTypeable    #-}
@@ -32,6 +33,7 @@ module Network.BitTorrent.Internal
 
          -- * Swarm
        , SwarmSession(SwarmSession, torrentMeta, clientSession)
+       , getSessionCount
        , newLeacher, newSeeder
        , enterSwarm, leaveSwarm , waitVacancy
 
@@ -46,7 +48,6 @@ module Network.BitTorrent.Internal
        , SessionException(..)
        , isSessionException
        , putSessionException
-       , sessionError
 
          -- ** Properties
        , bitfield, status
@@ -123,7 +124,7 @@ defaultThreadCount = 1000
 data ClientSession = ClientSession {
     -- | Our peer ID used in handshaked and discovery mechanism. The
     -- clientPeerID is unique 'ClientSession' identifier.
-    clientPeerID      :: PeerID
+    clientPeerID      :: !PeerID
 
     -- | Extensions we should try to use. Hovewer some particular peer
     -- might not support some extension, so we keep enableExtension in
@@ -186,24 +187,31 @@ defLeacherConns = defaultNumWant
 -- | Extensions are set globally by
 --   Swarm session are un
 data SwarmSession = SwarmSession {
-    torrentMeta       :: Torrent
-  , clientSession     :: ClientSession
+    torrentMeta       :: !Torrent
+  , clientSession     :: !ClientSession
 
     -- | Represent count of peers we _currently_ can connect to in the
     -- swarm. Used to bound number of concurrent threads.
-  , vacantPeers       :: MSem SessionCount
+  , vacantPeers       :: !(MSem SessionCount)
 
     -- | Modify this carefully updating global progress.
-  , clientBitfield    :: TVar  Bitfield
-  , connectedPeers    :: TVar (Set PeerSession)
+  , clientBitfield    :: !(TVar  Bitfield)
+  , connectedPeers    :: !(TVar (Set PeerSession))
   }
 
+-- INVARIANT:
+--   max_sessions_count - sizeof connectedPeers = value vacantPeers
+
 instance Eq SwarmSession where
   (==) = (==) `on` (tInfoHash . torrentMeta)
 
 instance Ord SwarmSession where
   compare = comparing (tInfoHash . torrentMeta)
 
+getSessionCount :: SwarmSession -> IO SessionCount
+getSessionCount SwarmSession {..} = do
+  S.size <$> readTVarIO connectedPeers
+
 newSwarmSession :: Int -> Bitfield -> ClientSession -> Torrent
                 -> IO SwarmSession
 newSwarmSession n bf cs @ ClientSession {..} t @ Torrent {..}
@@ -255,9 +263,9 @@ waitVacancy se =
 data PeerSession = PeerSession {
     -- | Used as unique 'PeerSession' identifier within one
     -- 'SwarmSession'.
-    connectedPeerAddr :: PeerAddr
+    connectedPeerAddr :: !PeerAddr
 
-  , swarmSession      :: SwarmSession
+  , swarmSession      :: !SwarmSession
 
     -- | Extensions such that both peer and client support.
   , enabledExtensions :: [Extension]
@@ -269,7 +277,7 @@ data PeerSession = PeerSession {
     --
     -- We should update timeout if we /receive/ any message within
     -- timeout interval to keep connection up.
-  , incomingTimeout     :: TimeoutKey
+  , incomingTimeout     :: !TimeoutKey
 
     -- | To send KA message appropriately we should know when was last
     -- time we sent a message to a peer. To do that we keep registered
@@ -279,17 +287,17 @@ data PeerSession = PeerSession {
     --
     -- We should update timeout if we /send/ any message within timeout
     -- to avoid reduntant KA messages.
-  , outcomingTimeout   :: TimeoutKey
+  , outcomingTimeout   :: !TimeoutKey
 
     -- TODO use dupChan for broadcasting
-  , broadcastMessages :: Chan   [Message]
-  , sessionState      :: IORef   SessionState
+  , broadcastMessages :: !(Chan   [Message])
+  , sessionState      :: !(IORef  SessionState)
   }
 
 data SessionState = SessionState {
-    _bitfield :: Bitfield
-  , _status   :: SessionStatus
-  }
+    _bitfield :: !Bitfield
+  , _status   :: !SessionStatus
+  } deriving (Show, Eq)
 
 $(makeLenses ''SessionState)
 
@@ -301,10 +309,16 @@ instance Ord PeerSession where
 
 instance (MonadIO m, MonadReader PeerSession m)
       => MonadState SessionState m where
-  get   = asks sessionState >>= liftIO . readIORef
-  put s = asks sessionState >>= \ref -> liftIO $ writeIORef ref s
+  get    = do
+    ref <- asks sessionState
+    st <- liftIO (readIORef ref)
+    liftIO $ print (completeness (_bitfield st))
+    return st
+
+  put !s = asks sessionState >>= \ref -> liftIO $ writeIORef ref s
 
-data SessionException = SessionException
+data SessionException = PeerDisconnected
+                      | ProtocolError Doc
                         deriving (Show, Typeable)
 
 instance Exception SessionException
@@ -315,10 +329,6 @@ isSessionException _ = return ()
 putSessionException :: SessionException -> IO ()
 putSessionException = print
 
-sessionError :: MonadIO m => Doc -> m ()
-sessionError msg
-  = liftIO $ throwIO $ userError $ render $ msg <+> "in session"
-
 -- TODO check if it connected yet peer
 withPeerSession :: SwarmSession -> PeerAddr
                 -> ((Socket, PeerSession) -> IO ())
@@ -342,7 +352,7 @@ withPeerSession ss @ SwarmSession {..} addr
       let enabled = decodeExts (enabledCaps caps (handshakeCaps phs))
       ps <- PeerSession addr ss enabled
          <$> registerTimeout (eventManager clientSession)
-                maxIncomingTime abortSession
+                maxIncomingTime (return ())
          <*> registerTimeout (eventManager clientSession)
                 maxOutcomingTime (sendKA sock)
          <*> newChan
@@ -350,9 +360,13 @@ withPeerSession ss @ SwarmSession {..} addr
            ; tc <- totalCount <$> readTVarIO clientBitfield
            ; newIORef (SessionState (haveNone tc) def)
            }
+
+      atomically $ modifyTVar' connectedPeers (S.insert ps)
+
       return (sock, ps)
 
-    closeSession (sock, _) = do
+    closeSession (sock, ps) = do
+      atomically $ modifyTVar' connectedPeers (S.delete ps)
       close sock
 
 getPieceCount :: (MonadReader PeerSession m) => m PieceCount
@@ -411,6 +425,3 @@ sendKA sock {- SwarmSession {..} -} = do
 --  let mgr = eventManager clientSession
 --  updateTimeout mgr
 --  print "Done.."
-
-abortSession :: IO ()
-abortSession = error "abortSession: not implemented"
diff --git a/src/Network/BitTorrent/Tracker/Protocol.hs b/src/Network/BitTorrent/Tracker/Protocol.hs
index f7599ba2..d4f3da3c 100644
--- a/src/Network/BitTorrent/Tracker/Protocol.hs
+++ b/src/Network/BitTorrent/Tracker/Protocol.hs
@@ -212,10 +212,11 @@ defaultPorts = [6881..6889]
 --   in fact only actively forms new connections if it has less than
 --   30 peers and will refuse connections if it has 55.
 --
---   So the default value is set to 25.
+--   So the default value is set to 50 because usually 30-50% of peers
+--   are not responding.
 --
 defaultNumWant :: Int
-defaultNumWant = 25
+defaultNumWant = 50
 
 mkHTTPRequest :: URI -> Request ByteString
 mkHTTPRequest uri = Request uri GET [] ""