summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--src/Network/BitTorrent/Exchange.hs54
-rw-r--r--src/Network/BitTorrent/Internal.hs15
2 files changed, 37 insertions, 32 deletions
diff --git a/src/Network/BitTorrent/Exchange.hs b/src/Network/BitTorrent/Exchange.hs
index 9b5a8535..0fd1d15a 100644
--- a/src/Network/BitTorrent/Exchange.hs
+++ b/src/Network/BitTorrent/Exchange.hs
@@ -11,6 +11,7 @@
11{-# LANGUAGE MultiParamTypeClasses #-} 11{-# LANGUAGE MultiParamTypeClasses #-}
12{-# LANGUAGE RecordWildCards #-} 12{-# LANGUAGE RecordWildCards #-}
13{-# LANGUAGE FlexibleContexts #-} 13{-# LANGUAGE FlexibleContexts #-}
14{-# LANGUAGE BangPatterns #-}
14module Network.BitTorrent.Exchange 15module Network.BitTorrent.Exchange
15 ( -- * Block 16 ( -- * Block
16 Block(..), BlockIx(..) 17 Block(..), BlockIx(..)
@@ -35,14 +36,17 @@ import Control.Concurrent
35import Control.Lens 36import Control.Lens
36import Control.Monad.Fork.Class 37import Control.Monad.Fork.Class
37import Control.Monad.Reader 38import Control.Monad.Reader
39import Control.Monad.State
38import Control.Monad.Trans.Resource 40import Control.Monad.Trans.Resource
39 41
42import Data.IORef
40import Data.Conduit as C 43import Data.Conduit as C
41import Data.Conduit.Cereal 44import Data.Conduit.Cereal
42import Data.Conduit.Network 45import Data.Conduit.Network
43import Data.Serialize as S 46import Data.Serialize as S
44import Text.PrettyPrint as PP hiding (($$)) 47import Text.PrettyPrint as PP hiding (($$))
45 48
49
46import Network 50import Network
47 51
48 52
@@ -75,10 +79,12 @@ runPeerWire sock p2p =
75awaitMessage :: P2P Message 79awaitMessage :: P2P Message
76awaitMessage = P2P (ReaderT (const go)) 80awaitMessage = P2P (ReaderT (const go))
77 where 81 where
78 go = await >>= maybe disconnect return 82 go = await >>= maybe (monadThrow PeerDisconnected) return
83{-# INLINE awaitMessage #-}
79 84
80yieldMessage :: Message -> P2P () 85yieldMessage :: Message -> P2P ()
81yieldMessage msg = P2P $ ReaderT $ \se -> C.yield msg 86yieldMessage msg = P2P $ ReaderT $ \se -> C.yield msg
87{-# INLINE yieldMessage #-}
82 88
83{----------------------------------------------------------------------- 89{-----------------------------------------------------------------------
84 P2P monad 90 P2P monad
@@ -97,6 +103,14 @@ newtype P2P a = P2P {
97 , MonadIO, MonadThrow, MonadActive 103 , MonadIO, MonadThrow, MonadActive
98 , MonadReader PeerSession 104 , MonadReader PeerSession
99 ) 105 )
106
107instance MonadState SessionState P2P where
108 {-# SPECIALIZE instance MonadState SessionState P2P #-}
109 get = asks sessionState >>= liftIO . readIORef
110 {-# INLINE get #-}
111 put !s = asks sessionState >>= \ref -> liftIO $ writeIORef ref s
112 {-# INLINE put #-}
113
100-- TODO instance for MonadFork 114-- TODO instance for MonadFork
101 115
102runSession :: SwarmSession -> PeerAddr -> P2P () -> IO () 116runSession :: SwarmSession -> PeerAddr -> P2P () -> IO ()
@@ -134,50 +148,50 @@ chainP2P :: SwarmSession -> PeerConnection -> P2P () -> IO ()
134-----------------------------------------------------------------------} 148-----------------------------------------------------------------------}
135 149
136-- | Terminate the current 'P2P' session. 150-- | Terminate the current 'P2P' session.
137disconnect :: MonadThrow m => m a 151disconnect :: P2P a
138disconnect = monadThrow PeerDisconnected 152disconnect = monadThrow PeerDisconnected
139 153
140-- TODO handle all protocol details here so we can hide this from 154-- TODO handle all protocol details here so we can hide this from
141-- public interface | 155-- public interface |
142protocolError :: MonadThrow m => Doc -> m a 156protocolError :: Doc -> P2P a
143protocolError = monadThrow . ProtocolError 157protocolError = monadThrow . ProtocolError
144 158
145{----------------------------------------------------------------------- 159{-----------------------------------------------------------------------
146 Helpers 160 Helpers
147-----------------------------------------------------------------------} 161-----------------------------------------------------------------------}
148 162
149-- | Count of client have pieces. 163getClientBF :: P2P Bitfield
150getHaveCount :: (MonadReader PeerSession m) => m PieceCount 164getClientBF = asks swarmSession >>= liftIO . getClientBitfield
151getHaveCount = undefined 165{-# INLINE getClientBF #-}
166
167-- | Count of client /have/ pieces.
168getHaveCount :: P2P PieceCount
169getHaveCount = haveCount <$> getClientBF
152{-# INLINE getHaveCount #-} 170{-# INLINE getHaveCount #-}
153 171
154-- | Count of client do not have pieces. 172-- | Count of client do not /have/ pieces.
155getWantCount :: (MonadReader PeerSession m) => m PieceCount 173getWantCount :: P2P PieceCount
156getWantCount = undefined 174getWantCount = totalCount <$> getClientBF
157{-# INLINE getWantCount #-} 175{-# INLINE getWantCount #-}
158 176
159-- | Count of both have and want pieces. 177-- | Count of both /have/ and /want/ pieces.
160getPieceCount :: (MonadReader PeerSession m) => m PieceCount 178getPieceCount :: P2P PieceCount
161getPieceCount = asks findPieceCount 179getPieceCount = asks findPieceCount
162{-# INLINE getPieceCount #-} 180{-# INLINE getPieceCount #-}
163 181
164-- for internal use only 182-- for internal use only
165emptyBF :: (MonadReader PeerSession m) => m Bitfield 183emptyBF :: P2P Bitfield
166emptyBF = liftM haveNone getPieceCount 184emptyBF = liftM haveNone getPieceCount
167 185
168fullBF :: (MonadReader PeerSession m) => m Bitfield 186fullBF :: P2P Bitfield
169fullBF = liftM haveAll getPieceCount 187fullBF = liftM haveAll getPieceCount
170 188
171singletonBF :: (MonadReader PeerSession m) => PieceIx -> m Bitfield 189singletonBF :: PieceIx -> P2P Bitfield
172singletonBF i = liftM (BF.singleton i) getPieceCount 190singletonBF i = liftM (BF.singleton i) getPieceCount
173 191
174adjustBF :: (MonadReader PeerSession m) => Bitfield -> m Bitfield 192adjustBF :: Bitfield -> P2P Bitfield
175adjustBF bf = (`adjustSize` bf) `liftM` getPieceCount 193adjustBF bf = (`adjustSize` bf) `liftM` getPieceCount
176 194
177getClientBF :: (MonadIO m, MonadReader PeerSession m) => m Bitfield
178getClientBF = asks swarmSession >>= liftIO . getClientBitfield
179
180
181 195
182peerWant :: P2P Bitfield 196peerWant :: P2P Bitfield
183peerWant = BF.difference <$> getClientBF <*> use bitfield 197peerWant = BF.difference <$> getClientBF <*> use bitfield
@@ -332,7 +346,7 @@ awaitEvent = awaitMessage >>= go
332 requireExtension ExtFast 346 requireExtension ExtFast
333 awaitEvent 347 awaitEvent
334 348
335 349-- TODO minimized number of peerOffer calls
336 350
337-- | 351-- |
338-- @ 352-- @
diff --git a/src/Network/BitTorrent/Internal.hs b/src/Network/BitTorrent/Internal.hs
index 8ce7afbf..5ac58d46 100644
--- a/src/Network/BitTorrent/Internal.hs
+++ b/src/Network/BitTorrent/Internal.hs
@@ -56,6 +56,7 @@ module Network.BitTorrent.Internal
56 -- * Peer 56 -- * Peer
57 , PeerSession( PeerSession, connectedPeerAddr 57 , PeerSession( PeerSession, connectedPeerAddr
58 , swarmSession, enabledExtensions 58 , swarmSession, enabledExtensions
59 , sessionState
59 ) 60 )
60 , SessionState 61 , SessionState
61 , withPeerSession 62 , withPeerSession
@@ -352,7 +353,8 @@ waitVacancy se =
352 Peer session 353 Peer session
353-----------------------------------------------------------------------} 354-----------------------------------------------------------------------}
354 355
355-- | Peer session contain all data necessary for peer to peer communication. 356-- | Peer session contain all data necessary for peer to peer
357-- communication.
356data PeerSession = PeerSession { 358data PeerSession = PeerSession {
357 -- | Used as unique 'PeerSession' identifier within one 359 -- | Used as unique 'PeerSession' identifier within one
358 -- 'SwarmSession'. 360 -- 'SwarmSession'.
@@ -409,17 +411,6 @@ instance Eq PeerSession where
409instance Ord PeerSession where 411instance Ord PeerSession where
410 compare = comparing connectedPeerAddr 412 compare = comparing connectedPeerAddr
411 413
412instance (MonadIO m, MonadReader PeerSession m)
413 => MonadState SessionState m where
414 get = do
415 ref <- asks sessionState
416 st <- liftIO (readIORef ref)
417 liftIO $ print (completeness (_bitfield st))
418 return st
419
420 put !s = asks sessionState >>= \ref -> liftIO $ writeIORef ref s
421
422
423-- | Exceptions used to interrupt the current P2P session. This 414-- | Exceptions used to interrupt the current P2P session. This
424-- exceptions will NOT affect other P2P sessions, DHT, peer <-> 415-- exceptions will NOT affect other P2P sessions, DHT, peer <->
425-- tracker, or any other session. 416-- tracker, or any other session.