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(..),
         createMsgChannel,
         writePieceToDisk,
         Piece(..),

-        pieceMapFromFile
+        pieceMapFromFile,
+        Stats(..),
+        getStats

        )
        where
 
 import           Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
        )
        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 (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)
 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)
 data Msg = ReadPiece PieceNum Integer (MVar Piece)
          | WritePiece Piece
          | VerifyPiece PieceNum (MVar Bool)

+         | GetStats (MVar Stats)

 
 type MsgChannel = Chan Msg
 
 
 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
 
 run' :: PieceMap -> MsgChannel -> Handle -> StateT Stats IO ()
 run' pieceMap c handle = do

+  stats <- get

   msg <- liftIO recvMsg
   msg <- liftIO recvMsg

-  liftIO $ sendResponse msg
+  liftIO $ sendResponse msg stats

   updateStats msg
   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
 
 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)
 
 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