tracker: refactor around Http and Udp (to be worked on) modules

[functorrent.git] / src / FuncTorrent / FileSystem.hs

diff --git a/src/FuncTorrent/FileSystem.hs b/src/FuncTorrent/FileSystem.hs
index fdf89fa94bae2415d9fa12d572cade889823eebf..3c625e69e0825af4f660ee783498fcaba1d8f782 100644 (file)
--- a/src/FuncTorrent/FileSystem.hs
+++ b/src/FuncTorrent/FileSystem.hs
@@ -6,14 +6,16 @@ module FuncTorrent.FileSystem
         createMsgChannel,
         writePieceToDisk,
         Piece(..),
-        pieceMapFromFile
+        pieceMapFromFile,
+        Stats(..),
+        getStats
        )
        where
 
 import           Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
-import           Control.Concurrent.MVar (MVar, putMVar)
+import           Control.Concurrent.MVar (MVar, newEmptyMVar, putMVar)
 import           Control.Monad (forever)
-import           Control.Monad.State (StateT, liftIO, runStateT, modify)
+import           Control.Monad.State (StateT, liftIO, get, runStateT, modify)
 import qualified Data.ByteString as BS
 import           Data.Map (traverseWithKey, (!))
 import           System.IO (Handle, IOMode (ReadWriteMode), withFile)
@@ -28,6 +30,7 @@ data Piece = Piece PieceNum BS.ByteString
 data Msg = ReadPiece PieceNum Integer (MVar Piece)
          | WritePiece Piece
          | VerifyPiece PieceNum (MVar Bool)
+         | GetStats (MVar Stats)
 
 type MsgChannel = Chan Msg
 
@@ -47,38 +50,41 @@ run pieceMap c handle = forever $ do
 
 run' :: PieceMap -> MsgChannel -> Handle -> StateT Stats IO ()
 run' pieceMap c handle = do
+  stats <- get
   msg <- liftIO recvMsg
-  liftIO $ sendResponse msg
+  liftIO $ sendResponse msg stats
   updateStats msg
-  where
-    recvMsg = readChan c
-    sendResponse msg =
-      case msg of
-      ReadPiece n len' var -> do
-        bs <- readPiece n len'
-        putMVar var (Piece n bs)
-      WritePiece (Piece n bs) ->
-        writePiece n bs
-      VerifyPiece n var -> do
-        isHashValid <- verifyPiece n
-        putMVar var isHashValid
-    readPiece n len' = do
-      let offset = pieceNumToOffset pieceMap n
-      readFileAtOffset handle offset len'
-    writePiece n piece = do
-      let offset = pieceNumToOffset pieceMap n
-      writeFileAtOffset handle offset piece
-    verifyPiece n = do
-      let offset = pieceNumToOffset pieceMap n
-          hash'  = hash (pieceMap ! n)
-          len'   = len (pieceMap ! n)
-      bs' <- readFileAtOffset handle offset len'
-      return $ verifyHash bs' hash'
-    updateStats (ReadPiece _ l _) =
-      modify (\st -> st {bytesRead = bytesRead st + l})
-    updateStats (WritePiece (Piece _ bs)) =
-      modify (\st -> st {bytesWritten = bytesWritten st + fromIntegral (BS.length bs)})
-    updateStats _ = modify id
+    where
+      recvMsg = readChan c
+      sendResponse msg stats =
+        case msg of
+          ReadPiece n len' var -> do
+            bs <- readPiece n len'
+            putMVar var (Piece n bs)
+          WritePiece (Piece n bs) ->
+            writePiece n bs
+          VerifyPiece n var -> do
+            isHashValid <- verifyPiece n
+            putMVar var isHashValid
+          GetStats var -> do
+            putMVar var stats
+      readPiece n len' = do
+        let offset = pieceNumToOffset pieceMap n
+        readFileAtOffset handle offset len'
+      writePiece n piece = do
+        let offset = pieceNumToOffset pieceMap n
+        writeFileAtOffset handle offset piece
+      verifyPiece n = do
+        let offset = pieceNumToOffset pieceMap n
+            hash'  = hash (pieceMap ! n)
+            len'   = len (pieceMap ! n)
+        bs' <- readFileAtOffset handle offset len'
+        return $ verifyHash bs' hash'
+      updateStats (ReadPiece _ l _) =
+        modify (\st -> st {bytesRead = bytesRead st + l})
+      updateStats (WritePiece (Piece _ bs)) =
+        modify (\st -> st {bytesWritten = bytesWritten st + fromIntegral (BS.length bs)})
+      updateStats _ = modify id
 
 pieceMapFromFile :: FilePath -> Integer -> PieceMap -> IO PieceMap
 pieceMapFromFile filePath fileLen pieceMap = do
@@ -98,3 +104,8 @@ writePieceToDisk :: MsgChannel -> PieceNum -> BS.ByteString -> IO ()
 writePieceToDisk c pieceNum bs =
   writeChan c $ WritePiece (Piece pieceNum bs)
 
+getStats :: MsgChannel -> IO (MVar Stats)
+getStats c = do
+  v <- newEmptyMVar
+  writeChan c $ GetStats v
+  return v