1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
|
-- |
-- Copyright : (c) Sam Truzjan 2013
-- License : BSD3
-- Maintainer : pxqr.sta@gmail.com
-- Stability : experimental
-- Portability : portable
--
-- This module implements mapping from single continious piece space
-- to file storage. Storage can be used in two modes:
--
-- * As in memory storage - in this case we don't touch filesystem.
--
-- * As ordinary mmaped file storage - when we need to store
-- data in the filesystem.
--
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE DeriveDataTypeable #-}
module System.Torrent.Storage
( -- * Storage
Storage
, StorageFailure (..)
-- * Construction
, Mode (..)
, def
, open
, close
, withStorage
-- * Query
, genPieceInfo
, getBitfield
-- * Modification
, writePiece
, readPiece
, hintRead
, unsafeReadPiece
-- * Streaming
, sourceStorage
, sinkStorage
) where
import Control.Applicative
import Control.Exception
import Control.Monad.Trans
import Data.ByteString.Lazy as BL
import Data.Conduit
import Data.Typeable
import Data.Torrent.Bitfield as BF
import Data.Torrent.Layout
import Data.Torrent.Piece
import System.Torrent.FileMap as FM
data StorageFailure
-- | Occurs on a write operation if the storage has been opened
-- using 'ReadOnly' mode.
= StorageIsRO
-- | Piece index is out of bounds.
| InvalidIndex PieceIx
-- | Piece size do not match with one passed to the 'open'
-- function.
| InvalidSize PieceSize
deriving (Show, Eq, Typeable)
instance Exception StorageFailure
-- TODO validation
data Storage = Storage
{ mode :: !Mode
, pieceLen :: {-# UNPACK #-} !PieceSize
, fileMap :: {-# UNPACK #-} !FileMap
}
-- ResourceT ?
open :: Mode -> PieceSize -> FileLayout FileSize -> IO Storage
open mode s l = Storage mode s <$> mmapFiles mode l
close :: Storage -> IO ()
close Storage {..} = unmapFiles fileMap
withStorage :: Mode -> PieceSize -> FileLayout FileSize
-> (Storage -> IO ()) -> IO ()
withStorage m s l = bracket (open m s l) close
totalPieces :: Storage -> PieceCount
totalPieces Storage {..} = FM.size fileMap `sizeInBase` pieceLen
isValidIx :: PieceIx -> Storage -> Bool
isValidIx i s = 0 <= i && i < totalPieces s
writePiece :: Piece BL.ByteString -> Storage -> IO ()
writePiece p @ Piece {..} s @ Storage {..}
| mode == ReadOnly = throwIO StorageIsRO
| isNotValidIx pieceIndex = throwIO (InvalidIndex pieceIndex)
| isNotValidSize pieceIndex (pieceSize p)
= throwIO (InvalidSize (pieceSize p))
| otherwise = writeBytes offset pieceData fileMap
where
isNotValidSize pix psize
| succ pix == pcount = psize /= lastPieceLen -- last piece may be shorter
| otherwise = psize /= pieceLen
where
lastPieceLen = fromIntegral (FM.size fileMap `rem` fromIntegral pieceLen)
{-# INLINE isNotValidSize #-}
isNotValidIx i = i < 0 || i >= pcount
{-# INLINE isNotValidIx #-}
pcount = totalPieces s
offset = fromIntegral pieceIndex * fromIntegral pieceLen
readPiece :: PieceIx -> Storage -> IO (Piece BL.ByteString)
readPiece pix s @ Storage {..}
| not (isValidIx pix s) = throwIO (InvalidIndex pix)
| otherwise = Piece pix <$> readBytes offset sz fileMap
where
offset = fromIntegral pix * fromIntegral pieceLen
sz = fromIntegral pieceLen
-- | Hint about the coming 'readPiece'. Ignores invalid indexes, for e.g.:
--
-- @forall s. hindRead (-1) s == return ()@
--
hintRead :: PieceIx -> Storage -> IO ()
hintRead _pix Storage {..} = return ()
unsafeReadPiece :: PieceIx -> Storage -> IO (Piece BL.ByteString)
unsafeReadPiece pix s @ Storage {..}
| not (isValidIx pix s) = throwIO (InvalidIndex pix)
| otherwise = return $ Piece pix (unsafeReadBytes offset sz fileMap)
where
offset = fromIntegral pix * fromIntegral pieceLen
sz = fromIntegral pieceLen
-- | Stream storage pieces from first to the last.
sourceStorage :: Storage -> Source IO (Piece BL.ByteString)
sourceStorage s = go 0
where
go pix
| pix < totalPieces s = do
piece <- liftIO $ readPiece pix s
liftIO $ hintRead (succ pix) s
yield piece
go (succ pix)
| otherwise = return ()
-- | Write stream of pieces to the storage. Fail if storage is 'ReadOnly'.
sinkStorage :: Storage -> Sink (Piece BL.ByteString) IO ()
sinkStorage s = do
awaitForever $ \ piece ->
liftIO $ writePiece piece s
-- | TODO examples of use
genPieceInfo :: Storage -> IO PieceInfo
genPieceInfo = undefined
-- | TODO examples of use
getBitfield :: Storage -> PieceInfo -> IO Bitfield
getBitfield = undefined
|
