]> git.rkrishnan.org Git - functorrent.git/blobdiff - src/FuncTorrent/FileSystem.hs
more hlint fixes
[functorrent.git] / src / FuncTorrent / FileSystem.hs
index 1b190ca88936b197ab3563e1fc8a1e8238629575..7c95c1e2ffcffe330681c626bbf21c95e9c63d8d 100644 (file)
@@ -1,25 +1,48 @@
+{-
+ - Copyright (C) 2015-2016 Ramakrishnan Muthukrishnan <ram@rkrishnan.org>
+ -
+ - This file is part of FuncTorrent.
+ -
+ - FuncTorrent is free software; you can redistribute it and/or modify
+ - it under the terms of the GNU General Public License as published by
+ - the Free Software Foundation; either version 3 of the License, or
+ - (at your option) any later version.
+ -
+ - FuncTorrent is distributed in the hope that it will be useful,
+ - but WITHOUT ANY WARRANTY; without even the implied warranty of
+ - MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ - GNU General Public License for more details.
+ -
+ - You should have received a copy of the GNU General Public License
+ - along with FuncTorrent; if not,  see <http://www.gnu.org/licenses/>
+ -}
+
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE FlexibleContexts #-}
+
 module FuncTorrent.FileSystem
-       (startThread,
+       (run,
         MsgChannel,
-        initFS,
-        Msg(..),
+        createMsgChannel,
+        writePieceToDisk,
         Piece(..),
-        pieceMapFromFile
+        pieceMapFromFile,
+        Stats(..),
+        getStats
        )
        where
 
-import           Control.Concurrent (ThreadId, forkIO)
-import           Control.Concurrent.Chan (Chan, newChan, readChan)
-import           Control.Concurrent.MVar (MVar, putMVar)
+import           Control.Concurrent.Chan (Chan, newChan, readChan, writeChan)
+import           Control.Concurrent.MVar (MVar, newEmptyMVar, putMVar)
 import           Control.Monad (forever)
-import           Data.Map (traverseWithKey)
-
+import           Control.Monad.State (StateT, liftIO, get, runStateT, modify)
 import qualified Data.ByteString as BS
-import           Data.Map ((!))
-import           System.IO (Handle, openFile, IOMode (ReadWriteMode))
+import           Data.Map (traverseWithKey, (!))
+import           System.IO (Handle, IOMode (ReadWriteMode), withFile)
+import           System.Directory (doesFileExist)
 
 import           FuncTorrent.PieceManager (PieceDlState(..), PieceData(..), PieceMap, pieceNumToOffset)
-import           FuncTorrent.Utils (readFileAtOffset, writeFileAtOffset, verifyHash)
+import           FuncTorrent.Utils (createDummyFile, readFileAtOffset, writeFileAtOffset, verifyHash)
 
 type PieceNum = Integer
 data Piece = Piece PieceNum BS.ByteString
@@ -27,51 +50,82 @@ 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
 
--- init :: FileName -> IO (Handle, MsgChannel)
-initFS :: FilePath -> IO (Handle, MsgChannel)
-initFS filepath = do
-  c <- newChan
-  h <- openFile filepath ReadWriteMode
-  return (h, c)
+data Stats = Stats { bytesRead :: Integer
+                   , bytesWritten :: Integer
+                   }
 
-startThread :: Handle -> MsgChannel -> PieceMap -> IO ThreadId
-startThread handle c pieceMap = do
-  forkIO $ forever $ recvMsg >>= sendResponse
-  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) -> do
-          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'
+createMsgChannel :: IO (Chan Msg)
+createMsgChannel = newChan
 
-pieceMapFromFile :: Handle -> PieceMap -> IO PieceMap
-pieceMapFromFile handle pieceMap = do
-  traverseWithKey f pieceMap
+run :: PieceMap -> MsgChannel -> Handle -> IO ()
+run pieceMap c handle = forever $ do
+  _ <- runStateT (run' pieceMap c handle) initialStats
+  return ()
+    where initialStats = Stats { bytesRead = 0
+                               , bytesWritten = 0 }
+
+run' :: PieceMap -> MsgChannel -> Handle -> StateT Stats IO ()
+run' pieceMap c handle = do
+  stats <- get
+  msg <- liftIO recvMsg
+  liftIO $ sendResponse msg stats
+  updateStats msg
+    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 ->
+            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
+  dfe <- doesFileExist filePath
+  if dfe
+    then traverseWithKey f pieceMap
+    else createDummyFile filePath (fromIntegral fileLen) >> return pieceMap
   where
     f k v = do
       let offset = pieceNumToOffset pieceMap k
-      isHashValid <- flip verifyHash (hash v) <$> readFileAtOffset handle offset (len v)
+      isHashValid <- flip verifyHash (hash v) <$> withFile filePath ReadWriteMode (\handle -> readFileAtOffset handle offset (len v))
       if isHashValid
         then return $ v { dlstate = Have }
         else return v
+
+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