]> git.rkrishnan.org Git - functorrent.git/commitdiff
Peer: Use Monad Transformers to simplify code
authorRamakrishnan Muthukrishnan <ram@rkrishnan.org>
Thu, 3 Sep 2015 11:02:49 +0000 (16:32 +0530)
committerRamakrishnan Muthukrishnan <ram@rkrishnan.org>
Thu, 3 Sep 2015 11:09:13 +0000 (16:39 +0530)
Wrap PState in a State Monad. Reorganization of the code in progress.

src/FuncTorrent/Peer.hs

index afc2028a73df4e56d7c0acad347d3fe6415777bf..65372d9500dfc939c7959c65345860ad6a71c5c5 100644 (file)
@@ -10,7 +10,7 @@ import System.IO (Handle, BufferMode(..), hSetBuffering)
 import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, empty)
 import qualified Data.ByteString.Char8 as BC (length)
 import Network (connectTo, PortID(..))
 import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, empty)
 import qualified Data.ByteString.Char8 as BC (length)
 import Network (connectTo, PortID(..))
-import Control.Monad (liftM, forM)
+import Control.Monad.State
 import Data.Bits
 import Data.Word (Word8)
 import Data.Map (Map, fromList, toList, (!), mapWithKey, adjust)
 import Data.Bits
 import Data.Word (Word8)
 import Data.Map (Map, fromList, toList, (!), mapWithKey, adjust)
@@ -21,13 +21,14 @@ import FuncTorrent.Utils (splitN, splitNum)
 import FuncTorrent.Fileops (createDummyFile, writeFileAtOffset)
 import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
 
 import FuncTorrent.Fileops (createDummyFile, writeFileAtOffset)
 import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
 
--- PeerState is a misnomer
-data PeerState = PeerState { handle :: Handle
-                           , peer :: Peer
-                           , meChoking :: Bool
-                           , meInterested :: Bool
-                           , heChoking :: Bool
-                           , heInterested :: Bool}
+data PState = PState { handle :: Handle
+                     , peer :: Peer
+                     , meChoking :: Bool
+                     , meInterested :: Bool
+                     , heChoking :: Bool
+                     , heInterested :: Bool}
+
+type PeerState = State PState
 
 data PieceDlState = Pending
                   | InProgress
 
 data PieceDlState = Pending
                   | InProgress
@@ -36,7 +37,7 @@ data PieceDlState = Pending
 
 -- todo - map with index to a new data structure (peers who have that piece amd state)
 data PieceData = PieceData { peers :: [Peer]        -- ^ list of peers who have this piece
 
 -- todo - map with index to a new data structure (peers who have that piece amd state)
 data PieceData = PieceData { peers :: [Peer]        -- ^ list of peers who have this piece
-                           , state :: PieceDlState  -- ^ state of the piece from download perspective.
+                           , dlstate :: PieceDlState  -- ^ state of the piece from download perspective.
                            , hash  :: ByteString    -- ^ piece hash
                            , len :: Integer }       -- ^ piece length
 
                            , hash  :: ByteString    -- ^ piece hash
                            , len :: Integer }       -- ^ piece length
 
@@ -49,7 +50,7 @@ type PieceMap = Map Integer PieceData
 mkPieceMap :: Integer -> ByteString -> [Integer] -> PieceMap
 mkPieceMap numPieces pieceHash pLengths = fromList kvs
   where kvs = [(i, PieceData { peers = []
 mkPieceMap :: Integer -> ByteString -> [Integer] -> PieceMap
 mkPieceMap numPieces pieceHash pLengths = fromList kvs
   where kvs = [(i, PieceData { peers = []
-                             , state = Pending
+                             , dlstate = Pending
                              , hash = h
                              , len = pLen })
               | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
                              , hash = h
                              , len = pLen })
               | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
@@ -57,7 +58,7 @@ mkPieceMap numPieces pieceHash pLengths = fromList kvs
 
 havePiece :: PieceMap -> Integer -> Bool
 havePiece pm index =
 
 havePiece :: PieceMap -> Integer -> Bool
 havePiece pm index =
-  state (pm ! index) == Have
+  dlstate (pm ! index) == Have
 
 connectToPeer :: Peer -> IO Handle
 connectToPeer (Peer _ ip port) = do
 
 connectToPeer :: Peer -> IO Handle
 connectToPeer (Peer _ ip port) = do
@@ -82,60 +83,69 @@ bitfieldToList bs = go bs 0
           in
            setBits ++ go bs' (pos + 1)
 
           in
            setBits ++ go bs' (pos + 1)
 
--- recvMsg :: Peer -> Handle -> Msg
-msgLoop :: PeerState -> PieceMap -> FilePath -> IO ()
-msgLoop pState@(PeerState { meInterested = False, heChoking = True }) pieceStatus file =
-  do
-    -- if me NOT Interested and she is Choking, tell her that
-    -- I am interested.
-    let h = handle pState
-    sendMsg h InterestedMsg
-    putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState)
-    msgLoop (pState { meInterested = True }) pieceStatus file
-msgLoop pState@(PeerState { meInterested = True, heChoking = False }) pieceStatus file =
-  -- if me Interested and she not Choking, send her a request
-  -- for a piece.
-  case pickPiece pieceStatus of
-   Nothing -> putStrLn "Nothing to download"
-   Just workPiece -> do
-     let pLen = len (pieceStatus ! workPiece)
-     putStrLn $ "piece length = " ++ show pLen
-     pBS <- downloadPiece (handle pState) workPiece pLen
-     if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
-       then
-       putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
-       else do
-       let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
-       putStrLn $ "Write into file at offset: " ++ show fileOffset
-       writeFileAtOffset file fileOffset pBS
-       msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus) file
-msgLoop pState pieceStatus file = do
-  msg <- getMsg (handle pState)
-  putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
-  case msg of
-   KeepAliveMsg -> do
-     sendMsg (handle pState) KeepAliveMsg
-     putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
-     msgLoop pState pieceStatus file
-   BitFieldMsg bss -> do
-     let pieceList = bitfieldToList (unpack bss)
-         pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList
-     putStrLn $ show (length pieceList) ++ " Pieces"
-     -- for each pieceIndex in pieceList, make an entry in the pieceStatus
-     -- map with pieceIndex as the key and modify the value to add the peer.
-     -- download each of the piece in order
-     msgLoop pState pieceStatus' file
-   UnChokeMsg ->
-     msgLoop (pState { heChoking = False }) pieceStatus file
-   _ ->
-     msgLoop pState pieceStatus file
+-- helper functions to manipulate PeerState
+toPeerState :: Handle -> Peer -> Bool -> Bool -> Bool -> Bool -> PState
+toPeerState h p meCh meIn heCh heIn =
+  PState { handle = h
+         , peer = p
+         , heChoking = heCh
+         , heInterested = heIn
+         , meChoking = meCh
+         , meInterested = meIn }
+
+-- -- recvMsg :: Peer -> Handle -> Msg
+-- msgLoop :: PeerState -> PieceMap -> FilePath -> IO ()
+-- msgLoop pState@(PeerState { meInterested = False, heChoking = True }) pieceStatus file = do
+--   -- if me NOT Interested and she is Choking, tell her that
+--   -- I am interested.
+--   let h = handle pState
+--   sendMsg h InterestedMsg
+--   putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState)
+--   msgLoop (pState { meInterested = True }) pieceStatus file
+-- msgLoop pState@(PeerState { meInterested = True, heChoking = False }) pieceStatus file =
+--   -- if me Interested and she not Choking, send her a request
+--   -- for a piece.
+--   case pickPiece pieceStatus of
+--    Nothing -> putStrLn "Nothing to download"
+--    Just workPiece -> do
+--      let pLen = len (pieceStatus ! workPiece)
+--      putStrLn $ "piece length = " ++ show pLen
+--      pBS <- downloadPiece (handle pState) workPiece pLen
+--      if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
+--        then
+--        putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
+--        else do
+--        let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
+--        putStrLn $ "Write into file at offset: " ++ show fileOffset
+--        writeFileAtOffset file fileOffset pBS
+--        msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus) file
+-- msgLoop pState pieceStatus file = do
+--   msg <- getMsg (handle pState)
+--   putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
+--   case msg of
+--    KeepAliveMsg -> do
+--      sendMsg (handle pState) KeepAliveMsg
+--      putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
+--      msgLoop pState pieceStatus file
+--    BitFieldMsg bss -> do
+--      let pieceList = bitfieldToList (unpack bss)
+--          pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList
+--      putStrLn $ show (length pieceList) ++ " Pieces"
+--      -- for each pieceIndex in pieceList, make an entry in the pieceStatus
+--      -- map with pieceIndex as the key and modify the value to add the peer.
+--      -- download each of the piece in order
+--      msgLoop pState pieceStatus' file
+--    UnChokeMsg ->
+--      msgLoop (pState { heChoking = False }) pieceStatus file
+--    _ ->
+--      msgLoop pState pieceStatus file
 
 -- simple algorithm to pick piece.
 -- pick the first piece from 0 that is not downloaded yet.
 pickPiece :: PieceMap -> Maybe Integer
 pickPiece m =
   let pieceList = toList m
 
 -- simple algorithm to pick piece.
 -- pick the first piece from 0 that is not downloaded yet.
 pickPiece :: PieceMap -> Maybe Integer
 pickPiece m =
   let pieceList = toList m
-      allPending = filter (\(_, v) -> state v == Pending) pieceList
+      allPending = filter (\(_, v) -> dlstate v == Pending) pieceList
   in
    case allPending of
     [] -> Nothing
   in
    case allPending of
     [] -> Nothing
@@ -151,12 +161,7 @@ handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
 handlePeerMsgs p m peerId = do
   h <- connectToPeer p
   doHandshake h p (infoHash m) peerId
 handlePeerMsgs p m peerId = do
   h <- connectToPeer p
   doHandshake h p (infoHash m) peerId
-  let state = PeerState { handle = h
-                        , peer = p
-                        , heInterested = False
-                        , heChoking = True
-                        , meInterested = False
-                        , meChoking = True }
+  let pstate = toPeerState h p False True False True
       pieceHash = pieces (info m)
       numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
       pLen = pieceLength (info m)
       pieceHash = pieces (info m)
       numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
       pLen = pieceLength (info m)
@@ -164,8 +169,21 @@ handlePeerMsgs p m peerId = do
       fileName = name (info m)
       pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
   createDummyFile fileName (fromIntegral fileLen)
       fileName = name (info m)
       pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
   createDummyFile fileName (fromIntegral fileLen)
-  msgLoop state pieceStatus fileName
-  
+  (r, _) <- runStateT (msgLoop pieceStatus fileName) pstate
+  return ()
+
+msgLoop :: PieceMap -> FilePath -> StateT PState IO ()
+msgLoop pieceStatus file =
+  StateT(\pState -> do
+             let h = handle pState
+             msg <- getMsg h
+             liftIO $ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
+             case msg of
+               KeepAliveMsg -> do
+                 sendMsg h KeepAliveMsg
+                 liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
+                 runStateT (msgLoop pieceStatus file) pState)
+
 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
 downloadPiece h index pieceLength = do
   let chunks = splitNum pieceLength 16384
 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
 downloadPiece h index pieceLength = do
   let chunks = splitNum pieceLength 16384