diff options
Diffstat (limited to 'Utility/Metered.hs')
| -rw-r--r-- | Utility/Metered.hs | 34 |
1 files changed, 22 insertions, 12 deletions
diff --git a/Utility/Metered.hs b/Utility/Metered.hs index 4618aec..f27eee2 100644 --- a/Utility/Metered.hs +++ b/Utility/Metered.hs @@ -99,33 +99,43 @@ offsetMeterUpdate base offset = \n -> base (offset `addBytesProcessed` n) {- This is like L.hGetContents, but after each chunk is read, a meter - is updated based on the size of the chunk. - + - All the usual caveats about using unsafeInterleaveIO apply to the + - meter updates, so use caution. + -} +hGetContentsMetered :: Handle -> MeterUpdate -> IO L.ByteString +hGetContentsMetered h = hGetUntilMetered h (const True) + +{- Reads from the Handle, updating the meter after each chunk. + - - Note that the meter update is run in unsafeInterleaveIO, which means that - it can be run at any time. It's even possible for updates to run out - of order, as different parts of the ByteString are consumed. - - - All the usual caveats about using unsafeInterleaveIO apply to the - - meter updates, so use caution. + - Stops at EOF, or when keepgoing evaluates to False. + - Closes the Handle at EOF, but otherwise leaves it open. -} -hGetContentsMetered :: Handle -> MeterUpdate -> IO L.ByteString -hGetContentsMetered h meterupdate = lazyRead zeroBytesProcessed +hGetUntilMetered :: Handle -> (Integer -> Bool) -> MeterUpdate -> IO L.ByteString +hGetUntilMetered h keepgoing meterupdate = lazyRead zeroBytesProcessed where lazyRead sofar = unsafeInterleaveIO $ loop sofar loop sofar = do - c <- S.hGetSome h defaultChunkSize + c <- S.hGet h defaultChunkSize if S.null c then do hClose h return $ L.empty else do - let sofar' = addBytesProcessed sofar $ - S.length c + let sofar' = addBytesProcessed sofar (S.length c) meterupdate sofar' - {- unsafeInterleaveIO causes this to be - - deferred until the data is read from the - - ByteString. -} - cs <- lazyRead sofar' - return $ L.append (L.fromChunks [c]) cs + if keepgoing (fromBytesProcessed sofar') + then do + {- unsafeInterleaveIO causes this to be + - deferred until the data is read from the + - ByteString. -} + cs <- lazyRead sofar' + return $ L.append (L.fromChunks [c]) cs + else return $ L.fromChunks [c] {- Same default chunk size Lazy ByteStrings use. -} defaultChunkSize :: Int |