-- | Peer proceeses
{-# LANGUAGE ScopedTypeVariables #-}
module Process.Peer (
-- * Types
PeerMessage(..)
-- * Interface
, peerChildren
)
where
import Control.Concurrent
import Control.Concurrent.CML
import Control.Monad.State
import Control.Monad.Reader
import Prelude hiding (catch, log)
import Data.Bits
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L
import qualified Data.Serialize.Get as G
import qualified Data.Map as M
import qualified Data.IntSet as IS
import Data.Maybe
import Data.Set as S hiding (map)
import Data.Time.Clock
import Data.Word
import System.IO
import System.Timeout
import PeerTypes
import Process
import Process.FS
import Process.PieceMgr
import qualified Data.Queue as Q
import RateCalc as RC
import Process.Status
import Supervisor
import Process.Timer as Timer
import Torrent
import Protocol.Wire
-- INTERFACE
----------------------------------------------------------------------
peerChildren :: Handle -> MgrChannel -> PieceMgrChannel
-> FSPChannel -> StatusChan -> PieceMap -> Int -> IO Children
peerChildren handle pMgrC pieceMgrC fsC statusC pm nPieces = do
queueC <- channel
senderC <- channel
receiverC <- channel
sendBWC <- channel
return [Worker $ senderP handle senderC,
Worker $ sendQueueP queueC senderC sendBWC,
Worker $ receiverP handle receiverC,
Worker $ peerP pMgrC pieceMgrC fsC pm nPieces handle
queueC receiverC sendBWC statusC]
-- INTERNAL FUNCTIONS
----------------------------------------------------------------------
data SPCF = SPCF { spMsgCh :: Channel B.ByteString
}
instance Logging SPCF where
logName _ = "Process.Peer.Sender"
-- | The raw sender process, it does nothing but send out what it syncs on.
senderP :: Handle -> Channel B.ByteString -> SupervisorChan -> IO ThreadId
senderP h ch supC = spawnP (SPCF ch) h (catchP (foreverP pgm)
(do t <- liftIO $ myThreadId
syncP =<< (sendP supC $ IAmDying t)
liftIO $ hClose h))
where
pgm :: Process SPCF Handle ()
pgm = do
m <- liftIO $ timeout defaultTimeout s
h <- get
case m of
Nothing -> putMsg (encodePacket KeepAlive)
Just m -> putMsg m
liftIO $ hFlush h
defaultTimeout = 120 * 10^6
putMsg m = liftIO $ B.hPut h m
s = sync $ receive ch (const True)
-- | Messages we can send to the Send Queue
data SendQueueMessage = SendQCancel PieceNum Block -- ^ Peer requested that we cancel a piece
| SendQMsg Message -- ^ We want to send the Message to the peer
| SendOChoke -- ^ We want to choke the peer
| SendQRequestPrune PieceNum Block -- ^ Prune SendQueue of this (pn, blk) pair
data SQCF = SQCF { sqInCh :: Channel SendQueueMessage
, sqOutCh :: Channel B.ByteString
, bandwidthCh :: BandwidthChannel
}
data SQST = SQST { outQueue :: Q.Queue Message
, bytesTransferred :: Integer
}
instance Logging SQCF where
logName _ = "Process.Peer.SendQueue"
-- | sendQueue Process, simple version.
-- TODO: Split into fast and slow.
sendQueueP :: Channel SendQueueMessage -> Channel B.ByteString -> BandwidthChannel
-> SupervisorChan
-> IO ThreadId
sendQueueP inC outC bandwC supC = spawnP (SQCF inC outC bandwC) (SQST Q.empty 0)
(catchP (foreverP pgm)
(defaultStopHandler supC))
where
pgm :: Process SQCF SQST ()
pgm = do
q <- gets outQueue
l <- gets bytesTransferred
-- Gather together events which may trigger
syncP =<< (chooseP $
concat [if Q.isEmpty q then [] else [sendEvent],
if l > 0 then [rateUpdateEvent] else [],
[queueEvent]])
rateUpdateEvent = do
l <- gets bytesTransferred
ev <- sendPC bandwidthCh l
wrapP ev (\() ->
modify (\s -> s { bytesTransferred = 0 }))
queueEvent = do
recvWrapPC sqInCh
(\m -> case m of
SendQMsg msg -> do debugP "Queueing event for sending"
modifyQ (Q.push msg)
SendQCancel n blk -> modifyQ (Q.filter (filterPiece n (blockOffset blk)))
SendOChoke -> do modifyQ (Q.filter filterAllPiece)
modifyQ (Q.push Choke)
SendQRequestPrune n blk ->
modifyQ (Q.filter (filterRequest n blk)))
modifyQ :: (Q.Queue Message -> Q.Queue Message) -> Process SQCF SQST ()
modifyQ f = modify (\s -> s { outQueue = f (outQueue s) })
sendEvent = do
Just (e, r) <- gets (Q.pop . outQueue)
let bs = encodePacket e
tEvt <- sendPC sqOutCh bs
wrapP tEvt (\() -> do debugP "Dequeued event"
modify (\s -> s { outQueue = r,
bytesTransferred =
bytesTransferred s + fromIntegral (B.length bs)}))
filterAllPiece (Piece _ _ _) = True
filterAllPiece _ = False
filterPiece n off m =
case m of Piece n off _ -> False
_ -> True
filterRequest n blk m =
case m of Request n blk -> False
_ -> True
data RPCF = RPCF { rpMsgC :: Channel (Message, Integer) }
instance Logging RPCF where
logName _ = "Process.Peer.Receiver"
receiverP :: Handle -> Channel (Message, Integer)
-> SupervisorChan -> IO ThreadId
receiverP h ch supC = spawnP (RPCF ch) h
(catchP (foreverP pgm)
(defaultStopHandler supC))
where
pgm = do debugP "Peer waiting for input"
readHeader ch
readHeader ch = do
h <- get
feof <- liftIO $ hIsEOF h
if feof
then do debugP "Handle Closed"
stopP
else do bs' <- liftIO $ B.hGet h 4
l <- conv bs'
readMessage l ch
readMessage l ch = do
if (l == 0)
then return ()
else do debugP $ "Reading off " ++ show l ++ " bytes"
h <- get
bs <- liftIO $ B.hGet h (fromIntegral l)
case G.runGet decodeMsg bs of
Left _ -> do warningP "Incorrect parse in receiver, dying!"
stopP
Right msg -> do sendPC rpMsgC (msg, fromIntegral l) >>= syncP
conv bs = do
case G.runGet G.getWord32be bs of
Left err -> do warningP $ "Incorrent parse in receiver, dying: " ++ show err
stopP
Right i -> return i
data PCF = PCF { inCh :: Channel (Message, Integer)
, outCh :: Channel SendQueueMessage
, peerMgrCh :: MgrChannel
, pieceMgrCh :: PieceMgrChannel
, fsCh :: FSPChannel
, peerCh :: PeerChannel
, sendBWCh :: BandwidthChannel
, timerCh :: Channel ()
, statCh :: StatusChan
, pieceMap :: PieceMap
}
instance Logging PCF where
logName _ = "Process.Peer"
data PST = PST { weChoke :: Bool -- ^ True if we are choking the peer
, weInterested :: Bool -- ^ True if we are interested in the peer
, blockQueue :: S.Set (PieceNum, Block) -- ^ Blocks queued at the peer
, peerChoke :: Bool -- ^ Is the peer choking us? True if yes
, peerInterested :: Bool -- ^ True if the peer is interested
, peerPieces :: IS.IntSet -- ^ List of pieces the peer has access to
, upRate :: Rate -- ^ Upload rate towards the peer (estimated)
, downRate :: Rate -- ^ Download rate from the peer (estimated)
, runningEndgame :: Bool -- ^ True if we are in endgame
}
peerP :: MgrChannel -> PieceMgrChannel -> FSPChannel -> PieceMap -> Int -> Handle
-> Channel SendQueueMessage -> Channel (Message, Integer) -> BandwidthChannel
-> StatusChan
-> SupervisorChan -> IO ThreadId
peerP pMgrC pieceMgrC fsC pm nPieces h outBound inBound sendBWC statC supC = do
ch <- channel
tch <- channel
ct <- getCurrentTime
spawnP (PCF inBound outBound pMgrC pieceMgrC fsC ch sendBWC tch statC pm)
(PST True False S.empty True False IS.empty (RC.new ct) (RC.new ct) False)
(cleanupP startup (defaultStopHandler supC) cleanup)
where startup = do
tid <- liftIO $ myThreadId
debugP "Syncing a connectBack"
asks peerCh >>= (\ch -> sendPC peerMgrCh $ Connect tid ch) >>= syncP
pieces <- getPiecesDone
syncP =<< (sendPC outCh $ SendQMsg $ BitField (constructBitField nPieces pieces))
-- Install the StatusP timer
c <- asks timerCh
Timer.register 30 () c
foreverP (eventLoop)
cleanup = do
t <- liftIO myThreadId
syncP =<< sendPC peerMgrCh (Disconnect t)
getPiecesDone = do
c <- liftIO $ channel
syncP =<< (sendPC pieceMgrCh $ GetDone c)
recvP c (const True) >>= syncP
eventLoop = do
syncP =<< chooseP [peerMsgEvent, chokeMgrEvent, upRateEvent, timerEvent]
chokeMgrEvent = do
evt <- recvPC peerCh
wrapP evt (\msg -> do
debugP "ChokeMgrEvent"
case msg of
PieceCompleted pn -> do
syncP =<< (sendPC outCh $ SendQMsg $ Have pn)
ChokePeer -> do syncP =<< sendPC outCh SendOChoke
debugP "Choke Peer"
modify (\s -> s {weChoke = True})
UnchokePeer -> do syncP =<< (sendPC outCh $ SendQMsg Unchoke)
debugP "UnchokePeer"
modify (\s -> s {weChoke = False})
PeerStats t retCh -> do
i <- gets peerInterested
ur <- gets upRate
dr <- gets downRate
let (up, nur) = RC.extractRate t ur
(down, ndr) = RC.extractRate t dr
infoP $ "Peer has rates up/down: " ++ show up ++ "/" ++ show down
sendP retCh (up, down, i) >>= syncP
modify (\s -> s { upRate = nur , downRate = ndr })
CancelBlock pn blk -> do
modify (\s -> s { blockQueue = S.delete (pn, blk) $ blockQueue s })
syncP =<< (sendPC outCh $ SendQRequestPrune pn blk))
timerEvent = do
evt <- recvPC timerCh
wrapP evt (\() -> do
debugP "TimerEvent"
tch <- asks timerCh
Timer.register 30 () tch
ur <- gets upRate
dr <- gets downRate
let (upCnt, nuRate) = RC.extractCount $ ur
(downCnt, ndRate) = RC.extractCount $ dr
debugP $ "Sending peerStats: " ++ show upCnt ++ ", " ++ show downCnt
(sendPC statCh $ PeerStat { peerUploaded = upCnt
, peerDownloaded = downCnt }) >>= syncP
modify (\s -> s { upRate = nuRate, downRate = ndRate }))
upRateEvent = do
evt <- recvPC sendBWCh
wrapP evt (\uploaded -> do
modify (\s -> s { upRate = RC.update uploaded $ upRate s}))
peerMsgEvent = do
evt <- recvPC inCh
wrapP evt (\(msg, sz) -> do
modify (\s -> s { downRate = RC.update sz $ downRate s})
case msg of
KeepAlive -> return ()
Choke -> do putbackBlocks
modify (\s -> s { peerChoke = True })
Unchoke -> do modify (\s -> s { peerChoke = False })
fillBlocks
Interested -> modify (\s -> s { peerInterested = True })
NotInterested -> modify (\s -> s { peerInterested = False })
Have pn -> haveMsg pn
BitField bf -> bitfieldMsg bf
Request pn blk -> requestMsg pn blk
Piece n os bs -> do pieceMsg n os bs
fillBlocks
Cancel pn blk -> cancelMsg pn blk
Port _ -> return ()) -- No DHT yet, silently ignore
putbackBlocks = do
blks <- gets blockQueue
syncP =<< sendPC pieceMgrCh (PutbackBlocks (S.toList blks))
modify (\s -> s { blockQueue = S.empty })
haveMsg :: PieceNum -> Process PCF PST ()
haveMsg pn = do
pm <- asks pieceMap
if M.member pn pm
then do modify (\s -> s { peerPieces = IS.insert pn $ peerPieces s})
considerInterest
else do warningP "Unknown Piece"
stopP
bitfieldMsg bf = do
pieces <- gets peerPieces
if IS.null pieces
-- TODO: Don't trust the bitfield
then do modify (\s -> s { peerPieces = createPeerPieces bf})
considerInterest
else do infoP "Got out of band Bitfield request, dying"
stopP
requestMsg :: PieceNum -> Block -> Process PCF PST ()
requestMsg pn blk = do
choking <- gets weChoke
unless (choking)
(do
bs <- readBlock pn blk -- TODO: Pushdown to send process
syncP =<< sendPC outCh (SendQMsg $ Piece pn (blockOffset blk) bs))
readBlock :: PieceNum -> Block -> Process PCF PST B.ByteString
readBlock pn blk = do
c <- liftIO $ channel
syncP =<< sendPC fsCh (ReadBlock pn blk c)
syncP =<< recvP c (const True)
pieceMsg :: PieceNum -> Int -> B.ByteString -> Process PCF PST ()
pieceMsg n os bs = do
let sz = B.length bs
blk = Block os sz
e = (n, blk)
q <- gets blockQueue
when (S.member e q)
(do storeBlock n (Block os sz) bs
modify (\s -> s { blockQueue = S.delete e (blockQueue s)}))
-- When e is not a member, the piece may be stray, so ignore it.
-- Perhaps print something here.
cancelMsg n blk =
syncP =<< sendPC outCh (SendQCancel n blk)
considerInterest = do
c <- liftIO channel
pcs <- gets peerPieces
syncP =<< sendPC pieceMgrCh (AskInterested pcs c)
interested <- syncP =<< recvP c (const True)
if interested
then do modify (\s -> s { weInterested = True })
syncP =<< sendPC outCh (SendQMsg Interested)
else modify (\s -> s { weInterested = False})
fillBlocks = do
choked <- gets peerChoke
unless choked checkWatermark
checkWatermark = do
q <- gets blockQueue
eg <- gets runningEndgame
let sz = S.size q
mark = if eg then endgameLoMark else loMark
when (sz < mark)
(do
toQueue <- grabBlocks (hiMark - sz)
debugP $ "Got " ++ show (length toQueue) ++ " blocks: " ++ show toQueue
queuePieces toQueue)
queuePieces toQueue = do
mapM_ pushPiece toQueue
modify (\s -> s { blockQueue = S.union (blockQueue s) (S.fromList toQueue) })
pushPiece (pn, blk) =
syncP =<< sendPC outCh (SendQMsg $ Request pn blk)
storeBlock n blk bs =
syncP =<< sendPC pieceMgrCh (StoreBlock n blk bs)
grabBlocks n = do
c <- liftIO $ channel
ps <- gets peerPieces
syncP =<< sendPC pieceMgrCh (GrabBlocks n ps c)
blks <- syncP =<< recvP c (const True)
case blks of
Leech blks -> return blks
Endgame blks ->
modify (\s -> s { runningEndgame = True }) >> return blks
loMark = 10
endgameLoMark = 1
hiMark = 15 -- These three values are chosen rather arbitrarily at the moment.
createPeerPieces :: L.ByteString -> IS.IntSet
createPeerPieces = IS.fromList . map fromIntegral . concat . decodeBytes 0 . L.unpack
where decodeByte :: Int -> Word8 -> [Maybe Int]
decodeByte soFar w =
let dBit n = if testBit w (7-n)
then Just (n+soFar)
else Nothing
in fmap dBit [0..7]
decodeBytes _ [] = []
decodeBytes soFar (w : ws) = catMaybes (decodeByte soFar w) : decodeBytes (soFar + 8) ws
disconnectPeer :: MgrChannel -> ThreadId -> IO ()
disconnectPeer c t = sync $ transmit c $ Disconnect t