5 files changed, 49 insertions, 40 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)
 
diff --git a/src/System/IO/MMap/Fixed.hs b/src/System/IO/MMap/Fixed.hs
index df6a6603..1e83c350 100644
--- a/src/System/IO/MMap/Fixed.hs
+++ b/src/System/IO/MMap/Fixed.hs
@@ -151,8 +151,8 @@ mallocTo fi s = do
 lookupRegion :: FixedOffset -> Fixed -> Maybe B.ByteString
 lookupRegion offset Fixed {..} =
   case intersecting imap $ IntervalCO offset (succ offset) of
-    [(i, (fptr, off))] -> let s = max 0 $ upperBound i - lowerBound i
-                          in  Just $ fromForeignPtr fptr off s
+    [(i, (fptr, off))] -> let s = upperBound i - lowerBound i
+                          in  Just $ fromForeignPtr fptr off (max 0 s)
     _         -> Nothing
 
 -- | Note: this is unsafe operation.
diff --git a/src/System/Torrent/Storage.hs b/src/System/Torrent/Storage.hs
index dd7258a0..c355d697 100644
--- a/src/System/Torrent/Storage.hs
+++ b/src/System/Torrent/Storage.hs
@@ -24,6 +24,7 @@ module System.Torrent.Storage
 
          -- * Construction
        , openStorage, closeStorage, withStorage
+       , getCompleteBitfield
 
          -- * Modification
        , getBlk, putBlk, selBlk
@@ -51,19 +52,20 @@ import Data.Bitfield as BF
 import Data.Torrent
 import Network.BitTorrent.Exchange.Protocol
 import System.IO.MMap.Fixed as Fixed
-import Debug.Trace
-
 
+-- TODO merge piece validation and Sessions.available into one transaction.
 data Storage = Storage {
     -- |
-    metainfo:: !Torrent
+    metainfo  :: !Torrent
 
-    -- |
-  , blocks  :: !(TVar Bitfield)
+    -- | Bitmask of complete and verified _pieces_.
+  , complete  :: !(TVar  Bitfield)
+
+    -- | Bitmask of complete _blocks_.
+  , blocks    :: !(TVar Bitfield)
     -- TODO use bytestring for fast serialization
     -- because we need to write this bitmap to disc periodically
 
-
   , blockSize :: !Int
 
     -- | Used to map linear block addresses to disjoint
@@ -76,19 +78,23 @@ ppStorage Storage {..} = pp <$> readTVarIO blocks
   where
     pp bf = int blockSize
 
+getCompleteBitfield :: Storage -> STM Bitfield
+getCompleteBitfield Storage {..} = readTVar complete
+
 {-----------------------------------------------------------------------
     Construction
 -----------------------------------------------------------------------}
 
 -- TODO doc args
-openStorage :: Torrent -> FilePath -> IO Storage
-openStorage t @ Torrent {..} contentPath = do
+openStorage :: Torrent -> FilePath -> Bitfield -> IO Storage
+openStorage t @ Torrent {..} contentPath bf = do
     let content_paths = contentLayout contentPath tInfo
     mapM_ (mkDir . fst)  content_paths
 
     let blockSize = defaultBlockSize `min` ciPieceLength tInfo
     print $ "content length " ++ show (contentLength tInfo)
-    Storage t <$> newTVarIO (haveNone (blockCount blockSize tInfo))
+    Storage t <$> newTVarIO bf
+              <*> newTVarIO (haveNone (blockCount blockSize tInfo))
               <*> pure blockSize
               <*> coalesceFiles content_paths
   where
@@ -103,8 +109,8 @@ closeStorage :: Storage -> IO ()
 closeStorage st = return ()
 
 
-withStorage :: Torrent -> FilePath -> (Storage -> IO a) -> IO a
-withStorage se path = bracket (openStorage se path) closeStorage
+withStorage :: Torrent -> FilePath -> Bitfield -> (Storage -> IO a) -> IO a
+withStorage se path bf = bracket (openStorage se path bf) closeStorage
 
 {-----------------------------------------------------------------------
     Modification
@@ -191,7 +197,9 @@ validatePiece pix st @ Storage {..} = {-# SCC validatePiece #-} do
     else do
       piece <- getPiece pix st
       if checkPiece (tInfo metainfo) pix piece
-        then return True
+        then do
+          atomically $ modifyTVar' complete (BF.have pix)
+          return True
         else do
           print $ "----------------------------- invalid " ++ show pix
 --          resetPiece pix st