From 3867719780293528e604452818b9d9a616938783 Mon Sep 17 00:00:00 2001 From: Sam Truzjan Date: Tue, 8 Apr 2014 03:56:29 +0400 Subject: Move bitfield to exchange subsystem --- src/Network/BitTorrent/Exchange/Bitfield.hs | 324 ++++++++++++++++++++++ src/Network/BitTorrent/Exchange/Connection.hs | 4 +- src/Network/BitTorrent/Exchange/Message.hs | 2 +- src/Network/BitTorrent/Exchange/Selection.hs | 2 +- src/Network/BitTorrent/Exchange/Session.hs | 2 +- src/Network/BitTorrent/Exchange/Session/Status.hs | 2 +- 6 files changed, 330 insertions(+), 6 deletions(-) create mode 100644 src/Network/BitTorrent/Exchange/Bitfield.hs (limited to 'src/Network/BitTorrent/Exchange') diff --git a/src/Network/BitTorrent/Exchange/Bitfield.hs b/src/Network/BitTorrent/Exchange/Bitfield.hs new file mode 100644 index 00000000..3f4931f3 --- /dev/null +++ b/src/Network/BitTorrent/Exchange/Bitfield.hs @@ -0,0 +1,324 @@ +-- | +-- Copyright : (c) Sam Truzjan 2013 +-- License : BSD3 +-- Maintainer : pxqr.sta@gmail.com +-- Stability : experimental +-- Portability : portable +-- +-- This modules provides all necessary machinery to work with +-- bitfields. Bitfields are used to keep track indices of complete +-- pieces either peer have or client have. +-- +-- There are also commonly used piece seletion algorithms +-- which used to find out which one next piece to download. +-- Selectors considered to be used in the following order: +-- +-- * Random first - at the start. +-- +-- * Rarest first selection - performed to avoid situation when +-- rarest piece is unaccessible. +-- +-- * /End game/ seletion - performed after a peer has requested all +-- the subpieces of the content. +-- +-- Note that BitTorrent applies the strict priority policy for +-- /subpiece/ or /blocks/ selection. +-- +{-# LANGUAGE CPP #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE RecordWildCards #-} +module Network.BitTorrent.Exchange.Bitfield + ( -- * Bitfield + PieceIx + , PieceCount + , Bitfield + + -- * Construction + , haveAll + , haveNone + , have + , singleton + , interval + , adjustSize + + -- * Query + -- ** Cardinality + , Network.BitTorrent.Exchange.Bitfield.null + , Network.BitTorrent.Exchange.Bitfield.full + , haveCount + , totalCount + , completeness + + -- ** Membership + , member + , notMember + , findMin + , findMax + , isSubsetOf + + -- ** Availability + , complement + , Frequency + , frequencies + , rarest + + -- * Combine + , insert + , union + , intersection + , difference + + -- * Conversion + , toList + , fromList + + -- * Serialization + , fromBitmap + , toBitmap + ) where + +import Control.Monad +import Control.Monad.ST +import Data.ByteString (ByteString) +import qualified Data.ByteString as B +import qualified Data.ByteString.Lazy as Lazy +import Data.Vector.Unboxed (Vector) +import qualified Data.Vector.Unboxed as V +import qualified Data.Vector.Unboxed.Mutable as VM +import Data.IntervalSet (IntSet) +import qualified Data.IntervalSet as S +import qualified Data.IntervalSet.ByteString as S +import Data.List (foldl') +import Data.Monoid +import Data.Ratio + +import Data.Torrent + +-- TODO cache some operations + +-- | Bitfields are represented just as integer sets but with +-- restriction: the each set should be within given interval (or +-- subset of the specified interval). Size is used to specify +-- interval, so bitfield of size 10 might contain only indices in +-- interval [0..9]. +-- +data Bitfield = Bitfield { + bfSize :: !PieceCount + , bfSet :: !IntSet + } deriving (Show, Read, Eq) + +-- Invariants: all elements of bfSet lie in [0..bfSize - 1]; + +instance Monoid Bitfield where + {-# SPECIALIZE instance Monoid Bitfield #-} + mempty = haveNone 0 + mappend = union + mconcat = unions + +{----------------------------------------------------------------------- + Construction +-----------------------------------------------------------------------} + +-- | The empty bitfield of the given size. +haveNone :: PieceCount -> Bitfield +haveNone s = Bitfield s S.empty + +-- | The full bitfield containing all piece indices for the given size. +haveAll :: PieceCount -> Bitfield +haveAll s = Bitfield s (S.interval 0 (s - 1)) + +-- | Insert the index in the set ignoring out of range indices. +have :: PieceIx -> Bitfield -> Bitfield +have ix Bitfield {..} + | 0 <= ix && ix < bfSize = Bitfield bfSize (S.insert ix bfSet) + | otherwise = Bitfield bfSize bfSet + +singleton :: PieceIx -> PieceCount -> Bitfield +singleton ix pc = have ix (haveNone pc) + +-- | Assign new size to bitfield. FIXME Normally, size should be only +-- decreased, otherwise exception raised. +adjustSize :: PieceCount -> Bitfield -> Bitfield +adjustSize s Bitfield {..} = Bitfield s bfSet + +-- | NOTE: for internal use only +interval :: PieceCount -> PieceIx -> PieceIx -> Bitfield +interval pc a b = Bitfield pc (S.interval a b) + +{----------------------------------------------------------------------- + Query +-----------------------------------------------------------------------} + +-- | Test if bitifield have no one index: peer do not have anything. +null :: Bitfield -> Bool +null Bitfield {..} = S.null bfSet + +-- | Test if bitfield have all pieces. +full :: Bitfield -> Bool +full Bitfield {..} = S.size bfSet == bfSize + +-- | Count of peer have pieces. +haveCount :: Bitfield -> PieceCount +haveCount = S.size . bfSet + +-- | Total count of pieces and its indices. +totalCount :: Bitfield -> PieceCount +totalCount = bfSize + +-- | Ratio of /have/ piece count to the /total/ piece count. +-- +-- > forall bf. 0 <= completeness bf <= 1 +-- +completeness :: Bitfield -> Ratio PieceCount +completeness b = haveCount b % totalCount b + +inRange :: PieceIx -> Bitfield -> Bool +inRange ix Bitfield {..} = 0 <= ix && ix < bfSize + +member :: PieceIx -> Bitfield -> Bool +member ix bf @ Bitfield {..} + | ix `inRange` bf = ix `S.member` bfSet + | otherwise = False + +notMember :: PieceIx -> Bitfield -> Bool +notMember ix bf @ Bitfield {..} + | ix `inRange` bf = ix `S.notMember` bfSet + | otherwise = True + +-- | Find first available piece index. +findMin :: Bitfield -> PieceIx +findMin = S.findMin . bfSet +{-# INLINE findMin #-} + +-- | Find last available piece index. +findMax :: Bitfield -> PieceIx +findMax = S.findMax . bfSet +{-# INLINE findMax #-} + +-- | Check if all pieces from first bitfield present if the second bitfield +isSubsetOf :: Bitfield -> Bitfield -> Bool +isSubsetOf a b = bfSet a `S.isSubsetOf` bfSet b +{-# INLINE isSubsetOf #-} + +-- | Resulting bitfield includes only missing pieces. +complement :: Bitfield -> Bitfield +complement Bitfield {..} = Bitfield + { bfSet = uni `S.difference` bfSet + , bfSize = bfSize + } + where + Bitfield _ uni = haveAll bfSize +{-# INLINE complement #-} + +{----------------------------------------------------------------------- +-- Availability +-----------------------------------------------------------------------} + +-- | Frequencies are needed in piece selection startegies which use +-- availability quantity to find out the optimal next piece index to +-- download. +type Frequency = Int + +-- TODO rename to availability +-- | How many times each piece index occur in the given bitfield set. +frequencies :: [Bitfield] -> Vector Frequency +frequencies [] = V.fromList [] +frequencies xs = runST $ do + v <- VM.new size + VM.set v 0 + forM_ xs $ \ Bitfield {..} -> do + forM_ (S.toList bfSet) $ \ x -> do + fr <- VM.read v x + VM.write v x (succ fr) + V.unsafeFreeze v + where + size = maximum (map bfSize xs) + +-- TODO it seems like this operation is veeery slow + +-- | Find least available piece index. If no piece available return +-- 'Nothing'. +rarest :: [Bitfield] -> Maybe PieceIx +rarest xs + | V.null freqMap = Nothing + | otherwise + = Just $ fst $ V.ifoldr' minIx (0, freqMap V.! 0) freqMap + where + freqMap = frequencies xs + + minIx :: PieceIx -> Frequency + -> (PieceIx, Frequency) + -> (PieceIx, Frequency) + minIx ix fr acc@(_, fra) + | fr < fra && fr > 0 = (ix, fr) + | otherwise = acc + + +{----------------------------------------------------------------------- + Combine +-----------------------------------------------------------------------} + +insert :: PieceIx -> Bitfield -> Bitfield +insert pix bf @ Bitfield {..} + | 0 <= pix && pix < bfSize = Bitfield + { bfSet = S.insert pix bfSet + , bfSize = bfSize + } + | otherwise = bf + +-- | Find indices at least one peer have. +union :: Bitfield -> Bitfield -> Bitfield +union a b = {-# SCC union #-} Bitfield { + bfSize = bfSize a `max` bfSize b + , bfSet = bfSet a `S.union` bfSet b + } + +-- | Find indices both peers have. +intersection :: Bitfield -> Bitfield -> Bitfield +intersection a b = {-# SCC intersection #-} Bitfield { + bfSize = bfSize a `min` bfSize b + , bfSet = bfSet a `S.intersection` bfSet b + } + +-- | Find indices which have first peer but do not have the second peer. +difference :: Bitfield -> Bitfield -> Bitfield +difference a b = {-# SCC difference #-} Bitfield { + bfSize = bfSize a -- FIXME is it reasonable? + , bfSet = bfSet a `S.difference` bfSet b + } + +-- | Find indices the any of the peers have. +unions :: [Bitfield] -> Bitfield +unions = {-# SCC unions #-} foldl' union (haveNone 0) + +{----------------------------------------------------------------------- + Serialization +-----------------------------------------------------------------------} + +-- | List all /have/ indexes. +toList :: Bitfield -> [PieceIx] +toList Bitfield {..} = S.toList bfSet + +-- | Make bitfield from list of /have/ indexes. +fromList :: PieceCount -> [PieceIx] -> Bitfield +fromList s ixs = Bitfield { + bfSize = s + , bfSet = S.splitGT (-1) $ S.splitLT s $ S.fromList ixs + } + +-- | Unpack 'Bitfield' from tightly packed bit array. Note resulting +-- size might be more than real bitfield size, use 'adjustSize'. +fromBitmap :: ByteString -> Bitfield +fromBitmap bs = {-# SCC fromBitmap #-} Bitfield { + bfSize = B.length bs * 8 + , bfSet = S.fromByteString bs + } +{-# INLINE fromBitmap #-} + +-- | Pack a 'Bitfield' to tightly packed bit array. +toBitmap :: Bitfield -> Lazy.ByteString +toBitmap Bitfield {..} = {-# SCC toBitmap #-} Lazy.fromChunks [intsetBM, alignment] + where + byteSize = bfSize `div` 8 + if bfSize `mod` 8 == 0 then 0 else 1 + alignment = B.replicate (byteSize - B.length intsetBM) 0 + intsetBM = S.toByteString bfSet diff --git a/src/Network/BitTorrent/Exchange/Connection.hs b/src/Network/BitTorrent/Exchange/Connection.hs index 9b7942ae..f208fa54 100644 --- a/src/Network/BitTorrent/Exchange/Connection.hs +++ b/src/Network/BitTorrent/Exchange/Connection.hs @@ -135,10 +135,10 @@ import Text.Show.Functions () import System.Log.FastLogger (ToLogStr(..)) import System.Timeout -import Data.Torrent.Bitfield as BF import Data.Torrent import Network.BitTorrent.Address -import Network.BitTorrent.Exchange.Message as Msg +import Network.BitTorrent.Exchange.Bitfield as BF +import Network.BitTorrent.Exchange.Message as Msg -- TODO handle port message? -- TODO handle limits? diff --git a/src/Network/BitTorrent/Exchange/Message.hs b/src/Network/BitTorrent/Exchange/Message.hs index a0cb5c91..f8b76186 100644 --- a/src/Network/BitTorrent/Exchange/Message.hs +++ b/src/Network/BitTorrent/Exchange/Message.hs @@ -117,10 +117,10 @@ import Network.Socket hiding (KeepAlive) import Text.PrettyPrint as PP hiding ((<>)) import Text.PrettyPrint.Class -import Data.Torrent.Bitfield import Data.Torrent hiding (Piece (..)) import qualified Data.Torrent as P (Piece (..)) import Network.BitTorrent.Address +import Network.BitTorrent.Exchange.Bitfield import Network.BitTorrent.Exchange.Block {----------------------------------------------------------------------- diff --git a/src/Network/BitTorrent/Exchange/Selection.hs b/src/Network/BitTorrent/Exchange/Selection.hs index 2724fabc..3701450b 100644 --- a/src/Network/BitTorrent/Exchange/Selection.hs +++ b/src/Network/BitTorrent/Exchange/Selection.hs @@ -22,7 +22,7 @@ module Network.BitTorrent.Exchange.Selection import Data.Ratio -import Data.Torrent.Bitfield +import Network.BitTorrent.Exchange.Bitfield type Selector = Bitfield -- ^ Indices of client /have/ pieces. diff --git a/src/Network/BitTorrent/Exchange/Session.hs b/src/Network/BitTorrent/Exchange/Session.hs index b68f17a0..4c6811d9 100644 --- a/src/Network/BitTorrent/Exchange/Session.hs +++ b/src/Network/BitTorrent/Exchange/Session.hs @@ -46,9 +46,9 @@ import System.Log.FastLogger (LogStr, ToLogStr (..)) import Data.BEncode as BE import Data.Torrent as Torrent -import Data.Torrent.Bitfield as BF import Network.BitTorrent.Internal.Types import Network.BitTorrent.Address +import Network.BitTorrent.Exchange.Bitfield as BF import Network.BitTorrent.Exchange.Block as Block import Network.BitTorrent.Exchange.Connection import Network.BitTorrent.Exchange.Message as Message diff --git a/src/Network/BitTorrent/Exchange/Session/Status.hs b/src/Network/BitTorrent/Exchange/Session/Status.hs index 63b91926..af3e94f5 100644 --- a/src/Network/BitTorrent/Exchange/Session/Status.hs +++ b/src/Network/BitTorrent/Exchange/Session/Status.hs @@ -29,7 +29,7 @@ import Data.Set as S import Data.Tuple import Data.Torrent -import Data.Torrent.Bitfield as BF +import Network.BitTorrent.Exchange.Bitfield as BF import Network.BitTorrent.Address import Network.BitTorrent.Exchange.Block as Block import System.Torrent.Storage (Storage, writePiece) -- cgit v1.2.3