diff --git a/Utility/Metered.hs b/Utility/Metered.hs index 4618aecfe4..447eab2e89 100644 --- a/Utility/Metered.hs +++ b/Utility/Metered.hs @@ -99,15 +99,23 @@ 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. - - - 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. -} hGetContentsMetered :: Handle -> MeterUpdate -> IO L.ByteString -hGetContentsMetered h meterupdate = lazyRead zeroBytesProcessed +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. + - + - Stops at EOF, or when keepgoing evaluates to False. + - Closes the Handle at EOF, but otherwise leaves it open. + -} +hGetUntilMetered :: Handle -> (Integer -> Bool) -> MeterUpdate -> IO L.ByteString +hGetUntilMetered h keepgoing meterupdate = lazyRead zeroBytesProcessed where lazyRead sofar = unsafeInterleaveIO $ loop sofar @@ -118,14 +126,16 @@ hGetContentsMetered h meterupdate = lazyRead zeroBytesProcessed 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