3 files changed, 27 insertions, 26 deletions

diff --git a/src/Network/BitTorrent/Exchange.hs b/src/Network/BitTorrent/Exchange.hs
index dc1b2752..71be3f88 100644
--- a/src/Network/BitTorrent/Exchange.hs
+++ b/src/Network/BitTorrent/Exchange.hs
@@ -71,6 +71,7 @@ module Network.BitTorrent.Exchange
        ) where
 
 import Control.Applicative
+import Control.Concurrent.STM
 import Control.Exception
 import Control.Lens
 import Control.Monad.Reader
@@ -80,7 +81,7 @@ import Control.Monad.Trans.Resource
 import Data.IORef
 import Data.Conduit as C
 import Data.Conduit.Cereal as S
-import Data.Conduit.Serialization.Binary as B
+--import Data.Conduit.Serialization.Binary as B
 import Data.Conduit.Network
 import Data.Serialize as S
 import Text.PrettyPrint as PP hiding (($$))
@@ -100,11 +101,11 @@ import System.Torrent.Storage
 type PeerWire = ConduitM Message Message IO
 
 runPeerWire :: Socket -> PeerWire () -> IO ()
-runPeerWire sock p2p =
+runPeerWire sock action =
   sourceSocket sock     $=
     S.conduitGet S.get  $=
 --    B.conduitDecode     $=
-      p2p               $=
+      action            $=
     S.conduitPut S.put  $$
 --    B.conduitEncode     $$
   sinkSocket sock
@@ -153,9 +154,9 @@ instance MonadState SessionState P2P where
   {-# INLINE put #-}
 
 runP2P :: (Socket, PeerSession) -> P2P () -> IO ()
-runP2P (sock, ses) p2p =
+runP2P (sock, ses) action =
   handle isIOException $
-    runPeerWire sock (runReaderT (unP2P p2p) ses)
+    runPeerWire sock (runReaderT (unP2P action) ses)
   where
     isIOException :: IOException -> IO ()
     isIOException _ = return ()
@@ -428,7 +429,10 @@ yieldEvent e = {-# SCC yieldEvent #-} do
     go e
     flushPending
   where
-    go (Available ixs) = asks swarmSession >>= liftIO . available ixs
+    go (Available ixs) = do
+      ses <- asks swarmSession
+      liftIO $ atomically $ available ixs ses
+
     go (Want      bix) = do
       offer <- peerOffer
       if ixPiece bix `BF.member` offer
diff --git a/src/Network/BitTorrent/Sessions.hs b/src/Network/BitTorrent/Sessions.hs
index 43a34df9..2118ccf0 100644
--- a/src/Network/BitTorrent/Sessions.hs
+++ b/src/Network/BitTorrent/Sessions.hs
@@ -202,9 +202,9 @@ discover swarm @ SwarmSession {..} = {-# SCC discover #-} do
     forever $ do
       addr <- getPeerAddr tses
       forkThrottle swarm $ do
-        initiatePeerSession swarm addr $ \conn -> do
+        initiatePeerSession swarm addr $ \pconn -> do
           print addr
-          runP2P conn p2p
+          runP2P pconn p2p
 
 registerSwarmSession :: SwarmSession -> STM ()
 registerSwarmSession ss @ SwarmSession {..} =
@@ -223,8 +223,7 @@ openSwarmSession cs @ ClientSession {..} loc @ TorrentLoc {..} = do
 
   ss <- SwarmSession t cs
     <$> MSem.new defLeecherConns
-    <*> newTVarIO bf
-    <*> openStorage t dataDirPath
+    <*> openStorage t dataDirPath bf
     <*> newTVarIO S.empty
     <*> newBroadcastTChanIO
 
@@ -232,7 +231,7 @@ openSwarmSession cs @ ClientSession {..} loc @ TorrentLoc {..} = do
     modifyTVar' currentProgress $ enqueuedProgress $ contentLength $ tInfo t
     registerSwarmSession ss
 
-  forkIO $ discover ss
+  _ <- forkIO $ discover ss
 
   return ss
 
@@ -246,8 +245,8 @@ closeSwarmSession se @ SwarmSession {..} = do
 
 getSwarm :: ClientSession -> InfoHash -> IO SwarmSession
 getSwarm cs @ ClientSession {..} ih = do
-  status <- torrentPresence cs ih
-  case status of
+  tstatus <- torrentPresence cs ih
+  case tstatus of
     Unknown        -> throw $ UnknownTorrent ih
     Active sw      -> return sw
     Registered loc -> openSwarmSession cs loc
@@ -342,9 +341,11 @@ openSession :: SwarmSession -> PeerAddr -> Handshake -> IO PeerSession
 openSession ss @ SwarmSession {..} addr Handshake {..} = do
   let clientCaps = encodeExts $ allowedExtensions $ clientSession
   let enabled    = decodeExts (enabledCaps clientCaps hsReserved)
+
+  bf <- getClientBitfield ss
   ps <- PeerSession addr ss enabled
     <$> atomically (dupTChan broadcastMessages)
-    <*> (newIORef . initialSessionState . totalCount =<< readTVarIO clientBitfield)
+    <*> newIORef (initialSessionState (totalCount bf))
     -- TODO we could implement more interesting throtling scheme
     -- using connected peer information
   registerPeerSession ps
@@ -408,7 +409,7 @@ listener cs action serverPort = bracket openListener close loop
       putStrLn "accepted"
       case addr of
         SockAddrInet port host -> do
-          forkIO $ do
+          _ <- forkIO $ do
             acceptPeerSession cs (PeerAddr Nothing host port) conn action
           return ()
         _                      -> return ()
diff --git a/src/Network/BitTorrent/Sessions/Types.lhs b/src/Network/BitTorrent/Sessions/Types.lhs
index 69411d4e..af2a6755 100644
--- a/src/Network/BitTorrent/Sessions/Types.lhs
+++ b/src/Network/BitTorrent/Sessions/Types.lhs
@@ -299,7 +299,7 @@ So if client is a leecher then max sessions count depends on the
 number of unchoke slots.
 
 > -- | Used to bound the number of simultaneous connections and, which
-> -- is the same, P2P sessions within the swarm session.
+> --   is the same, P2P sessions within the swarm session.
 > type SessionCount = Int
 
 However if client is a seeder then the value depends on .
@@ -322,8 +322,6 @@ Throttling* section.
 Client bitfield is used to keep track "the client have" piece set.
 Modify this carefully always updating global progress.
 
->   , clientBitfield    :: !(TVar  Bitfield)
-
 >   , storage           :: !Storage
 
 We keep set of the all connected peers for the each particular torrent
@@ -371,7 +369,7 @@ INVARIANT: max_sessions_count - sizeof connectedPeers = value vacantPeers
 >   compare = comparing (tInfoHash . torrentMeta)
 
 > getClientBitfield :: SwarmSession -> IO Bitfield
-> getClientBitfield = readTVarIO . clientBitfield
+> getClientBitfield SwarmSession {..} = atomically $ getCompleteBitfield storage
 
 Peer sessions
 ------------------------------------------------------------------------
@@ -458,16 +456,14 @@ messages & events we should send.
 2. Update downloaded stats     --/
 3. Signal to the all other peer about this.
 
-> available :: Bitfield -> SwarmSession -> IO ()
-> available bf se @ SwarmSession {..} = {-# SCC available #-} do
->     mark >> atomically broadcast
+> available :: Bitfield -> SwarmSession -> STM ()
+> available bf SwarmSession {..} = {-# SCC available #-} do
+>     updateProgress >> broadcast
 >   where
->     mark = do
+>     updateProgress = do
 >       let piLen = ciPieceLength $ tInfo $ torrentMeta
 >       let bytes = piLen * BF.haveCount bf
->       atomically $ do
->         modifyTVar' clientBitfield (BF.union bf)
->         modifyTVar' (currentProgress clientSession) (downloadedProgress bytes)
+>       modifyTVar' (currentProgress clientSession) (downloadedProgress bytes)
 >
 >     broadcast = mapM_ (writeTChan broadcastMessages . Have) (BF.toList bf)