module Process.PeerMgr (
-- * Types
Peer(..)
, PeerMgrMsg(..)
, PeerMgrChannel
-- * Interface
, start
)
where
import qualified Data.Map as M
import Control.Concurrent
import Control.Concurrent.CML
import Control.Monad.State
import Control.Monad.Reader
import Network
import System.IO
import System.Log.Logger
import PeerTypes
import Process
import Process.Peer as Peer
import Process.ChokeMgr hiding (start)
import Process.FS hiding (start)
import Process.PieceMgr hiding (start)
import Process.Status hiding (start)
import Protocol.Wire
import Supervisor
import Torrent hiding (infoHash)
data PeerMgrMsg = PeersFromTracker [Peer]
| NewIncoming (Handle, HostName, PortNumber)
type PeerMgrChannel = Channel PeerMgrMsg
data CF = CF { peerCh :: PeerMgrChannel
, pieceMgrCh :: PieceMgrChannel
, mgrCh :: Channel MgrMessage
, fsCh :: FSPChannel
, peerPool :: SupervisorChan
, chokeMgrCh :: ChokeMgrChannel
}
instance Logging CF where
logName _ = "Process.PeerMgr"
data ST = ST { peersInQueue :: [Peer]
, peers :: M.Map ThreadId (Channel PeerMessage)
, peerId :: PeerId
, infoHash :: InfoHash
}
start :: PeerMgrChannel -> PeerId -> InfoHash -> PieceMap -> PieceMgrChannel -> FSPChannel
-> ChokeMgrChannel -> StatusChan -> Int -> SupervisorChan
-> IO ThreadId
start ch pid ih pm pieceMgrC fsC chokeMgrC statC nPieces supC =
do mgrC <- channel
fakeChan <- channel
pool <- liftM snd $ oneForOne "PeerPool" [] fakeChan
spawnP (CF ch pieceMgrC mgrC fsC pool chokeMgrC)
(ST [] M.empty pid ih) (catchP (forever lp)
(defaultStopHandler supC))
where
lp = do chooseP [incomingPeers, peerEvent] >>= syncP
fillPeers
incomingPeers =
recvWrapPC peerCh (\msg ->
case msg of
PeersFromTracker ps -> do
debugP "Adding peers to queue"
modify (\s -> s { peersInQueue = ps ++ peersInQueue s })
NewIncoming conn@(h, _, _) -> do
sz <- liftM M.size $ gets peers
if sz < numPeers
then do debugP "New incoming peer, handling"
addIncoming conn
return ()
else do debugP "Already too many peers, closing!"
liftIO $ hClose h)
peerEvent =
recvWrapPC mgrCh (\msg -> case msg of
Connect tid c -> newPeer tid c
Disconnect tid -> removePeer tid)
newPeer tid c = do debugP $ "Adding new peer " ++ show tid
sendPC chokeMgrCh (AddPeer tid c) >>= syncP
modify (\s -> s { peers = M.insert tid c (peers s)})
removePeer tid = do debugP $ "Removing peer " ++ show tid
sendPC chokeMgrCh (RemovePeer tid) >>= syncP
modify (\s -> s { peers = M.delete tid (peers s)})
numPeers = 40
fillPeers = do
sz <- liftM M.size $ gets peers
when (sz < numPeers)
(do q <- gets peersInQueue
let (toAdd, rest) = splitAt (numPeers - sz) q
debugP $ "Filling with up to " ++ show (numPeers - sz) ++ " peers"
mapM_ addPeer toAdd
modify (\s -> s { peersInQueue = rest }))
addPeer (Peer hn prt) = do
pid <- gets peerId
ih <- gets infoHash
pool <- asks peerPool
pmC <- asks pieceMgrCh
fsC <- asks fsCh
mgrC <- asks mgrCh
liftIO $ connect (hn, prt, pid, ih, pm) pool pmC fsC statC mgrC nPieces
addIncoming conn = do
pid <- gets peerId
ih <- gets infoHash
pool <- asks peerPool
pieceMgrC <- asks pieceMgrCh
fsC <- asks fsCh
mgrC <- asks mgrCh
let ihTst = (== ih)
liftIO $ acceptor conn ih ihTst pool pid mgrC pieceMgrC fsC statC pm nPieces
type ConnectRecord = (HostName, PortID, PeerId, InfoHash, PieceMap)
connect :: ConnectRecord -> SupervisorChan -> PieceMgrChannel -> FSPChannel
-> StatusChan
-> MgrChannel -> Int
-> IO ThreadId
connect (host, port, pid, ih, pm) pool pieceMgrC fsC statC mgrC nPieces =
spawn (connector >> return ())
where connector =
do debugM "Process.PeerMgr.connect" $
"Connecting to " ++ show host ++ " (" ++ showPort port ++ ")"
h <- connectTo host port
debugM "Process.PeerMgr.connect" "Connected, initiating handShake"
r <- initiateHandshake h pid ih
debugM "Process.PeerMgr.connect" "Handshake run"
case r of
Left err -> do debugM "Process.PeerMgr.connect"
("Peer handshake failure at host " ++ host
++ " with error " ++ err)
return ()
Right (_caps, _rpid) ->
do debugM "Process.PeerMgr.connect" "entering peerP loop code"
supC <- channel -- TODO: Should be linked later on
children <- peerChildren h mgrC pieceMgrC fsC statC pm nPieces
sync $ transmit pool $ SpawnNew (Supervisor $ allForOne "PeerSup" children)
return ()
acceptor (h,hn,pn) ih ihTst pool pid mgrC pieceMgrC fsC statC pm nPieces =
spawn (connector >> return ())
where connector = do
debugLog "Handling incoming connection"
r <- receiveHandshake h pid ihTst ih
debugLog "RecvHandshake run"
case r of
Left err -> do debugLog ("Incoming Peer handshake failure with " ++ show hn ++ "("
++ show pn ++ "), error: " ++ err)
return ()
Right (_caps, _rpid) ->
do debugLog "entering peerP loop code"
supC <- channel -- TODO: Should be linked later on
children <- peerChildren h mgrC pieceMgrC fsC statC pm nPieces
sync $ transmit pool $ SpawnNew (Supervisor $ allForOne "PeerSup" children)
return ()
debugLog = debugM "Process.PeerMgr.acceptor"
showPort :: PortID -> String
showPort (PortNumber pn) = show pn
showPort _ = "N/A"